Skip to main content

cortiq_engine/
qtensor.rs

1//! QTensor — weight tensor with pluggable storage.
2//!
3//! Two backings, one interface:
4//! - `F32`   — owned dense floats (small models, tests). Every operation
5//!   is bit-identical to the historical `&[f32]` code paths.
6//! - `Mapped` — quantized bytes zero-copy from the CMF mmap (`q8_row` /
7//!   `q8_2f`). The matvec is fused: int8 rows × f32 activations, the
8//!   q8_2f column field folds into a pre-scale of the input
9//!   (`x'[i] = col[i]·x[i]`), so the inner loop is the same i8 dot as
10//!   q8_row. This is what lets a 15B file run in a few GB of RSS.
11//!
12//! Extension point: new dtypes = new match arm here, nothing else moves.
13
14use crate::pool::{Pool, matvec_rows, matvec_rows2};
15use cortiq_core::quant::{
16    GROUP_SIZE, Q1_TILE, Q2TP_CHUNK, Q4_TILE, Q4TP_NIB, f16_to_f32, q2tp_ladder, q2tp_sections,
17    q4tp_code, q4tp_ladder, q4tp_sections,
18};
19use cortiq_core::{CmfModel, TensorDtype};
20use std::sync::Arc;
21
22pub enum QTensor {
23    F32 {
24        data: Vec<f32>,
25        rows: usize,
26        cols: usize,
27    },
28    Mapped {
29        model: Arc<CmfModel>,
30        /// Index into the model's tensor directory.
31        idx: usize,
32        dtype: TensorDtype,
33        rows: usize,
34        cols: usize,
35        /// Per-row scales, dequantized to f32 up front (tiny).
36        row_scale: Vec<f32>,
37        /// q8_2f column field (θ), dequantized up front; empty for q8_row.
38        col_field: Vec<f32>,
39        /// Vbit only: byte offset of each row's packed data within the
40        /// tensor blob (`[rows + 1]`, computed once at load — the per-
41        /// matvec prefix scan over row bit-widths was O(rows) each call).
42        vbit_offsets: Vec<usize>,
43        /// q8-family decode repack (load-time, optional): rows in groups
44        /// of 4, interleaved in 16-byte units — one 64-byte line per
45        /// iteration feeds all 4 sdot lanes, ONE sequential weight
46        /// stream per worker instead of four (this is where llama.cpp's
47        /// repacked Q8 kernels get their bandwidth). Empty = off
48        /// (CMF_REPACK=0, non-SDOT arch, or an ineligible shape). Trades
49        /// an anonymous copy of the quants for mmap pages that go cold.
50        repack: Vec<u8>,
51    },
52}
53
54/// Load-time q8 repack gate (see `Mapped::repack`). OPT-IN
55/// (`CMF_REPACK=1`): the single-stream hypothesis LOST on Apple Silicon
56/// (M4, interleaved A/B: decode 101 vs 94 tok/s — four adjacent row
57/// streams per worker feed the prefetcher MORE memory-level parallelism
58/// than one); kept as an experiment flag for x86, where the tradeoff
59/// may land differently.
60fn repack_enabled() -> bool {
61    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
62    *ON.get_or_init(|| {
63        std::env::var("CMF_REPACK")
64            .map(|v| v == "1")
65            .unwrap_or(cfg!(target_os = "android"))
66    })
67}
68
69/// Interleave q8 rows for the decode kernel: group g holds rows
70/// 4g..4g+4 as [r0[c], r1[c], r2[c], r3[c]] per 16-byte chunk c. Only
71/// full groups are packed — tail rows keep reading the mmap layout.
72fn q8_repack(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
73    #[cfg(target_arch = "aarch64")]
74    let arch_ok = sdot_enabled();
75    #[cfg(not(target_arch = "aarch64"))]
76    let arch_ok = false;
77    if !arch_ok || !repack_enabled() || rows < 256 || cols % 16 != 0 {
78        return Vec::new();
79    }
80    q8_repack_layout(bytes, rows, cols)
81}
82
83/// The pure layout transform behind `q8_repack` (tested directly —
84/// the gate depends on arch and env).
85fn q8_repack_layout(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
86    let groups = rows / 4;
87    let mut rep = vec![0u8; groups * 4 * cols];
88    for g in 0..groups {
89        let dst = &mut rep[g * 4 * cols..(g + 1) * 4 * cols];
90        for c in 0..cols / 16 {
91            for lane in 0..4 {
92                let src = (g * 4 + lane) * cols + c * 16;
93                dst[c * 64 + lane * 16..c * 64 + lane * 16 + 16]
94                    .copy_from_slice(&bytes[src..src + 16]);
95            }
96        }
97    }
98    rep
99}
100
101/// Prefix-sum of vbit row payload offsets (absolute within the tensor
102/// bytes). `offsets[r]..offsets[r+1]` is row r's packed data.
103fn vbit_row_offsets(bytes: &[u8], rows: usize, cols: usize) -> Vec<usize> {
104    let ng = cols / GROUP_SIZE;
105    let bits = &bytes[..rows];
106    let mut offsets = Vec::with_capacity(rows + 1);
107    let mut off = rows + rows * ng * 2;
108    for r in 0..rows {
109        offsets.push(off);
110        off += (cols * bits[r] as usize).div_ceil(8);
111    }
112    offsets.push(off);
113    offsets
114}
115
116/// `CMF_X86_BLOCKED` / `CMF_GPU_LMHEAD` / `CMF_GPU_SPLIT`, read once. They
117/// used to be read from the environment on every large matvec and on every
118/// matmat in six places — microseconds each, but also a knob that could
119/// change under a running process, which is not a thing a kernel choice
120/// should be able to do mid-sequence.
121fn blocked_enabled() -> bool {
122    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
123    *ON.get_or_init(|| std::env::var("CMF_X86_BLOCKED").map(|v| v != "0").unwrap_or(true))
124}
125
126fn gpu_lmhead_enabled() -> bool {
127    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
128    *ON.get_or_init(|| std::env::var("CMF_GPU_LMHEAD").map(|v| v != "0").unwrap_or(true))
129}
130
131fn gpu_split_frac() -> f32 {
132    static F: std::sync::OnceLock<f32> = std::sync::OnceLock::new();
133    *F.get_or_init(|| {
134        std::env::var("CMF_GPU_SPLIT")
135            .ok()
136            .and_then(|v| v.parse::<f32>().ok())
137            .unwrap_or(0.5)
138            .clamp(0.0, 1.0)
139    })
140}
141
142impl QTensor {
143    pub fn from_f32(data: Vec<f32>, rows: usize, cols: usize) -> Self {
144        debug_assert_eq!(data.len(), rows * cols);
145        Self::F32 { data, rows, cols }
146    }
147
148    /// Wrap a directory tensor without dequantizing the payload.
149    /// Falls back to dequantized f32 for dtypes without a fused kernel.
150    pub fn from_model(model: &Arc<CmfModel>, name: &str) -> Result<Self, String> {
151        // Indexed lookup: the linear directory scan made pipeline build
152        // O(N²) on MoE/skills files with thousands of tensors.
153        let idx = model
154            .tensor_index(name)
155            .ok_or_else(|| format!("tensor '{name}' not found in CMF directory"))?;
156        let entry = &model.tensors[idx];
157        if entry.shape.len() != 2 {
158            return Err(format!("QTensor::from_model needs 2-D, got '{name}'"));
159        }
160        let (rows, cols) = (entry.shape[0], entry.shape[1]);
161        let bytes = model.entry_bytes(entry);
162
163        match entry.dtype {
164            TensorDtype::Q8Row | TensorDtype::Q8_2f => {
165                let n = rows * cols;
166                let scales_off = n;
167                let row_scale: Vec<f32> = (0..rows)
168                    .map(|o| {
169                        f16_to_f32(u16::from_le_bytes([
170                            bytes[scales_off + o * 2],
171                            bytes[scales_off + o * 2 + 1],
172                        ]))
173                    })
174                    .collect();
175                let col_field: Vec<f32> = if entry.dtype == TensorDtype::Q8_2f {
176                    let col_off = n + rows * 2;
177                    (0..cols)
178                        .map(|i| {
179                            f16_to_f32(u16::from_le_bytes([
180                                bytes[col_off + i * 2],
181                                bytes[col_off + i * 2 + 1],
182                            ]))
183                        })
184                        .collect()
185                } else {
186                    Vec::new()
187                };
188                Ok(Self::Mapped {
189                    model: model.clone(),
190                    idx,
191                    dtype: entry.dtype,
192                    rows,
193                    cols,
194                    row_scale,
195                    col_field,
196                    vbit_offsets: Vec::new(),
197                    repack: q8_repack(bytes, rows, cols),
198                })
199            }
200            // vbit: fused kernel unpacks variable-bit rows from mmap.
201            TensorDtype::Vbit if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
202                model: model.clone(),
203                idx,
204                dtype: entry.dtype,
205                rows,
206                cols,
207                row_scale: Vec::new(),
208                col_field: Vec::new(),
209                vbit_offsets: vbit_row_offsets(bytes, rows, cols),
210                repack: Vec::new(),
211            }),
212            // vbit_ro (§4.2): the offset table comes straight from the
213            // file — no load-time prefix scan; kernels are shared with
214            // legacy vbit (they consume absolute offsets either way).
215            TensorDtype::VbitRo if cols % GROUP_SIZE == 0 => {
216                let (_, off_off, packed_off) = cortiq_core::quant::vbit_ro_sections(rows, cols);
217                let offsets: Vec<usize> = (0..=rows)
218                    .map(|r| packed_off + cortiq_core::quant::vbit_ro_offset(bytes, off_off, r))
219                    .collect();
220                Ok(Self::Mapped {
221                    model: model.clone(),
222                    idx,
223                    dtype: entry.dtype,
224                    rows,
225                    cols,
226                    row_scale: Vec::new(),
227                    col_field: Vec::new(),
228                    vbit_offsets: offsets,
229                    repack: Vec::new(),
230                })
231            }
232            // q4_block: fused kernel reads nibbles straight from mmap —
233            // a 14B q4 file no longer explodes into ×8 f32 RAM.
234            // q4_tiled (§4.3): interleaved [scale][nibbles] tiles — one
235            // sequential memory stream (measured ×1.66 ARM / ×1.13 AVX2
236            // at kernel level over the split layout).
237            TensorDtype::Q4Tiled if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
238                model: model.clone(),
239                idx,
240                dtype: entry.dtype,
241                rows,
242                cols,
243                row_scale: Vec::new(),
244                col_field: Vec::new(),
245                vbit_offsets: Vec::new(),
246                repack: Vec::new(),
247            }),
248            // q4tp (§4.10): nibbles from mmap, scale from the row ladder —
249            // 7.3% less file than q4t at the same 4-bit grid.
250            TensorDtype::Q4TiledP if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
251                model: model.clone(),
252                idx,
253                dtype: entry.dtype,
254                rows,
255                cols,
256                row_scale: Vec::new(),
257                col_field: Vec::new(),
258                vbit_offsets: Vec::new(),
259                repack: Vec::new(),
260            }),
261            // q2tp: 2-bit chunks from mmap, scale from the same row ladder.
262            TensorDtype::Q2TiledP if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
263                model: model.clone(),
264                idx,
265                dtype: entry.dtype,
266                rows,
267                cols,
268                row_scale: Vec::new(),
269                col_field: Vec::new(),
270                vbit_offsets: Vec::new(),
271                repack: Vec::new(),
272            }),
273            TensorDtype::Q4Block if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
274                model: model.clone(),
275                idx,
276                dtype: entry.dtype,
277                rows,
278                cols,
279                row_scale: Vec::new(),
280                col_field: Vec::new(),
281                vbit_offsets: Vec::new(),
282                repack: Vec::new(),
283            }),
284            // q1: binary sign-bit tiles from mmap (1-bit-trained models).
285            TensorDtype::Q1 if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
286                model: model.clone(),
287                idx,
288                dtype: entry.dtype,
289                rows,
290                cols,
291                row_scale: Vec::new(),
292                col_field: Vec::new(),
293                vbit_offsets: Vec::new(),
294                repack: Vec::new(),
295            }),
296            // q1t (ternary + outlier overlay): fused per-row dequant kernel
297            // reads straight from mmap — a 12B q1t stays ~its file size in
298            // RAM instead of dequantizing to ~48 GB of f32.
299            TensorDtype::Q1T if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
300                model: model.clone(),
301                idx,
302                dtype: entry.dtype,
303                rows,
304                cols,
305                row_scale: Vec::new(),
306                col_field: Vec::new(),
307                vbit_offsets: Vec::new(),
308                repack: Vec::new(),
309            }),
310            // No fused kernel yet → dequantize once (correct, more RAM).
311            _ => {
312                let mut data = vec![0.0f32; rows * cols];
313                cortiq_core::quant::dequant_tensor(entry, bytes, &mut data)?;
314                Ok(Self::from_f32(data, rows, cols))
315            }
316        }
317    }
318
319    /// q1-mapped tensor? (GPU gates: the q1 CPU kernel is
320    /// compute-bound, so offload pays at much smaller shapes than q8.)
321    pub(crate) fn is_q1(&self) -> bool {
322        matches!(
323            self,
324            Self::Mapped {
325                dtype: TensorDtype::Q1,
326                ..
327            }
328        )
329    }
330
331    /// Owned-f32 view (data, rows, cols) — the GDN a/b gate projections
332    /// arrive dequantized (force-f16 in the converter → F32 in RAM).
333    pub(crate) fn f32_parts(&self) -> Option<(&[f32], usize, usize)> {
334        match self {
335            Self::F32 { data, rows, cols } => Some((data, *rows, *cols)),
336            _ => None,
337        }
338    }
339
340    /// (directory idx, rows, cols) of a q1-mapped tensor — the
341    /// whole-block GPU path resolves offsets itself.
342    /// (idx, rows, cols) of a mapped tensor the whole-token GPU graph can drive
343    /// — Q1, Q1T or Q4-block (it resolves the offset and picks the kernel by
344    /// dtype). Q4-block lets a precise down_proj/lm_head stay on-device.
345    /// Named `q1_parts` for historical reasons.
346    pub(crate) fn q1_parts(&self) -> Option<(usize, usize, usize)> {
347        match self {
348            #[cfg(target_os = "macos")]
349            Self::Mapped {
350                dtype: TensorDtype::Q1T,
351                ..
352            } if !crate::gpu::metal_q1t_enabled() => None,
353            Self::Mapped {
354                idx,
355                dtype:
356                    TensorDtype::Q1
357                    | TensorDtype::Q1T
358                    | TensorDtype::Q4Block
359                    | TensorDtype::Q4Tiled
360                    | TensorDtype::Q4TiledP
361                    | TensorDtype::Q2TiledP
362                    | TensorDtype::Q8Row
363                    | TensorDtype::Q8_2f,
364                rows,
365                cols,
366                ..
367            } => Some((*idx, *rows, *cols)),
368            _ => None,
369        }
370    }
371
372    /// (directory idx, rows, cols) of a q4_tiled mapped tensor. The
373    /// chunk-prefill graph takes it in the same 4-tuple slot as
374    /// `q8_row_parts` with an EMPTY row_scale — q4t carries its scales
375    /// inside the 18-byte tiles, and the empty slice is what tells the
376    /// encoder to reach for the q4t kernels.
377    pub(crate) fn q4t_parts(&self) -> Option<(usize, usize, usize)> {
378        match self {
379            Self::Mapped {
380                idx,
381                dtype: TensorDtype::Q4Tiled,
382                rows,
383                cols,
384                ..
385            } => Some((*idx, *rows, *cols)),
386            _ => None,
387        }
388    }
389
390    /// (directory idx, rows, cols) of a q4tp mapped tensor. Same empty-scale
391    /// slot as `q4t_parts` in the chunk graph — the encoder tells the two
392    /// apart by the tensor's dtype, not by the slot.
393    pub(crate) fn q4tp_parts(&self) -> Option<(usize, usize, usize)> {
394        match self {
395            Self::Mapped {
396                idx,
397                dtype: TensorDtype::Q4TiledP,
398                rows,
399                cols,
400                ..
401            } => Some((*idx, *rows, *cols)),
402            _ => None,
403        }
404    }
405
406    /// (directory idx, rows, cols, row_scale) of a plain q8_row mapped
407    /// tensor — the chunk-prefill GPU graph resolves offsets itself.
408    /// q8_2f is excluded on purpose: its column field would need a
409    /// prescale stage on the device.
410    pub(crate) fn q8_row_parts(&self) -> Option<(usize, usize, usize, &[f32])> {
411        match self {
412            Self::Mapped {
413                idx,
414                dtype: TensorDtype::Q8Row,
415                rows,
416                cols,
417                row_scale,
418                col_field,
419                ..
420            } if col_field.is_empty() => Some((*idx, *rows, *cols, row_scale)),
421            _ => None,
422        }
423    }
424
425    /// The layout this tensor is stored in, when it is mapped from a model.
426    /// The frames branch on it — a q2tp gate against a q4tp down is a real
427    /// combination in the 2-bit profile and needs a different kernel.
428    pub fn model_dtype(&self) -> Option<cortiq_core::TensorDtype> {
429        match self {
430            Self::Mapped { dtype, .. } => Some(*dtype),
431            _ => None,
432        }
433    }
434
435    /// The tensor's index in the model directory, when it is mapped from one.
436    /// The GPU frames bind by index rather than by name — a name lookup per
437    /// layer per token is not free, and the index is what the device cache is
438    /// keyed on anyway.
439    pub fn model_idx(&self) -> Option<usize> {
440        match self {
441            Self::Mapped { idx, .. } => Some(*idx),
442            _ => None,
443        }
444    }
445
446    /// The model this tensor is mapped from, when it is mapped at all. The
447    /// GPU frames need the container to reach the bytes; a QTensor already
448    /// holds it, and threading a second handle down every call site to say
449    /// the same thing invites the two to disagree.
450    pub fn model_arc(&self) -> Option<std::sync::Arc<cortiq_core::CmfModel>> {
451        match self {
452            Self::Mapped { model, .. } => Some(model.clone()),
453            _ => None,
454        }
455    }
456
457    pub fn rows(&self) -> usize {
458        match self {
459            Self::F32 { rows, .. } | Self::Mapped { rows, .. } => *rows,
460        }
461    }
462
463    /// Mapped q4t handle (model + directory index) — the fused GPU FFN
464    /// needs the raw file coordinates of its three projections.
465    pub(crate) fn mapped_q4t(&self) -> Option<(&Arc<CmfModel>, usize)> {
466        match self {
467            Self::Mapped {
468                model,
469                idx,
470                dtype: TensorDtype::Q4Tiled,
471                ..
472            } => Some((model, *idx)),
473            _ => None,
474        }
475    }
476
477    /// Same slot as `mapped_q4t` for a q4tp tensor — the fused DiT FFN picks
478    /// its kernels by which of the two answers.
479    pub fn mapped_q4tp(&self) -> Option<(&Arc<CmfModel>, usize)> {
480        match self {
481            Self::Mapped {
482                model,
483                idx,
484                dtype: TensorDtype::Q4TiledP,
485                ..
486            } => Some((model, *idx)),
487            _ => None,
488        }
489    }
490
491    /// (model, tensor idx) for a q2tp mapped weight — the 2-bit twin of
492    /// `mapped_q4tp`, used by the mixed MoE profile.
493    pub fn mapped_q2tp(&self) -> Option<(&Arc<CmfModel>, usize)> {
494        match self {
495            Self::Mapped {
496                model,
497                idx,
498                dtype: TensorDtype::Q2TiledP,
499                ..
500            } => Some((model, *idx)),
501            _ => None,
502        }
503    }
504
505    pub fn cols(&self) -> usize {
506        match self {
507            Self::F32 { cols, .. } | Self::Mapped { cols, .. } => *cols,
508        }
509    }
510
511    /// (model, tensor idx) for a q1 mapped weight — the wgpu token graph
512    /// keys its resident VRAM cache by idx. None for any other dtype/kind.
513    pub fn mapped_q1(&self) -> Option<(&std::sync::Arc<CmfModel>, usize)> {
514        match self {
515            Self::Mapped {
516                model,
517                idx,
518                dtype: TensorDtype::Q1,
519                ..
520            } => Some((model, *idx)),
521            _ => None,
522        }
523    }
524
525    /// (model, idx, kind, row_scale) for a graph-capable mapped weight. kind:
526    /// 0=q8_row (per-row scales), 1=q1, 2=q4_tiled, 3=q1t (tile-embedded, no
527    /// rs). None for dtypes the token graph does not handle (q8_2f/q4_block/vbit).
528    pub fn graph_weight(&self) -> Option<(&std::sync::Arc<CmfModel>, usize, u8, &[f32])> {
529        match self {
530            Self::Mapped {
531                model,
532                idx,
533                dtype: TensorDtype::Q8Row,
534                row_scale,
535                ..
536            } => Some((model, *idx, 0, row_scale.as_slice())),
537            Self::Mapped {
538                model,
539                idx,
540                dtype: TensorDtype::Q1,
541                ..
542            } => Some((model, *idx, 1, &[])),
543            // Q4Tiled is kind 5, NOT 2: both carried 2 historically, and
544            // the wgpu token graph fed 18B interleaved tiles to the
545            // split-layout q4b kernel — garbage output on q4t models
546            // (caught by an end-to-end answer check on real Vulkan).
547            Self::Mapped {
548                model,
549                idx,
550                dtype: TensorDtype::Q4Tiled,
551                ..
552            } => Some((model, *idx, 5, &[])),
553            // Kind 6, not 5: q4tp's nibble stride and scale planes differ,
554            // and feeding them to the q4t kernel is exactly the mistake that
555            // produced garbage when Q4Tiled shared kind 2 with Q4Block.
556            Self::Mapped {
557                model,
558                idx,
559                dtype: TensorDtype::Q4TiledP,
560                ..
561            } => Some((model, *idx, 6, &[])),
562            Self::Mapped {
563                model,
564                idx,
565                dtype: TensorDtype::Q4Block,
566                ..
567            } => Some((model, *idx, 2, &[])),
568            Self::Mapped {
569                model,
570                idx,
571                dtype: TensorDtype::Q1T,
572                ..
573            } => Some((model, *idx, 3, &[])),
574            _ => None,
575        }
576    }
577
578    /// Dense f32 view — only for owned tensors. Masked/sparse execution
579    /// paths require it; quantized weights don't support masks yet.
580    pub fn as_f32(&self) -> Option<&[f32]> {
581        match self {
582            Self::F32 { data, .. } => Some(data),
583            Self::Mapped { .. } => None,
584        }
585    }
586
587    fn quant_bytes(&self) -> &[u8] {
588        match self {
589            Self::Mapped { model, idx, .. } => model.entry_bytes(&model.tensors[*idx]),
590            Self::F32 { .. } => unreachable!("quant_bytes on F32"),
591        }
592    }
593
594    /// Dequantize one row into `dst` (embedding lookup).
595    pub fn row_f32(&self, r: usize, dst: &mut [f32]) {
596        let cols = self.cols();
597        debug_assert_eq!(dst.len(), cols);
598        match self {
599            Self::F32 { data, .. } => dst.copy_from_slice(&data[r * cols..(r + 1) * cols]),
600            Self::Mapped {
601                dtype,
602                row_scale,
603                col_field,
604                vbit_offsets,
605                ..
606            } => {
607                if *dtype == TensorDtype::Q4Tiled {
608                    let bytes = self.quant_bytes();
609                    let gpr = cols / GROUP_SIZE;
610                    for gi in 0..gpr {
611                        let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
612                        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
613                        for (k, &b) in tile[2..].iter().enumerate() {
614                            dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
615                            dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
616                        }
617                    }
618                    return;
619                }
620                if *dtype == TensorDtype::Q4TiledP {
621                    let bytes = self.quant_bytes();
622                    let gpr = cols / GROUP_SIZE;
623                    let v = Q4tpView::new(bytes, self.rows(), cols);
624                    let mut sc = vec![0f32; gpr];
625                    v.scales_into(r, gpr, &mut sc);
626                    for gi in 0..gpr {
627                        let tile = &v.nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
628                        let s = sc[gi];
629                        for (k, &b) in tile.iter().enumerate() {
630                            dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
631                            dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
632                        }
633                    }
634                    return;
635                }
636                if *dtype == TensorDtype::Q2TiledP {
637                    let bytes = self.quant_bytes();
638                    let gpr = cols / GROUP_SIZE;
639                    let v = Q4tpView::new_q2(bytes, self.rows(), cols);
640                    let mut sc = vec![0f32; gpr];
641                    v.scales_into(r, gpr, &mut sc);
642                    for gi in 0..gpr {
643                        let ch =
644                            &v.nib[(r * gpr + gi) * Q2TP_CHUNK..(r * gpr + gi + 1) * Q2TP_CHUNK];
645                        let s = sc[gi];
646                        for (k, &b) in ch.iter().enumerate() {
647                            for j in 0..4 {
648                                dst[gi * GROUP_SIZE + k * 4 + j] =
649                                    (((b >> (2 * j)) & 3) as f32 - 1.5) * s;
650                            }
651                        }
652                    }
653                    return;
654                }
655                if *dtype == TensorDtype::Q4Block {
656                    let (packed, scales) = q4_split(self.quant_bytes(), self.rows(), cols);
657                    let gpr = cols / GROUP_SIZE;
658                    for gi in 0..gpr {
659                        let g = r * gpr + gi;
660                        let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
661                        for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
662                            dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
663                            dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
664                        }
665                    }
666                    return;
667                }
668                if *dtype == TensorDtype::Q1 {
669                    let bytes = self.quant_bytes();
670                    let gpr = cols / GROUP_SIZE;
671                    for gi in 0..gpr {
672                        let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
673                        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
674                        for (j, &b) in tile[2..].iter().enumerate() {
675                            for k in 0..8 {
676                                dst[gi * GROUP_SIZE + j * 8 + k] =
677                                    (((b >> k) & 1) as f32 * 2.0 - 1.0) * s;
678                            }
679                        }
680                    }
681                    return;
682                }
683                if *dtype == TensorDtype::Q1T {
684                    let bytes = self.quant_bytes();
685                    let gpr = cols / GROUP_SIZE;
686                    let base_len = self.rows() * gpr * cortiq_core::quant::Q1T_TILE;
687                    for gi in 0..gpr {
688                        let off = (r * gpr + gi) * cortiq_core::quant::Q1T_TILE;
689                        let s = cortiq_core::quant::f16_to_f32(u16::from_le_bytes([
690                            bytes[off],
691                            bytes[off + 1],
692                        ]));
693                        let codes = &bytes[off + 2..off + cortiq_core::quant::Q1T_TILE];
694                        for k in 0..GROUP_SIZE {
695                            dst[gi * GROUP_SIZE + k] = match cortiq_core::quant::q1t_code(codes, k)
696                            {
697                                1 => s,
698                                2 => -s,
699                                _ => 0.0,
700                            };
701                        }
702                    }
703                    // Overlay
704                    let rows = self.rows();
705                    let entries = base_len + (rows + 1) * 4;
706                    if entries <= bytes.len() {
707                        let ptrs = &bytes[base_len..base_len + (rows + 1) * 4];
708                        let r0 = u32::from_le_bytes([
709                            ptrs[r * 4],
710                            ptrs[r * 4 + 1],
711                            ptrs[r * 4 + 2],
712                            ptrs[r * 4 + 3],
713                        ]) as usize;
714                        let r1 = u32::from_le_bytes([
715                            ptrs[(r + 1) * 4],
716                            ptrs[(r + 1) * 4 + 1],
717                            ptrs[(r + 1) * 4 + 2],
718                            ptrs[(r + 1) * 4 + 3],
719                        ]) as usize;
720                        let off = entries + r0 * 4;
721                        for i in 0..r1 - r0 {
722                            let item = &bytes[off + i * 4..off + i * 4 + 4];
723                            let c = u16::from_le_bytes([item[0], item[1]]) as usize;
724                            let v = cortiq_core::quant::f16_to_f32(u16::from_le_bytes([
725                                item[2], item[3],
726                            ]));
727                            if c < cols {
728                                dst[c] = v;
729                            }
730                        }
731                    }
732                    return;
733                }
734                if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
735                    let bytes = self.quant_bytes();
736                    let rows = self.rows();
737                    let ng = cols / GROUP_SIZE;
738                    let bits = &bytes[..rows];
739                    let sc_off = rows;
740                    // Precomputed at load — embedding lookup used to scan
741                    // the bit-widths of every preceding row (O(token_id)).
742                    let off = vbit_offsets[r];
743                    let b = bits[r] as usize;
744                    let l = ((1usize << (b - 1)) - 1) as f32;
745                    let data = &bytes[off..];
746                    let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
747                    for (i, d) in dst.iter_mut().enumerate() {
748                        while nbits < b {
749                            acc = (acc << 8) | data[idx] as u64;
750                            idx += 1;
751                            nbits += 8;
752                        }
753                        let u = ((acc >> (nbits - b)) & ((1u64 << b) - 1)) as f32;
754                        nbits -= b;
755                        let so = (r * ng + i / GROUP_SIZE) * 2;
756                        let sv = f16_to_f32(u16::from_le_bytes([
757                            bytes[sc_off + so],
758                            bytes[sc_off + so + 1],
759                        ]));
760                        *d = (u - l) * sv;
761                    }
762                    return;
763                }
764                let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
765                let s = row_scale[r];
766                match dtype {
767                    TensorDtype::Q8Row => {
768                        for (d, &b) in dst.iter_mut().zip(q) {
769                            *d = (b as i8) as f32 * s;
770                        }
771                    }
772                    TensorDtype::Q8_2f => {
773                        for (i, (d, &b)) in dst.iter_mut().zip(q).enumerate() {
774                            *d = (b as i8) as f32 * s * col_field[i];
775                        }
776                    }
777                    _ => unreachable!(),
778                }
779            }
780        }
781    }
782
783    /// Can this tensor's columns be read cheaply (for sparse down_proj)?
784    /// True for F32/Q8Row/Q8_2f (per-row scale, direct strided access);
785    /// false for group-packed q4/vbit (column access would unpack whole
786    /// groups — sparse execution falls back to f32 for those).
787    pub fn sparse_col_ok(&self) -> bool {
788        match self {
789            Self::F32 { .. } => true,
790            Self::Mapped { dtype, .. } => {
791                matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
792            }
793        }
794    }
795
796    /// down_proj [hidden, inter]: accumulate `w · col(c)` into `out`
797    /// [hidden] — reads ONLY column `c` (one neuron) from the mmap,
798    /// no full-matrix dequant. `out[k] += w · down[k, c]`.
799    pub fn add_col_scaled(&self, c: usize, w: f32, out: &mut [f32]) {
800        let inter = self.cols();
801        let hidden = self.rows();
802        debug_assert_eq!(out.len(), hidden);
803        match self {
804            Self::F32 { data, .. } => {
805                for (k, o) in out.iter_mut().enumerate() {
806                    *o += w * data[k * inter + c];
807                }
808            }
809            Self::Mapped {
810                dtype,
811                row_scale,
812                col_field,
813                ..
814            } => {
815                let q = self.quant_bytes();
816                let colf = if *dtype == TensorDtype::Q8_2f {
817                    col_field[c]
818                } else {
819                    1.0
820                };
821                let wc = w * colf;
822                for (k, o) in out.iter_mut().enumerate() {
823                    let b = q[k * inter + c] as i8 as f32;
824                    *o += wc * b * row_scale[k];
825                }
826            }
827        }
828    }
829
830    /// Dot of row `r` with `x` (gate/up active-neuron path). Reads only
831    /// row `r` from the mmap — no full dequant. q4/vbit dequant the row
832    /// into `scratch` first (rare for active-FFN weights).
833    pub fn row_dot(&self, r: usize, x: &[f32], scratch: &mut [f32]) -> f32 {
834        let cols = self.cols();
835        match self {
836            Self::F32 { data, .. } => {
837                let row = &data[r * cols..(r + 1) * cols];
838                row.iter().zip(x).map(|(w, v)| w * v).sum()
839            }
840            Self::Mapped {
841                dtype,
842                row_scale,
843                col_field,
844                ..
845            } => match dtype {
846                TensorDtype::Q8Row => {
847                    let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
848                    dot_i8_f32(q, x) * row_scale[r]
849                }
850                TensorDtype::Q8_2f => {
851                    let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
852                    dot_i8_col_f32(q, x, col_field) * row_scale[r]
853                }
854                _ => {
855                    self.row_f32(r, scratch);
856                    scratch.iter().zip(x).map(|(w, v)| w * v).sum()
857                }
858            },
859        }
860    }
861
862    /// `out = W · x` (row-major). F32 delegates to the historical
863    /// bit-exact path; Mapped runs the fused int8 kernel.
864    pub fn matvec(&self, x: &[f32], out: &mut [f32], pool: Option<&Pool>) {
865        match self {
866            // NOTE: `out.len()` DRIVES this arm — it computes that many rows,
867            // and `x.len()` is the stride. A short `out` is legitimate here,
868            // which is why the check below lives in the Mapped arm only.
869            Self::F32 { data, .. } => matvec_rows(pool, data, x, out),
870            Self::Mapped {
871                model,
872                idx,
873                dtype,
874                rows,
875                cols,
876                row_scale,
877                col_field,
878                vbit_offsets,
879                repack,
880            } => {
881                let _ = (model, idx);
882                // Every kernel below writes `rows` entries through a raw
883                // pointer, so a short `out` is an out-of-bounds WRITE, not a
884                // wrong answer: it scribbles on the allocator's metadata and
885                // the process aborts much later, somewhere innocent
886                // (`double free or corruption`, `corrupted double-linked
887                // list`). The debug_assert two of the kernels carried is
888                // compiled out of the release — exactly the build where it
889                // matters. Fail here instead, while the caller is still on
890                // the stack to be named.
891                assert!(
892                    out.len() >= *rows && x.len() >= *cols,
893                    "matvec {rows}x{cols}: out {} (need {rows}), x {} (need {cols})",
894                    out.len(),
895                    x.len(),
896                );
897                if *dtype == TensorDtype::Q4Block {
898                    // GPU route (wgpu q4b kernel) for large q4_block matvecs —
899                    // gives NVIDIA/AMD/Intel q4 models a GPU path. Probe keeps
900                    // the winner; Metal returns false → the CPU kernel below.
901                    if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
902                        let t0 = std::time::Instant::now();
903                        match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
904                            crate::gpu::ProbeArm::Gpu => {
905                                if crate::gpu::q4b_matvec(model, *idx, x, *rows, *cols, out) {
906                                    crate::gpu::probe_record(
907                                        crate::gpu::OpClass::Matvec,
908                                        true,
909                                        t0.elapsed(),
910                                    );
911                                    return;
912                                }
913                            }
914                            crate::gpu::ProbeArm::CpuTimed => {
915                                q4matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
916                                crate::gpu::probe_record(
917                                    crate::gpu::OpClass::Matvec,
918                                    false,
919                                    t0.elapsed(),
920                                );
921                                return;
922                            }
923                            crate::gpu::ProbeArm::Cpu => {}
924                        }
925                    }
926                    q4matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
927                    return;
928                }
929                if *dtype == TensorDtype::Q4Tiled {
930                    q4t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
931                    return;
932                }
933                if *dtype == TensorDtype::Q4TiledP {
934                    // GPU route for large q4tp matvecs — the lm_head class.
935                    // On a q4tp checkpoint the head is the biggest single
936                    // host matvec left in the decode step, and the batched
937                    // kernel at b=1 already exists on both backends. Probe
938                    // keeps the winner, same as q4_block above.
939                    if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
940                        let t0 = std::time::Instant::now();
941                        let cls = crate::gpu::matvec_class(*rows, *cols);
942                        match crate::gpu::probe_arm(cls) {
943                            crate::gpu::ProbeArm::Gpu => {
944                                if crate::gpu::q4tp_matvec(model, *idx, x, *rows, *cols, out) {
945                                    crate::gpu::probe_record(cls, true, t0.elapsed());
946                                    return;
947                                }
948                            }
949                            crate::gpu::ProbeArm::CpuTimed => {
950                                q4tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
951                                crate::gpu::probe_record(cls, false, t0.elapsed());
952                                return;
953                            }
954                            crate::gpu::ProbeArm::Cpu => {}
955                        }
956                    }
957                    q4tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
958                    return;
959                }
960                if *dtype == TensorDtype::Q2TiledP {
961                    q2tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
962                    return;
963                }
964                if *dtype == TensorDtype::Q1 {
965                    // GPU route for large q1 matvecs (out_proj / lm_head
966                    // class): the CPU q1 kernel is load-port-bound at
967                    // ~4 GB/s/core, the GPU one is bandwidth-bound — the
968                    // probe measures both arms and keeps the winner.
969                    if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
970                        let t0 = std::time::Instant::now();
971                        let arm = if crate::gpu::q1_force() {
972                            crate::gpu::ProbeArm::Gpu
973                        } else {
974                            crate::gpu::probe_arm(crate::gpu::OpClass::Matvec)
975                        };
976                        match arm {
977                            crate::gpu::ProbeArm::Gpu => {
978                                if crate::gpu::q1_matvec(model, *idx, x, *rows, *cols, out) {
979                                    crate::gpu::probe_record(
980                                        crate::gpu::OpClass::Matvec,
981                                        true,
982                                        t0.elapsed(),
983                                    );
984                                    return;
985                                }
986                            }
987                            crate::gpu::ProbeArm::CpuTimed => {
988                                q1_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
989                                crate::gpu::probe_record(
990                                    crate::gpu::OpClass::Matvec,
991                                    false,
992                                    t0.elapsed(),
993                                );
994                                return;
995                            }
996                            crate::gpu::ProbeArm::Cpu => {}
997                        }
998                    }
999                    q1_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1000                    return;
1001                }
1002                if *dtype == TensorDtype::Q1T {
1003                    // GPU route for large q1t matvecs: the ternary BASE dot runs
1004                    // on the GPU (load-port-bound on CPU, like q1), then the
1005                    // sparse overlay is added on the CPU. Probe keeps the winner.
1006                    if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1007                        let t0 = std::time::Instant::now();
1008                        match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
1009                            crate::gpu::ProbeArm::Gpu => {
1010                                if crate::gpu::q1t_matvec(model, *idx, x, *rows, *cols, out) {
1011                                    q1t_add_overlay(self.quant_bytes(), x, *rows, *cols, out, pool);
1012                                    crate::gpu::probe_record(
1013                                        crate::gpu::OpClass::Matvec,
1014                                        true,
1015                                        t0.elapsed(),
1016                                    );
1017                                    return;
1018                                }
1019                            }
1020                            crate::gpu::ProbeArm::CpuTimed => {
1021                                q1t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1022                                crate::gpu::probe_record(
1023                                    crate::gpu::OpClass::Matvec,
1024                                    false,
1025                                    t0.elapsed(),
1026                                );
1027                                return;
1028                            }
1029                            crate::gpu::ProbeArm::Cpu => {}
1030                        }
1031                    }
1032                    q1t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1033                    return;
1034                }
1035                if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1036                    vbitmatvec(self.quant_bytes(), vbit_offsets, x, *rows, *cols, out, pool);
1037                    return;
1038                }
1039                let xs = prescale(x, col_field, *dtype);
1040                // D5: large q8 matrices (lm_head-class) — hybrid
1041                // CPU∥GPU: split the rows, both sides compute
1042                // SIMULTANEOUSLY (same math, shared prescale).
1043                // GPU share: CMF_GPU_SPLIT (0..1, default 0.5).
1044                if *rows >= crate::gpu::min_rows()
1045                    && matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
1046                    && gpu_lmhead_enabled()
1047                    && crate::gpu::enabled_here()
1048                {
1049                    // Runtime probe: alternate the hybrid against the
1050                    // pure-CPU matvec, keep whichever is faster HERE.
1051                    let t0 = std::time::Instant::now();
1052                    match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
1053                        crate::gpu::ProbeArm::Gpu => {}
1054                        crate::gpu::ProbeArm::CpuTimed => {
1055                            qmatvec(
1056                                self.quant_bytes(),
1057                                repack,
1058                                row_scale,
1059                                x,
1060                                col_field,
1061                                *dtype,
1062                                *rows,
1063                                *cols,
1064                                out,
1065                                pool,
1066                            );
1067                            crate::gpu::probe_record(
1068                                crate::gpu::OpClass::Matvec,
1069                                false,
1070                                t0.elapsed(),
1071                            );
1072                            return;
1073                        }
1074                        crate::gpu::ProbeArm::Cpu => {
1075                            qmatvec(
1076                                self.quant_bytes(),
1077                                repack,
1078                                row_scale,
1079                                x,
1080                                col_field,
1081                                *dtype,
1082                                *rows,
1083                                *cols,
1084                                out,
1085                                pool,
1086                            );
1087                            return;
1088                        }
1089                    }
1090                    let frac = gpu_split_frac();
1091                    let cpu_rows = ((*rows as f32) * (1.0 - frac)) as usize;
1092                    let (out_cpu, out_gpu) = out.split_at_mut(cpu_rows);
1093                    let bytes = self.quant_bytes();
1094                    let ok = std::thread::scope(|sc| {
1095                        let g = sc.spawn(|| {
1096                            crate::gpu::q8_matvec_range(
1097                                model,
1098                                *idx,
1099                                cpu_rows,
1100                                &row_scale[cpu_rows..],
1101                                &xs,
1102                                *rows - cpu_rows,
1103                                *cols,
1104                                out_gpu,
1105                            )
1106                        });
1107                        if cpu_rows > 0 {
1108                            // Repack prefix covers the full groups of the
1109                            // CPU half (the split starts at row 0).
1110                            let rep_cpu = if repack.is_empty() {
1111                                &[][..]
1112                            } else {
1113                                &repack[..(cpu_rows / 4) * 4 * *cols]
1114                            };
1115                            qmatvec(
1116                                &bytes[..cpu_rows * *cols],
1117                                rep_cpu,
1118                                &row_scale[..cpu_rows],
1119                                x,
1120                                col_field,
1121                                *dtype,
1122                                cpu_rows,
1123                                *cols,
1124                                out_cpu,
1125                                pool,
1126                            );
1127                        }
1128                        g.join().unwrap_or(false)
1129                    });
1130                    if ok {
1131                        crate::gpu::probe_record(crate::gpu::OpClass::Matvec, true, t0.elapsed());
1132                        return;
1133                    }
1134                    // GPU failed — CPU finishes its half (rows rebased —
1135                    // group offsets don't line up, mmap layout only).
1136                    qmatvec(
1137                        &bytes[cpu_rows * *cols..(*rows) * *cols],
1138                        &[],
1139                        &row_scale[cpu_rows..],
1140                        x,
1141                        col_field,
1142                        *dtype,
1143                        *rows - cpu_rows,
1144                        *cols,
1145                        out_gpu,
1146                        pool,
1147                    );
1148                    return;
1149                }
1150                qmatvec(
1151                    self.quant_bytes(),
1152                    repack,
1153                    row_scale,
1154                    x,
1155                    col_field,
1156                    *dtype,
1157                    *rows,
1158                    *cols,
1159                    out,
1160                    pool,
1161                );
1162            }
1163        }
1164    }
1165
1166    /// Fused two-input matvec (MTP verify pair): weights streamed once.
1167    pub fn matvec2(
1168        &self,
1169        x1: &[f32],
1170        x2: &[f32],
1171        o1: &mut [f32],
1172        o2: &mut [f32],
1173        pool: Option<&Pool>,
1174    ) {
1175        match self {
1176            Self::F32 { data, .. } => matvec_rows2(pool, data, x1, x2, o1, o2),
1177            Self::Mapped {
1178                dtype,
1179                rows,
1180                cols,
1181                row_scale,
1182                col_field,
1183                vbit_offsets,
1184                ..
1185            } => {
1186                if *dtype == TensorDtype::Q4Block {
1187                    q4matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1188                    return;
1189                }
1190                if *dtype == TensorDtype::Q4Tiled {
1191                    q4t_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1192                    return;
1193                }
1194                if *dtype == TensorDtype::Q4TiledP {
1195                    q4tp_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1196                    return;
1197                }
1198                if *dtype == TensorDtype::Q2TiledP {
1199                    q2tp_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1200                    return;
1201                }
1202                if *dtype == TensorDtype::Q1 {
1203                    q1_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1204                    return;
1205                }
1206                if *dtype == TensorDtype::Q1T {
1207                    // Fused ternary pair: one row pass, the register
1208                    // unpack shared across both streams on ARM. (Q1T
1209                    // lacks a row_scale array — scales live inline in
1210                    // the tiles — so it must not fall through to the
1211                    // q8 qmatvec2 below.)
1212                    q1t_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1213                    return;
1214                }
1215                if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1216                    vbitmatvec2(
1217                        self.quant_bytes(),
1218                        vbit_offsets,
1219                        x1,
1220                        x2,
1221                        *rows,
1222                        *cols,
1223                        o1,
1224                        o2,
1225                        pool,
1226                    );
1227                    return;
1228                }
1229                qmatvec2(
1230                    self.quant_bytes(),
1231                    row_scale,
1232                    x1,
1233                    x2,
1234                    col_field,
1235                    *dtype,
1236                    *rows,
1237                    *cols,
1238                    o1,
1239                    o2,
1240                    pool,
1241                );
1242            }
1243        }
1244    }
1245}
1246
1247impl QTensor {
1248    /// Batched matvec (prefill-GEMM): xs — row-major [b, cols],
1249    /// out — row-major [b, rows]. Element-wise semantics are IDENTICAL
1250    /// to b matvec calls (same dot kernels in the same order); the win —
1251    /// the weight row streams from DRAM once per batch, not b times.
1252    pub fn matmat(&self, xs_all: &[f32], b: usize, out: &mut [f32], pool: Option<&Pool>) {
1253        let cols = self.cols();
1254        let rows = self.rows();
1255        debug_assert_eq!(xs_all.len(), b * cols);
1256        debug_assert_eq!(out.len(), b * rows);
1257        // GPTQ calibration: fold this layer's inputs into its Hessian. Only
1258        // Mapped tensors carry a directory name; the check is a relaxed
1259        // atomic load, free when not calibrating.
1260        if crate::gptq_capture::capturing() {
1261            if let Self::Mapped { model, idx, .. } = self {
1262                crate::gptq_capture::accumulate(&model.tensors[*idx].name, xs_all, b, cols);
1263            }
1264        }
1265        match self {
1266            Self::F32 { data, .. } => {
1267                let out_addr = SendMut(out.as_mut_ptr());
1268                let run = |start: usize, end: usize| {
1269                    for o in start..end {
1270                        let row = &data[o * cols..(o + 1) * cols];
1271                        for bi in 0..b {
1272                            let x = &xs_all[bi * cols..(bi + 1) * cols];
1273                            let mut acc = 0f32;
1274                            for j in 0..cols {
1275                                acc += row[j] * x[j];
1276                            }
1277                            unsafe { *out_addr.at(bi * rows + o) = acc };
1278                        }
1279                    }
1280                };
1281                dispatch_rows(pool, rows, &run);
1282            }
1283            Self::Mapped {
1284                dtype,
1285                row_scale,
1286                col_field,
1287                vbit_offsets,
1288                ..
1289            } => {
1290                if *dtype == TensorDtype::Q4Block {
1291                    q4matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1292                    return;
1293                }
1294                if *dtype == TensorDtype::Q4TiledP {
1295                    // GPU batched q4tp GEMM (dequant + f32nt mul_mm on the
1296                    // device); the probe keeps whichever beats the CPU arm.
1297                    // Narrow (prompt-encode) and wide (DiT) batches probe
1298                    // as separate classes — the regimes have opposite
1299                    // winners and one shared verdict locked the wrong arm.
1300                    // Kill switch (gpu::mm_kill): one grossly slow GPU op
1301                    // (a fair-condition op is ≤~100 ms even at 1024px)
1302                    // means the device is contended by another process
1303                    // (e.g. a simulator) — verdicts are per-process, so
1304                    // without the bail the whole render crawls behind
1305                    // someone else's queue.
1306                    if b >= 32
1307                        && b * rows * cols >= 128_000_000
1308                        && cols % 32 == 0
1309                        && !crate::gpu::mm_killed()
1310                        && crate::gpu::enabled_here()
1311                    {
1312                        let class = if b >= 128 {
1313                            crate::gpu::OpClass::MatmatWide
1314                        } else {
1315                            crate::gpu::OpClass::Matmat
1316                        };
1317                        if let Self::Mapped { model, idx, .. } = self {
1318                            let t0 = std::time::Instant::now();
1319                            match crate::gpu::probe_arm(class) {
1320                                crate::gpu::ProbeArm::Gpu => {
1321                                    if crate::gpu::q4tp_matmat(
1322                                        model, *idx, xs_all, b, rows, cols, out,
1323                                    ) {
1324                                        let el = t0.elapsed();
1325                                        // Work-proportional budget: ~8× the
1326                                        // fair-device estimate (+20 ms slack).
1327                                        // An absolute cap missed the worst
1328                                        // case — contended ops sit at
1329                                        // 100–240 ms each and still bury a
1330                                        // render whose fair op is 3–9 ms.
1331                                        // Cold ops (first PSO build, buffer
1332                                        // alloc) are exempt: a one-off
1333                                        // ~50 ms compile is not contention.
1334                                        let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
1335                                        let budget = std::time::Duration::from_secs_f64(
1336                                            flops / 1.5e12 * 8.0 + 0.020,
1337                                        );
1338                                        if el > budget && !crate::gpu::probe_was_cold() {
1339                                            tracing::warn!(
1340                                                "gpu q4tp matmat took {el:?} (budget {budget:?}) — \
1341                                                 device contended, CPU for the rest of the process"
1342                                            );
1343                                            crate::gpu::mm_kill();
1344                                        }
1345                                        crate::gpu::probe_record(class, true, el);
1346                                        return;
1347                                    }
1348                                }
1349                                crate::gpu::ProbeArm::CpuTimed => {
1350                                    q4tp_matmat(
1351                                        self.quant_bytes(),
1352                                        xs_all,
1353                                        b,
1354                                        rows,
1355                                        cols,
1356                                        out,
1357                                        pool,
1358                                    );
1359                                    crate::gpu::probe_record(class, false, t0.elapsed());
1360                                    return;
1361                                }
1362                                crate::gpu::ProbeArm::Cpu => {}
1363                            }
1364                        }
1365                    }
1366                    q4tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1367                    return;
1368                }
1369                if *dtype == TensorDtype::Q2TiledP {
1370                    // Without this arm a q2tp tensor falls through to the
1371                    // q8 fallback, which reads it at one BYTE per weight —
1372                    // a 2x overrun that killed pool workers mid-prefill
1373                    // while the dispatcher waited forever.
1374                    q2tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1375                    return;
1376                }
1377                if *dtype == TensorDtype::Q4Tiled {
1378                    // GPU batched q4t GEMM (dequant + f32nt mul_mm on the
1379                    // device); the probe keeps whichever beats the CPU arm.
1380                    // Narrow (prompt-encode) and wide (DiT) batches probe
1381                    // as separate classes — the regimes have opposite
1382                    // winners and one shared verdict locked the wrong arm.
1383                    // Kill switch (gpu::mm_kill): one grossly slow GPU op
1384                    // (a fair-condition op is ≤~100 ms even at 1024px)
1385                    // means the device is contended by another process
1386                    // (e.g. a simulator) — verdicts are per-process, so
1387                    // without the bail the whole render crawls behind
1388                    // someone else's queue.
1389                    if b >= 32
1390                        && b * rows * cols >= 128_000_000
1391                        && cols % 32 == 0
1392                        && !crate::gpu::mm_killed()
1393                        && crate::gpu::enabled_here()
1394                    {
1395                        let class = if b >= 128 {
1396                            crate::gpu::OpClass::MatmatWide
1397                        } else {
1398                            crate::gpu::OpClass::Matmat
1399                        };
1400                        if let Self::Mapped { model, idx, .. } = self {
1401                            let t0 = std::time::Instant::now();
1402                            match crate::gpu::probe_arm(class) {
1403                                crate::gpu::ProbeArm::Gpu => {
1404                                    if crate::gpu::q4t_matmat(
1405                                        model, *idx, xs_all, b, rows, cols, out,
1406                                    ) {
1407                                        let el = t0.elapsed();
1408                                        // Work-proportional budget: ~8× the
1409                                        // fair-device estimate (+20 ms slack).
1410                                        // An absolute cap missed the worst
1411                                        // case — contended ops sit at
1412                                        // 100–240 ms each and still bury a
1413                                        // render whose fair op is 3–9 ms.
1414                                        // Cold ops (first PSO build, buffer
1415                                        // alloc) are exempt: a one-off
1416                                        // ~50 ms compile is not contention.
1417                                        let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
1418                                        let budget = std::time::Duration::from_secs_f64(
1419                                            flops / 1.5e12 * 8.0 + 0.020,
1420                                        );
1421                                        if el > budget && !crate::gpu::probe_was_cold() {
1422                                            tracing::warn!(
1423                                                "gpu q4t matmat took {el:?} (budget {budget:?}) — \
1424                                                 device contended, CPU for the rest of the process"
1425                                            );
1426                                            crate::gpu::mm_kill();
1427                                        }
1428                                        crate::gpu::probe_record(class, true, el);
1429                                        return;
1430                                    }
1431                                }
1432                                crate::gpu::ProbeArm::CpuTimed => {
1433                                    q4t_matmat(
1434                                        self.quant_bytes(),
1435                                        xs_all,
1436                                        b,
1437                                        rows,
1438                                        cols,
1439                                        out,
1440                                        pool,
1441                                    );
1442                                    crate::gpu::probe_record(class, false, t0.elapsed());
1443                                    return;
1444                                }
1445                                crate::gpu::ProbeArm::Cpu => {}
1446                            }
1447                        }
1448                    }
1449                    q4t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1450                    return;
1451                }
1452                if *dtype == TensorDtype::Q1 {
1453                    // GPU batched q1 GEMM for wide prefill (q1_mul_mm on the
1454                    // device); the probe keeps whichever beats the CPU matmat.
1455                    if b >= 32
1456                        && b * rows * cols >= 128_000_000
1457                        && cols % 64 == 0
1458                        && crate::gpu::enabled_here()
1459                    {
1460                        if let Self::Mapped { model, idx, .. } = self {
1461                            let t0 = std::time::Instant::now();
1462                            match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
1463                                crate::gpu::ProbeArm::Gpu => {
1464                                    if crate::gpu::q1_matmat(
1465                                        model, *idx, xs_all, b, rows, cols, out,
1466                                    ) {
1467                                        crate::gpu::probe_record(
1468                                            crate::gpu::OpClass::Matmat,
1469                                            true,
1470                                            t0.elapsed(),
1471                                        );
1472                                        return;
1473                                    }
1474                                }
1475                                crate::gpu::ProbeArm::CpuTimed => {
1476                                    q1_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1477                                    crate::gpu::probe_record(
1478                                        crate::gpu::OpClass::Matmat,
1479                                        false,
1480                                        t0.elapsed(),
1481                                    );
1482                                    return;
1483                                }
1484                                crate::gpu::ProbeArm::Cpu => {}
1485                            }
1486                        }
1487                    }
1488                    q1_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1489                    return;
1490                }
1491                if *dtype == TensorDtype::Q1T {
1492                    // GPU batched GEMM for wide prefill (base + overlay on the
1493                    // device); probe keeps the winner vs the CPU matmat.
1494                    if b >= 32 && b * rows * cols >= 128_000_000 && crate::gpu::enabled_here() {
1495                        if let Self::Mapped { model, idx, .. } = self {
1496                            let t0 = std::time::Instant::now();
1497                            match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
1498                                crate::gpu::ProbeArm::Gpu => {
1499                                    if crate::gpu::q1t_matmat(
1500                                        model, *idx, xs_all, b, rows, cols, out,
1501                                    ) {
1502                                        crate::gpu::probe_record(
1503                                            crate::gpu::OpClass::Matmat,
1504                                            true,
1505                                            t0.elapsed(),
1506                                        );
1507                                        return;
1508                                    }
1509                                }
1510                                crate::gpu::ProbeArm::CpuTimed => {
1511                                    q1t_matmat(
1512                                        self.quant_bytes(),
1513                                        xs_all,
1514                                        b,
1515                                        rows,
1516                                        cols,
1517                                        out,
1518                                        pool,
1519                                    );
1520                                    crate::gpu::probe_record(
1521                                        crate::gpu::OpClass::Matmat,
1522                                        false,
1523                                        t0.elapsed(),
1524                                    );
1525                                    return;
1526                                }
1527                                crate::gpu::ProbeArm::Cpu => {}
1528                            }
1529                        }
1530                    }
1531                    q1t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1532                    return;
1533                }
1534                if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1535                    vbitmatmat(
1536                        self.quant_bytes(),
1537                        vbit_offsets,
1538                        xs_all,
1539                        b,
1540                        rows,
1541                        cols,
1542                        out,
1543                        pool,
1544                    );
1545                    return;
1546                }
1547                let pre: Vec<std::borrow::Cow<'_, [f32]>> = (0..b)
1548                    .map(|bi| prescale(&xs_all[bi * cols..(bi + 1) * cols], col_field, *dtype))
1549                    .collect();
1550                // D5: large prefill-batch GEMMs — on the GPU (threshold by
1551                // work volume: submission carries b×rows×cols MACs).
1552                // Runtime probe: the naive GEMM shader + sync readback
1553                // lose to the CPU GEMM on slow driver stacks — alternate
1554                // both arms and keep the winner.
1555                if b >= 8 && b * rows * cols >= 128_000_000 && crate::gpu::enabled_here() {
1556                    if let Self::Mapped { model, idx, .. } = self {
1557                        let t0 = std::time::Instant::now();
1558                        match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
1559                            crate::gpu::ProbeArm::Gpu
1560                                if crate::gpu::probe_deciding(crate::gpu::OpClass::Matmat)
1561                                    && !crate::gpu::q8_resident_or_upload(model, *idx) =>
1562                            {
1563                                // Cold weights during probing: the upload
1564                                // has started, the count runs on the CPU —
1565                                // the GPU arm samples on the next touch.
1566                                let q = self.quant_bytes();
1567                                qmatmat(q, row_scale, &pre, rows, cols, out, pool);
1568                                return;
1569                            }
1570                            crate::gpu::ProbeArm::Gpu => {
1571                                let flat: Vec<f32> =
1572                                    pre.iter().flat_map(|v| v.iter().copied()).collect();
1573                                if crate::gpu::q8_matmat(
1574                                    model, *idx, row_scale, &flat, b, rows, cols, out,
1575                                ) {
1576                                    crate::gpu::probe_record(
1577                                        crate::gpu::OpClass::Matmat,
1578                                        true,
1579                                        t0.elapsed(),
1580                                    );
1581                                    return;
1582                                }
1583                            }
1584                            crate::gpu::ProbeArm::CpuTimed => {
1585                                let q = self.quant_bytes();
1586                                qmatmat(q, row_scale, &pre, rows, cols, out, pool);
1587                                crate::gpu::probe_record(
1588                                    crate::gpu::OpClass::Matmat,
1589                                    false,
1590                                    t0.elapsed(),
1591                                );
1592                                return;
1593                            }
1594                            crate::gpu::ProbeArm::Cpu => {}
1595                        }
1596                    }
1597                }
1598                let q = self.quant_bytes();
1599                qmatmat(q, row_scale, &pre, rows, cols, out, pool);
1600            }
1601        }
1602    }
1603}
1604
1605impl QTensor {
1606    /// Multi-matrix job (roadmap §3 P0): N tensors sharing one input
1607    /// run under a SINGLE pool dispatch — QKV or gate+up cost one
1608    /// barrier instead of N. Per-row math is the exact same kernel as
1609    /// `matvec` (bit-identical outputs); only the dispatch is fused.
1610    /// Falls back to N sequential matvecs when the set is not a uniform
1611    /// q8-family/F32 group or there is no pool.
1612    pub fn matvec_many<const N: usize>(
1613        ts: [&QTensor; N],
1614        x: &[f32],
1615        mut outs: [&mut [f32]; N],
1616        pool: Option<&Pool>,
1617    ) {
1618        let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
1619        let uniform_q8 = ts.iter().all(|t| {
1620            matches!(
1621                t,
1622                Self::Mapped {
1623                    dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f,
1624                    ..
1625                }
1626            )
1627        });
1628        let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
1629        let uniform_q4 = ts.iter().all(|t| {
1630            matches!(
1631                t,
1632                Self::Mapped {
1633                    dtype: TensorDtype::Q4Block,
1634                    ..
1635                }
1636            )
1637        });
1638        let uniform_vbit = ts.iter().all(|t| {
1639            matches!(
1640                t,
1641                Self::Mapped {
1642                    dtype: TensorDtype::Vbit | TensorDtype::VbitRo,
1643                    ..
1644                }
1645            )
1646        });
1647        let uniform_q1 = ts.iter().all(|t| {
1648            matches!(
1649                t,
1650                Self::Mapped {
1651                    dtype: TensorDtype::Q1,
1652                    ..
1653                }
1654            )
1655        });
1656        let uniform_q1t = ts.iter().all(|t| {
1657            matches!(
1658                t,
1659                Self::Mapped {
1660                    dtype: TensorDtype::Q1T,
1661                    ..
1662                }
1663            )
1664        });
1665        // q4tp is the skeleton dtype of the big MoE files, and without an arm
1666        // here every projection that shares an input paid its own pool
1667        // barrier: DeepSeek-V4's attention step alone hands this function
1668        // wq_a, wkv and both compressors' pairs off the same hidden state.
1669        let uniform_q4tp = ts.iter().all(|t| {
1670            matches!(
1671                t,
1672                Self::Mapped {
1673                    dtype: TensorDtype::Q4TiledP,
1674                    ..
1675                }
1676            )
1677        }) && ts.iter().all(|t| t.cols() == ts[0].cols() && t.cols() % GROUP_SIZE == 0);
1678        let Some(pool) = pool else {
1679            for (t, o) in ts.iter().zip(outs.iter_mut()) {
1680                t.matvec(x, o, None);
1681            }
1682            return;
1683        };
1684        if total_rows < 256
1685            || !(uniform_q8
1686                || uniform_f32
1687                || uniform_q4
1688                || uniform_vbit
1689                || uniform_q1
1690                || uniform_q1t
1691                || uniform_q4tp)
1692        {
1693            for (t, o) in ts.iter().zip(outs.iter_mut()) {
1694                t.matvec(x, o, Some(pool));
1695            }
1696            return;
1697        }
1698
1699        if uniform_q4tp {
1700            // Every tensor's rows laid end to end in one virtual row space,
1701            // so the whole set is ONE dispatch. The per-row body is the
1702            // `q4tp_matvec` arm verbatim — same activation split, same
1703            // accumulation order — so the outputs are bit-identical to the
1704            // sequential calls this replaces.
1705            let cols = ts[0].cols();
1706            let gpr = cols / GROUP_SIZE;
1707            let views: Vec<Q4tpView> = ts
1708                .iter()
1709                .map(|t| Q4tpView::new(t.quant_bytes(), t.rows(), cols))
1710                .collect();
1711            let rows_of: Vec<usize> = ts.iter().map(|t| t.rows()).collect();
1712            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1713            // flat index -> (which tensor, which of its rows)
1714            let locate = |flat: usize| -> (usize, usize) {
1715                let mut acc = 0;
1716                for (i, &r) in rows_of.iter().enumerate() {
1717                    if flat < acc + r {
1718                        return (i, flat - acc);
1719                    }
1720                    acc += r;
1721                }
1722                (rows_of.len() - 1, 0)
1723            };
1724            let (views, outs_addr) = (&views, &outs_addr);
1725            if a8w8_enabled() {
1726                let act = split_act(x);
1727                let act = &act;
1728                let run = |start: usize, end: usize| {
1729                    let mut sc = vec![0f32; gpr];
1730                    for flat in start..end {
1731                        let (t, r) = locate(flat);
1732                        let v = &views[t];
1733                        v.scales_into(r, gpr, &mut sc);
1734                        let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, &sc) * act.sx;
1735                        for &(j, xv) in &act.outliers {
1736                            let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
1737                            acc += w * s * xv;
1738                        }
1739                        // SAFETY: one worker owns each (tensor, row) pair.
1740                        unsafe { *outs_addr[t].at(r) = acc };
1741                    }
1742                };
1743                pool.run_rows(total_rows, &run);
1744            } else {
1745                let run = |start: usize, end: usize| {
1746                    let mut sc = vec![0f32; gpr];
1747                    for flat in start..end {
1748                        let (t, r) = locate(flat);
1749                        let v = &views[t];
1750                        v.scales_into(r, gpr, &mut sc);
1751                        // SAFETY: one worker owns each (tensor, row) pair.
1752                        unsafe { *outs_addr[t].at(r) = q4tp_row_exact(v.nib, r, gpr, x, &sc) };
1753                    }
1754                };
1755                pool.run_rows(total_rows, &run);
1756            }
1757            return;
1758        }
1759
1760        if uniform_q1 {
1761            // One shared activation split + group sums (q1 has no col
1762            // field; the same input feeds every tensor).
1763            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1764            if a8w8_enabled() {
1765                let act = split_act(x);
1766                let gsum = q1_group_sums(&act.xq, ts[0].cols() / GROUP_SIZE);
1767                let (act, gsum) = (&act, &gsum);
1768                let closures: [_; N] = std::array::from_fn(|i| {
1769                    let (bytes, gpr, out) =
1770                        (ts[i].quant_bytes(), ts[i].cols() / GROUP_SIZE, outs_addr[i]);
1771                    move |s: usize, e: usize| q1_range_a8w8(bytes, gpr, act, gsum, out, s, e)
1772                });
1773                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1774                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1775                pool.run_many(&parts);
1776            } else {
1777                let closures: [_; N] = std::array::from_fn(|i| {
1778                    let (bytes, gpr, out) =
1779                        (ts[i].quant_bytes(), ts[i].cols() / GROUP_SIZE, outs_addr[i]);
1780                    move |s: usize, e: usize| q1_range_f32(bytes, gpr, x, out, s, e)
1781                });
1782                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1783                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1784                pool.run_many(&parts);
1785            }
1786            return;
1787        }
1788
1789        if uniform_q1t {
1790            // Q1T batched: one shared activation split + overlay decode,
1791            // all tensors' rows in ONE pool dispatch (saves N−1 dispatches
1792            // and N−1 redundant split_act calls per layer).
1793            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1794            const TILE: usize = cortiq_core::quant::Q1T_TILE;
1795            if a8w8_enabled() {
1796                let act = split_act(x);
1797                let act = &act;
1798                let x_ref = x;
1799                let closures: [_; N] = std::array::from_fn(|i| {
1800                    let bytes = ts[i].quant_bytes();
1801                    let (rows, cols) = (ts[i].rows(), ts[i].cols());
1802                    let gpr = cols / GROUP_SIZE;
1803                    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
1804                    let out = outs_addr[i];
1805                    move |s: usize, e: usize| {
1806                        q1t_range_a8w8(bytes, gpr, rp_off, ent_off, has_ov, act, x_ref, out, s, e)
1807                    }
1808                });
1809                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1810                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1811                pool.run_many(&parts);
1812            } else {
1813                let x_ref = x;
1814                let closures: [_; N] = std::array::from_fn(|i| {
1815                    let bytes = ts[i].quant_bytes();
1816                    let (rows, cols) = (ts[i].rows(), ts[i].cols());
1817                    let gpr = cols / GROUP_SIZE;
1818                    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
1819                    let out = outs_addr[i];
1820                    move |s: usize, e: usize| {
1821                        q1t_range_f32_batch(bytes, gpr, rp_off, ent_off, has_ov, x_ref, out, s, e)
1822                    }
1823                });
1824                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1825                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1826                pool.run_many(&parts);
1827            }
1828            return;
1829        }
1830
1831        if uniform_q4 || uniform_vbit {
1832            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1833            // q4/vbit share one activation split — no per-tensor col field.
1834            if a8w8_enabled() {
1835                let act = split_act(x);
1836                let act = &act;
1837                if uniform_q4 {
1838                    let closures: [_; N] = std::array::from_fn(|i| {
1839                        let (packed, scales) =
1840                            q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
1841                        let (gpr, cols, out) =
1842                            (ts[i].cols() / GROUP_SIZE, ts[i].cols(), outs_addr[i]);
1843                        move |s: usize, e: usize| {
1844                            q4_range_a8w8(packed, scales, gpr, cols, act, out, s, e)
1845                        }
1846                    });
1847                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1848                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1849                    pool.run_many(&parts);
1850                } else {
1851                    let closures: [_; N] = std::array::from_fn(|i| {
1852                        let Self::Mapped { vbit_offsets, .. } = ts[i] else {
1853                            unreachable!()
1854                        };
1855                        let (bytes, rows, cols, out) = (
1856                            ts[i].quant_bytes(),
1857                            ts[i].rows(),
1858                            ts[i].cols(),
1859                            outs_addr[i],
1860                        );
1861                        move |s: usize, e: usize| {
1862                            vbit_range_a8w8(bytes, vbit_offsets, x, act, rows, cols, out, s, e)
1863                        }
1864                    });
1865                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1866                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1867                    pool.run_many(&parts);
1868                }
1869                return;
1870            }
1871            if uniform_q4 {
1872                let closures: [_; N] = std::array::from_fn(|i| {
1873                    let (packed, scales) =
1874                        q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
1875                    let (gpr, out) = (ts[i].cols() / GROUP_SIZE, outs_addr[i]);
1876                    move |s: usize, e: usize| q4_range_f32(packed, scales, gpr, x, out, s, e)
1877                });
1878                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1879                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1880                pool.run_many(&parts);
1881            } else {
1882                let closures: [_; N] = std::array::from_fn(|i| {
1883                    let Self::Mapped { vbit_offsets, .. } = ts[i] else {
1884                        unreachable!()
1885                    };
1886                    let (bytes, rows, cols, out) = (
1887                        ts[i].quant_bytes(),
1888                        ts[i].rows(),
1889                        ts[i].cols(),
1890                        outs_addr[i],
1891                    );
1892                    move |s: usize, e: usize| {
1893                        vbit_range_f32(bytes, vbit_offsets, x, rows, cols, out, s, e)
1894                    }
1895                });
1896                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1897                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1898                pool.run_many(&parts);
1899            }
1900            return;
1901        }
1902
1903        if uniform_f32 {
1904            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1905            let closures: [_; N] = std::array::from_fn(|i| {
1906                let Self::F32 { data, cols, .. } = ts[i] else {
1907                    unreachable!()
1908                };
1909                let out = outs_addr[i];
1910                move |start: usize, end: usize| {
1911                    for o in start..end {
1912                        let row = &data[o * cols..(o + 1) * cols];
1913                        let mut sum = 0.0f32;
1914                        for j in 0..*cols {
1915                            sum += row[j] * x[j];
1916                        }
1917                        // SAFETY: disjoint (tensor, row) cells per worker.
1918                        unsafe { *out.at(o) = sum };
1919                    }
1920                }
1921            });
1922            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1923                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1924            pool.run_many(&parts);
1925            return;
1926        }
1927
1928        // Uniform q8-family: per-tensor prescale (q8_2f col fields
1929        // differ per tensor) + the shared range kernels.
1930        struct Ctx<'a> {
1931            bytes: &'a [u8],
1932            #[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
1933            rep: &'a [u8],
1934            row_scale: &'a [f32],
1935            cols: usize,
1936            xs: std::borrow::Cow<'a, [f32]>,
1937        }
1938        let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
1939            let Self::Mapped {
1940                dtype,
1941                cols,
1942                row_scale,
1943                col_field,
1944                repack,
1945                ..
1946            } = ts[i]
1947            else {
1948                unreachable!()
1949            };
1950            Ctx {
1951                bytes: ts[i].quant_bytes(),
1952                rep: repack,
1953                row_scale,
1954                cols: *cols,
1955                xs: prescale(x, col_field, *dtype),
1956            }
1957        });
1958        let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1959        #[cfg(target_arch = "aarch64")]
1960        if sdot_enabled() {
1961            let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
1962            let closures: [_; N] = std::array::from_fn(|i| {
1963                let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
1964                move |start: usize, end: usize| {
1965                    q8_range_sdot(c.bytes, c.rep, c.row_scale, act, c.cols, out, start, end)
1966                }
1967            });
1968            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1969                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1970            pool.run_many(&parts);
1971            return;
1972        }
1973        #[cfg(target_arch = "x86_64")]
1974        if avx2_a8w8_enabled() {
1975            let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
1976            let closures: [_; N] = std::array::from_fn(|i| {
1977                let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
1978                move |start: usize, end: usize| {
1979                    q8_range_avx2(c.bytes, c.row_scale, act, c.cols, out, start, end)
1980                }
1981            });
1982            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1983                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1984            pool.run_many(&parts);
1985            return;
1986        }
1987        let closures: [_; N] = std::array::from_fn(|i| {
1988            let (c, out) = (&ctxs[i], outs_addr[i]);
1989            move |start: usize, end: usize| {
1990                q8_range_f32(c.bytes, c.row_scale, &c.xs, c.cols, out, start, end)
1991            }
1992        });
1993        let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1994            std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1995        pool.run_many(&parts);
1996    }
1997}
1998
1999impl QTensor {
2000    /// Pair-input multi-matrix job: N tensors × 2 shared inputs under a
2001    /// single pool dispatch — the MTP/pair decode path publishes one job
2002    /// for Q/K/V (and one for gate+up) instead of one per tensor.
2003    /// Per-row math is exactly `matvec2`'s kernels; bit-identical.
2004    #[allow(clippy::needless_range_loop)]
2005    pub fn matvec2_many<const N: usize>(
2006        ts: [&QTensor; N],
2007        x1: &[f32],
2008        x2: &[f32],
2009        mut o1s: [&mut [f32]; N],
2010        mut o2s: [&mut [f32]; N],
2011        pool: Option<&Pool>,
2012    ) {
2013        let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
2014        let uniform_q8 = ts.iter().all(|t| {
2015            matches!(
2016                t,
2017                Self::Mapped {
2018                    dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f,
2019                    ..
2020                }
2021            )
2022        });
2023        let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
2024        let uniform_q4 = ts.iter().all(|t| {
2025            matches!(
2026                t,
2027                Self::Mapped {
2028                    dtype: TensorDtype::Q4Block,
2029                    ..
2030                }
2031            )
2032        });
2033        let uniform_vbit = ts.iter().all(|t| {
2034            matches!(
2035                t,
2036                Self::Mapped {
2037                    dtype: TensorDtype::Vbit | TensorDtype::VbitRo,
2038                    ..
2039                }
2040            )
2041        });
2042        let fusable = pool.is_some()
2043            && total_rows >= 256
2044            && (uniform_q8 || uniform_f32 || uniform_q4 || uniform_vbit);
2045        if !fusable {
2046            for i in 0..N {
2047                ts[i].matvec2(x1, x2, o1s[i], o2s[i], pool);
2048            }
2049            return;
2050        }
2051        let pool = pool.unwrap();
2052
2053        if uniform_q4 || uniform_vbit {
2054            let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
2055            let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
2056            // q4/vbit share activation splits — no per-tensor col field.
2057            if a8w8_enabled() {
2058                let a1 = split_act(x1);
2059                let a2 = split_act(x2);
2060                let (a1, a2) = (&a1, &a2);
2061                if uniform_q4 {
2062                    let closures: [_; N] = std::array::from_fn(|i| {
2063                        let (packed, scales) =
2064                            q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
2065                        let (gpr, cols, o1, o2) =
2066                            (ts[i].cols() / GROUP_SIZE, ts[i].cols(), p1[i], p2[i]);
2067                        move |s: usize, e: usize| {
2068                            q4_range2_a8w8(packed, scales, gpr, cols, a1, a2, o1, o2, s, e)
2069                        }
2070                    });
2071                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2072                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2073                    pool.run_many(&parts);
2074                } else {
2075                    let closures: [_; N] = std::array::from_fn(|i| {
2076                        let Self::Mapped { vbit_offsets, .. } = ts[i] else {
2077                            unreachable!()
2078                        };
2079                        let (bytes, rows, cols, o1, o2) = (
2080                            ts[i].quant_bytes(),
2081                            ts[i].rows(),
2082                            ts[i].cols(),
2083                            p1[i],
2084                            p2[i],
2085                        );
2086                        move |s: usize, e: usize| {
2087                            vbit_range2_a8w8(
2088                                bytes,
2089                                vbit_offsets,
2090                                x1,
2091                                x2,
2092                                a1,
2093                                a2,
2094                                rows,
2095                                cols,
2096                                o1,
2097                                o2,
2098                                s,
2099                                e,
2100                            )
2101                        }
2102                    });
2103                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2104                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2105                    pool.run_many(&parts);
2106                }
2107                return;
2108            }
2109            if uniform_q4 {
2110                let closures: [_; N] = std::array::from_fn(|i| {
2111                    let (packed, scales) =
2112                        q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
2113                    let (gpr, o1, o2) = (ts[i].cols() / GROUP_SIZE, p1[i], p2[i]);
2114                    move |s: usize, e: usize| {
2115                        q4_range2_f32(packed, scales, gpr, x1, x2, o1, o2, s, e)
2116                    }
2117                });
2118                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2119                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2120                pool.run_many(&parts);
2121            } else {
2122                let closures: [_; N] = std::array::from_fn(|i| {
2123                    let Self::Mapped { vbit_offsets, .. } = ts[i] else {
2124                        unreachable!()
2125                    };
2126                    let (bytes, rows, cols, o1, o2) = (
2127                        ts[i].quant_bytes(),
2128                        ts[i].rows(),
2129                        ts[i].cols(),
2130                        p1[i],
2131                        p2[i],
2132                    );
2133                    move |s: usize, e: usize| {
2134                        vbit_range2_f32(bytes, vbit_offsets, x1, x2, rows, cols, o1, o2, s, e)
2135                    }
2136                });
2137                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2138                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2139                pool.run_many(&parts);
2140            }
2141            return;
2142        }
2143
2144        if uniform_f32 {
2145            let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
2146            let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
2147            let closures: [_; N] = std::array::from_fn(|i| {
2148                let Self::F32 { data, cols, .. } = ts[i] else {
2149                    unreachable!()
2150                };
2151                let (o1, o2) = (p1[i], p2[i]);
2152                move |start: usize, end: usize| {
2153                    for o in start..end {
2154                        let row = &data[o * cols..(o + 1) * cols];
2155                        let (mut s1, mut s2) = (0.0f32, 0.0f32);
2156                        for j in 0..*cols {
2157                            s1 += row[j] * x1[j];
2158                            s2 += row[j] * x2[j];
2159                        }
2160                        // SAFETY: disjoint (tensor, row) cells per worker.
2161                        unsafe {
2162                            *o1.at(o) = s1;
2163                            *o2.at(o) = s2;
2164                        }
2165                    }
2166                }
2167            });
2168            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2169                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2170            pool.run_many(&parts);
2171            return;
2172        }
2173
2174        struct Ctx<'a> {
2175            bytes: &'a [u8],
2176            row_scale: &'a [f32],
2177            cols: usize,
2178            xs1: std::borrow::Cow<'a, [f32]>,
2179            xs2: std::borrow::Cow<'a, [f32]>,
2180        }
2181        let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
2182            let Self::Mapped {
2183                dtype,
2184                cols,
2185                row_scale,
2186                col_field,
2187                ..
2188            } = ts[i]
2189            else {
2190                unreachable!()
2191            };
2192            Ctx {
2193                bytes: ts[i].quant_bytes(),
2194                row_scale,
2195                cols: *cols,
2196                xs1: prescale(x1, col_field, *dtype),
2197                xs2: prescale(x2, col_field, *dtype),
2198            }
2199        });
2200        let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
2201        let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
2202        #[cfg(target_arch = "aarch64")]
2203        if sdot_enabled() {
2204            let acts: [(SplitAct, SplitAct); N] =
2205                std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
2206            let closures: [_; N] = std::array::from_fn(|i| {
2207                let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
2208                move |start: usize, end: usize| {
2209                    q8_range2_sdot(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
2210                }
2211            });
2212            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2213                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2214            pool.run_many(&parts);
2215            return;
2216        }
2217        #[cfg(target_arch = "x86_64")]
2218        if avx2_a8w8_enabled() {
2219            let acts: [(SplitAct, SplitAct); N] =
2220                std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
2221            let closures: [_; N] = std::array::from_fn(|i| {
2222                let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
2223                move |start: usize, end: usize| {
2224                    q8_range2_avx2(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
2225                }
2226            });
2227            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2228                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2229            pool.run_many(&parts);
2230            return;
2231        }
2232        let closures: [_; N] = std::array::from_fn(|i| {
2233            let (c, o1, o2) = (&ctxs[i], p1[i], p2[i]);
2234            move |start: usize, end: usize| {
2235                q8_range2_f32(
2236                    c.bytes,
2237                    c.row_scale,
2238                    &c.xs1,
2239                    &c.xs2,
2240                    c.cols,
2241                    o1,
2242                    o2,
2243                    start,
2244                    end,
2245                )
2246            }
2247        });
2248        let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2249            std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2250        pool.run_many(&parts);
2251    }
2252
2253    /// Fused gate+up matvec with SiLU·mul: for each row r, computes
2254    /// `silu(gate·x) * (up·x)` and writes to `out[r]`. ONE pool dispatch,
2255    /// no intermediate g/u buffers, no separate silu pass. Falls back
2256    /// (returns false) for unsupported dtype combos.
2257    pub fn matvec_silu_mul(
2258        gate: &QTensor,
2259        up: &QTensor,
2260        x: &[f32],
2261        out: &mut [f32],
2262        pool: Option<&Pool>,
2263    ) -> bool {
2264        let inter = gate.rows();
2265        debug_assert_eq!(up.rows(), inter);
2266        debug_assert_eq!(out.len(), inter);
2267        debug_assert_eq!(gate.cols(), up.cols());
2268        if !a8w8_enabled() {
2269            return false;
2270        }
2271        let act = split_act(x);
2272        let act = &act;
2273        let x_ref = x;
2274        let out_addr = SendMut(out.as_mut_ptr());
2275
2276        match (gate, up) {
2277            // Q4Block gate + Q4Block up (most common mobile q4 models)
2278            (
2279                Self::Mapped {
2280                    dtype: TensorDtype::Q4Block,
2281                    ..
2282                },
2283                Self::Mapped {
2284                    dtype: TensorDtype::Q4Block,
2285                    ..
2286                },
2287            ) => {
2288                let (gp, gs) = q4_split(gate.quant_bytes(), gate.rows(), gate.cols());
2289                let (up_p, up_s) = q4_split(up.quant_bytes(), up.rows(), up.cols());
2290                let gpr = gate.cols() / GROUP_SIZE;
2291                let cols = gate.cols();
2292                let run = move |start: usize, end: usize| {
2293                    for r in start..end {
2294                        let mut gv = dot_q4_row_i8(gp, gs, r * gpr, gpr, &act.xq) * act.sx;
2295                        let mut uv = dot_q4_row_i8(up_p, up_s, r * gpr, gpr, &act.xq) * act.sx;
2296                        for &(j, xv) in &act.outliers {
2297                            let flat = r * cols + j;
2298                            let gb = gp[flat / 2];
2299                            let gn = if flat & 1 == 0 { gb & 0x0F } else { gb >> 4 };
2300                            let gsc = f16_to_f32(u16::from_le_bytes([
2301                                gs[(flat / GROUP_SIZE) * 2],
2302                                gs[(flat / GROUP_SIZE) * 2 + 1],
2303                            ]));
2304                            gv += ((gn as i32 - 8) as f32) * gsc * xv;
2305                            let ub = up_p[flat / 2];
2306                            let un = if flat & 1 == 0 { ub & 0x0F } else { ub >> 4 };
2307                            let usc = f16_to_f32(u16::from_le_bytes([
2308                                up_s[(flat / GROUP_SIZE) * 2],
2309                                up_s[(flat / GROUP_SIZE) * 2 + 1],
2310                            ]));
2311                            uv += ((un as i32 - 8) as f32) * usc * xv;
2312                        }
2313                        let silu_g = gv / (1.0 + (-gv).exp());
2314                        // SAFETY: disjoint row ranges per worker.
2315                        unsafe { *out_addr.at(r) = silu_g * uv };
2316                    }
2317                };
2318                dispatch_rows(pool, inter, &run);
2319                true
2320            }
2321            // Q4Tiled gate + Q4Tiled up — one row pass, both tile
2322            // streams sequential, silu·mul fused (same per-row math as
2323            // `q4t_matvec`).
2324            (
2325                Self::Mapped {
2326                    dtype: TensorDtype::Q4Tiled,
2327                    ..
2328                },
2329                Self::Mapped {
2330                    dtype: TensorDtype::Q4Tiled,
2331                    ..
2332                },
2333            ) => {
2334                let g_bytes = gate.quant_bytes();
2335                let u_bytes = up.quant_bytes();
2336                let gpr = gate.cols() / GROUP_SIZE;
2337                let run = move |start: usize, end: usize| {
2338                    for r in start..end {
2339                        let mut gv = dot_q4t_row_i8(g_bytes, r, gpr, &act.xq) * act.sx;
2340                        let mut uv = dot_q4t_row_i8(u_bytes, r, gpr, &act.xq) * act.sx;
2341                        for &(j, xv) in &act.outliers {
2342                            let (w, s) = q4t_outlier(g_bytes, r, gpr, j);
2343                            gv += w * s * xv;
2344                            let (w, s) = q4t_outlier(u_bytes, r, gpr, j);
2345                            uv += w * s * xv;
2346                        }
2347                        let silu_g = gv / (1.0 + (-gv).exp());
2348                        // SAFETY: disjoint row ranges per worker.
2349                        unsafe { *out_addr.at(r) = silu_g * uv };
2350                    }
2351                };
2352                dispatch_rows(pool, inter, &run);
2353                true
2354            }
2355            // Q4TiledP gate + Q4TiledP up — the same fused row pass, with
2356            // each row's two ladders built once and spent on both streams.
2357            (
2358                Self::Mapped {
2359                    dtype: TensorDtype::Q4TiledP,
2360                    ..
2361                },
2362                Self::Mapped {
2363                    dtype: TensorDtype::Q4TiledP,
2364                    ..
2365                },
2366            ) => {
2367                let cols = gate.cols();
2368                let gpr = cols / GROUP_SIZE;
2369                let gv_view = Q4tpView::new(gate.quant_bytes(), inter, cols);
2370                let uv_view = Q4tpView::new(up.quant_bytes(), inter, cols);
2371                let run = |start: usize, end: usize| {
2372                    let (mut gsc, mut usc) = (vec![0f32; gpr], vec![0f32; gpr]);
2373                    for r in start..end {
2374                        gv_view.scales_into(r, gpr, &mut gsc);
2375                        uv_view.scales_into(r, gpr, &mut usc);
2376                        let mut gv = dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, &gsc) * act.sx;
2377                        let mut uv = dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, &usc) * act.sx;
2378                        for &(j, xv) in &act.outliers {
2379                            let (w, s) = q4tp_outlier(gv_view.nib, r, gpr, j, &gsc);
2380                            gv += w * s * xv;
2381                            let (w, s) = q4tp_outlier(uv_view.nib, r, gpr, j, &usc);
2382                            uv += w * s * xv;
2383                        }
2384                        let silu_g = gv / (1.0 + (-gv).exp());
2385                        // SAFETY: disjoint row ranges per worker.
2386                        unsafe { *out_addr.at(r) = silu_g * uv };
2387                    }
2388                };
2389                dispatch_rows(pool, inter, &run);
2390                true
2391            }
2392            // Q1T gate + Q1T up
2393            (
2394                Self::Mapped {
2395                    dtype: TensorDtype::Q1T,
2396                    ..
2397                },
2398                Self::Mapped {
2399                    dtype: TensorDtype::Q1T,
2400                    ..
2401                },
2402            ) => {
2403                const TILE: usize = cortiq_core::quant::Q1T_TILE;
2404                let g_bytes = gate.quant_bytes();
2405                let u_bytes = up.quant_bytes();
2406                let gpr = gate.cols() / GROUP_SIZE;
2407                let (g_rp, g_ent, g_ov) = q1t_overlay(g_bytes, inter * gpr * TILE, inter);
2408                let (u_rp, u_ent, u_ov) = q1t_overlay(u_bytes, inter * gpr * TILE, inter);
2409                let run = move |start: usize, end: usize| {
2410                    for r in start..end {
2411                        let mut gv = q1t_dot_row_i8(g_bytes, r, gpr, &act.xq) * act.sx;
2412                        let mut uv = q1t_dot_row_i8(u_bytes, r, gpr, &act.xq) * act.sx;
2413                        for &(j, xv) in &act.outliers {
2414                            gv += q1t_base_weight(g_bytes, r, gpr, j) * xv;
2415                            uv += q1t_base_weight(u_bytes, r, gpr, j) * xv;
2416                        }
2417                        gv += q1t_row_outlier_correction(g_bytes, r, g_rp, g_ent, g_ov, x_ref);
2418                        uv += q1t_row_outlier_correction(u_bytes, r, u_rp, u_ent, u_ov, x_ref);
2419                        let silu_g = gv / (1.0 + (-gv).exp());
2420                        // SAFETY: disjoint row ranges per worker.
2421                        unsafe { *out_addr.at(r) = silu_g * uv };
2422                    }
2423                };
2424                dispatch_rows(pool, inter, &run);
2425                true
2426            }
2427            _ => false,
2428        }
2429    }
2430
2431    /// Every routed expert's fused gate/up/SiLU under ONE pool dispatch.
2432    ///
2433    /// The per-expert path pays a pool barrier per expert per stage: at 9
2434    /// experts over 40 layers that is ~720 barriers a token, and a decode
2435    /// profile of Qwen3.6-35B-A3B showed the pool parked in
2436    /// `psynch_cvwait` about twice as long as it spent computing. Laying
2437    /// every expert's rows end-to-end in one virtual row space collapses
2438    /// the stage to a single dispatch. The per-row body is the
2439    /// single-expert q4tp arm verbatim, so outputs are bit-identical.
2440    ///
2441    /// `false` = something is outside the fused q4tp kernel (dtype, shape,
2442    /// or the `CMF_SDOT=0` exact contract); the caller walks the ordinary
2443    /// per-expert path.
2444    pub fn moe_gate_up_many(
2445        pairs: &[(&QTensor, &QTensor)],
2446        x: &[f32],
2447        outs: &mut [Vec<f32>],
2448        pool: Option<&Pool>,
2449    ) -> bool {
2450        if pairs.is_empty() || pairs.len() != outs.len() || !a8w8_enabled() {
2451            return false;
2452        }
2453        let inter = pairs[0].0.rows();
2454        let cols = pairs[0].0.cols();
2455        if cols % GROUP_SIZE != 0 {
2456            return false;
2457        }
2458        let gpr = cols / GROUP_SIZE;
2459        let mut views = Vec::with_capacity(pairs.len() * 2);
2460        for ((g, u), o) in pairs.iter().zip(outs.iter()) {
2461            let both_q4tp = matches!(
2462                g,
2463                Self::Mapped {
2464                    dtype: TensorDtype::Q4TiledP,
2465                    ..
2466                }
2467            ) && matches!(
2468                u,
2469                Self::Mapped {
2470                    dtype: TensorDtype::Q4TiledP,
2471                    ..
2472                }
2473            );
2474            if !both_q4tp
2475                || g.rows() != inter
2476                || u.rows() != inter
2477                || g.cols() != cols
2478                || u.cols() != cols
2479                || o.len() != inter
2480            {
2481                return false;
2482            }
2483            views.push(Q4tpView::new(g.quant_bytes(), inter, cols));
2484            views.push(Q4tpView::new(u.quant_bytes(), inter, cols));
2485        }
2486        let act = split_act(x);
2487        let act = &act;
2488        let ptrs: Vec<SendMut> = outs.iter_mut().map(|o| SendMut(o.as_mut_ptr())).collect();
2489        let (views, ptrs) = (&views, &ptrs);
2490        let run = |start: usize, end: usize| {
2491            let (mut gsc, mut usc) = (vec![0f32; gpr], vec![0f32; gpr]);
2492            for flat in start..end {
2493                let (e, r) = (flat / inter, flat % inter);
2494                let gv_view = &views[e * 2];
2495                let uv_view = &views[e * 2 + 1];
2496                gv_view.scales_into(r, gpr, &mut gsc);
2497                uv_view.scales_into(r, gpr, &mut usc);
2498                let mut gv = dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, &gsc) * act.sx;
2499                let mut uv = dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, &usc) * act.sx;
2500                for &(j, xv) in &act.outliers {
2501                    let (w, s) = q4tp_outlier(gv_view.nib, r, gpr, j, &gsc);
2502                    gv += w * s * xv;
2503                    let (w, s) = q4tp_outlier(uv_view.nib, r, gpr, j, &usc);
2504                    uv += w * s * xv;
2505                }
2506                let silu_g = gv / (1.0 + (-gv).exp());
2507                // SAFETY: one worker owns each (expert, row) pair.
2508                unsafe { *ptrs[e].at(r) = silu_g * uv };
2509            }
2510        };
2511        dispatch_rows(pool, pairs.len() * inter, &run);
2512        true
2513    }
2514
2515    /// Every routed expert's down projection, weighted and summed into
2516    /// `out`, under ONE pool dispatch.
2517    ///
2518    /// Partitioned by OUTPUT row rather than by expert: each row is owned
2519    /// by a single worker, so the experts are summed in the caller's order
2520    /// — the same sequence of f32 adds the serial `out[i] += w·eo[i]` loop
2521    /// performs, hence bit-identical. Partitioning by expert instead would
2522    /// race on the shared accumulator.
2523    pub fn moe_down_many(
2524        downs: &[&QTensor],
2525        gs: &[Vec<f32>],
2526        weights: &[f32],
2527        out: &mut [f32],
2528        pool: Option<&Pool>,
2529    ) -> bool {
2530        if downs.is_empty()
2531            || downs.len() != gs.len()
2532            || downs.len() != weights.len()
2533            || !a8w8_enabled()
2534        {
2535            return false;
2536        }
2537        let rows = out.len();
2538        let cols = downs[0].cols();
2539        if cols % GROUP_SIZE != 0 {
2540            return false;
2541        }
2542        let gpr = cols / GROUP_SIZE;
2543        let mut views = Vec::with_capacity(downs.len());
2544        for (d, g) in downs.iter().zip(gs.iter()) {
2545            if !matches!(
2546                d,
2547                Self::Mapped {
2548                    dtype: TensorDtype::Q4TiledP,
2549                    ..
2550                }
2551            ) || d.rows() != rows
2552                || d.cols() != cols
2553                || g.len() != cols
2554            {
2555                return false;
2556            }
2557            views.push(Q4tpView::new(d.quant_bytes(), rows, cols));
2558        }
2559        // One int8 split per expert — the activation vectors differ.
2560        let acts: Vec<SplitAct> = gs.iter().map(|g| split_act(g)).collect();
2561        // Partitioned by OUTPUT row, with the experts folded inside: each
2562        // row is owned by one worker, so they are summed in the caller's
2563        // order — the same f32 sequence the serial `out[i] += w·eo[i]`
2564        // loop produces. Partitioning by expert instead would either race
2565        // on the accumulator or need a scratch plane and a second pass;
2566        // measured, that variant was a wash, so this keeps the simpler
2567        // shape.
2568        let out_addr = SendMut(out.as_mut_ptr());
2569        let (views, acts, weights) = (&views, &acts, &weights);
2570        let run = |start: usize, end: usize| {
2571            let mut sc = vec![0f32; gpr];
2572            for r in start..end {
2573                let mut acc = 0f32;
2574                for (e, v) in views.iter().enumerate() {
2575                    v.scales_into(r, gpr, &mut sc);
2576                    let a = &acts[e];
2577                    let mut d = dot_q4tp_row_i8(v.nib, r, gpr, &a.xq, &sc) * a.sx;
2578                    for &(j, xv) in &a.outliers {
2579                        let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
2580                        d += w * s * xv;
2581                    }
2582                    acc += weights[e] * d;
2583                }
2584                // SAFETY: disjoint row ranges per worker.
2585                unsafe { *out_addr.at(r) = acc };
2586            }
2587        };
2588        dispatch_rows(pool, rows, &run);
2589        true
2590    }
2591}
2592
2593/// Batched q8 kernel: same math as qmatvec, the row makes a single
2594/// pass from memory for the whole batch.
2595/// Accelerate CBLAS — the Apple AMX matrix units, the same engine
2596/// llama.cpp's `-ngl 0` prefill rides via ggml-blas.
2597#[cfg(target_os = "macos")]
2598mod accel_blas {
2599    #[link(name = "Accelerate", kind = "framework")]
2600    unsafe extern "C" {
2601        pub fn cblas_sgemm(
2602            order: i32,
2603            trans_a: i32,
2604            trans_b: i32,
2605            m: i32,
2606            n: i32,
2607            k: i32,
2608            alpha: f32,
2609            a: *const f32,
2610            lda: i32,
2611            b: *const f32,
2612            ldb: i32,
2613            beta: f32,
2614            c: *mut f32,
2615            ldc: i32,
2616        );
2617    }
2618}
2619
2620#[cfg(target_os = "macos")]
2621pub(crate) fn accel_gemm_enabled() -> bool {
2622    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2623    *ON.get_or_init(|| std::env::var("CMF_ACCEL").map(|v| v != "0").unwrap_or(true))
2624}
2625
2626/// Off macOS the "accel" GEMM is the portable NEON micro-kernel below —
2627/// same entry point, so the batched-attention path opens on mobile.
2628#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
2629pub(crate) fn accel_gemm_enabled() -> bool {
2630    true
2631}
2632
2633/// Portable NEON f32 GEMM (row-major, optional Bᵀ): a 4×8 fmla
2634/// micro-kernel with A broadcast against B panels — the mobile stand-in
2635/// for Accelerate in the batched causal attention (QKᵀ and P·V). Not a
2636/// BLAS: shapes here are the attention panels (m ≤ heads·chunk,
2637/// k = head_dim or context), and the goal is removing the per-position
2638/// quadratic wall, not peak GEMM.
2639#[cfg(target_arch = "aarch64")]
2640#[allow(clippy::too_many_arguments)]
2641pub(crate) fn neon_gemm_rm(
2642    m: usize,
2643    n: usize,
2644    k: usize,
2645    alpha: f32,
2646    a: &[f32],
2647    lda: usize,
2648    b_mat: &[f32],
2649    ldb: usize,
2650    b_rows_are_n: bool,
2651    c: &mut [f32],
2652    ldc: usize,
2653) {
2654    debug_assert!(a.len() >= (m - 1) * lda + k);
2655    debug_assert!(c.len() >= (m - 1) * ldc + n);
2656    // SAFETY: bounds asserted above; NEON is baseline on aarch64.
2657    unsafe {
2658        use core::arch::aarch64::*;
2659        let mut i = 0usize;
2660        while i < m {
2661            let mi = (m - i).min(4);
2662            let mut j = 0usize;
2663            while j < n {
2664                let nj = (n - j).min(8);
2665                if mi == 4 && nj == 8 {
2666                    let (mut c0a, mut c0b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
2667                    let (mut c1a, mut c1b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
2668                    let (mut c2a, mut c2b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
2669                    let (mut c3a, mut c3b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
2670                    for p in 0..k {
2671                        let (b0, b1) = if b_rows_are_n {
2672                            // B is [n, k]: column p of Bᵀ = element p of
2673                            // eight consecutive B rows — gathered.
2674                            let base = b_mat.as_ptr().add(j * ldb + p);
2675                            let g = |o: usize| *base.add(o * ldb);
2676                            ([g(0), g(1), g(2), g(3)], [g(4), g(5), g(6), g(7)])
2677                        } else {
2678                            let base = b_mat.as_ptr().add(p * ldb + j);
2679                            (
2680                                [*base, *base.add(1), *base.add(2), *base.add(3)],
2681                                [*base.add(4), *base.add(5), *base.add(6), *base.add(7)],
2682                            )
2683                        };
2684                        let bv0 = vld1q_f32(b0.as_ptr());
2685                        let bv1 = vld1q_f32(b1.as_ptr());
2686                        let a0 = vdupq_n_f32(*a.as_ptr().add(i * lda + p));
2687                        let a1 = vdupq_n_f32(*a.as_ptr().add((i + 1) * lda + p));
2688                        let a2 = vdupq_n_f32(*a.as_ptr().add((i + 2) * lda + p));
2689                        let a3 = vdupq_n_f32(*a.as_ptr().add((i + 3) * lda + p));
2690                        c0a = vfmaq_f32(c0a, a0, bv0);
2691                        c0b = vfmaq_f32(c0b, a0, bv1);
2692                        c1a = vfmaq_f32(c1a, a1, bv0);
2693                        c1b = vfmaq_f32(c1b, a1, bv1);
2694                        c2a = vfmaq_f32(c2a, a2, bv0);
2695                        c2b = vfmaq_f32(c2b, a2, bv1);
2696                        c3a = vfmaq_f32(c3a, a3, bv0);
2697                        c3b = vfmaq_f32(c3b, a3, bv1);
2698                    }
2699                    let al = vdupq_n_f32(alpha);
2700                    for (r, (ca, cb)) in [(c0a, c0b), (c1a, c1b), (c2a, c2b), (c3a, c3b)]
2701                        .iter()
2702                        .enumerate()
2703                    {
2704                        let dst = c.as_mut_ptr().add((i + r) * ldc + j);
2705                        vst1q_f32(dst, vmulq_f32(*ca, al));
2706                        vst1q_f32(dst.add(4), vmulq_f32(*cb, al));
2707                    }
2708                } else {
2709                    for r in 0..mi {
2710                        for q in 0..nj {
2711                            let mut acc = 0f32;
2712                            for p in 0..k {
2713                                let bv = if b_rows_are_n {
2714                                    b_mat[(j + q) * ldb + p]
2715                                } else {
2716                                    b_mat[p * ldb + j + q]
2717                                };
2718                                acc += a[(i + r) * lda + p] * bv;
2719                            }
2720                            c[(i + r) * ldc + j + q] = acc * alpha;
2721                        }
2722                    }
2723                }
2724                j += nj;
2725            }
2726            i += mi;
2727        }
2728    }
2729}
2730
2731/// Off-macOS aarch64: the batched attention rides the NEON micro-GEMM.
2732#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
2733#[allow(clippy::too_many_arguments)]
2734pub(crate) fn sgemm_rm(
2735    m: usize,
2736    n: usize,
2737    k: usize,
2738    alpha: f32,
2739    a: &[f32],
2740    lda: usize,
2741    b_mat: &[f32],
2742    ldb: usize,
2743    b_rows_are_n: bool,
2744    c: &mut [f32],
2745    ldc: usize,
2746) {
2747    neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
2748}
2749
2750/// Row-major f32 GEMM, exposed for offline tools (the AWNP pass builds a
2751/// per-layer projection and applies it to every expert; a naive triple loop
2752/// would turn a two-minute job into half an hour).
2753#[allow(clippy::too_many_arguments)]
2754pub fn sgemm_public(
2755    m: usize,
2756    n: usize,
2757    k: usize,
2758    alpha: f32,
2759    a: &[f32],
2760    lda: usize,
2761    b_mat: &[f32],
2762    ldb: usize,
2763    b_rows_are_n: bool,
2764    c: &mut [f32],
2765    ldc: usize,
2766) {
2767    #[cfg(any(target_os = "macos", target_arch = "aarch64"))]
2768    {
2769        sgemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
2770    }
2771    // x86 without Accelerate has no sgemm_rm: the specialized paths there are
2772    // quantized kernels, not an f32 GEMM. Only the offline AWNP pass reaches
2773    // this, so correctness matters and throughput does not — a triple loop is
2774    // the honest fallback rather than a reason to make the tool macOS-only.
2775    #[cfg(not(any(target_os = "macos", target_arch = "aarch64")))]
2776    {
2777        for i in 0..m {
2778            for j in 0..n {
2779                let mut acc = 0f32;
2780                for p in 0..k {
2781                    let bv = if b_rows_are_n {
2782                        b_mat[j * ldb + p]
2783                    } else {
2784                        b_mat[p * ldb + j]
2785                    };
2786                    acc += a[i * lda + p] * bv;
2787                }
2788                c[i * ldc + j] = alpha * acc;
2789            }
2790        }
2791    }
2792}
2793
2794/// Row-major f32 GEMM on Accelerate: C[m,n] = alpha·A[m,k] × B(ᵀ).
2795/// `b_rows_are_n` = true multiplies by Bᵀ where B is stored [n, k].
2796#[cfg(target_os = "macos")]
2797#[allow(clippy::too_many_arguments)]
2798pub(crate) fn sgemm_rm(
2799    m: usize,
2800    n: usize,
2801    k: usize,
2802    alpha: f32,
2803    a: &[f32],
2804    lda: usize,
2805    b_mat: &[f32],
2806    ldb: usize,
2807    b_rows_are_n: bool,
2808    c: &mut [f32],
2809    ldc: usize,
2810) {
2811    debug_assert!(a.len() >= (m - 1) * lda + k);
2812    debug_assert!(c.len() >= (m - 1) * ldc + n);
2813    // Test hook: route the attention GEMMs through the portable NEON
2814    // micro-kernel ON APPLE SILICON — how the mobile batched attend is
2815    // measured without a phone in the loop. (Intel macOS has no NEON —
2816    // the hook is a no-op there, Accelerate continues below.)
2817    #[cfg(target_arch = "aarch64")]
2818    if std::env::var("CMF_FORCE_NEON_GEMM")
2819        .map(|v| v == "1")
2820        .unwrap_or(false)
2821    {
2822        return neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
2823    }
2824    unsafe {
2825        accel_blas::cblas_sgemm(
2826            101, // RowMajor
2827            111, // NoTrans A
2828            if b_rows_are_n { 112 } else { 111 },
2829            m as i32,
2830            n as i32,
2831            k as i32,
2832            alpha,
2833            a.as_ptr(),
2834            lda as i32,
2835            b_mat.as_ptr(),
2836            ldb as i32,
2837            0.0,
2838            c.as_mut_ptr(),
2839            ldc as i32,
2840        );
2841    }
2842}
2843
2844/// Prefill GEMM through Accelerate (macOS): dequantize q8 rows into
2845/// f32 tiles (scale folded in, pool-parallel) and multiply each tile
2846/// on the AMX with one row-major sgemm. Tiles live in cache, weights
2847/// stream once. Numerics are f32-GEMM (not the int8 dot): prefill
2848/// logits shift within f32 rounding — tolerance-class, like every
2849/// reduction-order change; decode (M=1) never takes this path.
2850#[cfg(target_os = "macos")]
2851fn qmatmat_accel(
2852    q: &[u8],
2853    row_scale: &[f32],
2854    pre: &[std::borrow::Cow<'_, [f32]>],
2855    rows: usize,
2856    cols: usize,
2857    out: &mut [f32],
2858    pool: Option<&Pool>,
2859) {
2860    // NOTE: double-buffering the dequant against the sgemm (a scoped
2861    // thread driving the pool on tile k+1 while the caller multiplies
2862    // tile k) was tried and LOST ~6%: Accelerate's sgemm is itself
2863    // multithreaded, and the dequant workers just steal its cores.
2864    const TR: usize = 2048;
2865    let b = pre.len();
2866    thread_local! {
2867        static XPANEL: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
2868        static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
2869    }
2870    XPANEL.with(|xp| {
2871        WTILE.with(|wt| {
2872            let mut xpanel = xp.borrow_mut();
2873            xpanel.clear();
2874            for x in pre {
2875                xpanel.extend_from_slice(x);
2876            }
2877            let mut wtile = wt.borrow_mut();
2878            wtile.resize(TR * cols, 0.0);
2879            let mut r0 = 0usize;
2880            while r0 < rows {
2881                let tr = TR.min(rows - r0);
2882                // Dequant the tile (scale folded) — pool-parallel.
2883                let wt_addr = SendMut(wtile.as_mut_ptr());
2884                let run = |start: usize, end: usize| {
2885                    for r in start..end {
2886                        let row = &q[(r0 + r) * cols..(r0 + r + 1) * cols];
2887                        let s = row_scale[r0 + r];
2888                        // SAFETY: workers cover disjoint r ranges.
2889                        let dst =
2890                            unsafe { std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols) };
2891                        for (d, &v) in dst.iter_mut().zip(row) {
2892                            *d = (v as i8) as f32 * s;
2893                        }
2894                    }
2895                };
2896                dispatch_rows(pool, tr, &run);
2897                // C[b, tr] (at column r0 of out[b, rows]) = X · Wtileᵀ
2898                unsafe {
2899                    accel_blas::cblas_sgemm(
2900                        101, // RowMajor
2901                        111, // NoTrans A
2902                        112, // Trans B
2903                        b as i32,
2904                        tr as i32,
2905                        cols as i32,
2906                        1.0,
2907                        xpanel.as_ptr(),
2908                        cols as i32,
2909                        wtile.as_ptr(),
2910                        cols as i32,
2911                        0.0,
2912                        out.as_mut_ptr().add(r0),
2913                        rows as i32,
2914                    );
2915                }
2916                r0 += tr;
2917            }
2918        })
2919    });
2920}
2921
2922fn qmatmat(
2923    q: &[u8],
2924    row_scale: &[f32],
2925    pre: &[std::borrow::Cow<'_, [f32]>],
2926    rows: usize,
2927    cols: usize,
2928    out: &mut [f32],
2929    pool: Option<&Pool>,
2930) {
2931    let b = pre.len();
2932    debug_assert_eq!(out.len(), b * rows);
2933    // Big prefill batches ride the AMX (roadmap PR3): the row×batch
2934    // SDOT loop below peaks near the CPU's dot throughput, an order
2935    // below the matrix units. Small tensors and tiny test models stay
2936    // on the exact integer path.
2937    #[cfg(target_os = "macos")]
2938    if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
2939        qmatmat_accel(q, row_scale, pre, rows, cols, out, pool);
2940        return;
2941    }
2942    #[cfg(target_arch = "aarch64")]
2943    if sdot_enabled() {
2944        let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
2945        let out_addr = SendMut(out.as_mut_ptr());
2946        // Blocked 2×4 (mobile prefill: no AMX to fall back on — this
2947        // path IS the ARM prefill GEMM off Apple silicon).
2948        let blocked_ok = blocked_enabled();
2949        let use_i8mm = i8mm_enabled();
2950        if blocked_ok {
2951            let run = |start: usize, end: usize| {
2952                let mut o = start;
2953                while o < end {
2954                    if o + 2 <= end {
2955                        let r0 = &q[o * cols..(o + 1) * cols];
2956                        let r1 = &q[(o + 1) * cols..(o + 2) * cols];
2957                        let mut bi = 0usize;
2958                        while bi + 4 <= acts.len() {
2959                            let xs = [
2960                                acts[bi].xq.as_slice(),
2961                                acts[bi + 1].xq.as_slice(),
2962                                acts[bi + 2].xq.as_slice(),
2963                                acts[bi + 3].xq.as_slice(),
2964                            ];
2965                            let d = if use_i8mm {
2966                                unsafe { dot_i8_smmla_2x4(r0, r1, xs) }
2967                            } else {
2968                                unsafe { dot_i8_sdot_2x4(r0, r1, xs) }
2969                            };
2970                            for (r, row) in [r0, r1].into_iter().enumerate() {
2971                                for k in 0..4 {
2972                                    let act = &acts[bi + k];
2973                                    let mut v = d[r][k] as f32 * act.sx;
2974                                    for &(j, xv) in &act.outliers {
2975                                        v += (row[j] as i8) as f32 * xv;
2976                                    }
2977                                    unsafe {
2978                                        *out_addr.at((bi + k) * rows + o + r) = v * row_scale[o + r]
2979                                    };
2980                                }
2981                            }
2982                            bi += 4;
2983                        }
2984                        while bi < acts.len() {
2985                            for (r, row) in [r0, r1].into_iter().enumerate() {
2986                                let v = row_dot_sdot(row, &acts[bi]) * row_scale[o + r];
2987                                unsafe { *out_addr.at(bi * rows + o + r) = v };
2988                            }
2989                            bi += 1;
2990                        }
2991                        o += 2;
2992                    } else {
2993                        let row = &q[o * cols..(o + 1) * cols];
2994                        for (bi, act) in acts.iter().enumerate() {
2995                            let v = row_dot_sdot(row, act) * row_scale[o];
2996                            unsafe { *out_addr.at(bi * rows + o) = v };
2997                        }
2998                        o += 1;
2999                    }
3000                }
3001            };
3002            dispatch_rows(pool, rows, &run);
3003            return;
3004        }
3005        let run = |start: usize, end: usize| {
3006            for o in start..end {
3007                let row = &q[o * cols..(o + 1) * cols];
3008                for (bi, act) in acts.iter().enumerate() {
3009                    let v = row_dot_sdot(row, act) * row_scale[o];
3010                    unsafe { *out_addr.at(bi * rows + o) = v };
3011                }
3012            }
3013        };
3014        dispatch_rows(pool, rows, &run);
3015        return;
3016    }
3017    // x86 A8W8 batch. Non-VNNI parts take the BLOCKED 2×4 kernel
3018    // (roadmap P0: two weight rows' abs() stay in registers across four
3019    // activation streams); VNNI machines keep the per-row bias-trick
3020    // dot, which is already throughput-bound there.
3021    #[cfg(target_arch = "x86_64")]
3022    if avx2_a8w8_enabled() {
3023        let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
3024        let out_addr = SendMut(out.as_mut_ptr());
3025        // CMF_X86_BLOCKED=0 forces the per-row path (paired in-process
3026        // A/B on noisy shared-vCPU hosts).
3027        let blocked_ok = blocked_enabled();
3028        if !avx512vnni_enabled() && blocked_ok {
3029            let run = |start: usize, end: usize| {
3030                let mut o = start;
3031                while o < end {
3032                    if o + 2 <= end {
3033                        let r0 = &q[o * cols..(o + 1) * cols];
3034                        let r1 = &q[(o + 1) * cols..(o + 2) * cols];
3035                        let mut bi = 0usize;
3036                        while bi + 4 <= acts.len() {
3037                            let xs = [
3038                                acts[bi].xq.as_slice(),
3039                                acts[bi + 1].xq.as_slice(),
3040                                acts[bi + 2].xq.as_slice(),
3041                                acts[bi + 3].xq.as_slice(),
3042                            ];
3043                            let d = unsafe { dot_i8_i8_avx2_2x4(r0, r1, xs) };
3044                            for (r, row) in [r0, r1].into_iter().enumerate() {
3045                                for k in 0..4 {
3046                                    let act = &acts[bi + k];
3047                                    let mut v = d[r][k] as f32 * act.sx;
3048                                    for &(j, xv) in &act.outliers {
3049                                        v += (row[j] as i8) as f32 * xv;
3050                                    }
3051                                    unsafe {
3052                                        *out_addr.at((bi + k) * rows + o + r) = v * row_scale[o + r]
3053                                    };
3054                                }
3055                            }
3056                            bi += 4;
3057                        }
3058                        while bi < acts.len() {
3059                            for (r, row) in [r0, r1].into_iter().enumerate() {
3060                                let v = row_dot_avx2(row, &acts[bi]) * row_scale[o + r];
3061                                unsafe { *out_addr.at(bi * rows + o + r) = v };
3062                            }
3063                            bi += 1;
3064                        }
3065                        o += 2;
3066                    } else {
3067                        let row = &q[o * cols..(o + 1) * cols];
3068                        for (bi, act) in acts.iter().enumerate() {
3069                            let v = row_dot_avx2(row, act) * row_scale[o];
3070                            unsafe { *out_addr.at(bi * rows + o) = v };
3071                        }
3072                        o += 1;
3073                    }
3074                }
3075            };
3076            dispatch_rows(pool, rows, &run);
3077            return;
3078        }
3079        let run = |start: usize, end: usize| {
3080            for o in start..end {
3081                let row = &q[o * cols..(o + 1) * cols];
3082                for (bi, act) in acts.iter().enumerate() {
3083                    let v = row_dot_avx2(row, act) * row_scale[o];
3084                    unsafe { *out_addr.at(bi * rows + o) = v };
3085                }
3086            }
3087        };
3088        dispatch_rows(pool, rows, &run);
3089        return;
3090    }
3091    let out_addr = SendMut(out.as_mut_ptr());
3092    let run = |start: usize, end: usize| {
3093        for o in start..end {
3094            let row = &q[o * cols..(o + 1) * cols];
3095            for (bi, x) in pre.iter().enumerate() {
3096                let mut acc = 0f32;
3097                for j in 0..cols {
3098                    acc += (row[j] as i8) as f32 * x[j];
3099                }
3100                unsafe { *out_addr.at(bi * rows + o) = acc * row_scale[o] };
3101            }
3102        }
3103    };
3104    dispatch_rows(pool, rows, &run);
3105}
3106
3107/// Split rows across pool workers (shared qmatvec pattern). Self-balancing
3108/// — see `Pool::run_rows` for why a static 1/n split is wrong here.
3109fn dispatch_rows(pool: Option<&Pool>, rows: usize, run: &(dyn Fn(usize, usize) + Sync)) {
3110    match pool {
3111        Some(pool) if rows >= 256 => pool.run_rows(rows, run),
3112        _ => run(0, rows),
3113    }
3114}
3115
3116/// Split a q4_block blob into (packed nibbles, f16 group scales).
3117fn q4_split(bytes: &[u8], rows: usize, cols: usize) -> (&[u8], &[u8]) {
3118    let groups = rows * cols / GROUP_SIZE;
3119    bytes.split_at(groups * 16)
3120}
3121
3122/// SIMD unpack for the dominant vbit width B=4 (94% of rows on the
3123/// log2-shape calibration): 16 packed bytes -> 32 centered i8 values.
3124/// vbit packs MSB-first, so the HIGH nibble is the even element
3125/// (opposite of q4_block's lo-first interleave). Centering is u-7.
3126#[inline]
3127fn vbit_fill4(data: &[u8], buf: &mut [u8]) {
3128    #[cfg(target_arch = "aarch64")]
3129    unsafe {
3130        return vbit_fill4_neon(data, buf);
3131    }
3132    #[cfg(target_arch = "x86_64")]
3133    if avx2_enabled() {
3134        return unsafe { vbit_fill4_avx2(data, buf) };
3135    }
3136    #[allow(unreachable_code)]
3137    for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
3138        let u = unpack8::<4>(&data[blk * 4..]);
3139        for k in 0..8 {
3140            chunk[k] = (u[k] - 7) as i8 as u8;
3141        }
3142    }
3143}
3144
3145#[cfg(target_arch = "aarch64")]
3146#[target_feature(enable = "neon")]
3147unsafe fn vbit_fill4_neon(data: &[u8], buf: &mut [u8]) {
3148    // SAFETY: buf.len() is a multiple of GROUP_SIZE=32; data holds
3149    // buf.len()/2 packed bytes (validated at load).
3150    unsafe {
3151        use core::arch::aarch64::*;
3152        let n = buf.len();
3153        let mask = vdupq_n_u8(0x0F);
3154        let seven = vdupq_n_s8(7);
3155        let mut g = 0usize;
3156        while g * 32 + 32 <= n {
3157            let b = vld1q_u8(data.as_ptr().add(g * 16));
3158            let hi = vshrq_n_u8::<4>(b);
3159            let lo = vandq_u8(b, mask);
3160            let z0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(hi, lo)), seven);
3161            let z1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(hi, lo)), seven);
3162            vst1q_u8(buf.as_mut_ptr().add(g * 32), vreinterpretq_u8_s8(z0));
3163            vst1q_u8(buf.as_mut_ptr().add(g * 32 + 16), vreinterpretq_u8_s8(z1));
3164            g += 1;
3165        }
3166    }
3167}
3168
3169#[cfg(target_arch = "x86_64")]
3170#[target_feature(enable = "avx2")]
3171unsafe fn vbit_fill4_avx2(data: &[u8], buf: &mut [u8]) {
3172    // SAFETY: see vbit_fill4_neon.
3173    unsafe {
3174        use core::arch::x86_64::*;
3175        let n = buf.len();
3176        let mask = _mm_set1_epi8(0x0F);
3177        let seven = _mm256_set1_epi8(7);
3178        let mut g = 0usize;
3179        while g * 32 + 32 <= n {
3180            let b = _mm_loadu_si128(data.as_ptr().add(g * 16) as *const __m128i);
3181            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), mask);
3182            let lo = _mm_and_si128(b, mask);
3183            let z = _mm256_sub_epi8(
3184                _mm256_set_m128i(_mm_unpackhi_epi8(hi, lo), _mm_unpacklo_epi8(hi, lo)),
3185                seven,
3186            );
3187            _mm256_storeu_si256(buf.as_mut_ptr().add(g * 32) as *mut __m256i, z);
3188            g += 1;
3189        }
3190    }
3191}
3192
3193/// Unpack 8 MSB-first B-bit values from exactly B bytes (fixed shifts —
3194/// no serial bit-buffer, auto-vectorizable). Every 32-value group starts
3195/// byte-aligned (32·B/8 is integral for B∈3..8), so groups decompose
3196/// into 4 such blocks.
3197#[inline(always)]
3198fn unpack8<const B: usize>(data: &[u8]) -> [i32; 8] {
3199    let mut acc = 0u64;
3200    for i in 0..B {
3201        acc = (acc << 8) | data[i] as u64;
3202    }
3203    let mask = (1u64 << B) - 1;
3204    let mut out = [0i32; 8];
3205    for (k, o) in out.iter_mut().enumerate() {
3206        *o = ((acc >> ((7 - k) * B)) & mask) as i32;
3207    }
3208    out
3209}
3210
3211/// Fused vbit matvec straight from the mapped bytes (spec §3, P13
3212/// FIG.3): [u8 bits: rows][f16 scales: rows·cols/32][bit-packed rows,
3213/// MSB-first, byte-padded]. Row data offsets are precomputed at load
3214/// (`vbit_row_offsets`) — the per-call prefix scan was O(rows) pure
3215/// overhead on every matvec.
3216#[allow(clippy::too_many_arguments)]
3217fn vbitmatvec(
3218    bytes: &[u8],
3219    offsets: &[usize],
3220    x: &[f32],
3221    rows: usize,
3222    cols: usize,
3223    out: &mut [f32],
3224    pool: Option<&Pool>,
3225) {
3226    debug_assert_eq!(out.len(), rows);
3227    debug_assert_eq!(offsets.len(), rows + 1);
3228
3229    // SDOT path: unpack the row to centered i8 once, then per-group
3230    // int8 dot against the quantized activations — same A8W8 contract
3231    // as q8 (bounded noise; CMF_SDOT=0 keeps the exact scalar path).
3232    if a8w8_enabled() {
3233        let act = split_act(x);
3234        let out_addr = SendMut(out.as_mut_ptr());
3235        let run = move |start: usize, end: usize| {
3236            vbit_range_a8w8(bytes, offsets, x, &act, rows, cols, out_addr, start, end)
3237        };
3238        dispatch_rows(pool, rows, &run);
3239        return;
3240    }
3241
3242    let out_addr = SendMut(out.as_mut_ptr());
3243    let run = move |start: usize, end: usize| {
3244        vbit_range_f32(bytes, offsets, x, rows, cols, out_addr, start, end)
3245    };
3246    dispatch_rows(pool, rows, &run);
3247}
3248
3249/// One vbit row range via the A8W8 int8 path — kernel body of
3250/// `vbitmatvec`, extracted so multi-matrix jobs can drive it for
3251/// several tensors in one dispatch (b=8 rows go exact f32).
3252#[allow(clippy::too_many_arguments)]
3253fn vbit_range_a8w8(
3254    bytes: &[u8],
3255    offsets: &[usize],
3256    x: &[f32],
3257    act: &SplitAct,
3258    rows: usize,
3259    cols: usize,
3260    out: SendMut,
3261    start: usize,
3262    end: usize,
3263) {
3264    let ng = cols / GROUP_SIZE;
3265    let bits = &bytes[..rows];
3266    let sc_off = rows;
3267    let row_dot = |r: usize| -> f32 {
3268        let b = bits[r] as usize;
3269        let l = (1i32 << (b - 1)) - 1;
3270        let mask = (1u64 << b) - 1;
3271        let data = &bytes[offsets[r]..offsets[r + 1]];
3272        if b == 8 {
3273            // u−L reaches 128 → does not fit i8; exact f32 path.
3274            let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
3275            let mut dot = 0f32;
3276            for g in 0..ng {
3277                let so = (r * ng + g) * 2;
3278                let sgf = f16_to_f32(u16::from_le_bytes([
3279                    bytes[sc_off + so],
3280                    bytes[sc_off + so + 1],
3281                ]));
3282                let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
3283                let mut gd = 0f32;
3284                for &xv in xg.iter() {
3285                    if nbits < 8 {
3286                        acc = (acc << 8) | data[idx] as u64;
3287                        idx += 1;
3288                        nbits += 8;
3289                    }
3290                    let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
3291                    nbits -= 8;
3292                    gd += (u - l) as f32 * xv;
3293                }
3294                dot += gd * sgf;
3295            }
3296            return dot;
3297        }
3298        // Per-worker scratch: this closure runs for every row of the
3299        // tensor (lm_head ≈ 150k rows/token) — a heap allocation per
3300        // row was measurable pure overhead.
3301        thread_local! {
3302            static VBIT_SCRATCH: std::cell::RefCell<Vec<u8>> =
3303                const { std::cell::RefCell::new(Vec::new()) };
3304        }
3305        #[inline(always)]
3306        fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
3307            for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
3308                let u = unpack8::<B>(&data[blk * B..]);
3309                for k in 0..8 {
3310                    chunk[k] = (u[k] - l) as i8 as u8;
3311                }
3312            }
3313        }
3314        let _ = mask;
3315        VBIT_SCRATCH.with(|scratch| {
3316            let mut buf = scratch.borrow_mut();
3317            buf.resize(cols, 0);
3318            match b {
3319                3 => fill::<3>(data, l, &mut buf),
3320                4 => vbit_fill4(data, &mut buf),
3321                5 => fill::<5>(data, l, &mut buf),
3322                6 => fill::<6>(data, l, &mut buf),
3323                _ => unreachable!(),
3324            }
3325            let mut dot = 0f32;
3326            for g in 0..ng {
3327                let so = (r * ng + g) * 2;
3328                let s = f16_to_f32(u16::from_le_bytes([
3329                    bytes[sc_off + so],
3330                    bytes[sc_off + so + 1],
3331                ]));
3332                let d = dot_i8_i8(
3333                    &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
3334                    &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
3335                ) as f32
3336                    * act.sx;
3337                dot += d * s;
3338            }
3339            for &(j, xv) in &act.outliers {
3340                let so = (r * ng + j / GROUP_SIZE) * 2;
3341                let s = f16_to_f32(u16::from_le_bytes([
3342                    bytes[sc_off + so],
3343                    bytes[sc_off + so + 1],
3344                ]));
3345                // xq is zeroed at outlier slots — add the exact term.
3346                dot += (buf[j] as i8) as f32 * s * xv;
3347            }
3348            dot
3349        })
3350    };
3351    for r in start..end {
3352        // SAFETY: disjoint row ranges per worker.
3353        unsafe { *out.at(r) = row_dot(r) };
3354    }
3355}
3356
3357/// Exact scalar vbit row range (same extraction, non-SDOT path).
3358#[allow(clippy::too_many_arguments)]
3359fn vbit_range_f32(
3360    bytes: &[u8],
3361    offsets: &[usize],
3362    x: &[f32],
3363    rows: usize,
3364    cols: usize,
3365    out: SendMut,
3366    start: usize,
3367    end: usize,
3368) {
3369    let ng = cols / GROUP_SIZE;
3370    let bits = &bytes[..rows];
3371    let sc_off = rows;
3372    // Per-bit-width specialized inner loops: the compiler unrolls the
3373    // constant shifts (the generic bit-buffer loop was branch-bound —
3374    // 5.6 vs 13.2 tok/s q4 on the 0.8B).
3375    #[inline(always)]
3376    fn dot_row<const B: usize>(
3377        data: &[u8],
3378        bytes: &[u8],
3379        sc_off: usize,
3380        r: usize,
3381        ng: usize,
3382        x: &[f32],
3383    ) -> f32 {
3384        let l = ((1i32 << (B - 1)) - 1) as f32;
3385        let gbytes = GROUP_SIZE * B / 8;
3386        let mut dot = 0f32;
3387        for g in 0..ng {
3388            let so = (r * ng + g) * 2;
3389            let s = f16_to_f32(u16::from_le_bytes([
3390                bytes[sc_off + so],
3391                bytes[sc_off + so + 1],
3392            ]));
3393            let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
3394            let gd0 = &data[g * gbytes..(g + 1) * gbytes];
3395            let mut gd = 0f32;
3396            for blk in 0..GROUP_SIZE / 8 {
3397                let u = unpack8::<B>(&gd0[blk * B..]);
3398                let xb = &xg[blk * 8..blk * 8 + 8];
3399                for k in 0..8 {
3400                    gd += (u[k] as f32 - l) * xb[k];
3401                }
3402            }
3403            dot += gd * s;
3404        }
3405        dot
3406    }
3407    for r in start..end {
3408        let data = &bytes[offsets[r]..offsets[r + 1]];
3409        let v = match bits[r] {
3410            3 => dot_row::<3>(data, bytes, sc_off, r, ng, x),
3411            4 => dot_row::<4>(data, bytes, sc_off, r, ng, x),
3412            5 => dot_row::<5>(data, bytes, sc_off, r, ng, x),
3413            6 => dot_row::<6>(data, bytes, sc_off, r, ng, x),
3414            8 => dot_row::<8>(data, bytes, sc_off, r, ng, x),
3415            b => unreachable!("vbit bit-width {b} (validated at load)"),
3416        };
3417        // SAFETY: disjoint row ranges per worker.
3418        unsafe { *out.at(r) = v };
3419    }
3420}
3421
3422/// Fused two-input vbit matvec: each row is unpacked from the mmap ONCE
3423/// and dotted against BOTH activations (MTP verify / pair prefill used
3424/// to run two full matvecs — double weight traffic and double unpack).
3425/// Per-input math is identical to `vbitmatvec` → same accuracy contract.
3426#[allow(clippy::too_many_arguments)]
3427fn vbitmatvec2(
3428    bytes: &[u8],
3429    offsets: &[usize],
3430    x1: &[f32],
3431    x2: &[f32],
3432    rows: usize,
3433    cols: usize,
3434    o1: &mut [f32],
3435    o2: &mut [f32],
3436    pool: Option<&Pool>,
3437) {
3438    debug_assert_eq!(o1.len(), rows);
3439    debug_assert_eq!(o2.len(), rows);
3440
3441    if a8w8_enabled() {
3442        let a1 = split_act(x1);
3443        let a2 = split_act(x2);
3444        let p1 = SendMut(o1.as_mut_ptr());
3445        let p2 = SendMut(o2.as_mut_ptr());
3446        let run = move |start: usize, end: usize| {
3447            vbit_range2_a8w8(
3448                bytes, offsets, x1, x2, &a1, &a2, rows, cols, p1, p2, start, end,
3449            )
3450        };
3451        dispatch_rows(pool, rows, &run);
3452        return;
3453    }
3454
3455    let p1 = SendMut(o1.as_mut_ptr());
3456    let p2 = SendMut(o2.as_mut_ptr());
3457    let run = move |start: usize, end: usize| {
3458        vbit_range2_f32(bytes, offsets, x1, x2, rows, cols, p1, p2, start, end)
3459    };
3460    dispatch_rows(pool, rows, &run);
3461}
3462
3463/// Two-input vbit row range via the A8W8 int8 path — kernel body of
3464/// `vbitmatvec2`, extracted for pair multi-matrix jobs (b=8 rows go
3465/// exact f32 for both lanes, bits streamed once).
3466#[allow(clippy::too_many_arguments)]
3467fn vbit_range2_a8w8(
3468    bytes: &[u8],
3469    offsets: &[usize],
3470    x1: &[f32],
3471    x2: &[f32],
3472    a1: &SplitAct,
3473    a2: &SplitAct,
3474    rows: usize,
3475    cols: usize,
3476    p1: SendMut,
3477    p2: SendMut,
3478    start: usize,
3479    end: usize,
3480) {
3481    let ng = cols / GROUP_SIZE;
3482    let bits = &bytes[..rows];
3483    let sc_off = rows;
3484    let row_dots = |r: usize| -> (f32, f32) {
3485        let b = bits[r] as usize;
3486        let l = (1i32 << (b - 1)) - 1;
3487        let data = &bytes[offsets[r]..offsets[r + 1]];
3488        if b == 8 {
3489            // u−L reaches 128 → does not fit i8; exact f32 path,
3490            // bits still streamed once for both lanes.
3491            let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
3492            let (mut d1, mut d2) = (0f32, 0f32);
3493            for g in 0..ng {
3494                let so = (r * ng + g) * 2;
3495                let sgf = f16_to_f32(u16::from_le_bytes([
3496                    bytes[sc_off + so],
3497                    bytes[sc_off + so + 1],
3498                ]));
3499                let (mut g1, mut g2) = (0f32, 0f32);
3500                for k in 0..GROUP_SIZE {
3501                    if nbits < 8 {
3502                        acc = (acc << 8) | data[idx] as u64;
3503                        idx += 1;
3504                        nbits += 8;
3505                    }
3506                    let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
3507                    nbits -= 8;
3508                    let w = (u - l) as f32;
3509                    g1 += w * x1[g * GROUP_SIZE + k];
3510                    g2 += w * x2[g * GROUP_SIZE + k];
3511                }
3512                d1 += g1 * sgf;
3513                d2 += g2 * sgf;
3514            }
3515            return (d1, d2);
3516        }
3517        thread_local! {
3518            static VBIT_SCRATCH2: std::cell::RefCell<Vec<u8>> =
3519                const { std::cell::RefCell::new(Vec::new()) };
3520        }
3521        #[inline(always)]
3522        fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
3523            for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
3524                let u = unpack8::<B>(&data[blk * B..]);
3525                for k in 0..8 {
3526                    chunk[k] = (u[k] - l) as i8 as u8;
3527                }
3528            }
3529        }
3530        VBIT_SCRATCH2.with(|scratch| {
3531            let mut buf = scratch.borrow_mut();
3532            buf.resize(cols, 0);
3533            match b {
3534                3 => fill::<3>(data, l, &mut buf),
3535                4 => vbit_fill4(data, &mut buf),
3536                5 => fill::<5>(data, l, &mut buf),
3537                6 => fill::<6>(data, l, &mut buf),
3538                _ => unreachable!(),
3539            }
3540            let (mut d1, mut d2) = (0f32, 0f32);
3541            for g in 0..ng {
3542                let so = (r * ng + g) * 2;
3543                let s = f16_to_f32(u16::from_le_bytes([
3544                    bytes[sc_off + so],
3545                    bytes[sc_off + so + 1],
3546                ]));
3547                let wg = &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
3548                let v1 = dot_i8_i8(wg, &a1.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a1.sx;
3549                let v2 = dot_i8_i8(wg, &a2.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a2.sx;
3550                d1 += v1 * s;
3551                d2 += v2 * s;
3552            }
3553            for &(j, xv) in &a1.outliers {
3554                let so = (r * ng + j / GROUP_SIZE) * 2;
3555                let s = f16_to_f32(u16::from_le_bytes([
3556                    bytes[sc_off + so],
3557                    bytes[sc_off + so + 1],
3558                ]));
3559                d1 += (buf[j] as i8) as f32 * s * xv;
3560            }
3561            for &(j, xv) in &a2.outliers {
3562                let so = (r * ng + j / GROUP_SIZE) * 2;
3563                let s = f16_to_f32(u16::from_le_bytes([
3564                    bytes[sc_off + so],
3565                    bytes[sc_off + so + 1],
3566                ]));
3567                d2 += (buf[j] as i8) as f32 * s * xv;
3568            }
3569            (d1, d2)
3570        })
3571    };
3572    for r in start..end {
3573        let (v1, v2) = row_dots(r);
3574        // SAFETY: disjoint row ranges per worker.
3575        unsafe {
3576            *p1.at(r) = v1;
3577            *p2.at(r) = v2;
3578        }
3579    }
3580}
3581
3582/// Two-input exact scalar vbit row range (same extraction) —
3583/// per-bit-width specialized, two accumulators per row; per-lane
3584/// accumulation order matches `vbitmatvec` exactly.
3585#[allow(clippy::too_many_arguments)]
3586fn vbit_range2_f32(
3587    bytes: &[u8],
3588    offsets: &[usize],
3589    x1: &[f32],
3590    x2: &[f32],
3591    rows: usize,
3592    cols: usize,
3593    p1: SendMut,
3594    p2: SendMut,
3595    start: usize,
3596    end: usize,
3597) {
3598    let ng = cols / GROUP_SIZE;
3599    let bits = &bytes[..rows];
3600    let sc_off = rows;
3601    #[inline(always)]
3602    #[allow(clippy::too_many_arguments)]
3603    fn dot_row2<const B: usize>(
3604        data: &[u8],
3605        bytes: &[u8],
3606        sc_off: usize,
3607        r: usize,
3608        ng: usize,
3609        x1: &[f32],
3610        x2: &[f32],
3611    ) -> (f32, f32) {
3612        let l = ((1i32 << (B - 1)) - 1) as f32;
3613        let gbytes = GROUP_SIZE * B / 8;
3614        let (mut d1, mut d2) = (0f32, 0f32);
3615        for g in 0..ng {
3616            let so = (r * ng + g) * 2;
3617            let s = f16_to_f32(u16::from_le_bytes([
3618                bytes[sc_off + so],
3619                bytes[sc_off + so + 1],
3620            ]));
3621            let x1g = &x1[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
3622            let x2g = &x2[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
3623            let gd0 = &data[g * gbytes..(g + 1) * gbytes];
3624            let (mut g1, mut g2) = (0f32, 0f32);
3625            for blk in 0..GROUP_SIZE / 8 {
3626                let u = unpack8::<B>(&gd0[blk * B..]);
3627                for k in 0..8 {
3628                    let w = u[k] as f32 - l;
3629                    g1 += w * x1g[blk * 8 + k];
3630                    g2 += w * x2g[blk * 8 + k];
3631                }
3632            }
3633            d1 += g1 * s;
3634            d2 += g2 * s;
3635        }
3636        (d1, d2)
3637    }
3638    for r in start..end {
3639        let data = &bytes[offsets[r]..offsets[r + 1]];
3640        let (v1, v2) = match bits[r] {
3641            3 => dot_row2::<3>(data, bytes, sc_off, r, ng, x1, x2),
3642            4 => dot_row2::<4>(data, bytes, sc_off, r, ng, x1, x2),
3643            5 => dot_row2::<5>(data, bytes, sc_off, r, ng, x1, x2),
3644            6 => dot_row2::<6>(data, bytes, sc_off, r, ng, x1, x2),
3645            8 => dot_row2::<8>(data, bytes, sc_off, r, ng, x1, x2),
3646            b => unreachable!("vbit bit-width {b} (validated at load)"),
3647        };
3648        // SAFETY: disjoint row ranges per worker.
3649        unsafe {
3650            *p1.at(r) = v1;
3651            *p2.at(r) = v2;
3652        }
3653    }
3654}
3655
3656// ───────────────────── q4_tiled kernels (§4.3) ─────────────────────
3657
3658/// One q4_tiled row dot on the A8W8 int8 path: per 32-group the tile
3659/// is ONE sequential read — [f16 scale][16B nibbles] — versus the two
3660/// distant streams of the split layout. Values/order identical to the
3661/// split kernels.
3662#[inline]
3663#[allow(unreachable_code)]
3664fn dot_q4t_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
3665    #[cfg(target_arch = "aarch64")]
3666    unsafe {
3667        return dot_q4t_row_sdot(bytes, r, gpr, xq);
3668    }
3669    #[cfg(target_arch = "x86_64")]
3670    unsafe {
3671        if vnni_tiles_enabled() {
3672            return dot_q4t_row_vnni(bytes, r, gpr, xq);
3673        }
3674        return dot_q4t_row_avx2(bytes, r, gpr, xq);
3675    }
3676    let mut acc = 0f32;
3677    for gi in 0..gpr {
3678        let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
3679        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
3680        let mut d = 0i32;
3681        for (k, &b) in tile[2..].iter().enumerate() {
3682            d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
3683                + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
3684        }
3685        acc += d as f32 * s;
3686    }
3687    acc
3688}
3689
3690#[cfg(target_arch = "aarch64")]
3691#[target_feature(enable = "neon,dotprod")]
3692unsafe fn dot_q4t_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
3693    // SAFETY: callers uphold slice-length contracts (18B tile per group,
3694    // xq.len() == gpr·GROUP_SIZE).
3695    unsafe {
3696        use core::arch::aarch64::*;
3697        use core::arch::asm;
3698        let lomask = vdupq_n_u8(0x0F);
3699        let eight = vdupq_n_s8(8);
3700        let mut acc = 0f32;
3701        for gi in 0..gpr {
3702            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
3703            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
3704            let b = vld1q_u8(t.add(2));
3705            let lo = vandq_u8(b, lomask);
3706            let hi = vshrq_n_u8::<4>(b);
3707            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
3708            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
3709            let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
3710            let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
3711            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
3712            asm!(
3713                "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
3714                "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
3715                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
3716                e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
3717                options(pure, nomem, nostack),
3718            );
3719            acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
3720        }
3721        acc
3722    }
3723}
3724
3725#[cfg(target_arch = "x86_64")]
3726#[target_feature(enable = "avx2")]
3727unsafe fn dot_q4t_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
3728    // SAFETY: see dot_q4t_row_sdot.
3729    unsafe {
3730        use core::arch::x86_64::*;
3731        let lomask = _mm_set1_epi8(0x0F);
3732        let eight = _mm256_set1_epi8(8);
3733        let ones = _mm256_set1_epi16(1);
3734        let mut acc = 0f32;
3735        for gi in 0..gpr {
3736            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
3737            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
3738            let b = _mm_loadu_si128(t.add(2) as *const __m128i);
3739            let lo = _mm_and_si128(b, lomask);
3740            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
3741            let w = _mm256_sub_epi8(
3742                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
3743                eight,
3744            );
3745            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
3746            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
3747            let d = _mm256_madd_epi16(p16, ones);
3748            let hi128 = _mm256_extracti128_si256::<1>(d);
3749            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
3750            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
3751            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
3752            acc += _mm_cvtsi128_si32(s32) as f32 * s;
3753        }
3754        acc
3755    }
3756}
3757
3758/// VNNI twin of `dot_q4t_row_avx2`: same unpack, `vpdpbusd` replaces
3759/// the maddubs+madd pair (see `dpbusd_hsum` — sums are bit-identical).
3760/// 256-bit VL encoding, so the VEX `vpsignb` stays usable.
3761#[cfg(target_arch = "x86_64")]
3762#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
3763unsafe fn dot_q4t_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
3764    // SAFETY: see dot_q4t_row_sdot.
3765    unsafe {
3766        use core::arch::x86_64::*;
3767        let lomask = _mm_set1_epi8(0x0F);
3768        let eight = _mm256_set1_epi8(8);
3769        let mut acc = 0f32;
3770        for gi in 0..gpr {
3771            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
3772            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
3773            let b = _mm_loadu_si128(t.add(2) as *const __m128i);
3774            let lo = _mm_and_si128(b, lomask);
3775            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
3776            let w = _mm256_sub_epi8(
3777                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
3778                eight,
3779            );
3780            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
3781            let d = dpbusd_hsum(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
3782            acc += d as f32 * s;
3783        }
3784        acc
3785    }
3786}
3787
3788/// One q4_tiled row against FOUR activation streams: the nibble unpack
3789/// and abs() happen once per group instead of once per (group,
3790/// activation) — the unpack is the dominant per-element cost of the
3791/// tiled format (roadmap P0 portable blocking, q4t leg).
3792#[cfg(target_arch = "x86_64")]
3793// `fma` is NOT implied by `avx2`: without it LLVM lowers _mm256_fmadd_ps
3794// to a libm call per lane — measured 2x slower than the reduction this
3795// kernel replaces. The runtime gate (`avx2_enabled`) already requires
3796// both features, so declaring it here is safe.
3797#[target_feature(enable = "avx2,fma")]
3798unsafe fn dot_q4t_row_1x4_avx2(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
3799    // SAFETY: callers uphold the 18B-tile and xq-length contracts.
3800    unsafe {
3801        use core::arch::x86_64::*;
3802        let lomask = _mm_set1_epi8(0x0F);
3803        let eight = _mm256_set1_epi8(8);
3804        let ones = _mm256_set1_epi16(1);
3805        // One f32 accumulator VECTOR per activation, reduced once at the
3806        // end. Folding each group's i32 lanes to a scalar inside the loop
3807        // costs an extracti128 + three shift/add + a movd — a cross-lane
3808        // dependency chain per (group, activation), 288 of them per row at
3809        // cols=2304. The per-group scale is what forces a float
3810        // accumulator; it does not force a horizontal sum.
3811        //
3812        // The four accumulators are NAMED, not an array: as `[__m256; 4]`
3813        // indexed by a loop variable LLVM keeps them in memory and every
3814        // group pays four 32-byte loads and stores. That alone made this
3815        // kernel 2x SLOWER than the per-group reduction it replaces
3816        // (measured on the EPYC box: 150 s vs 71 s for two 256² steps).
3817        let mut f0 = _mm256_setzero_ps();
3818        let mut f1 = _mm256_setzero_ps();
3819        let mut f2 = _mm256_setzero_ps();
3820        let mut f3 = _mm256_setzero_ps();
3821        for gi in 0..gpr {
3822            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
3823            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
3824            let sv = _mm256_set1_ps(s);
3825            let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
3826            let lo = _mm_and_si128(bb, lomask);
3827            let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
3828            let w = _mm256_sub_epi8(
3829                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
3830                eight,
3831            );
3832            let aw = _mm256_abs_epi8(w);
3833            let off = gi * GROUP_SIZE;
3834            let dot = |xq: &[i8]| {
3835                let x = _mm256_loadu_si256(xq.as_ptr().add(off) as *const __m256i);
3836                let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
3837                _mm256_cvtepi32_ps(_mm256_madd_epi16(p16, ones))
3838            };
3839            f0 = _mm256_fmadd_ps(dot(xs[0]), sv, f0);
3840            f1 = _mm256_fmadd_ps(dot(xs[1]), sv, f1);
3841            f2 = _mm256_fmadd_ps(dot(xs[2]), sv, f2);
3842            f3 = _mm256_fmadd_ps(dot(xs[3]), sv, f3);
3843        }
3844        [
3845            hsum256_ps(f0),
3846            hsum256_ps(f1),
3847            hsum256_ps(f2),
3848            hsum256_ps(f3),
3849        ]
3850    }
3851}
3852
3853/// Horizontal sum of eight f32 lanes — the one cross-lane reduction the
3854/// blocked kernels pay, once per row instead of once per group.
3855#[cfg(target_arch = "x86_64")]
3856#[target_feature(enable = "avx2")]
3857#[inline]
3858unsafe fn hsum256_ps(v: core::arch::x86_64::__m256) -> f32 {
3859    // SAFETY: pure register arithmetic on the caller's vector.
3860    unsafe {
3861        use core::arch::x86_64::*;
3862        let hi = _mm256_extractf128_ps::<1>(v);
3863        let s = _mm_add_ps(_mm256_castps256_ps128(v), hi);
3864        let s = _mm_add_ps(s, _mm_movehl_ps(s, s));
3865        let s = _mm_add_ss(s, _mm_shuffle_ps::<0x55>(s, s));
3866        _mm_cvtss_f32(s)
3867    }
3868}
3869
3870/// VNNI twin of `dot_q4t_row_1x4_avx2` (see `dpbusd_hsum`).
3871#[cfg(target_arch = "x86_64")]
3872#[target_feature(enable = "avx2,fma,avx512f,avx512bw,avx512vl,avx512vnni")]
3873unsafe fn dot_q4t_row_1x4_vnni(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
3874    // SAFETY: callers uphold the 18B-tile and xq-length contracts.
3875    unsafe {
3876        use core::arch::x86_64::*;
3877        let lomask = _mm_set1_epi8(0x0F);
3878        let eight = _mm256_set1_epi8(8);
3879        // Same shape as the AVX2 twin: accumulate in f32 vectors and pay
3880        // one cross-lane reduction per row, not per (group, activation).
3881        let mut f0 = _mm256_setzero_ps();
3882        let mut f1 = _mm256_setzero_ps();
3883        let mut f2 = _mm256_setzero_ps();
3884        let mut f3 = _mm256_setzero_ps();
3885        for gi in 0..gpr {
3886            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
3887            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
3888            let sv = _mm256_set1_ps(s);
3889            let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
3890            let lo = _mm_and_si128(bb, lomask);
3891            let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
3892            let w = _mm256_sub_epi8(
3893                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
3894                eight,
3895            );
3896            let aw = _mm256_abs_epi8(w);
3897            let off = gi * GROUP_SIZE;
3898            let dot = |xq: &[i8]| {
3899                let x = _mm256_loadu_si256(xq.as_ptr().add(off) as *const __m256i);
3900                _mm256_cvtepi32_ps(_mm256_dpbusd_epi32(
3901                    _mm256_setzero_si256(),
3902                    aw,
3903                    _mm256_sign_epi8(x, w),
3904                ))
3905            };
3906            f0 = _mm256_fmadd_ps(dot(xs[0]), sv, f0);
3907            f1 = _mm256_fmadd_ps(dot(xs[1]), sv, f1);
3908            f2 = _mm256_fmadd_ps(dot(xs[2]), sv, f2);
3909            f3 = _mm256_fmadd_ps(dot(xs[3]), sv, f3);
3910        }
3911        let acc = [
3912            hsum256_ps(f0),
3913            hsum256_ps(f1),
3914            hsum256_ps(f2),
3915            hsum256_ps(f3),
3916        ];
3917        acc
3918    }
3919}
3920
3921/// ARM twin of `dot_q4t_row_1x4_avx2`: one nibble unpack per group
3922/// serves FOUR activation streams. Per stream the group order and f32
3923/// accumulation match `dot_q4t_row_sdot` exactly — batch == matvec
3924/// bit-for-bit.
3925#[cfg(target_arch = "aarch64")]
3926#[target_feature(enable = "neon,dotprod")]
3927unsafe fn dot_q4t_row_1x4_sdot(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
3928    // SAFETY: callers uphold the 18B-tile and xq-length contracts.
3929    unsafe {
3930        use core::arch::aarch64::*;
3931        use core::arch::asm;
3932        let lomask = vdupq_n_u8(0x0F);
3933        let eight = vdupq_n_s8(8);
3934        let mut acc = [0f32; 4];
3935        for gi in 0..gpr {
3936            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
3937            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
3938            let b = vld1q_u8(t.add(2));
3939            let lo = vandq_u8(b, lomask);
3940            let hi = vshrq_n_u8::<4>(b);
3941            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
3942            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
3943            for (k, xq) in xs.iter().enumerate() {
3944                let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
3945                let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
3946                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
3947                asm!(
3948                    "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
3949                    "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
3950                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
3951                    e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
3952                    options(pure, nomem, nostack),
3953                );
3954                acc[k] += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
3955            }
3956        }
3957        acc
3958    }
3959}
3960
3961/// Exact-term correction for A8W8 outliers on a tiled row.
3962#[inline]
3963fn q4t_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
3964    let gi = j / GROUP_SIZE;
3965    let k = j % GROUP_SIZE;
3966    let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
3967    let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
3968    let byte = tile[2 + k / 2];
3969    let nib = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
3970    ((nib as i32 - 8) as f32, s)
3971}
3972
3973/// Exact scalar q4_tiled row (CMF_SDOT=0 contract) — same pairwise
3974/// accumulation shape as `q4_range_f32`.
3975#[inline]
3976fn q4t_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
3977    let mut acc = 0f32;
3978    for gi in 0..gpr {
3979        let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
3980        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
3981        let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
3982        let mut ga = 0f32;
3983        for (k, &b) in tile[2..].iter().enumerate() {
3984            ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
3985                + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
3986        }
3987        acc += ga * s;
3988    }
3989    acc
3990}
3991
3992/// Split view of a `q4tp` payload. The three planes are resolved once per
3993/// matvec instead of per row — `q4tp_sections` is cheap, but doing it inside
3994/// the row loop would put a division on the hot path for nothing.
3995struct Q4tpView<'a> {
3996    nib: &'a [u8],
3997    params: &'a [u8],
3998    codes: &'a [u8],
3999    stride: usize,
4000    /// q2tp reads the ladder with rung 0 = exact zero.
4001    zero_rung: bool,
4002}
4003
4004impl<'a> Q4tpView<'a> {
4005    fn new(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
4006        let (params_off, codes_off, stride) = q4tp_sections(rows, cols);
4007        Self {
4008            nib: &bytes[..params_off],
4009            params: &bytes[params_off..codes_off],
4010            codes: &bytes[codes_off..],
4011            stride,
4012            zero_rung: false,
4013        }
4014    }
4015
4016    /// The q2tp view: identical params/codes planes, 8 B weight chunks.
4017    fn new_q2(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
4018        let (params_off, codes_off, stride) = q2tp_sections(rows, cols);
4019        Self {
4020            nib: &bytes[..params_off],
4021            params: &bytes[params_off..codes_off],
4022            codes: &bytes[codes_off..],
4023            stride,
4024            zero_rung: true,
4025        }
4026    }
4027
4028    /// Expand row `r`'s per-tile scales into `out` (length `gpr`).
4029    ///
4030    /// Doing this once per row — rather than decoding a 5-bit code inside the
4031    /// tile loop — is what makes the format free at runtime. Random access to
4032    /// a packed 5-bit field costs a division, two bounds checks and a branch;
4033    /// the tile's actual work is two `sdot`s, so per-tile decoding dominated
4034    /// the kernel and cost 5x (measured: 1.4 vs 6.9 tok/s on Nanbeige-3B).
4035    /// Walking the plane sequentially with a bit accumulator is ~3 ops.
4036    /// Eight 5-bit codes are exactly five bytes, so a whole group of
4037    /// eight decodes from one little-endian word at fixed shifts. The
4038    /// bit-accumulator this replaces carried a data-dependent `while
4039    /// have < 5` refill whose branch sat in the innermost loop of every
4040    /// q4tp row; a decode profile put this function above the dot
4041    /// products it feeds. Same bitstream, same codes — just no branch
4042    /// and eight independent extractions.
4043    #[inline]
4044    fn scales_into(&self, r: usize, gpr: usize, out: &mut [f32]) {
4045        let tab = if self.zero_rung {
4046            q2tp_ladder(self.params, r)
4047        } else {
4048            q4tp_ladder(self.params, r)
4049        };
4050        let codes = &self.codes[r * self.stride..(r + 1) * self.stride];
4051        let out = &mut out[..gpr];
4052        let mut chunks = out.chunks_exact_mut(8);
4053        let mut ci = 0usize;
4054        for c in &mut chunks {
4055            let w = u64::from(codes[ci])
4056                | u64::from(codes[ci + 1]) << 8
4057                | u64::from(codes[ci + 2]) << 16
4058                | u64::from(codes[ci + 3]) << 24
4059                | u64::from(codes[ci + 4]) << 32;
4060            for (k, o) in c.iter_mut().enumerate() {
4061                *o = tab[((w >> (5 * k)) & 31) as usize];
4062            }
4063            ci += 5;
4064        }
4065        // Fewer than eight codes left: the shared total accessor, which
4066        // tolerates a 5-bit field whose spill byte is past the stride.
4067        let tail = &codes[ci..];
4068        for (k, o) in chunks.into_remainder().iter_mut().enumerate() {
4069            *o = tab[q4tp_code(tail, k)];
4070        }
4071    }
4072}
4073
4074#[inline]
4075fn dot_q4tp_row_i8(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
4076    #[cfg(target_arch = "aarch64")]
4077    unsafe {
4078        return dot_q4tp_row_sdot(nib, r, gpr, xq, scales);
4079    }
4080    #[cfg(target_arch = "x86_64")]
4081    unsafe {
4082        if vnni_tiles_enabled() {
4083            return dot_q4tp_row_vnni(nib, r, gpr, xq, scales);
4084        }
4085        return dot_q4tp_row_avx2(nib, r, gpr, xq, scales);
4086    }
4087    #[allow(unreachable_code)]
4088    {
4089        let mut acc = 0f32;
4090        for gi in 0..gpr {
4091            let tile = &nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
4092            let s = scales[gi];
4093            let mut d = 0i32;
4094            for (k, &b) in tile.iter().enumerate() {
4095                d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
4096                    + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
4097            }
4098            acc += d as f32 * s;
4099        }
4100        acc
4101    }
4102}
4103
4104/// q4tp twin of `dot_q4t_row_sdot`: identical nibble math, but the tile
4105/// stride is 16 B (no inline scale) and the scale is a ladder lookup.
4106#[cfg(target_arch = "aarch64")]
4107#[target_feature(enable = "neon,dotprod")]
4108unsafe fn dot_q4tp_row_sdot(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
4109    // SAFETY: callers uphold slice-length contracts (16B tile per group,
4110    // xq.len() == gpr·GROUP_SIZE, codes covering gpr 5-bit fields).
4111    unsafe {
4112        use core::arch::aarch64::*;
4113        use core::arch::asm;
4114        let lomask = vdupq_n_u8(0x0F);
4115        let eight = vdupq_n_s8(8);
4116        let mut acc = 0f32;
4117        for gi in 0..gpr {
4118            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
4119            let s = *scales.get_unchecked(gi);
4120            let b = vld1q_u8(t);
4121            let lo = vandq_u8(b, lomask);
4122            let hi = vshrq_n_u8::<4>(b);
4123            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
4124            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
4125            let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
4126            let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
4127            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
4128            asm!(
4129                "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
4130                "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
4131                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4132                e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
4133                options(pure, nomem, nostack),
4134            );
4135            acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
4136        }
4137        acc
4138    }
4139}
4140
4141#[cfg(target_arch = "x86_64")]
4142#[target_feature(enable = "avx2")]
4143unsafe fn dot_q4tp_row_avx2(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
4144    // SAFETY: see dot_q4tp_row_sdot.
4145    unsafe {
4146        use core::arch::x86_64::*;
4147        let lomask = _mm_set1_epi8(0x0F);
4148        let eight = _mm256_set1_epi8(8);
4149        let ones = _mm256_set1_epi16(1);
4150        let mut acc = 0f32;
4151        for gi in 0..gpr {
4152            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
4153            let s = *scales.get_unchecked(gi);
4154            let b = _mm_loadu_si128(t as *const __m128i);
4155            let lo = _mm_and_si128(b, lomask);
4156            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
4157            let w = _mm256_sub_epi8(
4158                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4159                eight,
4160            );
4161            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4162            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
4163            let d = _mm256_madd_epi16(p16, ones);
4164            let hi128 = _mm256_extracti128_si256::<1>(d);
4165            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
4166            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
4167            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
4168            acc += _mm_cvtsi128_si32(s32) as f32 * s;
4169        }
4170        acc
4171    }
4172}
4173
4174/// VNNI twin of `dot_q4tp_row_avx2` (see `dot_q4t_row_vnni` for why the
4175/// 256-bit VL encoding is the one to use here).
4176#[cfg(target_arch = "x86_64")]
4177#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
4178unsafe fn dot_q4tp_row_vnni(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
4179    // SAFETY: see dot_q4tp_row_sdot.
4180    unsafe {
4181        use core::arch::x86_64::*;
4182        let lomask = _mm_set1_epi8(0x0F);
4183        let eight = _mm256_set1_epi8(8);
4184        let mut acc = 0f32;
4185        for gi in 0..gpr {
4186            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
4187            let s = *scales.get_unchecked(gi);
4188            let b = _mm_loadu_si128(t as *const __m128i);
4189            let lo = _mm_and_si128(b, lomask);
4190            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
4191            let w = _mm256_sub_epi8(
4192                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4193                eight,
4194            );
4195            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4196            acc += dpbusd_hsum(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w)) as f32 * s;
4197        }
4198        acc
4199    }
4200}
4201
4202/// Exact scalar q4tp row — the `CMF_SDOT=0` contract, same pairwise
4203/// accumulation shape as `q4t_row_exact`.
4204#[inline]
4205fn q4tp_row_exact(nib: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
4206    let mut acc = 0f32;
4207    for gi in 0..gpr {
4208        let tile = &nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
4209        let s = scales[gi];
4210        let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
4211        let mut ga = 0f32;
4212        for (k, &b) in tile.iter().enumerate() {
4213            ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
4214                + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
4215        }
4216        acc += ga * s;
4217    }
4218    acc
4219}
4220
4221/// Single weight of a q4tp tensor — the a8w8 outlier path, which restores
4222/// activation outliers at full precision after the int8 pass.
4223#[inline]
4224fn q4tp_outlier(nib: &[u8], r: usize, gpr: usize, j: usize, scales: &[f32]) -> (f32, f32) {
4225    let (gi, k) = (j / GROUP_SIZE, j % GROUP_SIZE);
4226    let byte = nib[(r * gpr + gi) * Q4TP_NIB + k / 2];
4227    let n = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
4228    ((n as i32 - 8) as f32, scales[gi])
4229}
4230
4231/// Fused q4tp matvec (dispatch mirrors `q4t_matvec`).
4232fn q4tp_matvec(
4233    bytes: &[u8],
4234    x: &[f32],
4235    rows: usize,
4236    cols: usize,
4237    out: &mut [f32],
4238    pool: Option<&Pool>,
4239) {
4240    debug_assert_eq!(out.len(), rows);
4241    let gpr = cols / GROUP_SIZE;
4242    let v = Q4tpView::new(bytes, rows, cols);
4243    let out_addr = SendMut(out.as_mut_ptr());
4244    if a8w8_enabled() {
4245        let act = split_act(x);
4246        let run = |start: usize, end: usize| {
4247            // One scratch row of scales per worker — borrowed, not minted.
4248            with_krow(gpr, |sc| {
4249                for r in start..end {
4250                    v.scales_into(r, gpr, sc);
4251                    let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, sc) * act.sx;
4252                    for &(j, xv) in &act.outliers {
4253                        let (w, s) = q4tp_outlier(v.nib, r, gpr, j, sc);
4254                        acc += w * s * xv;
4255                    }
4256                    // SAFETY: disjoint row ranges per worker.
4257                    unsafe { *out_addr.at(r) = acc };
4258                }
4259            })
4260        };
4261        dispatch_rows(pool, rows, &run);
4262        return;
4263    }
4264    let run = |start: usize, end: usize| {
4265        with_krow(gpr, |sc| {
4266            for r in start..end {
4267                v.scales_into(r, gpr, sc);
4268                // SAFETY: disjoint row ranges per worker.
4269                unsafe { *out_addr.at(r) = q4tp_row_exact(v.nib, r, gpr, x, sc) };
4270            }
4271        })
4272    };
4273    dispatch_rows(pool, rows, &run);
4274}
4275
4276/// Fused two-input q4tp matvec — the SwiGLU gate/up pair. Weights and the
4277/// row ladder are read once and spent on both activation streams.
4278#[allow(clippy::too_many_arguments)]
4279fn q4tp_matvec2(
4280    bytes: &[u8],
4281    x1: &[f32],
4282    x2: &[f32],
4283    rows: usize,
4284    cols: usize,
4285    o1: &mut [f32],
4286    o2: &mut [f32],
4287    pool: Option<&Pool>,
4288) {
4289    let gpr = cols / GROUP_SIZE;
4290    let v = Q4tpView::new(bytes, rows, cols);
4291    let (p1, p2) = (SendMut(o1.as_mut_ptr()), SendMut(o2.as_mut_ptr()));
4292    let run = |start: usize, end: usize| {
4293        let mut sc = vec![0f32; gpr];
4294        for r in start..end {
4295            v.scales_into(r, gpr, &mut sc);
4296            // SAFETY: disjoint row ranges per worker.
4297            unsafe {
4298                *p1.at(r) = q4tp_row_exact(v.nib, r, gpr, x1, &sc);
4299                *p2.at(r) = q4tp_row_exact(v.nib, r, gpr, x2, &sc);
4300            }
4301        }
4302    };
4303    dispatch_rows(pool, rows, &run);
4304}
4305
4306/// Exact f32 dot of one q2tp row: 2-bit fields LSB-first, (c − 1.5)·s.
4307/// Scalar on purpose — the 2-bit class targets the GPU graph; the CPU
4308/// path exists for parity gates and small-machine fallback.
4309fn q2tp_row_exact(chunks: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
4310    let mut acc = 0f32;
4311    for gi in 0..gpr {
4312        let ch = &chunks[(r * gpr + gi) * Q2TP_CHUNK..(r * gpr + gi + 1) * Q2TP_CHUNK];
4313        let s = scales[gi];
4314        let xb = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
4315        let mut g = 0f32;
4316        for (k, &b) in ch.iter().enumerate() {
4317            g += ((b & 3) as f32 - 1.5) * xb[k * 4]
4318                + (((b >> 2) & 3) as f32 - 1.5) * xb[k * 4 + 1]
4319                + (((b >> 4) & 3) as f32 - 1.5) * xb[k * 4 + 2]
4320                + (((b >> 6) & 3) as f32 - 1.5) * xb[k * 4 + 3];
4321        }
4322        acc += s * g;
4323    }
4324    acc
4325}
4326
4327fn q2tp_matvec(
4328    bytes: &[u8],
4329    x: &[f32],
4330    rows: usize,
4331    cols: usize,
4332    out: &mut [f32],
4333    pool: Option<&Pool>,
4334) {
4335    debug_assert_eq!(out.len(), rows);
4336    let gpr = cols / GROUP_SIZE;
4337    let v = Q4tpView::new_q2(bytes, rows, cols);
4338    let out_addr = SendMut(out.as_mut_ptr());
4339    let run = |start: usize, end: usize| {
4340        with_krow(gpr, |sc| {
4341            for r in start..end {
4342                v.scales_into(r, gpr, sc);
4343                // SAFETY: disjoint row ranges per worker.
4344                unsafe { *out_addr.at(r) = q2tp_row_exact(v.nib, r, gpr, x, sc) };
4345            }
4346        })
4347    };
4348    dispatch_rows(pool, rows, &run);
4349}
4350
4351/// Fused two-input q2tp matvec — the SwiGLU gate/up pair.
4352#[allow(clippy::too_many_arguments)]
4353fn q2tp_matvec2(
4354    bytes: &[u8],
4355    x1: &[f32],
4356    x2: &[f32],
4357    rows: usize,
4358    cols: usize,
4359    o1: &mut [f32],
4360    o2: &mut [f32],
4361    pool: Option<&Pool>,
4362) {
4363    let gpr = cols / GROUP_SIZE;
4364    let v = Q4tpView::new_q2(bytes, rows, cols);
4365    let (p1, p2) = (SendMut(o1.as_mut_ptr()), SendMut(o2.as_mut_ptr()));
4366    let run = |start: usize, end: usize| {
4367        let mut sc = vec![0f32; gpr];
4368        for r in start..end {
4369            v.scales_into(r, gpr, &mut sc);
4370            // SAFETY: disjoint row ranges per worker.
4371            unsafe {
4372                *p1.at(r) = q2tp_row_exact(v.nib, r, gpr, x1, &sc);
4373                *p2.at(r) = q2tp_row_exact(v.nib, r, gpr, x2, &sc);
4374            }
4375        }
4376    };
4377    dispatch_rows(pool, rows, &run);
4378}
4379
4380/// Batched q2tp matmat: scalar row kernel over every batch column. CPU
4381/// prefill only — decode rides the graph, so plain and correct beats
4382/// clever here.
4383/// Test doors into the host 2-bit kernels: the stand's heap corruption
4384/// pointed at down-shaped tensors, and the private fns need a way to be
4385/// held to a reference without a model file around them.
4386pub fn q2tp_matvec_for_test(bytes: &[u8], x: &[f32], rows: usize, cols: usize, out: &mut [f32]) {
4387    q2tp_matvec(bytes, x, rows, cols, out, None);
4388}
4389
4390pub fn q2tp_matmat_for_test(
4391    bytes: &[u8],
4392    xs_all: &[f32],
4393    b: usize,
4394    rows: usize,
4395    cols: usize,
4396    out: &mut [f32],
4397) {
4398    q2tp_matmat(bytes, xs_all, b, rows, cols, out, None);
4399}
4400
4401fn q2tp_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    let v = Q4tpView::new_q2(bytes, rows, cols);
4413    let out_addr = SendMut(out.as_mut_ptr());
4414    let run = |start: usize, end: usize| {
4415        let mut sc = vec![0f32; gpr];
4416        for r in start..end {
4417            v.scales_into(r, gpr, &mut sc);
4418            for bi in 0..b {
4419                let x = &xs_all[bi * cols..(bi + 1) * cols];
4420                // SAFETY: disjoint row ranges per worker.
4421                unsafe { *out_addr.at(bi * rows + r) = q2tp_row_exact(v.nib, r, gpr, x, &sc) };
4422            }
4423        }
4424    };
4425    dispatch_rows(pool, rows, &run);
4426}
4427
4428/// Four batch columns against one q4tp row: the tile is unpacked ONCE and
4429/// spent on four activation streams, which is where a prefill batch stops
4430/// being weight-bandwidth-bound. Twin of `dot_q4t_row_1x4_sdot`.
4431#[cfg(target_arch = "aarch64")]
4432#[target_feature(enable = "neon,dotprod")]
4433unsafe fn dot_q4tp_row_1x4_sdot(
4434    nib: &[u8],
4435    r: usize,
4436    gpr: usize,
4437    xs: [&[i8]; 4],
4438    scales: &[f32],
4439) -> [f32; 4] {
4440    // SAFETY: see dot_q4tp_row_sdot; every xs[k] is gpr·GROUP_SIZE long.
4441    unsafe {
4442        use core::arch::aarch64::*;
4443        use core::arch::asm;
4444        let lomask = vdupq_n_u8(0x0F);
4445        let eight = vdupq_n_s8(8);
4446        // Named accumulators, NOT an array indexed by a loop variable: the
4447        // latter does not stay in registers (the same defect cost 2x in the
4448        // AVX2 q4t kernel and again in WGSL).
4449        let (mut f0, mut f1, mut f2, mut f3) = (0f32, 0f32, 0f32, 0f32);
4450        for gi in 0..gpr {
4451            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
4452            let s = *scales.get_unchecked(gi);
4453            let bb = vld1q_u8(t);
4454            let lo = vandq_u8(bb, lomask);
4455            let hi = vshrq_n_u8::<4>(bb);
4456            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
4457            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
4458            let mut d = [0f32; 4];
4459            for (k, dk) in d.iter_mut().enumerate() {
4460                let x0 = vld1q_s8(xs[k].as_ptr().add(gi * GROUP_SIZE));
4461                let x1 = vld1q_s8(xs[k].as_ptr().add(gi * GROUP_SIZE + 16));
4462                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
4463                asm!(
4464                    "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
4465                    "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
4466                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4467                    e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
4468                    options(pure, nomem, nostack),
4469                );
4470                *dk = vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
4471            }
4472            f0 += d[0];
4473            f1 += d[1];
4474            f2 += d[2];
4475            f3 += d[3];
4476        }
4477        [f0, f1, f2, f3]
4478    }
4479}
4480
4481/// Fused q4tp matmat — the same three arms `q4t_matmat` has. Shipping only
4482/// the scalar one made Nanbeige-3B decode at 1.2 tok/s against q4t's 5.9:
4483/// the format was fine, the missing arms were the whole regression.
4484fn q4tp_matmat(
4485    bytes: &[u8],
4486    xs_all: &[f32],
4487    b: usize,
4488    rows: usize,
4489    cols: usize,
4490    out: &mut [f32],
4491    pool: Option<&Pool>,
4492) {
4493    debug_assert_eq!(out.len(), b * rows);
4494    let gpr = cols / GROUP_SIZE;
4495    let v = Q4tpView::new(bytes, rows, cols);
4496
4497    // Wide batches ride the AMX through a dequant-tile sgemm, as in q4t.
4498    #[cfg(target_os = "macos")]
4499    if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
4500        dequant_matmat_accel(
4501            &|r, dst| {
4502                let mut sc = [0f32; 32];
4503                let mut scv;
4504                let s: &[f32] = if gpr <= 32 {
4505                    v.scales_into(r, gpr, &mut sc);
4506                    &sc[..gpr]
4507                } else {
4508                    scv = vec![0f32; gpr];
4509                    v.scales_into(r, gpr, &mut scv);
4510                    &scv
4511                };
4512                for gi in 0..gpr {
4513                    let tile = &v.nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
4514                    for (k, &bb) in tile.iter().enumerate() {
4515                        dst[gi * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s[gi];
4516                        dst[gi * GROUP_SIZE + k * 2 + 1] =
4517                            (((bb >> 4) & 0x0F) as f32 - 8.0) * s[gi];
4518                    }
4519                }
4520            },
4521            xs_all,
4522            b,
4523            rows,
4524            cols,
4525            out,
4526            pool,
4527        );
4528        return;
4529    }
4530
4531    let out_addr = SendMut(out.as_mut_ptr());
4532    if a8w8_enabled() {
4533        let acts: Vec<SplitAct> = (0..b)
4534            .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
4535            .collect();
4536        let acts = &acts;
4537        #[cfg(target_arch = "aarch64")]
4538        let blocked_ok = sdot_enabled() && blocked_enabled();
4539        #[cfg(not(target_arch = "aarch64"))]
4540        let blocked_ok = false;
4541        let run = |start: usize, end: usize| {
4542            let mut sc = vec![0f32; gpr];
4543            for r in start..end {
4544                v.scales_into(r, gpr, &mut sc);
4545                let mut bi = 0usize;
4546                #[cfg(target_arch = "aarch64")]
4547                if blocked_ok {
4548                    while bi + 4 <= acts.len() {
4549                        let xs = [
4550                            acts[bi].xq.as_slice(),
4551                            acts[bi + 1].xq.as_slice(),
4552                            acts[bi + 2].xq.as_slice(),
4553                            acts[bi + 3].xq.as_slice(),
4554                        ];
4555                        let d = unsafe { dot_q4tp_row_1x4_sdot(v.nib, r, gpr, xs, &sc) };
4556                        for k in 0..4 {
4557                            let act = &acts[bi + k];
4558                            let mut acc = d[k] * act.sx;
4559                            for &(j, xv) in &act.outliers {
4560                                let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
4561                                acc += w * s * xv;
4562                            }
4563                            // SAFETY: disjoint (bi, r) cells per worker.
4564                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
4565                        }
4566                        bi += 4;
4567                    }
4568                }
4569                let _ = blocked_ok;
4570                while bi < acts.len() {
4571                    let act = &acts[bi];
4572                    let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, &sc) * act.sx;
4573                    for &(j, xv) in &act.outliers {
4574                        let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
4575                        acc += w * s * xv;
4576                    }
4577                    // SAFETY: disjoint (bi, r) cells per worker range.
4578                    unsafe { *out_addr.at(bi * rows + r) = acc };
4579                    bi += 1;
4580                }
4581            }
4582        };
4583        dispatch_rows(pool, rows, &run);
4584        return;
4585    }
4586
4587    let run = |start: usize, end: usize| {
4588        let mut sc = vec![0f32; gpr];
4589        for r in start..end {
4590            v.scales_into(r, gpr, &mut sc);
4591            for bi in 0..b {
4592                let x = &xs_all[bi * cols..(bi + 1) * cols];
4593                // SAFETY: disjoint (bi, r) cells per worker range.
4594                unsafe { *out_addr.at(bi * rows + r) = q4tp_row_exact(v.nib, r, gpr, x, &sc) };
4595            }
4596        }
4597    };
4598    dispatch_rows(pool, rows, &run);
4599}
4600
4601/// Fused q4_tiled matvec (dispatch mirrors `q4matvec`).
4602fn q4t_matvec(
4603    bytes: &[u8],
4604    x: &[f32],
4605    rows: usize,
4606    cols: usize,
4607    out: &mut [f32],
4608    pool: Option<&Pool>,
4609) {
4610    debug_assert_eq!(out.len(), rows);
4611    let gpr = cols / GROUP_SIZE;
4612    let out_addr = SendMut(out.as_mut_ptr());
4613    if a8w8_enabled() {
4614        let act = split_act(x);
4615        let run = move |start: usize, end: usize| {
4616            for r in start..end {
4617                let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
4618                for &(j, xv) in &act.outliers {
4619                    let (w, s) = q4t_outlier(bytes, r, gpr, j);
4620                    acc += w * s * xv;
4621                }
4622                // SAFETY: disjoint row ranges per worker.
4623                unsafe { *out_addr.at(r) = acc };
4624            }
4625        };
4626        dispatch_rows(pool, rows, &run);
4627        return;
4628    }
4629    let run = move |start: usize, end: usize| {
4630        for r in start..end {
4631            // SAFETY: disjoint row ranges per worker.
4632            unsafe { *out_addr.at(r) = q4t_row_exact(bytes, r, gpr, x) };
4633        }
4634    };
4635    dispatch_rows(pool, rows, &run);
4636}
4637
4638/// Fused two-input q4_tiled matvec (weights read once per pair).
4639#[allow(clippy::too_many_arguments)]
4640fn q4t_matvec2(
4641    bytes: &[u8],
4642    x1: &[f32],
4643    x2: &[f32],
4644    rows: usize,
4645    cols: usize,
4646    o1: &mut [f32],
4647    o2: &mut [f32],
4648    pool: Option<&Pool>,
4649) {
4650    let gpr = cols / GROUP_SIZE;
4651    let p1 = SendMut(o1.as_mut_ptr());
4652    let p2 = SendMut(o2.as_mut_ptr());
4653    if a8w8_enabled() {
4654        let a1 = split_act(x1);
4655        let a2 = split_act(x2);
4656        let run = move |start: usize, end: usize| {
4657            for r in start..end {
4658                let mut v1 = dot_q4t_row_i8(bytes, r, gpr, &a1.xq) * a1.sx;
4659                let mut v2 = dot_q4t_row_i8(bytes, r, gpr, &a2.xq) * a2.sx;
4660                for &(j, xv) in &a1.outliers {
4661                    let (w, s) = q4t_outlier(bytes, r, gpr, j);
4662                    v1 += w * s * xv;
4663                }
4664                for &(j, xv) in &a2.outliers {
4665                    let (w, s) = q4t_outlier(bytes, r, gpr, j);
4666                    v2 += w * s * xv;
4667                }
4668                // SAFETY: disjoint row ranges per worker.
4669                unsafe {
4670                    *p1.at(r) = v1;
4671                    *p2.at(r) = v2;
4672                }
4673            }
4674        };
4675        dispatch_rows(pool, rows, &run);
4676        return;
4677    }
4678    let run = move |start: usize, end: usize| {
4679        for r in start..end {
4680            // SAFETY: disjoint row ranges per worker.
4681            unsafe {
4682                *p1.at(r) = q4t_row_exact(bytes, r, gpr, x1);
4683                *p2.at(r) = q4t_row_exact(bytes, r, gpr, x2);
4684            }
4685        }
4686    };
4687    dispatch_rows(pool, rows, &run);
4688}
4689
4690/// Batched q4_tiled matmat: each row's tiles stream once per microbatch.
4691#[allow(clippy::too_many_arguments)]
4692/// Prefill GEMM through Accelerate for group-quantized codecs: a
4693/// caller-supplied row dequantizer fills f32 tiles (pool-parallel) and
4694/// each tile rides the AMX with one sgemm — the generic sibling of
4695/// `qmatmat_accel` (q8). Numerics are f32-GEMM (tolerance class);
4696/// decode (b=1) never takes this path.
4697#[cfg(target_os = "macos")]
4698fn dequant_matmat_accel(
4699    dequant_row: &(dyn Fn(usize, &mut [f32]) + Sync),
4700    xs_all: &[f32],
4701    b: usize,
4702    rows: usize,
4703    cols: usize,
4704    out: &mut [f32],
4705    pool: Option<&Pool>,
4706) {
4707    const TR: usize = 2048;
4708    thread_local! {
4709        static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
4710    }
4711    WTILE.with(|wt| {
4712        let mut wtile = wt.borrow_mut();
4713        wtile.resize(TR * cols, 0.0);
4714        let mut r0 = 0usize;
4715        while r0 < rows {
4716            let tr = TR.min(rows - r0);
4717            let wt_addr = SendMut(wtile.as_mut_ptr());
4718            let run = |start: usize, end: usize| {
4719                for r in start..end {
4720                    // SAFETY: workers cover disjoint r ranges.
4721                    let dst = unsafe { std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols) };
4722                    dequant_row(r0 + r, dst);
4723                }
4724            };
4725            dispatch_rows(pool, tr, &run);
4726            unsafe {
4727                accel_blas::cblas_sgemm(
4728                    101, // RowMajor
4729                    111, // NoTrans A
4730                    112, // Trans B
4731                    b as i32,
4732                    tr as i32,
4733                    cols as i32,
4734                    1.0,
4735                    xs_all.as_ptr(),
4736                    cols as i32,
4737                    wtile.as_ptr(),
4738                    cols as i32,
4739                    0.0,
4740                    out.as_mut_ptr().add(r0),
4741                    rows as i32,
4742                );
4743            }
4744            r0 += tr;
4745        }
4746    });
4747}
4748
4749fn q4t_matmat(
4750    bytes: &[u8],
4751    xs_all: &[f32],
4752    b: usize,
4753    rows: usize,
4754    cols: usize,
4755    out: &mut [f32],
4756    pool: Option<&Pool>,
4757) {
4758    debug_assert_eq!(out.len(), b * rows);
4759    let gpr = cols / GROUP_SIZE;
4760    // Wide batches ride the AMX like q8's qmatmat: on Apple silicon
4761    // the dequant-tile sgemm is an order above the SDOT row loop for
4762    // prefill shapes (imagegen DiT forwards are exactly this).
4763    #[cfg(target_os = "macos")]
4764    if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
4765        dequant_matmat_accel(
4766            &|r, dst| {
4767                for gi in 0..gpr {
4768                    let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
4769                    let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
4770                    for (k, &bb) in tile[2..].iter().enumerate() {
4771                        dst[gi * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s;
4772                        dst[gi * GROUP_SIZE + k * 2 + 1] = (((bb >> 4) & 0x0F) as f32 - 8.0) * s;
4773                    }
4774                }
4775            },
4776            xs_all,
4777            b,
4778            rows,
4779            cols,
4780            out,
4781            pool,
4782        );
4783        return;
4784    }
4785    let out_addr = SendMut(out.as_mut_ptr());
4786    if a8w8_enabled() {
4787        let acts: Vec<SplitAct> = (0..b)
4788            .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
4789            .collect();
4790        let acts = &acts;
4791        #[cfg(target_arch = "x86_64")]
4792        let blocked_ok = avx2_enabled()
4793            && blocked_enabled();
4794        #[cfg(target_arch = "aarch64")]
4795        let blocked_ok = sdot_enabled()
4796            && blocked_enabled();
4797        #[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))]
4798        let blocked_ok = false;
4799        let run = move |start: usize, end: usize| {
4800            for r in start..end {
4801                let mut bi = 0usize;
4802                #[cfg(target_arch = "aarch64")]
4803                if blocked_ok {
4804                    while bi + 4 <= acts.len() {
4805                        let xs = [
4806                            acts[bi].xq.as_slice(),
4807                            acts[bi + 1].xq.as_slice(),
4808                            acts[bi + 2].xq.as_slice(),
4809                            acts[bi + 3].xq.as_slice(),
4810                        ];
4811                        let d = unsafe { dot_q4t_row_1x4_sdot(bytes, r, gpr, xs) };
4812                        for k in 0..4 {
4813                            let act = &acts[bi + k];
4814                            let mut acc = d[k] * act.sx;
4815                            for &(j, xv) in &act.outliers {
4816                                let (w, sc) = q4t_outlier(bytes, r, gpr, j);
4817                                acc += w * sc * xv;
4818                            }
4819                            // SAFETY: disjoint (bi, r) cells per worker.
4820                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
4821                        }
4822                        bi += 4;
4823                    }
4824                }
4825                #[cfg(target_arch = "x86_64")]
4826                if blocked_ok {
4827                    while bi + 4 <= acts.len() {
4828                        let xs = [
4829                            acts[bi].xq.as_slice(),
4830                            acts[bi + 1].xq.as_slice(),
4831                            acts[bi + 2].xq.as_slice(),
4832                            acts[bi + 3].xq.as_slice(),
4833                        ];
4834                        let d = unsafe {
4835                            if vnni_tiles_enabled() {
4836                                dot_q4t_row_1x4_vnni(bytes, r, gpr, xs)
4837                            } else {
4838                                dot_q4t_row_1x4_avx2(bytes, r, gpr, xs)
4839                            }
4840                        };
4841                        for k in 0..4 {
4842                            let act = &acts[bi + k];
4843                            let mut acc = d[k] * act.sx;
4844                            for &(j, xv) in &act.outliers {
4845                                let (w, sc) = q4t_outlier(bytes, r, gpr, j);
4846                                acc += w * sc * xv;
4847                            }
4848                            // SAFETY: disjoint (bi, r) cells per worker.
4849                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
4850                        }
4851                        bi += 4;
4852                    }
4853                }
4854                let _ = blocked_ok;
4855                while bi < acts.len() {
4856                    let act = &acts[bi];
4857                    let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
4858                    for &(j, xv) in &act.outliers {
4859                        let (w, s) = q4t_outlier(bytes, r, gpr, j);
4860                        acc += w * s * xv;
4861                    }
4862                    // SAFETY: disjoint (bi, r) cells per worker range.
4863                    unsafe { *out_addr.at(bi * rows + r) = acc };
4864                    bi += 1;
4865                }
4866            }
4867        };
4868        dispatch_rows(pool, rows, &run);
4869        return;
4870    }
4871    let run = move |start: usize, end: usize| {
4872        for r in start..end {
4873            for bi in 0..b {
4874                let x = &xs_all[bi * cols..(bi + 1) * cols];
4875                // SAFETY: disjoint (bi, r) cells per worker range.
4876                unsafe { *out_addr.at(bi * rows + r) = q4t_row_exact(bytes, r, gpr, x) };
4877            }
4878        }
4879    };
4880    dispatch_rows(pool, rows, &run);
4881}
4882
4883// ── q1 (dtype 12): binary weights, [f16 scale][4B sign bits] per
4884// 32-group tile. The kernel family mirrors q4_tiled: one sequential
4885// stream of 6-byte tiles, per-tile integer dot × scale, exact outlier
4886// correction (A8W8 contract), exact scalar path under CMF_SDOT=0. ──
4887
4888/// Per-32-group sums of the quantized activation — the ±1 identity's
4889/// shared half: `dot = −2·sdot(mask, x) − gsum[g]`, computed ONCE per
4890/// matvec and reused by every row.
4891fn q1_group_sums(xq: &[i8], gpr: usize) -> Vec<i32> {
4892    (0..gpr)
4893        .map(|gi| {
4894            xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE]
4895                .iter()
4896                .map(|&v| v as i32)
4897                .sum()
4898        })
4899        .collect()
4900}
4901
4902/// One q1 row via the A8W8 int8 path — mask-SDOT on ARM (no ±1
4903/// expansion at all), scalar bit loop elsewhere (AVX2 queued with the
4904/// x86 pass).
4905#[inline]
4906#[allow(unreachable_code)]
4907/// AVX2 q1 row via the same ±1 identity as the ARM sdot kernel: the
4908/// sign bits expand to a {0, −1} byte mask through shuffle+cmpeq, the
4909/// masked activation sums through maddubs(1, x&mask), and
4910/// `dot = −(2·masked_sum + Σx_group)` — bit-identical integer math.
4911#[cfg(target_arch = "x86_64")]
4912#[target_feature(enable = "avx2")]
4913unsafe fn dot_q1_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
4914    // SAFETY: callers uphold the 6B-tile and xq/gsum length contracts.
4915    unsafe {
4916        use core::arch::x86_64::*;
4917        // Byte j of the mask must replicate bits-byte j/8.
4918        let expand = _mm256_setr_epi8(
4919            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,
4920            3, 3, 3,
4921        );
4922        let bitsel = _mm256_setr_epi8(
4923            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
4924            -128, 1, 2, 4, 8, 16, 32, 64, -128,
4925        );
4926        let ones8 = _mm256_set1_epi8(1);
4927        let ones16 = _mm256_set1_epi16(1);
4928        let mut acc = 0f32;
4929        for gi in 0..gpr {
4930            let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
4931            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4932            let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
4933            let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
4934            let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
4935            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4936            let sel = _mm256_and_si256(x, mask);
4937            // Σ of selected i8 lanes: maddubs(1u8, sel_i8) pairs → madd.
4938            let p16 = _mm256_maddubs_epi16(ones8, sel);
4939            let d32 = _mm256_madd_epi16(p16, ones16);
4940            let hi128 = _mm256_extracti128_si256::<1>(d32);
4941            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
4942            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
4943            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
4944            let msum = _mm_cvtsi128_si32(s32);
4945            // The and-select keeps x UN-negated (unlike ARM's −1-mask
4946            // sdot): d = Σ_set − Σ_unset = 2·Σ_set − Σ_all.
4947            let d = 2 * msum - gsum[gi];
4948            acc += d as f32 * s;
4949        }
4950        acc
4951    }
4952}
4953
4954/// VNNI twin of `dot_q1_row_avx2`: the masked-select sum goes through
4955/// one `vpdpbusd(1u8, sel)` (see `dpbusd_hsum` — bit-identical).
4956#[cfg(target_arch = "x86_64")]
4957#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
4958unsafe fn dot_q1_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
4959    // SAFETY: callers uphold the 6B-tile and xq/gsum length contracts.
4960    unsafe {
4961        use core::arch::x86_64::*;
4962        let expand = _mm256_setr_epi8(
4963            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,
4964            3, 3, 3,
4965        );
4966        let bitsel = _mm256_setr_epi8(
4967            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
4968            -128, 1, 2, 4, 8, 16, 32, 64, -128,
4969        );
4970        let ones8 = _mm256_set1_epi8(1);
4971        let mut acc = 0f32;
4972        for gi in 0..gpr {
4973            let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
4974            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4975            let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
4976            let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
4977            let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
4978            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4979            let msum = dpbusd_hsum(ones8, _mm256_and_si256(x, mask));
4980            let d = 2 * msum - gsum[gi];
4981            acc += d as f32 * s;
4982        }
4983        acc
4984    }
4985}
4986
4987/// VNNI twin of `dot_q1_row_1x4_avx2` (see `dpbusd_hsum`).
4988#[cfg(target_arch = "x86_64")]
4989#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
4990unsafe fn dot_q1_row_1x4_vnni(
4991    bytes: &[u8],
4992    r: usize,
4993    gpr: usize,
4994    xs: [&[i8]; 4],
4995    gsums: [&[i32]; 4],
4996) -> [f32; 4] {
4997    // SAFETY: callers uphold the 6B-tile and xq/gsum length contracts.
4998    unsafe {
4999        use core::arch::x86_64::*;
5000        let expand = _mm256_setr_epi8(
5001            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,
5002            3, 3, 3,
5003        );
5004        let bitsel = _mm256_setr_epi8(
5005            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
5006            -128, 1, 2, 4, 8, 16, 32, 64, -128,
5007        );
5008        let ones8 = _mm256_set1_epi8(1);
5009        let mut acc = [0f32; 4];
5010        for gi in 0..gpr {
5011            let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
5012            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5013            let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
5014            let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
5015            let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
5016            for (k, xq) in xs.iter().enumerate() {
5017                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5018                let msum = dpbusd_hsum(ones8, _mm256_and_si256(x, mask));
5019                let d = 2 * msum - gsums[k][gi];
5020                acc[k] += d as f32 * s;
5021            }
5022        }
5023        acc
5024    }
5025}
5026
5027/// The blocked 1×4 flavor: the expanded bit mask serves four activation
5028/// streams per group (mask build once, four select+reduce chains).
5029#[cfg(target_arch = "x86_64")]
5030#[target_feature(enable = "avx2")]
5031unsafe fn dot_q1_row_1x4_avx2(
5032    bytes: &[u8],
5033    r: usize,
5034    gpr: usize,
5035    xs: [&[i8]; 4],
5036    gsums: [&[i32]; 4],
5037) -> [f32; 4] {
5038    // SAFETY: callers uphold the 6B-tile and xq/gsum length contracts.
5039    unsafe {
5040        use core::arch::x86_64::*;
5041        let expand = _mm256_setr_epi8(
5042            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,
5043            3, 3, 3,
5044        );
5045        let bitsel = _mm256_setr_epi8(
5046            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
5047            -128, 1, 2, 4, 8, 16, 32, 64, -128,
5048        );
5049        let ones8 = _mm256_set1_epi8(1);
5050        let ones16 = _mm256_set1_epi16(1);
5051        let mut acc = [0f32; 4];
5052        for gi in 0..gpr {
5053            let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
5054            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5055            let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
5056            let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
5057            let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
5058            for (k, xq) in xs.iter().enumerate() {
5059                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5060                let sel = _mm256_and_si256(x, mask);
5061                let p16 = _mm256_maddubs_epi16(ones8, sel);
5062                let d32 = _mm256_madd_epi16(p16, ones16);
5063                let hi128 = _mm256_extracti128_si256::<1>(d32);
5064                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
5065                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
5066                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
5067                let msum = _mm_cvtsi128_si32(s32);
5068                let d = 2 * msum - gsums[k][gi];
5069                acc[k] += d as f32 * s;
5070            }
5071        }
5072        acc
5073    }
5074}
5075
5076#[allow(unreachable_code)]
5077fn dot_q1_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
5078    #[cfg(target_arch = "aarch64")]
5079    unsafe {
5080        return dot_q1_row_sdot(bytes, r, gpr, xq, gsum);
5081    }
5082    #[cfg(target_arch = "x86_64")]
5083    if avx2_enabled() {
5084        unsafe {
5085            if vnni_tiles_enabled() {
5086                return dot_q1_row_vnni(bytes, r, gpr, xq, gsum);
5087            }
5088            return dot_q1_row_avx2(bytes, r, gpr, xq, gsum);
5089        }
5090    }
5091    let _ = gsum;
5092    let mut acc = 0f32;
5093    for gi in 0..gpr {
5094        let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
5095        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
5096        let mut d = 0i32;
5097        for (j, &b) in tile[2..].iter().enumerate() {
5098            for k in 0..8 {
5099                let w = ((b >> k) & 1) as i32 * 2 - 1;
5100                d += w * xq[gi * GROUP_SIZE + j * 8 + k] as i32;
5101            }
5102        }
5103        acc += d as f32 * s;
5104    }
5105    acc
5106}
5107
5108/// SDOT q1 row via the ±1 identity: the vtst mask (0xFF where the bit
5109/// is set, i.e. −1 as i8) feeds `sdot` DIRECTLY — no expansion to ±1
5110/// lanes at all — and `dot = −(2·sdot(mask, x) + Σx_group)`, with the
5111/// per-group activation sums shared across every row of the matvec.
5112/// Four tiles (128 weights) per iteration: integer dots reduce through
5113/// a vpaddq tree into ONE i32x4 that meets its four scales in a single
5114/// fused f32 multiply-add. Integer math throughout — bit-identical to
5115/// the scalar ±1 reference.
5116#[cfg(target_arch = "aarch64")]
5117#[target_feature(enable = "neon,dotprod")]
5118unsafe fn dot_q1_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
5119    // SAFETY: callers uphold slice-length contracts (6B tile per group,
5120    // xq.len() == gpr·GROUP_SIZE, gsum.len() == gpr).
5121    unsafe {
5122        use core::arch::aarch64::*;
5123        use core::arch::asm;
5124        const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
5125        let m = vld1q_u8(MASKS.as_ptr());
5126        // One tile's −Σ_set(x) as an UNREDUCED i32x4 (two mask-sdots).
5127        macro_rules! tile_dot {
5128            ($t:expr, $x:expr) => {{
5129                let v0 = vcombine_u8(vdup_n_u8(*$t.add(2)), vdup_n_u8(*$t.add(3)));
5130                let v1 = vcombine_u8(vdup_n_u8(*$t.add(4)), vdup_n_u8(*$t.add(5)));
5131                let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
5132                let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
5133                let x0 = vld1q_s8($x);
5134                let x1 = vld1q_s8($x.add(16));
5135                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5136                asm!(
5137                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
5138                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
5139                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5140                    w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
5141                    options(pure, nomem, nostack),
5142                );
5143                vaddq_s32(a0, a1)
5144            }};
5145        }
5146        // TBL unpack over PAIR loads: one vld1q covers two 6B tiles
5147        // ([s s b b b b][s s b b b b] + 4B slack), TBL replicates each
5148        // bit-byte across 8 lanes for vtst, and the four scales gather
5149        // through tbl2 into one fcvtl — the 16 ld1r broadcast loads and
5150        // 4 branchy software f16 conversions per 128 weights (the
5151        // measured load-port wall of this kernel) become 2 vector
5152        // loads + 9 table lookups. Integer math order is unchanged —
5153        // bit-identical results (FCVTL is exact on every f16).
5154        const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
5155        const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
5156        const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
5157        const IW11: [u8; 16] = [
5158            10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
5159        ];
5160        const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
5161        let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
5162        let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
5163        let isc = vld1_u8(ISC.as_ptr());
5164        // One tile's −Σ_set(x) from a TBL-unpacked pair load.
5165        macro_rules! tile_dot_tbl {
5166            ($ld:expr, $i0:expr, $i1:expr, $x:expr) => {{
5167                let w0 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i0), m));
5168                let w1 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i1), m));
5169                let x0 = vld1q_s8($x);
5170                let x1 = vld1q_s8($x.add(16));
5171                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5172                asm!(
5173                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
5174                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
5175                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5176                    w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
5177                    options(pure, nomem, nostack),
5178                );
5179                vaddq_s32(a0, a1)
5180            }};
5181        }
5182        let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
5183        let row_base = r * gpr * Q1_TILE;
5184        let abs_end = bytes.len();
5185        let xp = xq.as_ptr();
5186        let gp = gsum.as_ptr();
5187        let mut accv = vdupq_n_f32(0.0);
5188        let mut gi = 0;
5189        // The second pair load reads 4B past tile gi+3 — stay inside
5190        // the payload slice (only the file's final tiles fall back).
5191        while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
5192            let t0 = base.add(gi * Q1_TILE);
5193            let ld_a = vld1q_u8(t0);
5194            let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
5195            let d0 = tile_dot_tbl!(ld_a, iw00, iw01, xp.add(gi * GROUP_SIZE));
5196            let d1 = tile_dot_tbl!(ld_a, iw10, iw11, xp.add((gi + 1) * GROUP_SIZE));
5197            let d2 = tile_dot_tbl!(ld_b, iw00, iw01, xp.add((gi + 2) * GROUP_SIZE));
5198            let d3 = tile_dot_tbl!(ld_b, iw10, iw11, xp.add((gi + 3) * GROUP_SIZE));
5199            // [−Σ0, −Σ1, −Σ2, −Σ3] → dots = −(2·Σset_neg + gsum)
5200            let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
5201            let g = vld1q_s32(gp.add(gi));
5202            let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
5203            let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
5204            let scf: float32x4_t;
5205            asm!(
5206                "fcvtl {o:v}.4s, {i:v}.4h",
5207                o = out(vreg) scf, i = in(vreg) sc16,
5208                options(pure, nomem, nostack),
5209            );
5210            accv = vfmaq_f32(accv, vcvtq_f32_s32(dots), scf);
5211            gi += 4;
5212        }
5213        let mut acc = vaddvq_f32(accv);
5214        while gi < gpr {
5215            let t = base.add(gi * Q1_TILE);
5216            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5217            let d = vaddvq_s32(tile_dot!(t, xp.add(gi * GROUP_SIZE)));
5218            acc += (-(2 * d + *gp.add(gi))) as f32 * s;
5219            gi += 1;
5220        }
5221        acc
5222    }
5223}
5224
5225/// Blocked q1 1×4: one TBL unpack of the tile pair serves FOUR
5226/// activation streams (prefill amortization — the same idea as the
5227/// AVX2 twin; per stream the group order, fma order and tail match the
5228/// single-row kernel exactly, so batch == matvec bit-for-bit).
5229#[cfg(target_arch = "aarch64")]
5230#[target_feature(enable = "neon,dotprod")]
5231unsafe fn dot_q1_row_1x4_sdot(
5232    bytes: &[u8],
5233    r: usize,
5234    gpr: usize,
5235    xs: [&[i8]; 4],
5236    gs: [&[i32]; 4],
5237) -> [f32; 4] {
5238    // SAFETY: same slice-length contracts as `dot_q1_row_sdot`, ×4.
5239    unsafe {
5240        use core::arch::aarch64::*;
5241        use core::arch::asm;
5242        const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
5243        const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
5244        const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
5245        const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
5246        const IW11: [u8; 16] = [
5247            10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
5248        ];
5249        const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
5250        let m = vld1q_u8(MASKS.as_ptr());
5251        let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
5252        let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
5253        let isc = vld1_u8(ISC.as_ptr());
5254        macro_rules! sdot2 {
5255            ($w0:expr, $w1:expr, $x:expr) => {{
5256                let x0 = vld1q_s8($x);
5257                let x1 = vld1q_s8($x.add(16));
5258                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5259                asm!(
5260                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
5261                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
5262                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5263                    w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
5264                    options(pure, nomem, nostack),
5265                );
5266                vaddq_s32(a0, a1)
5267            }};
5268        }
5269        let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
5270        let row_base = r * gpr * Q1_TILE;
5271        let abs_end = bytes.len();
5272        let mut accv = [vdupq_n_f32(0.0); 4];
5273        let mut gi = 0;
5274        while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
5275            let t0 = base.add(gi * Q1_TILE);
5276            let ld_a = vld1q_u8(t0);
5277            let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
5278            // Unpack ONCE — eight ±mask vectors serve all four streams.
5279            let w00 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw00), m));
5280            let w01 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw01), m));
5281            let w10 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw10), m));
5282            let w11 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw11), m));
5283            let w20 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw00), m));
5284            let w21 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw01), m));
5285            let w30 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw10), m));
5286            let w31 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw11), m));
5287            let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
5288            let scf: float32x4_t;
5289            asm!(
5290                "fcvtl {o:v}.4s, {i:v}.4h",
5291                o = out(vreg) scf, i = in(vreg) sc16,
5292                options(pure, nomem, nostack),
5293            );
5294            for k in 0..4 {
5295                let xp = xs[k].as_ptr();
5296                let d0 = sdot2!(w00, w01, xp.add(gi * GROUP_SIZE));
5297                let d1 = sdot2!(w10, w11, xp.add((gi + 1) * GROUP_SIZE));
5298                let d2 = sdot2!(w20, w21, xp.add((gi + 2) * GROUP_SIZE));
5299                let d3 = sdot2!(w30, w31, xp.add((gi + 3) * GROUP_SIZE));
5300                let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
5301                let g = vld1q_s32(gs[k].as_ptr().add(gi));
5302                let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
5303                accv[k] = vfmaq_f32(accv[k], vcvtq_f32_s32(dots), scf);
5304            }
5305            gi += 4;
5306        }
5307        let mut acc = [
5308            vaddvq_f32(accv[0]),
5309            vaddvq_f32(accv[1]),
5310            vaddvq_f32(accv[2]),
5311            vaddvq_f32(accv[3]),
5312        ];
5313        while gi < gpr {
5314            let t = base.add(gi * Q1_TILE);
5315            let sc = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5316            let v0 = vcombine_u8(vdup_n_u8(*t.add(2)), vdup_n_u8(*t.add(3)));
5317            let v1 = vcombine_u8(vdup_n_u8(*t.add(4)), vdup_n_u8(*t.add(5)));
5318            let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
5319            let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
5320            for k in 0..4 {
5321                let d = vaddvq_s32(sdot2!(w0, w1, xs[k].as_ptr().add(gi * GROUP_SIZE)));
5322                acc[k] += (-(2 * d + *gs[k].as_ptr().add(gi))) as f32 * sc;
5323            }
5324            gi += 1;
5325        }
5326        acc
5327    }
5328}
5329
5330/// (weight ±1, scale) of one q1 element — the exact outlier term.
5331#[inline]
5332fn q1_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
5333    let gi = j / GROUP_SIZE;
5334    let k = j % GROUP_SIZE;
5335    let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
5336    let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
5337    let bit = (tile[2 + k / 8] >> (k % 8)) & 1;
5338    ((bit as i32 * 2 - 1) as f32, s)
5339}
5340
5341/// Exact scalar q1 row (CMF_SDOT=0 contract).
5342#[inline]
5343fn q1_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
5344    let mut acc = 0f32;
5345    for gi in 0..gpr {
5346        let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
5347        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
5348        let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
5349        let mut ga = 0f32;
5350        for (j, &b) in tile[2..].iter().enumerate() {
5351            for k in 0..8 {
5352                ga += (((b >> k) & 1) as f32 * 2.0 - 1.0) * xg[j * 8 + k];
5353            }
5354        }
5355        acc += ga * s;
5356    }
5357    acc
5358}
5359
5360/// One q1 row range via A8W8 (the body of `q1_matvec`'s hot loop,
5361/// extracted so multi-matrix jobs drive the same kernel).
5362#[allow(clippy::too_many_arguments)]
5363fn q1_range_a8w8(
5364    bytes: &[u8],
5365    gpr: usize,
5366    act: &SplitAct,
5367    gsum: &[i32],
5368    out: SendMut,
5369    start: usize,
5370    end: usize,
5371) {
5372    for r in start..end {
5373        let mut acc = dot_q1_row_i8(bytes, r, gpr, &act.xq, gsum) * act.sx;
5374        for &(j, xv) in &act.outliers {
5375            let (w, s) = q1_outlier(bytes, r, gpr, j);
5376            acc += w * s * xv;
5377        }
5378        // SAFETY: disjoint row ranges per worker.
5379        unsafe { *out.at(r) = acc };
5380    }
5381}
5382
5383/// Exact-scalar q1 row range (CMF_SDOT=0 contract).
5384fn q1_range_f32(bytes: &[u8], gpr: usize, x: &[f32], out: SendMut, start: usize, end: usize) {
5385    for r in start..end {
5386        // SAFETY: disjoint row ranges per worker.
5387        unsafe { *out.at(r) = q1_row_exact(bytes, r, gpr, x) };
5388    }
5389}
5390
5391/// q1t per-row overlay locator. After the base (`base_len`) come
5392/// `[u32 row_ptr[rows+1]]` then `[(u16 col, f16 val)]` grouped by row (row
5393/// `r`'s entries are `[row_ptr[r], row_ptr[r+1])`). Returns
5394/// `(row_ptr offset, entries offset, present)`.
5395fn q1t_overlay(bytes: &[u8], base_len: usize, rows: usize) -> (usize, usize, bool) {
5396    let entries = base_len + (rows + 1) * 4;
5397    (base_len, entries, entries <= bytes.len())
5398}
5399
5400/// Read `row_ptr[r]` from the overlay's prefix-sum table.
5401#[inline]
5402fn q1t_rowptr(bytes: &[u8], rp_off: usize, r: usize) -> usize {
5403    let o = rp_off + r * 4;
5404    u32::from_le_bytes([bytes[o], bytes[o + 1], bytes[o + 2], bytes[o + 3]]) as usize
5405}
5406
5407/// Byte → the 5 ternary signs it packs `{−1,0,+1}` as f32, precomputed so
5408/// decoding a q1t code is a table load, not the base-3 divide/modulo per
5409/// weight (division is ~20–40× the cost of a load). Built at compile time.
5410const SIGN5: [[f32; 5]; 256] = {
5411    let mut lut = [[0.0f32; 5]; 256];
5412    let pow3 = [1u16, 3, 9, 27, 81];
5413    let mut byte = 0usize;
5414    while byte < 256 {
5415        let mut i = 0usize;
5416        while i < 5 {
5417            let code = (byte as u16 / pow3[i]) % 3;
5418            lut[byte][i] = if code == 1 {
5419                1.0
5420            } else if code == 2 {
5421                -1.0
5422            } else {
5423                0.0
5424            };
5425            i += 1;
5426        }
5427        byte += 1;
5428    }
5429    lut
5430};
5431
5432/// Same table, as i8 signs — the operand for the int8 SDOT base kernel.
5433const SIGN5_I8: [[i8; 5]; 256] = {
5434    let mut lut = [[0i8; 5]; 256];
5435    let pow3 = [1u16, 3, 9, 27, 81];
5436    let mut byte = 0usize;
5437    while byte < 256 {
5438        let mut i = 0usize;
5439        while i < 5 {
5440            let code = (byte as u16 / pow3[i]) % 3;
5441            lut[byte][i] = if code == 1 {
5442                1
5443            } else if code == 2 {
5444                -1
5445            } else {
5446                0
5447            };
5448            i += 1;
5449        }
5450        byte += 1;
5451    }
5452    lut
5453};
5454
5455/// The same 5 i8 signs packed into a u64 (`[s0 s1 s2 s3 s4 0 0 0]`, LE) so the
5456/// group unpack is 7 unaligned u64 stores at offsets 0,5,10,…,30 instead of
5457/// six 5-byte copies + LUT indexing — each store's trailing zeros are fixed by
5458/// the next store, and the last one runs 6 B past the 32nd weight (the unpack
5459/// buffer is padded to 40). This is the decode/prefill hot inner op.
5460const SIGN5_U64: [u64; 256] = {
5461    let mut lut = [0u64; 256];
5462    let pow3 = [1u16, 3, 9, 27, 81];
5463    let mut byte = 0usize;
5464    while byte < 256 {
5465        let mut v = 0u64;
5466        let mut i = 0usize;
5467        while i < 5 {
5468            let code = (byte as u16 / pow3[i]) % 3;
5469            let s: u8 = if code == 1 {
5470                1
5471            } else if code == 2 {
5472                0xFF
5473            } else {
5474                0
5475            };
5476            v |= (s as u64) << (i * 8);
5477            i += 1;
5478        }
5479        lut[byte] = v;
5480        byte += 1;
5481    }
5482    lut
5483};
5484
5485/// Ternary base weight at `(row r, col j)` = `sign(code)·s_group`. Used to add
5486/// back activation-outlier columns, whose `x` was zeroed for the int8 bulk dot
5487/// (`split_act`). At a weight-outlier position the code is 0, so this is 0 and
5488/// the overlay correction owns that column — no double counting.
5489#[inline]
5490fn q1t_base_weight(bytes: &[u8], r: usize, gpr: usize, j: usize) -> f32 {
5491    const TILE: usize = cortiq_core::quant::Q1T_TILE;
5492    let off = (r * gpr + j / GROUP_SIZE) * TILE;
5493    let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
5494    let within = j % GROUP_SIZE;
5495    SIGN5[bytes[off + 2 + within / 5] as usize][within % 5] * s
5496}
5497
5498/// One 32-group int8 dot via two SDOTs. Bit-exact vs the scalar i8 sum
5499/// (integer accumulation is order-independent).
5500#[cfg(target_arch = "aarch64")]
5501#[target_feature(enable = "neon,dotprod")]
5502#[inline]
5503unsafe fn sdot32_i8(w: *const i8, x: *const i8) -> i32 {
5504    // SAFETY: caller guarantees 32 readable i8 at each pointer.
5505    unsafe {
5506        use core::arch::aarch64::*;
5507        use core::arch::asm;
5508        let w0 = vld1q_s8(w);
5509        let w1 = vld1q_s8(w.add(16));
5510        let x0 = vld1q_s8(x);
5511        let x1 = vld1q_s8(x.add(16));
5512        let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5513        asm!(
5514            "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
5515            "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
5516            a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5517            w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
5518            options(pure, nomem, nostack),
5519        );
5520        vaddvq_s32(vaddq_s32(a0, a1))
5521    }
5522}
5523
5524/// One 32-group int8 dot via AVX2: signed·signed as `maddubs(|w|, sign(x,w))`
5525/// then `madd` and a horizontal reduce (the same idiom as `dot_q4t_row_avx2`).
5526#[cfg(target_arch = "x86_64")]
5527#[target_feature(enable = "avx2")]
5528#[inline]
5529unsafe fn i8dot32_avx2(w: *const i8, x: *const i8) -> i32 {
5530    // SAFETY: caller guarantees 32 readable i8 at each pointer.
5531    unsafe {
5532        use core::arch::x86_64::*;
5533        let wv = _mm256_loadu_si256(w as *const __m256i);
5534        let xv = _mm256_loadu_si256(x as *const __m256i);
5535        let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
5536        let d = _mm256_madd_epi16(p16, _mm256_set1_epi16(1));
5537        let hi128 = _mm256_extracti128_si256::<1>(d);
5538        let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
5539        let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
5540        let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
5541        _mm_cvtsi128_si32(s32)
5542    }
5543}
5544
5545/// Unpack one q1t group's base-3 codes into 32 i8 signs via 7 unaligned u64
5546/// stores (see `SIGN5_U64`). `dst` MUST have ≥ 40 bytes: the 7th store writes
5547/// `dst[30..38]`. Stores go in order so each one's trailing zeros are
5548/// overwritten by the next; the final 6 padding bytes are unused by the dot.
5549#[inline]
5550fn q1t_unpack_group_i8(codes: *const u8, dst: &mut [i8]) {
5551    debug_assert!(dst.len() >= 40);
5552    // SAFETY: codes points at 7 readable bytes; dst has ≥ 40 bytes so every
5553    // 8-byte store at offset bi*5 (bi ≤ 6 → ≤ 30) stays in bounds.
5554    unsafe {
5555        let p = dst.as_mut_ptr();
5556        for bi in 0..7 {
5557            core::ptr::write_unaligned(
5558                p.add(bi * 5) as *mut u64,
5559                SIGN5_U64[*codes.add(bi) as usize],
5560            );
5561        }
5562    }
5563}
5564
5565/// One 32-group int8 dot, arch-dispatched (the matmat inner loop, where the
5566/// row's signs are unpacked once and dotted against every batch input).
5567/// Callers are gated by `a8w8_enabled()`, so the target-feature arms are
5568/// reachable; the scalar arm is a non-SIMD-arch fallback.
5569#[inline]
5570fn q1t_i8dot32(w: *const i8, x: *const i8) -> i32 {
5571    #[cfg(target_arch = "aarch64")]
5572    unsafe {
5573        return sdot32_i8(w, x);
5574    }
5575    #[cfg(target_arch = "x86_64")]
5576    unsafe {
5577        return i8dot32_avx2(w, x);
5578    }
5579    #[allow(unreachable_code)]
5580    unsafe {
5581        let mut s = 0i32;
5582        for k in 0..GROUP_SIZE {
5583            s += *w.add(k) as i32 * *x.add(k) as i32;
5584        }
5585        s
5586    }
5587}
5588
5589#[inline]
5590unsafe fn q1t_unpack_reg_u64s(codes: *const u8) -> (u64, u64, u64, u64) {
5591    let (s0, s1, s2, s3, s4, s5, s6) = unsafe {
5592        (
5593            SIGN5_U64[*codes as usize],
5594            SIGN5_U64[*codes.add(1) as usize],
5595            SIGN5_U64[*codes.add(2) as usize],
5596            SIGN5_U64[*codes.add(3) as usize],
5597            SIGN5_U64[*codes.add(4) as usize],
5598            SIGN5_U64[*codes.add(5) as usize],
5599            SIGN5_U64[*codes.add(6) as usize],
5600        )
5601    };
5602
5603    let u0 = s0 | (s1 << 40);
5604    let u1 = (s1 >> 24) | (s2 << 16) | (s3 << 56);
5605    let u2 = (s3 >> 8) | (s4 << 32);
5606    let u3 = (s4 >> 32) | (s5 << 8) | (s6 << 48);
5607
5608    (u0, u1, u2, u3)
5609}
5610
5611/// One q1t row's int8 base dot: `Σ_group s·dot(signs, xq)` (before the shared
5612/// `sx`). Direct register unpacking (zero stack stores/loads, no STLF stalls).
5613/// ARM SDOT.
5614#[cfg(target_arch = "aarch64")]
5615#[target_feature(enable = "neon,dotprod")]
5616unsafe fn q1t_dot_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
5617    use core::arch::aarch64::*;
5618    use core::arch::asm;
5619    unsafe {
5620        const TILE: usize = cortiq_core::quant::Q1T_TILE;
5621        let mut acc = 0f32;
5622        let bytes_ptr = bytes.as_ptr();
5623        let xq_ptr = xq.as_ptr();
5624        let row_off = r * gpr * TILE;
5625
5626        let gpr2 = gpr & !1;
5627        let mut gi = 0;
5628        while gi < gpr2 {
5629            let off0 = row_off + gi * TILE;
5630            let off1 = off0 + TILE;
5631            let s0 = f16_to_f32(u16::from_le_bytes([
5632                *bytes_ptr.add(off0),
5633                *bytes_ptr.add(off0 + 1),
5634            ]));
5635            let s1 = f16_to_f32(u16::from_le_bytes([
5636                *bytes_ptr.add(off1),
5637                *bytes_ptr.add(off1 + 1),
5638            ]));
5639
5640            let (u0_0, u1_0, u2_0, u3_0) = q1t_unpack_reg_u64s(bytes_ptr.add(off0 + 2));
5641            let (u0_1, u1_1, u2_1, u3_1) = q1t_unpack_reg_u64s(bytes_ptr.add(off1 + 2));
5642
5643            let w0_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_0), vcreate_u64(u1_0)));
5644            let w1_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_0), vcreate_u64(u3_0)));
5645            let w0_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_1), vcreate_u64(u1_1)));
5646            let w1_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_1), vcreate_u64(u3_1)));
5647
5648            let x0_0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE));
5649            let x1_0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE + 16));
5650            let x0_1 = vld1q_s8(xq_ptr.add((gi + 1) * GROUP_SIZE));
5651            let x1_1 = vld1q_s8(xq_ptr.add((gi + 1) * GROUP_SIZE + 16));
5652
5653            let (mut a0_0, mut a1_0) = (vdupq_n_s32(0), vdupq_n_s32(0));
5654            let (mut a0_1, mut a1_1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5655            asm!(
5656                "sdot {a0_0:v}.4s, {w0_0:v}.16b, {x0_0:v}.16b",
5657                "sdot {a1_0:v}.4s, {w1_0:v}.16b, {x1_0:v}.16b",
5658                "sdot {a0_1:v}.4s, {w0_1:v}.16b, {x0_1:v}.16b",
5659                "sdot {a1_1:v}.4s, {w1_1:v}.16b, {x1_1:v}.16b",
5660                a0_0 = inout(vreg) a0_0, a1_0 = inout(vreg) a1_0,
5661                a0_1 = inout(vreg) a0_1, a1_1 = inout(vreg) a1_1,
5662                w0_0 = in(vreg) w0_0, x0_0 = in(vreg) x0_0, w1_0 = in(vreg) w1_0, x1_0 = in(vreg) x1_0,
5663                w0_1 = in(vreg) w0_1, x0_1 = in(vreg) x0_1, w1_1 = in(vreg) w1_1, x1_1 = in(vreg) x1_1,
5664                options(pure, nomem, nostack),
5665            );
5666            let d0 = vaddvq_s32(vaddq_s32(a0_0, a1_0));
5667            let d1 = vaddvq_s32(vaddq_s32(a0_1, a1_1));
5668            acc += d0 as f32 * s0 + d1 as f32 * s1;
5669            gi += 2;
5670        }
5671
5672        if gi < gpr {
5673            let off = row_off + gi * TILE;
5674            let s = f16_to_f32(u16::from_le_bytes([
5675                *bytes_ptr.add(off),
5676                *bytes_ptr.add(off + 1),
5677            ]));
5678            let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
5679            let w0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0), vcreate_u64(u1)));
5680            let w1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2), vcreate_u64(u3)));
5681            let x0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE));
5682            let x1 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE + 16));
5683            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5684            asm!(
5685                "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
5686                "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
5687                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5688                w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
5689                options(pure, nomem, nostack),
5690            );
5691            let d = vaddvq_s32(vaddq_s32(a0, a1));
5692            acc += d as f32 * s;
5693        }
5694        acc
5695    }
5696}
5697
5698/// x86 AVX2 mirror of `q1t_dot_row_sdot` (maddubs int8 dot per group).
5699#[cfg(target_arch = "x86_64")]
5700#[target_feature(enable = "avx2")]
5701unsafe fn q1t_dot_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
5702    use core::arch::x86_64::*;
5703    unsafe {
5704        const TILE: usize = cortiq_core::quant::Q1T_TILE;
5705        let mut acc = 0f32;
5706        let bytes_ptr = bytes.as_ptr();
5707        let xq_ptr = xq.as_ptr();
5708        let row_off = r * gpr * TILE;
5709
5710        let ones = _mm256_set1_epi16(1);
5711        for gi in 0..gpr {
5712            let off = row_off + gi * TILE;
5713            let s = f16_to_f32(u16::from_le_bytes([
5714                *bytes_ptr.add(off),
5715                *bytes_ptr.add(off + 1),
5716            ]));
5717            let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
5718            let wv = _mm256_set_epi64x(u3 as i64, u2 as i64, u1 as i64, u0 as i64);
5719            let xv = _mm256_loadu_si256(xq_ptr.add(gi * GROUP_SIZE) as *const __m256i);
5720            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
5721            let d256 = _mm256_madd_epi16(p16, ones);
5722            let d128 = _mm_add_epi32(
5723                _mm256_castsi256_si128(d256),
5724                _mm256_extracti128_si256(d256, 1),
5725            );
5726            let d64 = _mm_add_epi32(d128, _mm_shuffle_epi32(d128, 0xee));
5727            let d32 = _mm_cvtsi128_si32(_mm_add_epi32(d64, _mm_shuffle_epi32(d64, 0x55)));
5728            acc += d32 as f32 * s;
5729        }
5730        acc
5731    }
5732}
5733
5734/// VNNI twin of `q1t_dot_row_avx2` (see `dpbusd_hsum`).
5735#[cfg(target_arch = "x86_64")]
5736#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
5737unsafe fn q1t_dot_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
5738    use core::arch::x86_64::*;
5739    // SAFETY: same tile/xq contracts as `q1t_dot_row_avx2`.
5740    unsafe {
5741        const TILE: usize = cortiq_core::quant::Q1T_TILE;
5742        let mut acc = 0f32;
5743        let bytes_ptr = bytes.as_ptr();
5744        let xq_ptr = xq.as_ptr();
5745        let row_off = r * gpr * TILE;
5746        for gi in 0..gpr {
5747            let off = row_off + gi * TILE;
5748            let s = f16_to_f32(u16::from_le_bytes([
5749                *bytes_ptr.add(off),
5750                *bytes_ptr.add(off + 1),
5751            ]));
5752            let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
5753            let wv = _mm256_set_epi64x(u3 as i64, u2 as i64, u1 as i64, u0 as i64);
5754            let xv = _mm256_loadu_si256(xq_ptr.add(gi * GROUP_SIZE) as *const __m256i);
5755            let d = dpbusd_hsum(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
5756            acc += d as f32 * s;
5757        }
5758        acc
5759    }
5760}
5761
5762/// Per-row int8 base dot, dispatched once per row (matvec decode hot path).
5763/// Callers are gated by `a8w8_enabled()`, so the target-feature kernels are
5764/// reachable.
5765#[inline]
5766fn q1t_dot_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
5767    #[cfg(target_arch = "aarch64")]
5768    unsafe {
5769        return q1t_dot_row_sdot(bytes, r, gpr, xq);
5770    }
5771    #[cfg(target_arch = "x86_64")]
5772    unsafe {
5773        if vnni_tiles_enabled() {
5774            return q1t_dot_row_vnni(bytes, r, gpr, xq);
5775        }
5776        return q1t_dot_row_avx2(bytes, r, gpr, xq);
5777    }
5778    #[allow(unreachable_code)]
5779    {
5780        const TILE: usize = cortiq_core::quant::Q1T_TILE;
5781        let mut acc = 0f32;
5782        let mut sg = [0i8; GROUP_SIZE + 8]; // +8 slack for the u64-store unpack
5783        for gi in 0..gpr {
5784            let off = (r * gpr + gi) * TILE;
5785            let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
5786            q1t_unpack_group_i8(bytes.as_ptr().wrapping_add(off + 2), &mut sg);
5787            let mut d = 0i32;
5788            for k in 0..GROUP_SIZE {
5789                d += sg[k] as i32 * xq[gi * GROUP_SIZE + k] as i32;
5790            }
5791            acc += d as f32 * s;
5792        }
5793        acc
5794    }
5795}
5796
5797/// Σ over a row's outliers of `value·x[col]` — the correction that adds the
5798/// overlay's exact weights on top of the base dot. INVARIANT: the encoder
5799/// writes ternary code 0 at every outlier position (`quantize_q1t`), so the
5800/// base contributes nothing there and this is a plain `value·x`, not
5801/// `(value − base)·x` — no scattered per-outlier scale read. Row `r`'s entries
5802/// are the contiguous slice `[row_ptr[r], row_ptr[r+1])`, so no binary search.
5803fn q1t_row_outlier_correction(
5804    bytes: &[u8],
5805    r: usize,
5806    rp_off: usize,
5807    entries_off: usize,
5808    has_ov: bool,
5809    x: &[f32],
5810) -> f32 {
5811    if !has_ov {
5812        return 0.0;
5813    }
5814    let (c0, c1) = (
5815        q1t_rowptr(bytes, rp_off, r),
5816        q1t_rowptr(bytes, rp_off, r + 1),
5817    );
5818    let mut corr = 0f32;
5819    for p in c0..c1 {
5820        let e = entries_off + p * 4;
5821        let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
5822        let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
5823        corr += val * x[col];
5824    }
5825    corr
5826}
5827
5828/// Dequantize one q1t row into `buf[..cols]` via the sign LUT (no division),
5829/// then apply the row's outliers (its `[row_ptr[r], row_ptr[r+1])` slice).
5830/// Used by the batched (prefill) path where the decode amortizes over the batch.
5831fn q1t_dequant_row(
5832    bytes: &[u8],
5833    r: usize,
5834    gpr: usize,
5835    rp_off: usize,
5836    entries_off: usize,
5837    has_ov: bool,
5838    buf: &mut [f32],
5839) {
5840    const TILE: usize = cortiq_core::quant::Q1T_TILE;
5841    for g in 0..gpr {
5842        let off = (r * gpr + g) * TILE;
5843        let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
5844        let codes = &bytes[off + 2..off + TILE];
5845        let bc = g * GROUP_SIZE;
5846        // 6 full bytes (30 codes) + a 7th byte holding the last 2.
5847        for bi in 0..6 {
5848            let lut = &SIGN5[codes[bi] as usize];
5849            let d = &mut buf[bc + bi * 5..bc + bi * 5 + 5];
5850            for i in 0..5 {
5851                d[i] = lut[i] * s;
5852            }
5853        }
5854        let lut = &SIGN5[codes[6] as usize];
5855        buf[bc + 30] = lut[0] * s;
5856        buf[bc + 31] = lut[1] * s;
5857    }
5858    if !has_ov {
5859        return;
5860    }
5861    let (c0, c1) = (
5862        q1t_rowptr(bytes, rp_off, r),
5863        q1t_rowptr(bytes, rp_off, r + 1),
5864    );
5865    for p in c0..c1 {
5866        let e = entries_off + p * 4;
5867        let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
5868        buf[col] = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
5869    }
5870}
5871
5872/// Add the sparse outlier overlay onto a base dot already in `out` (the GPU
5873/// computes the ternary base; the overlay stays on the CPU — its entries are
5874/// few and its per-row gather doesn't vectorize on the GPU). Row-parallel.
5875fn q1t_add_overlay(
5876    bytes: &[u8],
5877    x: &[f32],
5878    rows: usize,
5879    cols: usize,
5880    out: &mut [f32],
5881    pool: Option<&Pool>,
5882) {
5883    const TILE: usize = cortiq_core::quant::Q1T_TILE;
5884    let gpr = cols / GROUP_SIZE;
5885    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
5886    if !has_ov {
5887        return;
5888    }
5889    let out_addr = SendMut(out.as_mut_ptr());
5890    let run = move |start: usize, end: usize| {
5891        for r in start..end {
5892            let corr = q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
5893            // SAFETY: disjoint rows; add onto the base the GPU already wrote.
5894            unsafe { *out_addr.at(r) += corr };
5895        }
5896    };
5897    dispatch_rows(pool, rows, &run);
5898}
5899
5900/// Q1T row range via the A8W8 int8 path — shared activation split,
5901/// per-row: base SDOT dot + outlier correction + overlay.
5902#[allow(clippy::too_many_arguments)]
5903fn q1t_range_a8w8(
5904    bytes: &[u8],
5905    gpr: usize,
5906    rp_off: usize,
5907    ent_off: usize,
5908    has_ov: bool,
5909    act: &SplitAct,
5910    x: &[f32],
5911    out: SendMut,
5912    start: usize,
5913    end: usize,
5914) {
5915    for r in start..end {
5916        let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
5917        for &(j, xv) in &act.outliers {
5918            acc += q1t_base_weight(bytes, r, gpr, j) * xv;
5919        }
5920        acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
5921        // SAFETY: disjoint row ranges per worker.
5922        unsafe { *out.at(r) = acc };
5923    }
5924}
5925
5926/// Q1T row range via the f32 path (no SDOT) — for matvec_many batched
5927/// dispatch when a8w8 is unavailable.
5928#[allow(clippy::too_many_arguments)]
5929fn q1t_range_f32_batch(
5930    bytes: &[u8],
5931    gpr: usize,
5932    rp_off: usize,
5933    ent_off: usize,
5934    has_ov: bool,
5935    x: &[f32],
5936    out: SendMut,
5937    start: usize,
5938    end: usize,
5939) {
5940    const TILE: usize = cortiq_core::quant::Q1T_TILE;
5941    let mut sg = [0f32; GROUP_SIZE];
5942    for r in start..end {
5943        let mut acc = 0f32;
5944        for g in 0..gpr {
5945            let off = (r * gpr + g) * TILE;
5946            let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
5947            let codes = &bytes[off + 2..off + TILE];
5948            let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
5949            for bi in 0..6 {
5950                sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
5951            }
5952            let lut = &SIGN5[codes[6] as usize];
5953            sg[30] = lut[0];
5954            sg[31] = lut[1];
5955            let mut gsum = 0f32;
5956            for k in 0..GROUP_SIZE {
5957                gsum += sg[k] * xg[k];
5958            }
5959            acc += s * gsum;
5960        }
5961        acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
5962        // SAFETY: disjoint row ranges per worker.
5963        unsafe { *out.at(r) = acc };
5964    }
5965}
5966
5967/// Ternary (q1t) matvec — decode+dot straight from mmap, one group at a time:
5968/// no per-ROW buffer, no division (the sign LUT), and a tiny per-group sign
5969/// buffer so the 32-wide dot vectorizes. This is the decode hot path.
5970fn q1t_matvec(
5971    bytes: &[u8],
5972    x: &[f32],
5973    rows: usize,
5974    cols: usize,
5975    out: &mut [f32],
5976    pool: Option<&Pool>,
5977) {
5978    debug_assert_eq!(out.len(), rows);
5979    const TILE: usize = cortiq_core::quant::Q1T_TILE;
5980    let gpr = cols / GROUP_SIZE;
5981    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
5982    let out_addr = SendMut(out.as_mut_ptr());
5983    // int8 SDOT base dot (ARM dotprod): ~4× the f32 arithmetic. x → i8 once
5984    // (`split_act`), activation outliers added back exactly in f32, weight
5985    // overlay on top. ARM SDOT / x86 AVX2; CMF_SDOT=0 keeps the exact f32 path.
5986    if a8w8_enabled() {
5987        let act = split_act(x);
5988        let act = &act;
5989        let run = move |start: usize, end: usize| {
5990            for r in start..end {
5991                let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
5992                for &(j, xv) in &act.outliers {
5993                    acc += q1t_base_weight(bytes, r, gpr, j) * xv;
5994                }
5995                acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
5996                // SAFETY: disjoint row ranges per worker.
5997                unsafe { *out_addr.at(r) = acc };
5998            }
5999        };
6000        dispatch_rows(pool, rows, &run);
6001        return;
6002    }
6003    let run = move |start: usize, end: usize| {
6004        // Per-group signs, unpacked contiguously so the dot below is a clean
6005        // 32-wide reduction the autovectorizer turns into f32x4 FMAs — the
6006        // 5-values-per-byte base-3 layout won't SIMD in place.
6007        let mut sg = [0f32; GROUP_SIZE];
6008        for r in start..end {
6009            let mut acc = 0f32;
6010            for g in 0..gpr {
6011                let off = (r * gpr + g) * TILE;
6012                let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
6013                let codes = &bytes[off + 2..off + TILE];
6014                let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
6015                for bi in 0..6 {
6016                    sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
6017                }
6018                let lut = &SIGN5[codes[6] as usize];
6019                sg[30] = lut[0];
6020                sg[31] = lut[1];
6021                let mut gsum = 0f32;
6022                for k in 0..GROUP_SIZE {
6023                    gsum += sg[k] * xg[k];
6024                }
6025                acc += s * gsum;
6026            }
6027            acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
6028            unsafe { *out_addr.at(r) = acc };
6029        }
6030    };
6031    dispatch_rows(pool, rows, &run);
6032}
6033
6034/// Fused-pair twin of `q1t_dot_row_sdot`: ONE register unpack of the
6035/// ternary codes serves BOTH activation streams (the unpack chain is
6036/// the dominant per-row cost — MTP verify pairs paid it twice). Per
6037/// stream the group order and f32 accumulation match the single-row
6038/// kernel exactly, so pair == 2×matvec bit-for-bit.
6039#[cfg(target_arch = "aarch64")]
6040#[target_feature(enable = "neon,dotprod")]
6041unsafe fn q1t_dot_row_sdot2(bytes: &[u8], r: usize, gpr: usize, xa: &[i8], xb: &[i8]) -> [f32; 2] {
6042    use core::arch::aarch64::*;
6043    use core::arch::asm;
6044    // SAFETY: same slice-length contracts as `q1t_dot_row_sdot`, ×2.
6045    unsafe {
6046        const TILE: usize = cortiq_core::quant::Q1T_TILE;
6047        let bytes_ptr = bytes.as_ptr();
6048        let row_off = r * gpr * TILE;
6049        let xp = [xa.as_ptr(), xb.as_ptr()];
6050        let mut acc = [0f32; 2];
6051        macro_rules! sdot2 {
6052            ($w0:expr, $w1:expr, $x:expr) => {{
6053                let x0 = vld1q_s8($x);
6054                let x1 = vld1q_s8($x.add(16));
6055                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6056                asm!(
6057                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
6058                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
6059                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6060                    w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
6061                    options(pure, nomem, nostack),
6062                );
6063                vaddvq_s32(vaddq_s32(a0, a1))
6064            }};
6065        }
6066        let gpr2 = gpr & !1;
6067        let mut gi = 0;
6068        while gi < gpr2 {
6069            let off0 = row_off + gi * TILE;
6070            let off1 = off0 + TILE;
6071            let s0 = f16_to_f32(u16::from_le_bytes([
6072                *bytes_ptr.add(off0),
6073                *bytes_ptr.add(off0 + 1),
6074            ]));
6075            let s1 = f16_to_f32(u16::from_le_bytes([
6076                *bytes_ptr.add(off1),
6077                *bytes_ptr.add(off1 + 1),
6078            ]));
6079            let (u0_0, u1_0, u2_0, u3_0) = q1t_unpack_reg_u64s(bytes_ptr.add(off0 + 2));
6080            let (u0_1, u1_1, u2_1, u3_1) = q1t_unpack_reg_u64s(bytes_ptr.add(off1 + 2));
6081            let w0_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_0), vcreate_u64(u1_0)));
6082            let w1_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_0), vcreate_u64(u3_0)));
6083            let w0_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_1), vcreate_u64(u1_1)));
6084            let w1_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_1), vcreate_u64(u3_1)));
6085            for k in 0..2 {
6086                let d0 = sdot2!(w0_0, w1_0, xp[k].add(gi * GROUP_SIZE));
6087                let d1 = sdot2!(w0_1, w1_1, xp[k].add((gi + 1) * GROUP_SIZE));
6088                acc[k] += d0 as f32 * s0 + d1 as f32 * s1;
6089            }
6090            gi += 2;
6091        }
6092        if gi < gpr {
6093            let off = row_off + gi * TILE;
6094            let s = f16_to_f32(u16::from_le_bytes([
6095                *bytes_ptr.add(off),
6096                *bytes_ptr.add(off + 1),
6097            ]));
6098            let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
6099            let w0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0), vcreate_u64(u1)));
6100            let w1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2), vcreate_u64(u3)));
6101            for k in 0..2 {
6102                let d = sdot2!(w0, w1, xp[k].add(gi * GROUP_SIZE));
6103                acc[k] += d as f32 * s;
6104            }
6105        }
6106        acc
6107    }
6108}
6109
6110/// Fused Q1T pair matvec: ONE pass over the rows serves both
6111/// activation streams — on ARM the ternary register unpack happens
6112/// once per tile pair (`q1t_dot_row_sdot2`); elsewhere the second dot
6113/// rides the row's L1-warm tile bytes. Per stream the math matches
6114/// `q1t_matvec` exactly.
6115fn q1t_matvec2(
6116    bytes: &[u8],
6117    x1: &[f32],
6118    x2: &[f32],
6119    rows: usize,
6120    cols: usize,
6121    o1: &mut [f32],
6122    o2: &mut [f32],
6123    pool: Option<&Pool>,
6124) {
6125    debug_assert_eq!(o1.len(), rows);
6126    debug_assert_eq!(o2.len(), rows);
6127    const TILE: usize = cortiq_core::quant::Q1T_TILE;
6128    let gpr = cols / GROUP_SIZE;
6129    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
6130    let out1 = SendMut(o1.as_mut_ptr());
6131    let out2 = SendMut(o2.as_mut_ptr());
6132    if a8w8_enabled() {
6133        let a1 = split_act(x1);
6134        let a2 = split_act(x2);
6135        let (a1, a2) = (&a1, &a2);
6136        let run = move |start: usize, end: usize| {
6137            for r in start..end {
6138                #[cfg(target_arch = "aarch64")]
6139                // a8w8 on aarch64 ⇔ sdot_enabled(), so the kernel's
6140                // target features are present.
6141                let ds = unsafe { q1t_dot_row_sdot2(bytes, r, gpr, &a1.xq, &a2.xq) };
6142                #[cfg(not(target_arch = "aarch64"))]
6143                let ds = [
6144                    q1t_dot_row_i8(bytes, r, gpr, &a1.xq),
6145                    q1t_dot_row_i8(bytes, r, gpr, &a2.xq),
6146                ];
6147                let mut acc1 = ds[0] * a1.sx;
6148                for &(j, xv) in &a1.outliers {
6149                    acc1 += q1t_base_weight(bytes, r, gpr, j) * xv;
6150                }
6151                acc1 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x1);
6152                let mut acc2 = ds[1] * a2.sx;
6153                for &(j, xv) in &a2.outliers {
6154                    acc2 += q1t_base_weight(bytes, r, gpr, j) * xv;
6155                }
6156                acc2 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x2);
6157                // SAFETY: disjoint row ranges per worker.
6158                unsafe {
6159                    *out1.at(r) = acc1;
6160                    *out2.at(r) = acc2;
6161                }
6162            }
6163        };
6164        dispatch_rows(pool, rows, &run);
6165        return;
6166    }
6167    let run = move |start: usize, end: usize| {
6168        // Exact path (CMF_SDOT=0): unpack the sign LUT once per group,
6169        // dot both streams — same op order per stream as `q1t_matvec`.
6170        let mut sg = [0f32; GROUP_SIZE];
6171        for r in start..end {
6172            let mut acc1 = 0f32;
6173            let mut acc2 = 0f32;
6174            for g in 0..gpr {
6175                let off = (r * gpr + g) * TILE;
6176                let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
6177                let codes = &bytes[off + 2..off + TILE];
6178                for bi in 0..6 {
6179                    sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
6180                }
6181                let lut = &SIGN5[codes[6] as usize];
6182                sg[30] = lut[0];
6183                sg[31] = lut[1];
6184                let xg1 = &x1[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
6185                let xg2 = &x2[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
6186                let mut gsum1 = 0f32;
6187                for k in 0..GROUP_SIZE {
6188                    gsum1 += sg[k] * xg1[k];
6189                }
6190                acc1 += s * gsum1;
6191                let mut gsum2 = 0f32;
6192                for k in 0..GROUP_SIZE {
6193                    gsum2 += sg[k] * xg2[k];
6194                }
6195                acc2 += s * gsum2;
6196            }
6197            acc1 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x1);
6198            acc2 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x2);
6199            // SAFETY: disjoint row ranges per worker.
6200            unsafe {
6201                *out1.at(r) = acc1;
6202                *out2.at(r) = acc2;
6203            }
6204        }
6205    };
6206    dispatch_rows(pool, rows, &run);
6207}
6208
6209/// Ternary (q1t) matmat (prefill) — dequant each row once, dot the whole
6210/// batch against it (amortizes the per-row decode).
6211fn q1t_matmat(
6212    bytes: &[u8],
6213    xs: &[f32],
6214    b: usize,
6215    rows: usize,
6216    cols: usize,
6217    out: &mut [f32],
6218    pool: Option<&Pool>,
6219) {
6220    debug_assert_eq!(out.len(), b * rows);
6221    const TILE: usize = cortiq_core::quant::Q1T_TILE;
6222    let gpr = cols / GROUP_SIZE;
6223    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
6224    let out_addr = SendMut(out.as_mut_ptr());
6225    // int8 prefill (ARM SDOT / x86 AVX2): quantize the B inputs once, unpack
6226    // each weight row's signs to i8 ONCE, then int8-dot against every input —
6227    // the row sign-decode amortizes over the whole batch. CMF_SDOT=0 → f32.
6228    if a8w8_enabled() {
6229        let acts: Vec<SplitAct> = (0..b)
6230            .map(|bi| split_act(&xs[bi * cols..(bi + 1) * cols]))
6231            .collect();
6232        let acts = &acts;
6233        let run = move |start: usize, end: usize| {
6234            let mut sg = vec![0i8; cols + 8]; // row signs, i8 (+8 unpack slack)
6235            let mut sc = vec![0f32; gpr]; // per-group scales
6236            let mut accs = vec![0f32; b]; // per-batch accumulators, reused per row
6237            for r in start..end {
6238                for g in 0..gpr {
6239                    let off = (r * gpr + g) * TILE;
6240                    sc[g] = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
6241                    q1t_unpack_group_i8(
6242                        bytes.as_ptr().wrapping_add(off + 2),
6243                        &mut sg[g * GROUP_SIZE..],
6244                    );
6245                }
6246                for bi in 0..b {
6247                    let act = &acts[bi];
6248                    let mut isum = 0f32;
6249                    for g in 0..gpr {
6250                        let d = q1t_i8dot32(
6251                            sg.as_ptr().wrapping_add(g * GROUP_SIZE),
6252                            act.xq.as_ptr().wrapping_add(g * GROUP_SIZE),
6253                        );
6254                        isum += d as f32 * sc[g];
6255                    }
6256                    let mut acc = isum * act.sx;
6257                    for &(j, xv) in &act.outliers {
6258                        acc += q1t_base_weight(bytes, r, gpr, j) * xv;
6259                    }
6260                    accs[bi] = acc;
6261                }
6262                // Overlay ONCE per row for the whole batch: read each (col, val)
6263                // from mmap a single time (was b× — the re-read dominated prefill)
6264                // and fan it out over the batch via the cached inputs.
6265                if has_ov {
6266                    let (c0, c1) = (
6267                        q1t_rowptr(bytes, rp_off, r),
6268                        q1t_rowptr(bytes, rp_off, r + 1),
6269                    );
6270                    for p in c0..c1 {
6271                        let e = ent_off + p * 4;
6272                        let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
6273                        let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
6274                        for bi in 0..b {
6275                            accs[bi] += val * xs[bi * cols + col];
6276                        }
6277                    }
6278                }
6279                for bi in 0..b {
6280                    unsafe { *out_addr.at(bi * rows + r) = accs[bi] };
6281                }
6282            }
6283        };
6284        dispatch_rows(pool, rows, &run);
6285        return;
6286    }
6287    let run = move |start: usize, end: usize| {
6288        let mut buf = vec![0f32; cols];
6289        for r in start..end {
6290            q1t_dequant_row(bytes, r, gpr, rp_off, ent_off, has_ov, &mut buf);
6291            for bi in 0..b {
6292                let xr = &xs[bi * cols..(bi + 1) * cols];
6293                let mut acc = 0f32;
6294                for j in 0..cols {
6295                    acc += buf[j] * xr[j];
6296                }
6297                unsafe { *out_addr.at(bi * rows + r) = acc };
6298            }
6299        }
6300    };
6301    dispatch_rows(pool, rows, &run);
6302}
6303
6304fn q1_matvec(
6305    bytes: &[u8],
6306    x: &[f32],
6307    rows: usize,
6308    cols: usize,
6309    out: &mut [f32],
6310    pool: Option<&Pool>,
6311) {
6312    debug_assert_eq!(out.len(), rows);
6313    let gpr = cols / GROUP_SIZE;
6314    let out_addr = SendMut(out.as_mut_ptr());
6315    if a8w8_enabled() {
6316        let act = split_act(x);
6317        let gsum = q1_group_sums(&act.xq, gpr);
6318        let (act, gsum) = (&act, &gsum);
6319        let run = move |start: usize, end: usize| {
6320            q1_range_a8w8(bytes, gpr, act, gsum, out_addr, start, end)
6321        };
6322        dispatch_rows(pool, rows, &run);
6323        return;
6324    }
6325    let run = move |start: usize, end: usize| q1_range_f32(bytes, gpr, x, out_addr, start, end);
6326    dispatch_rows(pool, rows, &run);
6327}
6328
6329/// Fused two-input q1 matvec (weights read once per pair).
6330#[allow(clippy::too_many_arguments)]
6331fn q1_matvec2(
6332    bytes: &[u8],
6333    x1: &[f32],
6334    x2: &[f32],
6335    rows: usize,
6336    cols: usize,
6337    o1: &mut [f32],
6338    o2: &mut [f32],
6339    pool: Option<&Pool>,
6340) {
6341    let gpr = cols / GROUP_SIZE;
6342    let p1 = SendMut(o1.as_mut_ptr());
6343    let p2 = SendMut(o2.as_mut_ptr());
6344    if a8w8_enabled() {
6345        let a1 = split_act(x1);
6346        let a2 = split_act(x2);
6347        let g1 = q1_group_sums(&a1.xq, gpr);
6348        let g2 = q1_group_sums(&a2.xq, gpr);
6349        let (a1, a2, g1, g2) = (&a1, &a2, &g1, &g2);
6350        let run = move |start: usize, end: usize| {
6351            for r in start..end {
6352                let mut v1 = dot_q1_row_i8(bytes, r, gpr, &a1.xq, g1) * a1.sx;
6353                let mut v2 = dot_q1_row_i8(bytes, r, gpr, &a2.xq, g2) * a2.sx;
6354                for &(j, xv) in &a1.outliers {
6355                    let (w, s) = q1_outlier(bytes, r, gpr, j);
6356                    v1 += w * s * xv;
6357                }
6358                for &(j, xv) in &a2.outliers {
6359                    let (w, s) = q1_outlier(bytes, r, gpr, j);
6360                    v2 += w * s * xv;
6361                }
6362                // SAFETY: disjoint row ranges per worker.
6363                unsafe {
6364                    *p1.at(r) = v1;
6365                    *p2.at(r) = v2;
6366                }
6367            }
6368        };
6369        dispatch_rows(pool, rows, &run);
6370        return;
6371    }
6372    let run = move |start: usize, end: usize| {
6373        for r in start..end {
6374            // SAFETY: disjoint row ranges per worker.
6375            unsafe {
6376                *p1.at(r) = q1_row_exact(bytes, r, gpr, x1);
6377                *p2.at(r) = q1_row_exact(bytes, r, gpr, x2);
6378            }
6379        }
6380    };
6381    dispatch_rows(pool, rows, &run);
6382}
6383
6384/// Batched q1 matmat: each row's tiles stream once per microbatch.
6385#[allow(clippy::too_many_arguments)]
6386fn q1_matmat(
6387    bytes: &[u8],
6388    xs_all: &[f32],
6389    b: usize,
6390    rows: usize,
6391    cols: usize,
6392    out: &mut [f32],
6393    pool: Option<&Pool>,
6394) {
6395    debug_assert_eq!(out.len(), b * rows);
6396    let gpr = cols / GROUP_SIZE;
6397    let out_addr = SendMut(out.as_mut_ptr());
6398    if a8w8_enabled() {
6399        let acts: Vec<(SplitAct, Vec<i32>)> = (0..b)
6400            .map(|bi| {
6401                let act = split_act(&xs_all[bi * cols..(bi + 1) * cols]);
6402                let gsum = q1_group_sums(&act.xq, gpr);
6403                (act, gsum)
6404            })
6405            .collect();
6406        let acts = &acts;
6407        #[cfg(target_arch = "x86_64")]
6408        let blocked_ok = avx2_enabled()
6409            && blocked_enabled();
6410        #[cfg(target_arch = "aarch64")]
6411        let blocked_ok = sdot_enabled()
6412            && blocked_enabled();
6413        let run = move |start: usize, end: usize| {
6414            for r in start..end {
6415                let mut bi = 0usize;
6416                // Blocked 1×4: the unpacked bit mask serves four
6417                // activation streams per group.
6418                #[cfg(target_arch = "aarch64")]
6419                if blocked_ok {
6420                    while bi + 4 <= acts.len() {
6421                        let xs = [
6422                            acts[bi].0.xq.as_slice(),
6423                            acts[bi + 1].0.xq.as_slice(),
6424                            acts[bi + 2].0.xq.as_slice(),
6425                            acts[bi + 3].0.xq.as_slice(),
6426                        ];
6427                        let gs = [
6428                            acts[bi].1.as_slice(),
6429                            acts[bi + 1].1.as_slice(),
6430                            acts[bi + 2].1.as_slice(),
6431                            acts[bi + 3].1.as_slice(),
6432                        ];
6433                        let d = unsafe { dot_q1_row_1x4_sdot(bytes, r, gpr, xs, gs) };
6434                        for k in 0..4 {
6435                            let (act, _) = &acts[bi + k];
6436                            let mut acc = d[k] * act.sx;
6437                            for &(j, xv) in &act.outliers {
6438                                let (w, sc) = q1_outlier(bytes, r, gpr, j);
6439                                acc += w * sc * xv;
6440                            }
6441                            // SAFETY: disjoint (bi, r) cells per worker.
6442                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
6443                        }
6444                        bi += 4;
6445                    }
6446                }
6447                #[cfg(target_arch = "x86_64")]
6448                if blocked_ok {
6449                    while bi + 4 <= acts.len() {
6450                        let xs = [
6451                            acts[bi].0.xq.as_slice(),
6452                            acts[bi + 1].0.xq.as_slice(),
6453                            acts[bi + 2].0.xq.as_slice(),
6454                            acts[bi + 3].0.xq.as_slice(),
6455                        ];
6456                        let gs = [
6457                            acts[bi].1.as_slice(),
6458                            acts[bi + 1].1.as_slice(),
6459                            acts[bi + 2].1.as_slice(),
6460                            acts[bi + 3].1.as_slice(),
6461                        ];
6462                        let d = unsafe {
6463                            if vnni_tiles_enabled() {
6464                                dot_q1_row_1x4_vnni(bytes, r, gpr, xs, gs)
6465                            } else {
6466                                dot_q1_row_1x4_avx2(bytes, r, gpr, xs, gs)
6467                            }
6468                        };
6469                        for k in 0..4 {
6470                            let (act, _) = &acts[bi + k];
6471                            let mut acc = d[k] * act.sx;
6472                            for &(j, xv) in &act.outliers {
6473                                let (w, sc) = q1_outlier(bytes, r, gpr, j);
6474                                acc += w * sc * xv;
6475                            }
6476                            // SAFETY: disjoint (bi, r) cells per worker.
6477                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
6478                        }
6479                        bi += 4;
6480                    }
6481                }
6482                while bi < acts.len() {
6483                    let (act, gsum) = &acts[bi];
6484                    let mut acc = dot_q1_row_i8(bytes, r, gpr, &act.xq, gsum) * act.sx;
6485                    for &(j, xv) in &act.outliers {
6486                        let (w, s) = q1_outlier(bytes, r, gpr, j);
6487                        acc += w * s * xv;
6488                    }
6489                    // SAFETY: disjoint (bi, r) cells per worker range.
6490                    unsafe { *out_addr.at(bi * rows + r) = acc };
6491                    bi += 1;
6492                }
6493            }
6494        };
6495        dispatch_rows(pool, rows, &run);
6496        return;
6497    }
6498    let run = move |start: usize, end: usize| {
6499        for r in start..end {
6500            for bi in 0..b {
6501                let x = &xs_all[bi * cols..(bi + 1) * cols];
6502                // SAFETY: disjoint (bi, r) cells per worker range.
6503                unsafe { *out_addr.at(bi * rows + r) = q1_row_exact(bytes, r, gpr, x) };
6504            }
6505        }
6506    };
6507    dispatch_rows(pool, rows, &run);
6508}
6509
6510/// Fused q4_block matvec straight from the mapped bytes. SDOT path when
6511/// dotprod is available (port of vmfcore `dot_q4_block_sdot`, measured
6512/// +23% on q4 decode): nibbles → centered i8, int8×int8 `sdot` per
6513/// 32-group, exact outlier correction — the same A8W8 contract as q8.
6514/// `CMF_SDOT=0` keeps the exact scalar path.
6515fn q4matvec(
6516    bytes: &[u8],
6517    x: &[f32],
6518    rows: usize,
6519    cols: usize,
6520    out: &mut [f32],
6521    pool: Option<&Pool>,
6522) {
6523    debug_assert_eq!(out.len(), rows);
6524    let (packed, scales) = q4_split(bytes, rows, cols);
6525    let gpr = cols / GROUP_SIZE;
6526    let out_addr = SendMut(out.as_mut_ptr());
6527
6528    if a8w8_enabled() {
6529        let act = split_act(x);
6530        let run = move |start: usize, end: usize| {
6531            q4_range_a8w8(packed, scales, gpr, cols, &act, out_addr, start, end)
6532        };
6533        dispatch_rows(pool, rows, &run);
6534        return;
6535    }
6536
6537    let run =
6538        move |start: usize, end: usize| q4_range_f32(packed, scales, gpr, x, out_addr, start, end);
6539    dispatch_rows(pool, rows, &run);
6540}
6541
6542/// One q4 row via the A8W8 int8 path — SDOT on ARM, AVX2 maddubs on
6543/// x86 (scalar fallback is unreachable: callers gate on a8w8_enabled).
6544#[inline]
6545#[allow(unreachable_code)]
6546/// One UNPACKED q4 row (centered i8 in `buf`) against four activation
6547/// streams: the 32-byte weight chunk and its abs() load once per group,
6548/// the per-group f16 scale decodes once — four maddubs+reduce chains
6549/// instead of four full (load, abs, dot) rounds.
6550#[cfg(target_arch = "x86_64")]
6551#[target_feature(enable = "avx2")]
6552unsafe fn dot_q4b_row_1x4_avx2(
6553    buf: &[u8],
6554    scales: &[u8],
6555    g0: usize,
6556    gpr: usize,
6557    xs: [&[i8]; 4],
6558) -> [f32; 4] {
6559    // SAFETY: callers uphold buffer contracts (buf.len() == gpr·32).
6560    unsafe {
6561        use core::arch::x86_64::*;
6562        let ones = _mm256_set1_epi16(1);
6563        let mut acc = [0f32; 4];
6564        for gi in 0..gpr {
6565            let s = f16_to_f32(u16::from_le_bytes([
6566                scales[(g0 + gi) * 2],
6567                scales[(g0 + gi) * 2 + 1],
6568            ]));
6569            let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6570            let aw = _mm256_abs_epi8(w);
6571            for (k, xq) in xs.iter().enumerate() {
6572                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6573                let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
6574                let d = _mm256_madd_epi16(p16, ones);
6575                let hi128 = _mm256_extracti128_si256::<1>(d);
6576                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
6577                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
6578                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
6579                acc[k] += _mm_cvtsi128_si32(s32) as f32 * s;
6580            }
6581        }
6582        acc
6583    }
6584}
6585
6586/// VNNI twin of `dot_q4b_row_1x4_avx2` (see `dpbusd_hsum`).
6587#[cfg(target_arch = "x86_64")]
6588#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
6589unsafe fn dot_q4b_row_1x4_vnni(
6590    buf: &[u8],
6591    scales: &[u8],
6592    g0: usize,
6593    gpr: usize,
6594    xs: [&[i8]; 4],
6595) -> [f32; 4] {
6596    // SAFETY: callers uphold buffer contracts (buf.len() == gpr·32).
6597    unsafe {
6598        use core::arch::x86_64::*;
6599        let mut acc = [0f32; 4];
6600        for gi in 0..gpr {
6601            let s = f16_to_f32(u16::from_le_bytes([
6602                scales[(g0 + gi) * 2],
6603                scales[(g0 + gi) * 2 + 1],
6604            ]));
6605            let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6606            let aw = _mm256_abs_epi8(w);
6607            for (k, xq) in xs.iter().enumerate() {
6608                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6609                let d = dpbusd_hsum(aw, _mm256_sign_epi8(x, w));
6610                acc[k] += d as f32 * s;
6611            }
6612        }
6613        acc
6614    }
6615}
6616
6617/// The vbit flavor of the blocked 1×4: the per-activation A8W8 scale
6618/// folds in PER GROUP as `(d·sx)·s` — bit-matching the single-matvec
6619/// accumulation order (the q4_block flavor applies sx once at the end,
6620/// matching ITS single path; the two conventions are historical and
6621/// each blocked leg must mirror its own).
6622#[cfg(target_arch = "x86_64")]
6623#[target_feature(enable = "avx2")]
6624unsafe fn dot_q4b_row_1x4_sx_avx2(
6625    buf: &[u8],
6626    scales: &[u8],
6627    g0: usize,
6628    gpr: usize,
6629    xs: [&[i8]; 4],
6630    sxs: [f32; 4],
6631) -> [f32; 4] {
6632    // SAFETY: callers uphold buffer contracts (buf.len() == gpr·32).
6633    unsafe {
6634        use core::arch::x86_64::*;
6635        let ones = _mm256_set1_epi16(1);
6636        let mut acc = [0f32; 4];
6637        for gi in 0..gpr {
6638            let s = f16_to_f32(u16::from_le_bytes([
6639                scales[(g0 + gi) * 2],
6640                scales[(g0 + gi) * 2 + 1],
6641            ]));
6642            let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6643            let aw = _mm256_abs_epi8(w);
6644            for (k, xq) in xs.iter().enumerate() {
6645                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6646                let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
6647                let d = _mm256_madd_epi16(p16, ones);
6648                let hi128 = _mm256_extracti128_si256::<1>(d);
6649                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
6650                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
6651                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
6652                acc[k] += (_mm_cvtsi128_si32(s32) as f32 * sxs[k]) * s;
6653            }
6654        }
6655        acc
6656    }
6657}
6658
6659/// VNNI twin of `dot_q4b_row_1x4_sx_avx2` (see `dpbusd_hsum`; the
6660/// per-group `(d·sx)·s` fold mirrors the vbit single path).
6661#[cfg(target_arch = "x86_64")]
6662#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
6663unsafe fn dot_q4b_row_1x4_sx_vnni(
6664    buf: &[u8],
6665    scales: &[u8],
6666    g0: usize,
6667    gpr: usize,
6668    xs: [&[i8]; 4],
6669    sxs: [f32; 4],
6670) -> [f32; 4] {
6671    // SAFETY: callers uphold buffer contracts (buf.len() == gpr·32).
6672    unsafe {
6673        use core::arch::x86_64::*;
6674        let mut acc = [0f32; 4];
6675        for gi in 0..gpr {
6676            let s = f16_to_f32(u16::from_le_bytes([
6677                scales[(g0 + gi) * 2],
6678                scales[(g0 + gi) * 2 + 1],
6679            ]));
6680            let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6681            let aw = _mm256_abs_epi8(w);
6682            for (k, xq) in xs.iter().enumerate() {
6683                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6684                let d = dpbusd_hsum(aw, _mm256_sign_epi8(x, w));
6685                acc[k] += (d as f32 * sxs[k]) * s;
6686            }
6687        }
6688        acc
6689    }
6690}
6691
6692#[allow(unreachable_code)]
6693fn dot_q4_row_i8(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
6694    #[cfg(target_arch = "aarch64")]
6695    unsafe {
6696        return dot_q4_row_sdot(packed, scales, g0, gpr, xq);
6697    }
6698    #[cfg(target_arch = "x86_64")]
6699    unsafe {
6700        return dot_q4_row_avx2(packed, scales, g0, gpr, xq);
6701    }
6702    let mut acc = 0f32;
6703    for gi in 0..gpr {
6704        let g = g0 + gi;
6705        let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
6706        let mut d = 0i32;
6707        for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
6708            d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
6709                + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
6710        }
6711        acc += d as f32 * s;
6712    }
6713    acc
6714}
6715
6716/// Two-activation q4 row via the A8W8 int8 path (see `dot_q4_row_i8`).
6717#[inline]
6718#[allow(unreachable_code)]
6719fn dot_q4_row_i8_2(
6720    packed: &[u8],
6721    scales: &[u8],
6722    g0: usize,
6723    gpr: usize,
6724    xq1: &[i8],
6725    xq2: &[i8],
6726) -> (f32, f32) {
6727    #[cfg(target_arch = "aarch64")]
6728    unsafe {
6729        return dot_q4_row_sdot2(packed, scales, g0, gpr, xq1, xq2);
6730    }
6731    #[cfg(target_arch = "x86_64")]
6732    unsafe {
6733        return dot_q4_row_avx2_2(packed, scales, g0, gpr, xq1, xq2);
6734    }
6735    (
6736        dot_q4_row_i8(packed, scales, g0, gpr, xq1),
6737        dot_q4_row_i8(packed, scales, g0, gpr, xq2),
6738    )
6739}
6740
6741/// One q4 row range via SDOT (kernel body of `q4matvec`, extracted so
6742/// multi-matrix jobs can drive it for several tensors in one dispatch).
6743#[allow(clippy::too_many_arguments)]
6744fn q4_range_a8w8(
6745    packed: &[u8],
6746    scales: &[u8],
6747    gpr: usize,
6748    cols: usize,
6749    act: &SplitAct,
6750    out: SendMut,
6751    start: usize,
6752    end: usize,
6753) {
6754    for r in start..end {
6755        let mut acc = dot_q4_row_i8(packed, scales, r * gpr, gpr, &act.xq) * act.sx;
6756        // xq is zeroed at outlier slots — add the exact terms.
6757        for &(j, xv) in &act.outliers {
6758            let flat = r * cols + j;
6759            let byte = packed[flat / 2];
6760            let nib = if flat & 1 == 0 {
6761                byte & 0x0F
6762            } else {
6763                byte >> 4
6764            };
6765            let s = f16_to_f32(u16::from_le_bytes([
6766                scales[(flat / GROUP_SIZE) * 2],
6767                scales[(flat / GROUP_SIZE) * 2 + 1],
6768            ]));
6769            acc += ((nib as i32 - 8) as f32) * s * xv;
6770        }
6771        // SAFETY: disjoint row ranges per worker.
6772        unsafe { *out.at(r) = acc };
6773    }
6774}
6775
6776/// Two-input q4 row range via the A8W8 int8 path — kernel body of
6777/// `q4matvec2`, extracted for pair multi-matrix jobs.
6778#[allow(clippy::too_many_arguments)]
6779fn q4_range2_a8w8(
6780    packed: &[u8],
6781    scales: &[u8],
6782    gpr: usize,
6783    cols: usize,
6784    a1: &SplitAct,
6785    a2: &SplitAct,
6786    p1: SendMut,
6787    p2: SendMut,
6788    start: usize,
6789    end: usize,
6790) {
6791    for r in start..end {
6792        let (s1, s2) = dot_q4_row_i8_2(packed, scales, r * gpr, gpr, &a1.xq, &a2.xq);
6793        let mut acc1 = s1 * a1.sx;
6794        let mut acc2 = s2 * a2.sx;
6795        // xq is zeroed at outlier slots — add the exact terms.
6796        let fix = |outliers: &[(usize, f32)], acc: &mut f32| {
6797            for &(j, xv) in outliers {
6798                let flat = r * cols + j;
6799                let byte = packed[flat / 2];
6800                let nib = if flat & 1 == 0 {
6801                    byte & 0x0F
6802                } else {
6803                    byte >> 4
6804                };
6805                let s = f16_to_f32(u16::from_le_bytes([
6806                    scales[(flat / GROUP_SIZE) * 2],
6807                    scales[(flat / GROUP_SIZE) * 2 + 1],
6808                ]));
6809                *acc += ((nib as i32 - 8) as f32) * s * xv;
6810            }
6811        };
6812        fix(&a1.outliers, &mut acc1);
6813        fix(&a2.outliers, &mut acc2);
6814        // SAFETY: disjoint row ranges per worker.
6815        unsafe {
6816            *p1.at(r) = acc1;
6817            *p2.at(r) = acc2;
6818        }
6819    }
6820}
6821
6822/// Exact scalar q4 row range (same extraction, non-SDOT path).
6823fn q4_range_f32(
6824    packed: &[u8],
6825    scales: &[u8],
6826    gpr: usize,
6827    x: &[f32],
6828    out: SendMut,
6829    start: usize,
6830    end: usize,
6831) {
6832    for r in start..end {
6833        let mut acc = 0f32;
6834        for gi in 0..gpr {
6835            let g = r * gpr + gi;
6836            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
6837            let pk = &packed[g * 16..(g + 1) * 16];
6838            let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
6839            let mut ga = 0f32;
6840            for (k, &b) in pk.iter().enumerate() {
6841                ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
6842                    + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
6843            }
6844            acc += ga * s;
6845        }
6846        // SAFETY: disjoint row ranges per worker.
6847        unsafe { *out.at(r) = acc };
6848    }
6849}
6850
6851/// Fused two-input q4 matvec: nibbles are unpacked ONCE per group and
6852/// dotted against both activations (was: two full matvecs — double
6853/// weight traffic). Per-lane math matches `q4matvec` exactly.
6854#[allow(clippy::too_many_arguments)]
6855fn q4matvec2(
6856    bytes: &[u8],
6857    x1: &[f32],
6858    x2: &[f32],
6859    rows: usize,
6860    cols: usize,
6861    o1: &mut [f32],
6862    o2: &mut [f32],
6863    pool: Option<&Pool>,
6864) {
6865    debug_assert_eq!(o1.len(), rows);
6866    debug_assert_eq!(o2.len(), rows);
6867    let (packed, scales) = q4_split(bytes, rows, cols);
6868    let gpr = cols / GROUP_SIZE;
6869
6870    if a8w8_enabled() {
6871        let a1 = split_act(x1);
6872        let a2 = split_act(x2);
6873        let p1 = SendMut(o1.as_mut_ptr());
6874        let p2 = SendMut(o2.as_mut_ptr());
6875        let run = move |start: usize, end: usize| {
6876            q4_range2_a8w8(packed, scales, gpr, cols, &a1, &a2, p1, p2, start, end)
6877        };
6878        dispatch_rows(pool, rows, &run);
6879        return;
6880    }
6881
6882    let p1 = SendMut(o1.as_mut_ptr());
6883    let p2 = SendMut(o2.as_mut_ptr());
6884    let run = move |start: usize, end: usize| {
6885        q4_range2_f32(packed, scales, gpr, x1, x2, p1, p2, start, end)
6886    };
6887    dispatch_rows(pool, rows, &run);
6888}
6889
6890/// Two-input exact scalar q4 row range (same extraction).
6891#[allow(clippy::too_many_arguments)]
6892fn q4_range2_f32(
6893    packed: &[u8],
6894    scales: &[u8],
6895    gpr: usize,
6896    x1: &[f32],
6897    x2: &[f32],
6898    p1: SendMut,
6899    p2: SendMut,
6900    start: usize,
6901    end: usize,
6902) {
6903    for r in start..end {
6904        let (mut acc1, mut acc2) = (0f32, 0f32);
6905        for gi in 0..gpr {
6906            let g = r * gpr + gi;
6907            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
6908            let pk = &packed[g * 16..(g + 1) * 16];
6909            let x1g = &x1[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
6910            let x2g = &x2[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
6911            let (mut g1, mut g2) = (0f32, 0f32);
6912            for (k, &b) in pk.iter().enumerate() {
6913                let wl = (b & 0x0F) as f32 - 8.0;
6914                let wh = ((b >> 4) & 0x0F) as f32 - 8.0;
6915                g1 += wl * x1g[k * 2] + wh * x1g[k * 2 + 1];
6916                g2 += wl * x2g[k * 2] + wh * x2g[k * 2 + 1];
6917            }
6918            acc1 += g1 * s;
6919            acc2 += g2 * s;
6920        }
6921        // SAFETY: disjoint row ranges per worker.
6922        unsafe {
6923            *p1.at(r) = acc1;
6924            *p2.at(r) = acc2;
6925        }
6926    }
6927}
6928
6929thread_local! {
6930    /// Per-worker decoded-row scratch for the batched q4/vbit kernels
6931    /// (centered i8 for SDOT, f32 for the exact/scalar paths).
6932    static ROW_I8: std::cell::RefCell<Vec<u8>> = const { std::cell::RefCell::new(Vec::new()) };
6933    static ROW_F32: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
6934}
6935
6936/// Batched q4 matmat: each weight row is unpacked from the mmap ONCE
6937/// and dotted against ALL b activations (prefill used to fall back to b
6938/// full matvecs — b× weight traffic and b× nibble decode). Per-position
6939/// math matches `q4matvec` exactly: same group order, same accumulation.
6940/// `out` is row-major [b, rows] like `qmatmat`.
6941#[allow(clippy::too_many_arguments)]
6942fn q4matmat(
6943    bytes: &[u8],
6944    xs_all: &[f32],
6945    b: usize,
6946    rows: usize,
6947    cols: usize,
6948    out: &mut [f32],
6949    pool: Option<&Pool>,
6950) {
6951    debug_assert_eq!(xs_all.len(), b * cols);
6952    debug_assert_eq!(out.len(), b * rows);
6953    let (packed, scales) = q4_split(bytes, rows, cols);
6954    let gpr = cols / GROUP_SIZE;
6955    let gscale = |g: usize| f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
6956
6957    if a8w8_enabled() {
6958        let acts: Vec<SplitAct> = (0..b)
6959            .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
6960            .collect();
6961        let acts = &acts;
6962        let out_addr = SendMut(out.as_mut_ptr());
6963        let run = move |start: usize, end: usize| {
6964            ROW_I8.with(|rb| {
6965                let mut buf = rb.borrow_mut();
6966                buf.resize(cols, 0);
6967                for r in start..end {
6968                    // Unpack the row's nibbles to centered i8 once
6969                    // (element 2k = low nibble, 2k+1 = high — flat order,
6970                    // same as dot_q4_row_sdot's zip).
6971                    for gi in 0..gpr {
6972                        let g = r * gpr + gi;
6973                        for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
6974                            buf[gi * GROUP_SIZE + k * 2] = ((bt & 0x0F) as i32 - 8) as i8 as u8;
6975                            buf[gi * GROUP_SIZE + k * 2 + 1] =
6976                                (((bt >> 4) & 0x0F) as i32 - 8) as i8 as u8;
6977                        }
6978                    }
6979                    let mut bi = 0usize;
6980                    #[cfg(target_arch = "x86_64")]
6981                    if avx2_enabled()
6982                        && blocked_enabled()
6983                    {
6984                        while bi + 4 <= acts.len() {
6985                            let xs = [
6986                                acts[bi].xq.as_slice(),
6987                                acts[bi + 1].xq.as_slice(),
6988                                acts[bi + 2].xq.as_slice(),
6989                                acts[bi + 3].xq.as_slice(),
6990                            ];
6991                            let d = unsafe {
6992                                if vnni_tiles_enabled() {
6993                                    dot_q4b_row_1x4_vnni(&buf, scales, r * gpr, gpr, xs)
6994                                } else {
6995                                    dot_q4b_row_1x4_avx2(&buf, scales, r * gpr, gpr, xs)
6996                                }
6997                            };
6998                            for k in 0..4 {
6999                                let act = &acts[bi + k];
7000                                let mut acc = d[k] * act.sx;
7001                                for &(j, xv) in &act.outliers {
7002                                    acc += (buf[j] as i8) as f32
7003                                        * gscale((r * cols + j) / GROUP_SIZE)
7004                                        * xv;
7005                                }
7006                                // SAFETY: disjoint (bi, r) cells per worker.
7007                                unsafe { *out_addr.at((bi + k) * rows + r) = acc };
7008                            }
7009                            bi += 4;
7010                        }
7011                    }
7012                    while bi < acts.len() {
7013                        let act = &acts[bi];
7014                        let mut acc = 0f32;
7015                        for gi in 0..gpr {
7016                            let d = dot_i8_i8(
7017                                &buf[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
7018                                &act.xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
7019                            );
7020                            acc += d as f32 * gscale(r * gpr + gi);
7021                        }
7022                        acc *= act.sx;
7023                        // xq is zeroed at outlier slots — exact terms.
7024                        for &(j, xv) in &act.outliers {
7025                            acc += (buf[j] as i8) as f32 * gscale((r * cols + j) / GROUP_SIZE) * xv;
7026                        }
7027                        // SAFETY: disjoint (bi, r) cells per worker row range.
7028                        unsafe { *out_addr.at(bi * rows + r) = acc };
7029                        bi += 1;
7030                    }
7031                }
7032            })
7033        };
7034        dispatch_rows(pool, rows, &run);
7035        return;
7036    }
7037
7038    let out_addr = SendMut(out.as_mut_ptr());
7039    let run = move |start: usize, end: usize| {
7040        ROW_F32.with(|rb| {
7041            let mut buf = rb.borrow_mut();
7042            buf.resize(cols, 0.0);
7043            for r in start..end {
7044                // Decode raw (nib − 8) values once; scales stay per-group
7045                // so the accumulation order matches q4matvec bit-for-bit.
7046                for gi in 0..gpr {
7047                    let g = r * gpr + gi;
7048                    for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
7049                        buf[gi * GROUP_SIZE + k * 2] = (bt & 0x0F) as f32 - 8.0;
7050                        buf[gi * GROUP_SIZE + k * 2 + 1] = ((bt >> 4) & 0x0F) as f32 - 8.0;
7051                    }
7052                }
7053                for bi in 0..b {
7054                    let x = &xs_all[bi * cols..(bi + 1) * cols];
7055                    let mut acc = 0f32;
7056                    for gi in 0..gpr {
7057                        let mut ga = 0f32;
7058                        // Pairwise (lo + hi) addition, matching
7059                        // q4matvec's `ga += lo·x + hi·x` shape exactly —
7060                        // a flat one-per-element loop rounds differently
7061                        // and broke bit-parity on the scalar (x86) path.
7062                        for k in 0..GROUP_SIZE / 2 {
7063                            let e = gi * GROUP_SIZE + k * 2;
7064                            ga += buf[e] * x[e] + buf[e + 1] * x[e + 1];
7065                        }
7066                        acc += ga * gscale(r * gpr + gi);
7067                    }
7068                    // SAFETY: disjoint (bi, r) cells per worker row range.
7069                    unsafe { *out_addr.at(bi * rows + r) = acc };
7070                }
7071            }
7072        })
7073    };
7074    dispatch_rows(pool, rows, &run);
7075}
7076
7077/// Batched vbit matmat: each variable-bit row is decoded from the mmap
7078/// ONCE for the whole microbatch. Same per-position math as
7079/// `vbitmatvec` (SDOT A8W8 with exact outliers / exact f32 for b=8 rows
7080/// and the scalar path).
7081#[allow(clippy::too_many_arguments)]
7082fn vbitmatmat(
7083    bytes: &[u8],
7084    offsets: &[usize],
7085    xs_all: &[f32],
7086    b: usize,
7087    rows: usize,
7088    cols: usize,
7089    out: &mut [f32],
7090    pool: Option<&Pool>,
7091) {
7092    debug_assert_eq!(xs_all.len(), b * cols);
7093    debug_assert_eq!(out.len(), b * rows);
7094    debug_assert_eq!(offsets.len(), rows + 1);
7095    let ng = cols / GROUP_SIZE;
7096    let bits = &bytes[..rows];
7097    let sc_off = rows;
7098    let gscale = |r: usize, g: usize| {
7099        let so = (r * ng + g) * 2;
7100        f16_to_f32(u16::from_le_bytes([
7101            bytes[sc_off + so],
7102            bytes[sc_off + so + 1],
7103        ]))
7104    };
7105
7106    // Decode row r's raw (u − L) values into `dst` (f32, unscaled).
7107    let decode_f32 = |r: usize, dst: &mut [f32]| {
7108        let bw = bits[r] as usize;
7109        let l = ((1i32 << (bw - 1)) - 1) as f32;
7110        let data = &bytes[offsets[r]..offsets[r + 1]];
7111        let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
7112        for d in dst.iter_mut() {
7113            while nbits < bw {
7114                acc = (acc << 8) | data[idx] as u64;
7115                idx += 1;
7116                nbits += 8;
7117            }
7118            let u = ((acc >> (nbits - bw)) & ((1u64 << bw) - 1)) as f32;
7119            nbits -= bw;
7120            *d = u - l;
7121        }
7122    };
7123
7124    if a8w8_enabled() {
7125        let acts: Vec<SplitAct> = (0..b)
7126            .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
7127            .collect();
7128        let acts = &acts;
7129        let out_addr = SendMut(out.as_mut_ptr());
7130        let run = move |start: usize, end: usize| {
7131            for r in start..end {
7132                let bw = bits[r] as usize;
7133                if bw == 8 {
7134                    // u−L reaches 128 → no i8 path; decode once, exact
7135                    // f32 dots for every position (same as vbitmatvec).
7136                    ROW_F32.with(|rb| {
7137                        let mut buf = rb.borrow_mut();
7138                        buf.resize(cols, 0.0);
7139                        decode_f32(r, &mut buf);
7140                        for bi in 0..b {
7141                            let x = &xs_all[bi * cols..(bi + 1) * cols];
7142                            let mut dot = 0f32;
7143                            for g in 0..ng {
7144                                let mut gd = 0f32;
7145                                for k in 0..GROUP_SIZE {
7146                                    gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
7147                                }
7148                                dot += gd * gscale(r, g);
7149                            }
7150                            // SAFETY: disjoint (bi, r) cells per worker range.
7151                            unsafe { *out_addr.at(bi * rows + r) = dot };
7152                        }
7153                    });
7154                    continue;
7155                }
7156                let l = (1i32 << (bw - 1)) - 1;
7157                let data = &bytes[offsets[r]..offsets[r + 1]];
7158                ROW_I8.with(|rb| {
7159                    let mut buf = rb.borrow_mut();
7160                    buf.resize(cols, 0);
7161                    #[inline(always)]
7162                    fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
7163                        for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
7164                            let u = unpack8::<B>(&data[blk * B..]);
7165                            for k in 0..8 {
7166                                chunk[k] = (u[k] - l) as i8 as u8;
7167                            }
7168                        }
7169                    }
7170                    match bw {
7171                        3 => fill::<3>(data, l, &mut buf),
7172                        4 => vbit_fill4(data, &mut buf),
7173                        5 => fill::<5>(data, l, &mut buf),
7174                        6 => fill::<6>(data, l, &mut buf),
7175                        _ => unreachable!("vbit bit-width {bw} (validated at load)"),
7176                    }
7177                    let mut bi = 0usize;
7178                    // The vbit scale table shares q4_block's layout
7179                    // (contiguous f16 per (row·ng + g)), so the same
7180                    // blocked 1×4 kernel serves the decoded row.
7181                    #[cfg(target_arch = "x86_64")]
7182                    if avx2_enabled()
7183                        && blocked_enabled()
7184                    {
7185                        while bi + 4 <= acts.len() {
7186                            let xs = [
7187                                acts[bi].xq.as_slice(),
7188                                acts[bi + 1].xq.as_slice(),
7189                                acts[bi + 2].xq.as_slice(),
7190                                acts[bi + 3].xq.as_slice(),
7191                            ];
7192                            let sxs = [
7193                                acts[bi].sx,
7194                                acts[bi + 1].sx,
7195                                acts[bi + 2].sx,
7196                                acts[bi + 3].sx,
7197                            ];
7198                            let d = unsafe {
7199                                if vnni_tiles_enabled() {
7200                                    dot_q4b_row_1x4_sx_vnni(
7201                                        &buf,
7202                                        &bytes[sc_off..],
7203                                        r * ng,
7204                                        ng,
7205                                        xs,
7206                                        sxs,
7207                                    )
7208                                } else {
7209                                    dot_q4b_row_1x4_sx_avx2(
7210                                        &buf,
7211                                        &bytes[sc_off..],
7212                                        r * ng,
7213                                        ng,
7214                                        xs,
7215                                        sxs,
7216                                    )
7217                                }
7218                            };
7219                            for k in 0..4 {
7220                                let act = &acts[bi + k];
7221                                let mut dot = d[k];
7222                                for &(j, xv) in &act.outliers {
7223                                    dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
7224                                }
7225                                // SAFETY: disjoint (bi, r) cells per worker.
7226                                unsafe { *out_addr.at((bi + k) * rows + r) = dot };
7227                            }
7228                            bi += 4;
7229                        }
7230                    }
7231                    while bi < acts.len() {
7232                        let act = &acts[bi];
7233                        let mut dot = 0f32;
7234                        for g in 0..ng {
7235                            let d = dot_i8_i8(
7236                                &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
7237                                &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
7238                            ) as f32
7239                                * act.sx;
7240                            dot += d * gscale(r, g);
7241                        }
7242                        for &(j, xv) in &act.outliers {
7243                            dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
7244                        }
7245                        // SAFETY: disjoint (bi, r) cells per worker range.
7246                        unsafe { *out_addr.at(bi * rows + r) = dot };
7247                        bi += 1;
7248                    }
7249                });
7250            }
7251        };
7252        dispatch_rows(pool, rows, &run);
7253        return;
7254    }
7255
7256    let out_addr = SendMut(out.as_mut_ptr());
7257    let run = move |start: usize, end: usize| {
7258        ROW_F32.with(|rb| {
7259            let mut buf = rb.borrow_mut();
7260            buf.resize(cols, 0.0);
7261            for r in start..end {
7262                decode_f32(r, &mut buf);
7263                for bi in 0..b {
7264                    let x = &xs_all[bi * cols..(bi + 1) * cols];
7265                    let mut dot = 0f32;
7266                    for g in 0..ng {
7267                        let mut gd = 0f32;
7268                        for k in 0..GROUP_SIZE {
7269                            gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
7270                        }
7271                        dot += gd * gscale(r, g);
7272                    }
7273                    // SAFETY: disjoint (bi, r) cells per worker range.
7274                    unsafe { *out_addr.at(bi * rows + r) = dot };
7275                }
7276            }
7277        })
7278    };
7279    dispatch_rows(pool, rows, &run);
7280}
7281
7282/// Build a GPU batch job for a q8-family mapped tensor (primary
7283/// shard): prescaled input + directory coordinates. None → not
7284/// GPU-eligible, caller stays on the CPU.
7285pub(crate) fn gpu_batch_job<'a>(
7286    t: &'a QTensor,
7287    x: &[f32],
7288) -> Option<(std::sync::Arc<CmfModel>, crate::gpu::BatchJob<'a>)> {
7289    match t {
7290        QTensor::Mapped {
7291            model,
7292            idx,
7293            dtype: dt @ (TensorDtype::Q8Row | TensorDtype::Q8_2f),
7294            rows,
7295            cols,
7296            row_scale,
7297            col_field,
7298            ..
7299        } => Some((
7300            model.clone(),
7301            crate::gpu::BatchJob {
7302                idx: *idx,
7303                rows: *rows,
7304                cols: *cols,
7305                row_scale,
7306                xs: prescale(x, col_field, *dt).into_owned(),
7307                layout: crate::gpu::BatchLayout::Q8,
7308            },
7309        )),
7310        // q1: raw f32 activations, tile-embedded scales.
7311        QTensor::Mapped {
7312            model,
7313            idx,
7314            dtype: TensorDtype::Q1,
7315            rows,
7316            cols,
7317            ..
7318        } => Some((
7319            model.clone(),
7320            crate::gpu::BatchJob {
7321                idx: *idx,
7322                rows: *rows,
7323                cols: *cols,
7324                row_scale: &[],
7325                xs: x.to_vec(),
7326                layout: crate::gpu::BatchLayout::Q1,
7327            },
7328        )),
7329        _ => None,
7330    }
7331}
7332
7333thread_local! {
7334    static PRESCALE_BUF1: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
7335    static PRESCALE_BUF2: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
7336}
7337
7338pub(crate) fn prescale<'a>(
7339    x: &'a [f32],
7340    col_field: &[f32],
7341    dtype: TensorDtype,
7342) -> std::borrow::Cow<'a, [f32]> {
7343    if dtype == TensorDtype::Q8_2f {
7344        x.iter().zip(col_field).map(|(a, c)| a * c).collect()
7345    } else {
7346        std::borrow::Cow::Borrowed(x)
7347    }
7348}
7349
7350/// θ col-field fold for q8_2f activations. Borrowed pass-through for
7351/// every other dtype, using thread-local buffers to eliminate per-matvec allocations.
7352pub(crate) fn prescale_with<R, F: FnOnce(&[f32]) -> R>(
7353    x: &[f32],
7354    col_field: &[f32],
7355    dtype: TensorDtype,
7356    buf_id: u8,
7357    f: F,
7358) -> R {
7359    if dtype == TensorDtype::Q8_2f {
7360        if buf_id == 1 {
7361            PRESCALE_BUF1.with(|b| {
7362                let mut buf = b.borrow_mut();
7363                buf.clear();
7364                buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
7365                f(&buf)
7366            })
7367        } else {
7368            PRESCALE_BUF2.with(|b| {
7369                let mut buf = b.borrow_mut();
7370                buf.clear();
7371                buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
7372                f(&buf)
7373            })
7374        }
7375    } else {
7376        f(x)
7377    }
7378}
7379
7380// ───────────────────── x86-64 AVX2 kernels (roadmap этап 2) ─────────────────────
7381
7382/// AVX2+FMA available? Default ON when the CPU supports both;
7383/// `CMF_AVX2=0` disables (falls back to the autovectorized loops).
7384#[cfg(target_arch = "x86_64")]
7385pub(crate) fn avx2_enabled() -> bool {
7386    use std::sync::OnceLock;
7387    static ON: OnceLock<bool> = OnceLock::new();
7388    *ON.get_or_init(|| {
7389        std::env::var("CMF_AVX2").map(|v| v != "0").unwrap_or(true)
7390            && std::arch::is_x86_feature_detected!("avx2")
7391            && std::arch::is_x86_feature_detected!("fma")
7392    })
7393}
7394
7395/// AVX2 A8W8 allowed? The quantized-activation contract is switched by
7396/// the SAME env as the ARM SDOT path: `CMF_SDOT=0` keeps exact kernels
7397/// (the golden-parity exact gate relies on it) — AVX2 f32 kernels stay
7398/// active either way, they are exact (regrouped sums only).
7399#[cfg(target_arch = "x86_64")]
7400fn avx2_a8w8_enabled() -> bool {
7401    use std::sync::OnceLock;
7402    static ON: OnceLock<bool> = OnceLock::new();
7403    *ON.get_or_init(|| {
7404        avx2_enabled() && std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true)
7405    })
7406}
7407
7408/// A8W8 quantized-activation path available on THIS machine? One
7409/// switch across architectures: ARM dotprod (CMF_SDOT) or x86 AVX2
7410/// (CMF_AVX2 + the same CMF_SDOT exact-contract override).
7411#[inline]
7412pub(crate) fn a8w8_enabled() -> bool {
7413    #[cfg(target_arch = "aarch64")]
7414    {
7415        sdot_enabled()
7416    }
7417    #[cfg(target_arch = "x86_64")]
7418    {
7419        avx2_a8w8_enabled()
7420    }
7421    #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
7422    {
7423        false
7424    }
7425}
7426
7427/// int8·int8 dot dispatch: SDOT on ARM; AVX-512 VNNI (vpdpbusd) or AVX2
7428/// maddubs on x86. Callers are gated by `a8w8_enabled()`.
7429#[inline]
7430#[allow(unreachable_code)]
7431fn dot_i8_i8(w: &[u8], xq: &[i8]) -> i32 {
7432    #[cfg(target_arch = "aarch64")]
7433    unsafe {
7434        return dot_i8_sdot(w, xq);
7435    }
7436    #[cfg(target_arch = "x86_64")]
7437    unsafe {
7438        if avx512vnni_enabled() {
7439            return dot_i8_i8_vnni(w, xq);
7440        }
7441        return dot_i8_i8_avx2(w, xq);
7442    }
7443    w.iter()
7444        .zip(xq)
7445        .map(|(&a, &b)| (a as i8) as i32 * b as i32)
7446        .sum()
7447}
7448
7449/// AVX-512 VNNI available? (F+BW+VL+VNNI; `CMF_AVX512=0` falls back to
7450/// AVX2.) VL matters: short 32-byte groups (q4/vbit) ride the 256-bit
7451/// `vpdpbusd` encoding.
7452#[cfg(target_arch = "x86_64")]
7453fn avx512vnni_enabled() -> bool {
7454    use std::sync::OnceLock;
7455    static ON: OnceLock<bool> = OnceLock::new();
7456    *ON.get_or_init(|| {
7457        std::env::var("CMF_AVX512")
7458            .map(|v| v != "0")
7459            .unwrap_or(true)
7460            && std::arch::is_x86_feature_detected!("avx512f")
7461            && std::arch::is_x86_feature_detected!("avx512bw")
7462            && std::arch::is_x86_feature_detected!("avx512vl")
7463            && std::arch::is_x86_feature_detected!("avx512vnni")
7464    })
7465}
7466
7467/// Grouped-codec VNNI arms (the q4t/q4b/q1/q1t tile kernels): default
7468/// ON where AVX-512 VNNI exists (`CMF_VNNI_TILES=0` opt-out). Measured
7469/// on Ryzen 7950X (Zen4, 3 alternating process pairs, blocked GEMM
7470/// 4864×896 b=256): q4t 63→68 GF/s (+8%), q1 53→56 (+6%), q4b 72→75
7471/// (+4%) — consistent, no leg regressed. The tile kernels keep a
7472/// horizontal reduce per 32-weight group, so the `vpdpbusd` saving is
7473/// smaller than the long-dot q8 win (+13%), but it is real and free.
7474#[cfg(target_arch = "x86_64")]
7475fn vnni_tiles_enabled() -> bool {
7476    use std::sync::OnceLock;
7477    static ON: OnceLock<bool> = OnceLock::new();
7478    *ON.get_or_init(|| {
7479        std::env::var("CMF_VNNI_TILES")
7480            .map(|v| v != "0")
7481            .unwrap_or(true)
7482            && avx512vnni_enabled()
7483    })
7484}
7485
7486/// One 256-bit u8×i8 dot → i32 via `vpdpbusd` into a fresh accumulator
7487/// plus the same horizontal reduce the AVX2 kernels use. Products are
7488/// bounded (|w| ≤ 8 or ≤ 1), so maddubs never saturated — the i32 sum
7489/// is bit-identical to the maddubs+madd pair it replaces.
7490#[cfg(target_arch = "x86_64")]
7491#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
7492#[inline]
7493unsafe fn dpbusd_hsum(aw: core::arch::x86_64::__m256i, xs: core::arch::x86_64::__m256i) -> i32 {
7494    // SAFETY: pure register math.
7495    unsafe {
7496        use core::arch::x86_64::*;
7497        let d = _mm256_dpbusd_epi32(_mm256_setzero_si256(), aw, xs);
7498        let hi128 = _mm256_extracti128_si256::<1>(d);
7499        let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
7500        let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7501        let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7502        _mm_cvtsi128_si32(s32)
7503    }
7504}
7505
7506/// int8·int8 via AVX-512 VNNI: `vpdpbusd` fuses the maddubs+madd+add
7507/// triple into one u8×i8 dot-accumulate. AVX-512 has no vpsignb, so the
7508/// |w|·sign(x,w) trick becomes |w| × (x negated where w<0) via a mask
7509/// subtract — w==0 lanes contribute 0 through |w|=0 either way.
7510#[cfg(target_arch = "x86_64")]
7511#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
7512unsafe fn dot_i8_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
7513    // SAFETY: callers uphold slice-length contracts (see call sites).
7514    unsafe {
7515        use core::arch::x86_64::*;
7516        let n = w.len();
7517        let mut j = 0usize;
7518        let mut total: i32;
7519        // 4 independent accumulators: vpdpbusd is its own loop-carried
7520        // dependency (~5-cycle latency) — a single-acc loop runs
7521        // latency-bound and LOSES to the AVX2 maddubs kernel, measured
7522        // on Granite Rapids.
7523        {
7524            #[inline(always)]
7525            unsafe fn step(
7526                w: *const u8,
7527                x: *const i8,
7528                acc: core::arch::x86_64::__m512i,
7529            ) -> core::arch::x86_64::__m512i {
7530                unsafe {
7531                    use core::arch::x86_64::*;
7532                    let wv = _mm512_loadu_si512(w as *const _);
7533                    let xv = _mm512_loadu_si512(x as *const _);
7534                    let aw = _mm512_abs_epi8(wv);
7535                    let neg = _mm512_movepi8_mask(wv);
7536                    let sx = _mm512_mask_sub_epi8(xv, neg, _mm512_setzero_si512(), xv);
7537                    _mm512_dpbusd_epi32(acc, aw, sx)
7538                }
7539            }
7540            let (mut a0, mut a1, mut a2, mut a3) = (
7541                _mm512_setzero_si512(),
7542                _mm512_setzero_si512(),
7543                _mm512_setzero_si512(),
7544                _mm512_setzero_si512(),
7545            );
7546            while j + 256 <= n {
7547                a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
7548                a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), a1);
7549                a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), a2);
7550                a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), a3);
7551                j += 256;
7552            }
7553            while j + 64 <= n {
7554                a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
7555                j += 64;
7556            }
7557            let s01 = _mm512_add_epi32(a0, a1);
7558            let s23 = _mm512_add_epi32(a2, a3);
7559            total = _mm512_reduce_add_epi32(_mm512_add_epi32(s01, s23));
7560        }
7561        // 32-wide (q4/vbit groups are exactly 32 bytes).
7562        if j + 32 <= n {
7563            let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
7564            let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
7565            let d = _mm256_dpbusd_epi32(
7566                _mm256_setzero_si256(),
7567                _mm256_abs_epi8(wv),
7568                _mm256_sign_epi8(xv, wv),
7569            );
7570            let hi128 = _mm256_extracti128_si256::<1>(d);
7571            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
7572            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7573            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7574            total += _mm_cvtsi128_si32(s32);
7575            j += 32;
7576        }
7577        while j < n {
7578            total += (w[j] as i8) as i32 * xq[j] as i32;
7579            j += 1;
7580        }
7581        total
7582    }
7583}
7584
7585/// i8 row · f32 x via AVX2/FMA (x86 mirror of `dot_i8_f32_neon`).
7586#[cfg(target_arch = "x86_64")]
7587#[target_feature(enable = "avx2,fma")]
7588unsafe fn dot_i8_f32_avx2(w: &[u8], x: &[f32]) -> f32 {
7589    // SAFETY: callers uphold slice-length contracts (see call sites).
7590    unsafe {
7591        use core::arch::x86_64::*;
7592        let n = x.len();
7593        let wp = w.as_ptr();
7594        let xp = x.as_ptr();
7595        let (mut a0, mut a1) = (_mm256_setzero_ps(), _mm256_setzero_ps());
7596        let mut j = 0usize;
7597        while j + 16 <= n {
7598            let wb = _mm_loadu_si128(wp.add(j) as *const __m128i);
7599            let lo = _mm256_cvtepi8_epi32(wb);
7600            let hi = _mm256_cvtepi8_epi32(_mm_srli_si128::<8>(wb));
7601            a0 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(lo), _mm256_loadu_ps(xp.add(j)), a0);
7602            a1 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(hi), _mm256_loadu_ps(xp.add(j + 8)), a1);
7603            j += 16;
7604        }
7605        let acc = _mm256_add_ps(a0, a1);
7606        let hi128 = _mm256_extractf128_ps::<1>(acc);
7607        let s128 = _mm_add_ps(_mm256_castps256_ps128(acc), hi128);
7608        let s64 = _mm_add_ps(s128, _mm_movehl_ps(s128, s128));
7609        let s32 = _mm_add_ss(s64, _mm_shuffle_ps::<1>(s64, s64));
7610        let mut sum = _mm_cvtss_f32(s32);
7611        while j < n {
7612            sum += (*wp.add(j) as i8) as f32 * *xp.add(j);
7613            j += 1;
7614        }
7615        sum
7616    }
7617}
7618
7619/// int8(weight)·int8(activation) → i32 via AVX2 maddubs — the x86
7620/// analogue of the SDOT A8W8 path. `maddubs` takes u8×i8, so the
7621/// standard sign trick applies: |w| × sign(x, w) ≡ w × x per lane.
7622/// Pair saturation is safe: |w|≤128, |x|≤127 → 2·128·127 < 32767.
7623#[cfg(target_arch = "x86_64")]
7624#[target_feature(enable = "avx2")]
7625unsafe fn dot_i8_i8_avx2(w: &[u8], xq: &[i8]) -> i32 {
7626    // SAFETY: callers uphold slice-length contracts (see call sites).
7627    unsafe {
7628        use core::arch::x86_64::*;
7629        let n = w.len();
7630        let ones = _mm256_set1_epi16(1);
7631        let mut acc = _mm256_setzero_si256();
7632        let mut j = 0usize;
7633        while j + 32 <= n {
7634            let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
7635            let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
7636            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
7637            acc = _mm256_add_epi32(acc, _mm256_madd_epi16(p16, ones));
7638            j += 32;
7639        }
7640        let hi128 = _mm256_extracti128_si256::<1>(acc);
7641        let s128 = _mm_add_epi32(_mm256_castsi256_si128(acc), hi128);
7642        let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7643        let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7644        let mut s = _mm_cvtsi128_si32(s32);
7645        while j < n {
7646            s += (w[j] as i8) as i32 * xq[j] as i32;
7647            j += 1;
7648        }
7649        s
7650    }
7651}
7652
7653/// smmla 2×4: one instruction covers a 2-row × 2-activation × 8-deep
7654/// tile (32 MACs vs sdot's 16) — the weight pair loads once per 8-k
7655/// slice as a combined 2×8 register and meets two activation pairs.
7656#[cfg(target_arch = "aarch64")]
7657#[target_feature(enable = "neon,i8mm")]
7658unsafe fn dot_i8_smmla_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
7659    // SAFETY: callers uphold slice-length contracts.
7660    unsafe {
7661        use core::arch::aarch64::*;
7662        use core::arch::asm;
7663        let n = w0.len();
7664        let w0p = w0.as_ptr() as *const i8;
7665        let w1p = w1.as_ptr() as *const i8;
7666        // acc01 holds [c(r0,x0) c(r0,x1) c(r1,x0) c(r1,x1)]; acc23 the
7667        // same for x2/x3.
7668        let mut acc01 = vdupq_n_s32(0);
7669        let mut acc23 = vdupq_n_s32(0);
7670        let mut i = 0usize;
7671        while i + 8 <= n {
7672            let wa = vcombine_s8(vld1_s8(w0p.add(i)), vld1_s8(w1p.add(i)));
7673            let xb01 = vcombine_s8(
7674                vld1_s8(xs[0].as_ptr().add(i)),
7675                vld1_s8(xs[1].as_ptr().add(i)),
7676            );
7677            let xb23 = vcombine_s8(
7678                vld1_s8(xs[2].as_ptr().add(i)),
7679                vld1_s8(xs[3].as_ptr().add(i)),
7680            );
7681            asm!(
7682                "smmla {a01:v}.4s, {w:v}.16b, {x01:v}.16b",
7683                "smmla {a23:v}.4s, {w:v}.16b, {x23:v}.16b",
7684                a01 = inout(vreg) acc01, a23 = inout(vreg) acc23,
7685                w = in(vreg) wa, x01 = in(vreg) xb01, x23 = in(vreg) xb23,
7686                options(pure, nomem, nostack),
7687            );
7688            i += 8;
7689        }
7690        let mut out = [[0i32; 4]; 2];
7691        let a01: [i32; 4] = core::mem::transmute(acc01);
7692        let a23: [i32; 4] = core::mem::transmute(acc23);
7693        out[0][0] = a01[0];
7694        out[0][1] = a01[1];
7695        out[1][0] = a01[2];
7696        out[1][1] = a01[3];
7697        out[0][2] = a23[0];
7698        out[0][3] = a23[1];
7699        out[1][2] = a23[2];
7700        out[1][3] = a23[3];
7701        if i < n {
7702            for (k, x) in xs.iter().enumerate() {
7703                for j in i..n {
7704                    out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
7705                    out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
7706                }
7707            }
7708        }
7709        out
7710    }
7711}
7712
7713/// ARM twin of the x86 blocked prefill GEMM: two weight rows stay in
7714/// registers across four activation streams, eight sdot accumulators.
7715/// (The per-row form re-read each W row once per activation.)
7716#[cfg(target_arch = "aarch64")]
7717#[target_feature(enable = "neon,dotprod")]
7718unsafe fn dot_i8_sdot_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
7719    // SAFETY: callers uphold slice-length contracts.
7720    unsafe {
7721        use core::arch::aarch64::*;
7722        use core::arch::asm;
7723        let n = w0.len();
7724        let w0p = w0.as_ptr() as *const i8;
7725        let w1p = w1.as_ptr() as *const i8;
7726        let mut acc = [[vdupq_n_s32(0); 4]; 2];
7727        let mut i = 0usize;
7728        while i + 16 <= n {
7729            let wv0 = vld1q_s8(w0p.add(i));
7730            let wv1 = vld1q_s8(w1p.add(i));
7731            for (k, x) in xs.iter().enumerate() {
7732                let xv = vld1q_s8(x.as_ptr().add(i));
7733                let (mut a0, mut a1) = (acc[0][k], acc[1][k]);
7734                asm!(
7735                    "sdot {a0:v}.4s, {w0:v}.16b, {x:v}.16b",
7736                    "sdot {a1:v}.4s, {w1:v}.16b, {x:v}.16b",
7737                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
7738                    w0 = in(vreg) wv0, w1 = in(vreg) wv1, x = in(vreg) xv,
7739                    options(pure, nomem, nostack),
7740                );
7741                acc[0][k] = a0;
7742                acc[1][k] = a1;
7743            }
7744            i += 16;
7745        }
7746        let mut out = [[0i32; 4]; 2];
7747        for r in 0..2 {
7748            for k in 0..4 {
7749                out[r][k] = vaddvq_s32(acc[r][k]);
7750            }
7751        }
7752        if i < n {
7753            for (k, x) in xs.iter().enumerate() {
7754                for j in i..n {
7755                    out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
7756                    out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
7757                }
7758            }
7759        }
7760        out
7761    }
7762}
7763
7764/// Blocked 2 weight rows × 4 activations for the prefill GEMM
7765/// (roadmap P0: packed panels + multi-row accumulators). The two rows'
7766/// abs() live in registers across all four activation streams; the
7767/// sign-fixup is recomputed per pair (the price of the maddubs trick).
7768/// Returns raw i8·i8 dots; the caller applies scales and outliers.
7769#[cfg(target_arch = "x86_64")]
7770#[target_feature(enable = "avx2")]
7771unsafe fn dot_i8_i8_avx2_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
7772    // SAFETY: callers uphold slice-length contracts.
7773    unsafe {
7774        use core::arch::x86_64::*;
7775        let n = w0.len();
7776        let ones = _mm256_set1_epi16(1);
7777        let mut acc = [[_mm256_setzero_si256(); 4]; 2];
7778        let mut j = 0usize;
7779        while j + 32 <= n {
7780            let wv0 = _mm256_loadu_si256(w0.as_ptr().add(j) as *const __m256i);
7781            let wv1 = _mm256_loadu_si256(w1.as_ptr().add(j) as *const __m256i);
7782            let aw0 = _mm256_abs_epi8(wv0);
7783            let aw1 = _mm256_abs_epi8(wv1);
7784            for (k, x) in xs.iter().enumerate() {
7785                let xv = _mm256_loadu_si256(x.as_ptr().add(j) as *const __m256i);
7786                let p0 = _mm256_maddubs_epi16(aw0, _mm256_sign_epi8(xv, wv0));
7787                acc[0][k] = _mm256_add_epi32(acc[0][k], _mm256_madd_epi16(p0, ones));
7788                let p1 = _mm256_maddubs_epi16(aw1, _mm256_sign_epi8(xv, wv1));
7789                acc[1][k] = _mm256_add_epi32(acc[1][k], _mm256_madd_epi16(p1, ones));
7790            }
7791            j += 32;
7792        }
7793        let mut out = [[0i32; 4]; 2];
7794        for r in 0..2 {
7795            for k in 0..4 {
7796                let a = acc[r][k];
7797                let hi128 = _mm256_extracti128_si256::<1>(a);
7798                let s128 = _mm_add_epi32(_mm256_castsi256_si128(a), hi128);
7799                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7800                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7801                out[r][k] = _mm_cvtsi128_si32(s32);
7802            }
7803        }
7804        if j < n {
7805            for (k, x) in xs.iter().enumerate() {
7806                for i in j..n {
7807                    out[0][k] += (w0[i] as i8) as i32 * x[i] as i32;
7808                    out[1][k] += (w1[i] as i8) as i32 * x[i] as i32;
7809                }
7810            }
7811        }
7812        out
7813    }
7814}
7815
7816/// AVX2/VNNI q8 row dot with exact outlier correction (x86 mirror of
7817/// `row_dot_sdot` — same A8W8 contract). With AVX-512 VNNI the row goes
7818/// through the bias trick: Σ(w+128)·x via pure `vpdpbusd` (no per-lane
7819/// sign fixups), corrected by −128·Σx with Σx precomputed per split.
7820#[cfg(target_arch = "x86_64")]
7821#[inline]
7822fn row_dot_avx2(row: &[u8], act: &SplitAct) -> f32 {
7823    let dot = if avx512vnni_enabled() && row.len() >= 64 {
7824        (unsafe { dot_u8p128_i8_vnni(row, &act.xq) }) - 128 * act.xsum
7825    } else {
7826        unsafe { dot_i8_i8_avx2(row, &act.xq) }
7827    };
7828    let mut acc = dot as f32 * act.sx;
7829    for &(j, xv) in &act.outliers {
7830        acc += (row[j] as i8) as f32 * xv;
7831    }
7832    acc
7833}
7834
7835/// Σ (w[i]+128)·x[i] via pure `vpdpbusd` — the caller subtracts
7836/// 128·Σx. Four independent accumulators (dpbusd is ~5-cycle latency;
7837/// a single-acc loop runs latency-bound, measured on Granite Rapids).
7838#[cfg(target_arch = "x86_64")]
7839#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
7840unsafe fn dot_u8p128_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
7841    // SAFETY: callers uphold slice-length contracts (see call sites).
7842    unsafe {
7843        use core::arch::x86_64::*;
7844        let n = w.len();
7845        let flip = _mm512_set1_epi8(-128); // XOR 0x80: i8 w → u8 (w+128)
7846        #[inline(always)]
7847        unsafe fn step(
7848            w: *const u8,
7849            x: *const i8,
7850            flip: core::arch::x86_64::__m512i,
7851            acc: core::arch::x86_64::__m512i,
7852        ) -> core::arch::x86_64::__m512i {
7853            unsafe {
7854                use core::arch::x86_64::*;
7855                let wv = _mm512_xor_si512(_mm512_loadu_si512(w as *const _), flip);
7856                _mm512_dpbusd_epi32(acc, wv, _mm512_loadu_si512(x as *const _))
7857            }
7858        }
7859        let (mut a0, mut a1, mut a2, mut a3) = (
7860            _mm512_setzero_si512(),
7861            _mm512_setzero_si512(),
7862            _mm512_setzero_si512(),
7863            _mm512_setzero_si512(),
7864        );
7865        let mut j = 0usize;
7866        while j + 256 <= n {
7867            a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
7868            a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), flip, a1);
7869            a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), flip, a2);
7870            a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), flip, a3);
7871            j += 256;
7872        }
7873        while j + 64 <= n {
7874            a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
7875            j += 64;
7876        }
7877        let mut total = _mm512_reduce_add_epi32(_mm512_add_epi32(
7878            _mm512_add_epi32(a0, a1),
7879            _mm512_add_epi32(a2, a3),
7880        ));
7881        // Scalar tail: (w as i8) + 128 ≡ (w as u8) ^ 0x80.
7882        while j < n {
7883            total += ((w[j] ^ 0x80) as i32) * xq[j] as i32;
7884            j += 1;
7885        }
7886        total
7887    }
7888}
7889
7890/// One q4 row via AVX2: nibbles → centered i8 (unpacklo/hi restores the
7891/// writer's flat order, same as the NEON vzip pair), maddubs against
7892/// the pre-quantized activation group, × the group's f16 scale. Pair
7893/// saturation safe: |w|≤8, |x|≤127 → 2·8·127 ≪ 32767. Mirror of
7894/// `dot_q4_row_sdot`.
7895#[cfg(target_arch = "x86_64")]
7896#[target_feature(enable = "avx2")]
7897unsafe fn dot_q4_row_avx2(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
7898    // SAFETY: callers uphold slice-length contracts (16 packed bytes and
7899    // 2 scale bytes per group; xq.len() == gpr·GROUP_SIZE).
7900    unsafe {
7901        use core::arch::x86_64::*;
7902        let lomask = _mm_set1_epi8(0x0F);
7903        let eight = _mm256_set1_epi8(8);
7904        let ones = _mm256_set1_epi16(1);
7905        let mut acc = 0f32;
7906        for gi in 0..gpr {
7907            let g = g0 + gi;
7908            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
7909            let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
7910            let lo = _mm_and_si128(b, lomask);
7911            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
7912            let w = _mm256_sub_epi8(
7913                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
7914                eight,
7915            );
7916            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7917            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
7918            let d = _mm256_madd_epi16(p16, ones);
7919            let hi128 = _mm256_extracti128_si256::<1>(d);
7920            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
7921            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7922            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7923            acc += _mm_cvtsi128_si32(s32) as f32 * s;
7924        }
7925        acc
7926    }
7927}
7928
7929/// Two-activation q4 row via AVX2: nibbles unpacked ONCE per group,
7930/// both activations dotted against the same centered i8 register.
7931#[cfg(target_arch = "x86_64")]
7932#[target_feature(enable = "avx2")]
7933unsafe fn dot_q4_row_avx2_2(
7934    packed: &[u8],
7935    scales: &[u8],
7936    g0: usize,
7937    gpr: usize,
7938    xq1: &[i8],
7939    xq2: &[i8],
7940) -> (f32, f32) {
7941    // SAFETY: callers uphold slice-length contracts (see dot_q4_row_avx2).
7942    unsafe {
7943        use core::arch::x86_64::*;
7944        let lomask = _mm_set1_epi8(0x0F);
7945        let eight = _mm256_set1_epi8(8);
7946        let ones = _mm256_set1_epi16(1);
7947        let (mut acc1, mut acc2) = (0f32, 0f32);
7948        #[inline(always)]
7949        unsafe fn hsum(d: core::arch::x86_64::__m256i) -> i32 {
7950            unsafe {
7951                use core::arch::x86_64::*;
7952                let hi128 = _mm256_extracti128_si256::<1>(d);
7953                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
7954                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7955                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7956                _mm_cvtsi128_si32(s32)
7957            }
7958        }
7959        for gi in 0..gpr {
7960            let g = g0 + gi;
7961            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
7962            let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
7963            let lo = _mm_and_si128(b, lomask);
7964            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
7965            let w = _mm256_sub_epi8(
7966                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
7967                eight,
7968            );
7969            let aw = _mm256_abs_epi8(w);
7970            let x1 = _mm256_loadu_si256(xq1.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7971            let x2 = _mm256_loadu_si256(xq2.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7972            let d1 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x1, w)), ones);
7973            let d2 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x2, w)), ones);
7974            acc1 += hsum(d1) as f32 * s;
7975            acc2 += hsum(d2) as f32 * s;
7976        }
7977        (acc1, acc2)
7978    }
7979}
7980
7981/// One q8 row range via AVX2 (x86 mirror of `q8_range_sdot`).
7982#[cfg(target_arch = "x86_64")]
7983fn q8_range_avx2(
7984    q: &[u8],
7985    row_scale: &[f32],
7986    act: &SplitAct,
7987    cols: usize,
7988    out_addr: SendMut,
7989    start: usize,
7990    end: usize,
7991) {
7992    for o in start..end {
7993        let v = row_dot_avx2(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
7994        // SAFETY: disjoint row ranges per worker.
7995        unsafe { *out_addr.at(o) = v };
7996    }
7997}
7998
7999/// Two-input q8 row range via AVX2 (x86 mirror of `q8_range2_sdot`).
8000#[cfg(target_arch = "x86_64")]
8001#[allow(clippy::too_many_arguments)]
8002fn q8_range2_avx2(
8003    q: &[u8],
8004    row_scale: &[f32],
8005    a1: &SplitAct,
8006    a2: &SplitAct,
8007    cols: usize,
8008    p1: SendMut,
8009    p2: SendMut,
8010    start: usize,
8011    end: usize,
8012) {
8013    for o in start..end {
8014        let row = &q[o * cols..(o + 1) * cols];
8015        // SAFETY: disjoint row ranges per worker.
8016        unsafe {
8017            *p1.at(o) = row_dot_avx2(row, a1) * row_scale[o];
8018            *p2.at(o) = row_dot_avx2(row, a2) * row_scale[o];
8019        }
8020    }
8021}
8022
8023// ───────────────────── A8W8 SDOT path (port of vmfcore, ×1.78 decode) ─────────────────────
8024
8025/// ARMv8.6 i8mm (smmla): 32 int8 MACs per instruction vs sdot's 16 —
8026/// yet MEASURED 2.4× SLOWER than the blocked sdot on Apple silicon
8027/// (108 vs 264 GF/s): the on-the-fly vcombine packing and the two-
8028/// accumulator dependency chain swamp the MAC advantage, and Apple's
8029/// four SIMD pipes already keep sdot fed. OPT-IN (CMF_I8MM=1) for
8030/// field trials on Cortex-A710/X-class parts with two pipes, where the
8031/// balance may differ; a pre-interleaved weight layout (repack infra)
8032/// is the known path if it ever earns its keep.
8033#[cfg(target_arch = "aarch64")]
8034fn i8mm_enabled() -> bool {
8035    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8036    *ON.get_or_init(|| {
8037        std::env::var("CMF_I8MM").map(|v| v == "1").unwrap_or(false)
8038            && std::arch::is_aarch64_feature_detected!("i8mm")
8039    })
8040}
8041
8042/// SDOT enabled? Default ON when the CPU has ARMv8.2 dotprod;
8043/// `CMF_SDOT=0` disables (falls back to i8×f32 NEON).
8044/// (On non-ARM release builds only the test tolerance switch calls it.)
8045#[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
8046fn sdot_enabled() -> bool {
8047    use std::sync::OnceLock;
8048    static ON: OnceLock<bool> = OnceLock::new();
8049    *ON.get_or_init(|| {
8050        let want = std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true);
8051        if !want {
8052            return false;
8053        }
8054
8055        #[cfg(target_arch = "aarch64")]
8056        {
8057            if std::arch::is_aarch64_feature_detected!("dotprod") {
8058                return true;
8059            }
8060            #[cfg(target_os = "android")]
8061            {
8062                if let Ok(cpuinfo) = std::fs::read_to_string("/proc/cpuinfo") {
8063                    if cpuinfo.lines().any(|l| {
8064                        (l.starts_with("Features") || l.starts_with("features"))
8065                            && l.contains("asimddp")
8066                    }) {
8067                        return true;
8068                    }
8069                }
8070            }
8071            false
8072        }
8073        #[cfg(not(target_arch = "aarch64"))]
8074        {
8075            false
8076        }
8077    })
8078}
8079
8080/// Two-field activation split (≡ vmfcore `q8_split_prep`): outlier
8081/// channels (>8·rms) are computed exactly in f32; the bulk (outliers
8082/// zeroed → clean absmax) goes through int8 SDOT. Computed ONCE per
8083/// matvec, shared by all rows/workers.
8084struct SplitAct {
8085    xq: Vec<i8>,
8086    sx: f32,
8087    outliers: Vec<(usize, f32)>,
8088    /// Σ xq — the VNNI bias-trick correction (`(w+128)·x` sums need
8089    /// `−128·Σx`); one i32 per split, computed once per matvec.
8090    #[cfg_attr(not(target_arch = "x86_64"), allow(dead_code))]
8091    xsum: i32,
8092}
8093
8094thread_local! {
8095    /// Recycled xq buffers: split_act runs for every matvec (~200/token)
8096    /// and its hidden-size allocation was steady-state heap churn.
8097    static XQ_FREE: std::cell::RefCell<Vec<Vec<i8>>> =
8098        const { std::cell::RefCell::new(Vec::new()) };
8099}
8100
8101impl Drop for SplitAct {
8102    fn drop(&mut self) {
8103        let buf = std::mem::take(&mut self.xq);
8104        if buf.capacity() > 0 {
8105            XQ_FREE.with(|f| {
8106                let mut f = f.borrow_mut();
8107                if f.len() < 16 {
8108                    f.push(buf);
8109                }
8110            });
8111        }
8112    }
8113}
8114
8115thread_local! {
8116    /// One scratch row per WORKER, kept for the life of the thread.
8117    ///
8118    /// The kernels take a row of group scales per dispatch, and a fresh
8119    /// `vec![0f32; gpr]` inside the closure is one allocation per worker per
8120    /// dispatch — on the release checkpoint about six thousand a token, a
8121    /// quarter of everything the benchmark counts.
8122    static KROW: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
8123}
8124
8125/// Borrow `n` floats of the calling worker's scratch. Nothing inside a
8126/// kernel body borrows it again, which is what keeps the RefCell honest.
8127#[inline]
8128fn with_krow<R>(n: usize, f: impl FnOnce(&mut [f32]) -> R) -> R {
8129    KROW.with(|s| {
8130        let mut b = s.borrow_mut();
8131        if b.len() < n {
8132            b.resize(n, 0.0);
8133        }
8134        f(&mut b[..n])
8135    })
8136}
8137
8138fn split_act(x: &[f32]) -> SplitAct {
8139    let n = x.len();
8140    let rms = (x.iter().map(|&v| (v * v) as f64).sum::<f64>() / n.max(1) as f64).sqrt() as f32;
8141    let thr = 8.0 * rms;
8142    // One pass: collect outliers and the bulk absmax (outliers excluded —
8143    // identical to the old zero-then-fold over a copied buffer, minus the
8144    // full-vector copy).
8145    let mut outliers: Vec<(usize, f32)> = Vec::new();
8146    let mut amax = 0f32;
8147    for (j, &v) in x.iter().enumerate() {
8148        let a = v.abs();
8149        if a > thr {
8150            outliers.push((j, v));
8151        } else if a > amax {
8152            amax = a;
8153        }
8154    }
8155    let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
8156    let inv = 1.0 / sx;
8157    let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
8158    xq.clear();
8159    xq.reserve(n);
8160    if outliers.is_empty() {
8161        xq.extend(
8162            x.iter()
8163                .map(|&v| (v * inv).round().clamp(-127.0, 127.0) as i8),
8164        );
8165    } else {
8166        // Outlier slots quantize to 0 (their exact term is added later).
8167        xq.extend(x.iter().map(|&v| {
8168            if v.abs() > thr {
8169                0
8170            } else {
8171                (v * inv).round().clamp(-127.0, 127.0) as i8
8172            }
8173        }));
8174    }
8175    let xsum = xq.iter().map(|&v| v as i32).sum();
8176    SplitAct {
8177        xq,
8178        sx,
8179        outliers,
8180        xsum,
8181    }
8182}
8183
8184fn split_act_q8_2f(x: &[f32], col: &[f32]) -> SplitAct {
8185    let n = x.len();
8186    let rms = (x
8187        .iter()
8188        .zip(col)
8189        .map(|(&a, &c)| {
8190            let v = a * c;
8191            (v * v) as f64
8192        })
8193        .sum::<f64>()
8194        / n.max(1) as f64)
8195        .sqrt() as f32;
8196    let thr = 8.0 * rms;
8197
8198    let mut outliers = Vec::new();
8199    let mut amax = 0f32;
8200    for (j, (&a, &c)) in x.iter().zip(col).enumerate() {
8201        let v = a * c;
8202        let s = v.abs();
8203        if s > thr {
8204            outliers.push((j, v));
8205        } else if s > amax {
8206            amax = s;
8207        }
8208    }
8209
8210    let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
8211    let inv = 1.0 / sx;
8212    let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
8213    xq.clear();
8214    xq.reserve(n);
8215    if outliers.is_empty() {
8216        xq.extend(
8217            x.iter()
8218                .zip(col)
8219                .map(|(&a, &c)| ((a * c) * inv).round().clamp(-127.0, 127.0) as i8),
8220        );
8221    } else {
8222        xq.extend(x.iter().zip(col).map(|(&a, &c)| {
8223            let v = a * c;
8224            if v.abs() > thr {
8225                0
8226            } else {
8227                (v * inv).round().clamp(-127.0, 127.0) as i8
8228            }
8229        }));
8230    }
8231    let xsum = xq.iter().map(|&v| v as i32).sum();
8232    SplitAct {
8233        xq,
8234        sx,
8235        outliers,
8236        xsum,
8237    }
8238}
8239
8240/// int8(weight)·int8(activation) → i32 via `sdot` (inline asm — the
8241/// vdotq intrinsic is unstable; port of vmfcore `dot_i8_sdot`).
8242#[cfg(target_arch = "aarch64")]
8243#[target_feature(enable = "neon,dotprod")]
8244unsafe fn dot_i8_sdot(w: &[u8], xq: &[i8]) -> i32 {
8245    // SAFETY: callers uphold slice-length contracts (see call sites).
8246    unsafe {
8247        use core::arch::aarch64::*;
8248        use core::arch::asm;
8249        let wp = w.as_ptr() as *const i8;
8250        let n = w.len();
8251        let (mut a0, mut a1, mut a2, mut a3) = (
8252            vdupq_n_s32(0),
8253            vdupq_n_s32(0),
8254            vdupq_n_s32(0),
8255            vdupq_n_s32(0),
8256        );
8257        let mut i = 0;
8258        while i + 64 <= n {
8259            let (w0, x0) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
8260            let (w1, x1) = (vld1q_s8(wp.add(i + 16)), vld1q_s8(xq.as_ptr().add(i + 16)));
8261            let (w2, x2) = (vld1q_s8(wp.add(i + 32)), vld1q_s8(xq.as_ptr().add(i + 32)));
8262            let (w3, x3) = (vld1q_s8(wp.add(i + 48)), vld1q_s8(xq.as_ptr().add(i + 48)));
8263            asm!(
8264                "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
8265                "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
8266                "sdot {a2:v}.4s, {w2:v}.16b, {x2:v}.16b",
8267                "sdot {a3:v}.4s, {w3:v}.16b, {x3:v}.16b",
8268                a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
8269                w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
8270                w2 = in(vreg) w2, x2 = in(vreg) x2, w3 = in(vreg) w3, x3 = in(vreg) x3,
8271                options(pure, nomem, nostack),
8272            );
8273            i += 64;
8274        }
8275        while i + 16 <= n {
8276            let (wv, xv) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
8277            asm!("sdot {a:v}.4s, {w:v}.16b, {x:v}.16b",
8278                 a = inout(vreg) a0, w = in(vreg) wv, x = in(vreg) xv, options(pure, nomem, nostack));
8279            i += 16;
8280        }
8281        let mut s = vaddvq_s32(vaddq_s32(vaddq_s32(a0, a1), vaddq_s32(a2, a3)));
8282        while i < n {
8283            s += (*wp.add(i)) as i32 * xq[i] as i32;
8284            i += 1;
8285        }
8286        s
8287    }
8288}
8289
8290/// Row-blocked SDOT: 4 output rows per pass — the activation chunk is
8291/// loaded once and reused, 4 independent accumulators hide sdot latency
8292/// (port of vmfcore `dot_i8_sdot_4rows`).
8293#[cfg(target_arch = "aarch64")]
8294#[target_feature(enable = "neon,dotprod")]
8295unsafe fn dot_i8_sdot_4rows(w0: &[u8], w1: &[u8], w2: &[u8], w3: &[u8], xq: &[i8]) -> [i32; 4] {
8296    // SAFETY: callers uphold slice-length contracts (see call sites).
8297    unsafe {
8298        use core::arch::aarch64::*;
8299        use core::arch::asm;
8300        let n = xq.len();
8301        let px = xq.as_ptr();
8302        let (p0, p1, p2, p3) = (
8303            w0.as_ptr() as *const i8,
8304            w1.as_ptr() as *const i8,
8305            w2.as_ptr() as *const i8,
8306            w3.as_ptr() as *const i8,
8307        );
8308        let (mut a0, mut a1, mut a2, mut a3) = (
8309            vdupq_n_s32(0),
8310            vdupq_n_s32(0),
8311            vdupq_n_s32(0),
8312            vdupq_n_s32(0),
8313        );
8314        let mut i = 0;
8315        while i + 16 <= n {
8316            let x = vld1q_s8(px.add(i));
8317            let v0 = vld1q_s8(p0.add(i));
8318            let v1 = vld1q_s8(p1.add(i));
8319            let v2 = vld1q_s8(p2.add(i));
8320            let v3 = vld1q_s8(p3.add(i));
8321            asm!(
8322                "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
8323                "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
8324                "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
8325                "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
8326                a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
8327                v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
8328                options(pure, nomem, nostack),
8329            );
8330            i += 16;
8331        }
8332        let mut r = [
8333            vaddvq_s32(a0),
8334            vaddvq_s32(a1),
8335            vaddvq_s32(a2),
8336            vaddvq_s32(a3),
8337        ];
8338        while i < n {
8339            let xi = *px.add(i) as i32;
8340            r[0] += (*p0.add(i)) as i32 * xi;
8341            r[1] += (*p1.add(i)) as i32 * xi;
8342            r[2] += (*p2.add(i)) as i32 * xi;
8343            r[3] += (*p3.add(i)) as i32 * xi;
8344            i += 1;
8345        }
8346        r
8347    }
8348}
8349
8350/// 4 interleaved rows in one pass: the repacked group is [r0[c], r1[c],
8351/// r2[c], r3[c]] per 16-byte chunk, so each iteration reads ONE 64-byte
8352/// line plus the shared activation chunk — a single sequential weight
8353/// stream per worker. Per-row accumulation is the same one-accumulator
8354/// scheme as `dot_i8_sdot_4rows`; integer sums are exact, so outputs
8355/// are bit-identical to the mmap-layout kernel.
8356#[cfg(target_arch = "aarch64")]
8357#[target_feature(enable = "neon,dotprod")]
8358unsafe fn dot_i8_sdot_4rows_il(g: &[u8], xq: &[i8]) -> [i32; 4] {
8359    // SAFETY: callers uphold slice-length contracts (g.len() == 4·n,
8360    // n % 16 == 0 — guaranteed by the repack gate).
8361    unsafe {
8362        use core::arch::aarch64::*;
8363        use core::arch::asm;
8364        let n = xq.len();
8365        let px = xq.as_ptr();
8366        let pg = g.as_ptr() as *const i8;
8367        let (mut a0, mut a1, mut a2, mut a3) = (
8368            vdupq_n_s32(0),
8369            vdupq_n_s32(0),
8370            vdupq_n_s32(0),
8371            vdupq_n_s32(0),
8372        );
8373        let mut i = 0;
8374        while i + 16 <= n {
8375            let x = vld1q_s8(px.add(i));
8376            let base = pg.add(4 * i);
8377            let v0 = vld1q_s8(base);
8378            let v1 = vld1q_s8(base.add(16));
8379            let v2 = vld1q_s8(base.add(32));
8380            let v3 = vld1q_s8(base.add(48));
8381            asm!(
8382                "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
8383                "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
8384                "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
8385                "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
8386                a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
8387                v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
8388                options(pure, nomem, nostack),
8389            );
8390            i += 16;
8391        }
8392        [
8393            vaddvq_s32(a0),
8394            vaddvq_s32(a1),
8395            vaddvq_s32(a2),
8396            vaddvq_s32(a3),
8397        ]
8398    }
8399}
8400
8401/// One q8 row range via SDOT (4-row blocks + tail) — the body of
8402/// `qmatvec`'s hot loop, extracted so multi-matrix jobs can drive the
8403/// SAME kernel for several tensors under one pool dispatch. `rep` — the
8404/// load-time interleaved repack (empty = mmap layout only); rows outside
8405/// full 4-row groups always come from the mmap layout.
8406#[cfg(target_arch = "aarch64")]
8407fn q8_range_sdot(
8408    q: &[u8],
8409    rep: &[u8],
8410    row_scale: &[f32],
8411    act: &SplitAct,
8412    cols: usize,
8413    out_addr: SendMut,
8414    start: usize,
8415    end: usize,
8416) {
8417    let mut o = start;
8418    // Leading rows to the group boundary (repack path only): the pool
8419    // splits row ranges arbitrarily, groups are absolute.
8420    if !rep.is_empty() {
8421        while o < end && o % 4 != 0 {
8422            let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
8423            unsafe { *out_addr.at(o) = v };
8424            o += 1;
8425        }
8426    }
8427    while o + 4 <= end {
8428        let r = if rep.is_empty() {
8429            unsafe {
8430                dot_i8_sdot_4rows(
8431                    &q[o * cols..(o + 1) * cols],
8432                    &q[(o + 1) * cols..(o + 2) * cols],
8433                    &q[(o + 2) * cols..(o + 3) * cols],
8434                    &q[(o + 3) * cols..(o + 4) * cols],
8435                    &act.xq,
8436                )
8437            }
8438        } else {
8439            unsafe { dot_i8_sdot_4rows_il(&rep[o * cols..(o + 4) * cols], &act.xq) }
8440        };
8441        for k in 0..4 {
8442            let mut acc = r[k] as f32 * act.sx;
8443            for &(j, xv) in &act.outliers {
8444                acc += (q[(o + k) * cols + j] as i8) as f32 * xv;
8445            }
8446            // SAFETY: disjoint row ranges per worker.
8447            unsafe { *out_addr.at(o + k) = acc * row_scale[o + k] };
8448        }
8449        o += 4;
8450    }
8451    while o < end {
8452        let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
8453        unsafe { *out_addr.at(o) = v };
8454        o += 1;
8455    }
8456}
8457
8458/// Two-input q8 row range via SDOT — `qmatvec2`'s hot loop, extracted
8459/// for the fused pair multi-matrix job (`matvec2_many`).
8460#[cfg(target_arch = "aarch64")]
8461#[allow(clippy::too_many_arguments)]
8462fn q8_range2_sdot(
8463    q: &[u8],
8464    row_scale: &[f32],
8465    a1: &SplitAct,
8466    a2: &SplitAct,
8467    cols: usize,
8468    p1: SendMut,
8469    p2: SendMut,
8470    start: usize,
8471    end: usize,
8472) {
8473    for o in start..end {
8474        let row = &q[o * cols..(o + 1) * cols];
8475        // SAFETY: disjoint row ranges per worker.
8476        unsafe {
8477            *p1.at(o) = row_dot_sdot(row, a1) * row_scale[o];
8478            *p2.at(o) = row_dot_sdot(row, a2) * row_scale[o];
8479        }
8480    }
8481}
8482
8483/// Two-input q8 row range, f32 kernel (non-SDOT) — same extraction.
8484#[allow(clippy::too_many_arguments)]
8485fn q8_range2_f32(
8486    q: &[u8],
8487    row_scale: &[f32],
8488    x1: &[f32],
8489    x2: &[f32],
8490    cols: usize,
8491    p1: SendMut,
8492    p2: SendMut,
8493    start: usize,
8494    end: usize,
8495) {
8496    for o in start..end {
8497        let row = &q[o * cols..(o + 1) * cols];
8498        // SAFETY: disjoint row ranges per worker.
8499        unsafe {
8500            *p1.at(o) = dot_i8_f32(row, x1) * row_scale[o];
8501            *p2.at(o) = dot_i8_f32(row, x2) * row_scale[o];
8502        }
8503    }
8504}
8505
8506/// Scalar/NEON-f32 q8 row range (non-SDOT platforms) — same extraction.
8507fn q8_range_f32(
8508    q: &[u8],
8509    row_scale: &[f32],
8510    xs: &[f32],
8511    cols: usize,
8512    out_addr: SendMut,
8513    start: usize,
8514    end: usize,
8515) {
8516    for o in start..end {
8517        let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
8518        // SAFETY: disjoint row ranges per worker.
8519        unsafe { *out_addr.at(o) = v };
8520    }
8521}
8522
8523/// SDOT row dot with exact outlier correction:
8524/// `dot = sdot(w, xq)·sx + Σ_outl w[j]·x[j]` (then × row_scale by caller).
8525#[cfg(target_arch = "aarch64")]
8526#[inline]
8527fn row_dot_sdot(row: &[u8], act: &SplitAct) -> f32 {
8528    let mut acc = unsafe { dot_i8_sdot(row, &act.xq) } as f32 * act.sx;
8529    for &(j, xv) in &act.outliers {
8530        acc += (row[j] as i8) as f32 * xv;
8531    }
8532    acc
8533}
8534
8535/// One q4 row via SDOT: each 32-group's nibbles unpack to centered i8
8536/// (nib−8 ∈ [−8,7]), int8×int8 `sdot` against the pre-quantized
8537/// activation group, × the group's f16 scale. Returns Σ_g dot_g·s_g;
8538/// the caller multiplies by the activation scale and adds the exact
8539/// outlier terms (port of vmfcore `dot_q4_block_sdot`, +23% measured).
8540/// Nibble order matches the writer: element 2k = low nibble, 2k+1 = high
8541/// → zip(lo,hi) restores flat order.
8542#[cfg(target_arch = "aarch64")]
8543#[target_feature(enable = "neon,dotprod")]
8544unsafe fn dot_q4_row_sdot(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
8545    // SAFETY: callers uphold slice-length contracts (16 packed bytes and
8546    // 2 scale bytes per group; xq.len() == gpr·GROUP_SIZE).
8547    unsafe {
8548        use core::arch::aarch64::*;
8549        use core::arch::asm;
8550        let lomask = vdupq_n_u8(0x0F);
8551        let eight = vdupq_n_s8(8);
8552        let mut acc = 0f32;
8553        for gi in 0..gpr {
8554            let g = g0 + gi;
8555            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
8556            let b = vld1q_u8(packed.as_ptr().add(g * 16));
8557            let lo = vandq_u8(b, lomask);
8558            let hi = vshrq_n_u8::<4>(b);
8559            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
8560            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
8561            let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
8562            let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
8563            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
8564            asm!(
8565                "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
8566                "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
8567                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
8568                e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
8569                options(pure, nomem, nostack),
8570            );
8571            acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
8572        }
8573        acc
8574    }
8575}
8576
8577/// Two-activation q4 row via SDOT: the nibble unpack (the expensive
8578/// part) happens ONCE per group; both pre-quantized activations are
8579/// dotted against the same centered i8 registers. Per-lane math matches
8580/// `dot_q4_row_sdot` exactly.
8581#[cfg(target_arch = "aarch64")]
8582#[target_feature(enable = "neon,dotprod")]
8583unsafe fn dot_q4_row_sdot2(
8584    packed: &[u8],
8585    scales: &[u8],
8586    g0: usize,
8587    gpr: usize,
8588    xq1: &[i8],
8589    xq2: &[i8],
8590) -> (f32, f32) {
8591    // SAFETY: callers uphold slice-length contracts (16 packed bytes and
8592    // 2 scale bytes per group; xq*.len() == gpr·GROUP_SIZE).
8593    unsafe {
8594        use core::arch::aarch64::*;
8595        use core::arch::asm;
8596        let lomask = vdupq_n_u8(0x0F);
8597        let eight = vdupq_n_s8(8);
8598        let (mut acc1, mut acc2) = (0f32, 0f32);
8599        for gi in 0..gpr {
8600            let g = g0 + gi;
8601            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
8602            let b = vld1q_u8(packed.as_ptr().add(g * 16));
8603            let lo = vandq_u8(b, lomask);
8604            let hi = vshrq_n_u8::<4>(b);
8605            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
8606            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
8607            let x10 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE));
8608            let x11 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE + 16));
8609            let x20 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE));
8610            let x21 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE + 16));
8611            let (mut a0, mut a1, mut b0, mut b1) = (
8612                vdupq_n_s32(0),
8613                vdupq_n_s32(0),
8614                vdupq_n_s32(0),
8615                vdupq_n_s32(0),
8616            );
8617            asm!(
8618                "sdot {a0:v}.4s, {e0:v}.16b, {x10:v}.16b",
8619                "sdot {a1:v}.4s, {e1:v}.16b, {x11:v}.16b",
8620                "sdot {b0:v}.4s, {e0:v}.16b, {x20:v}.16b",
8621                "sdot {b1:v}.4s, {e1:v}.16b, {x21:v}.16b",
8622                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
8623                b0 = inout(vreg) b0, b1 = inout(vreg) b1,
8624                e0 = in(vreg) e0, e1 = in(vreg) e1,
8625                x10 = in(vreg) x10, x11 = in(vreg) x11,
8626                x20 = in(vreg) x20, x21 = in(vreg) x21,
8627                options(pure, nomem, nostack),
8628            );
8629            acc1 += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
8630            acc2 += vaddvq_s32(vaddq_s32(b0, b1)) as f32 * s;
8631        }
8632        (acc1, acc2)
8633    }
8634}
8635
8636// ───────────────────── fused int8 kernels ─────────────────────
8637
8638/// `acc += w · row` where the row is centered i8 — NEON widen+fma on
8639/// aarch64, scalar elsewhere. The KV-cache q8 value path rides on this.
8640#[inline]
8641pub(crate) fn axpy_i8_f32(acc: &mut [f32], row: &[i8], w: f32) {
8642    #[cfg(target_arch = "aarch64")]
8643    unsafe {
8644        return axpy_i8_f32_neon(acc, row, w);
8645    }
8646    #[cfg(target_arch = "x86_64")]
8647    if avx2_enabled() {
8648        return unsafe { axpy_i8_f32_avx2(acc, row, w) };
8649    }
8650    #[allow(unreachable_code)]
8651    {
8652        for (a, &b) in acc.iter_mut().zip(row) {
8653            *a += w * b as f32;
8654        }
8655    }
8656}
8657
8658/// i8→f32 axpy via AVX2/FMA (x86 mirror of `axpy_i8_f32_neon`).
8659#[cfg(target_arch = "x86_64")]
8660#[target_feature(enable = "avx2,fma")]
8661unsafe fn axpy_i8_f32_avx2(acc: &mut [f32], row: &[i8], w: f32) {
8662    // SAFETY: callers uphold slice-length contracts (see call sites).
8663    unsafe {
8664        use core::arch::x86_64::*;
8665        let n = acc.len().min(row.len());
8666        let ap = acc.as_mut_ptr();
8667        let rp = row.as_ptr();
8668        let wv = _mm256_set1_ps(w);
8669        let mut j = 0usize;
8670        while j + 16 <= n {
8671            let rb = _mm_loadu_si128(rp.add(j) as *const __m128i);
8672            let lo = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(rb));
8673            let hi = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(rb)));
8674            let v0 = _mm256_fmadd_ps(wv, lo, _mm256_loadu_ps(ap.add(j)));
8675            let v1 = _mm256_fmadd_ps(wv, hi, _mm256_loadu_ps(ap.add(j + 8)));
8676            _mm256_storeu_ps(ap.add(j), v0);
8677            _mm256_storeu_ps(ap.add(j + 8), v1);
8678            j += 16;
8679        }
8680        while j < n {
8681            *ap.add(j) += w * (*rp.add(j)) as f32;
8682            j += 1;
8683        }
8684    }
8685}
8686
8687#[cfg(target_arch = "aarch64")]
8688#[target_feature(enable = "neon")]
8689unsafe fn axpy_i8_f32_neon(acc: &mut [f32], row: &[i8], w: f32) {
8690    // SAFETY: callers uphold slice-length contracts (see call sites).
8691    unsafe {
8692        use core::arch::aarch64::*;
8693        let n = acc.len().min(row.len());
8694        let ap = acc.as_mut_ptr();
8695        let rp = row.as_ptr();
8696        let wv = vdupq_n_f32(w);
8697        let mut j = 0usize;
8698        while j + 16 <= n {
8699            let rb = vld1q_s8(rp.add(j));
8700            let lo = vmovl_s8(vget_low_s8(rb));
8701            let hi = vmovl_s8(vget_high_s8(rb));
8702            for (off, half) in [(0, lo), (8, hi)] {
8703                let f0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(half)));
8704                let f1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(half)));
8705                let o = j + off;
8706                vst1q_f32(ap.add(o), vfmaq_f32(vld1q_f32(ap.add(o)), wv, f0));
8707                vst1q_f32(ap.add(o + 4), vfmaq_f32(vld1q_f32(ap.add(o + 4)), wv, f1));
8708            }
8709            j += 16;
8710        }
8711        while j < n {
8712            *ap.add(j) += w * (*rp.add(j)) as f32;
8713            j += 1;
8714        }
8715    }
8716}
8717
8718/// i8 row · f32 x. NEON on aarch64 (ported from vmfcore `dot_i8_f32_neon`,
8719/// ≈9× scalar), scalar elsewhere.
8720#[inline]
8721pub(crate) fn dot_i8_f32(w: &[u8], x: &[f32]) -> f32 {
8722    #[cfg(target_arch = "aarch64")]
8723    unsafe {
8724        return dot_i8_f32_neon(w, x);
8725    }
8726    #[cfg(target_arch = "x86_64")]
8727    if avx2_enabled() {
8728        return unsafe { dot_i8_f32_avx2(w, x) };
8729    }
8730    #[allow(unreachable_code)]
8731    {
8732        let mut sum = 0.0f32;
8733        for (j, &b) in w.iter().enumerate() {
8734            sum += (b as i8) as f32 * x[j];
8735        }
8736        sum
8737    }
8738}
8739
8740/// i8 row · (x ⊙ col_field) — the q8_2f row dot with the θ col-field
8741/// folded into the product (no prescaled copy of x). NEON on aarch64,
8742/// scalar elsewhere. Used by the active-neuron path `row_dot`.
8743#[inline]
8744fn dot_i8_col_f32(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
8745    #[cfg(target_arch = "aarch64")]
8746    unsafe {
8747        return dot_i8_col_f32_neon(w, x, col);
8748    }
8749    #[allow(unreachable_code)]
8750    {
8751        let mut sum = 0.0f32;
8752        for (j, &b) in w.iter().enumerate() {
8753            sum += (b as i8) as f32 * x[j] * col[j];
8754        }
8755        sum
8756    }
8757}
8758
8759#[cfg(target_arch = "aarch64")]
8760#[target_feature(enable = "neon")]
8761unsafe fn dot_i8_col_f32_neon(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
8762    // SAFETY: callers uphold slice-length contracts (see call sites).
8763    unsafe {
8764        use core::arch::aarch64::*;
8765        let n = x.len();
8766        let wp = w.as_ptr() as *const i8;
8767        let xp = x.as_ptr();
8768        let cp = col.as_ptr();
8769        let (mut a0, mut a1, mut a2, mut a3) = (
8770            vdupq_n_f32(0.0),
8771            vdupq_n_f32(0.0),
8772            vdupq_n_f32(0.0),
8773            vdupq_n_f32(0.0),
8774        );
8775        let mut j = 0usize;
8776        while j + 16 <= n {
8777            let wb = vld1q_s8(wp.add(j));
8778            let lo = vmovl_s8(vget_low_s8(wb));
8779            let hi = vmovl_s8(vget_high_s8(wb));
8780            let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
8781            let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
8782            let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
8783            let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
8784            a0 = vfmaq_f32(
8785                a0,
8786                w0,
8787                vmulq_f32(vld1q_f32(xp.add(j)), vld1q_f32(cp.add(j))),
8788            );
8789            a1 = vfmaq_f32(
8790                a1,
8791                w1,
8792                vmulq_f32(vld1q_f32(xp.add(j + 4)), vld1q_f32(cp.add(j + 4))),
8793            );
8794            a2 = vfmaq_f32(
8795                a2,
8796                w2,
8797                vmulq_f32(vld1q_f32(xp.add(j + 8)), vld1q_f32(cp.add(j + 8))),
8798            );
8799            a3 = vfmaq_f32(
8800                a3,
8801                w3,
8802                vmulq_f32(vld1q_f32(xp.add(j + 12)), vld1q_f32(cp.add(j + 12))),
8803            );
8804            j += 16;
8805        }
8806        let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
8807        while j < n {
8808            sum += (*wp.add(j)) as f32 * *xp.add(j) * *cp.add(j);
8809            j += 1;
8810        }
8811        sum
8812    }
8813}
8814
8815#[cfg(target_arch = "aarch64")]
8816#[target_feature(enable = "neon")]
8817unsafe fn dot_i8_f32_neon(w: &[u8], x: &[f32]) -> f32 {
8818    // SAFETY: callers uphold slice-length contracts (see call sites).
8819    unsafe {
8820        use core::arch::aarch64::*;
8821        let n = x.len();
8822        let wp = w.as_ptr() as *const i8;
8823        let xp = x.as_ptr();
8824        let (mut a0, mut a1, mut a2, mut a3) = (
8825            vdupq_n_f32(0.0),
8826            vdupq_n_f32(0.0),
8827            vdupq_n_f32(0.0),
8828            vdupq_n_f32(0.0),
8829        );
8830        let mut j = 0usize;
8831        while j + 16 <= n {
8832            let wb = vld1q_s8(wp.add(j));
8833            let lo = vmovl_s8(vget_low_s8(wb));
8834            let hi = vmovl_s8(vget_high_s8(wb));
8835            let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
8836            let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
8837            let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
8838            let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
8839            a0 = vfmaq_f32(a0, w0, vld1q_f32(xp.add(j)));
8840            a1 = vfmaq_f32(a1, w1, vld1q_f32(xp.add(j + 4)));
8841            a2 = vfmaq_f32(a2, w2, vld1q_f32(xp.add(j + 8)));
8842            a3 = vfmaq_f32(a3, w3, vld1q_f32(xp.add(j + 12)));
8843            j += 16;
8844        }
8845        let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
8846        while j < n {
8847            sum += (*wp.add(j)) as f32 * *xp.add(j);
8848            j += 1;
8849        }
8850        sum
8851    }
8852}
8853
8854#[allow(clippy::too_many_arguments)]
8855fn qmatvec(
8856    q: &[u8],
8857    rep: &[u8],
8858    row_scale: &[f32],
8859    x: &[f32],
8860    col_field: &[f32],
8861    dtype: TensorDtype,
8862    rows: usize,
8863    cols: usize,
8864    out: &mut [f32],
8865    pool: Option<&Pool>,
8866) {
8867    debug_assert_eq!(out.len(), rows);
8868    #[cfg(not(target_arch = "aarch64"))]
8869    let _ = rep;
8870
8871    #[cfg(target_arch = "aarch64")]
8872    if sdot_enabled() {
8873        let act = if dtype == TensorDtype::Q8_2f {
8874            split_act_q8_2f(x, col_field)
8875        } else {
8876            split_act(x)
8877        };
8878        let out_addr = SendMut(out.as_mut_ptr());
8879        let run_range = |start: usize, end: usize| {
8880            q8_range_sdot(q, rep, row_scale, &act, cols, out_addr, start, end)
8881        };
8882        match pool {
8883            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
8884            _ => run_range(0, rows),
8885        }
8886        return;
8887    }
8888    // x86 A8W8 via AVX2 maddubs — same quantized-activation contract as
8889    // the SDOT path (CMF_AVX2=0 keeps the exact i8×f32 loop).
8890    #[cfg(target_arch = "x86_64")]
8891    if avx2_a8w8_enabled() {
8892        let act = if dtype == TensorDtype::Q8_2f {
8893            split_act_q8_2f(x, col_field)
8894        } else {
8895            split_act(x)
8896        };
8897        let out_addr = SendMut(out.as_mut_ptr());
8898        let run_range = |start: usize, end: usize| {
8899            q8_range_avx2(q, row_scale, &act, cols, out_addr, start, end)
8900        };
8901        match pool {
8902            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
8903            _ => run_range(0, rows),
8904        }
8905        return;
8906    }
8907
8908    prescale_with(x, col_field, dtype, 1, |xs| {
8909        let out_addr = SendMut(out.as_mut_ptr());
8910        let run_range = move |start: usize, end: usize| {
8911            for o in start..end {
8912                let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
8913                // SAFETY: disjoint row ranges per worker.
8914                unsafe { *out_addr.at(o) = v };
8915            }
8916        };
8917        match pool {
8918            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
8919            _ => run_range(0, rows),
8920        }
8921    });
8922}
8923
8924#[allow(clippy::too_many_arguments)]
8925fn qmatvec2(
8926    q: &[u8],
8927    row_scale: &[f32],
8928    x1: &[f32],
8929    x2: &[f32],
8930    col_field: &[f32],
8931    dtype: TensorDtype,
8932    rows: usize,
8933    cols: usize,
8934    o1: &mut [f32],
8935    o2: &mut [f32],
8936    pool: Option<&Pool>,
8937) {
8938    #[cfg(target_arch = "aarch64")]
8939    if sdot_enabled() {
8940        let a1s = if dtype == TensorDtype::Q8_2f {
8941            split_act_q8_2f(x1, col_field)
8942        } else {
8943            split_act(x1)
8944        };
8945        let a2s = if dtype == TensorDtype::Q8_2f {
8946            split_act_q8_2f(x2, col_field)
8947        } else {
8948            split_act(x2)
8949        };
8950        let p1 = SendMut(o1.as_mut_ptr());
8951        let p2 = SendMut(o2.as_mut_ptr());
8952        let run_range = |start: usize, end: usize| {
8953            q8_range2_sdot(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
8954        };
8955        match pool {
8956            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
8957            _ => run_range(0, rows),
8958        }
8959        return;
8960    }
8961    #[cfg(target_arch = "x86_64")]
8962    if avx2_a8w8_enabled() {
8963        let a1s = if dtype == TensorDtype::Q8_2f {
8964            split_act_q8_2f(x1, col_field)
8965        } else {
8966            split_act(x1)
8967        };
8968        let a2s = if dtype == TensorDtype::Q8_2f {
8969            split_act_q8_2f(x2, col_field)
8970        } else {
8971            split_act(x2)
8972        };
8973        let p1 = SendMut(o1.as_mut_ptr());
8974        let p2 = SendMut(o2.as_mut_ptr());
8975        let run_range = |start: usize, end: usize| {
8976            q8_range2_avx2(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
8977        };
8978        match pool {
8979            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
8980            _ => run_range(0, rows),
8981        }
8982        return;
8983    }
8984
8985    prescale_with(x1, col_field, dtype, 1, |x1s| {
8986        prescale_with(x2, col_field, dtype, 2, |x2s| {
8987            let p1 = SendMut(o1.as_mut_ptr());
8988            let p2 = SendMut(o2.as_mut_ptr());
8989            let run_range = move |start: usize, end: usize| {
8990                for o in start..end {
8991                    let row = &q[o * cols..(o + 1) * cols];
8992                    let s1 = dot_i8_f32(row, x1s) * row_scale[o];
8993                    let s2 = dot_i8_f32(row, x2s) * row_scale[o];
8994                    // SAFETY: disjoint row ranges per worker.
8995                    unsafe {
8996                        *p1.at(o) = s1;
8997                        *p2.at(o) = s2;
8998                    }
8999                }
9000            };
9001            match pool {
9002                Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
9003                _ => run_range(0, rows),
9004            }
9005        });
9006    });
9007}
9008
9009#[derive(Clone, Copy)]
9010struct SendMut(*mut f32);
9011unsafe impl Send for SendMut {}
9012unsafe impl Sync for SendMut {}
9013
9014impl SendMut {
9015    #[inline]
9016    fn at(self, i: usize) -> *mut f32 {
9017        unsafe { self.0.add(i) }
9018    }
9019}
9020
9021#[cfg(test)]
9022mod tests {
9023    use super::*;
9024
9025    #[test]
9026    fn f32_matvec_matches_matvec_rows_bitexact() {
9027        let (rows, cols) = (300, 40);
9028        let w: Vec<f32> = (0..rows * cols).map(|i| (i as f32 * 0.017).sin()).collect();
9029        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.05).cos()).collect();
9030        let qt = QTensor::from_f32(w.clone(), rows, cols);
9031
9032        let mut a = vec![0.0f32; rows];
9033        matvec_rows(None, &w, &x, &mut a);
9034        let mut b = vec![0.0f32; rows];
9035        qt.matvec(&x, &mut b, None);
9036        assert_eq!(a, b);
9037    }
9038
9039    #[test]
9040    fn sdot_kernel_exact_on_grid() {
9041        // Activations already on the i8 grid (±1 with amax=1 → sx=1/127,
9042        // xq=±127 dequantizes EXACTLY) → the SDOT path must match the
9043        // exact f32 dot to float rounding. This isolates kernel
9044        // correctness from quantization noise.
9045        eprintln!("sdot_enabled = {}", sdot_enabled());
9046        let (rows, cols) = (9, 80); // odd rows → exercises 4-row + tail
9047        let w: Vec<u8> = (0..rows * cols)
9048            .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
9049            .collect();
9050        let scales: Vec<f32> = (0..rows).map(|o| 0.005 + o as f32 * 0.001).collect();
9051        let x: Vec<f32> = (0..cols)
9052            .map(|i| match i % 3 {
9053                0 => 1.0,
9054                1 => -1.0,
9055                _ => 0.0,
9056            })
9057            .collect();
9058        let mut a = vec![0.0f32; rows];
9059        qmatvec(
9060            &w,
9061            &[],
9062            &scales,
9063            &x,
9064            &[],
9065            TensorDtype::Q8Row,
9066            rows,
9067            cols,
9068            &mut a,
9069            None,
9070        );
9071        for o in 0..rows {
9072            let mut acc = 0.0f32;
9073            for j in 0..cols {
9074                acc += (w[o * cols + j] as i8) as f32 * x[j];
9075            }
9076            let expect = acc * scales[o];
9077            assert!(
9078                (a[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
9079                "row {o}: {} vs {expect}",
9080                a[o]
9081            );
9082        }
9083    }
9084
9085    #[test]
9086    fn q1_tbl_fast_path_matches_reference() {
9087        // gpr = 8 exercises the TBL pair-load fast loop, and the LAST
9088        // row's final 4-tile window trips the 4B-overread guard (the
9089        // payload ends exactly at the last tile) — both paths must
9090        // agree with the dequant reference.
9091        let (rows, cols) = (5, 256);
9092        let gpr = cols / GROUP_SIZE;
9093        let mut bytes = Vec::new();
9094        for t in 0..rows * gpr {
9095            let s = 0.007 + (t % 11) as f32 * 0.004;
9096            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
9097            for j in 0..4 {
9098                bytes.push(((t * 53 + j * 89 + 7) % 249) as u8);
9099            }
9100        }
9101        let x: Vec<f32> = (0..cols)
9102            .map(|i| if (i * 5) % 7 < 3 { 1.0 } else { -1.0 })
9103            .collect();
9104        let mut w = vec![0.0f32; rows * cols];
9105        cortiq_core::quant::dequant_q1(&bytes, &mut w);
9106        let mut got = vec![0.0f32; rows];
9107        q1_matvec(&bytes, &x, rows, cols, &mut got, None);
9108        for o in 0..rows {
9109            let expect: f32 = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
9110            assert!(
9111                (got[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
9112                "row {o}: {} vs {expect}",
9113                got[o]
9114            );
9115        }
9116        // Blocked 1×4 batch (b=5: one quad + remainder) must equal the
9117        // single-matvec path bit-for-bit.
9118        let b = 5usize;
9119        let mut xs_all = Vec::new();
9120        for bi in 0..b {
9121            xs_all.extend(x.iter().map(|v| if bi % 2 == 0 { *v } else { -*v }));
9122        }
9123        let mut mm = vec![0.0f32; b * rows];
9124        q1_matmat(&bytes, &xs_all, b, rows, cols, &mut mm, None);
9125        for bi in 0..b {
9126            let mut single = vec![0.0f32; rows];
9127            q1_matvec(
9128                &bytes,
9129                &xs_all[bi * cols..(bi + 1) * cols],
9130                rows,
9131                cols,
9132                &mut single,
9133                None,
9134            );
9135            assert_eq!(&mm[bi * rows..(bi + 1) * rows], &single[..], "stream {bi}");
9136        }
9137    }
9138
9139    #[test]
9140    fn q1_kernels_match_exact_reference() {
9141        // Synthetic q1 payload: 6-byte tiles [f16 scale][4B bits].
9142        let (rows, cols) = (7, 96);
9143        let gpr = cols / GROUP_SIZE;
9144        let mut bytes = Vec::new();
9145        for t in 0..rows * gpr {
9146            let s = 0.01 + (t % 13) as f32 * 0.003;
9147            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
9148            for j in 0..4 {
9149                bytes.push(((t * 31 + j * 97) % 251) as u8);
9150            }
9151        }
9152        // On-grid activations (±1, amax 1) → the SDOT path is exact.
9153        let x: Vec<f32> = (0..cols)
9154            .map(|i| if i % 3 == 0 { 1.0 } else { -1.0 })
9155            .collect();
9156        // Reference through the core dequant.
9157        let mut w = vec![0.0f32; rows * cols];
9158        cortiq_core::quant::dequant_q1(&bytes, &mut w);
9159        let mut expect = vec![0.0f32; rows];
9160        for o in 0..rows {
9161            expect[o] = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
9162        }
9163        let mut got = vec![0.0f32; rows];
9164        q1_matvec(&bytes, &x, rows, cols, &mut got, None);
9165        for o in 0..rows {
9166            assert!(
9167                (got[o] - expect[o]).abs() < 1e-3 * expect[o].abs().max(1e-3),
9168                "row {o}: {} vs {}",
9169                got[o],
9170                expect[o]
9171            );
9172        }
9173        // Pair and batch paths agree with the single path.
9174        let x2: Vec<f32> = x.iter().map(|v| -v).collect();
9175        let (mut a1, mut a2) = (vec![0.0f32; rows], vec![0.0f32; rows]);
9176        q1_matvec2(&bytes, &x, &x2, rows, cols, &mut a1, &mut a2, None);
9177        assert_eq!(a1, got);
9178        let mut xs = x.clone();
9179        xs.extend_from_slice(&x2);
9180        let mut mm = vec![0.0f32; 2 * rows];
9181        q1_matmat(&bytes, &xs, 2, rows, cols, &mut mm, None);
9182        assert_eq!(&mm[..rows], got.as_slice());
9183        assert_eq!(&mm[rows..], a2.as_slice());
9184    }
9185
9186    #[test]
9187    fn repack_is_bit_identical() {
9188        // The interleaved-repack kernel must produce EXACTLY the same
9189        // bits as the mmap-layout kernel: integer accumulation is order-
9190        // exact, the f32 epilogue is identical. Odd rows exercise the
9191        // tail; direct range calls exercise unaligned pool splits.
9192        let (rows, cols) = (267, 96); // 66 groups + 3 tail rows, cols % 16 == 0
9193        let w: Vec<u8> = (0..rows * cols)
9194            .map(|i| (((i * 89) % 253) as i32 - 126) as i8 as u8)
9195            .collect();
9196        let scales: Vec<f32> = (0..rows).map(|o| 0.003 + o as f32 * 0.0007).collect();
9197        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.37).sin() * 2.0).collect();
9198        let rep = q8_repack_layout(&w, rows, cols);
9199        // Group interleave round-trips.
9200        for g in 0..rows / 4 {
9201            for c in 0..cols / 16 {
9202                for lane in 0..4 {
9203                    assert_eq!(
9204                        &rep[g * 4 * cols + c * 64 + lane * 16
9205                            ..g * 4 * cols + c * 64 + lane * 16 + 16],
9206                        &w[(g * 4 + lane) * cols + c * 16..(g * 4 + lane) * cols + c * 16 + 16],
9207                    );
9208                }
9209            }
9210        }
9211        let mut a = vec![0.0f32; rows];
9212        qmatvec(
9213            &w,
9214            &[],
9215            &scales,
9216            &x,
9217            &[],
9218            TensorDtype::Q8Row,
9219            rows,
9220            cols,
9221            &mut a,
9222            None,
9223        );
9224        let mut b = vec![0.0f32; rows];
9225        qmatvec(
9226            &w,
9227            &rep,
9228            &scales,
9229            &x,
9230            &[],
9231            TensorDtype::Q8Row,
9232            rows,
9233            cols,
9234            &mut b,
9235            None,
9236        );
9237        assert_eq!(a, b, "full-range repack output diverged");
9238
9239        #[cfg(target_arch = "aarch64")]
9240        if sdot_enabled() {
9241            // Unaligned range split (pool workers get arbitrary bounds).
9242            let act = split_act(&x);
9243            let mut c1 = vec![0.0f32; rows];
9244            let mut c2 = vec![0.0f32; rows];
9245            q8_range_sdot(
9246                &w,
9247                &[],
9248                &scales,
9249                &act,
9250                cols,
9251                SendMut(c1.as_mut_ptr()),
9252                3,
9253                rows - 2,
9254            );
9255            q8_range_sdot(
9256                &w,
9257                &rep,
9258                &scales,
9259                &act,
9260                cols,
9261                SendMut(c2.as_mut_ptr()),
9262                3,
9263                rows - 2,
9264            );
9265            assert_eq!(c1, c2, "unaligned-range repack output diverged");
9266        }
9267    }
9268
9269    #[test]
9270    fn sdot_a8w8_noise_is_bounded() {
9271        // Off-grid activations: A8 quantization noise must stay small in
9272        // relative L2 over the whole output (realistic accuracy contract;
9273        // vmfcore measured argmax-identical decode on real models).
9274        let (rows, cols) = (16, 512);
9275        let w: Vec<u8> = (0..rows * cols)
9276            .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
9277            .collect();
9278        let scales = vec![0.01f32; rows];
9279        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.21).sin()).collect();
9280        let mut a = vec![0.0f32; rows];
9281        qmatvec(
9282            &w,
9283            &[],
9284            &scales,
9285            &x,
9286            &[],
9287            TensorDtype::Q8Row,
9288            rows,
9289            cols,
9290            &mut a,
9291            None,
9292        );
9293        let (mut num, mut den) = (0f64, 0f64);
9294        for o in 0..rows {
9295            let mut acc = 0.0f32;
9296            for j in 0..cols {
9297                acc += (w[o * cols + j] as i8) as f32 * x[j];
9298            }
9299            let expect = acc * scales[o];
9300            num += ((a[o] - expect) as f64).powi(2);
9301            den += (expect as f64).powi(2);
9302        }
9303        let rel = (num / den.max(1e-12)).sqrt();
9304        assert!(rel < 0.05, "A8W8 relative L2 error too high: {rel}");
9305    }
9306
9307    #[test]
9308    fn i8_dot_neon_matches_scalar() {
9309        let n = 100;
9310        let w: Vec<u8> = (0..n).map(|i| ((i * 37 + 11) % 251) as u8).collect();
9311        let x: Vec<f32> = (0..n).map(|i| (i as f32 * 0.13).sin()).collect();
9312        let mut scalar = 0.0f32;
9313        for j in 0..n {
9314            scalar += (w[j] as i8) as f32 * x[j];
9315        }
9316        let fast = dot_i8_f32(&w, &x);
9317        assert!((scalar - fast).abs() < 1e-3 * scalar.abs().max(1.0));
9318    }
9319
9320    /// Fused vbit matvec must match full dequant_vbit + dense matvec.
9321    #[test]
9322    fn vbitmatvec_matches_full_dequant() {
9323        let (rows, cols) = (6, 64);
9324        let ng = cols / GROUP_SIZE;
9325        // Hand-craft: bits per row, f16 scales, packed rows.
9326        let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
9327        let mut bytes = bits.clone();
9328        for g in 0..rows * ng {
9329            let s = 0.02 + 0.001 * g as f32;
9330            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
9331        }
9332        for r in 0..rows {
9333            let b = bits[r] as usize;
9334            let (mut acc, mut nb) = (0u64, 0usize);
9335            let mut rowbytes = Vec::new();
9336            for i in 0..cols {
9337                let v = ((i * 7 + r * 13) % (1 << b)) as u64;
9338                acc = (acc << b) | v;
9339                nb += b;
9340                while nb >= 8 {
9341                    nb -= 8;
9342                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
9343                }
9344            }
9345            if nb > 0 {
9346                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
9347            }
9348            bytes.extend_from_slice(&rowbytes);
9349        }
9350        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
9351
9352        let mut reference = vec![0f32; rows * cols];
9353        cortiq_core::quant::dequant_vbit(&bytes, rows, cols, &mut reference).unwrap();
9354        let mut expect = vec![0f32; rows];
9355        for r in 0..rows {
9356            expect[r] = reference[r * cols..(r + 1) * cols]
9357                .iter()
9358                .zip(&x)
9359                .map(|(w, xv)| w * xv)
9360                .sum();
9361        }
9362        let mut got = vec![0f32; rows];
9363        let offsets = vbit_row_offsets(&bytes, rows, cols);
9364        vbitmatvec(&bytes, &offsets, &x, rows, cols, &mut got, None);
9365        // SDOT path quantizes activations to i8 (A8W8): bounded noise,
9366        // same contract as q8 (exact path is pinned by CMF_SDOT=0 in
9367        // the golden-parity gate).
9368        let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
9369        let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
9370        for r in 0..rows {
9371            assert!(
9372                (got[r] - expect[r]).abs() < tol * scale,
9373                "row {r}: {} vs {}",
9374                got[r],
9375                expect[r]
9376            );
9377        }
9378    }
9379
9380    /// Fused q4 matvec must match the reference full-dequant + dense
9381    /// matvec bit-for-bit in structure (same f32 math, group order).
9382    /// vbit matmat: the blocked 1×4 leg must match the per-row path
9383    /// (paired env toggle; larger shape so both code paths engage).
9384    #[test]
9385    #[cfg(target_arch = "x86_64")]
9386    fn vbit_matmat_blocked_matches_per_row() {
9387        let (rows, cols, b) = (64usize, 128usize, 9usize);
9388        let ng = cols / GROUP_SIZE;
9389        let bits: Vec<u8> = (0..rows).map(|r| [3u8, 4, 5, 6][r % 4]).collect();
9390        let mut bytes = bits.clone();
9391        for g in 0..rows * ng {
9392            let sc = 0.02 + 0.0005 * g as f32;
9393            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
9394        }
9395        for r in 0..rows {
9396            let bw = bits[r] as usize;
9397            let (mut acc, mut nb) = (0u64, 0usize);
9398            let mut rowbytes = Vec::new();
9399            for i in 0..cols {
9400                let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
9401                acc = (acc << bw) | v;
9402                nb += bw;
9403                while nb >= 8 {
9404                    nb -= 8;
9405                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
9406                }
9407            }
9408            if nb > 0 {
9409                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
9410            }
9411            bytes.extend_from_slice(&rowbytes);
9412        }
9413        let x: Vec<f32> = (0..b * cols)
9414            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
9415            .collect();
9416        let offsets = vbit_row_offsets(&bytes, rows, cols);
9417        let mut y_a = vec![0f32; b * rows];
9418        let mut y_b = vec![0f32; b * rows];
9419        unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
9420        vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_a, None);
9421        unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
9422        vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_b, None);
9423        unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
9424        let max_d = y_a
9425            .iter()
9426            .zip(&y_b)
9427            .map(|(p, q)| (p - q).abs())
9428            .fold(0.0f32, f32::max);
9429        assert!(max_d < 1e-4, "vbit blocked ≠ per-row: max|Δ| = {max_d}");
9430    }
9431
9432    /// q4t blocked 1×4 (SDOT on ARM, AVX2 on x86) must equal the
9433    /// per-row path exactly: same nibble unpack, same group order,
9434    /// same f32 accumulation — batch == matvec bit-for-bit. b=9 covers
9435    /// two full 1×4 blocks plus a remainder through the single-row
9436    /// kernel. (Both paths produce identical output, so the shared
9437    /// CMF_X86_BLOCKED env var racing with other tests cannot flip
9438    /// the verdict — worst case both sides take the same path.)
9439    #[test]
9440    fn q4t_matmat_blocked_matches_per_row() {
9441        let (rows, cols, b) = (16usize, 64usize, 9usize);
9442        let gpr = cols / GROUP_SIZE;
9443        let mut bytes = vec![0u8; rows * gpr * Q4_TILE];
9444        for r in 0..rows {
9445            for g in 0..gpr {
9446                let t = (r * gpr + g) * Q4_TILE;
9447                let sc = 0.02 + 0.001 * (r * gpr + g) as f32;
9448                bytes[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
9449                for k in 0..16 {
9450                    bytes[t + 2 + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
9451                }
9452            }
9453        }
9454        let x: Vec<f32> = (0..b * cols)
9455            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
9456            .collect();
9457        let mut y_blk = vec![0f32; b * rows];
9458        let mut y_row = vec![0f32; b * rows];
9459        unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
9460        q4t_matmat(&bytes, &x, b, rows, cols, &mut y_blk, None);
9461        unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
9462        q4t_matmat(&bytes, &x, b, rows, cols, &mut y_row, None);
9463        unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
9464        assert_eq!(y_blk, y_row, "q4t blocked 1x4 ≠ per-row");
9465    }
9466
9467    /// The wide-batch Accelerate arm of q4t_matmat vs a brute-force
9468    /// f32 dequant matmul: both are f32 GEMMs, so only reduction
9469    /// order differs — tight tolerance.
9470    /// A synthetic q4tp payload: random nibbles plus a per-row ladder whose
9471    /// span varies row to row, so the codes actually exercise the full 0..31
9472    /// range rather than clustering on one rung.
9473    fn synth_q4tp(rows: usize, cols: usize) -> Vec<u8> {
9474        use cortiq_core::quant::{f32_to_f16, q4tp_code_stride, q4tp_put_code};
9475        let gpr = cols / GROUP_SIZE;
9476        let stride = q4tp_code_stride(gpr);
9477        let (params_off, codes_off, _) = q4tp_sections(rows, cols);
9478        let mut b = vec![0u8; codes_off + rows * stride];
9479        for r in 0..rows {
9480            for g in 0..gpr {
9481                let t = (r * gpr + g) * Q4TP_NIB;
9482                for k in 0..16 {
9483                    b[t + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
9484                }
9485            }
9486            let lo = -6.0 - 0.03 * (r % 17) as f32;
9487            let step = 0.01 + 0.004 * (r % 11) as f32;
9488            let p = params_off + r * 4;
9489            b[p..p + 2].copy_from_slice(&f32_to_f16(lo).to_le_bytes());
9490            b[p + 2..p + 4].copy_from_slice(&f32_to_f16(step).to_le_bytes());
9491            let crow = &mut b[codes_off + r * stride..codes_off + (r + 1) * stride];
9492            for g in 0..gpr {
9493                q4tp_put_code(crow, g, (r * 5 + g * 3) % 32);
9494            }
9495        }
9496        b
9497    }
9498
9499    /// The same weights re-expressed as q4_tiled, so the proven kernel can
9500    /// be the reference: each tile stores the ladder scale its code selects.
9501    /// Only the f16 rounding of that scale separates the two payloads.
9502    fn q4tp_as_q4t(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
9503        let gpr = cols / GROUP_SIZE;
9504        let v = Q4tpView::new(bytes, rows, cols);
9505        let mut out = vec![0u8; rows * gpr * Q4_TILE];
9506        let mut sc = vec![0f32; gpr];
9507        for r in 0..rows {
9508            v.scales_into(r, gpr, &mut sc);
9509            for g in 0..gpr {
9510                let t = (r * gpr + g) * Q4_TILE;
9511                let s = sc[g];
9512                out[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
9513                let src = (r * gpr + g) * Q4TP_NIB;
9514                out[t + 2..t + Q4_TILE].copy_from_slice(&v.nib[src..src + Q4TP_NIB]);
9515            }
9516        }
9517        out
9518    }
9519
9520    /// The exact (`CMF_SDOT=0`) path must reproduce `dequant_q4tp` to f32
9521    /// rounding — that scalar routine is the format's definition, and the
9522    /// kernels re-derive the scale from the ladder independently. Call the
9523    /// row kernel directly: `matmat` picks the int8 arm when a8w8 is on,
9524    /// so routing through it would test the other path by accident.
9525    #[test]
9526    fn q4tp_exact_path_matches_dequant_reference() {
9527        let (rows, cols) = (256usize, 512usize);
9528        let gpr = cols / GROUP_SIZE;
9529        let bytes = synth_q4tp(rows, cols);
9530        let mut w = vec![0f32; rows * cols];
9531        cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
9532
9533        let x: Vec<f32> = (0..cols)
9534            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
9535            .collect();
9536        let v = Q4tpView::new(&bytes, rows, cols);
9537        let mut sc = vec![0f32; gpr];
9538        for r in 0..rows {
9539            v.scales_into(r, gpr, &mut sc);
9540            let got = q4tp_row_exact(v.nib, r, gpr, &x, &sc);
9541            let want: f32 = (0..cols).map(|c| w[r * cols + c] * x[c]).sum();
9542            // These dot products cancel down to ~1e-3 from terms of ~5e-2, so
9543            // the meaningful yardstick is the summed magnitude, not the result:
9544            // against the result any reordering of a 512-term f32 sum "fails".
9545            let mag: f32 = (0..cols).map(|c| (w[r * cols + c] * x[c]).abs()).sum();
9546            assert!(
9547                (got - want).abs() <= 1e-5 * mag,
9548                "row {r}: kernel {got} vs dequant {want}"
9549            );
9550        }
9551    }
9552
9553    /// The int8 (a8w8) path can't be checked against an f32 reference — the
9554    /// activation quantization dominates. Check it against the q4t kernel it
9555    /// was ported from instead, on payloads holding the same weights: that
9556    /// isolates exactly what the port could break (16 B stride, ladder
9557    /// lookup, nibble unpack) from what it deliberately shares.
9558    #[test]
9559    fn q4tp_matvec_matches_the_q4t_kernel_it_was_ported_from() {
9560        let (rows, cols) = (256usize, 512usize);
9561        let bytes = synth_q4tp(rows, cols);
9562        let twin = q4tp_as_q4t(&bytes, rows, cols);
9563        let x: Vec<f32> = (0..cols)
9564            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
9565            .collect();
9566
9567        let mut got = vec![0f32; rows];
9568        q4tp_matvec(&bytes, &x, rows, cols, &mut got, None);
9569        let mut want = vec![0f32; rows];
9570        q4t_matvec(&twin, &x, rows, cols, &mut want, None);
9571
9572        // Scale is f16 in the twin and f32 here, so allow that rounding on
9573        // top of the summed magnitude (same cancellation argument as above).
9574        let mut w = vec![0f32; rows * cols];
9575        cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
9576        for r in 0..rows {
9577            let mag: f32 = (0..cols).map(|c| (w[r * cols + c] * x[c]).abs()).sum();
9578            assert!(
9579                (got[r] - want[r]).abs() <= 1e-3 * mag,
9580                "row {r}: q4tp {} vs q4t {}",
9581                got[r],
9582                want[r]
9583            );
9584        }
9585    }
9586
9587    /// `matmat` carries three arms (Accelerate, blocked int8 1x4, scalar).
9588    /// Batch 5 crosses the blocked kernel's stride, so this exercises the
9589    /// 1x4 path AND its scalar tail in one run — the blocked kernel is new
9590    /// code and its four accumulators are exactly what tends to go wrong.
9591    #[test]
9592    fn q4tp_matmat_matches_the_q4t_kernel_it_was_ported_from() {
9593        let (rows, cols, b) = (256usize, 512usize, 5usize);
9594        let bytes = synth_q4tp(rows, cols);
9595        let twin = q4tp_as_q4t(&bytes, rows, cols);
9596        let xs: Vec<f32> = (0..b * cols)
9597            .map(|i| ((i * 29 + 11) % 89) as f32 / 89.0 - 0.5)
9598            .collect();
9599
9600        let mut got = vec![0f32; b * rows];
9601        q4tp_matmat(&bytes, &xs, b, rows, cols, &mut got, None);
9602        let mut want = vec![0f32; b * rows];
9603        q4t_matmat(&twin, &xs, b, rows, cols, &mut want, None);
9604
9605        let mut w = vec![0f32; rows * cols];
9606        cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
9607        for t in 0..b {
9608            for r in 0..rows {
9609                let mag: f32 = (0..cols)
9610                    .map(|c| (w[r * cols + c] * xs[t * cols + c]).abs())
9611                    .sum();
9612                let (g, wa) = (got[t * rows + r], want[t * rows + r]);
9613                assert!(
9614                    (g - wa).abs() <= 1e-3 * mag,
9615                    "batch {t} row {r}: q4tp {g} vs q4t {wa}"
9616                );
9617            }
9618        }
9619    }
9620
9621    #[test]
9622    fn q4tp_matvec2_matches_the_single_stream_kernel() {
9623        let (rows, cols) = (128usize, 256usize);
9624        let gpr = cols / GROUP_SIZE;
9625        let bytes = synth_q4tp(rows, cols);
9626        let xs: Vec<f32> = (0..2 * cols)
9627            .map(|i| ((i * 29 + 11) % 89) as f32 / 89.0 - 0.5)
9628            .collect();
9629
9630        let (mut o1, mut o2) = (vec![0f32; rows], vec![0f32; rows]);
9631        q4tp_matvec2(
9632            &bytes,
9633            &xs[..cols],
9634            &xs[cols..],
9635            rows,
9636            cols,
9637            &mut o1,
9638            &mut o2,
9639            None,
9640        );
9641
9642        // matvec2 takes the exact path for both streams, so the single-row
9643        // kernel is an exact reference — no tolerance for path differences.
9644        let v = Q4tpView::new(&bytes, rows, cols);
9645        let mut sc = vec![0f32; gpr];
9646        for r in 0..rows {
9647            v.scales_into(r, gpr, &mut sc);
9648            assert_eq!(o1[r], q4tp_row_exact(v.nib, r, gpr, &xs[..cols], &sc));
9649            assert_eq!(o2[r], q4tp_row_exact(v.nib, r, gpr, &xs[cols..], &sc));
9650        }
9651    }
9652
9653    /// q4tp must not COST speed — it exists to save bytes, and a format that
9654    /// trades 7% of a file for a slower model is a bad trade. This guard is
9655    /// here because correctness tests happily passed while `q4tp_matmat` was
9656    /// missing its int8 and Accelerate arms and the model ran 5x slower.
9657    /// Measured on M-series: 0.97-1.04x, i.e. parity (16 B tiles are better
9658    /// aligned than q4t's 18 B, which pays for the scale indirection).
9659    #[test]
9660    fn q4tp_matvec_keeps_pace_with_q4t() {
9661        let (rows, cols) = (4096usize, 3072usize);
9662        let bytes = synth_q4tp(rows, cols);
9663        let twin = q4tp_as_q4t(&bytes, rows, cols);
9664        let x: Vec<f32> = (0..cols).map(|i| (i % 97) as f32 / 97.0 - 0.5).collect();
9665        let mut o = vec![0f32; rows];
9666        let n = 12;
9667        let mut best = (f64::MAX, f64::MAX);
9668        // Interleaved A/B, minimum statistic: this machine throttles, and a
9669        // mean over a thermal ramp reliably indicts whichever ran second.
9670        for _ in 0..3 {
9671            let t0 = std::time::Instant::now();
9672            for _ in 0..n {
9673                q4t_matvec(&twin, &x, rows, cols, &mut o, None);
9674            }
9675            best.0 = best.0.min(t0.elapsed().as_secs_f64());
9676            let t0 = std::time::Instant::now();
9677            for _ in 0..n {
9678                q4tp_matvec(&bytes, &x, rows, cols, &mut o, None);
9679            }
9680            best.1 = best.1.min(t0.elapsed().as_secs_f64());
9681        }
9682        let ratio = best.1 / best.0;
9683        println!(
9684            "q4t {:.3} ms | q4tp {:.3} ms | {ratio:.2}x",
9685            best.0 * 1e3 / n as f64,
9686            best.1 * 1e3 / n as f64
9687        );
9688        assert!(ratio < 2.0, "q4tp matvec {ratio:.2}x slower than q4t");
9689    }
9690
9691    #[cfg(target_os = "macos")]
9692    #[test]
9693    fn q4t_matmat_accel_matches_dequant_reference() {
9694        if !accel_gemm_enabled() {
9695            return; // CMF_ACCEL=0
9696        }
9697        let (rows, cols, b) = (512usize, 1024usize, 8usize); // ≥500K → accel arm
9698        let gpr = cols / GROUP_SIZE;
9699        let mut bytes = vec![0u8; rows * gpr * Q4_TILE];
9700        for r in 0..rows {
9701            for g in 0..gpr {
9702                let t = (r * gpr + g) * Q4_TILE;
9703                let sc = 0.02 + 0.0005 * ((r * gpr + g) % 64) as f32;
9704                bytes[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
9705                for k in 0..16 {
9706                    bytes[t + 2 + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
9707                }
9708            }
9709        }
9710        let x: Vec<f32> = (0..b * cols)
9711            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
9712            .collect();
9713        let mut got = vec![0f32; b * rows];
9714        q4t_matmat(&bytes, &x, b, rows, cols, &mut got, None);
9715        // Brute-force reference off the same tiles.
9716        let mut w = vec![0f32; rows * cols];
9717        for r in 0..rows {
9718            for g in 0..gpr {
9719                let t = (r * gpr + g) * Q4_TILE;
9720                let s = f16_to_f32(u16::from_le_bytes([bytes[t], bytes[t + 1]]));
9721                for (k, &bb) in bytes[t + 2..t + Q4_TILE].iter().enumerate() {
9722                    w[r * cols + g * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s;
9723                    w[r * cols + g * GROUP_SIZE + k * 2 + 1] =
9724                        (((bb >> 4) & 0x0F) as f32 - 8.0) * s;
9725                }
9726            }
9727        }
9728        for bi in 0..b {
9729            for r in 0..rows {
9730                let want: f32 = (0..cols).map(|j| x[bi * cols + j] * w[r * cols + j]).sum();
9731                let d = (got[bi * rows + r] - want).abs();
9732                assert!(
9733                    d <= want.abs().max(1.0) * 1e-4,
9734                    "accel q4t GEMM diverged at ({bi},{r}): {} vs {want}",
9735                    got[bi * rows + r]
9736                );
9737            }
9738        }
9739    }
9740
9741    #[test]
9742    fn q4matvec_matches_full_dequant() {
9743        let (rows, cols) = (8, 64);
9744        let groups = rows * cols / GROUP_SIZE;
9745        // Hand-craft a q4_block blob: nibbles then f16 scales.
9746        let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
9747        for i in 0..groups * 16 {
9748            bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
9749        }
9750        for g in 0..groups {
9751            let s = 0.01 + 0.003 * g as f32;
9752            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
9753        }
9754        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
9755
9756        let mut reference = vec![0.0f32; rows * cols];
9757        cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
9758        let mut expect = vec![0.0f32; rows];
9759        for r in 0..rows {
9760            expect[r] = reference[r * cols..(r + 1) * cols]
9761                .iter()
9762                .zip(&x)
9763                .map(|(w, xv)| w * xv)
9764                .sum();
9765        }
9766
9767        let mut got = vec![0.0f32; rows];
9768        q4matvec(&bytes, &x, rows, cols, &mut got, None);
9769        // SDOT path quantizes activations to i8 (A8W8): bounded noise,
9770        // same contract as q8/vbit (exact path is pinned by CMF_SDOT=0
9771        // in the golden-parity gate).
9772        let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
9773        let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
9774        for r in 0..rows {
9775            assert!(
9776                (got[r] - expect[r]).abs() < tol * scale,
9777                "row {r}: {} vs {}",
9778                got[r],
9779                expect[r]
9780            );
9781        }
9782    }
9783
9784    /// Fused two-input vbit matvec must equal two single matvecs exactly
9785    /// (same per-lane accumulation order on both scalar and SDOT paths).
9786    #[test]
9787    fn vbitmatvec2_equals_two_singles() {
9788        let (rows, cols) = (6, 64);
9789        let ng = cols / GROUP_SIZE;
9790        let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
9791        let mut bytes = bits.clone();
9792        for g in 0..rows * ng {
9793            let s = 0.02 + 0.001 * g as f32;
9794            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
9795        }
9796        for r in 0..rows {
9797            let b = bits[r] as usize;
9798            let (mut acc, mut nb) = (0u64, 0usize);
9799            let mut rowbytes = Vec::new();
9800            for i in 0..cols {
9801                let v = ((i * 7 + r * 13) % (1 << b)) as u64;
9802                acc = (acc << b) | v;
9803                nb += b;
9804                while nb >= 8 {
9805                    nb -= 8;
9806                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
9807                }
9808            }
9809            if nb > 0 {
9810                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
9811            }
9812            bytes.extend_from_slice(&rowbytes);
9813        }
9814        let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
9815        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).cos()).collect();
9816        let offsets = vbit_row_offsets(&bytes, rows, cols);
9817
9818        let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
9819        vbitmatvec(&bytes, &offsets, &x1, rows, cols, &mut a1, None);
9820        vbitmatvec(&bytes, &offsets, &x2, rows, cols, &mut a2, None);
9821        let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
9822        vbitmatvec2(
9823            &bytes, &offsets, &x1, &x2, rows, cols, &mut b1, &mut b2, None,
9824        );
9825        assert_eq!(a1, b1, "fused vbit lane 1 must be bit-identical");
9826        assert_eq!(a2, b2, "fused vbit lane 2 must be bit-identical");
9827    }
9828
9829    /// Fused two-input q4 matvec must equal two single matvecs exactly.
9830    #[test]
9831    fn q4matvec2_equals_two_singles() {
9832        let (rows, cols) = (8, 128);
9833        let groups = rows * cols / GROUP_SIZE;
9834        let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
9835        for i in 0..groups * 16 {
9836            bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
9837        }
9838        for g in 0..groups {
9839            let s = 0.01 + 0.003 * g as f32;
9840            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
9841        }
9842        // Include an outlier channel so the SDOT correction path is
9843        // exercised in the pair kernel too.
9844        let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
9845        x1[9] = 250.0;
9846        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
9847
9848        let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
9849        q4matvec(&bytes, &x1, rows, cols, &mut a1, None);
9850        q4matvec(&bytes, &x2, rows, cols, &mut a2, None);
9851        let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
9852        q4matvec2(&bytes, &x1, &x2, rows, cols, &mut b1, &mut b2, None);
9853        assert_eq!(a1, b1, "fused q4 lane 1 must be bit-identical");
9854        assert_eq!(a2, b2, "fused q4 lane 2 must be bit-identical");
9855    }
9856
9857    /// Multi-matrix job must equal separate matvecs exactly — same
9858    /// kernels, only the dispatch is fused.
9859    #[test]
9860    fn matvec_many_equals_separate_matvecs() {
9861        use crate::pool::Pool;
9862        let (r1, r2, cols) = (300, 200, 64);
9863        let mk = |salt: usize, rows: usize| {
9864            QTensor::from_f32(
9865                (0..rows * cols)
9866                    .map(|i| ((i * 7 + salt) % 97) as f32 / 97.0 - 0.5)
9867                    .collect(),
9868                rows,
9869                cols,
9870            )
9871        };
9872        let (a, b) = (mk(1, r1), mk(5, r2));
9873        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).sin()).collect();
9874        let pool = Pool::new(3);
9875
9876        let (mut ea, mut eb) = (vec![0f32; r1], vec![0f32; r2]);
9877        a.matvec(&x, &mut ea, Some(&pool));
9878        b.matvec(&x, &mut eb, Some(&pool));
9879        let (mut ga, mut gb) = (vec![0f32; r1], vec![0f32; r2]);
9880        QTensor::matvec_many([&a, &b], &x, [&mut ga, &mut gb], Some(&pool));
9881        assert_eq!(ea, ga, "fused multi-matrix lane 1 must be bit-identical");
9882        assert_eq!(eb, gb, "fused multi-matrix lane 2 must be bit-identical");
9883    }
9884
9885    /// Batched q4/vbit matmat must equal per-position matvec calls
9886    /// exactly (the fallback it replaced) — same kernels, same order.
9887    #[test]
9888    fn batched_matmat_equals_per_position_matvec() {
9889        let (rows, cols, b) = (8, 64, 5);
9890        // q4 blob.
9891        let groups = rows * cols / GROUP_SIZE;
9892        let mut q4 = Vec::new();
9893        for i in 0..groups * 16 {
9894            q4.push((((i * 7 + 3) % 256) & 0xFF) as u8);
9895        }
9896        for g in 0..groups {
9897            q4.extend_from_slice(
9898                &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
9899            );
9900        }
9901        // vbit blob (mixed widths incl. 8).
9902        let ng = cols / GROUP_SIZE;
9903        let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4, 5, 3];
9904        let mut vb = bits.clone();
9905        for g in 0..rows * ng {
9906            vb.extend_from_slice(
9907                &cortiq_core::quant::f32_to_f16(0.02 + 0.001 * g as f32).to_le_bytes(),
9908            );
9909        }
9910        for r in 0..rows {
9911            let bw = bits[r] as usize;
9912            let (mut acc, mut nb) = (0u64, 0usize);
9913            let mut rowbytes = Vec::new();
9914            for i in 0..cols {
9915                let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
9916                acc = (acc << bw) | v;
9917                nb += bw;
9918                while nb >= 8 {
9919                    nb -= 8;
9920                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
9921                }
9922            }
9923            if nb > 0 {
9924                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
9925            }
9926            vb.extend_from_slice(&rowbytes);
9927        }
9928        let offsets = vbit_row_offsets(&vb, rows, cols);
9929
9930        let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
9931
9932        // q4: batch vs singles.
9933        let mut got = vec![0f32; b * rows];
9934        q4matmat(&q4, &xs, b, rows, cols, &mut got, None);
9935        for bi in 0..b {
9936            let mut expect = vec![0f32; rows];
9937            q4matvec(
9938                &q4,
9939                &xs[bi * cols..(bi + 1) * cols],
9940                rows,
9941                cols,
9942                &mut expect,
9943                None,
9944            );
9945            assert_eq!(
9946                &got[bi * rows..(bi + 1) * rows],
9947                &expect[..],
9948                "q4 batch pos {bi}"
9949            );
9950        }
9951
9952        // vbit: batch vs singles.
9953        let mut got = vec![0f32; b * rows];
9954        vbitmatmat(&vb, &offsets, &xs, b, rows, cols, &mut got, None);
9955        for bi in 0..b {
9956            let mut expect = vec![0f32; rows];
9957            vbitmatvec(
9958                &vb,
9959                &offsets,
9960                &xs[bi * cols..(bi + 1) * cols],
9961                rows,
9962                cols,
9963                &mut expect,
9964                None,
9965            );
9966            assert_eq!(
9967                &got[bi * rows..(bi + 1) * rows],
9968                &expect[..],
9969                "vbit batch pos {bi}"
9970            );
9971        }
9972    }
9973
9974    /// q4_tiled kernels must produce BIT-identical outputs to the q4
9975    /// split kernels on the same values (same ints, same order — only
9976    /// the byte placement differs).
9977    #[test]
9978    fn q4_tiled_matches_q4_block_bitexact() {
9979        let (rows, cols, b) = (8usize, 128usize, 3usize);
9980        let groups = rows * cols / GROUP_SIZE;
9981        let mut split = Vec::with_capacity(groups * 18);
9982        for i in 0..groups * 16 {
9983            split.push((((i * 7 + 3) % 256) & 0xFF) as u8);
9984        }
9985        for g in 0..groups {
9986            split.extend_from_slice(
9987                &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
9988            );
9989        }
9990        // Re-tile: [scale][nibbles] per group.
9991        let (packed, scales) = split.split_at(groups * 16);
9992        let mut tiled = Vec::with_capacity(groups * Q4_TILE);
9993        for g in 0..groups {
9994            tiled.extend_from_slice(&scales[g * 2..g * 2 + 2]);
9995            tiled.extend_from_slice(&packed[g * 16..(g + 1) * 16]);
9996        }
9997
9998        let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
9999        x1[9] = 250.0; // exercise the outlier path
10000        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
10001
10002        let (mut a, mut t) = (vec![0f32; rows], vec![0f32; rows]);
10003        q4matvec(&split, &x1, rows, cols, &mut a, None);
10004        q4t_matvec(&tiled, &x1, rows, cols, &mut t, None);
10005        assert_eq!(a, t, "q4t matvec must match q4 bit-for-bit");
10006
10007        let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
10008        let (mut t1, mut t2) = (vec![0f32; rows], vec![0f32; rows]);
10009        q4matvec2(&split, &x1, &x2, rows, cols, &mut a1, &mut a2, None);
10010        q4t_matvec2(&tiled, &x1, &x2, rows, cols, &mut t1, &mut t2, None);
10011        assert_eq!(a1, t1);
10012        assert_eq!(a2, t2);
10013
10014        let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
10015        let (mut am, mut tm) = (vec![0f32; b * rows], vec![0f32; b * rows]);
10016        q4matmat(&split, &xs, b, rows, cols, &mut am, None);
10017        q4t_matmat(&tiled, &xs, b, rows, cols, &mut tm, None);
10018        assert_eq!(am, tm, "q4t matmat must match q4 bit-for-bit");
10019    }
10020
10021    /// q4 SDOT outlier correction: a single huge activation channel
10022    /// (>8·rms → outlier, zeroed in xq) must still contribute its EXACT
10023    /// term. On-grid bulk (±1/0 → xq dequantizes exactly) isolates the
10024    /// correction from A8W8 noise. cols must exceed 64: at n=64 the
10025    /// 8·rms threshold equals sqrt(v²+rest) ≥ v, so a single outlier
10026    /// can never qualify (8² = n).
10027    #[test]
10028    fn q4matvec_sdot_outlier_exact() {
10029        let (rows, cols) = (4, 128);
10030        let groups = rows * cols / GROUP_SIZE;
10031        let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
10032        for i in 0..groups * 16 {
10033            bytes.push(((i * 11 + 5) % 256) as u8);
10034        }
10035        for g in 0..groups {
10036            let s = 0.02 + 0.002 * g as f32;
10037            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
10038        }
10039        let mut x: Vec<f32> = (0..cols)
10040            .map(|i| match i % 3 {
10041                0 => 1.0,
10042                1 => -1.0,
10043                _ => 0.0,
10044            })
10045            .collect();
10046        x[17] = 300.0; // ≫ 8·rms → outlier channel
10047
10048        let mut reference = vec![0.0f32; rows * cols];
10049        cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
10050        let mut expect = vec![0.0f32; rows];
10051        for r in 0..rows {
10052            expect[r] = reference[r * cols..(r + 1) * cols]
10053                .iter()
10054                .zip(&x)
10055                .map(|(w, xv)| w * xv)
10056                .sum();
10057        }
10058        let mut got = vec![0.0f32; rows];
10059        q4matvec(&bytes, &x, rows, cols, &mut got, None);
10060        let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
10061        for r in 0..rows {
10062            assert!(
10063                (got[r] - expect[r]).abs() < 2e-3 * scale,
10064                "row {r}: {} vs {} (outlier term must be exact)",
10065                got[r],
10066                expect[r]
10067            );
10068        }
10069    }
10070
10071    /// The fused q1t matvec must equal the reference (dequant_q1t → dot),
10072    /// including the ternary zero level and the binary-searched outlier
10073    /// overlay. Guards the mmap kernel that makes a 12B q1t runnable.
10074    #[test]
10075    fn q1t_matvec_matches_reference() {
10076        use cortiq_core::quant::{dequant_q1t, f32_to_f16};
10077        let (rows, cols) = (3usize, 64usize); // gpr = 2
10078        let gpr = cols / GROUP_SIZE;
10079        let scales = [0.5f32, 0.3, 0.7, 0.2, 0.6, 0.15];
10080        // Overlay (must be sorted by flat index): a few spikes across rows.
10081        let outliers: [(u32, f32); 3] = [(5, 9.0), (70, -4.5), (150, 3.25)];
10082        let is_out = |flat: usize| outliers.iter().any(|&(i, _)| i as usize == flat);
10083        let mut bytes = Vec::new();
10084        for r in 0..rows {
10085            for g in 0..gpr {
10086                bytes.extend_from_slice(&f32_to_f16(scales[r * gpr + g]).to_le_bytes());
10087                let mut c = [0u8; 7];
10088                for k in 0..GROUP_SIZE {
10089                    // Encoder invariant: code 0 at outlier positions.
10090                    let code = if is_out(r * cols + g * GROUP_SIZE + k) {
10091                        0
10092                    } else {
10093                        ((k + r * 3 + g) % 3) as u8 // 0,1,2
10094                    };
10095                    cortiq_core::quant::q1t_pack(&mut c, k, code);
10096                }
10097                bytes.extend_from_slice(&c);
10098            }
10099        }
10100        // Per-row overlay: [u32 row_ptr[rows+1]] then [(u16 col, f16 val)] by
10101        // row (outliers are sorted by flat index → already grouped by row).
10102        let mut row_ptr = vec![0u32; rows + 1];
10103        for &(idx, _) in &outliers {
10104            row_ptr[idx as usize / cols + 1] += 1;
10105        }
10106        for r in 0..rows {
10107            row_ptr[r + 1] += row_ptr[r];
10108        }
10109        for &p in &row_ptr {
10110            bytes.extend_from_slice(&p.to_le_bytes());
10111        }
10112        for &(idx, v) in &outliers {
10113            bytes.extend_from_slice(&((idx as usize % cols) as u16).to_le_bytes());
10114            bytes.extend_from_slice(&f32_to_f16(v).to_le_bytes());
10115        }
10116
10117        let mut refw = vec![0f32; rows * cols];
10118        dequant_q1t(&bytes, rows, cols, &mut refw);
10119        // On-grid activations (±1, amax 1) so the int8 SDOT path reconstructs
10120        // x exactly and matches the f32 reference (same trick as the q1 test).
10121        let x: Vec<f32> = (0..cols)
10122            .map(|j| if j % 3 == 0 { 1.0 } else { -1.0 })
10123            .collect();
10124        let mut expect = vec![0f32; rows];
10125        for r in 0..rows {
10126            let mut a = 0.0f32;
10127            for j in 0..cols {
10128                a += refw[r * cols + j] * x[j];
10129            }
10130            expect[r] = a;
10131        }
10132        let tol = |e: f32| 1e-3 * e.abs().max(1e-3);
10133        let mut got = vec![0f32; rows];
10134        q1t_matvec(&bytes, &x, rows, cols, &mut got, None);
10135        for r in 0..rows {
10136            assert!(
10137                (got[r] - expect[r]).abs() < tol(expect[r]),
10138                "row {r}: {} vs {}",
10139                got[r],
10140                expect[r]
10141            );
10142        }
10143        // matmat (b=2, f32 decode path) must agree too.
10144        let x2: Vec<f32> = x.iter().chain(x.iter().map(|v| v)).copied().collect();
10145        let mut gm = vec![0f32; 2 * rows];
10146        q1t_matmat(&bytes, &x2, 2, rows, cols, &mut gm, None);
10147        for r in 0..rows {
10148            assert!((gm[r] - expect[r]).abs() < tol(expect[r]));
10149            assert!((gm[rows + r] - expect[r]).abs() < tol(expect[r]));
10150        }
10151        // Fused pair (q1t_matvec2) must equal two single matvecs
10152        // bit-for-bit: same unpack, same group order, same f32
10153        // accumulation per stream. Distinct x2 exercises both lanes.
10154        let xb: Vec<f32> = (0..cols)
10155            .map(|j| if j % 5 == 0 { -1.0 } else { 1.0 })
10156            .collect();
10157        let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
10158        q1t_matvec(&bytes, &x, rows, cols, &mut s1, None);
10159        q1t_matvec(&bytes, &xb, rows, cols, &mut s2, None);
10160        let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
10161        q1t_matvec2(&bytes, &x, &xb, rows, cols, &mut p1, &mut p2, None);
10162        assert_eq!(p1, s1, "q1t pair lane 1 ≠ single matvec");
10163        assert_eq!(p2, s2, "q1t pair lane 2 ≠ single matvec");
10164    }
10165
10166    /// Pair == 2×matvec with an ODD group count (the kernel's tail
10167    /// group) and no overlay section.
10168    #[test]
10169    fn q1t_matvec2_odd_gpr_matches_singles() {
10170        use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_pack};
10171        let (rows, cols) = (5usize, 96usize); // gpr = 3 → paired + tail
10172        let gpr = cols / GROUP_SIZE;
10173        let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE);
10174        for r in 0..rows {
10175            for g in 0..gpr {
10176                bytes.extend_from_slice(&f32_to_f16(0.1 + 0.05 * (r + g) as f32).to_le_bytes());
10177                let mut c = [0u8; 7];
10178                for k in 0..GROUP_SIZE {
10179                    q1t_pack(&mut c, k, ((k * 7 + r * 5 + g * 3) % 3) as u8);
10180                }
10181                bytes.extend_from_slice(&c);
10182            }
10183        }
10184        let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.31).sin()).collect();
10185        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).cos()).collect();
10186        let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
10187        q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
10188        q1t_matvec(&bytes, &x2, rows, cols, &mut s2, None);
10189        let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
10190        q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
10191        assert_eq!(p1, s1, "odd-gpr pair lane 1 ≠ single");
10192        assert_eq!(p2, s2, "odd-gpr pair lane 2 ≠ single");
10193    }
10194
10195    // Speed A/B: fused pair (one unpack, two streams) vs two single
10196    // matvecs. Single-threaded, FFN-sized, min-of paired in-process.
10197    //   cargo test -p cortiq-engine --release q1t_matvec2_speed -- --ignored --nocapture
10198    #[test]
10199    #[ignore]
10200    fn q1t_matvec2_speed() {
10201        use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_pack};
10202        use std::time::Instant;
10203        let (rows, cols) = (8192usize, 4096usize);
10204        let gpr = cols / GROUP_SIZE;
10205        let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE);
10206        for r in 0..rows {
10207            for g in 0..gpr {
10208                let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
10209                bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
10210                let mut c = [0u8; 7];
10211                for k in 0..GROUP_SIZE {
10212                    q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
10213                }
10214                bytes.extend_from_slice(&c);
10215            }
10216        }
10217        let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.31).sin()).collect();
10218        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).cos()).collect();
10219        let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
10220        let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
10221        // Warm both paths once.
10222        q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
10223        q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
10224        let (mut t_pair, mut t_two) = (f64::MAX, f64::MAX);
10225        for _ in 0..8 {
10226            let t0 = Instant::now();
10227            q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
10228            t_pair = t_pair.min(t0.elapsed().as_secs_f64() * 1000.0);
10229            let t1 = Instant::now();
10230            q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
10231            q1t_matvec(&bytes, &x2, rows, cols, &mut s2, None);
10232            t_two = t_two.min(t1.elapsed().as_secs_f64() * 1000.0);
10233        }
10234        assert_eq!(p1, s1);
10235        assert_eq!(p2, s2);
10236        println!("q1t pair {rows}x{cols}: fused {t_pair:.2} ms | two singles {t_two:.2} ms");
10237    }
10238
10239    // Speed A/B: the base-3-division decode (what the packing commit left in
10240    // place) vs the fused sign-LUT matvec. Both single-threaded, same bytes.
10241    //   cargo test -p cortiq-engine q1t_matvec_speed -- --ignored --nocapture
10242    #[test]
10243    #[ignore]
10244    fn q1t_matvec_speed() {
10245        use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_code, q1t_pack};
10246        use std::time::Instant;
10247        let (rows, cols) = (8192usize, 4096usize); // FFN-sized
10248        let gpr = cols / GROUP_SIZE;
10249        let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE + 16);
10250        for r in 0..rows {
10251            for g in 0..gpr {
10252                let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
10253                bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
10254                let mut c = [0u8; 7];
10255                for k in 0..GROUP_SIZE {
10256                    q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
10257                }
10258                bytes.extend_from_slice(&c);
10259            }
10260        }
10261        let (n, stride) = (rows * cols, 40usize); // ~2.5% outliers, per-row overlay
10262        let mut row_ptr = vec![0u32; rows + 1];
10263        let mut idx = 0usize;
10264        while idx < n {
10265            row_ptr[idx / cols + 1] += 1;
10266            idx += stride;
10267        }
10268        for r in 0..rows {
10269            row_ptr[r + 1] += row_ptr[r];
10270        }
10271        for &p in &row_ptr {
10272            bytes.extend_from_slice(&p.to_le_bytes());
10273        }
10274        let mut idx = 0usize;
10275        while idx < n {
10276            bytes.extend_from_slice(&((idx % cols) as u16).to_le_bytes());
10277            bytes.extend_from_slice(&f32_to_f16((idx % 13) as f32 * 0.1 - 0.6).to_le_bytes());
10278            idx += stride;
10279        }
10280        // On-grid ±1 so the fast path's int8 SDOT is exact vs the f32 "slow"
10281        // reference (the A/B is a timing check; values must still agree).
10282        let x: Vec<f32> = (0..cols)
10283            .map(|j| if j % 3 == 0 { 1.0 } else { -1.0 })
10284            .collect();
10285        let (rp_off, ent_off, has_ov) = q1t_overlay(&bytes, rows * gpr * Q1T_TILE, rows);
10286
10287        // "before": base-3 division decode into a buffer, then dot.
10288        let slow = |out: &mut [f32]| {
10289            let mut buf = vec![0f32; cols];
10290            for r in 0..rows {
10291                for g in 0..gpr {
10292                    let off = (r * gpr + g) * Q1T_TILE;
10293                    let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
10294                    let codes = &bytes[off + 2..off + Q1T_TILE];
10295                    for k in 0..GROUP_SIZE {
10296                        buf[g * GROUP_SIZE + k] = match q1t_code(codes, k) {
10297                            1 => s,
10298                            2 => -s,
10299                            _ => 0.0,
10300                        };
10301                    }
10302                }
10303                out[r] = q1t_row_outlier_correction(&bytes, r, rp_off, ent_off, has_ov, &x)
10304                    + (0..cols).map(|j| buf[j] * x[j]).sum::<f32>();
10305            }
10306        };
10307        let iters = 5;
10308        let mut a = vec![0f32; rows];
10309        slow(&mut a); // warm
10310        let t = Instant::now();
10311        for _ in 0..iters {
10312            slow(&mut a);
10313        }
10314        let slow_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
10315
10316        let mut b = vec![0f32; rows];
10317        q1t_matvec(&bytes, &x, rows, cols, &mut b, None); // warm
10318        let t = Instant::now();
10319        for _ in 0..iters {
10320            q1t_matvec(&bytes, &x, rows, cols, &mut b, None);
10321        }
10322        let fast_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
10323
10324        for r in 0..rows {
10325            assert!((a[r] - b[r]).abs() < 1e-2, "mismatch row {r}");
10326        }
10327        println!(
10328            "q1t matvec {rows}x{cols} (1 thread): div-decode {slow_ms:.2} ms  fused-LUT {fast_ms:.2} ms  => {:.2}x",
10329            slow_ms / fast_ms
10330        );
10331    }
10332}