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::cell::UnsafeCell;
21use std::sync::Arc;
22
23pub enum QTensor {
24    F32 {
25        data: Vec<f32>,
26        rows: usize,
27        cols: usize,
28    },
29    Mapped {
30        model: Arc<CmfModel>,
31        /// Index into the model's tensor directory.
32        idx: usize,
33        dtype: TensorDtype,
34        rows: usize,
35        cols: usize,
36        /// Per-row scales, dequantized to f32 up front (tiny).
37        row_scale: Vec<f32>,
38        /// q8_2f column field (θ), dequantized up front; empty for q8_row.
39        col_field: Vec<f32>,
40        /// Vbit only: byte offset of each row's packed data within the
41        /// tensor blob (`[rows + 1]`, computed once at load — the per-
42        /// matvec prefix scan over row bit-widths was O(rows) each call).
43        vbit_offsets: Vec<usize>,
44        /// q8-family decode repack (load-time, optional): rows in groups
45        /// of 4, interleaved in 16-byte units — one 64-byte line per
46        /// iteration feeds all 4 sdot lanes, ONE sequential weight
47        /// stream per worker instead of four (this is where llama.cpp's
48        /// repacked Q8 kernels get their bandwidth). Empty = off
49        /// (CMF_REPACK=0, non-SDOT arch, or an ineligible shape). Trades
50        /// an anonymous copy of the quants for mmap pages that go cold.
51        repack: Vec<u8>,
52    },
53    /// Binary low-rank (`blr`): ±1 paths, signs from the mmap, scales as
54    /// f32 (docs/BLR_NVG.md §3). Every device-graph accessor answers `None`
55    /// for it (whole-token graphs decline to per-op/host paths); the CPU
56    /// kernels live in `crate::blr`, the per-op wgpu kernels in
57    /// `gpu_wgpu::blr`.
58    Blr(Box<crate::blr::BlrTensor>),
59}
60
61/// Load-time q8 repack gate (see `Mapped::repack`). OPT-IN
62/// (`CMF_REPACK=1`): the single-stream hypothesis LOST on Apple Silicon
63/// (M4, interleaved A/B: decode 101 vs 94 tok/s — four adjacent row
64/// streams per worker feed the prefetcher MORE memory-level parallelism
65/// than one); kept as an experiment flag for x86, where the tradeoff
66/// may land differently.
67fn repack_enabled() -> bool {
68    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
69    *ON.get_or_init(|| {
70        std::env::var("CMF_REPACK")
71            .map(|v| v == "1")
72            .unwrap_or(cfg!(target_os = "android"))
73    })
74}
75
76/// Interleave q8 rows for the decode kernel: group g holds rows
77/// 4g..4g+4 as [r0[c], r1[c], r2[c], r3[c]] per 16-byte chunk c. Only
78/// full groups are packed — tail rows keep reading the mmap layout.
79fn q8_repack(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
80    #[cfg(target_arch = "aarch64")]
81    let arch_ok = sdot_enabled();
82    #[cfg(not(target_arch = "aarch64"))]
83    let arch_ok = false;
84    if !arch_ok || !repack_enabled() || rows < 256 || cols % 16 != 0 {
85        return Vec::new();
86    }
87    q8_repack_layout(bytes, rows, cols)
88}
89
90/// The pure layout transform behind `q8_repack` (tested directly —
91/// the gate depends on arch and env).
92fn q8_repack_layout(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
93    let groups = rows / 4;
94    let mut rep = vec![0u8; groups * 4 * cols];
95    for g in 0..groups {
96        let dst = &mut rep[g * 4 * cols..(g + 1) * 4 * cols];
97        for c in 0..cols / 16 {
98            for lane in 0..4 {
99                let src = (g * 4 + lane) * cols + c * 16;
100                dst[c * 64 + lane * 16..c * 64 + lane * 16 + 16]
101                    .copy_from_slice(&bytes[src..src + 16]);
102            }
103        }
104    }
105    rep
106}
107
108/// Prefix-sum of vbit row payload offsets (absolute within the tensor
109/// bytes). `offsets[r]..offsets[r+1]` is row r's packed data.
110fn vbit_row_offsets(bytes: &[u8], rows: usize, cols: usize) -> Vec<usize> {
111    let ng = cols / GROUP_SIZE;
112    let bits = &bytes[..rows];
113    let mut offsets = Vec::with_capacity(rows + 1);
114    let mut off = rows + rows * ng * 2;
115    for r in 0..rows {
116        offsets.push(off);
117        off += (cols * bits[r] as usize).div_ceil(8);
118    }
119    offsets.push(off);
120    offsets
121}
122
123/// `CMF_X86_BLOCKED` / `CMF_GPU_LMHEAD` / `CMF_GPU_SPLIT`, read once. They
124/// used to be read from the environment on every large matvec and on every
125/// matmat in six places — microseconds each, but also a knob that could
126/// change under a running process, which is not a thing a kernel choice
127/// should be able to do mid-sequence.
128fn blocked_enabled() -> bool {
129    use std::sync::atomic::Ordering::Relaxed;
130    match BLOCKED_OVERRIDE.load(Relaxed) {
131        1 => false,
132        2 => true,
133        _ => {
134            static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
135            *ON.get_or_init(|| {
136                std::env::var("CMF_X86_BLOCKED")
137                    .map(|v| v != "0")
138                    .unwrap_or(true)
139            })
140        }
141    }
142}
143
144static BLOCKED_OVERRIDE: std::sync::atomic::AtomicU8 = std::sync::atomic::AtomicU8::new(0);
145
146/// Force the blocked GEMM on or off, ignoring the environment; `None`
147/// restores it. For tests that need to run BOTH paths and compare them:
148/// `blocked_enabled` caches its answer for the life of the process, which
149/// is right when the environment is the only input, but leaves a test that
150/// flips `CMF_X86_BLOCKED` between two calls comparing a path against
151/// itself — or against whatever a test running in parallel latched first.
152pub fn set_blocked_override(on: Option<bool>) {
153    let v = match on {
154        None => 0,
155        Some(false) => 1,
156        Some(true) => 2,
157    };
158    BLOCKED_OVERRIDE.store(v, std::sync::atomic::Ordering::Relaxed);
159}
160
161fn gpu_lmhead_enabled() -> bool {
162    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
163    *ON.get_or_init(|| {
164        std::env::var("CMF_GPU_LMHEAD")
165            .map(|v| v != "0")
166            .unwrap_or(true)
167    })
168}
169
170fn gpu_split_frac() -> f32 {
171    // MiMo's banked verification projects every row on the device. Plain
172    // decode must not quantize half of its O/head activation on the CPU.
173    if FULL_GPU_Q8.get() {
174        return 1.0;
175    }
176    static F: std::sync::OnceLock<f32> = std::sync::OnceLock::new();
177    *F.get_or_init(|| {
178        std::env::var("CMF_GPU_SPLIT")
179            .ok()
180            .and_then(|v| v.parse::<f32>().ok())
181            .unwrap_or(0.5)
182            .clamp(0.0, 1.0)
183    })
184}
185
186impl QTensor {
187    pub fn from_f32(data: Vec<f32>, rows: usize, cols: usize) -> Self {
188        debug_assert_eq!(data.len(), rows * cols);
189        Self::F32 { data, rows, cols }
190    }
191
192    /// Wrap a directory tensor without dequantizing the payload.
193    /// Falls back to dequantized f32 for dtypes without a fused kernel.
194    pub fn from_model(model: &Arc<CmfModel>, name: &str) -> Result<Self, String> {
195        // Indexed lookup: the linear directory scan made pipeline build
196        // O(N²) on MoE/skills files with thousands of tensors.
197        let idx = model
198            .tensor_index(name)
199            .ok_or_else(|| format!("tensor '{name}' not found in CMF directory"))?;
200        let entry = &model.tensors[idx];
201        if entry.shape.len() != 2 {
202            return Err(format!("QTensor::from_model needs 2-D, got '{name}'"));
203        }
204        let (rows, cols) = (entry.shape[0], entry.shape[1]);
205        let bytes = model.entry_bytes(entry);
206
207        match entry.dtype {
208            TensorDtype::Q8Row | TensorDtype::Q8_2f => {
209                let n = rows * cols;
210                let scales_off = n;
211                let row_scale: Vec<f32> = (0..rows)
212                    .map(|o| {
213                        f16_to_f32(u16::from_le_bytes([
214                            bytes[scales_off + o * 2],
215                            bytes[scales_off + o * 2 + 1],
216                        ]))
217                    })
218                    .collect();
219                let col_field: Vec<f32> = if entry.dtype == TensorDtype::Q8_2f {
220                    let col_off = n + rows * 2;
221                    (0..cols)
222                        .map(|i| {
223                            f16_to_f32(u16::from_le_bytes([
224                                bytes[col_off + i * 2],
225                                bytes[col_off + i * 2 + 1],
226                            ]))
227                        })
228                        .collect()
229                } else {
230                    Vec::new()
231                };
232                Ok(Self::Mapped {
233                    model: model.clone(),
234                    idx,
235                    dtype: entry.dtype,
236                    rows,
237                    cols,
238                    row_scale,
239                    col_field,
240                    vbit_offsets: Vec::new(),
241                    repack: q8_repack(bytes, rows, cols),
242                })
243            }
244            // vbit: fused kernel unpacks variable-bit rows from mmap.
245            TensorDtype::Vbit if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
246                model: model.clone(),
247                idx,
248                dtype: entry.dtype,
249                rows,
250                cols,
251                row_scale: Vec::new(),
252                col_field: Vec::new(),
253                vbit_offsets: vbit_row_offsets(bytes, rows, cols),
254                repack: Vec::new(),
255            }),
256            // vbit_ro (§4.2): the offset table comes straight from the
257            // file — no load-time prefix scan; kernels are shared with
258            // legacy vbit (they consume absolute offsets either way).
259            TensorDtype::VbitRo if cols % GROUP_SIZE == 0 => {
260                let (_, off_off, packed_off) = cortiq_core::quant::vbit_ro_sections(rows, cols);
261                let offsets: Vec<usize> = (0..=rows)
262                    .map(|r| packed_off + cortiq_core::quant::vbit_ro_offset(bytes, off_off, r))
263                    .collect();
264                Ok(Self::Mapped {
265                    model: model.clone(),
266                    idx,
267                    dtype: entry.dtype,
268                    rows,
269                    cols,
270                    row_scale: Vec::new(),
271                    col_field: Vec::new(),
272                    vbit_offsets: offsets,
273                    repack: Vec::new(),
274                })
275            }
276            // q4_block: fused kernel reads nibbles straight from mmap —
277            // a 14B q4 file no longer explodes into ×8 f32 RAM.
278            // q4_tiled (§4.3): interleaved [scale][nibbles] tiles — one
279            // sequential memory stream (measured ×1.66 ARM / ×1.13 AVX2
280            // at kernel level over the split layout).
281            TensorDtype::Q4Tiled if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
282                model: model.clone(),
283                idx,
284                dtype: entry.dtype,
285                rows,
286                cols,
287                row_scale: Vec::new(),
288                col_field: Vec::new(),
289                vbit_offsets: Vec::new(),
290                repack: Vec::new(),
291            }),
292            // q4tp (§4.10): nibbles from mmap, scale from the row ladder —
293            // 7.3% less file than q4t at the same 4-bit grid.
294            TensorDtype::Q4TiledP if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
295                model: model.clone(),
296                idx,
297                dtype: entry.dtype,
298                rows,
299                cols,
300                row_scale: Vec::new(),
301                col_field: Vec::new(),
302                vbit_offsets: Vec::new(),
303                repack: Vec::new(),
304            }),
305            // q2tp: 2-bit chunks from mmap, scale from the same row ladder.
306            TensorDtype::Q2TiledP if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
307                model: model.clone(),
308                idx,
309                dtype: entry.dtype,
310                rows,
311                cols,
312                row_scale: Vec::new(),
313                col_field: Vec::new(),
314                vbit_offsets: Vec::new(),
315                repack: Vec::new(),
316            }),
317            TensorDtype::Q4Block if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
318                model: model.clone(),
319                idx,
320                dtype: entry.dtype,
321                rows,
322                cols,
323                row_scale: Vec::new(),
324                col_field: Vec::new(),
325                vbit_offsets: Vec::new(),
326                repack: Vec::new(),
327            }),
328            // q1: binary sign-bit tiles from mmap (1-bit-trained models).
329            // q1g: the same signs as planes with one scale per 128-group.
330            TensorDtype::Q1 if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
331                model: model.clone(),
332                idx,
333                dtype: entry.dtype,
334                rows,
335                cols,
336                row_scale: Vec::new(),
337                col_field: Vec::new(),
338                vbit_offsets: Vec::new(),
339                repack: Vec::new(),
340            }),
341            TensorDtype::Q1G if cols % cortiq_core::quant::Q1G_GROUP == 0 => Ok(Self::Mapped {
342                model: model.clone(),
343                idx,
344                dtype: entry.dtype,
345                rows,
346                cols,
347                row_scale: Vec::new(),
348                col_field: Vec::new(),
349                vbit_offsets: Vec::new(),
350                repack: Vec::new(),
351            }),
352            // q1t (ternary + outlier overlay): fused per-row dequant kernel
353            // reads straight from mmap — a 12B q1t stays ~its file size in
354            // RAM instead of dequantizing to ~48 GB of f32.
355            TensorDtype::Q1T if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
356                model: model.clone(),
357                idx,
358                dtype: entry.dtype,
359                rows,
360                cols,
361                row_scale: Vec::new(),
362                col_field: Vec::new(),
363                vbit_offsets: Vec::new(),
364                repack: Vec::new(),
365            }),
366            // blr: factored kernels over the packed signs. CMF_BLR_DENSE=1
367            // dequantizes to f32 instead (A/B parity runs).
368            TensorDtype::BinaryLowRank if std::env::var("CMF_BLR_DENSE").as_deref() != Ok("1") => {
369                Ok(Self::Blr(Box::new(crate::blr::BlrTensor::from_model(
370                    model, idx,
371                )?)))
372            }
373            // No fused kernel yet → dequantize once (correct, more RAM).
374            _ => {
375                let mut data = vec![0.0f32; rows * cols];
376                cortiq_core::quant::dequant_tensor(entry, bytes, &mut data)?;
377                Ok(Self::from_f32(data, rows, cols))
378            }
379        }
380    }
381
382    /// Binary low-rank tensor? Fused multi-tensor kernels refuse it and fall
383    /// back to per-tensor calls.
384    pub fn is_blr(&self) -> bool {
385        matches!(self, Self::Blr(_))
386    }
387
388    /// (model, idx) of a `blr` tensor, for the wgpu graphs only (their
389    /// `blr` stage kernels, kind [`crate::gpu::GRAPH_KIND_BLR`]). Kept apart
390    /// from [`graph_weight`](Self::graph_weight), which stays None for
391    /// `blr`: every other graph consumer keys its kernels on that answer.
392    /// A file with a Prism contract keeps the per-op path.
393    pub fn graph_blr(&self) -> Option<(&Arc<CmfModel>, usize)> {
394        match self {
395            Self::Blr(t) if !crate::prism::has_contract(&t.model) => Some((&t.model, t.idx)),
396            _ => None,
397        }
398    }
399
400    /// [`graph_weight`](Self::graph_weight), then the Prism descriptor
401    /// form, then a `blr` tensor as kind [`crate::gpu::GRAPH_KIND_BLR`] —
402    /// the wgpu token and batch graph builders' lookup.
403    pub(crate) fn graph_weight_or_blr(
404        &self,
405    ) -> Option<(&std::sync::Arc<CmfModel>, usize, u8, &[f32])> {
406        self.graph_weight()
407            .or_else(|| self.graph_weight_descriptor())
408            .or_else(|| {
409                self.graph_blr()
410                    .map(|(m, i)| (m, i, crate::gpu::GRAPH_KIND_BLR, &[][..]))
411            })
412    }
413
414    /// Binary-weight-mapped tensor, q1 or q1g? (GPU gates: the binary CPU
415    /// kernels are compute-bound, so offload pays at much smaller shapes
416    /// than q8.)
417    pub(crate) fn is_q1(&self) -> bool {
418        matches!(
419            self,
420            Self::Mapped {
421                dtype: TensorDtype::Q1 | TensorDtype::Q1G,
422                ..
423            }
424        )
425    }
426
427    /// The binary-weight view (q1 tiles or q1g planes) the shared CPU row
428    /// drivers take; None for every other storage.
429    fn bin_w(&self) -> Option<BinW<'_>> {
430        match self {
431            Self::Mapped {
432                dtype: TensorDtype::Q1,
433                ..
434            } => Some(BinW::Q1(self.quant_bytes())),
435            Self::Mapped {
436                dtype: TensorDtype::Q1G,
437                rows,
438                cols,
439                ..
440            } => Some(BinW::q1g(self.quant_bytes(), *rows, *cols)),
441            _ => None,
442        }
443    }
444
445    /// Owned-f32 view (data, rows, cols) — the GDN a/b gate projections
446    /// arrive dequantized (force-f16 in the converter → F32 in RAM).
447    pub(crate) fn f32_parts(&self) -> Option<(&[f32], usize, usize)> {
448        match self {
449            Self::F32 { data, rows, cols } => Some((data, *rows, *cols)),
450            _ => None,
451        }
452    }
453
454    /// (directory idx, rows, cols) of a q1-mapped tensor — the
455    /// whole-block GPU path resolves offsets itself.
456    /// (idx, rows, cols) of a mapped tensor the whole-token GPU graph can drive
457    /// — Q1, Q1T or Q4-block (it resolves the offset and picks the kernel by
458    /// dtype). Q4-block lets a precise down_proj/lm_head stay on-device.
459    /// Named `q1_parts` for historical reasons.
460    pub(crate) fn q1_parts(&self) -> Option<(usize, usize, usize)> {
461        if self.has_prism_contract() {
462            return None;
463        }
464        match self {
465            #[cfg(target_os = "macos")]
466            Self::Mapped {
467                dtype: TensorDtype::Q1T,
468                ..
469            } if !crate::gpu::metal_q1t_enabled() => None,
470            Self::Mapped {
471                idx,
472                dtype:
473                    TensorDtype::Q1
474                    | TensorDtype::Q1G
475                    | TensorDtype::Q1T
476                    | TensorDtype::Q4Block
477                    | TensorDtype::Q4Tiled
478                    // Q2TiledP deliberately absent: the Metal graph has no
479                    // q2tp kernel, and advertising it here made the block
480                    // plan truncate mid-run at the first q2tp layer.
481                    | TensorDtype::Q4TiledP
482                    | TensorDtype::Q8Row
483                    | TensorDtype::Q8_2f,
484                rows,
485                cols,
486                ..
487            } => Some((*idx, *rows, *cols)),
488            _ => None,
489        }
490    }
491
492    /// `(directory idx, rows, cols)` for the native Metal token graph.  The
493    /// historical q1 graph gate intentionally refuses every Prism tensor so
494    /// an untransformed q2 payload cannot slip into the resident path.  The
495    /// Metal2 graph is descriptor-aware and admits only the production
496    /// q2tp-affine forward targets; ordinary q1/q4 callers retain the old
497    /// `q1_parts` behaviour.
498    #[cfg(target_os = "macos")]
499    pub(crate) fn metal_graph_parts(&self) -> Option<(usize, usize, usize)> {
500        if let Some((model, idx, kind, _)) = self.graph_weight_descriptor() {
501            let name = &model.tensors[idx].name;
502            let forward = kind == 9 && crate::prism::is_forward_weight(model, name);
503            let affine = kind == 9 && crate::prism::is_affine_target(model, name);
504            if forward && affine {
505                let e = model.tensors.get(idx)?;
506                return Some((idx, *e.shape.first()?, *e.shape.get(1)?));
507            }
508        }
509        self.q1_parts()
510    }
511
512    /// (directory idx, rows, cols) of a q4_tiled mapped tensor. The
513    /// chunk-prefill graph takes it in the same 4-tuple slot as
514    /// `q8_row_parts` with an EMPTY row_scale — q4t carries its scales
515    /// inside the 18-byte tiles, and the empty slice is what tells the
516    /// encoder to reach for the q4t kernels.
517    pub(crate) fn q4t_parts(&self) -> Option<(usize, usize, usize)> {
518        if self.has_prism_contract() {
519            return None;
520        }
521        match self {
522            Self::Mapped {
523                idx,
524                dtype: TensorDtype::Q4Tiled,
525                rows,
526                cols,
527                ..
528            } => Some((*idx, *rows, *cols)),
529            _ => None,
530        }
531    }
532
533    /// (directory idx, rows, cols) of a q4tp mapped tensor. Same empty-scale
534    /// slot as `q4t_parts` in the chunk graph — the encoder tells the two
535    /// apart by the tensor's dtype, not by the slot.
536    pub(crate) fn q4tp_parts(&self) -> Option<(usize, usize, usize)> {
537        if self.has_prism_contract() {
538            return None;
539        }
540        match self {
541            Self::Mapped {
542                idx,
543                dtype: TensorDtype::Q4TiledP,
544                rows,
545                cols,
546                ..
547            } => Some((*idx, *rows, *cols)),
548            _ => None,
549        }
550    }
551
552    /// (directory idx, rows, cols, row_scale) of a plain q8_row mapped
553    /// tensor — the chunk-prefill GPU graph resolves offsets itself.
554    /// q8_2f is excluded on purpose: its column field would need a
555    /// prescale stage on the device.
556    pub(crate) fn q8_row_parts(&self) -> Option<(usize, usize, usize, &[f32])> {
557        if self.has_prism_contract() {
558            return None;
559        }
560        match self {
561            Self::Mapped {
562                idx,
563                dtype: TensorDtype::Q8Row,
564                rows,
565                cols,
566                row_scale,
567                col_field,
568                ..
569            } if col_field.is_empty() => Some((*idx, *rows, *cols, row_scale)),
570            _ => None,
571        }
572    }
573
574    /// (directory idx, rows, cols, row_scale, col_field) of a q8_2f mapped
575    /// tensor — the chunk-prefill graph stages x·col on the device first.
576    pub(crate) fn q8_2f_parts(&self) -> Option<(usize, usize, usize, &[f32], &[f32])> {
577        if self.has_prism_contract() {
578            return None;
579        }
580        match self {
581            Self::Mapped {
582                idx,
583                dtype: TensorDtype::Q8_2f,
584                rows,
585                cols,
586                row_scale,
587                col_field,
588                ..
589            } if col_field.len() == *cols => Some((*idx, *rows, *cols, row_scale, col_field)),
590            _ => None,
591        }
592    }
593
594    /// The layout this tensor is stored in, when it is mapped from a model.
595    /// The frames branch on it — a q2tp gate against a q4tp down is a real
596    /// combination in the 2-bit profile and needs a different kernel.
597    pub fn model_dtype(&self) -> Option<cortiq_core::TensorDtype> {
598        match self {
599            Self::Mapped { dtype, .. } => Some(*dtype),
600            _ => None,
601        }
602    }
603
604    /// The tensor's index in the model directory, when it is mapped from one.
605    /// The GPU frames bind by index rather than by name — a name lookup per
606    /// layer per token is not free, and the index is what the device cache is
607    /// keyed on anyway.
608    pub fn model_idx(&self) -> Option<usize> {
609        match self {
610            Self::Mapped { idx, .. } => Some(*idx),
611            _ => None,
612        }
613    }
614
615    /// The model this tensor is mapped from, when it is mapped at all. The
616    /// GPU frames need the container to reach the bytes; a QTensor already
617    /// holds it, and threading a second handle down every call site to say
618    /// the same thing invites the two to disagree.
619    pub fn model_arc(&self) -> Option<std::sync::Arc<cortiq_core::CmfModel>> {
620        match self {
621            Self::Mapped { model, .. } => Some(model.clone()),
622            _ => None,
623        }
624    }
625
626    /// Whether this mapped tensor belongs to the Prism/Bonsai transform
627    /// contract.  Device graphs do not carry the descriptor, so callers use
628    /// this conservative predicate to stay on the descriptor-aware CPU path
629    /// instead of silently executing an unrotated matrix.
630    pub(crate) fn has_prism_contract(&self) -> bool {
631        matches!(self, Self::Mapped { model, .. } if crate::prism::has_contract(model))
632    }
633
634    pub fn rows(&self) -> usize {
635        match self {
636            Self::F32 { rows, .. } | Self::Mapped { rows, .. } => *rows,
637            Self::Blr(t) => t.rows,
638        }
639    }
640
641    /// Mapped q4t handle (model + directory index) — the fused GPU FFN
642    /// needs the raw file coordinates of its three projections.
643    pub(crate) fn mapped_q4t(&self) -> Option<(&Arc<CmfModel>, usize)> {
644        if self.has_prism_contract() {
645            return None;
646        }
647        match self {
648            Self::Mapped {
649                model,
650                idx,
651                dtype: TensorDtype::Q4Tiled,
652                ..
653            } => Some((model, *idx)),
654            _ => None,
655        }
656    }
657
658    /// Same slot as `mapped_q4t` for a q4tp tensor — the fused DiT FFN picks
659    /// its kernels by which of the two answers.
660    pub fn mapped_q4tp(&self) -> Option<(&Arc<CmfModel>, usize)> {
661        if self.has_prism_contract() {
662            return None;
663        }
664        match self {
665            Self::Mapped {
666                model,
667                idx,
668                dtype: TensorDtype::Q4TiledP,
669                ..
670            } => Some((model, *idx)),
671            _ => None,
672        }
673    }
674
675    /// (model, tensor idx) for a mapped weight in ANY codec the fused device
676    /// paths can run — four-bit tiled or either int8 layout.
677    ///
678    /// The fused DiT chains asked for `mapped_q4tp` by name, so an eight-bit
679    /// container never reached them and rendered through per-op GEMMs even
680    /// after those kernels learned its codec. The gate is what the codec has
681    /// a device GEMM for, not which codec it is.
682    pub fn mapped_device_gemm(&self) -> Option<(&Arc<CmfModel>, usize)> {
683        if self.has_prism_contract() {
684            return None;
685        }
686        match self {
687            Self::Mapped {
688                model,
689                idx,
690                dtype: TensorDtype::Q4TiledP | TensorDtype::Q8Row | TensorDtype::Q8_2f,
691                ..
692            } => Some((model, *idx)),
693            _ => None,
694        }
695    }
696
697    /// (model, tensor idx) for a q8_2f mapped weight — the eight-bit slot of
698    /// the MoE graph ladder (int8 body, f16 row scales, f16 column field).
699    pub fn mapped_q8_2f(&self) -> Option<(&Arc<CmfModel>, usize)> {
700        if self.has_prism_contract() {
701            return None;
702        }
703        match self {
704            Self::Mapped {
705                model,
706                idx,
707                dtype: TensorDtype::Q8_2f,
708                ..
709            } => Some((model, *idx)),
710            _ => None,
711        }
712    }
713
714    /// (model, tensor idx) for a q2tp mapped weight — the 2-bit twin of
715    /// `mapped_q4tp`, used by the mixed MoE profile.
716    pub fn mapped_q2tp(&self) -> Option<(&Arc<CmfModel>, usize)> {
717        if self.has_prism_contract() {
718            return None;
719        }
720        match self {
721            Self::Mapped {
722                model,
723                idx,
724                dtype: TensorDtype::Q2TiledP,
725                ..
726            } => Some((model, *idx)),
727            _ => None,
728        }
729    }
730
731    pub fn cols(&self) -> usize {
732        match self {
733            Self::F32 { cols, .. } | Self::Mapped { cols, .. } => *cols,
734            Self::Blr(t) => t.cols,
735        }
736    }
737
738    /// (model, tensor idx) for a q1 mapped weight — the wgpu token graph
739    /// keys its resident VRAM cache by idx. None for any other dtype/kind.
740    pub fn mapped_q1(&self) -> Option<(&std::sync::Arc<CmfModel>, usize)> {
741        if self.has_prism_contract() {
742            return None;
743        }
744        match self {
745            Self::Mapped {
746                model,
747                idx,
748                dtype: TensorDtype::Q1,
749                ..
750            } => Some((model, *idx)),
751            _ => None,
752        }
753    }
754
755    /// (model, idx, kind, row_scale) for a graph-capable mapped weight.
756    /// kind: 0=q8_row (per-row scales), 1=q1, 2=q4_block, 3=q1t
757    /// (tile-embedded, no rs), 5=q4_tiled, 6=q4tp, 7=q8_2f (both scale
758    /// planes live inside the tensor), 9=q2tp, 10=q1g (sign + scale
759    /// planes). None only for `vbit`.
760    ///
761    /// The old comment here claimed q4_block was unhandled while the arm
762    /// right below mapped it, and it named q8_2f as unhandled after that
763    /// stopped being true — a stale comment on this function is how a
764    /// model silently loses the graph, so it is worth keeping honest.
765    pub fn graph_weight(&self) -> Option<(&std::sync::Arc<CmfModel>, usize, u8, &[f32])> {
766        if self.has_prism_contract() {
767            return None;
768        }
769        self.graph_weight_descriptor()
770    }
771
772    /// Descriptor-aware graph handle used only by the Prism token graph.
773    /// Ordinary graph callers continue to use [`graph_weight`] and therefore
774    /// remain fail-closed until they provide the same explicit transform
775    /// contract.
776    pub(crate) fn graph_weight_descriptor(
777        &self,
778    ) -> Option<(&std::sync::Arc<CmfModel>, usize, u8, &[f32])> {
779        match self {
780            Self::Mapped {
781                model,
782                idx,
783                dtype: TensorDtype::Q8Row,
784                row_scale,
785                ..
786            } => Some((model, *idx, 0, row_scale.as_slice())),
787            Self::Mapped {
788                model,
789                idx,
790                dtype: TensorDtype::Q1,
791                ..
792            } => Some((model, *idx, 1, &[])),
793            // Kind 10, never 1: q1g's sign and scale planes are a different
794            // walk from q1's 6-byte tiles (the kind-2/5 mixup taught that).
795            Self::Mapped {
796                model,
797                idx,
798                dtype: TensorDtype::Q1G,
799                ..
800            } => Some((model, *idx, 10, &[])),
801            // Q4Tiled is kind 5, NOT 2: both carried 2 historically, and
802            // the wgpu token graph fed 18B interleaved tiles to the
803            // split-layout q4b kernel — garbage output on q4t models
804            // (caught by an end-to-end answer check on real Vulkan).
805            Self::Mapped {
806                model,
807                idx,
808                dtype: TensorDtype::Q4Tiled,
809                ..
810            } => Some((model, *idx, 5, &[])),
811            // Kind 6, not 5: q4tp's nibble stride and scale planes differ,
812            // and feeding them to the q4t kernel is exactly the mistake that
813            // produced garbage when Q4Tiled shared kind 2 with Q4Block.
814            Self::Mapped {
815                model,
816                idx,
817                dtype: TensorDtype::Q4TiledP,
818                ..
819            } => Some((model, *idx, 6, &[])),
820            Self::Mapped {
821                model,
822                idx,
823                dtype: TensorDtype::Q4Block,
824                ..
825            } => Some((model, *idx, 2, &[])),
826            // q8_2f carries BOTH scale planes after the int8 body (rows
827            // f16, then cols f16), so the graph takes the whole tensor
828            // and the kernel reads them where they lie — no host-side
829            // prescale, which is what the per-op path does instead.
830            Self::Mapped {
831                model,
832                idx,
833                dtype: TensorDtype::Q8_2f,
834                ..
835            } => Some((model, *idx, 7, &[])),
836            Self::Mapped {
837                model,
838                idx,
839                dtype: TensorDtype::Q1T,
840                ..
841            } => Some((model, *idx, 3, &[])),
842            // Kind 9: the 2-bit plane on the q4tp ladder (dense FFN gate/up
843            // of the q2tp profile). Its own kernel — 8 bytes a group where
844            // q4tp has 16, and rung 0 is the exact zero.
845            Self::Mapped {
846                model,
847                idx,
848                dtype: TensorDtype::Q2TiledP,
849                ..
850            } => Some((model, *idx, 9, &[])),
851            _ => None,
852        }
853    }
854
855    /// Dense f32 view — only for owned tensors. Masked/sparse execution
856    /// paths require it; quantized weights don't support masks yet.
857    pub fn as_f32(&self) -> Option<&[f32]> {
858        match self {
859            Self::F32 { data, .. } => Some(data),
860            Self::Mapped { .. } | Self::Blr(_) => None,
861        }
862    }
863
864    fn quant_bytes(&self) -> &[u8] {
865        match self {
866            Self::Mapped { model, idx, .. } => model.entry_bytes(&model.tensors[*idx]),
867            Self::Blr(t) => t.bytes(),
868            Self::F32 { .. } => unreachable!("quant_bytes on F32"),
869        }
870    }
871
872    /// Dequantize one row into `dst` (embedding lookup).
873    pub fn row_f32(&self, r: usize, dst: &mut [f32]) {
874        let cols = self.cols();
875        debug_assert_eq!(dst.len(), cols);
876        match self {
877            Self::F32 { data, .. } => dst.copy_from_slice(&data[r * cols..(r + 1) * cols]),
878            Self::Blr(t) => t.row_f32(r, dst),
879            Self::Mapped {
880                model,
881                idx,
882                dtype,
883                row_scale,
884                col_field,
885                vbit_offsets,
886                ..
887            } => {
888                if *dtype == TensorDtype::Q4Tiled {
889                    let bytes = self.quant_bytes();
890                    let gpr = cols / GROUP_SIZE;
891                    for gi in 0..gpr {
892                        let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
893                        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
894                        for (k, &b) in tile[2..].iter().enumerate() {
895                            dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
896                            dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
897                        }
898                    }
899                    if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
900                        crate::prism::inverse_embedding(model, dst);
901                    }
902                    return;
903                }
904                if *dtype == TensorDtype::Q4TiledP {
905                    let bytes = self.quant_bytes();
906                    let gpr = cols / GROUP_SIZE;
907                    let v = Q4tpView::new(bytes, self.rows(), cols);
908                    let mut sc = vec![0f32; gpr];
909                    v.scales_into(r, gpr, &mut sc);
910                    for gi in 0..gpr {
911                        let tile = &v.nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
912                        let s = sc[gi];
913                        for (k, &b) in tile.iter().enumerate() {
914                            dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
915                            dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
916                        }
917                    }
918                    if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
919                        crate::prism::inverse_embedding(model, dst);
920                    }
921                    return;
922                }
923                if *dtype == TensorDtype::Q2TiledP {
924                    let bytes = self.quant_bytes();
925                    let gpr = cols / GROUP_SIZE;
926                    let v = Q4tpView::new_q2(bytes, self.rows(), cols);
927                    let mut sc = vec![0f32; gpr];
928                    v.scales_into(r, gpr, &mut sc);
929                    for gi in 0..gpr {
930                        let ch =
931                            &v.nib[(r * gpr + gi) * Q2TP_CHUNK..(r * gpr + gi + 1) * Q2TP_CHUNK];
932                        let s = sc[gi];
933                        for (k, &b) in ch.iter().enumerate() {
934                            for j in 0..4 {
935                                let center = if crate::prism::is_affine_target(
936                                    model,
937                                    &model.tensors[*idx].name,
938                                ) {
939                                    1.0
940                                } else {
941                                    1.5
942                                };
943                                dst[gi * GROUP_SIZE + k * 4 + j] =
944                                    (((b >> (2 * j)) & 3) as f32 - center) * s;
945                            }
946                        }
947                    }
948                    if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
949                        crate::prism::inverse_embedding(model, dst);
950                    }
951                    return;
952                }
953                if *dtype == TensorDtype::Q4Block {
954                    let (packed, scales) = q4_split(self.quant_bytes(), self.rows(), cols);
955                    let gpr = cols / GROUP_SIZE;
956                    for gi in 0..gpr {
957                        let g = r * gpr + gi;
958                        let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
959                        for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
960                            dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
961                            dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
962                        }
963                    }
964                    if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
965                        crate::prism::inverse_embedding(model, dst);
966                    }
967                    return;
968                }
969                if *dtype == TensorDtype::Q1 {
970                    let bytes = self.quant_bytes();
971                    let gpr = cols / GROUP_SIZE;
972                    for gi in 0..gpr {
973                        let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
974                        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
975                        for (j, &b) in tile[2..].iter().enumerate() {
976                            for k in 0..8 {
977                                dst[gi * GROUP_SIZE + j * 8 + k] =
978                                    (((b >> k) & 1) as f32 * 2.0 - 1.0) * s;
979                            }
980                        }
981                    }
982                    if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
983                        crate::prism::inverse_embedding(model, dst);
984                    }
985                    return;
986                }
987                if *dtype == TensorDtype::Q1G {
988                    // The tied Bonsai embedding is q1g too: one row = cols/8
989                    // sign bytes + cols/128 scales, same ±s values as q1.
990                    let bytes = self.quant_bytes();
991                    let sc_off = cortiq_core::quant::q1g_scale_offset(self.rows(), cols);
992                    let gpr = cols / GROUP_SIZE;
993                    for gi in 0..gpr {
994                        let t = r * gpr + gi;
995                        let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, t));
996                        for (j, &b) in bytes[t * 4..t * 4 + 4].iter().enumerate() {
997                            for k in 0..8 {
998                                dst[gi * GROUP_SIZE + j * 8 + k] =
999                                    (((b >> k) & 1) as f32 * 2.0 - 1.0) * s;
1000                            }
1001                        }
1002                    }
1003                    if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
1004                        crate::prism::inverse_embedding(model, dst);
1005                    }
1006                    return;
1007                }
1008                if *dtype == TensorDtype::Q1T {
1009                    let bytes = self.quant_bytes();
1010                    let gpr = cols / GROUP_SIZE;
1011                    let base_len = self.rows() * gpr * cortiq_core::quant::Q1T_TILE;
1012                    for gi in 0..gpr {
1013                        let off = (r * gpr + gi) * cortiq_core::quant::Q1T_TILE;
1014                        let s = cortiq_core::quant::f16_to_f32(u16::from_le_bytes([
1015                            bytes[off],
1016                            bytes[off + 1],
1017                        ]));
1018                        let codes = &bytes[off + 2..off + cortiq_core::quant::Q1T_TILE];
1019                        for k in 0..GROUP_SIZE {
1020                            dst[gi * GROUP_SIZE + k] = match cortiq_core::quant::q1t_code(codes, k)
1021                            {
1022                                1 => s,
1023                                2 => -s,
1024                                _ => 0.0,
1025                            };
1026                        }
1027                    }
1028                    // Overlay
1029                    let rows = self.rows();
1030                    let entries = base_len + (rows + 1) * 4;
1031                    if entries <= bytes.len() {
1032                        let ptrs = &bytes[base_len..base_len + (rows + 1) * 4];
1033                        let r0 = u32::from_le_bytes([
1034                            ptrs[r * 4],
1035                            ptrs[r * 4 + 1],
1036                            ptrs[r * 4 + 2],
1037                            ptrs[r * 4 + 3],
1038                        ]) as usize;
1039                        let r1 = u32::from_le_bytes([
1040                            ptrs[(r + 1) * 4],
1041                            ptrs[(r + 1) * 4 + 1],
1042                            ptrs[(r + 1) * 4 + 2],
1043                            ptrs[(r + 1) * 4 + 3],
1044                        ]) as usize;
1045                        let off = entries + r0 * 4;
1046                        for i in 0..r1 - r0 {
1047                            let item = &bytes[off + i * 4..off + i * 4 + 4];
1048                            let c = u16::from_le_bytes([item[0], item[1]]) as usize;
1049                            let v = cortiq_core::quant::f16_to_f32(u16::from_le_bytes([
1050                                item[2], item[3],
1051                            ]));
1052                            if c < cols {
1053                                dst[c] = v;
1054                            }
1055                        }
1056                    }
1057                    if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
1058                        crate::prism::inverse_embedding(model, dst);
1059                    }
1060                    return;
1061                }
1062                if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1063                    let bytes = self.quant_bytes();
1064                    let rows = self.rows();
1065                    let ng = cols / GROUP_SIZE;
1066                    let bits = &bytes[..rows];
1067                    let sc_off = rows;
1068                    // Precomputed at load — embedding lookup used to scan
1069                    // the bit-widths of every preceding row (O(token_id)).
1070                    let off = vbit_offsets[r];
1071                    let b = bits[r] as usize;
1072                    let l = ((1usize << (b - 1)) - 1) as f32;
1073                    let data = &bytes[off..];
1074                    let (mut acc, mut nbits, mut byte_idx) = (0u64, 0usize, 0usize);
1075                    for (i, d) in dst.iter_mut().enumerate() {
1076                        while nbits < b {
1077                            acc = (acc << 8) | data[byte_idx] as u64;
1078                            byte_idx += 1;
1079                            nbits += 8;
1080                        }
1081                        let u = ((acc >> (nbits - b)) & ((1u64 << b) - 1)) as f32;
1082                        nbits -= b;
1083                        let so = (r * ng + i / GROUP_SIZE) * 2;
1084                        let sv = f16_to_f32(u16::from_le_bytes([
1085                            bytes[sc_off + so],
1086                            bytes[sc_off + so + 1],
1087                        ]));
1088                        *d = (u - l) * sv;
1089                    }
1090                    if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
1091                        crate::prism::inverse_embedding(model, dst);
1092                    }
1093                    return;
1094                }
1095                let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
1096                let s = row_scale[r];
1097                match dtype {
1098                    TensorDtype::Q8Row => {
1099                        for (d, &b) in dst.iter_mut().zip(q) {
1100                            *d = (b as i8) as f32 * s;
1101                        }
1102                    }
1103                    TensorDtype::Q8_2f => {
1104                        for (i, (d, &b)) in dst.iter_mut().zip(q).enumerate() {
1105                            *d = (b as i8) as f32 * s * col_field[i];
1106                        }
1107                    }
1108                    _ => unreachable!(),
1109                }
1110                if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
1111                    crate::prism::inverse_embedding(model, dst);
1112                }
1113            }
1114        }
1115    }
1116
1117    /// Can this tensor's columns be read cheaply (for sparse down_proj)?
1118    /// True for F32/Q8Row/Q8_2f (per-row scale, direct strided access);
1119    /// false for group-packed q4/vbit (column access would unpack whole
1120    /// groups — sparse execution falls back to f32 for those).
1121    pub fn sparse_col_ok(&self) -> bool {
1122        match self {
1123            Self::F32 { .. } => true,
1124            Self::Blr(_) => false,
1125            Self::Mapped { dtype, .. } => {
1126                matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
1127            }
1128        }
1129    }
1130
1131    /// down_proj [hidden, inter]: accumulate `w · col(c)` into `out`
1132    /// [hidden] — reads ONLY column `c` (one neuron) from the mmap,
1133    /// no full-matrix dequant. `out[k] += w · down[k, c]`.
1134    pub fn add_col_scaled(&self, c: usize, w: f32, out: &mut [f32]) {
1135        let inter = self.cols();
1136        let hidden = self.rows();
1137        debug_assert_eq!(out.len(), hidden);
1138        match self {
1139            Self::F32 { data, .. } => {
1140                for (k, o) in out.iter_mut().enumerate() {
1141                    *o += w * data[k * inter + c];
1142                }
1143            }
1144            Self::Blr(t) => t.add_col_scaled(c, w, out),
1145            Self::Mapped {
1146                dtype,
1147                row_scale,
1148                col_field,
1149                ..
1150            } => {
1151                let q = self.quant_bytes();
1152                let colf = if *dtype == TensorDtype::Q8_2f {
1153                    col_field[c]
1154                } else {
1155                    1.0
1156                };
1157                let wc = w * colf;
1158                for (k, o) in out.iter_mut().enumerate() {
1159                    let b = q[k * inter + c] as i8 as f32;
1160                    *o += wc * b * row_scale[k];
1161                }
1162            }
1163        }
1164    }
1165
1166    /// Touch the head of row `r` so the DRAM latency of the next
1167    /// neuron's weights overlaps the current one's arithmetic.
1168    ///
1169    /// Scattered rows are what per-token sparsity reads, and a 2 KB
1170    /// stride is past what the hardware prefetcher follows: without this
1171    /// every row starts with a cold miss that nothing hides. One touch
1172    /// per 512 bytes is enough — the rest of the row is a sequential run
1173    /// the prefetcher does pick up.
1174    #[inline]
1175    pub fn prefetch_row(&self, r: usize) {
1176        let Self::Mapped { dtype, .. } = self else {
1177            return;
1178        };
1179        if !matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f) {
1180            return;
1181        }
1182        let cols = self.cols();
1183        let q = self.quant_bytes();
1184        let (a, b) = (r * cols, (r + 1) * cols);
1185        if b > q.len() {
1186            return;
1187        }
1188        let mut j = a;
1189        while j < b {
1190            unsafe { std::ptr::read_volatile(q.as_ptr().add(j)) };
1191            j += 512;
1192        }
1193    }
1194
1195    /// `out += w · row(r)` — the transposed twin of `add_col_scaled`.
1196    ///
1197    /// A neuron's `down` weights are a COLUMN of `[hidden, inter]`, and a
1198    /// column is strided: reading one costs a cache line per element, so
1199    /// per-neuron dynamic sparsity saves arithmetic and no bytes. Stored
1200    /// transposed (`down_proj.t.weight`, `[inter, hidden]`) the same
1201    /// weights are a contiguous ROW, and this accumulate reads exactly
1202    /// the neurons the token asked for.
1203    pub fn add_row_scaled(&self, r: usize, w: f32, out: &mut [f32], scratch: &mut [f32]) {
1204        let cols = self.cols();
1205        debug_assert_eq!(out.len(), cols);
1206        match self {
1207            Self::F32 { data, .. } => {
1208                let row = &data[r * cols..(r + 1) * cols];
1209                for (o, v) in out.iter_mut().zip(row) {
1210                    *o += w * v;
1211                }
1212            }
1213            Self::Blr(t) => {
1214                let tmp = &mut scratch[..cols];
1215                t.row_f32(r, tmp);
1216                for (o, v) in out.iter_mut().zip(tmp.iter()) {
1217                    *o += w * v;
1218                }
1219            }
1220            Self::Mapped {
1221                dtype,
1222                row_scale,
1223                col_field,
1224                ..
1225            } => match dtype {
1226                TensorDtype::Q8Row => {
1227                    let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
1228                    let ws = w * row_scale[r];
1229                    let row: &[i8] =
1230                        unsafe { std::slice::from_raw_parts(q.as_ptr() as *const i8, q.len()) };
1231                    axpy_i8_f32(out, row, ws);
1232                }
1233                TensorDtype::Q8_2f => {
1234                    let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
1235                    let ws = w * row_scale[r];
1236                    for ((o, b), c) in out.iter_mut().zip(q).zip(col_field) {
1237                        *o += ws * c * (*b as i8 as f32);
1238                    }
1239                }
1240                _ => {
1241                    self.row_f32(r, scratch);
1242                    for (o, v) in out.iter_mut().zip(scratch.iter()) {
1243                        *o += w * v;
1244                    }
1245                }
1246            },
1247        }
1248    }
1249
1250    /// Dot of row `r` with `x` (gate/up active-neuron path). Reads only
1251    /// row `r` from the mmap — no full dequant. q4/vbit dequant the row
1252    /// into `scratch` first (rare for active-FFN weights).
1253    pub fn row_dot(&self, r: usize, x: &[f32], scratch: &mut [f32]) -> f32 {
1254        let cols = self.cols();
1255        match self {
1256            Self::F32 { data, .. } => {
1257                let row = &data[r * cols..(r + 1) * cols];
1258                row.iter().zip(x).map(|(w, v)| w * v).sum()
1259            }
1260            Self::Blr(t) => {
1261                let tmp = &mut scratch[..cols];
1262                t.row_f32(r, tmp);
1263                tmp.iter().zip(x).map(|(w, v)| w * v).sum()
1264            }
1265            Self::Mapped {
1266                model,
1267                idx,
1268                dtype,
1269                row_scale,
1270                col_field,
1271                ..
1272            } => {
1273                let prism_forward =
1274                    crate::prism::is_forward_weight(model, &model.tensors[*idx].name);
1275                if prism_forward {
1276                    let transformed = crate::prism::forward(model, &x[..cols]);
1277                    let gpr = cols / GROUP_SIZE;
1278                    match dtype {
1279                        TensorDtype::Q2TiledP => {
1280                            let v = Q4tpView::new_q2(self.quant_bytes(), self.rows(), cols);
1281                            let mut sc = vec![0f32; gpr];
1282                            v.scales_into(r, gpr, &mut sc);
1283                            if crate::prism::is_affine_target(model, &model.tensors[*idx].name) {
1284                                return q2tp_affine_row_exact(v.nib, r, gpr, &transformed, &sc);
1285                            }
1286                            return q2tp_row_exact(v.nib, r, gpr, &transformed, &sc);
1287                        }
1288                        _ => {
1289                            self.row_f32(r, scratch);
1290                            return scratch.iter().zip(&transformed).map(|(w, v)| w * v).sum();
1291                        }
1292                    }
1293                }
1294                match dtype {
1295                    TensorDtype::Q8Row => {
1296                        let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
1297                        dot_i8_f32(q, x) * row_scale[r]
1298                    }
1299                    TensorDtype::Q8_2f => {
1300                        let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
1301                        dot_i8_col_f32(q, x, col_field) * row_scale[r]
1302                    }
1303                    _ => {
1304                        self.row_f32(r, scratch);
1305                        scratch.iter().zip(x).map(|(w, v)| w * v).sum()
1306                    }
1307                }
1308            }
1309        }
1310    }
1311
1312    /// `out = W · x` (row-major). F32 delegates to the historical
1313    /// bit-exact path; Mapped runs the fused int8 kernel.
1314    pub fn matvec(&self, x: &[f32], out: &mut [f32], pool: Option<&Pool>) {
1315        match self {
1316            // NOTE: `out.len()` DRIVES this arm — it computes that many rows,
1317            // and `x.len()` is the stride. A short `out` is legitimate here,
1318            // which is why the check below lives in the Mapped arm only.
1319            Self::F32 { data, .. } => {
1320                if !crate::f32_backend::matvec(data, x, out) {
1321                    matvec_rows(pool, data, x, out);
1322                }
1323            }
1324            Self::Blr(t) => {
1325                if crate::gpu::enabled_here()
1326                    && crate::blr::gpu_wanted()
1327                    && crate::gpu::blr_matmat(t, &x[..t.cols], 1, &mut out[..t.rows])
1328                {
1329                    return;
1330                }
1331                t.matvec(x, out, pool)
1332            }
1333            Self::Mapped {
1334                model,
1335                idx,
1336                dtype,
1337                rows,
1338                cols,
1339                row_scale,
1340                col_field,
1341                vbit_offsets,
1342                repack,
1343            } => {
1344                let _ = (model, idx);
1345                // Every kernel below writes `rows` entries through a raw
1346                // pointer, so a short `out` is an out-of-bounds WRITE, not a
1347                // wrong answer: it scribbles on the allocator's metadata and
1348                // the process aborts much later, somewhere innocent
1349                // (`double free or corruption`, `corrupted double-linked
1350                // list`). The debug_assert two of the kernels carried is
1351                // compiled out of the release — exactly the build where it
1352                // matters. Fail here instead, while the caller is still on
1353                // the stack to be named.
1354                assert!(
1355                    out.len() >= *rows && x.len() >= *cols,
1356                    "matvec {rows}x{cols}: out {} (need {rows}), x {} (need {cols})",
1357                    out.len(),
1358                    x.len(),
1359                );
1360                let prism_forward =
1361                    crate::prism::is_forward_weight(model, &model.tensors[*idx].name);
1362                if *dtype == TensorDtype::Q2TiledP
1363                    && std::env::var("CMF_Q2TP_TRACE").as_deref() == Ok("1")
1364                {
1365                    use std::sync::atomic::{AtomicUsize, Ordering};
1366                    static N: AtomicUsize = AtomicUsize::new(0);
1367                    let n = N.fetch_add(1, Ordering::Relaxed);
1368                    if n < 128 {
1369                        eprintln!(
1370                            "q2tp-dispatch #{n} name={} prism={} rows={} cols={} gpu={} optin={} layer={}",
1371                            model.tensors[*idx].name,
1372                            prism_forward,
1373                            rows,
1374                            cols,
1375                            crate::gpu::enabled_here(),
1376                            crate::gpu::q2tp_gpu_opt_in(),
1377                            crate::gpu::cur_layer(),
1378                        );
1379                    }
1380                }
1381                // Prism stores every manifest-listed forward matrix in the
1382                // signed-Hadamard basis.  The q2tp WGSL path receives that
1383                // transformed vector and an explicit affine bit; codecs
1384                // without a descriptor-aware kernel remain on CPU below.
1385                if prism_forward {
1386                    let transformed = crate::prism::forward(model, &x[..*cols]);
1387                    match dtype {
1388                        TensorDtype::Q4Block => {
1389                            q4matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1390                        }
1391                        TensorDtype::Q4Tiled => {
1392                            q4t_matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1393                        }
1394                        TensorDtype::Q4TiledP => {
1395                            q4tp_matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1396                        }
1397                        TensorDtype::Q2TiledP => {
1398                            let affine =
1399                                crate::prism::is_affine_target(model, &model.tensors[*idx].name);
1400                            if *rows * *cols >= 8_388_608
1401                                && crate::gpu::enabled_here()
1402                                && crate::gpu::q2tp_gpu_opt_in()
1403                            {
1404                                let gpu_ok = if affine {
1405                                    crate::gpu::q2tp_affine_matvec(
1406                                        model,
1407                                        *idx,
1408                                        &transformed,
1409                                        *rows,
1410                                        *cols,
1411                                        out,
1412                                    )
1413                                } else {
1414                                    crate::gpu::q2tp_matvec(
1415                                        model,
1416                                        *idx,
1417                                        &transformed,
1418                                        *rows,
1419                                        *cols,
1420                                        out,
1421                                    )
1422                                };
1423                                if gpu_ok {
1424                                    return;
1425                                }
1426                            }
1427                            if affine {
1428                                q2tp_affine_matvec(
1429                                    self.quant_bytes(),
1430                                    &transformed,
1431                                    *rows,
1432                                    *cols,
1433                                    out,
1434                                    pool,
1435                                )
1436                            } else {
1437                                q2tp_matvec(
1438                                    self.quant_bytes(),
1439                                    &transformed,
1440                                    *rows,
1441                                    *cols,
1442                                    out,
1443                                    pool,
1444                                )
1445                            }
1446                        }
1447                        TensorDtype::Q1 | TensorDtype::Q1G => bin_matvec(
1448                            self.bin_w().expect("binary dtype"),
1449                            &transformed,
1450                            *rows,
1451                            *cols,
1452                            out,
1453                            pool,
1454                        ),
1455                        TensorDtype::Q1T => {
1456                            q1t_matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1457                        }
1458                        TensorDtype::Vbit | TensorDtype::VbitRo => vbitmatvec(
1459                            self.quant_bytes(),
1460                            vbit_offsets,
1461                            &transformed,
1462                            *rows,
1463                            *cols,
1464                            out,
1465                            pool,
1466                        ),
1467                        TensorDtype::Q8Row | TensorDtype::Q8_2f => qmatvec(
1468                            self.quant_bytes(),
1469                            repack,
1470                            row_scale,
1471                            &transformed,
1472                            col_field,
1473                            *dtype,
1474                            *rows,
1475                            *cols,
1476                            out,
1477                            pool,
1478                        ),
1479                        _ => unreachable!("unsupported mapped Prism dtype {dtype:?}"),
1480                    }
1481                    return;
1482                }
1483                if *dtype == TensorDtype::Q4Block {
1484                    // GPU route (wgpu q4b kernel) for large q4_block matvecs —
1485                    // gives NVIDIA/AMD/Intel q4 models a GPU path. Probe keeps
1486                    // the winner; Metal returns false → the CPU kernel below.
1487                    if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1488                        let t0 = std::time::Instant::now();
1489                        match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
1490                            crate::gpu::ProbeArm::Gpu => {
1491                                if crate::gpu::q4b_matvec(model, *idx, x, *rows, *cols, out) {
1492                                    crate::gpu::probe_record(
1493                                        crate::gpu::OpClass::Matvec,
1494                                        true,
1495                                        t0.elapsed(),
1496                                    );
1497                                    return;
1498                                }
1499                            }
1500                            crate::gpu::ProbeArm::CpuTimed => {
1501                                q4matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1502                                crate::gpu::probe_record(
1503                                    crate::gpu::OpClass::Matvec,
1504                                    false,
1505                                    t0.elapsed(),
1506                                );
1507                                return;
1508                            }
1509                            crate::gpu::ProbeArm::Cpu => {}
1510                        }
1511                    }
1512                    q4matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1513                    return;
1514                }
1515                if *dtype == TensorDtype::Q4Tiled {
1516                    // GPU route for large q4t matvecs — the lm_head class,
1517                    // same shape as the q4tp arm below. The probe keeps the
1518                    // winner; a backend without the kernel refuses and the
1519                    // CPU path stays.
1520                    if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1521                        let t0 = std::time::Instant::now();
1522                        let cls = crate::gpu::matvec_class(*rows, *cols);
1523                        match crate::gpu::probe_arm(cls) {
1524                            crate::gpu::ProbeArm::Gpu => {
1525                                if crate::gpu::q4t_matvec(model, *idx, x, *rows, *cols, out) {
1526                                    crate::gpu::probe_record(cls, true, t0.elapsed());
1527                                    return;
1528                                }
1529                            }
1530                            crate::gpu::ProbeArm::CpuTimed => {
1531                                q4t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1532                                crate::gpu::probe_record(cls, false, t0.elapsed());
1533                                return;
1534                            }
1535                            crate::gpu::ProbeArm::Cpu => {}
1536                        }
1537                    }
1538                    q4t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1539                    return;
1540                }
1541                if *dtype == TensorDtype::Q4TiledP {
1542                    // GPU route for large q4tp matvecs — the lm_head class.
1543                    // On a q4tp checkpoint the head is the biggest single
1544                    // host matvec left in the decode step, and the batched
1545                    // kernel at b=1 already exists on both backends. Probe
1546                    // keeps the winner, same as q4_block above.
1547                    if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1548                        let t0 = std::time::Instant::now();
1549                        let cls = crate::gpu::matvec_class(*rows, *cols);
1550                        match crate::gpu::probe_arm(cls) {
1551                            crate::gpu::ProbeArm::Gpu => {
1552                                if crate::gpu::q4tp_matvec(model, *idx, x, *rows, *cols, out) {
1553                                    crate::gpu::probe_record(cls, true, t0.elapsed());
1554                                    return;
1555                                }
1556                            }
1557                            crate::gpu::ProbeArm::CpuTimed => {
1558                                q4tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1559                                crate::gpu::probe_record(cls, false, t0.elapsed());
1560                                return;
1561                            }
1562                            crate::gpu::ProbeArm::Cpu => {}
1563                        }
1564                    }
1565                    q4tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1566                    return;
1567                }
1568                if *dtype == TensorDtype::Q2TiledP {
1569                    q2tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1570                    return;
1571                }
1572                if *dtype == TensorDtype::Q1 || *dtype == TensorDtype::Q1G {
1573                    // GPU route for large q1 matvecs (out_proj / lm_head
1574                    // class): the CPU q1 kernel is load-port-bound at
1575                    // ~4 GB/s/core, the GPU one is bandwidth-bound — the
1576                    // probe measures both arms and keeps the winner.
1577                    let w = self.bin_w().expect("binary dtype");
1578                    if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1579                        let t0 = std::time::Instant::now();
1580                        let arm = if crate::gpu::q1_force() {
1581                            crate::gpu::ProbeArm::Gpu
1582                        } else {
1583                            crate::gpu::probe_arm(crate::gpu::OpClass::Matvec)
1584                        };
1585                        match arm {
1586                            crate::gpu::ProbeArm::Gpu => {
1587                                let ok = if *dtype == TensorDtype::Q1G {
1588                                    crate::gpu::q1g_matvec(model, *idx, x, *rows, *cols, out)
1589                                } else {
1590                                    crate::gpu::q1_matvec(model, *idx, x, *rows, *cols, out)
1591                                };
1592                                if ok {
1593                                    crate::gpu::probe_record(
1594                                        crate::gpu::OpClass::Matvec,
1595                                        true,
1596                                        t0.elapsed(),
1597                                    );
1598                                    return;
1599                                }
1600                            }
1601                            crate::gpu::ProbeArm::CpuTimed => {
1602                                bin_matvec(w, x, *rows, *cols, out, pool);
1603                                crate::gpu::probe_record(
1604                                    crate::gpu::OpClass::Matvec,
1605                                    false,
1606                                    t0.elapsed(),
1607                                );
1608                                return;
1609                            }
1610                            crate::gpu::ProbeArm::Cpu => {}
1611                        }
1612                    }
1613                    bin_matvec(w, x, *rows, *cols, out, pool);
1614                    return;
1615                }
1616                if *dtype == TensorDtype::Q1T {
1617                    // GPU route for large q1t matvecs: the ternary BASE dot runs
1618                    // on the GPU (load-port-bound on CPU, like q1), then the
1619                    // sparse overlay is added on the CPU. Probe keeps the winner.
1620                    if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1621                        let t0 = std::time::Instant::now();
1622                        match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
1623                            crate::gpu::ProbeArm::Gpu => {
1624                                if crate::gpu::q1t_matvec(model, *idx, x, *rows, *cols, out) {
1625                                    q1t_add_overlay(self.quant_bytes(), x, *rows, *cols, out, pool);
1626                                    crate::gpu::probe_record(
1627                                        crate::gpu::OpClass::Matvec,
1628                                        true,
1629                                        t0.elapsed(),
1630                                    );
1631                                    return;
1632                                }
1633                            }
1634                            crate::gpu::ProbeArm::CpuTimed => {
1635                                q1t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1636                                crate::gpu::probe_record(
1637                                    crate::gpu::OpClass::Matvec,
1638                                    false,
1639                                    t0.elapsed(),
1640                                );
1641                                return;
1642                            }
1643                            crate::gpu::ProbeArm::Cpu => {}
1644                        }
1645                    }
1646                    q1t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1647                    return;
1648                }
1649                if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1650                    vbitmatvec(self.quant_bytes(), vbit_offsets, x, *rows, *cols, out, pool);
1651                    return;
1652                }
1653                let xs = prescale(x, col_field, *dtype);
1654                // D5: large q8 matrices (lm_head-class) — hybrid
1655                // CPU∥GPU: split the rows, both sides compute
1656                // SIMULTANEOUSLY (same math, shared prescale).
1657                // GPU share: CMF_GPU_SPLIT (0..1, default 0.5).
1658                if *rows >= crate::gpu::min_rows()
1659                    && matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
1660                    && gpu_lmhead_enabled()
1661                    && crate::gpu::enabled_here()
1662                {
1663                    // Runtime probe: alternate the hybrid against the
1664                    // pure-CPU matvec, keep whichever is faster HERE.
1665                    let t0 = std::time::Instant::now();
1666                    match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
1667                        crate::gpu::ProbeArm::Gpu => {}
1668                        crate::gpu::ProbeArm::CpuTimed => {
1669                            qmatvec(
1670                                self.quant_bytes(),
1671                                repack,
1672                                row_scale,
1673                                x,
1674                                col_field,
1675                                *dtype,
1676                                *rows,
1677                                *cols,
1678                                out,
1679                                pool,
1680                            );
1681                            crate::gpu::probe_record(
1682                                crate::gpu::OpClass::Matvec,
1683                                false,
1684                                t0.elapsed(),
1685                            );
1686                            return;
1687                        }
1688                        crate::gpu::ProbeArm::Cpu => {
1689                            qmatvec(
1690                                self.quant_bytes(),
1691                                repack,
1692                                row_scale,
1693                                x,
1694                                col_field,
1695                                *dtype,
1696                                *rows,
1697                                *cols,
1698                                out,
1699                                pool,
1700                            );
1701                            return;
1702                        }
1703                    }
1704                    let frac = gpu_split_frac();
1705                    let cpu_rows = ((*rows as f32) * (1.0 - frac)) as usize;
1706                    let (out_cpu, out_gpu) = out.split_at_mut(cpu_rows);
1707                    let bytes = self.quant_bytes();
1708                    let ok = std::thread::scope(|sc| {
1709                        let g = sc.spawn(|| {
1710                            crate::gpu::q8_matvec_range(
1711                                model,
1712                                *idx,
1713                                cpu_rows,
1714                                &row_scale[cpu_rows..],
1715                                &xs,
1716                                *rows - cpu_rows,
1717                                *cols,
1718                                out_gpu,
1719                            )
1720                        });
1721                        if cpu_rows > 0 {
1722                            // Repack prefix covers the full groups of the
1723                            // CPU half (the split starts at row 0).
1724                            let rep_cpu = if repack.is_empty() {
1725                                &[][..]
1726                            } else {
1727                                &repack[..(cpu_rows / 4) * 4 * *cols]
1728                            };
1729                            qmatvec(
1730                                &bytes[..cpu_rows * *cols],
1731                                rep_cpu,
1732                                &row_scale[..cpu_rows],
1733                                x,
1734                                col_field,
1735                                *dtype,
1736                                cpu_rows,
1737                                *cols,
1738                                out_cpu,
1739                                pool,
1740                            );
1741                        }
1742                        g.join().unwrap_or(false)
1743                    });
1744                    if ok {
1745                        crate::gpu::probe_record(crate::gpu::OpClass::Matvec, true, t0.elapsed());
1746                        return;
1747                    }
1748                    // GPU failed — CPU finishes its half (rows rebased —
1749                    // group offsets don't line up, mmap layout only).
1750                    qmatvec(
1751                        &bytes[cpu_rows * *cols..(*rows) * *cols],
1752                        &[],
1753                        &row_scale[cpu_rows..],
1754                        x,
1755                        col_field,
1756                        *dtype,
1757                        *rows - cpu_rows,
1758                        *cols,
1759                        out_gpu,
1760                        pool,
1761                    );
1762                    return;
1763                }
1764                qmatvec(
1765                    self.quant_bytes(),
1766                    repack,
1767                    row_scale,
1768                    x,
1769                    col_field,
1770                    *dtype,
1771                    *rows,
1772                    *cols,
1773                    out,
1774                    pool,
1775                );
1776            }
1777        }
1778    }
1779
1780    /// Fused two-input matvec (MTP verify pair): weights streamed once.
1781    pub fn matvec2(
1782        &self,
1783        x1: &[f32],
1784        x2: &[f32],
1785        o1: &mut [f32],
1786        o2: &mut [f32],
1787        pool: Option<&Pool>,
1788    ) {
1789        match self {
1790            Self::F32 { data, .. } => matvec_rows2(pool, data, x1, x2, o1, o2),
1791            Self::Blr(_) => {
1792                self.matvec(x1, o1, pool);
1793                self.matvec(x2, o2, pool);
1794            }
1795            Self::Mapped {
1796                model,
1797                idx,
1798                dtype,
1799                rows,
1800                cols,
1801                row_scale,
1802                col_field,
1803                vbit_offsets,
1804                ..
1805            } => {
1806                if crate::prism::is_forward_weight(model, &model.tensors[*idx].name) {
1807                    let tx1 = crate::prism::forward(model, &x1[..*cols]);
1808                    let tx2 = crate::prism::forward(model, &x2[..*cols]);
1809                    match dtype {
1810                        TensorDtype::Q4Block => {
1811                            q4matvec2(self.quant_bytes(), &tx1, &tx2, *rows, *cols, o1, o2, pool)
1812                        }
1813                        TensorDtype::Q4Tiled => {
1814                            q4t_matvec2(self.quant_bytes(), &tx1, &tx2, *rows, *cols, o1, o2, pool)
1815                        }
1816                        TensorDtype::Q4TiledP => {
1817                            q4tp_matvec2(self.quant_bytes(), &tx1, &tx2, *rows, *cols, o1, o2, pool)
1818                        }
1819                        TensorDtype::Q2TiledP => {
1820                            if crate::prism::is_affine_target(model, &model.tensors[*idx].name) {
1821                                q2tp_affine_matvec2(
1822                                    self.quant_bytes(),
1823                                    &tx1,
1824                                    &tx2,
1825                                    *rows,
1826                                    *cols,
1827                                    o1,
1828                                    o2,
1829                                    pool,
1830                                )
1831                            } else {
1832                                q2tp_matvec2(
1833                                    self.quant_bytes(),
1834                                    &tx1,
1835                                    &tx2,
1836                                    *rows,
1837                                    *cols,
1838                                    o1,
1839                                    o2,
1840                                    pool,
1841                                )
1842                            }
1843                        }
1844                        TensorDtype::Q1 | TensorDtype::Q1G => bin_matvec2(
1845                            self.bin_w().expect("binary dtype"),
1846                            &tx1,
1847                            &tx2,
1848                            *rows,
1849                            *cols,
1850                            o1,
1851                            o2,
1852                            pool,
1853                        ),
1854                        TensorDtype::Q1T => {
1855                            q1t_matvec2(self.quant_bytes(), &tx1, &tx2, *rows, *cols, o1, o2, pool)
1856                        }
1857                        TensorDtype::Vbit | TensorDtype::VbitRo => vbitmatvec2(
1858                            self.quant_bytes(),
1859                            vbit_offsets,
1860                            &tx1,
1861                            &tx2,
1862                            *rows,
1863                            *cols,
1864                            o1,
1865                            o2,
1866                            pool,
1867                        ),
1868                        TensorDtype::Q8Row | TensorDtype::Q8_2f => qmatvec2(
1869                            self.quant_bytes(),
1870                            row_scale,
1871                            &tx1,
1872                            &tx2,
1873                            col_field,
1874                            *dtype,
1875                            *rows,
1876                            *cols,
1877                            o1,
1878                            o2,
1879                            pool,
1880                        ),
1881                        _ => unreachable!("unsupported mapped Prism dtype {dtype:?}"),
1882                    }
1883                    return;
1884                }
1885                if *dtype == TensorDtype::Q4Block {
1886                    q4matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1887                    return;
1888                }
1889                if *dtype == TensorDtype::Q4Tiled {
1890                    q4t_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1891                    return;
1892                }
1893                if *dtype == TensorDtype::Q4TiledP {
1894                    q4tp_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1895                    return;
1896                }
1897                if *dtype == TensorDtype::Q2TiledP {
1898                    q2tp_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1899                    return;
1900                }
1901                if let Some(w) = self.bin_w() {
1902                    bin_matvec2(w, x1, x2, *rows, *cols, o1, o2, pool);
1903                    return;
1904                }
1905                if *dtype == TensorDtype::Q1T {
1906                    // Fused ternary pair: one row pass, the register
1907                    // unpack shared across both streams on ARM. (Q1T
1908                    // lacks a row_scale array — scales live inline in
1909                    // the tiles — so it must not fall through to the
1910                    // q8 qmatvec2 below.)
1911                    q1t_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1912                    return;
1913                }
1914                if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1915                    vbitmatvec2(
1916                        self.quant_bytes(),
1917                        vbit_offsets,
1918                        x1,
1919                        x2,
1920                        *rows,
1921                        *cols,
1922                        o1,
1923                        o2,
1924                        pool,
1925                    );
1926                    return;
1927                }
1928                qmatvec2(
1929                    self.quant_bytes(),
1930                    row_scale,
1931                    x1,
1932                    x2,
1933                    col_field,
1934                    *dtype,
1935                    *rows,
1936                    *cols,
1937                    o1,
1938                    o2,
1939                    pool,
1940                );
1941            }
1942        }
1943    }
1944}
1945
1946impl QTensor {
1947    /// Batched matvec (prefill-GEMM): xs — row-major [b, cols],
1948    /// out — row-major [b, rows]. Element-wise semantics are IDENTICAL
1949    /// to b matvec calls (same dot kernels in the same order); the win —
1950    /// the weight row streams from DRAM once per batch, not b times.
1951    /// `(model, index)` when this is a memory-mapped q4tp tensor — the
1952    /// identity a device-resident chain needs to hand `tp_matmat` the
1953    /// weight without going through this struct's own dispatch.
1954    pub fn q4tp_mapped(&self) -> Option<(&std::sync::Arc<CmfModel>, usize)> {
1955        if self.has_prism_contract() {
1956            return None;
1957        }
1958        match self {
1959            Self::Mapped {
1960                model, idx, dtype, ..
1961            } if *dtype == TensorDtype::Q4TiledP => Some((model, *idx)),
1962            _ => None,
1963        }
1964    }
1965
1966    pub fn matmat(&self, xs_all: &[f32], b: usize, out: &mut [f32], pool: Option<&Pool>) {
1967        let cols = self.cols();
1968        let rows = self.rows();
1969        debug_assert_eq!(xs_all.len(), b * cols);
1970        debug_assert_eq!(out.len(), b * rows);
1971        let _prof = crate::cpuprof::time(crate::cpuprof::Slot::Matmat);
1972        // GPTQ calibration: fold this layer's inputs into its Hessian. Only
1973        // Mapped tensors carry a directory name; the check is a relaxed
1974        // atomic load, free when not calibrating.
1975        if crate::gptq_capture::capturing() {
1976            if let Self::Mapped { model, idx, .. } = self {
1977                crate::gptq_capture::accumulate(&model.tensors[*idx].name, xs_all, b, cols);
1978            }
1979            if let Self::Blr(t) = self {
1980                crate::gptq_capture::accumulate(t.name(), xs_all, b, cols);
1981            }
1982        }
1983        match self {
1984            Self::Blr(t) => {
1985                if crate::gpu::enabled_here()
1986                    && crate::blr::gpu_wanted()
1987                    && crate::gpu::blr_matmat(t, &xs_all[..b * cols], b, &mut out[..b * rows])
1988                {
1989                    return;
1990                }
1991                t.matmat(xs_all, b, out, pool)
1992            }
1993            Self::F32 { data, .. } => {
1994                if crate::f32_backend::matmat(data, xs_all, b, rows, cols, out) {
1995                    return;
1996                }
1997                let out_addr = SendMut(out.as_mut_ptr());
1998                // Over the batch, four positions a row at a time: four
1999                // independent chains, each still summing j = 0..cols in the
2000                // scalar order (bit-identical to the one-at-a-time loop and
2001                // to `matvec_rows`). The batch is the wide axis of a prefill;
2002                // the rows may be few (Spark-X2.5's 16-row g_proj spent
2003                // 17 ms a layer here split over its rows).
2004                let run = |start: usize, end: usize| {
2005                    let mut bi = start;
2006                    while bi < end {
2007                        let n = (end - bi).min(4);
2008                        for o in 0..rows {
2009                            let row = &data[o * cols..(o + 1) * cols];
2010                            if n == 4 {
2011                                let x = |k: usize| &xs_all[(bi + k) * cols..(bi + k + 1) * cols];
2012                                let (x0, x1, x2, x3) = (x(0), x(1), x(2), x(3));
2013                                let (mut a0, mut a1, mut a2, mut a3) = (0f32, 0f32, 0f32, 0f32);
2014                                for j in 0..cols {
2015                                    let w = row[j];
2016                                    a0 += w * x0[j];
2017                                    a1 += w * x1[j];
2018                                    a2 += w * x2[j];
2019                                    a3 += w * x3[j];
2020                                }
2021                                for (k, acc) in [a0, a1, a2, a3].into_iter().enumerate() {
2022                                    unsafe { *out_addr.at((bi + k) * rows + o) = acc };
2023                                }
2024                            } else {
2025                                for k in 0..n {
2026                                    let x = &xs_all[(bi + k) * cols..(bi + k + 1) * cols];
2027                                    let mut acc = 0f32;
2028                                    for j in 0..cols {
2029                                        acc += row[j] * x[j];
2030                                    }
2031                                    unsafe { *out_addr.at((bi + k) * rows + o) = acc };
2032                                }
2033                            }
2034                        }
2035                        bi += n;
2036                    }
2037                };
2038                dispatch_rows(pool, b, &run);
2039            }
2040            Self::Mapped {
2041                model,
2042                idx,
2043                dtype,
2044                row_scale,
2045                col_field,
2046                vbit_offsets,
2047                ..
2048            } => {
2049                if crate::prism::is_forward_weight(model, &model.tensors[*idx].name) {
2050                    let mut transformed = Vec::with_capacity(xs_all.len());
2051                    for bi in 0..b {
2052                        transformed.extend_from_slice(&crate::prism::forward(
2053                            model,
2054                            &xs_all[bi * cols..(bi + 1) * cols],
2055                        ));
2056                    }
2057                    match dtype {
2058                        TensorDtype::Q4Block => {
2059                            q4matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
2060                        }
2061                        TensorDtype::Q4Tiled => {
2062                            q4t_matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
2063                        }
2064                        TensorDtype::Q4TiledP => {
2065                            q4tp_matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
2066                        }
2067                        TensorDtype::Q2TiledP => {
2068                            let affine =
2069                                crate::prism::is_affine_target(model, &model.tensors[*idx].name);
2070                            // Affine Prism Q2TP has a descriptor-aware GPU
2071                            // kernel for short/tail batches too.  Unlike the
2072                            // ordinary Q2TP path, don't force b<32 back to a
2073                            // scalar CPU matmat: prefill chunks and the final
2074                            // tail both need to stay on the tested GPU arm.
2075                            let gpu_batch_ok = if affine {
2076                                b >= 2
2077                            } else {
2078                                b >= 32 && b * rows * cols >= 128_000_000
2079                            };
2080                            if gpu_batch_ok
2081                                && cols % 32 == 0
2082                                && crate::gpu::enabled_here()
2083                                && crate::gpu::q2tp_gpu_opt_in()
2084                            {
2085                                let gpu_ok = if affine {
2086                                    crate::gpu::q2tp_affine_matmat(
2087                                        model,
2088                                        *idx,
2089                                        &transformed,
2090                                        b,
2091                                        rows,
2092                                        cols,
2093                                        out,
2094                                    )
2095                                } else {
2096                                    crate::gpu::q2tp_matmat(
2097                                        model,
2098                                        *idx,
2099                                        &transformed,
2100                                        b,
2101                                        rows,
2102                                        cols,
2103                                        out,
2104                                    )
2105                                };
2106                                if gpu_ok {
2107                                    return;
2108                                }
2109                            }
2110                            // A one-token Prism decode is the other short
2111                            // case.  Use the descriptor-aware matvec kernel
2112                            // before falling back to the exact CPU path.
2113                            if affine
2114                                && b == 1
2115                                && cols % 32 == 0
2116                                && crate::gpu::enabled_here()
2117                                && crate::gpu::q2tp_gpu_opt_in()
2118                                && crate::gpu::q2tp_affine_matvec(
2119                                    model,
2120                                    *idx,
2121                                    &transformed[..cols],
2122                                    rows,
2123                                    cols,
2124                                    &mut out[..rows],
2125                                )
2126                            {
2127                                return;
2128                            }
2129                            if affine {
2130                                q2tp_affine_matmat(
2131                                    self.quant_bytes(),
2132                                    &transformed,
2133                                    b,
2134                                    rows,
2135                                    cols,
2136                                    out,
2137                                    pool,
2138                                )
2139                            } else {
2140                                q2tp_matmat(
2141                                    self.quant_bytes(),
2142                                    &transformed,
2143                                    b,
2144                                    rows,
2145                                    cols,
2146                                    out,
2147                                    pool,
2148                                )
2149                            }
2150                        }
2151                        TensorDtype::Q1 | TensorDtype::Q1G => bin_matmat(
2152                            self.bin_w().expect("binary dtype"),
2153                            &transformed,
2154                            b,
2155                            rows,
2156                            cols,
2157                            out,
2158                            pool,
2159                        ),
2160                        TensorDtype::Q1T => {
2161                            q1t_matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
2162                        }
2163                        TensorDtype::Vbit | TensorDtype::VbitRo => vbitmatmat(
2164                            self.quant_bytes(),
2165                            vbit_offsets,
2166                            &transformed,
2167                            b,
2168                            rows,
2169                            cols,
2170                            out,
2171                            pool,
2172                        ),
2173                        TensorDtype::Q8Row | TensorDtype::Q8_2f => {
2174                            let pre: Vec<std::borrow::Cow<'_, [f32]>> = (0..b)
2175                                .map(|bi| {
2176                                    prescale(
2177                                        &transformed[bi * cols..(bi + 1) * cols],
2178                                        col_field,
2179                                        *dtype,
2180                                    )
2181                                })
2182                                .collect();
2183                            qmatmat(self.quant_bytes(), row_scale, &pre, rows, cols, out, pool)
2184                        }
2185                        _ => unreachable!("unsupported mapped Prism dtype {dtype:?}"),
2186                    }
2187                    return;
2188                }
2189                if *dtype == TensorDtype::Q4Block {
2190                    q4matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2191                    return;
2192                }
2193                if *dtype == TensorDtype::Q4TiledP {
2194                    // GPU batched q4tp GEMM (dequant + f32nt mul_mm on the
2195                    // device); the probe keeps whichever beats the CPU arm.
2196                    // Narrow (prompt-encode) and wide (DiT) batches probe
2197                    // as separate classes — the regimes have opposite
2198                    // winners and one shared verdict locked the wrong arm.
2199                    // Kill switch (gpu::mm_kill): one grossly slow GPU op
2200                    // (a fair-condition op is ≤~100 ms even at 1024px)
2201                    // means the device is contended by another process
2202                    // (e.g. a simulator) — verdicts are per-process, so
2203                    // without the bail the whole render crawls behind
2204                    // someone else's queue.
2205                    // Row-exact batches stay on the host: the device GEMM
2206                    // is an f32 dequant-sgemm, not the host matvec's sum.
2207                    if b >= 32
2208                        && b * rows * cols >= 128_000_000
2209                        && cols % 32 == 0
2210                        && !row_exact()
2211                        && !crate::gpu::mm_killed()
2212                        && crate::gpu::enabled_here()
2213                    {
2214                        let class = if b >= 128 {
2215                            crate::gpu::OpClass::MatmatWide
2216                        } else {
2217                            crate::gpu::OpClass::Matmat
2218                        };
2219                        if let Self::Mapped { model, idx, .. } = self {
2220                            // In-process A/B (`CMF_MM_AB=1`). Three
2221                            // wall-clock A/Bs on a shared stand disagreed
2222                            // with each other by 25% on the same change,
2223                            // because the machine drifts between processes
2224                            // and interleaving whole renders does not fix
2225                            // that. Here both arms run back to back on the
2226                            // SAME data inside one call, so whatever the
2227                            // machine is doing, it does to both — and the
2228                            // disagreement between their outputs falls out
2229                            // for free. Doubles the work; a diagnostic,
2230                            // not a mode.
2231                            if crate::mm_ab::on() {
2232                                let mut g = vec![0f32; b * rows];
2233                                let t = std::time::Instant::now();
2234                                let took = crate::gpu::q4tp_matmat(
2235                                    model, *idx, xs_all, b, rows, cols, &mut g,
2236                                );
2237                                let dg = t.elapsed();
2238                                let t = std::time::Instant::now();
2239                                q4tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2240                                let dc = t.elapsed();
2241                                crate::mm_ab::record(b, rows, cols, took, dg, dc, &g, out);
2242                                return;
2243                            }
2244                            let t0 = std::time::Instant::now();
2245                            // A cold call takes the device arm: its sample
2246                            // is discarded either way, and the upload is
2247                            // what the next step needs.
2248                            let resident = crate::gpu::weight_is_resident(model, *idx);
2249                            match crate::gpu::probe_arm_cold_prefers_gpu(class, resident) {
2250                                crate::gpu::ProbeArm::Gpu => {
2251                                    if crate::gpu::q4tp_matmat(
2252                                        model, *idx, xs_all, b, rows, cols, out,
2253                                    ) {
2254                                        let el = t0.elapsed();
2255                                        // Work-proportional budget: ~8× the
2256                                        // fair-device estimate (+20 ms slack).
2257                                        // An absolute cap missed the worst
2258                                        // case — contended ops sit at
2259                                        // 100–240 ms each and still bury a
2260                                        // render whose fair op is 3–9 ms.
2261                                        // Cold ops (first PSO build, buffer
2262                                        // alloc) are exempt: a one-off
2263                                        // ~50 ms compile is not contention.
2264                                        let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
2265                                        let budget = std::time::Duration::from_secs_f64(
2266                                            flops / 1.5e12 * 8.0 + 0.020,
2267                                        );
2268                                        crate::gpu::mm_budget_check(
2269                                            "q4tp matmat",
2270                                            el,
2271                                            budget,
2272                                            crate::gpu::probe_was_cold() || !resident,
2273                                        );
2274                                        crate::gpu::probe_record(class, true, el);
2275                                        return;
2276                                    }
2277                                }
2278                                crate::gpu::ProbeArm::CpuTimed => {
2279                                    q4tp_matmat(
2280                                        self.quant_bytes(),
2281                                        xs_all,
2282                                        b,
2283                                        rows,
2284                                        cols,
2285                                        out,
2286                                        pool,
2287                                    );
2288                                    crate::gpu::probe_record(class, false, t0.elapsed());
2289                                    return;
2290                                }
2291                                crate::gpu::ProbeArm::Cpu => {}
2292                            }
2293                        }
2294                    }
2295                    q4tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2296                    return;
2297                }
2298                if *dtype == TensorDtype::Q2TiledP {
2299                    // Same device arm as q4tp, behind the same probe:
2300                    // the planes differ, the dispatch does not. Without
2301                    // this a q2tp file ran its widest projections on the
2302                    // host while the 4-bit one had the card, which is a
2303                    // codec paying for its size twice.
2304                    if b >= 32
2305                        && b * rows * cols >= 128_000_000
2306                        && cols % 32 == 0
2307                        && !crate::gpu::mm_killed()
2308                        && crate::gpu::enabled_here()
2309                    {
2310                        let class = if b >= 128 {
2311                            crate::gpu::OpClass::MatmatWide
2312                        } else {
2313                            crate::gpu::OpClass::Matmat
2314                        };
2315                        if let Self::Mapped { model, idx, .. } = self {
2316                            let t0 = std::time::Instant::now();
2317                            match crate::gpu::probe_arm(class) {
2318                                crate::gpu::ProbeArm::Gpu => {
2319                                    if crate::gpu::q2tp_matmat(
2320                                        model, *idx, xs_all, b, rows, cols, out,
2321                                    ) {
2322                                        crate::gpu::probe_record(class, true, t0.elapsed());
2323                                        return;
2324                                    }
2325                                }
2326                                crate::gpu::ProbeArm::CpuTimed => {
2327                                    q2tp_matmat(
2328                                        self.quant_bytes(),
2329                                        xs_all,
2330                                        b,
2331                                        rows,
2332                                        cols,
2333                                        out,
2334                                        pool,
2335                                    );
2336                                    crate::gpu::probe_record(class, false, t0.elapsed());
2337                                    return;
2338                                }
2339                                crate::gpu::ProbeArm::Cpu => {}
2340                            }
2341                        }
2342                    }
2343                    // Without a host arm a q2tp tensor falls through to
2344                    // the q8 fallback, which reads it at one BYTE per
2345                    // weight — a 2x overrun that killed pool workers
2346                    // mid-prefill while the dispatcher waited forever.
2347                    q2tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2348                    return;
2349                }
2350                if *dtype == TensorDtype::Q4Tiled {
2351                    // GPU batched q4t GEMM (dequant + f32nt mul_mm on the
2352                    // device); the probe keeps whichever beats the CPU arm.
2353                    // Narrow (prompt-encode) and wide (DiT) batches probe
2354                    // as separate classes — the regimes have opposite
2355                    // winners and one shared verdict locked the wrong arm.
2356                    // Kill switch (gpu::mm_kill): one grossly slow GPU op
2357                    // (a fair-condition op is ≤~100 ms even at 1024px)
2358                    // means the device is contended by another process
2359                    // (e.g. a simulator) — verdicts are per-process, so
2360                    // without the bail the whole render crawls behind
2361                    // someone else's queue.
2362                    if b >= 32
2363                        && b * rows * cols >= 128_000_000
2364                        && cols % 32 == 0
2365                        && !crate::gpu::mm_killed()
2366                        && crate::gpu::enabled_here()
2367                    {
2368                        let class = if b >= 128 {
2369                            crate::gpu::OpClass::MatmatWide
2370                        } else {
2371                            crate::gpu::OpClass::Matmat
2372                        };
2373                        if let Self::Mapped { model, idx, .. } = self {
2374                            let t0 = std::time::Instant::now();
2375                            match crate::gpu::probe_arm(class) {
2376                                crate::gpu::ProbeArm::Gpu => {
2377                                    if crate::gpu::q4t_matmat(
2378                                        model, *idx, xs_all, b, rows, cols, out,
2379                                    ) {
2380                                        let el = t0.elapsed();
2381                                        // Work-proportional budget: ~8× the
2382                                        // fair-device estimate (+20 ms slack).
2383                                        // An absolute cap missed the worst
2384                                        // case — contended ops sit at
2385                                        // 100–240 ms each and still bury a
2386                                        // render whose fair op is 3–9 ms.
2387                                        // Cold ops (first PSO build, buffer
2388                                        // alloc) are exempt: a one-off
2389                                        // ~50 ms compile is not contention.
2390                                        let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
2391                                        let budget = std::time::Duration::from_secs_f64(
2392                                            flops / 1.5e12 * 8.0 + 0.020,
2393                                        );
2394                                        crate::gpu::mm_budget_check(
2395                                            "q4t matmat",
2396                                            el,
2397                                            budget,
2398                                            crate::gpu::probe_was_cold(),
2399                                        );
2400                                        crate::gpu::probe_record(class, true, el);
2401                                        return;
2402                                    }
2403                                }
2404                                crate::gpu::ProbeArm::CpuTimed => {
2405                                    q4t_matmat(
2406                                        self.quant_bytes(),
2407                                        xs_all,
2408                                        b,
2409                                        rows,
2410                                        cols,
2411                                        out,
2412                                        pool,
2413                                    );
2414                                    crate::gpu::probe_record(class, false, t0.elapsed());
2415                                    return;
2416                                }
2417                                crate::gpu::ProbeArm::Cpu => {}
2418                            }
2419                        }
2420                    }
2421                    q4t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2422                    return;
2423                }
2424                if *dtype == TensorDtype::Q1 || *dtype == TensorDtype::Q1G {
2425                    // GPU batched q1 GEMM for wide prefill (q1_mul_mm on the
2426                    // device); the probe keeps whichever beats the CPU matmat.
2427                    // q1g has the twin `q1g_mul_mm` (cols % 128 by format).
2428                    let w = self.bin_w().expect("binary dtype");
2429                    if b >= 32
2430                        && b * rows * cols >= 128_000_000
2431                        && cols % 64 == 0
2432                        && crate::gpu::enabled_here()
2433                    {
2434                        if let Self::Mapped { model, idx, .. } = self {
2435                            let t0 = std::time::Instant::now();
2436                            match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
2437                                crate::gpu::ProbeArm::Gpu => {
2438                                    let ok = if *dtype == TensorDtype::Q1G {
2439                                        crate::gpu::q1g_matmat(
2440                                            model, *idx, xs_all, b, rows, cols, out,
2441                                        )
2442                                    } else {
2443                                        crate::gpu::q1_matmat(
2444                                            model, *idx, xs_all, b, rows, cols, out,
2445                                        )
2446                                    };
2447                                    if ok {
2448                                        crate::gpu::probe_record(
2449                                            crate::gpu::OpClass::Matmat,
2450                                            true,
2451                                            t0.elapsed(),
2452                                        );
2453                                        return;
2454                                    }
2455                                }
2456                                crate::gpu::ProbeArm::CpuTimed => {
2457                                    bin_matmat(w, xs_all, b, rows, cols, out, pool);
2458                                    crate::gpu::probe_record(
2459                                        crate::gpu::OpClass::Matmat,
2460                                        false,
2461                                        t0.elapsed(),
2462                                    );
2463                                    return;
2464                                }
2465                                crate::gpu::ProbeArm::Cpu => {}
2466                            }
2467                        }
2468                    }
2469                    bin_matmat(w, xs_all, b, rows, cols, out, pool);
2470                    return;
2471                }
2472                if *dtype == TensorDtype::Q1T {
2473                    // GPU batched GEMM for wide prefill (base + overlay on the
2474                    // device); probe keeps the winner vs the CPU matmat.
2475                    if b >= 32 && b * rows * cols >= 128_000_000 && crate::gpu::enabled_here() {
2476                        if let Self::Mapped { model, idx, .. } = self {
2477                            let t0 = std::time::Instant::now();
2478                            match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
2479                                crate::gpu::ProbeArm::Gpu => {
2480                                    if crate::gpu::q1t_matmat(
2481                                        model, *idx, xs_all, b, rows, cols, out,
2482                                    ) {
2483                                        crate::gpu::probe_record(
2484                                            crate::gpu::OpClass::Matmat,
2485                                            true,
2486                                            t0.elapsed(),
2487                                        );
2488                                        return;
2489                                    }
2490                                }
2491                                crate::gpu::ProbeArm::CpuTimed => {
2492                                    q1t_matmat(
2493                                        self.quant_bytes(),
2494                                        xs_all,
2495                                        b,
2496                                        rows,
2497                                        cols,
2498                                        out,
2499                                        pool,
2500                                    );
2501                                    crate::gpu::probe_record(
2502                                        crate::gpu::OpClass::Matmat,
2503                                        false,
2504                                        t0.elapsed(),
2505                                    );
2506                                    return;
2507                                }
2508                                crate::gpu::ProbeArm::Cpu => {}
2509                            }
2510                        }
2511                    }
2512                    q1t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2513                    return;
2514                }
2515                if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
2516                    vbitmatmat(
2517                        self.quant_bytes(),
2518                        vbit_offsets,
2519                        xs_all,
2520                        b,
2521                        rows,
2522                        cols,
2523                        out,
2524                        pool,
2525                    );
2526                    return;
2527                }
2528                let pre: Vec<std::borrow::Cow<'_, [f32]>> = (0..b)
2529                    .map(|bi| prescale(&xs_all[bi * cols..(bi + 1) * cols], col_field, *dtype))
2530                    .collect();
2531                // MiMo verification is a 2–4 row decode panel, not a wide
2532                // prompt GEMM. Keep q8 projections on the same device as
2533                // decode; the generic b>=8 gate otherwise silently moves
2534                // every projection back to CPU. The short wgpu matmat uses
2535                // the same 64-lane reduction as its single-token matvec.
2536                if row_exact()
2537                    && (1..=4).contains(&b)
2538                    && matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
2539                    && crate::gpu::enabled_here()
2540                    && crate::gpu::wgpu_active()
2541                {
2542                    let flat: Vec<f32> = pre.iter().flat_map(|v| v.iter().copied()).collect();
2543                    if crate::gpu::q8_matmat(model, *idx, row_scale, &flat, b, rows, cols, out) {
2544                        return;
2545                    }
2546                }
2547                // D5: large prefill-batch GEMMs — on the GPU (threshold by
2548                // work volume: submission carries b×rows×cols MACs).
2549                // Runtime probe: the naive GEMM shader + sync readback
2550                // lose to the CPU GEMM on slow driver stacks — alternate
2551                // both arms and keep the winner.
2552                if b >= 8 && b * rows * cols >= 128_000_000 && crate::gpu::enabled_here() {
2553                    if let Self::Mapped { model, idx, .. } = self {
2554                        let t0 = std::time::Instant::now();
2555                        match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
2556                            crate::gpu::ProbeArm::Gpu
2557                                if crate::gpu::probe_deciding(crate::gpu::OpClass::Matmat)
2558                                    && !crate::gpu::q8_resident_or_upload(model, *idx) =>
2559                            {
2560                                // Cold weights during probing: the upload
2561                                // has started, the count runs on the CPU —
2562                                // the GPU arm samples on the next touch.
2563                                let q = self.quant_bytes();
2564                                qmatmat(q, row_scale, &pre, rows, cols, out, pool);
2565                                return;
2566                            }
2567                            crate::gpu::ProbeArm::Gpu => {
2568                                // The two-field codec's own entry first, as
2569                                // `device_matmat` takes it: the column field
2570                                // folds into the weight plane and the GEMM
2571                                // runs on the matrix units. The plain int8
2572                                // entry below is the scalar f32 GEMM — on a
2573                                // Spark-X2.5 4B prefill it ran q|k|v|o at
2574                                // host speed.
2575                                if *dtype == TensorDtype::Q8_2f
2576                                    && std::env::var("CMF_Q8_2F_DEV").as_deref() != Ok("0")
2577                                    && crate::gpu::q8_matmat_2f(
2578                                        model, *idx, row_scale, col_field, xs_all, b, rows,
2579                                        cols, out,
2580                                    )
2581                                {
2582                                    crate::gpu::probe_record(
2583                                        crate::gpu::OpClass::Matmat,
2584                                        true,
2585                                        t0.elapsed(),
2586                                    );
2587                                    return;
2588                                }
2589                                let flat: Vec<f32> =
2590                                    pre.iter().flat_map(|v| v.iter().copied()).collect();
2591                                if crate::gpu::q8_matmat(
2592                                    model, *idx, row_scale, &flat, b, rows, cols, out,
2593                                ) {
2594                                    crate::gpu::probe_record(
2595                                        crate::gpu::OpClass::Matmat,
2596                                        true,
2597                                        t0.elapsed(),
2598                                    );
2599                                    return;
2600                                }
2601                            }
2602                            crate::gpu::ProbeArm::CpuTimed => {
2603                                let q = self.quant_bytes();
2604                                qmatmat(q, row_scale, &pre, rows, cols, out, pool);
2605                                crate::gpu::probe_record(
2606                                    crate::gpu::OpClass::Matmat,
2607                                    false,
2608                                    t0.elapsed(),
2609                                );
2610                                return;
2611                            }
2612                            crate::gpu::ProbeArm::Cpu => {}
2613                        }
2614                    }
2615                }
2616                let q = self.quant_bytes();
2617                qmatmat(q, row_scale, &pre, rows, cols, out, pool);
2618            }
2619        }
2620    }
2621}
2622
2623impl QTensor {
2624    /// The device GEMM this tensor would take, run once on the caller's
2625    /// data — the startup parity probe's arm, and the one place that knows
2626    /// which entry point each codec has.
2627    ///
2628    /// It exists because the probe used to look for a `q4tp` weight by
2629    /// name AND dtype, and a container packed any other way was declared
2630    /// "host path" for the whole render even though its codec had a device
2631    /// GEMM of its own. A gate that only recognizes one codec is a gate
2632    /// that silently downgrades every other one.
2633    pub fn device_matmat(&self, xs: &[f32], b: usize, out: &mut [f32]) -> bool {
2634        let (rows, cols) = (self.rows(), self.cols());
2635        let Self::Mapped {
2636            model,
2637            idx,
2638            dtype,
2639            row_scale,
2640            col_field,
2641            ..
2642        } = self
2643        else {
2644            return false;
2645        };
2646        if crate::prism::has_contract(model) {
2647            return false;
2648        }
2649        match *dtype {
2650            TensorDtype::Q4TiledP => crate::gpu::q4tp_matmat(model, *idx, xs, b, rows, cols, out),
2651            // The two-field codec folds its column field into the
2652            // activation, which leaves a plain per-row int8 GEMM — the
2653            // same kernel `q8_row` uses, on both backends.
2654            TensorDtype::Q8Row | TensorDtype::Q8_2f => {
2655                // The field belongs to the weight; only a backend that cannot
2656                // apply it there makes a scaled copy of the activation.
2657                if *dtype == TensorDtype::Q8_2f
2658                    && std::env::var("CMF_Q8_2F_DEV").as_deref() != Ok("0")
2659                    && crate::gpu::q8_matmat_2f(
2660                        model, *idx, row_scale, col_field, xs, b, rows, cols, out,
2661                    )
2662                {
2663                    return true;
2664                }
2665                let flat: Vec<f32> = (0..b)
2666                    .flat_map(|bi| {
2667                        prescale(&xs[bi * cols..(bi + 1) * cols], col_field, *dtype).into_owned()
2668                    })
2669                    .collect();
2670                crate::gpu::q8_matmat(model, *idx, row_scale, &flat, b, rows, cols, out)
2671            }
2672            _ => false,
2673        }
2674    }
2675
2676    /// Multi-matrix job (roadmap §3 P0): N tensors sharing one input
2677    /// run under a SINGLE pool dispatch — QKV or gate+up cost one
2678    /// barrier instead of N. Per-row math is the exact same kernel as
2679    /// `matvec` (bit-identical outputs); only the dispatch is fused.
2680    /// Falls back to N sequential matvecs when the set is not a uniform
2681    /// q8-family/F32 group or there is no pool.
2682    pub fn matvec_many<const N: usize>(
2683        ts: [&QTensor; N],
2684        x: &[f32],
2685        mut outs: [&mut [f32]; N],
2686        pool: Option<&Pool>,
2687    ) {
2688        let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
2689        if ts.iter().any(|t| t.has_prism_contract() || t.is_blr()) {
2690            // The fused range kernels have no transform descriptor.  Let
2691            // each tensor's ordinary matvec dispatch perform the explicit
2692            // signed FWHT (and retain CPU fallback for mixed q2tp/q4tp).
2693            // blr has no row-range kernel either: per-tensor calls.
2694            for (t, o) in ts.iter().zip(outs.iter_mut()) {
2695                t.matvec(x, o, pool);
2696            }
2697            return;
2698        }
2699        let uniform_q8 = ts.iter().all(|t| {
2700            matches!(
2701                t,
2702                Self::Mapped {
2703                    dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f,
2704                    ..
2705                }
2706            )
2707        });
2708        let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
2709        if uniform_f32 && crate::f32_backend::active() {
2710            for (t, o) in ts.iter().zip(outs.iter_mut()) {
2711                t.matvec(x, o, pool);
2712            }
2713            return;
2714        }
2715        let uniform_q4 = ts.iter().all(|t| {
2716            matches!(
2717                t,
2718                Self::Mapped {
2719                    dtype: TensorDtype::Q4Block,
2720                    ..
2721                }
2722            )
2723        });
2724        let uniform_vbit = ts.iter().all(|t| {
2725            matches!(
2726                t,
2727                Self::Mapped {
2728                    dtype: TensorDtype::Vbit | TensorDtype::VbitRo,
2729                    ..
2730                }
2731            )
2732        });
2733        // q1 and q1g may mix freely: the shared activation split and group
2734        // sums are layout-independent, only the row kernel differs.
2735        let uniform_q1 =
2736            ts.iter().all(|t| t.bin_w().is_some()) && ts.iter().all(|t| t.cols() == ts[0].cols());
2737        let uniform_q1t = ts.iter().all(|t| {
2738            matches!(
2739                t,
2740                Self::Mapped {
2741                    dtype: TensorDtype::Q1T,
2742                    ..
2743                }
2744            )
2745        });
2746        // q4tp is the skeleton dtype of the big MoE files, and without an arm
2747        // here every projection that shares an input paid its own pool
2748        // barrier: DeepSeek-V4's attention step alone hands this function
2749        // wq_a, wkv and both compressors' pairs off the same hidden state.
2750        let uniform_q4tp = ts.iter().all(|t| {
2751            matches!(
2752                t,
2753                Self::Mapped {
2754                    dtype: TensorDtype::Q4TiledP,
2755                    ..
2756                }
2757            )
2758        }) && ts
2759            .iter()
2760            .all(|t| t.cols() == ts[0].cols() && t.cols() % GROUP_SIZE == 0);
2761        let Some(pool) = pool else {
2762            for (t, o) in ts.iter().zip(outs.iter_mut()) {
2763                t.matvec(x, o, None);
2764            }
2765            return;
2766        };
2767        if total_rows < 256
2768            || !(uniform_q8
2769                || uniform_f32
2770                || uniform_q4
2771                || uniform_vbit
2772                || uniform_q1
2773                || uniform_q1t
2774                || uniform_q4tp)
2775        {
2776            for (t, o) in ts.iter().zip(outs.iter_mut()) {
2777                t.matvec(x, o, Some(pool));
2778            }
2779            return;
2780        }
2781
2782        if uniform_q4tp {
2783            // Every tensor's rows laid end to end in one virtual row space,
2784            // so the whole set is ONE dispatch. The per-row body is the
2785            // `q4tp_matvec` arm verbatim — same activation split, same
2786            // accumulation order — so the outputs are bit-identical to the
2787            // sequential calls this replaces.
2788            let cols = ts[0].cols();
2789            let gpr = cols / GROUP_SIZE;
2790            let views: [Q4tpView; N] =
2791                std::array::from_fn(|i| Q4tpView::new(ts[i].quant_bytes(), ts[i].rows(), cols));
2792            let rows_of: [usize; N] = std::array::from_fn(|i| ts[i].rows());
2793            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2794            // flat index -> (which tensor, which of its rows)
2795            let locate = |flat: usize| -> (usize, usize) {
2796                let mut acc = 0;
2797                for (i, &r) in rows_of.iter().enumerate() {
2798                    if flat < acc + r {
2799                        return (i, flat - acc);
2800                    }
2801                    acc += r;
2802                }
2803                (rows_of.len() - 1, 0)
2804            };
2805            let (views, outs_addr) = (&views, &outs_addr);
2806            if a8w8_enabled() {
2807                let act = split_act(x);
2808                let act = &act;
2809                let run = |start: usize, end: usize| {
2810                    with_krow(gpr, |sc| {
2811                        for flat in start..end {
2812                            let (t, r) = locate(flat);
2813                            let v = &views[t];
2814                            v.scales_into(r, gpr, sc);
2815                            let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, sc) * act.sx;
2816                            for &(j, xv) in &act.outliers {
2817                                let (w, s) = q4tp_outlier(v.nib, r, gpr, j, sc);
2818                                acc += w * s * xv;
2819                            }
2820                            // SAFETY: one worker owns each (tensor, row) pair.
2821                            unsafe { *outs_addr[t].at(r) = acc };
2822                        }
2823                    });
2824                };
2825                pool.run_rows(total_rows, &run);
2826            } else {
2827                let run = |start: usize, end: usize| {
2828                    with_krow(gpr, |sc| {
2829                        for flat in start..end {
2830                            let (t, r) = locate(flat);
2831                            let v = &views[t];
2832                            v.scales_into(r, gpr, sc);
2833                            // SAFETY: one worker owns each (tensor, row) pair.
2834                            unsafe { *outs_addr[t].at(r) = q4tp_row_exact(v.nib, r, gpr, x, sc) };
2835                        }
2836                    });
2837                };
2838                pool.run_rows(total_rows, &run);
2839            }
2840            return;
2841        }
2842
2843        if uniform_q1 {
2844            // One shared activation split + group sums (q1 has no col
2845            // field; the same input feeds every tensor).
2846            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2847            if a8w8_enabled() {
2848                let act = split_act(x);
2849                let gsum = q1_group_sums(&act.xq, ts[0].cols() / GROUP_SIZE);
2850                let (act, gsum) = (&act, &gsum);
2851                let closures: [_; N] = std::array::from_fn(|i| {
2852                    let (w, gpr, out) = (
2853                        ts[i].bin_w().expect("uniform_q1 checked"),
2854                        ts[i].cols() / GROUP_SIZE,
2855                        outs_addr[i],
2856                    );
2857                    move |s: usize, e: usize| bin_range_a8w8(w, gpr, act, gsum, out, s, e)
2858                });
2859                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2860                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2861                pool.run_many(&parts);
2862            } else {
2863                let closures: [_; N] = std::array::from_fn(|i| {
2864                    let (w, gpr, out) = (
2865                        ts[i].bin_w().expect("uniform_q1 checked"),
2866                        ts[i].cols() / GROUP_SIZE,
2867                        outs_addr[i],
2868                    );
2869                    move |s: usize, e: usize| bin_range_f32(w, gpr, x, out, s, e)
2870                });
2871                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2872                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2873                pool.run_many(&parts);
2874            }
2875            return;
2876        }
2877
2878        if uniform_q1t {
2879            // Q1T batched: one shared activation split + overlay decode,
2880            // all tensors' rows in ONE pool dispatch (saves N−1 dispatches
2881            // and N−1 redundant split_act calls per layer).
2882            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2883            const TILE: usize = cortiq_core::quant::Q1T_TILE;
2884            if a8w8_enabled() {
2885                let act = split_act(x);
2886                let act = &act;
2887                let x_ref = x;
2888                let closures: [_; N] = std::array::from_fn(|i| {
2889                    let bytes = ts[i].quant_bytes();
2890                    let (rows, cols) = (ts[i].rows(), ts[i].cols());
2891                    let gpr = cols / GROUP_SIZE;
2892                    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
2893                    let out = outs_addr[i];
2894                    move |s: usize, e: usize| {
2895                        q1t_range_a8w8(bytes, gpr, rp_off, ent_off, has_ov, act, x_ref, out, s, e)
2896                    }
2897                });
2898                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2899                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2900                pool.run_many(&parts);
2901            } else {
2902                let x_ref = x;
2903                let closures: [_; N] = std::array::from_fn(|i| {
2904                    let bytes = ts[i].quant_bytes();
2905                    let (rows, cols) = (ts[i].rows(), ts[i].cols());
2906                    let gpr = cols / GROUP_SIZE;
2907                    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
2908                    let out = outs_addr[i];
2909                    move |s: usize, e: usize| {
2910                        q1t_range_f32_batch(bytes, gpr, rp_off, ent_off, has_ov, x_ref, out, s, e)
2911                    }
2912                });
2913                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2914                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2915                pool.run_many(&parts);
2916            }
2917            return;
2918        }
2919
2920        if uniform_q4 || uniform_vbit {
2921            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2922            // q4/vbit share one activation split — no per-tensor col field.
2923            if a8w8_enabled() {
2924                let act = split_act(x);
2925                let act = &act;
2926                if uniform_q4 {
2927                    let closures: [_; N] = std::array::from_fn(|i| {
2928                        let (packed, scales) =
2929                            q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
2930                        let (gpr, cols, out) =
2931                            (ts[i].cols() / GROUP_SIZE, ts[i].cols(), outs_addr[i]);
2932                        move |s: usize, e: usize| {
2933                            q4_range_a8w8(packed, scales, gpr, cols, act, out, s, e)
2934                        }
2935                    });
2936                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2937                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2938                    pool.run_many(&parts);
2939                } else {
2940                    let closures: [_; N] = std::array::from_fn(|i| {
2941                        let Self::Mapped { vbit_offsets, .. } = ts[i] else {
2942                            unreachable!()
2943                        };
2944                        let (bytes, rows, cols, out) = (
2945                            ts[i].quant_bytes(),
2946                            ts[i].rows(),
2947                            ts[i].cols(),
2948                            outs_addr[i],
2949                        );
2950                        move |s: usize, e: usize| {
2951                            vbit_range_a8w8(bytes, vbit_offsets, x, act, rows, cols, out, s, e)
2952                        }
2953                    });
2954                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2955                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2956                    pool.run_many(&parts);
2957                }
2958                return;
2959            }
2960            if uniform_q4 {
2961                let closures: [_; N] = std::array::from_fn(|i| {
2962                    let (packed, scales) =
2963                        q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
2964                    let (gpr, out) = (ts[i].cols() / GROUP_SIZE, outs_addr[i]);
2965                    move |s: usize, e: usize| q4_range_f32(packed, scales, gpr, x, out, s, e)
2966                });
2967                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2968                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2969                pool.run_many(&parts);
2970            } else {
2971                let closures: [_; N] = std::array::from_fn(|i| {
2972                    let Self::Mapped { vbit_offsets, .. } = ts[i] else {
2973                        unreachable!()
2974                    };
2975                    let (bytes, rows, cols, out) = (
2976                        ts[i].quant_bytes(),
2977                        ts[i].rows(),
2978                        ts[i].cols(),
2979                        outs_addr[i],
2980                    );
2981                    move |s: usize, e: usize| {
2982                        vbit_range_f32(bytes, vbit_offsets, x, rows, cols, out, s, e)
2983                    }
2984                });
2985                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2986                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2987                pool.run_many(&parts);
2988            }
2989            return;
2990        }
2991
2992        if uniform_f32 {
2993            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2994            let closures: [_; N] = std::array::from_fn(|i| {
2995                let Self::F32 { data, cols, .. } = ts[i] else {
2996                    unreachable!()
2997                };
2998                let out = outs_addr[i];
2999                move |start: usize, end: usize| {
3000                    for o in start..end {
3001                        let row = &data[o * cols..(o + 1) * cols];
3002                        let mut sum = 0.0f32;
3003                        for j in 0..*cols {
3004                            sum += row[j] * x[j];
3005                        }
3006                        // SAFETY: disjoint (tensor, row) cells per worker.
3007                        unsafe { *out.at(o) = sum };
3008                    }
3009                }
3010            });
3011            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3012                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3013            pool.run_many(&parts);
3014            return;
3015        }
3016
3017        // Uniform q8-family: per-tensor prescale (q8_2f col fields
3018        // differ per tensor) + the shared range kernels.
3019        struct Ctx<'a> {
3020            bytes: &'a [u8],
3021            #[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
3022            rep: &'a [u8],
3023            row_scale: &'a [f32],
3024            cols: usize,
3025            xs: std::borrow::Cow<'a, [f32]>,
3026        }
3027        let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
3028            let Self::Mapped {
3029                dtype,
3030                cols,
3031                row_scale,
3032                col_field,
3033                repack,
3034                ..
3035            } = ts[i]
3036            else {
3037                unreachable!()
3038            };
3039            Ctx {
3040                bytes: ts[i].quant_bytes(),
3041                rep: repack,
3042                row_scale,
3043                cols: *cols,
3044                xs: prescale(x, col_field, *dtype),
3045            }
3046        });
3047        let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
3048        #[cfg(target_arch = "aarch64")]
3049        if sdot_enabled() {
3050            let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
3051            let closures: [_; N] = std::array::from_fn(|i| {
3052                let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
3053                move |start: usize, end: usize| {
3054                    q8_range_sdot(c.bytes, c.rep, c.row_scale, act, c.cols, out, start, end)
3055                }
3056            });
3057            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3058                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3059            pool.run_many(&parts);
3060            return;
3061        }
3062        #[cfg(target_arch = "x86_64")]
3063        if avx2_a8w8_enabled() {
3064            let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
3065            let closures: [_; N] = std::array::from_fn(|i| {
3066                let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
3067                move |start: usize, end: usize| {
3068                    q8_range_avx2(c.bytes, c.row_scale, act, c.cols, out, start, end)
3069                }
3070            });
3071            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3072                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3073            pool.run_many(&parts);
3074            return;
3075        }
3076        let closures: [_; N] = std::array::from_fn(|i| {
3077            let (c, out) = (&ctxs[i], outs_addr[i]);
3078            move |start: usize, end: usize| {
3079                q8_range_f32(c.bytes, c.row_scale, &c.xs, c.cols, out, start, end)
3080            }
3081        });
3082        let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3083            std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3084        pool.run_many(&parts);
3085    }
3086}
3087
3088impl QTensor {
3089    /// Pair-input multi-matrix job: N tensors × 2 shared inputs under a
3090    /// single pool dispatch — the MTP/pair decode path publishes one job
3091    /// for Q/K/V (and one for gate+up) instead of one per tensor.
3092    /// Per-row math is exactly `matvec2`'s kernels; bit-identical.
3093    #[allow(clippy::needless_range_loop)]
3094    pub fn matvec2_many<const N: usize>(
3095        ts: [&QTensor; N],
3096        x1: &[f32],
3097        x2: &[f32],
3098        mut o1s: [&mut [f32]; N],
3099        mut o2s: [&mut [f32]; N],
3100        pool: Option<&Pool>,
3101    ) {
3102        let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
3103        if ts.iter().any(|t| t.has_prism_contract() || t.is_blr()) {
3104            for i in 0..N {
3105                ts[i].matvec2(x1, x2, o1s[i], o2s[i], pool);
3106            }
3107            return;
3108        }
3109        let uniform_q8 = ts.iter().all(|t| {
3110            matches!(
3111                t,
3112                Self::Mapped {
3113                    dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f,
3114                    ..
3115                }
3116            )
3117        });
3118        let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
3119        let uniform_q4 = ts.iter().all(|t| {
3120            matches!(
3121                t,
3122                Self::Mapped {
3123                    dtype: TensorDtype::Q4Block,
3124                    ..
3125                }
3126            )
3127        });
3128        let uniform_vbit = ts.iter().all(|t| {
3129            matches!(
3130                t,
3131                Self::Mapped {
3132                    dtype: TensorDtype::Vbit | TensorDtype::VbitRo,
3133                    ..
3134                }
3135            )
3136        });
3137        let fusable = pool.is_some()
3138            && total_rows >= 256
3139            && (uniform_q8 || uniform_f32 || uniform_q4 || uniform_vbit);
3140        if !fusable {
3141            for i in 0..N {
3142                ts[i].matvec2(x1, x2, o1s[i], o2s[i], pool);
3143            }
3144            return;
3145        }
3146        let pool = pool.unwrap();
3147
3148        if uniform_q4 || uniform_vbit {
3149            let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
3150            let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
3151            // q4/vbit share activation splits — no per-tensor col field.
3152            if a8w8_enabled() {
3153                let a1 = split_act(x1);
3154                let a2 = split_act(x2);
3155                let (a1, a2) = (&a1, &a2);
3156                if uniform_q4 {
3157                    let closures: [_; N] = std::array::from_fn(|i| {
3158                        let (packed, scales) =
3159                            q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
3160                        let (gpr, cols, o1, o2) =
3161                            (ts[i].cols() / GROUP_SIZE, ts[i].cols(), p1[i], p2[i]);
3162                        move |s: usize, e: usize| {
3163                            q4_range2_a8w8(packed, scales, gpr, cols, a1, a2, o1, o2, s, e)
3164                        }
3165                    });
3166                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3167                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3168                    pool.run_many(&parts);
3169                } else {
3170                    let closures: [_; N] = std::array::from_fn(|i| {
3171                        let Self::Mapped { vbit_offsets, .. } = ts[i] else {
3172                            unreachable!()
3173                        };
3174                        let (bytes, rows, cols, o1, o2) = (
3175                            ts[i].quant_bytes(),
3176                            ts[i].rows(),
3177                            ts[i].cols(),
3178                            p1[i],
3179                            p2[i],
3180                        );
3181                        move |s: usize, e: usize| {
3182                            vbit_range2_a8w8(
3183                                bytes,
3184                                vbit_offsets,
3185                                x1,
3186                                x2,
3187                                a1,
3188                                a2,
3189                                rows,
3190                                cols,
3191                                o1,
3192                                o2,
3193                                s,
3194                                e,
3195                            )
3196                        }
3197                    });
3198                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3199                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3200                    pool.run_many(&parts);
3201                }
3202                return;
3203            }
3204            if uniform_q4 {
3205                let closures: [_; N] = std::array::from_fn(|i| {
3206                    let (packed, scales) =
3207                        q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
3208                    let (gpr, o1, o2) = (ts[i].cols() / GROUP_SIZE, p1[i], p2[i]);
3209                    move |s: usize, e: usize| {
3210                        q4_range2_f32(packed, scales, gpr, x1, x2, o1, o2, s, e)
3211                    }
3212                });
3213                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3214                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3215                pool.run_many(&parts);
3216            } else {
3217                let closures: [_; N] = std::array::from_fn(|i| {
3218                    let Self::Mapped { vbit_offsets, .. } = ts[i] else {
3219                        unreachable!()
3220                    };
3221                    let (bytes, rows, cols, o1, o2) = (
3222                        ts[i].quant_bytes(),
3223                        ts[i].rows(),
3224                        ts[i].cols(),
3225                        p1[i],
3226                        p2[i],
3227                    );
3228                    move |s: usize, e: usize| {
3229                        vbit_range2_f32(bytes, vbit_offsets, x1, x2, rows, cols, o1, o2, s, e)
3230                    }
3231                });
3232                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3233                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3234                pool.run_many(&parts);
3235            }
3236            return;
3237        }
3238
3239        if uniform_f32 {
3240            let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
3241            let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
3242            let closures: [_; N] = std::array::from_fn(|i| {
3243                let Self::F32 { data, cols, .. } = ts[i] else {
3244                    unreachable!()
3245                };
3246                let (o1, o2) = (p1[i], p2[i]);
3247                move |start: usize, end: usize| {
3248                    for o in start..end {
3249                        let row = &data[o * cols..(o + 1) * cols];
3250                        let (mut s1, mut s2) = (0.0f32, 0.0f32);
3251                        for j in 0..*cols {
3252                            s1 += row[j] * x1[j];
3253                            s2 += row[j] * x2[j];
3254                        }
3255                        // SAFETY: disjoint (tensor, row) cells per worker.
3256                        unsafe {
3257                            *o1.at(o) = s1;
3258                            *o2.at(o) = s2;
3259                        }
3260                    }
3261                }
3262            });
3263            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3264                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3265            pool.run_many(&parts);
3266            return;
3267        }
3268
3269        struct Ctx<'a> {
3270            bytes: &'a [u8],
3271            row_scale: &'a [f32],
3272            cols: usize,
3273            xs1: std::borrow::Cow<'a, [f32]>,
3274            xs2: std::borrow::Cow<'a, [f32]>,
3275        }
3276        let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
3277            let Self::Mapped {
3278                dtype,
3279                cols,
3280                row_scale,
3281                col_field,
3282                ..
3283            } = ts[i]
3284            else {
3285                unreachable!()
3286            };
3287            Ctx {
3288                bytes: ts[i].quant_bytes(),
3289                row_scale,
3290                cols: *cols,
3291                xs1: prescale(x1, col_field, *dtype),
3292                xs2: prescale(x2, col_field, *dtype),
3293            }
3294        });
3295        let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
3296        let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
3297        #[cfg(target_arch = "aarch64")]
3298        if sdot_enabled() {
3299            let acts: [(SplitAct, SplitAct); N] =
3300                std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
3301            let closures: [_; N] = std::array::from_fn(|i| {
3302                let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
3303                move |start: usize, end: usize| {
3304                    q8_range2_sdot(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
3305                }
3306            });
3307            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3308                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3309            pool.run_many(&parts);
3310            return;
3311        }
3312        #[cfg(target_arch = "x86_64")]
3313        if avx2_a8w8_enabled() {
3314            let acts: [(SplitAct, SplitAct); N] =
3315                std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
3316            let closures: [_; N] = std::array::from_fn(|i| {
3317                let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
3318                move |start: usize, end: usize| {
3319                    q8_range2_avx2(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
3320                }
3321            });
3322            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3323                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3324            pool.run_many(&parts);
3325            return;
3326        }
3327        let closures: [_; N] = std::array::from_fn(|i| {
3328            let (c, o1, o2) = (&ctxs[i], p1[i], p2[i]);
3329            move |start: usize, end: usize| {
3330                q8_range2_f32(
3331                    c.bytes,
3332                    c.row_scale,
3333                    &c.xs1,
3334                    &c.xs2,
3335                    c.cols,
3336                    o1,
3337                    o2,
3338                    start,
3339                    end,
3340                )
3341            }
3342        });
3343        let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3344            std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3345        pool.run_many(&parts);
3346    }
3347
3348    /// Fused gate+up matvec with SiLU·mul: for each row r, computes
3349    /// `silu(gate·x) * (up·x)` and writes to `out[r]`. ONE pool dispatch,
3350    /// no intermediate g/u buffers, no separate silu pass. Falls back
3351    /// (returns false) for unsupported dtype combos.
3352    pub fn matvec_silu_mul(
3353        gate: &QTensor,
3354        up: &QTensor,
3355        x: &[f32],
3356        out: &mut [f32],
3357        pool: Option<&Pool>,
3358    ) -> bool {
3359        Self::matvec_silu_mul_limited(gate, up, x, out, 0.0, pool)
3360    }
3361
3362    /// Fused gate+up+SiLU with the GLM asymmetrical clamp.  `limit == 0`
3363    /// preserves the historical unclamped helper; a positive limit clamps
3364    /// `up` to both sides and `gate` only from above, matching the GLM
3365    /// SwiGLU reference.  Keeping the limit in the row kernel avoids the two
3366    /// intermediate vectors and the extra combine pass on the Q2TP experts.
3367    pub fn matvec_silu_mul_limited(
3368        gate: &QTensor,
3369        up: &QTensor,
3370        x: &[f32],
3371        out: &mut [f32],
3372        limit: f32,
3373        pool: Option<&Pool>,
3374    ) -> bool {
3375        if gate.has_prism_contract() || up.has_prism_contract() || gate.is_blr() || up.is_blr() {
3376            // The fused gate/up kernels consume x directly.  Prism requires
3377            // a per-matrix signed FWHT, so the caller must use two ordinary
3378            // descriptor-aware matvecs instead of an unrotated fast path.
3379            return false;
3380        }
3381        let inter = gate.rows();
3382        debug_assert_eq!(up.rows(), inter);
3383        debug_assert_eq!(out.len(), inter);
3384        debug_assert_eq!(gate.cols(), up.cols());
3385        if !a8w8_enabled() {
3386            return false;
3387        }
3388        let act = split_act(x);
3389        let act = &act;
3390        let x_ref = x;
3391        let out_addr = SendMut(out.as_mut_ptr());
3392
3393        match (gate, up) {
3394            // Q4Block gate + Q4Block up (most common mobile q4 models)
3395            (
3396                Self::Mapped {
3397                    dtype: TensorDtype::Q4Block,
3398                    ..
3399                },
3400                Self::Mapped {
3401                    dtype: TensorDtype::Q4Block,
3402                    ..
3403                },
3404            ) => {
3405                let (gp, gs) = q4_split(gate.quant_bytes(), gate.rows(), gate.cols());
3406                let (up_p, up_s) = q4_split(up.quant_bytes(), up.rows(), up.cols());
3407                let gpr = gate.cols() / GROUP_SIZE;
3408                let cols = gate.cols();
3409                let run = move |start: usize, end: usize| {
3410                    for r in start..end {
3411                        let mut gv = dot_q4_row_i8(gp, gs, r * gpr, gpr, &act.xq) * act.sx;
3412                        let mut uv = dot_q4_row_i8(up_p, up_s, r * gpr, gpr, &act.xq) * act.sx;
3413                        for &(j, xv) in &act.outliers {
3414                            let flat = r * cols + j;
3415                            let gb = gp[flat / 2];
3416                            let gn = if flat & 1 == 0 { gb & 0x0F } else { gb >> 4 };
3417                            let gsc = f16_to_f32(u16::from_le_bytes([
3418                                gs[(flat / GROUP_SIZE) * 2],
3419                                gs[(flat / GROUP_SIZE) * 2 + 1],
3420                            ]));
3421                            gv += ((gn as i32 - 8) as f32) * gsc * xv;
3422                            let ub = up_p[flat / 2];
3423                            let un = if flat & 1 == 0 { ub & 0x0F } else { ub >> 4 };
3424                            let usc = f16_to_f32(u16::from_le_bytes([
3425                                up_s[(flat / GROUP_SIZE) * 2],
3426                                up_s[(flat / GROUP_SIZE) * 2 + 1],
3427                            ]));
3428                            uv += ((un as i32 - 8) as f32) * usc * xv;
3429                        }
3430                        // SAFETY: disjoint row ranges per worker.
3431                        unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3432                    }
3433                };
3434                dispatch_rows(pool, inter, &run);
3435                true
3436            }
3437            // Q4Tiled gate + Q4Tiled up — one row pass, both tile
3438            // streams sequential, silu·mul fused (same per-row math as
3439            // `q4t_matvec`).
3440            (
3441                Self::Mapped {
3442                    dtype: TensorDtype::Q4Tiled,
3443                    ..
3444                },
3445                Self::Mapped {
3446                    dtype: TensorDtype::Q4Tiled,
3447                    ..
3448                },
3449            ) => {
3450                let g_bytes = gate.quant_bytes();
3451                let u_bytes = up.quant_bytes();
3452                let gpr = gate.cols() / GROUP_SIZE;
3453                let run = move |start: usize, end: usize| {
3454                    for r in start..end {
3455                        let mut gv = dot_q4t_row_i8(g_bytes, r, gpr, &act.xq) * act.sx;
3456                        let mut uv = dot_q4t_row_i8(u_bytes, r, gpr, &act.xq) * act.sx;
3457                        for &(j, xv) in &act.outliers {
3458                            let (w, s) = q4t_outlier(g_bytes, r, gpr, j);
3459                            gv += w * s * xv;
3460                            let (w, s) = q4t_outlier(u_bytes, r, gpr, j);
3461                            uv += w * s * xv;
3462                        }
3463                        // SAFETY: disjoint row ranges per worker.
3464                        unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3465                    }
3466                };
3467                dispatch_rows(pool, inter, &run);
3468                true
3469            }
3470            // Q4TiledP gate + Q4TiledP up — the same fused row pass, with
3471            // each row's two ladders built once and spent on both streams.
3472            (
3473                Self::Mapped {
3474                    dtype: TensorDtype::Q4TiledP,
3475                    ..
3476                },
3477                Self::Mapped {
3478                    dtype: TensorDtype::Q4TiledP,
3479                    ..
3480                },
3481            ) => {
3482                let cols = gate.cols();
3483                let gpr = cols / GROUP_SIZE;
3484                let gv_view = Q4tpView::new(gate.quant_bytes(), inter, cols);
3485                let uv_view = Q4tpView::new(up.quant_bytes(), inter, cols);
3486                let run = |start: usize, end: usize| {
3487                    with_krows(gpr, |gsc, usc| {
3488                        for r in start..end {
3489                            gv_view.scales_into(r, gpr, gsc);
3490                            uv_view.scales_into(r, gpr, usc);
3491                            let mut gv =
3492                                dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsc) * act.sx;
3493                            let mut uv =
3494                                dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, usc) * act.sx;
3495                            for &(j, xv) in &act.outliers {
3496                                let (w, s) = q4tp_outlier(gv_view.nib, r, gpr, j, gsc);
3497                                gv += w * s * xv;
3498                                let (w, s) = q4tp_outlier(uv_view.nib, r, gpr, j, usc);
3499                                uv += w * s * xv;
3500                            }
3501                            // SAFETY: disjoint row ranges per worker.
3502                            unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3503                        }
3504                    });
3505                };
3506                dispatch_rows(pool, inter, &run);
3507                true
3508            }
3509            // Q1 gate + Q1 up — one row pass over both sign streams,
3510            // silu·mul fused (the per-row math of `q1_range_a8w8`); the
3511            // activation group sums are shared by both streams. Without
3512            // this arm a q1 dense FFN paid two dispatches + a combine
3513            // loop — the exact barrier this function exists to remove.
3514            // q1g (and a q1/q1g mix) rides the same arm through `BinW`.
3515            (
3516                Self::Mapped {
3517                    dtype: TensorDtype::Q1 | TensorDtype::Q1G,
3518                    ..
3519                },
3520                Self::Mapped {
3521                    dtype: TensorDtype::Q1 | TensorDtype::Q1G,
3522                    ..
3523                },
3524            ) => {
3525                let gw = gate.bin_w().expect("binary dtype matched");
3526                let uw = up.bin_w().expect("binary dtype matched");
3527                let gpr = gate.cols() / GROUP_SIZE;
3528                let gsum = q1_group_sums(&act.xq, gpr);
3529                let gsum = &gsum;
3530                let run = move |start: usize, end: usize| {
3531                    for r in start..end {
3532                        let mut gv = gw.dot_i8(r, gpr, &act.xq, gsum) * act.sx;
3533                        let mut uv = uw.dot_i8(r, gpr, &act.xq, gsum) * act.sx;
3534                        for &(j, xv) in &act.outliers {
3535                            let (w, s) = gw.outlier(r, gpr, j);
3536                            gv += w * s * xv;
3537                            let (w, s) = uw.outlier(r, gpr, j);
3538                            uv += w * s * xv;
3539                        }
3540                        // SAFETY: disjoint row ranges per worker.
3541                        unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3542                    }
3543                };
3544                dispatch_rows(pool, inter, &run);
3545                true
3546            }
3547            // Q2TiledP gate + Q2TiledP up — the 2-bit expert pair (MoE
3548            // FFNs of the W2 class): one row pass, both ladders built
3549            // once, integer code dots with shared group sums.
3550            (
3551                Self::Mapped {
3552                    dtype: TensorDtype::Q2TiledP,
3553                    ..
3554                },
3555                Self::Mapped {
3556                    dtype: TensorDtype::Q2TiledP,
3557                    ..
3558                },
3559            ) => {
3560                let cols = gate.cols();
3561                let gpr = cols / GROUP_SIZE;
3562                let gv_view = Q4tpView::new_q2(gate.quant_bytes(), inter, cols);
3563                let uv_view = Q4tpView::new_q2(up.quant_bytes(), inter, cols);
3564                let gsum = q1_group_sums(&act.xq, gpr);
3565                let gsum = &gsum;
3566                let run = move |start: usize, end: usize| {
3567                    with_krows(gpr, |gsc, usc| {
3568                        for r in start..end {
3569                            gv_view.scales_into(r, gpr, gsc);
3570                            uv_view.scales_into(r, gpr, usc);
3571                            let mut gv =
3572                                dot_q2tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsum, gsc) * act.sx;
3573                            let mut uv =
3574                                dot_q2tp_row_i8(uv_view.nib, r, gpr, &act.xq, gsum, usc) * act.sx;
3575                            for &(j, xv) in &act.outliers {
3576                                let (w, s) = q2tp_outlier(gv_view.nib, r, gpr, j, gsc);
3577                                gv += w * s * xv;
3578                                let (w, s) = q2tp_outlier(uv_view.nib, r, gpr, j, usc);
3579                                uv += w * s * xv;
3580                            }
3581                            // SAFETY: disjoint row ranges per worker.
3582                            unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3583                        }
3584                    });
3585                };
3586                dispatch_rows(pool, inter, &run);
3587                true
3588            }
3589            // Q8Row gate + Q8Row up — one row pass over both i8 streams.
3590            // Q8_2f stays out on purpose: its column field prescales the
3591            // activations PER TENSOR, which breaks this fn's shared
3592            // split_act contract — it keeps the two-dispatch path.
3593            (
3594                Self::Mapped {
3595                    dtype: TensorDtype::Q8Row,
3596                    row_scale: g_rs,
3597                    ..
3598                },
3599                Self::Mapped {
3600                    dtype: TensorDtype::Q8Row,
3601                    row_scale: u_rs,
3602                    ..
3603                },
3604            ) => {
3605                let g_bytes = gate.quant_bytes();
3606                let u_bytes = up.quant_bytes();
3607                let cols = gate.cols();
3608                let run = move |start: usize, end: usize| {
3609                    for r in start..end {
3610                        let gv = q8_row_dot(&g_bytes[r * cols..(r + 1) * cols], act) * g_rs[r];
3611                        let uv = q8_row_dot(&u_bytes[r * cols..(r + 1) * cols], act) * u_rs[r];
3612                        // SAFETY: disjoint row ranges per worker.
3613                        unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3614                    }
3615                };
3616                dispatch_rows(pool, inter, &run);
3617                true
3618            }
3619            // Q1T gate + Q1T up
3620            (
3621                Self::Mapped {
3622                    dtype: TensorDtype::Q1T,
3623                    ..
3624                },
3625                Self::Mapped {
3626                    dtype: TensorDtype::Q1T,
3627                    ..
3628                },
3629            ) => {
3630                const TILE: usize = cortiq_core::quant::Q1T_TILE;
3631                let g_bytes = gate.quant_bytes();
3632                let u_bytes = up.quant_bytes();
3633                let gpr = gate.cols() / GROUP_SIZE;
3634                let (g_rp, g_ent, g_ov) = q1t_overlay(g_bytes, inter * gpr * TILE, inter);
3635                let (u_rp, u_ent, u_ov) = q1t_overlay(u_bytes, inter * gpr * TILE, inter);
3636                let run = move |start: usize, end: usize| {
3637                    for r in start..end {
3638                        let mut gv = q1t_dot_row_i8(g_bytes, r, gpr, &act.xq) * act.sx;
3639                        let mut uv = q1t_dot_row_i8(u_bytes, r, gpr, &act.xq) * act.sx;
3640                        for &(j, xv) in &act.outliers {
3641                            gv += q1t_base_weight(g_bytes, r, gpr, j) * xv;
3642                            uv += q1t_base_weight(u_bytes, r, gpr, j) * xv;
3643                        }
3644                        gv += q1t_row_outlier_correction(g_bytes, r, g_rp, g_ent, g_ov, x_ref);
3645                        uv += q1t_row_outlier_correction(u_bytes, r, u_rp, u_ent, u_ov, x_ref);
3646                        // SAFETY: disjoint row ranges per worker.
3647                        unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3648                    }
3649                };
3650                dispatch_rows(pool, inter, &run);
3651                true
3652            }
3653            _ => false,
3654        }
3655    }
3656
3657    /// Every routed expert's fused gate/up/SiLU under ONE pool dispatch.
3658    ///
3659    /// The per-expert path pays a pool barrier per expert per stage: at 9
3660    /// experts over 40 layers that is ~720 barriers a token, and a decode
3661    /// profile of Qwen3.6-35B-A3B showed the pool parked in
3662    /// `psynch_cvwait` about twice as long as it spent computing. Laying
3663    /// every expert's rows end-to-end in one virtual row space collapses
3664    /// the stage to a single dispatch. The per-row body is the
3665    /// single-expert q4tp arm verbatim, so outputs are bit-identical.
3666    ///
3667    /// `false` = something is outside the fused q4tp kernel (dtype, shape,
3668    /// or a transformed tensor); the caller walks the ordinary per-expert
3669    /// path. Float activations use the same exact scalar rows, still fused
3670    /// under one pool dispatch.
3671    pub fn moe_gate_up_many(
3672        pairs: &[(&QTensor, &QTensor)],
3673        x: &[f32],
3674        outs: &mut [Vec<f32>],
3675        pool: Option<&Pool>,
3676    ) -> bool {
3677        Self::moe_gate_up_many_limited(pairs, x, outs, 0.0, pool)
3678    }
3679
3680    /// Batched gate/up/SiLU with the optional GLM clamp.  The public legacy
3681    /// helper above keeps its historical unclamped semantics; callers that
3682    /// implement a reference with a positive SwiGLU limit use this variant.
3683    pub fn moe_gate_up_many_limited(
3684        pairs: &[(&QTensor, &QTensor)],
3685        x: &[f32],
3686        outs: &mut [Vec<f32>],
3687        limit: f32,
3688        pool: Option<&Pool>,
3689    ) -> bool {
3690        if pairs.is_empty() || pairs.len() != outs.len() {
3691            return false;
3692        }
3693        if pairs.iter().any(|(g, u)| g.is_blr() || u.is_blr()) {
3694            return false;
3695        }
3696        if !a8w8_enabled() {
3697            if limit > 0.0 {
3698                // The exact-row fallback applies no SwiGLU clamp; the
3699                // caller's per-expert path carries it instead.
3700                return false;
3701            }
3702            let groups = vec![vec![0]; pairs.len()];
3703            return Self::moe_gate_up_rows(pairs, &groups, x, outs, pool);
3704        }
3705        let inter = pairs[0].0.rows();
3706        let cols = pairs[0].0.cols();
3707        if cols % GROUP_SIZE != 0 {
3708            return false;
3709        }
3710        let gpr = cols / GROUP_SIZE;
3711        // Uniform layout across every routed pair: q4tp, or the 2-bit
3712        // profile's q2tp gate/up (the W2 class). Mixed sets refuse.
3713        let q2 = matches!(
3714            pairs[0].0,
3715            Self::Mapped {
3716                dtype: TensorDtype::Q2TiledP,
3717                ..
3718            }
3719        );
3720        let want = if q2 {
3721            TensorDtype::Q2TiledP
3722        } else {
3723            TensorDtype::Q4TiledP
3724        };
3725        let mut views = Vec::with_capacity(pairs.len() * 2);
3726        for ((g, u), o) in pairs.iter().zip(outs.iter()) {
3727            let both = matches!(g, Self::Mapped { dtype, .. } if *dtype == want)
3728                && matches!(u, Self::Mapped { dtype, .. } if *dtype == want);
3729            if !both
3730                || g.rows() != inter
3731                || u.rows() != inter
3732                || g.cols() != cols
3733                || u.cols() != cols
3734                || o.len() != inter
3735            {
3736                return false;
3737            }
3738            let mk = if q2 { Q4tpView::new_q2 } else { Q4tpView::new };
3739            views.push(mk(g.quant_bytes(), inter, cols));
3740            views.push(mk(u.quant_bytes(), inter, cols));
3741        }
3742        let act = split_act(x);
3743        let gsum = if q2 {
3744            q1_group_sums(&act.xq, gpr)
3745        } else {
3746            Vec::new()
3747        };
3748        let (act, gsum) = (&act, &gsum);
3749        let ptrs: Vec<SendMut> = outs.iter_mut().map(|o| SendMut(o.as_mut_ptr())).collect();
3750        let (views, ptrs) = (&views, &ptrs);
3751        let run = |start: usize, end: usize| {
3752            with_krows(gpr, |gsc, usc| {
3753                for flat in start..end {
3754                    let (e, r) = (flat / inter, flat % inter);
3755                    let gv_view = &views[e * 2];
3756                    let uv_view = &views[e * 2 + 1];
3757                    gv_view.scales_into(r, gpr, gsc);
3758                    uv_view.scales_into(r, gpr, usc);
3759                    let (mut gv, mut uv) = if q2 {
3760                        (
3761                            dot_q2tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsum, gsc) * act.sx,
3762                            dot_q2tp_row_i8(uv_view.nib, r, gpr, &act.xq, gsum, usc) * act.sx,
3763                        )
3764                    } else {
3765                        (
3766                            dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsc) * act.sx,
3767                            dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, usc) * act.sx,
3768                        )
3769                    };
3770                    for &(j, xv) in &act.outliers {
3771                        let (og, ou) = if q2 {
3772                            (
3773                                q2tp_outlier(gv_view.nib, r, gpr, j, gsc),
3774                                q2tp_outlier(uv_view.nib, r, gpr, j, usc),
3775                            )
3776                        } else {
3777                            (
3778                                q4tp_outlier(gv_view.nib, r, gpr, j, gsc),
3779                                q4tp_outlier(uv_view.nib, r, gpr, j, usc),
3780                            )
3781                        };
3782                        gv += og.0 * og.1 * xv;
3783                        uv += ou.0 * ou.1 * xv;
3784                    }
3785                    // SAFETY: one worker owns each (expert, row) pair.
3786                    unsafe { *ptrs[e].at(r) = silu_mul_limited(gv, uv, limit) };
3787                }
3788            });
3789        };
3790        dispatch_rows(pool, pairs.len() * inter, &run);
3791        true
3792    }
3793
3794    /// Every routed expert's down projection, weighted and summed into
3795    /// `out`, under ONE pool dispatch.
3796    ///
3797    /// Partitioned by OUTPUT row rather than by expert: each row is owned
3798    /// by a single worker, so the experts are summed in the caller's order
3799    /// — the same sequence of f32 adds the serial `out[i] += w·eo[i]` loop
3800    /// performs, hence bit-identical. Partitioning by expert instead would
3801    /// race on the shared accumulator.
3802    pub fn moe_down_many(
3803        downs: &[&QTensor],
3804        gs: &[Vec<f32>],
3805        weights: &[f32],
3806        out: &mut [f32],
3807        pool: Option<&Pool>,
3808    ) -> bool {
3809        if downs.is_empty() || downs.len() != gs.len() || downs.len() != weights.len() {
3810            return false;
3811        }
3812        if downs.iter().any(|d| d.is_blr()) {
3813            return false;
3814        }
3815        if !a8w8_enabled() {
3816            let mut terms = vec![vec![0.0; out.len()]; downs.len()];
3817            if !Self::moe_down_rows(downs, &vec![1; downs.len()], gs, &mut terms, pool) {
3818                return false;
3819            }
3820            out.fill(0.0);
3821            for (row, &w) in terms.iter().zip(weights) {
3822                for (o, &v) in out.iter_mut().zip(row) {
3823                    *o += w * v;
3824                }
3825            }
3826            return true;
3827        }
3828        let rows = out.len();
3829        let cols = downs[0].cols();
3830        if cols % GROUP_SIZE != 0 {
3831            return false;
3832        }
3833        let gpr = cols / GROUP_SIZE;
3834        let mut views = Vec::with_capacity(downs.len());
3835        for (d, g) in downs.iter().zip(gs.iter()) {
3836            if !matches!(
3837                d,
3838                Self::Mapped {
3839                    dtype: TensorDtype::Q4TiledP,
3840                    ..
3841                }
3842            ) || d.rows() != rows
3843                || d.cols() != cols
3844                || g.len() != cols
3845            {
3846                return false;
3847            }
3848            views.push(Q4tpView::new(d.quant_bytes(), rows, cols));
3849        }
3850        // One int8 split per expert — the activation vectors differ.
3851        let acts: Vec<SplitAct> = gs.iter().map(|g| split_act(g)).collect();
3852        // Partitioned by OUTPUT row, with the experts folded inside: each
3853        // row is owned by one worker, so they are summed in the caller's
3854        // order — the same f32 sequence the serial `out[i] += w·eo[i]`
3855        // loop produces. Partitioning by expert instead would either race
3856        // on the accumulator or need a scratch plane and a second pass;
3857        // measured, that variant was a wash, so this keeps the simpler
3858        // shape.
3859        let out_addr = SendMut(out.as_mut_ptr());
3860        let (views, acts, weights) = (&views, &acts, &weights);
3861        let run = |start: usize, end: usize| {
3862            with_krow(gpr, |sc| {
3863                for r in start..end {
3864                    let mut acc = 0f32;
3865                    for (e, v) in views.iter().enumerate() {
3866                        v.scales_into(r, gpr, sc);
3867                        let a = &acts[e];
3868                        let mut d = dot_q4tp_row_i8(v.nib, r, gpr, &a.xq, sc) * a.sx;
3869                        for &(j, xv) in &a.outliers {
3870                            let (w, s) = q4tp_outlier(v.nib, r, gpr, j, sc);
3871                            d += w * s * xv;
3872                        }
3873                        acc += weights[e] * d;
3874                    }
3875                    // SAFETY: disjoint row ranges per worker.
3876                    unsafe { *out_addr.at(r) = acc };
3877                }
3878            });
3879        };
3880        dispatch_rows(pool, rows, &run);
3881        true
3882    }
3883
3884    /// `moe_gate_up_many` for SEVERAL tokens at once, decode-exact: expert
3885    /// `e` (`pairs[e]`, q4tp) serves the tokens `groups[e]` (row indices
3886    /// into `xs`, each `cols` wide). Every (expert, token) output is
3887    /// bit-identical to `moe_gate_up_many` run on that token alone — the
3888    /// same int8 activation split, VNNI dots, outlier terms and inline
3889    /// SiLU — while each weight row is read once for all the tokens routed
3890    /// to its expert (the speculative verify's expert sharing). `outs` is
3891    /// flat in (expert, token-of-group) order. False = not covered (not
3892    /// q4tp): the caller takes the per-token path. With float activations,
3893    /// the exact scalar row kernel replaces the int8 dot without changing
3894    /// the shared dispatch or route-order reduction.
3895    pub fn moe_gate_up_rows(
3896        pairs: &[(&QTensor, &QTensor)],
3897        groups: &[Vec<usize>],
3898        xs: &[f32],
3899        outs: &mut [Vec<f32>],
3900        pool: Option<&Pool>,
3901    ) -> bool {
3902        if pairs.is_empty() || pairs.len() != groups.len() {
3903            return false;
3904        }
3905        if pairs.iter().any(|(g, u)| g.is_blr() || u.is_blr()) {
3906            return false;
3907        }
3908        let inter = pairs[0].0.rows();
3909        let cols = pairs[0].0.cols();
3910        let n_pairs: usize = groups.iter().map(|g| g.len()).sum();
3911        if cols == 0 || cols % GROUP_SIZE != 0 || outs.len() != n_pairs || xs.len() % cols != 0 {
3912            return false;
3913        }
3914        let b = xs.len() / cols;
3915        let gpr = cols / GROUP_SIZE;
3916        let mut views = Vec::with_capacity(pairs.len() * 2);
3917        for (g, u) in pairs {
3918            let q4tp = |t: &QTensor| {
3919                matches!(
3920                    t,
3921                    Self::Mapped {
3922                        dtype: TensorDtype::Q4TiledP,
3923                        ..
3924                    }
3925                )
3926            };
3927            if g.has_prism_contract()
3928                || u.has_prism_contract()
3929                || !q4tp(g)
3930                || !q4tp(u)
3931                || g.rows() != inter
3932                || u.rows() != inter
3933                || g.cols() != cols
3934                || u.cols() != cols
3935            {
3936                return false;
3937            }
3938            views.push(Q4tpView::new(g.quant_bytes(), inter, cols));
3939            views.push(Q4tpView::new(u.quant_bytes(), inter, cols));
3940        }
3941        if outs.iter().any(|o| o.len() != inter) || groups.iter().flatten().any(|&t| t >= b) {
3942            return false;
3943        }
3944        let quantized = a8w8_enabled();
3945        let acts: Vec<SplitAct> = if quantized {
3946            (0..b)
3947                .map(|t| split_act(&xs[t * cols..(t + 1) * cols]))
3948                .collect()
3949        } else {
3950            Vec::new()
3951        };
3952        let mut offs = Vec::with_capacity(groups.len());
3953        let mut o = 0usize;
3954        for g in groups {
3955            offs.push(o);
3956            o += g.len();
3957        }
3958        let ptrs: Vec<SendMut> = outs.iter_mut().map(|o| SendMut(o.as_mut_ptr())).collect();
3959        let (views, ptrs, acts, offs) = (&views, &ptrs, &acts, &offs);
3960        let run = |start: usize, end: usize| {
3961            let (mut gsc, mut usc) = (vec![0f32; gpr], vec![0f32; gpr]);
3962            for flat in start..end {
3963                let (e, r) = (flat / inter, flat % inter);
3964                let (gv_view, uv_view) = (&views[e * 2], &views[e * 2 + 1]);
3965                gv_view.scales_into(r, gpr, &mut gsc);
3966                uv_view.scales_into(r, gpr, &mut usc);
3967                for (k, &t) in groups[e].iter().enumerate() {
3968                    if !quantized {
3969                        let x = &xs[t * cols..(t + 1) * cols];
3970                        let gv = q4tp_row_exact(gv_view.nib, r, gpr, x, &gsc);
3971                        let uv = q4tp_row_exact(uv_view.nib, r, gpr, x, &usc);
3972                        unsafe { *ptrs[offs[e] + k].at(r) = (gv / (1.0 + (-gv).exp())) * uv };
3973                        continue;
3974                    }
3975                    let act = &acts[t];
3976                    let mut gv = dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, &gsc) * act.sx;
3977                    let mut uv = dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, &usc) * act.sx;
3978                    for &(j, xv) in &act.outliers {
3979                        let og = q4tp_outlier(gv_view.nib, r, gpr, j, &gsc);
3980                        let ou = q4tp_outlier(uv_view.nib, r, gpr, j, &usc);
3981                        gv += og.0 * og.1 * xv;
3982                        uv += ou.0 * ou.1 * xv;
3983                    }
3984                    let silu_g = gv / (1.0 + (-gv).exp());
3985                    // SAFETY: one worker owns each (expert, row) cell of
3986                    // every output of the expert's group.
3987                    unsafe { *ptrs[offs[e] + k].at(r) = silu_g * uv };
3988                }
3989            }
3990        };
3991        dispatch_rows(pool, pairs.len() * inter, &run);
3992        true
3993    }
3994
3995    /// The per-(expert, token) down terms `moe_down_many` weights and sums,
3996    /// for SEVERAL tokens: `outs[p][o] = down_e[o] · gs[p]` (int8 split of
3997    /// `gs[p]`, VNNI dot, outlier terms — bit-identical to that kernel's
3998    /// `d`), each down row read once for its expert's whole group. The
3999    /// caller sums `w·d` per token in its route order, which reproduces
4000    /// `moe_down_many`'s f32 sequence exactly. Layout as `moe_gate_up_rows`.
4001    pub fn moe_down_rows(
4002        downs: &[&QTensor],
4003        group_lens: &[usize],
4004        gs: &[Vec<f32>],
4005        outs: &mut [Vec<f32>],
4006        pool: Option<&Pool>,
4007    ) -> bool {
4008        if downs.is_empty() || downs.len() != group_lens.len() {
4009            return false;
4010        }
4011        if downs.iter().any(|d| d.is_blr()) {
4012            return false;
4013        }
4014        let rows = downs[0].rows();
4015        let cols = downs[0].cols();
4016        let n_pairs: usize = group_lens.iter().sum();
4017        if cols == 0 || cols % GROUP_SIZE != 0 || gs.len() != n_pairs || outs.len() != n_pairs {
4018            return false;
4019        }
4020        let gpr = cols / GROUP_SIZE;
4021        let mut views = Vec::with_capacity(downs.len());
4022        for d in downs {
4023            if d.has_prism_contract()
4024                || !matches!(
4025                    d,
4026                    Self::Mapped {
4027                        dtype: TensorDtype::Q4TiledP,
4028                        ..
4029                    }
4030                )
4031                || d.rows() != rows
4032                || d.cols() != cols
4033            {
4034                return false;
4035            }
4036            views.push(Q4tpView::new(d.quant_bytes(), rows, cols));
4037        }
4038        if gs.iter().any(|g| g.len() != cols) || outs.iter().any(|o| o.len() != rows) {
4039            return false;
4040        }
4041        let quantized = a8w8_enabled();
4042        let acts: Vec<SplitAct> = if quantized {
4043            gs.iter().map(|g| split_act(g)).collect()
4044        } else {
4045            Vec::new()
4046        };
4047        let mut offs = Vec::with_capacity(group_lens.len());
4048        let mut o = 0usize;
4049        for &l in group_lens {
4050            offs.push(o);
4051            o += l;
4052        }
4053        let ptrs: Vec<SendMut> = outs.iter_mut().map(|o| SendMut(o.as_mut_ptr())).collect();
4054        let (views, ptrs, acts, offs) = (&views, &ptrs, &acts, &offs);
4055        let run = |start: usize, end: usize| {
4056            let mut sc = vec![0f32; gpr];
4057            for flat in start..end {
4058                let (e, r) = (flat / rows, flat % rows);
4059                let v = &views[e];
4060                v.scales_into(r, gpr, &mut sc);
4061                for k in 0..group_lens[e] {
4062                    if !quantized {
4063                        let d = q4tp_row_exact(v.nib, r, gpr, &gs[offs[e] + k], &sc);
4064                        unsafe { *ptrs[offs[e] + k].at(r) = d };
4065                        continue;
4066                    }
4067                    let a = &acts[offs[e] + k];
4068                    let mut d = dot_q4tp_row_i8(v.nib, r, gpr, &a.xq, &sc) * a.sx;
4069                    for &(j, xv) in &a.outliers {
4070                        let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
4071                        d += w * s * xv;
4072                    }
4073                    // SAFETY: one worker owns each (expert, row) cell.
4074                    unsafe { *ptrs[offs[e] + k].at(r) = d };
4075                }
4076            }
4077        };
4078        dispatch_rows(pool, downs.len() * rows, &run);
4079        true
4080    }
4081}
4082
4083/// Batched q8 kernel: same math as qmatvec, the row makes a single
4084/// pass from memory for the whole batch.
4085/// Accelerate CBLAS — the Apple AMX matrix units, the same engine
4086/// llama.cpp's `-ngl 0` prefill rides via ggml-blas.
4087#[cfg(target_os = "macos")]
4088mod accel_blas {
4089    #[link(name = "Accelerate", kind = "framework")]
4090    unsafe extern "C" {
4091        pub fn cblas_sgemm(
4092            order: i32,
4093            trans_a: i32,
4094            trans_b: i32,
4095            m: i32,
4096            n: i32,
4097            k: i32,
4098            alpha: f32,
4099            a: *const f32,
4100            lda: i32,
4101            b: *const f32,
4102            ldb: i32,
4103            beta: f32,
4104            c: *mut f32,
4105            ldc: i32,
4106        );
4107    }
4108}
4109
4110#[cfg(target_os = "macos")]
4111pub(crate) fn accel_gemm_enabled() -> bool {
4112    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4113    *ON.get_or_init(|| std::env::var("CMF_ACCEL").map(|v| v != "0").unwrap_or(true))
4114}
4115
4116/// Off macOS the "accel" GEMM is the portable NEON micro-kernel below —
4117/// same entry point, so the batched-attention path opens on mobile.
4118#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
4119pub(crate) fn accel_gemm_enabled() -> bool {
4120    true
4121}
4122
4123/// Portable NEON f32 GEMM (row-major, optional Bᵀ): a 4×8 fmla
4124/// micro-kernel with A broadcast against B panels — the mobile stand-in
4125/// for Accelerate in the batched causal attention (QKᵀ and P·V). Not a
4126/// BLAS: shapes here are the attention panels (m ≤ heads·chunk,
4127/// k = head_dim or context), and the goal is removing the per-position
4128/// quadratic wall, not peak GEMM.
4129#[cfg(target_arch = "aarch64")]
4130#[allow(clippy::too_many_arguments)]
4131pub(crate) fn neon_gemm_rm(
4132    m: usize,
4133    n: usize,
4134    k: usize,
4135    alpha: f32,
4136    a: &[f32],
4137    lda: usize,
4138    b_mat: &[f32],
4139    ldb: usize,
4140    b_rows_are_n: bool,
4141    c: &mut [f32],
4142    ldc: usize,
4143) {
4144    debug_assert!(a.len() >= (m - 1) * lda + k);
4145    debug_assert!(c.len() >= (m - 1) * ldc + n);
4146    // SAFETY: bounds asserted above; NEON is baseline on aarch64.
4147    unsafe {
4148        use core::arch::aarch64::*;
4149        let mut i = 0usize;
4150        while i < m {
4151            let mi = (m - i).min(4);
4152            let mut j = 0usize;
4153            while j < n {
4154                let nj = (n - j).min(8);
4155                if mi == 4 && nj == 8 {
4156                    let (mut c0a, mut c0b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
4157                    let (mut c1a, mut c1b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
4158                    let (mut c2a, mut c2b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
4159                    let (mut c3a, mut c3b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
4160                    for p in 0..k {
4161                        let (b0, b1) = if b_rows_are_n {
4162                            // B is [n, k]: column p of Bᵀ = element p of
4163                            // eight consecutive B rows — gathered.
4164                            let base = b_mat.as_ptr().add(j * ldb + p);
4165                            let g = |o: usize| *base.add(o * ldb);
4166                            ([g(0), g(1), g(2), g(3)], [g(4), g(5), g(6), g(7)])
4167                        } else {
4168                            let base = b_mat.as_ptr().add(p * ldb + j);
4169                            (
4170                                [*base, *base.add(1), *base.add(2), *base.add(3)],
4171                                [*base.add(4), *base.add(5), *base.add(6), *base.add(7)],
4172                            )
4173                        };
4174                        let bv0 = vld1q_f32(b0.as_ptr());
4175                        let bv1 = vld1q_f32(b1.as_ptr());
4176                        let a0 = vdupq_n_f32(*a.as_ptr().add(i * lda + p));
4177                        let a1 = vdupq_n_f32(*a.as_ptr().add((i + 1) * lda + p));
4178                        let a2 = vdupq_n_f32(*a.as_ptr().add((i + 2) * lda + p));
4179                        let a3 = vdupq_n_f32(*a.as_ptr().add((i + 3) * lda + p));
4180                        c0a = vfmaq_f32(c0a, a0, bv0);
4181                        c0b = vfmaq_f32(c0b, a0, bv1);
4182                        c1a = vfmaq_f32(c1a, a1, bv0);
4183                        c1b = vfmaq_f32(c1b, a1, bv1);
4184                        c2a = vfmaq_f32(c2a, a2, bv0);
4185                        c2b = vfmaq_f32(c2b, a2, bv1);
4186                        c3a = vfmaq_f32(c3a, a3, bv0);
4187                        c3b = vfmaq_f32(c3b, a3, bv1);
4188                    }
4189                    let al = vdupq_n_f32(alpha);
4190                    for (r, (ca, cb)) in [(c0a, c0b), (c1a, c1b), (c2a, c2b), (c3a, c3b)]
4191                        .iter()
4192                        .enumerate()
4193                    {
4194                        let dst = c.as_mut_ptr().add((i + r) * ldc + j);
4195                        vst1q_f32(dst, vmulq_f32(*ca, al));
4196                        vst1q_f32(dst.add(4), vmulq_f32(*cb, al));
4197                    }
4198                } else {
4199                    for r in 0..mi {
4200                        for q in 0..nj {
4201                            let mut acc = 0f32;
4202                            for p in 0..k {
4203                                let bv = if b_rows_are_n {
4204                                    b_mat[(j + q) * ldb + p]
4205                                } else {
4206                                    b_mat[p * ldb + j + q]
4207                                };
4208                                acc += a[(i + r) * lda + p] * bv;
4209                            }
4210                            c[(i + r) * ldc + j + q] = acc * alpha;
4211                        }
4212                    }
4213                }
4214                j += nj;
4215            }
4216            i += mi;
4217        }
4218    }
4219}
4220
4221/// Off-macOS aarch64: the batched attention rides the NEON micro-GEMM.
4222#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
4223#[allow(clippy::too_many_arguments)]
4224pub(crate) fn sgemm_rm(
4225    m: usize,
4226    n: usize,
4227    k: usize,
4228    alpha: f32,
4229    a: &[f32],
4230    lda: usize,
4231    b_mat: &[f32],
4232    ldb: usize,
4233    b_rows_are_n: bool,
4234    c: &mut [f32],
4235    ldc: usize,
4236) {
4237    neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
4238}
4239
4240/// Row-major f32 GEMM, exposed for offline tools (the AWNP pass builds a
4241/// per-layer projection and applies it to every expert; a naive triple loop
4242/// would turn a two-minute job into half an hour).
4243#[allow(clippy::too_many_arguments)]
4244pub fn sgemm_public(
4245    m: usize,
4246    n: usize,
4247    k: usize,
4248    alpha: f32,
4249    a: &[f32],
4250    lda: usize,
4251    b_mat: &[f32],
4252    ldb: usize,
4253    b_rows_are_n: bool,
4254    c: &mut [f32],
4255    ldc: usize,
4256) {
4257    #[cfg(any(target_os = "macos", target_arch = "aarch64"))]
4258    {
4259        sgemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
4260    }
4261    // x86 without Accelerate has no sgemm_rm (the specialized paths there are
4262    // quantized kernels): the packed AVX/FMA/AVX-512 kernels of
4263    // matrixmultiply, single-threaded — offline callers split rows across
4264    // their own threads.
4265    #[cfg(not(any(target_os = "macos", target_arch = "aarch64")))]
4266    {
4267        if m == 0 || n == 0 {
4268            return;
4269        }
4270        assert!(a.len() >= (m - 1) * lda + k && c.len() >= (m - 1) * ldc + n);
4271        let (rsb, csb) = if b_rows_are_n { (1, ldb as isize) } else { (ldb as isize, 1) };
4272        assert!(k == 0 || b_mat.len() > if b_rows_are_n { (n - 1) * ldb + k - 1 } else { (k - 1) * ldb + n - 1 });
4273        // SAFETY: the asserts above bound every element the strides reach.
4274        unsafe {
4275            matrixmultiply::sgemm(
4276                m,
4277                k,
4278                n,
4279                alpha,
4280                a.as_ptr(),
4281                lda as isize,
4282                1,
4283                b_mat.as_ptr(),
4284                rsb,
4285                csb,
4286                0.0,
4287                c.as_mut_ptr(),
4288                ldc as isize,
4289                1,
4290            );
4291        }
4292    }
4293}
4294
4295/// Row-major f32 GEMM on Accelerate: C[m,n] = alpha·A[m,k] × B(ᵀ).
4296/// `b_rows_are_n` = true multiplies by Bᵀ where B is stored [n, k].
4297#[cfg(target_os = "macos")]
4298#[allow(clippy::too_many_arguments)]
4299pub(crate) fn sgemm_rm(
4300    m: usize,
4301    n: usize,
4302    k: usize,
4303    alpha: f32,
4304    a: &[f32],
4305    lda: usize,
4306    b_mat: &[f32],
4307    ldb: usize,
4308    b_rows_are_n: bool,
4309    c: &mut [f32],
4310    ldc: usize,
4311) {
4312    debug_assert!(a.len() >= (m - 1) * lda + k);
4313    debug_assert!(c.len() >= (m - 1) * ldc + n);
4314    // Test hook: route the attention GEMMs through the portable NEON
4315    // micro-kernel ON APPLE SILICON — how the mobile batched attend is
4316    // measured without a phone in the loop. (Intel macOS has no NEON —
4317    // the hook is a no-op there, Accelerate continues below.)
4318    #[cfg(target_arch = "aarch64")]
4319    if std::env::var("CMF_FORCE_NEON_GEMM")
4320        .map(|v| v == "1")
4321        .unwrap_or(false)
4322    {
4323        return neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
4324    }
4325    unsafe {
4326        accel_blas::cblas_sgemm(
4327            101, // RowMajor
4328            111, // NoTrans A
4329            if b_rows_are_n { 112 } else { 111 },
4330            m as i32,
4331            n as i32,
4332            k as i32,
4333            alpha,
4334            a.as_ptr(),
4335            lda as i32,
4336            b_mat.as_ptr(),
4337            ldb as i32,
4338            0.0,
4339            c.as_mut_ptr(),
4340            ldc as i32,
4341        );
4342    }
4343}
4344
4345/// Prefill GEMM through Accelerate (macOS): dequantize q8 rows into
4346/// f32 tiles (scale folded in, pool-parallel) and multiply each tile
4347/// on the AMX with one row-major sgemm. Tiles live in cache, weights
4348/// stream once. Numerics are f32-GEMM (not the int8 dot): prefill
4349/// logits shift within f32 rounding — tolerance-class, like every
4350/// reduction-order change; decode (M=1) never takes this path.
4351#[cfg(target_os = "macos")]
4352fn qmatmat_accel(
4353    q: &[u8],
4354    row_scale: &[f32],
4355    pre: &[std::borrow::Cow<'_, [f32]>],
4356    rows: usize,
4357    cols: usize,
4358    out: &mut [f32],
4359    pool: Option<&Pool>,
4360) {
4361    // NOTE: double-buffering the dequant against the sgemm (a scoped
4362    // thread driving the pool on tile k+1 while the caller multiplies
4363    // tile k) was tried and LOST ~6%: Accelerate's sgemm is itself
4364    // multithreaded, and the dequant workers just steal its cores.
4365    const TR: usize = 2048;
4366    let b = pre.len();
4367    thread_local! {
4368        static XPANEL: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
4369        static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
4370    }
4371    XPANEL.with(|xp| {
4372        WTILE.with(|wt| {
4373            let mut xpanel = xp.borrow_mut();
4374            xpanel.clear();
4375            for x in pre {
4376                xpanel.extend_from_slice(x);
4377            }
4378            let mut wtile = wt.borrow_mut();
4379            wtile.resize(TR * cols, 0.0);
4380            let mut r0 = 0usize;
4381            while r0 < rows {
4382                let tr = TR.min(rows - r0);
4383                // Dequant the tile (scale folded) — pool-parallel.
4384                let wt_addr = SendMut(wtile.as_mut_ptr());
4385                let run = |start: usize, end: usize| {
4386                    for r in start..end {
4387                        let row = &q[(r0 + r) * cols..(r0 + r + 1) * cols];
4388                        let s = row_scale[r0 + r];
4389                        // SAFETY: workers cover disjoint r ranges.
4390                        let dst =
4391                            unsafe { std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols) };
4392                        for (d, &v) in dst.iter_mut().zip(row) {
4393                            *d = (v as i8) as f32 * s;
4394                        }
4395                    }
4396                };
4397                dispatch_rows(pool, tr, &run);
4398                // C[b, tr] (at column r0 of out[b, rows]) = X · Wtileᵀ
4399                unsafe {
4400                    accel_blas::cblas_sgemm(
4401                        101, // RowMajor
4402                        111, // NoTrans A
4403                        112, // Trans B
4404                        b as i32,
4405                        tr as i32,
4406                        cols as i32,
4407                        1.0,
4408                        xpanel.as_ptr(),
4409                        cols as i32,
4410                        wtile.as_ptr(),
4411                        cols as i32,
4412                        0.0,
4413                        out.as_mut_ptr().add(r0),
4414                        rows as i32,
4415                    );
4416                }
4417                r0 += tr;
4418            }
4419        })
4420    });
4421}
4422
4423fn qmatmat(
4424    q: &[u8],
4425    row_scale: &[f32],
4426    pre: &[std::borrow::Cow<'_, [f32]>],
4427    rows: usize,
4428    cols: usize,
4429    out: &mut [f32],
4430    pool: Option<&Pool>,
4431) {
4432    let b = pre.len();
4433    debug_assert_eq!(out.len(), b * rows);
4434    // Big prefill batches ride the AMX (roadmap PR3): the row×batch
4435    // SDOT loop below peaks near the CPU's dot throughput, an order
4436    // below the matrix units. Small tensors and tiny test models stay
4437    // on the exact integer path.
4438    #[cfg(target_os = "macos")]
4439    if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
4440        qmatmat_accel(q, row_scale, pre, rows, cols, out, pool);
4441        return;
4442    }
4443    #[cfg(target_arch = "aarch64")]
4444    if sdot_enabled() {
4445        let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
4446        let out_addr = SendMut(out.as_mut_ptr());
4447        // Blocked 2×4 (mobile prefill: no AMX to fall back on — this
4448        // path IS the ARM prefill GEMM off Apple silicon).
4449        let blocked_ok = blocked_enabled();
4450        let use_i8mm = i8mm_enabled();
4451        if blocked_ok {
4452            let run = |start: usize, end: usize| {
4453                let mut o = start;
4454                while o < end {
4455                    if o + 2 <= end {
4456                        let r0 = &q[o * cols..(o + 1) * cols];
4457                        let r1 = &q[(o + 1) * cols..(o + 2) * cols];
4458                        let mut bi = 0usize;
4459                        while bi + 4 <= acts.len() {
4460                            let xs = [
4461                                acts[bi].xq.as_slice(),
4462                                acts[bi + 1].xq.as_slice(),
4463                                acts[bi + 2].xq.as_slice(),
4464                                acts[bi + 3].xq.as_slice(),
4465                            ];
4466                            let d = if use_i8mm {
4467                                unsafe { dot_i8_smmla_2x4(r0, r1, xs) }
4468                            } else {
4469                                unsafe { dot_i8_sdot_2x4(r0, r1, xs) }
4470                            };
4471                            for (r, row) in [r0, r1].into_iter().enumerate() {
4472                                for k in 0..4 {
4473                                    let act = &acts[bi + k];
4474                                    let mut v = d[r][k] as f32 * act.sx;
4475                                    for &(j, xv) in &act.outliers {
4476                                        v += (row[j] as i8) as f32 * xv;
4477                                    }
4478                                    unsafe {
4479                                        *out_addr.at((bi + k) * rows + o + r) = v * row_scale[o + r]
4480                                    };
4481                                }
4482                            }
4483                            bi += 4;
4484                        }
4485                        while bi < acts.len() {
4486                            for (r, row) in [r0, r1].into_iter().enumerate() {
4487                                let v = row_dot_sdot(row, &acts[bi]) * row_scale[o + r];
4488                                unsafe { *out_addr.at(bi * rows + o + r) = v };
4489                            }
4490                            bi += 1;
4491                        }
4492                        o += 2;
4493                    } else {
4494                        let row = &q[o * cols..(o + 1) * cols];
4495                        for (bi, act) in acts.iter().enumerate() {
4496                            let v = row_dot_sdot(row, act) * row_scale[o];
4497                            unsafe { *out_addr.at(bi * rows + o) = v };
4498                        }
4499                        o += 1;
4500                    }
4501                }
4502            };
4503            dispatch_rows(pool, rows, &run);
4504            return;
4505        }
4506        let run = |start: usize, end: usize| {
4507            for o in start..end {
4508                let row = &q[o * cols..(o + 1) * cols];
4509                for (bi, act) in acts.iter().enumerate() {
4510                    let v = row_dot_sdot(row, act) * row_scale[o];
4511                    unsafe { *out_addr.at(bi * rows + o) = v };
4512                }
4513            }
4514        };
4515        dispatch_rows(pool, rows, &run);
4516        return;
4517    }
4518    // x86 A8W8 batch. Non-VNNI parts take the BLOCKED 2×4 kernel
4519    // (roadmap P0: two weight rows' abs() stay in registers across four
4520    // activation streams); VNNI machines keep the per-row bias-trick
4521    // dot, which is already throughput-bound there.
4522    #[cfg(target_arch = "x86_64")]
4523    if avx2_a8w8_enabled() {
4524        let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
4525        let out_addr = SendMut(out.as_mut_ptr());
4526        // CMF_X86_BLOCKED=0 forces the per-row path (paired in-process
4527        // A/B on noisy shared-vCPU hosts).
4528        let blocked_ok = blocked_enabled();
4529        if !avx512vnni_enabled() && blocked_ok && !row_exact() {
4530            let run = |start: usize, end: usize| {
4531                let mut o = start;
4532                while o < end {
4533                    if o + 2 <= end {
4534                        let r0 = &q[o * cols..(o + 1) * cols];
4535                        let r1 = &q[(o + 1) * cols..(o + 2) * cols];
4536                        let mut bi = 0usize;
4537                        while bi + 4 <= acts.len() {
4538                            let xs = [
4539                                acts[bi].xq.as_slice(),
4540                                acts[bi + 1].xq.as_slice(),
4541                                acts[bi + 2].xq.as_slice(),
4542                                acts[bi + 3].xq.as_slice(),
4543                            ];
4544                            let d = unsafe { dot_i8_i8_avx2_2x4(r0, r1, xs) };
4545                            for (r, row) in [r0, r1].into_iter().enumerate() {
4546                                for k in 0..4 {
4547                                    let act = &acts[bi + k];
4548                                    let mut v = d[r][k] as f32 * act.sx;
4549                                    for &(j, xv) in &act.outliers {
4550                                        v += (row[j] as i8) as f32 * xv;
4551                                    }
4552                                    unsafe {
4553                                        *out_addr.at((bi + k) * rows + o + r) = v * row_scale[o + r]
4554                                    };
4555                                }
4556                            }
4557                            bi += 4;
4558                        }
4559                        while bi < acts.len() {
4560                            for (r, row) in [r0, r1].into_iter().enumerate() {
4561                                let v = row_dot_avx2(row, &acts[bi]) * row_scale[o + r];
4562                                unsafe { *out_addr.at(bi * rows + o + r) = v };
4563                            }
4564                            bi += 1;
4565                        }
4566                        o += 2;
4567                    } else {
4568                        let row = &q[o * cols..(o + 1) * cols];
4569                        for (bi, act) in acts.iter().enumerate() {
4570                            let v = row_dot_avx2(row, act) * row_scale[o];
4571                            unsafe { *out_addr.at(bi * rows + o) = v };
4572                        }
4573                        o += 1;
4574                    }
4575                }
4576            };
4577            dispatch_rows(pool, rows, &run);
4578            return;
4579        }
4580        let run = |start: usize, end: usize| {
4581            for o in start..end {
4582                let row = &q[o * cols..(o + 1) * cols];
4583                for (bi, act) in acts.iter().enumerate() {
4584                    let v = row_dot_avx2(row, act) * row_scale[o];
4585                    unsafe { *out_addr.at(bi * rows + o) = v };
4586                }
4587            }
4588        };
4589        dispatch_rows(pool, rows, &run);
4590        return;
4591    }
4592    let out_addr = SendMut(out.as_mut_ptr());
4593    let run = |start: usize, end: usize| {
4594        for o in start..end {
4595            let row = &q[o * cols..(o + 1) * cols];
4596            for (bi, x) in pre.iter().enumerate() {
4597                let mut acc = 0f32;
4598                for j in 0..cols {
4599                    acc += (row[j] as i8) as f32 * x[j];
4600                }
4601                unsafe { *out_addr.at(bi * rows + o) = acc * row_scale[o] };
4602            }
4603        }
4604    };
4605    dispatch_rows(pool, rows, &run);
4606}
4607
4608/// Split rows across pool workers (shared qmatvec pattern). Self-balancing
4609/// — see `Pool::run_rows` for why a static 1/n split is wrong here.
4610fn dispatch_rows(pool: Option<&Pool>, rows: usize, run: &(dyn Fn(usize, usize) + Sync)) {
4611    match pool {
4612        Some(pool) if rows >= 256 => pool.run_rows(rows, run),
4613        _ => run(0, rows),
4614    }
4615}
4616
4617/// Split a q4_block blob into (packed nibbles, f16 group scales).
4618fn q4_split(bytes: &[u8], rows: usize, cols: usize) -> (&[u8], &[u8]) {
4619    let groups = rows * cols / GROUP_SIZE;
4620    bytes.split_at(groups * 16)
4621}
4622
4623/// SIMD unpack for the dominant vbit width B=4 (94% of rows on the
4624/// log2-shape calibration): 16 packed bytes -> 32 centered i8 values.
4625/// vbit packs MSB-first, so the HIGH nibble is the even element
4626/// (opposite of q4_block's lo-first interleave). Centering is u-7.
4627#[inline]
4628fn vbit_fill4(data: &[u8], buf: &mut [u8]) {
4629    #[cfg(target_arch = "aarch64")]
4630    unsafe {
4631        return vbit_fill4_neon(data, buf);
4632    }
4633    #[cfg(target_arch = "x86_64")]
4634    if avx2_enabled() {
4635        return unsafe { vbit_fill4_avx2(data, buf) };
4636    }
4637    #[allow(unreachable_code)]
4638    for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
4639        let u = unpack8::<4>(&data[blk * 4..]);
4640        for k in 0..8 {
4641            chunk[k] = (u[k] - 7) as i8 as u8;
4642        }
4643    }
4644}
4645
4646#[cfg(target_arch = "aarch64")]
4647#[target_feature(enable = "neon")]
4648unsafe fn vbit_fill4_neon(data: &[u8], buf: &mut [u8]) {
4649    // SAFETY: buf.len() is a multiple of GROUP_SIZE=32; data holds
4650    // buf.len()/2 packed bytes (validated at load).
4651    unsafe {
4652        use core::arch::aarch64::*;
4653        let n = buf.len();
4654        let mask = vdupq_n_u8(0x0F);
4655        let seven = vdupq_n_s8(7);
4656        let mut g = 0usize;
4657        while g * 32 + 32 <= n {
4658            let b = vld1q_u8(data.as_ptr().add(g * 16));
4659            let hi = vshrq_n_u8::<4>(b);
4660            let lo = vandq_u8(b, mask);
4661            let z0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(hi, lo)), seven);
4662            let z1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(hi, lo)), seven);
4663            vst1q_u8(buf.as_mut_ptr().add(g * 32), vreinterpretq_u8_s8(z0));
4664            vst1q_u8(buf.as_mut_ptr().add(g * 32 + 16), vreinterpretq_u8_s8(z1));
4665            g += 1;
4666        }
4667    }
4668}
4669
4670#[cfg(target_arch = "x86_64")]
4671#[target_feature(enable = "avx2")]
4672unsafe fn vbit_fill4_avx2(data: &[u8], buf: &mut [u8]) {
4673    // SAFETY: see vbit_fill4_neon.
4674    unsafe {
4675        use core::arch::x86_64::*;
4676        let n = buf.len();
4677        let mask = _mm_set1_epi8(0x0F);
4678        let seven = _mm256_set1_epi8(7);
4679        let mut g = 0usize;
4680        while g * 32 + 32 <= n {
4681            let b = _mm_loadu_si128(data.as_ptr().add(g * 16) as *const __m128i);
4682            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), mask);
4683            let lo = _mm_and_si128(b, mask);
4684            let z = _mm256_sub_epi8(
4685                _mm256_set_m128i(_mm_unpackhi_epi8(hi, lo), _mm_unpacklo_epi8(hi, lo)),
4686                seven,
4687            );
4688            _mm256_storeu_si256(buf.as_mut_ptr().add(g * 32) as *mut __m256i, z);
4689            g += 1;
4690        }
4691    }
4692}
4693
4694/// Unpack 8 MSB-first B-bit values from exactly B bytes (fixed shifts —
4695/// no serial bit-buffer, auto-vectorizable). Every 32-value group starts
4696/// byte-aligned (32·B/8 is integral for B∈3..8), so groups decompose
4697/// into 4 such blocks.
4698#[inline(always)]
4699fn unpack8<const B: usize>(data: &[u8]) -> [i32; 8] {
4700    let mut acc = 0u64;
4701    for i in 0..B {
4702        acc = (acc << 8) | data[i] as u64;
4703    }
4704    let mask = (1u64 << B) - 1;
4705    let mut out = [0i32; 8];
4706    for (k, o) in out.iter_mut().enumerate() {
4707        *o = ((acc >> ((7 - k) * B)) & mask) as i32;
4708    }
4709    out
4710}
4711
4712/// Fused vbit matvec straight from the mapped bytes (spec §3, P13
4713/// FIG.3): [u8 bits: rows][f16 scales: rows·cols/32][bit-packed rows,
4714/// MSB-first, byte-padded]. Row data offsets are precomputed at load
4715/// (`vbit_row_offsets`) — the per-call prefix scan was O(rows) pure
4716/// overhead on every matvec.
4717#[allow(clippy::too_many_arguments)]
4718fn vbitmatvec(
4719    bytes: &[u8],
4720    offsets: &[usize],
4721    x: &[f32],
4722    rows: usize,
4723    cols: usize,
4724    out: &mut [f32],
4725    pool: Option<&Pool>,
4726) {
4727    debug_assert_eq!(out.len(), rows);
4728    debug_assert_eq!(offsets.len(), rows + 1);
4729
4730    // SDOT path: unpack the row to centered i8 once, then per-group
4731    // int8 dot against the quantized activations — same A8W8 contract
4732    // as q8 (bounded noise; CMF_SDOT=0 keeps the exact scalar path).
4733    if a8w8_enabled() {
4734        let act = split_act(x);
4735        let out_addr = SendMut(out.as_mut_ptr());
4736        let run = move |start: usize, end: usize| {
4737            vbit_range_a8w8(bytes, offsets, x, &act, rows, cols, out_addr, start, end)
4738        };
4739        dispatch_rows(pool, rows, &run);
4740        return;
4741    }
4742
4743    let out_addr = SendMut(out.as_mut_ptr());
4744    let run = move |start: usize, end: usize| {
4745        vbit_range_f32(bytes, offsets, x, rows, cols, out_addr, start, end)
4746    };
4747    dispatch_rows(pool, rows, &run);
4748}
4749
4750/// One vbit row range via the A8W8 int8 path — kernel body of
4751/// `vbitmatvec`, extracted so multi-matrix jobs can drive it for
4752/// several tensors in one dispatch (b=8 rows go exact f32).
4753#[allow(clippy::too_many_arguments)]
4754fn vbit_range_a8w8(
4755    bytes: &[u8],
4756    offsets: &[usize],
4757    x: &[f32],
4758    act: &SplitAct,
4759    rows: usize,
4760    cols: usize,
4761    out: SendMut,
4762    start: usize,
4763    end: usize,
4764) {
4765    let ng = cols / GROUP_SIZE;
4766    let bits = &bytes[..rows];
4767    let sc_off = rows;
4768    let row_dot = |r: usize| -> f32 {
4769        let b = bits[r] as usize;
4770        let l = (1i32 << (b - 1)) - 1;
4771        let mask = (1u64 << b) - 1;
4772        let data = &bytes[offsets[r]..offsets[r + 1]];
4773        if b == 8 {
4774            // u−L reaches 128 → does not fit i8; exact f32 path.
4775            let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
4776            let mut dot = 0f32;
4777            for g in 0..ng {
4778                let so = (r * ng + g) * 2;
4779                let sgf = f16_to_f32(u16::from_le_bytes([
4780                    bytes[sc_off + so],
4781                    bytes[sc_off + so + 1],
4782                ]));
4783                let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
4784                let mut gd = 0f32;
4785                for &xv in xg.iter() {
4786                    if nbits < 8 {
4787                        acc = (acc << 8) | data[idx] as u64;
4788                        idx += 1;
4789                        nbits += 8;
4790                    }
4791                    let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
4792                    nbits -= 8;
4793                    gd += (u - l) as f32 * xv;
4794                }
4795                dot += gd * sgf;
4796            }
4797            return dot;
4798        }
4799        // Per-worker scratch: this closure runs for every row of the
4800        // tensor (lm_head ≈ 150k rows/token) — a heap allocation per
4801        // row was measurable pure overhead.
4802        thread_local! {
4803            static VBIT_SCRATCH: std::cell::RefCell<Vec<u8>> =
4804                const { std::cell::RefCell::new(Vec::new()) };
4805        }
4806        #[inline(always)]
4807        fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
4808            for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
4809                let u = unpack8::<B>(&data[blk * B..]);
4810                for k in 0..8 {
4811                    chunk[k] = (u[k] - l) as i8 as u8;
4812                }
4813            }
4814        }
4815        let _ = mask;
4816        VBIT_SCRATCH.with(|scratch| {
4817            let mut buf = scratch.borrow_mut();
4818            buf.resize(cols, 0);
4819            match b {
4820                3 => fill::<3>(data, l, &mut buf),
4821                4 => vbit_fill4(data, &mut buf),
4822                5 => fill::<5>(data, l, &mut buf),
4823                6 => fill::<6>(data, l, &mut buf),
4824                _ => unreachable!(),
4825            }
4826            let mut dot = 0f32;
4827            for g in 0..ng {
4828                let so = (r * ng + g) * 2;
4829                let s = f16_to_f32(u16::from_le_bytes([
4830                    bytes[sc_off + so],
4831                    bytes[sc_off + so + 1],
4832                ]));
4833                let d = dot_i8_i8(
4834                    &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
4835                    &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
4836                ) as f32
4837                    * act.sx;
4838                dot += d * s;
4839            }
4840            for &(j, xv) in &act.outliers {
4841                let so = (r * ng + j / GROUP_SIZE) * 2;
4842                let s = f16_to_f32(u16::from_le_bytes([
4843                    bytes[sc_off + so],
4844                    bytes[sc_off + so + 1],
4845                ]));
4846                // xq is zeroed at outlier slots — add the exact term.
4847                dot += (buf[j] as i8) as f32 * s * xv;
4848            }
4849            dot
4850        })
4851    };
4852    for r in start..end {
4853        // SAFETY: disjoint row ranges per worker.
4854        unsafe { *out.at(r) = row_dot(r) };
4855    }
4856}
4857
4858/// Exact scalar vbit row range (same extraction, non-SDOT path).
4859#[allow(clippy::too_many_arguments)]
4860fn vbit_range_f32(
4861    bytes: &[u8],
4862    offsets: &[usize],
4863    x: &[f32],
4864    rows: usize,
4865    cols: usize,
4866    out: SendMut,
4867    start: usize,
4868    end: usize,
4869) {
4870    let ng = cols / GROUP_SIZE;
4871    let bits = &bytes[..rows];
4872    let sc_off = rows;
4873    // Per-bit-width specialized inner loops: the compiler unrolls the
4874    // constant shifts (the generic bit-buffer loop was branch-bound —
4875    // 5.6 vs 13.2 tok/s q4 on the 0.8B).
4876    #[inline(always)]
4877    fn dot_row<const B: usize>(
4878        data: &[u8],
4879        bytes: &[u8],
4880        sc_off: usize,
4881        r: usize,
4882        ng: usize,
4883        x: &[f32],
4884    ) -> f32 {
4885        let l = ((1i32 << (B - 1)) - 1) as f32;
4886        let gbytes = GROUP_SIZE * B / 8;
4887        let mut dot = 0f32;
4888        for g in 0..ng {
4889            let so = (r * ng + g) * 2;
4890            let s = f16_to_f32(u16::from_le_bytes([
4891                bytes[sc_off + so],
4892                bytes[sc_off + so + 1],
4893            ]));
4894            let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
4895            let gd0 = &data[g * gbytes..(g + 1) * gbytes];
4896            let mut gd = 0f32;
4897            for blk in 0..GROUP_SIZE / 8 {
4898                let u = unpack8::<B>(&gd0[blk * B..]);
4899                let xb = &xg[blk * 8..blk * 8 + 8];
4900                for k in 0..8 {
4901                    gd += (u[k] as f32 - l) * xb[k];
4902                }
4903            }
4904            dot += gd * s;
4905        }
4906        dot
4907    }
4908    for r in start..end {
4909        let data = &bytes[offsets[r]..offsets[r + 1]];
4910        let v = match bits[r] {
4911            3 => dot_row::<3>(data, bytes, sc_off, r, ng, x),
4912            4 => dot_row::<4>(data, bytes, sc_off, r, ng, x),
4913            5 => dot_row::<5>(data, bytes, sc_off, r, ng, x),
4914            6 => dot_row::<6>(data, bytes, sc_off, r, ng, x),
4915            8 => dot_row::<8>(data, bytes, sc_off, r, ng, x),
4916            b => unreachable!("vbit bit-width {b} (validated at load)"),
4917        };
4918        // SAFETY: disjoint row ranges per worker.
4919        unsafe { *out.at(r) = v };
4920    }
4921}
4922
4923/// Fused two-input vbit matvec: each row is unpacked from the mmap ONCE
4924/// and dotted against BOTH activations (MTP verify / pair prefill used
4925/// to run two full matvecs — double weight traffic and double unpack).
4926/// Per-input math is identical to `vbitmatvec` → same accuracy contract.
4927#[allow(clippy::too_many_arguments)]
4928fn vbitmatvec2(
4929    bytes: &[u8],
4930    offsets: &[usize],
4931    x1: &[f32],
4932    x2: &[f32],
4933    rows: usize,
4934    cols: usize,
4935    o1: &mut [f32],
4936    o2: &mut [f32],
4937    pool: Option<&Pool>,
4938) {
4939    debug_assert_eq!(o1.len(), rows);
4940    debug_assert_eq!(o2.len(), rows);
4941
4942    if a8w8_enabled() {
4943        let a1 = split_act(x1);
4944        let a2 = split_act(x2);
4945        let p1 = SendMut(o1.as_mut_ptr());
4946        let p2 = SendMut(o2.as_mut_ptr());
4947        let run = move |start: usize, end: usize| {
4948            vbit_range2_a8w8(
4949                bytes, offsets, x1, x2, &a1, &a2, rows, cols, p1, p2, start, end,
4950            )
4951        };
4952        dispatch_rows(pool, rows, &run);
4953        return;
4954    }
4955
4956    let p1 = SendMut(o1.as_mut_ptr());
4957    let p2 = SendMut(o2.as_mut_ptr());
4958    let run = move |start: usize, end: usize| {
4959        vbit_range2_f32(bytes, offsets, x1, x2, rows, cols, p1, p2, start, end)
4960    };
4961    dispatch_rows(pool, rows, &run);
4962}
4963
4964/// Two-input vbit row range via the A8W8 int8 path — kernel body of
4965/// `vbitmatvec2`, extracted for pair multi-matrix jobs (b=8 rows go
4966/// exact f32 for both lanes, bits streamed once).
4967#[allow(clippy::too_many_arguments)]
4968fn vbit_range2_a8w8(
4969    bytes: &[u8],
4970    offsets: &[usize],
4971    x1: &[f32],
4972    x2: &[f32],
4973    a1: &SplitAct,
4974    a2: &SplitAct,
4975    rows: usize,
4976    cols: usize,
4977    p1: SendMut,
4978    p2: SendMut,
4979    start: usize,
4980    end: usize,
4981) {
4982    let ng = cols / GROUP_SIZE;
4983    let bits = &bytes[..rows];
4984    let sc_off = rows;
4985    let row_dots = |r: usize| -> (f32, f32) {
4986        let b = bits[r] as usize;
4987        let l = (1i32 << (b - 1)) - 1;
4988        let data = &bytes[offsets[r]..offsets[r + 1]];
4989        if b == 8 {
4990            // u−L reaches 128 → does not fit i8; exact f32 path,
4991            // bits still streamed once for both lanes.
4992            let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
4993            let (mut d1, mut d2) = (0f32, 0f32);
4994            for g in 0..ng {
4995                let so = (r * ng + g) * 2;
4996                let sgf = f16_to_f32(u16::from_le_bytes([
4997                    bytes[sc_off + so],
4998                    bytes[sc_off + so + 1],
4999                ]));
5000                let (mut g1, mut g2) = (0f32, 0f32);
5001                for k in 0..GROUP_SIZE {
5002                    if nbits < 8 {
5003                        acc = (acc << 8) | data[idx] as u64;
5004                        idx += 1;
5005                        nbits += 8;
5006                    }
5007                    let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
5008                    nbits -= 8;
5009                    let w = (u - l) as f32;
5010                    g1 += w * x1[g * GROUP_SIZE + k];
5011                    g2 += w * x2[g * GROUP_SIZE + k];
5012                }
5013                d1 += g1 * sgf;
5014                d2 += g2 * sgf;
5015            }
5016            return (d1, d2);
5017        }
5018        thread_local! {
5019            static VBIT_SCRATCH2: std::cell::RefCell<Vec<u8>> =
5020                const { std::cell::RefCell::new(Vec::new()) };
5021        }
5022        #[inline(always)]
5023        fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
5024            for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
5025                let u = unpack8::<B>(&data[blk * B..]);
5026                for k in 0..8 {
5027                    chunk[k] = (u[k] - l) as i8 as u8;
5028                }
5029            }
5030        }
5031        VBIT_SCRATCH2.with(|scratch| {
5032            let mut buf = scratch.borrow_mut();
5033            buf.resize(cols, 0);
5034            match b {
5035                3 => fill::<3>(data, l, &mut buf),
5036                4 => vbit_fill4(data, &mut buf),
5037                5 => fill::<5>(data, l, &mut buf),
5038                6 => fill::<6>(data, l, &mut buf),
5039                _ => unreachable!(),
5040            }
5041            let (mut d1, mut d2) = (0f32, 0f32);
5042            for g in 0..ng {
5043                let so = (r * ng + g) * 2;
5044                let s = f16_to_f32(u16::from_le_bytes([
5045                    bytes[sc_off + so],
5046                    bytes[sc_off + so + 1],
5047                ]));
5048                let wg = &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
5049                let v1 = dot_i8_i8(wg, &a1.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a1.sx;
5050                let v2 = dot_i8_i8(wg, &a2.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a2.sx;
5051                d1 += v1 * s;
5052                d2 += v2 * s;
5053            }
5054            for &(j, xv) in &a1.outliers {
5055                let so = (r * ng + j / GROUP_SIZE) * 2;
5056                let s = f16_to_f32(u16::from_le_bytes([
5057                    bytes[sc_off + so],
5058                    bytes[sc_off + so + 1],
5059                ]));
5060                d1 += (buf[j] as i8) as f32 * s * xv;
5061            }
5062            for &(j, xv) in &a2.outliers {
5063                let so = (r * ng + j / GROUP_SIZE) * 2;
5064                let s = f16_to_f32(u16::from_le_bytes([
5065                    bytes[sc_off + so],
5066                    bytes[sc_off + so + 1],
5067                ]));
5068                d2 += (buf[j] as i8) as f32 * s * xv;
5069            }
5070            (d1, d2)
5071        })
5072    };
5073    for r in start..end {
5074        let (v1, v2) = row_dots(r);
5075        // SAFETY: disjoint row ranges per worker.
5076        unsafe {
5077            *p1.at(r) = v1;
5078            *p2.at(r) = v2;
5079        }
5080    }
5081}
5082
5083/// Two-input exact scalar vbit row range (same extraction) —
5084/// per-bit-width specialized, two accumulators per row; per-lane
5085/// accumulation order matches `vbitmatvec` exactly.
5086#[allow(clippy::too_many_arguments)]
5087fn vbit_range2_f32(
5088    bytes: &[u8],
5089    offsets: &[usize],
5090    x1: &[f32],
5091    x2: &[f32],
5092    rows: usize,
5093    cols: usize,
5094    p1: SendMut,
5095    p2: SendMut,
5096    start: usize,
5097    end: usize,
5098) {
5099    let ng = cols / GROUP_SIZE;
5100    let bits = &bytes[..rows];
5101    let sc_off = rows;
5102    #[inline(always)]
5103    #[allow(clippy::too_many_arguments)]
5104    fn dot_row2<const B: usize>(
5105        data: &[u8],
5106        bytes: &[u8],
5107        sc_off: usize,
5108        r: usize,
5109        ng: usize,
5110        x1: &[f32],
5111        x2: &[f32],
5112    ) -> (f32, f32) {
5113        let l = ((1i32 << (B - 1)) - 1) as f32;
5114        let gbytes = GROUP_SIZE * B / 8;
5115        let (mut d1, mut d2) = (0f32, 0f32);
5116        for g in 0..ng {
5117            let so = (r * ng + g) * 2;
5118            let s = f16_to_f32(u16::from_le_bytes([
5119                bytes[sc_off + so],
5120                bytes[sc_off + so + 1],
5121            ]));
5122            let x1g = &x1[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
5123            let x2g = &x2[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
5124            let gd0 = &data[g * gbytes..(g + 1) * gbytes];
5125            let (mut g1, mut g2) = (0f32, 0f32);
5126            for blk in 0..GROUP_SIZE / 8 {
5127                let u = unpack8::<B>(&gd0[blk * B..]);
5128                for k in 0..8 {
5129                    let w = u[k] as f32 - l;
5130                    g1 += w * x1g[blk * 8 + k];
5131                    g2 += w * x2g[blk * 8 + k];
5132                }
5133            }
5134            d1 += g1 * s;
5135            d2 += g2 * s;
5136        }
5137        (d1, d2)
5138    }
5139    for r in start..end {
5140        let data = &bytes[offsets[r]..offsets[r + 1]];
5141        let (v1, v2) = match bits[r] {
5142            3 => dot_row2::<3>(data, bytes, sc_off, r, ng, x1, x2),
5143            4 => dot_row2::<4>(data, bytes, sc_off, r, ng, x1, x2),
5144            5 => dot_row2::<5>(data, bytes, sc_off, r, ng, x1, x2),
5145            6 => dot_row2::<6>(data, bytes, sc_off, r, ng, x1, x2),
5146            8 => dot_row2::<8>(data, bytes, sc_off, r, ng, x1, x2),
5147            b => unreachable!("vbit bit-width {b} (validated at load)"),
5148        };
5149        // SAFETY: disjoint row ranges per worker.
5150        unsafe {
5151            *p1.at(r) = v1;
5152            *p2.at(r) = v2;
5153        }
5154    }
5155}
5156
5157// ───────────────────── q4_tiled kernels (§4.3) ─────────────────────
5158
5159/// One q4_tiled row dot on the A8W8 int8 path: per 32-group the tile
5160/// is ONE sequential read — [f16 scale][16B nibbles] — versus the two
5161/// distant streams of the split layout. Values/order identical to the
5162/// split kernels.
5163#[inline]
5164#[allow(unreachable_code)]
5165fn dot_q4t_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
5166    #[cfg(target_arch = "aarch64")]
5167    unsafe {
5168        return dot_q4t_row_sdot(bytes, r, gpr, xq);
5169    }
5170    #[cfg(target_arch = "x86_64")]
5171    unsafe {
5172        if vnni_tiles_enabled() {
5173            return dot_q4t_row_vnni(bytes, r, gpr, xq);
5174        }
5175        return dot_q4t_row_avx2(bytes, r, gpr, xq);
5176    }
5177    let mut acc = 0f32;
5178    for gi in 0..gpr {
5179        let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
5180        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
5181        let mut d = 0i32;
5182        for (k, &b) in tile[2..].iter().enumerate() {
5183            d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
5184                + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
5185        }
5186        acc += d as f32 * s;
5187    }
5188    acc
5189}
5190
5191#[cfg(target_arch = "aarch64")]
5192#[target_feature(enable = "neon,dotprod")]
5193unsafe fn dot_q4t_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
5194    // SAFETY: callers uphold slice-length contracts (18B tile per group,
5195    // xq.len() == gpr·GROUP_SIZE).
5196    unsafe {
5197        use core::arch::aarch64::*;
5198        use core::arch::asm;
5199        let lomask = vdupq_n_u8(0x0F);
5200        let eight = vdupq_n_s8(8);
5201        let mut acc = 0f32;
5202        for gi in 0..gpr {
5203            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
5204            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5205            let b = vld1q_u8(t.add(2));
5206            let lo = vandq_u8(b, lomask);
5207            let hi = vshrq_n_u8::<4>(b);
5208            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
5209            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
5210            let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
5211            let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
5212            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5213            asm!(
5214                "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
5215                "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
5216                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5217                e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
5218                options(pure, nomem, nostack),
5219            );
5220            acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
5221        }
5222        acc
5223    }
5224}
5225
5226#[cfg(target_arch = "x86_64")]
5227#[target_feature(enable = "avx2")]
5228unsafe fn dot_q4t_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
5229    // SAFETY: see dot_q4t_row_sdot.
5230    unsafe {
5231        use core::arch::x86_64::*;
5232        let lomask = _mm_set1_epi8(0x0F);
5233        let eight = _mm256_set1_epi8(8);
5234        let ones = _mm256_set1_epi16(1);
5235        let mut acc = 0f32;
5236        for gi in 0..gpr {
5237            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
5238            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5239            let b = _mm_loadu_si128(t.add(2) as *const __m128i);
5240            let lo = _mm_and_si128(b, lomask);
5241            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
5242            let w = _mm256_sub_epi8(
5243                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5244                eight,
5245            );
5246            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5247            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
5248            let d = _mm256_madd_epi16(p16, ones);
5249            let hi128 = _mm256_extracti128_si256::<1>(d);
5250            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
5251            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
5252            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
5253            acc += _mm_cvtsi128_si32(s32) as f32 * s;
5254        }
5255        acc
5256    }
5257}
5258
5259/// VNNI twin of `dot_q4t_row_avx2`: same unpack, `vpdpbusd` replaces
5260/// the maddubs+madd pair (see `dpbusd_hsum` — sums are bit-identical).
5261/// 256-bit VL encoding, so the VEX `vpsignb` stays usable.
5262#[cfg(target_arch = "x86_64")]
5263#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
5264unsafe fn dot_q4t_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
5265    // SAFETY: see dot_q4t_row_sdot.
5266    unsafe {
5267        use core::arch::x86_64::*;
5268        let lomask = _mm_set1_epi8(0x0F);
5269        let eight = _mm256_set1_epi8(8);
5270        let mut acc = 0f32;
5271        for gi in 0..gpr {
5272            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
5273            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5274            let b = _mm_loadu_si128(t.add(2) as *const __m128i);
5275            let lo = _mm_and_si128(b, lomask);
5276            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
5277            let w = _mm256_sub_epi8(
5278                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5279                eight,
5280            );
5281            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5282            let d = dpbusd_hsum(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
5283            acc += d as f32 * s;
5284        }
5285        acc
5286    }
5287}
5288
5289/// One q4_tiled row against FOUR activation streams: the nibble unpack
5290/// and abs() happen once per group instead of once per (group,
5291/// activation) — the unpack is the dominant per-element cost of the
5292/// tiled format (roadmap P0 portable blocking, q4t leg).
5293#[cfg(target_arch = "x86_64")]
5294// `fma` is NOT implied by `avx2`: without it LLVM lowers _mm256_fmadd_ps
5295// to a libm call per lane — measured 2x slower than the reduction this
5296// kernel replaces. The runtime gate (`avx2_enabled`) already requires
5297// both features, so declaring it here is safe.
5298#[target_feature(enable = "avx2,fma")]
5299unsafe fn dot_q4t_row_1x4_avx2(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
5300    // SAFETY: callers uphold the 18B-tile and xq-length contracts.
5301    unsafe {
5302        use core::arch::x86_64::*;
5303        let lomask = _mm_set1_epi8(0x0F);
5304        let eight = _mm256_set1_epi8(8);
5305        let ones = _mm256_set1_epi16(1);
5306        // One f32 accumulator VECTOR per activation, reduced once at the
5307        // end. Folding each group's i32 lanes to a scalar inside the loop
5308        // costs an extracti128 + three shift/add + a movd — a cross-lane
5309        // dependency chain per (group, activation), 288 of them per row at
5310        // cols=2304. The per-group scale is what forces a float
5311        // accumulator; it does not force a horizontal sum.
5312        //
5313        // The four accumulators are NAMED, not an array: as `[__m256; 4]`
5314        // indexed by a loop variable LLVM keeps them in memory and every
5315        // group pays four 32-byte loads and stores. That alone made this
5316        // kernel 2x SLOWER than the per-group reduction it replaces
5317        // (measured on the EPYC box: 150 s vs 71 s for two 256² steps).
5318        let mut f0 = _mm256_setzero_ps();
5319        let mut f1 = _mm256_setzero_ps();
5320        let mut f2 = _mm256_setzero_ps();
5321        let mut f3 = _mm256_setzero_ps();
5322        for gi in 0..gpr {
5323            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
5324            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5325            let sv = _mm256_set1_ps(s);
5326            let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
5327            let lo = _mm_and_si128(bb, lomask);
5328            let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
5329            let w = _mm256_sub_epi8(
5330                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5331                eight,
5332            );
5333            let aw = _mm256_abs_epi8(w);
5334            let off = gi * GROUP_SIZE;
5335            let dot = |xq: &[i8]| {
5336                let x = _mm256_loadu_si256(xq.as_ptr().add(off) as *const __m256i);
5337                let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
5338                _mm256_cvtepi32_ps(_mm256_madd_epi16(p16, ones))
5339            };
5340            f0 = _mm256_fmadd_ps(dot(xs[0]), sv, f0);
5341            f1 = _mm256_fmadd_ps(dot(xs[1]), sv, f1);
5342            f2 = _mm256_fmadd_ps(dot(xs[2]), sv, f2);
5343            f3 = _mm256_fmadd_ps(dot(xs[3]), sv, f3);
5344        }
5345        [
5346            hsum256_ps(f0),
5347            hsum256_ps(f1),
5348            hsum256_ps(f2),
5349            hsum256_ps(f3),
5350        ]
5351    }
5352}
5353
5354/// Horizontal sum of eight f32 lanes — the one cross-lane reduction the
5355/// blocked kernels pay, once per row instead of once per group.
5356#[cfg(target_arch = "x86_64")]
5357#[target_feature(enable = "avx2")]
5358#[inline]
5359unsafe fn hsum256_ps(v: core::arch::x86_64::__m256) -> f32 {
5360    // SAFETY: pure register arithmetic on the caller's vector.
5361    unsafe {
5362        use core::arch::x86_64::*;
5363        let hi = _mm256_extractf128_ps::<1>(v);
5364        let s = _mm_add_ps(_mm256_castps256_ps128(v), hi);
5365        let s = _mm_add_ps(s, _mm_movehl_ps(s, s));
5366        let s = _mm_add_ss(s, _mm_shuffle_ps::<0x55>(s, s));
5367        _mm_cvtss_f32(s)
5368    }
5369}
5370
5371/// VNNI twin of `dot_q4t_row_1x4_avx2` (see `dpbusd_hsum`).
5372#[cfg(target_arch = "x86_64")]
5373#[target_feature(enable = "avx2,fma,avx512f,avx512bw,avx512vl,avx512vnni")]
5374unsafe fn dot_q4t_row_1x4_vnni(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
5375    // SAFETY: callers uphold the 18B-tile and xq-length contracts.
5376    unsafe {
5377        use core::arch::x86_64::*;
5378        let lomask = _mm_set1_epi8(0x0F);
5379        let eight = _mm256_set1_epi8(8);
5380        // Same shape as the AVX2 twin: accumulate in f32 vectors and pay
5381        // one cross-lane reduction per row, not per (group, activation).
5382        let mut f0 = _mm256_setzero_ps();
5383        let mut f1 = _mm256_setzero_ps();
5384        let mut f2 = _mm256_setzero_ps();
5385        let mut f3 = _mm256_setzero_ps();
5386        for gi in 0..gpr {
5387            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
5388            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5389            let sv = _mm256_set1_ps(s);
5390            let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
5391            let lo = _mm_and_si128(bb, lomask);
5392            let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
5393            let w = _mm256_sub_epi8(
5394                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5395                eight,
5396            );
5397            let aw = _mm256_abs_epi8(w);
5398            let off = gi * GROUP_SIZE;
5399            let dot = |xq: &[i8]| {
5400                let x = _mm256_loadu_si256(xq.as_ptr().add(off) as *const __m256i);
5401                _mm256_cvtepi32_ps(_mm256_dpbusd_epi32(
5402                    _mm256_setzero_si256(),
5403                    aw,
5404                    _mm256_sign_epi8(x, w),
5405                ))
5406            };
5407            f0 = _mm256_fmadd_ps(dot(xs[0]), sv, f0);
5408            f1 = _mm256_fmadd_ps(dot(xs[1]), sv, f1);
5409            f2 = _mm256_fmadd_ps(dot(xs[2]), sv, f2);
5410            f3 = _mm256_fmadd_ps(dot(xs[3]), sv, f3);
5411        }
5412        let acc = [
5413            hsum256_ps(f0),
5414            hsum256_ps(f1),
5415            hsum256_ps(f2),
5416            hsum256_ps(f3),
5417        ];
5418        acc
5419    }
5420}
5421
5422/// ARM twin of `dot_q4t_row_1x4_avx2`: one nibble unpack per group
5423/// serves FOUR activation streams. Per stream the group order and f32
5424/// accumulation match `dot_q4t_row_sdot` exactly — batch == matvec
5425/// bit-for-bit.
5426#[cfg(target_arch = "aarch64")]
5427#[target_feature(enable = "neon,dotprod")]
5428unsafe fn dot_q4t_row_1x4_sdot(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
5429    // SAFETY: callers uphold the 18B-tile and xq-length contracts.
5430    unsafe {
5431        use core::arch::aarch64::*;
5432        use core::arch::asm;
5433        let lomask = vdupq_n_u8(0x0F);
5434        let eight = vdupq_n_s8(8);
5435        let mut acc = [0f32; 4];
5436        for gi in 0..gpr {
5437            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
5438            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5439            let b = vld1q_u8(t.add(2));
5440            let lo = vandq_u8(b, lomask);
5441            let hi = vshrq_n_u8::<4>(b);
5442            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
5443            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
5444            for (k, xq) in xs.iter().enumerate() {
5445                let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
5446                let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
5447                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5448                asm!(
5449                    "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
5450                    "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
5451                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5452                    e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
5453                    options(pure, nomem, nostack),
5454                );
5455                acc[k] += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
5456            }
5457        }
5458        acc
5459    }
5460}
5461
5462/// Exact-term correction for A8W8 outliers on a tiled row.
5463#[inline]
5464fn q4t_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
5465    let gi = j / GROUP_SIZE;
5466    let k = j % GROUP_SIZE;
5467    let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
5468    let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
5469    let byte = tile[2 + k / 2];
5470    let nib = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
5471    ((nib as i32 - 8) as f32, s)
5472}
5473
5474/// Exact scalar q4_tiled row (CMF_SDOT=0 contract) — same pairwise
5475/// accumulation shape as `q4_range_f32`.
5476#[inline]
5477fn q4t_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
5478    let mut acc = 0f32;
5479    for gi in 0..gpr {
5480        let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
5481        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
5482        let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
5483        let mut ga = 0f32;
5484        for (k, &b) in tile[2..].iter().enumerate() {
5485            ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
5486                + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
5487        }
5488        acc += ga * s;
5489    }
5490    acc
5491}
5492
5493/// Split view of a `q4tp` payload. The three planes are resolved once per
5494/// matvec instead of per row — `q4tp_sections` is cheap, but doing it inside
5495/// the row loop would put a division on the hot path for nothing.
5496struct Q4tpView<'a> {
5497    nib: &'a [u8],
5498    params: &'a [u8],
5499    codes: &'a [u8],
5500    stride: usize,
5501    /// q2tp reads the ladder with rung 0 = exact zero.
5502    zero_rung: bool,
5503}
5504
5505impl<'a> Q4tpView<'a> {
5506    fn new(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
5507        let (params_off, codes_off, stride) = q4tp_sections(rows, cols);
5508        Self {
5509            nib: &bytes[..params_off],
5510            params: &bytes[params_off..codes_off],
5511            codes: &bytes[codes_off..],
5512            stride,
5513            zero_rung: false,
5514        }
5515    }
5516
5517    /// The q2tp view: identical params/codes planes, 8 B weight chunks.
5518    fn new_q2(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
5519        let (params_off, codes_off, stride) = q2tp_sections(rows, cols);
5520        Self {
5521            nib: &bytes[..params_off],
5522            params: &bytes[params_off..codes_off],
5523            codes: &bytes[codes_off..],
5524            stride,
5525            zero_rung: true,
5526        }
5527    }
5528
5529    /// Expand row `r`'s per-tile scales into `out` (length `gpr`).
5530    ///
5531    /// Doing this once per row — rather than decoding a 5-bit code inside the
5532    /// tile loop — is what makes the format free at runtime. Random access to
5533    /// a packed 5-bit field costs a division, two bounds checks and a branch;
5534    /// the tile's actual work is two `sdot`s, so per-tile decoding dominated
5535    /// the kernel and cost 5x (measured: 1.4 vs 6.9 tok/s on Nanbeige-3B).
5536    /// Walking the plane sequentially with a bit accumulator is ~3 ops.
5537    /// Eight 5-bit codes are exactly five bytes, so a whole group of
5538    /// eight decodes from one little-endian word at fixed shifts. The
5539    /// bit-accumulator this replaces carried a data-dependent `while
5540    /// have < 5` refill whose branch sat in the innermost loop of every
5541    /// q4tp row; a decode profile put this function above the dot
5542    /// products it feeds. Same bitstream, same codes — just no branch
5543    /// and eight independent extractions.
5544    #[inline]
5545    fn scales_into(&self, r: usize, gpr: usize, out: &mut [f32]) {
5546        let tab = if self.zero_rung {
5547            q2tp_ladder(self.params, r)
5548        } else {
5549            q4tp_ladder(self.params, r)
5550        };
5551        let codes = &self.codes[r * self.stride..(r + 1) * self.stride];
5552        let out = &mut out[..gpr];
5553        let mut chunks = out.chunks_exact_mut(8);
5554        let mut ci = 0usize;
5555        for c in &mut chunks {
5556            let w = u64::from(codes[ci])
5557                | u64::from(codes[ci + 1]) << 8
5558                | u64::from(codes[ci + 2]) << 16
5559                | u64::from(codes[ci + 3]) << 24
5560                | u64::from(codes[ci + 4]) << 32;
5561            for (k, o) in c.iter_mut().enumerate() {
5562                *o = tab[((w >> (5 * k)) & 31) as usize];
5563            }
5564            ci += 5;
5565        }
5566        // Fewer than eight codes left: the shared total accessor, which
5567        // tolerates a 5-bit field whose spill byte is past the stride.
5568        let tail = &codes[ci..];
5569        for (k, o) in chunks.into_remainder().iter_mut().enumerate() {
5570            *o = tab[q4tp_code(tail, k)];
5571        }
5572    }
5573}
5574
5575#[inline]
5576fn dot_q4tp_row_i8(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
5577    #[cfg(target_arch = "aarch64")]
5578    unsafe {
5579        return dot_q4tp_row_sdot(nib, r, gpr, xq, scales);
5580    }
5581    #[cfg(target_arch = "x86_64")]
5582    unsafe {
5583        if vnni_tiles_enabled() {
5584            return dot_q4tp_row_vnni(nib, r, gpr, xq, scales);
5585        }
5586        return dot_q4tp_row_avx2(nib, r, gpr, xq, scales);
5587    }
5588    #[allow(unreachable_code)]
5589    {
5590        let mut acc = 0f32;
5591        for gi in 0..gpr {
5592            let tile = &nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
5593            let s = scales[gi];
5594            let mut d = 0i32;
5595            for (k, &b) in tile.iter().enumerate() {
5596                d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
5597                    + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
5598            }
5599            acc += d as f32 * s;
5600        }
5601        acc
5602    }
5603}
5604
5605/// q4tp twin of `dot_q4t_row_sdot`: identical nibble math, but the tile
5606/// stride is 16 B (no inline scale) and the scale is a ladder lookup.
5607#[cfg(target_arch = "aarch64")]
5608#[target_feature(enable = "neon,dotprod")]
5609unsafe fn dot_q4tp_row_sdot(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
5610    // SAFETY: callers uphold slice-length contracts (16B tile per group,
5611    // xq.len() == gpr·GROUP_SIZE, codes covering gpr 5-bit fields).
5612    unsafe {
5613        use core::arch::aarch64::*;
5614        use core::arch::asm;
5615        let lomask = vdupq_n_u8(0x0F);
5616        let eight = vdupq_n_s8(8);
5617        let mut acc = 0f32;
5618        for gi in 0..gpr {
5619            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5620            let s = *scales.get_unchecked(gi);
5621            let b = vld1q_u8(t);
5622            let lo = vandq_u8(b, lomask);
5623            let hi = vshrq_n_u8::<4>(b);
5624            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
5625            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
5626            let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
5627            let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
5628            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5629            asm!(
5630                "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
5631                "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
5632                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5633                e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
5634                options(pure, nomem, nostack),
5635            );
5636            acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
5637        }
5638        acc
5639    }
5640}
5641
5642#[cfg(target_arch = "x86_64")]
5643#[target_feature(enable = "avx2")]
5644unsafe fn dot_q4tp_row_avx2(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
5645    // SAFETY: see dot_q4tp_row_sdot.
5646    unsafe {
5647        use core::arch::x86_64::*;
5648        let lomask = _mm_set1_epi8(0x0F);
5649        let eight = _mm256_set1_epi8(8);
5650        let ones = _mm256_set1_epi16(1);
5651        let mut acc = 0f32;
5652        for gi in 0..gpr {
5653            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5654            let s = *scales.get_unchecked(gi);
5655            let b = _mm_loadu_si128(t as *const __m128i);
5656            let lo = _mm_and_si128(b, lomask);
5657            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
5658            let w = _mm256_sub_epi8(
5659                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5660                eight,
5661            );
5662            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5663            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
5664            let d = _mm256_madd_epi16(p16, ones);
5665            let hi128 = _mm256_extracti128_si256::<1>(d);
5666            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
5667            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
5668            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
5669            acc += _mm_cvtsi128_si32(s32) as f32 * s;
5670        }
5671        acc
5672    }
5673}
5674
5675/// VNNI twin of `dot_q4tp_row_avx2` (see `dot_q4t_row_vnni` for why the
5676/// 256-bit VL encoding is the one to use here).
5677#[cfg(target_arch = "x86_64")]
5678#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
5679unsafe fn dot_q4tp_row_vnni(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
5680    // SAFETY: see dot_q4tp_row_sdot.
5681    unsafe {
5682        use core::arch::x86_64::*;
5683        let lomask = _mm_set1_epi8(0x0F);
5684        let eight = _mm256_set1_epi8(8);
5685        let mut acc = 0f32;
5686        for gi in 0..gpr {
5687            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5688            let s = *scales.get_unchecked(gi);
5689            let b = _mm_loadu_si128(t as *const __m128i);
5690            let lo = _mm_and_si128(b, lomask);
5691            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
5692            let w = _mm256_sub_epi8(
5693                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5694                eight,
5695            );
5696            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5697            acc += dpbusd_hsum(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w)) as f32 * s;
5698        }
5699        acc
5700    }
5701}
5702
5703/// Exact scalar q4tp row — the `CMF_SDOT=0` contract, same pairwise
5704/// accumulation shape as `q4t_row_exact`.
5705#[inline]
5706fn q4tp_row_exact(nib: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
5707    #[cfg(target_arch = "x86_64")]
5708    if avx2_enabled() {
5709        // Keep the scalar pair/group reduction order, not a vector sum.
5710        return unsafe { q4tp_row_float_avx2(nib, r, gpr, x, scales) };
5711    }
5712    q4tp_row_float_scalar(nib, r, gpr, x, scales)
5713}
5714
5715#[inline]
5716fn q4tp_row_float_scalar(nib: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
5717    let mut acc = 0f32;
5718    for gi in 0..gpr {
5719        let tile = &nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
5720        let s = scales[gi];
5721        let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
5722        let mut ga = 0f32;
5723        for (k, &b) in tile.iter().enumerate() {
5724            ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
5725                + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
5726        }
5727        acc += ga * s;
5728    }
5729    acc
5730}
5731
5732/// Vectorize unpack, conversion and multiplication, but preserve every
5733/// pair addition and the scalar accumulation order. No activation rounding
5734/// or FMA: bit-identical to the float scalar row, including its group scale.
5735#[cfg(target_arch = "x86_64")]
5736#[target_feature(enable = "avx2")]
5737unsafe fn q4tp_row_float_avx2(nib: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
5738    // SAFETY: caller checks AVX2 and provides the same complete 32-element
5739    // groups as the scalar row. Loads/stores are explicitly unaligned.
5740    unsafe {
5741        use core::arch::x86_64::*;
5742        let mask = _mm_set1_epi8(15);
5743        let eight = _mm_set1_epi8(8);
5744        let order = _mm256_setr_epi32(0, 1, 4, 5, 2, 3, 6, 7);
5745        let mut acc = 0.0f32;
5746        for gi in 0..gpr {
5747            let packed = _mm_loadu_si128(nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB).cast());
5748            let lo = _mm_and_si128(packed, mask);
5749            let hi = _mm_and_si128(_mm_srli_epi16::<4>(packed), mask);
5750            let w0 = _mm_sub_epi8(_mm_unpacklo_epi8(lo, hi), eight);
5751            let w1 = _mm_sub_epi8(_mm_unpackhi_epi8(lo, hi), eight);
5752            let xp = x.as_ptr().add(gi * GROUP_SIZE);
5753            let a = _mm256_mul_ps(
5754                _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(w0)),
5755                _mm256_loadu_ps(xp),
5756            );
5757            let b = _mm256_mul_ps(
5758                _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(w0))),
5759                _mm256_loadu_ps(xp.add(8)),
5760            );
5761            let c = _mm256_mul_ps(
5762                _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(w1)),
5763                _mm256_loadu_ps(xp.add(16)),
5764            );
5765            let d = _mm256_mul_ps(
5766                _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(w1))),
5767                _mm256_loadu_ps(xp.add(24)),
5768            );
5769            let mut pairs = [0.0f32; 16];
5770            _mm256_storeu_ps(
5771                pairs.as_mut_ptr(),
5772                _mm256_permutevar8x32_ps(_mm256_hadd_ps(a, b), order),
5773            );
5774            _mm256_storeu_ps(
5775                pairs.as_mut_ptr().add(8),
5776                _mm256_permutevar8x32_ps(_mm256_hadd_ps(c, d), order),
5777            );
5778            let mut ga = 0.0f32;
5779            for v in pairs {
5780                ga += v;
5781            }
5782            acc += ga * scales[gi];
5783        }
5784        acc
5785    }
5786}
5787
5788/// Single weight of a q4tp tensor — the a8w8 outlier path, which restores
5789/// activation outliers at full precision after the int8 pass.
5790#[inline]
5791fn q4tp_outlier(nib: &[u8], r: usize, gpr: usize, j: usize, scales: &[f32]) -> (f32, f32) {
5792    let (gi, k) = (j / GROUP_SIZE, j % GROUP_SIZE);
5793    let byte = nib[(r * gpr + gi) * Q4TP_NIB + k / 2];
5794    let n = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
5795    ((n as i32 - 8) as f32, scales[gi])
5796}
5797
5798/// Fused q4tp matvec (dispatch mirrors `q4t_matvec`).
5799fn q4tp_matvec(
5800    bytes: &[u8],
5801    x: &[f32],
5802    rows: usize,
5803    cols: usize,
5804    out: &mut [f32],
5805    pool: Option<&Pool>,
5806) {
5807    debug_assert_eq!(out.len(), rows);
5808    let gpr = cols / GROUP_SIZE;
5809    let v = Q4tpView::new(bytes, rows, cols);
5810    let out_addr = SendMut(out.as_mut_ptr());
5811    if a8w8_enabled() {
5812        let act = split_act(x);
5813        let run = |start: usize, end: usize| {
5814            // One scratch row of scales per worker — borrowed, not minted.
5815            with_krow(gpr, |sc| {
5816                for r in start..end {
5817                    v.scales_into(r, gpr, sc);
5818                    let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, sc) * act.sx;
5819                    for &(j, xv) in &act.outliers {
5820                        let (w, s) = q4tp_outlier(v.nib, r, gpr, j, sc);
5821                        acc += w * s * xv;
5822                    }
5823                    // SAFETY: disjoint row ranges per worker.
5824                    unsafe { *out_addr.at(r) = acc };
5825                }
5826            })
5827        };
5828        dispatch_rows(pool, rows, &run);
5829        return;
5830    }
5831    let run = |start: usize, end: usize| {
5832        with_krow(gpr, |sc| {
5833            for r in start..end {
5834                v.scales_into(r, gpr, sc);
5835                // SAFETY: disjoint row ranges per worker.
5836                unsafe { *out_addr.at(r) = q4tp_row_exact(v.nib, r, gpr, x, sc) };
5837            }
5838        })
5839    };
5840    dispatch_rows(pool, rows, &run);
5841}
5842
5843/// Fused two-input q4tp matvec — the SwiGLU gate/up pair. Weights and the
5844/// row ladder are read once and spent on both activation streams.
5845#[allow(clippy::too_many_arguments)]
5846fn q4tp_matvec2(
5847    bytes: &[u8],
5848    x1: &[f32],
5849    x2: &[f32],
5850    rows: usize,
5851    cols: usize,
5852    o1: &mut [f32],
5853    o2: &mut [f32],
5854    pool: Option<&Pool>,
5855) {
5856    let gpr = cols / GROUP_SIZE;
5857    let v = Q4tpView::new(bytes, rows, cols);
5858    let (p1, p2) = (SendMut(o1.as_mut_ptr()), SendMut(o2.as_mut_ptr()));
5859    let run = |start: usize, end: usize| {
5860        let mut sc = vec![0f32; gpr];
5861        for r in start..end {
5862            v.scales_into(r, gpr, &mut sc);
5863            // SAFETY: disjoint row ranges per worker.
5864            unsafe {
5865                *p1.at(r) = q4tp_row_exact(v.nib, r, gpr, x1, &sc);
5866                *p2.at(r) = q4tp_row_exact(v.nib, r, gpr, x2, &sc);
5867            }
5868        }
5869    };
5870    dispatch_rows(pool, rows, &run);
5871}
5872
5873/// One q2tp outlier weight at column `j` of row `r`: the 2-bit code and
5874/// its group scale, mirrored on `q4tp_outlier`.
5875#[inline]
5876fn q2tp_outlier(chunks: &[u8], r: usize, gpr: usize, j: usize, scales: &[f32]) -> (f32, f32) {
5877    let (gi, k) = (j / GROUP_SIZE, j % GROUP_SIZE);
5878    let byte = chunks[(r * gpr + gi) * Q2TP_CHUNK + k / 4];
5879    let c = (byte >> (2 * (k % 4))) & 3;
5880    (c as f32 - 1.5, scales[gi])
5881}
5882
5883#[cfg(target_arch = "x86_64")]
5884const Q2TP_DECODE_U32: [u32; 256] = {
5885    let mut tab = [0u32; 256];
5886    let mut b = 0usize;
5887    while b < 256 {
5888        tab[b] = ((b as u32) & 3)
5889            | ((((b as u32) >> 2) & 3) << 8)
5890            | ((((b as u32) >> 4) & 3) << 16)
5891            | ((((b as u32) >> 6) & 3) << 24);
5892        b += 1;
5893    }
5894    tab
5895};
5896
5897/// Eight packed q2tp bytes against 32 signed activation bytes. `maddubs`
5898/// exactly computes unsigned 2-bit code × signed i8; its pair sums cannot
5899/// saturate (2 × 3 × 127 < i16::MAX), and the second madd widens to i32.
5900#[cfg(target_arch = "x86_64")]
5901#[target_feature(enable = "avx2")]
5902unsafe fn q2tp_code_dot_avx2(ch: &[u8], x: &[i8]) -> i32 {
5903    use core::arch::x86_64::*;
5904    debug_assert!(ch.len() >= Q2TP_CHUNK && x.len() >= GROUP_SIZE);
5905    let codes = _mm256_setr_epi32(
5906        Q2TP_DECODE_U32[ch[0] as usize] as i32,
5907        Q2TP_DECODE_U32[ch[1] as usize] as i32,
5908        Q2TP_DECODE_U32[ch[2] as usize] as i32,
5909        Q2TP_DECODE_U32[ch[3] as usize] as i32,
5910        Q2TP_DECODE_U32[ch[4] as usize] as i32,
5911        Q2TP_DECODE_U32[ch[5] as usize] as i32,
5912        Q2TP_DECODE_U32[ch[6] as usize] as i32,
5913        Q2TP_DECODE_U32[ch[7] as usize] as i32,
5914    );
5915    let xv = unsafe { _mm256_loadu_si256(x.as_ptr().cast()) };
5916    let pair = _mm256_maddubs_epi16(codes, xv);
5917    let quad = _mm256_madd_epi16(pair, _mm256_set1_epi16(1));
5918    let sum128 = _mm_add_epi32(
5919        _mm256_castsi256_si128(quad),
5920        _mm256_extracti128_si256(quad, 1),
5921    );
5922    let sum64 = _mm_hadd_epi32(sum128, sum128);
5923    _mm_cvtsi128_si32(_mm_hadd_epi32(sum64, sum64))
5924}
5925
5926#[cfg(target_arch = "x86_64")]
5927#[inline]
5928fn q2tp_avx2_enabled() -> bool {
5929    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
5930    *ON.get_or_init(|| std::arch::is_x86_feature_detected!("avx2"))
5931}
5932
5933/// Integer dot of one q2tp row against pre-quantized activations:
5934/// Σ_g s_g · (Σ c·xq − 1.5·Σ xq). The half-integer grid (c − 1.5)
5935/// becomes exact integer math through the group sums — the same trick
5936/// every a8w8 kernel in this file rides. The codes decode into a
5937/// 32-byte scratch in natural order and the dot itself is the shared
5938/// SDOT primitive; elsewhere a scalar integer loop.
5939#[inline]
5940fn dot_q2tp_row_i8(
5941    chunks: &[u8],
5942    r: usize,
5943    gpr: usize,
5944    xq: &[i8],
5945    gsum: &[i32],
5946    scales: &[f32],
5947) -> f32 {
5948    let mut acc = 0f32;
5949    let base = r * gpr * Q2TP_CHUNK;
5950    #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
5951    let mut codes = [0i8; GROUP_SIZE];
5952    #[cfg(target_arch = "x86_64")]
5953    // Cache the process-wide AVX2 decision; probing it for every 32-weight
5954    // group adds a branch to the hottest Q2TP row loop while retaining the
5955    // established table-load AVX2 arithmetic (and without coupling this path
5956    // to the separate FMA-gated A8W8 switch).
5957    let avx2 = q2tp_avx2_enabled();
5958    for gi in 0..gpr {
5959        let ch = &chunks[base + gi * Q2TP_CHUNK..base + (gi + 1) * Q2TP_CHUNK];
5960        let xg = &xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
5961        #[cfg(target_arch = "aarch64")]
5962        // NEON: the byte's four 2-bit fields land in four lane vectors
5963        // (shift+mask), vld4 de-interleaves xq to match (xj[k] =
5964        // xq[4k+j]), widening MACs accumulate exactly in i32. A scalar
5965        // decode here cost as much as the dot it fed — the profile put
5966        // it at the top of the whole W2 decode.
5967        let dot = unsafe {
5968            use core::arch::aarch64::*;
5969            let b = vld1_u8(ch.as_ptr());
5970            let three = vdup_n_u8(3);
5971            let c0 = vreinterpret_s8_u8(vand_u8(b, three));
5972            let c1 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 2), three));
5973            let c2 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 4), three));
5974            let c3 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 6), three));
5975            let x4 = vld4_s8(xg.as_ptr());
5976            let mut acc4 = vdupq_n_s32(0);
5977            acc4 = vpadalq_s16(acc4, vmull_s8(c0, x4.0));
5978            acc4 = vpadalq_s16(acc4, vmull_s8(c1, x4.1));
5979            acc4 = vpadalq_s16(acc4, vmull_s8(c2, x4.2));
5980            acc4 = vpadalq_s16(acc4, vmull_s8(c3, x4.3));
5981            vaddvq_s32(acc4)
5982        };
5983        #[cfg(target_arch = "x86_64")]
5984        let dot: i32 = if avx2 {
5985            // SAFETY: the runtime feature check gates the target-feature body;
5986            // the group slices above are exactly 8 and 32 bytes long.
5987            unsafe { q2tp_code_dot_avx2(ch, xg) }
5988        } else {
5989            ch.iter()
5990                .enumerate()
5991                .map(|(k, &b)| {
5992                    ((b & 3) as i32) * xg[k * 4] as i32
5993                        + (((b >> 2) & 3) as i32) * xg[k * 4 + 1] as i32
5994                        + (((b >> 4) & 3) as i32) * xg[k * 4 + 2] as i32
5995                        + (((b >> 6) & 3) as i32) * xg[k * 4 + 3] as i32
5996                })
5997                .sum()
5998        };
5999        #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
6000        let dot: i32 = {
6001            for (k, &b) in ch.iter().enumerate() {
6002                codes[k * 4] = (b & 3) as i8;
6003                codes[k * 4 + 1] = ((b >> 2) & 3) as i8;
6004                codes[k * 4 + 2] = ((b >> 4) & 3) as i8;
6005                codes[k * 4 + 3] = ((b >> 6) & 3) as i8;
6006            }
6007            codes
6008                .iter()
6009                .zip(xg)
6010                .map(|(&c, &x)| c as i32 * x as i32)
6011                .sum()
6012        };
6013        acc += scales[gi] * (dot as f32 - 1.5 * gsum[gi] as f32);
6014    }
6015    acc
6016}
6017
6018/// Exact f32 dot of one q2tp row: 2-bit fields LSB-first, (c − 1.5)·s.
6019/// Scalar on purpose — the 2-bit class targets the GPU graph; the CPU
6020/// path exists for parity gates and small-machine fallback.
6021fn q2tp_row_exact(chunks: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
6022    q2tp_row_exact_center(chunks, r, gpr, x, scales, 1.5)
6023}
6024
6025/// Fused Prism affine row: the derived correction is applied inside the
6026/// decoded code, avoiding a second accumulated dot and avoiding cancellation
6027/// between `B=(c-1.5)s` and `+.5s` for long 5120/17408 rows.
6028#[inline]
6029fn q2tp_affine_row_exact(chunks: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
6030    q2tp_row_exact_center(chunks, r, gpr, x, scales, 1.0)
6031}
6032
6033#[inline]
6034fn q2tp_row_exact_center(
6035    chunks: &[u8],
6036    r: usize,
6037    gpr: usize,
6038    x: &[f32],
6039    scales: &[f32],
6040    center: f32,
6041) -> f32 {
6042    let mut acc = 0f32;
6043    for gi in 0..gpr {
6044        let ch = &chunks[(r * gpr + gi) * Q2TP_CHUNK..(r * gpr + gi + 1) * Q2TP_CHUNK];
6045        let s = scales[gi];
6046        let xb = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
6047        let mut g = 0f32;
6048        for (k, &b) in ch.iter().enumerate() {
6049            g += ((b & 3) as f32 - center) * xb[k * 4]
6050                + (((b >> 2) & 3) as f32 - center) * xb[k * 4 + 1]
6051                + (((b >> 4) & 3) as f32 - center) * xb[k * 4 + 2]
6052                + (((b >> 6) & 3) as f32 - center) * xb[k * 4 + 3];
6053        }
6054        acc += s * g;
6055    }
6056    acc
6057}
6058
6059fn q2tp_matvec(
6060    bytes: &[u8],
6061    x: &[f32],
6062    rows: usize,
6063    cols: usize,
6064    out: &mut [f32],
6065    pool: Option<&Pool>,
6066) {
6067    q2tp_matvec_mode(bytes, x, rows, cols, out, pool, false);
6068}
6069
6070fn q2tp_affine_matvec(
6071    bytes: &[u8],
6072    x: &[f32],
6073    rows: usize,
6074    cols: usize,
6075    out: &mut [f32],
6076    pool: Option<&Pool>,
6077) {
6078    q2tp_matvec_mode(bytes, x, rows, cols, out, pool, true);
6079}
6080
6081fn q2tp_matvec_mode(
6082    bytes: &[u8],
6083    x: &[f32],
6084    rows: usize,
6085    cols: usize,
6086    out: &mut [f32],
6087    pool: Option<&Pool>,
6088    affine: bool,
6089) {
6090    debug_assert_eq!(out.len(), rows);
6091    let gpr = cols / GROUP_SIZE;
6092    let v = Q4tpView::new_q2(bytes, rows, cols);
6093    let out_addr = SendMut(out.as_mut_ptr());
6094    // a8w8 fast path (CMF_SDOT=0 keeps the exact scalar walk): integer
6095    // code dots + group sums, exact outlier correction — the same
6096    // contract as every sibling kernel; measured 2-bit rows were the
6097    // only scalar holdout in the family.
6098    if !affine && a8w8_enabled() {
6099        let act = split_act(x);
6100        let gsum = q1_group_sums(&act.xq, gpr);
6101        let (act, gsum) = (&act, &gsum);
6102        let run = move |start: usize, end: usize| {
6103            with_krow(gpr, |sc| {
6104                for r in start..end {
6105                    v.scales_into(r, gpr, sc);
6106                    let mut acc = dot_q2tp_row_i8(v.nib, r, gpr, &act.xq, gsum, sc) * act.sx;
6107                    for &(j, xv) in &act.outliers {
6108                        let (w, s) = q2tp_outlier(v.nib, r, gpr, j, sc);
6109                        acc += w * s * xv;
6110                    }
6111                    // SAFETY: disjoint row ranges per worker.
6112                    unsafe { *out_addr.at(r) = acc };
6113                }
6114            })
6115        };
6116        dispatch_rows(pool, rows, &run);
6117        return;
6118    }
6119    let run = |start: usize, end: usize| {
6120        with_krow(gpr, |sc| {
6121            for r in start..end {
6122                v.scales_into(r, gpr, sc);
6123                // SAFETY: disjoint row ranges per worker.
6124                unsafe {
6125                    *out_addr.at(r) = if affine {
6126                        q2tp_affine_row_exact(v.nib, r, gpr, x, sc)
6127                    } else {
6128                        q2tp_row_exact(v.nib, r, gpr, x, sc)
6129                    }
6130                };
6131            }
6132        })
6133    };
6134    dispatch_rows(pool, rows, &run);
6135}
6136
6137/// Fused two-input q2tp matvec — the SwiGLU gate/up pair.
6138#[allow(clippy::too_many_arguments)]
6139fn q2tp_matvec2(
6140    bytes: &[u8],
6141    x1: &[f32],
6142    x2: &[f32],
6143    rows: usize,
6144    cols: usize,
6145    o1: &mut [f32],
6146    o2: &mut [f32],
6147    pool: Option<&Pool>,
6148) {
6149    q2tp_matvec2_mode(bytes, x1, x2, rows, cols, o1, o2, pool, false);
6150}
6151
6152#[allow(clippy::too_many_arguments)]
6153fn q2tp_affine_matvec2(
6154    bytes: &[u8],
6155    x1: &[f32],
6156    x2: &[f32],
6157    rows: usize,
6158    cols: usize,
6159    o1: &mut [f32],
6160    o2: &mut [f32],
6161    pool: Option<&Pool>,
6162) {
6163    q2tp_matvec2_mode(bytes, x1, x2, rows, cols, o1, o2, pool, true);
6164}
6165
6166#[allow(clippy::too_many_arguments)]
6167fn q2tp_matvec2_mode(
6168    bytes: &[u8],
6169    x1: &[f32],
6170    x2: &[f32],
6171    rows: usize,
6172    cols: usize,
6173    o1: &mut [f32],
6174    o2: &mut [f32],
6175    pool: Option<&Pool>,
6176    affine: bool,
6177) {
6178    let gpr = cols / GROUP_SIZE;
6179    let v = Q4tpView::new_q2(bytes, rows, cols);
6180    let (p1, p2) = (SendMut(o1.as_mut_ptr()), SendMut(o2.as_mut_ptr()));
6181    let run = |start: usize, end: usize| {
6182        let mut sc = vec![0f32; gpr];
6183        for r in start..end {
6184            v.scales_into(r, gpr, &mut sc);
6185            // SAFETY: disjoint row ranges per worker.
6186            unsafe {
6187                *p1.at(r) = if affine {
6188                    q2tp_affine_row_exact(v.nib, r, gpr, x1, &sc)
6189                } else {
6190                    q2tp_row_exact(v.nib, r, gpr, x1, &sc)
6191                };
6192                *p2.at(r) = if affine {
6193                    q2tp_affine_row_exact(v.nib, r, gpr, x2, &sc)
6194                } else {
6195                    q2tp_row_exact(v.nib, r, gpr, x2, &sc)
6196                };
6197            }
6198        }
6199    };
6200    dispatch_rows(pool, rows, &run);
6201}
6202
6203/// Batched q2tp matmat: scalar row kernel over every batch column. CPU
6204/// prefill only — decode rides the graph, so plain and correct beats
6205/// clever here.
6206/// Test doors into the host 2-bit kernels: the stand's heap corruption
6207/// pointed at down-shaped tensors, and the private fns need a way to be
6208/// held to a reference without a model file around them.
6209pub fn q2tp_matvec_for_test(bytes: &[u8], x: &[f32], rows: usize, cols: usize, out: &mut [f32]) {
6210    // The facade IS the reference: encoder oracles hold requant output
6211    // to the exact scalar walk. The production dispatch may take the i8
6212    // fast path, whose error scale is the ACTIVATIONS' — a different
6213    // claim than the encoder correctness these tests pin.
6214    let gpr = cols / GROUP_SIZE;
6215    let v = Q4tpView::new_q2(bytes, rows, cols);
6216    with_krow(gpr, |sc| {
6217        for r in 0..rows {
6218            v.scales_into(r, gpr, sc);
6219            out[r] = q2tp_row_exact(v.nib, r, gpr, x, sc);
6220        }
6221    });
6222}
6223
6224/// Test door for the descriptor-specific fused affine decode.  Production
6225/// callers select this through a validated Prism header, never by dtype alone.
6226pub fn q2tp_affine_matvec_for_test(
6227    bytes: &[u8],
6228    x: &[f32],
6229    rows: usize,
6230    cols: usize,
6231    out: &mut [f32],
6232) {
6233    q2tp_affine_matvec(bytes, x, rows, cols, out, None);
6234}
6235
6236pub fn q2tp_matmat_for_test(
6237    bytes: &[u8],
6238    xs_all: &[f32],
6239    b: usize,
6240    rows: usize,
6241    cols: usize,
6242    out: &mut [f32],
6243) {
6244    q2tp_matmat(bytes, xs_all, b, rows, cols, out, None);
6245}
6246
6247fn q2tp_matmat(
6248    bytes: &[u8],
6249    xs_all: &[f32],
6250    b: usize,
6251    rows: usize,
6252    cols: usize,
6253    out: &mut [f32],
6254    pool: Option<&Pool>,
6255) {
6256    q2tp_matmat_mode(bytes, xs_all, b, rows, cols, out, pool, false);
6257}
6258
6259fn q2tp_affine_matmat(
6260    bytes: &[u8],
6261    xs_all: &[f32],
6262    b: usize,
6263    rows: usize,
6264    cols: usize,
6265    out: &mut [f32],
6266    pool: Option<&Pool>,
6267) {
6268    q2tp_matmat_mode(bytes, xs_all, b, rows, cols, out, pool, true);
6269}
6270
6271fn q2tp_matmat_mode(
6272    bytes: &[u8],
6273    xs_all: &[f32],
6274    b: usize,
6275    rows: usize,
6276    cols: usize,
6277    out: &mut [f32],
6278    pool: Option<&Pool>,
6279    affine: bool,
6280) {
6281    debug_assert_eq!(out.len(), b * rows);
6282    let gpr = cols / GROUP_SIZE;
6283    let v = Q4tpView::new_q2(bytes, rows, cols);
6284    let out_addr = SendMut(out.as_mut_ptr());
6285    let run = |start: usize, end: usize| {
6286        let mut sc = vec![0f32; gpr];
6287        for r in start..end {
6288            v.scales_into(r, gpr, &mut sc);
6289            for bi in 0..b {
6290                let x = &xs_all[bi * cols..(bi + 1) * cols];
6291                // SAFETY: disjoint row ranges per worker.
6292                unsafe {
6293                    *out_addr.at(bi * rows + r) = if affine {
6294                        q2tp_affine_row_exact(v.nib, r, gpr, x, &sc)
6295                    } else {
6296                        q2tp_row_exact(v.nib, r, gpr, x, &sc)
6297                    }
6298                };
6299            }
6300        }
6301    };
6302    dispatch_rows(pool, rows, &run);
6303}
6304
6305/// The pre-vectorised shape, kept for A/B (`CMF_Q4TP_V1=1`): the
6306/// horizontal add lands once per group per column instead of once per
6307/// row. Same weights, same activations — only the reduction differs.
6308/// It is also the row-exact batch kernel: per group and column it forms
6309/// `int dot as f32 * scale` and adds it to a scalar running sum, which is
6310/// `dot_q4tp_row_sdot` step for step (Rust never contracts to an fma),
6311/// so each column is bit-identical to that column's matvec.
6312#[cfg(target_arch = "aarch64")]
6313#[target_feature(enable = "neon,dotprod")]
6314unsafe fn dot_q4tp_row_1x4_sdot_v1(
6315    nib: &[u8],
6316    r: usize,
6317    gpr: usize,
6318    xs: [&[i8]; 4],
6319    scales: &[f32],
6320) -> [f32; 4] {
6321    unsafe {
6322        use core::arch::aarch64::*;
6323        use core::arch::asm;
6324        let lomask = vdupq_n_u8(0x0F);
6325        let eight = vdupq_n_s8(8);
6326        let (mut f0, mut f1, mut f2, mut f3) = (0f32, 0f32, 0f32, 0f32);
6327        for gi in 0..gpr {
6328            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
6329            let s = *scales.get_unchecked(gi);
6330            let bb = vld1q_u8(t);
6331            let lo = vandq_u8(bb, lomask);
6332            let hi = vshrq_n_u8::<4>(bb);
6333            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
6334            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
6335            let mut d = [0f32; 4];
6336            for (k, dk) in d.iter_mut().enumerate() {
6337                let x0 = vld1q_s8(xs[k].as_ptr().add(gi * GROUP_SIZE));
6338                let x1 = vld1q_s8(xs[k].as_ptr().add(gi * GROUP_SIZE + 16));
6339                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6340                asm!(
6341                    "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
6342                    "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
6343                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6344                    e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
6345                    options(pure, nomem, nostack),
6346                );
6347                *dk = vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
6348            }
6349            f0 += d[0];
6350            f1 += d[1];
6351            f2 += d[2];
6352            f3 += d[3];
6353        }
6354        [f0, f1, f2, f3]
6355    }
6356}
6357
6358/// Which q4tp batch kernel to run: 1 = the previous one, 2 = the tuned
6359/// one, 0 = decide from the CPU. An atomic rather than a `OnceLock` so a
6360/// benchmark can alternate the two inside one process, where the machine's
6361/// mood — a shared box drifts ±25% between runs — is the same for both.
6362/// What the two mean is per-architecture: on x86 the blocked AVX-512 path
6363/// against the per-column one, on ARM the two reduction shapes.
6364#[allow(dead_code)]
6365static Q4TP_ALT: std::sync::atomic::AtomicU8 = std::sync::atomic::AtomicU8::new(0);
6366
6367/// Tests that store `Q4TP_ALT` hold this, so one test's kernel pick does
6368/// not leak into another's bit-exact comparison running in parallel.
6369#[cfg(test)]
6370static Q4TP_ALT_TEST_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
6371
6372/// Blocking pays on x86 only with 512-bit VNNI. With AVX2 alone, four
6373/// columns sharing an unpack still measured slower than the per-column
6374/// path (23.2 ms against 19.4 on a 48-thread EPYC), because that path
6375/// already dequantizes the row once — so the blocked kernel bought a
6376/// second unpack-free pass at the price of half the vector width.
6377#[cfg(target_arch = "x86_64")]
6378fn q4tp_blocked_x86() -> bool {
6379    match Q4TP_ALT.load(std::sync::atomic::Ordering::Relaxed) {
6380        1 => false,
6381        // A forced ON still asks the CPU. The switch exists so a bench can
6382        // pick a kernel, not so it can promise instructions the machine
6383        // does not have — CI caught that as a SIGILL on a runner without
6384        // AVX-512, where the parity test had turned the path on by hand.
6385        2 => avx512vnni_enabled(),
6386        // Deliberately not cached back into the switch: both gates below
6387        // hold their own `OnceLock`, and latching their answer here would
6388        // make a test's override outlive the test that set it.
6389        _ => blocked_enabled() && avx512vnni_enabled(),
6390    }
6391}
6392
6393/// `CMF_Q4TP_V1=1` picks the old reduction shape (A/B only).
6394#[cfg(target_arch = "aarch64")]
6395#[allow(dead_code)]
6396fn q4tp_v1() -> bool {
6397    match Q4TP_ALT.load(std::sync::atomic::Ordering::Relaxed) {
6398        1 => true,
6399        2 => false,
6400        _ => {
6401            static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6402            *ON.get_or_init(|| std::env::var("CMF_Q4TP_V1").is_ok_and(|v| v != "0"))
6403        }
6404    }
6405}
6406
6407/// Two weight rows against eight columns. The activation load is the
6408/// same for both rows, so it is paid once for twice the arithmetic, and
6409/// sixteen accumulator chains run where eight did — which is what a kernel
6410/// retiring 0.29 instructions a cycle is short of. Register pressure is
6411/// the limit: sixteen `zmm` accumulators, two weight tiles, one
6412/// activation, of thirty-two.
6413///
6414/// Four rows by four columns spends the same sixteen accumulators the
6415/// other way and measured worse — 1488 GFLOP/s against 1644 — so the
6416/// unpack, which four rows pay twice as often, costs more than the extra
6417/// sharing of one activation load buys.
6418#[cfg(target_arch = "x86_64")]
6419#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
6420unsafe fn dot_q4tp_2x8_avx512(
6421    nib: &[u8],
6422    r0: usize,
6423    gpr: usize,
6424    xs: [&[i8]; 8],
6425    sc0: &[f32],
6426    sc1: &[f32],
6427) -> [[f32; 8]; 2] {
6428    // SAFETY: as dot_q4tp_row_1x8_avx512, two adjacent rows at once; the
6429    // caller guarantees r0 + 1 < rows and the ISA.
6430    unsafe {
6431        use core::arch::x86_64::*;
6432        let lomask = _mm256_set1_epi8(0x0F);
6433        let eight = _mm256_set1_epi8(8);
6434        let zero = _mm512_setzero_si512();
6435        let mut v0 = [_mm512_setzero_ps(); 8];
6436        let mut v1 = [_mm512_setzero_ps(); 8];
6437        let pairs = gpr / 2;
6438        let unpack = |r: usize, gi: usize| -> (__m512i, __mmask64) {
6439            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
6440            let bb = _mm256_loadu_si256(t as *const __m256i);
6441            let lo = _mm256_and_si256(bb, lomask);
6442            let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
6443            let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
6444            let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
6445            let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
6446            let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
6447            (_mm512_abs_epi8(w), _mm512_movepi8_mask(w))
6448        };
6449        for gp in 0..pairs {
6450            let gi = gp * 2;
6451            let (wa0, neg0) = unpack(r0, gi);
6452            let (wa1, neg1) = unpack(r0 + 1, gi);
6453            let off = gi * GROUP_SIZE;
6454            let sv = |sc: &[f32]| {
6455                _mm512_insertf32x8::<1>(
6456                    _mm512_castps256_ps512(_mm256_set1_ps(*sc.get_unchecked(gi))),
6457                    _mm256_set1_ps(*sc.get_unchecked(gi + 1)),
6458                )
6459            };
6460            let s0 = sv(sc0);
6461            let s1 = sv(sc1);
6462            for k in 0..8 {
6463                let xv = _mm512_loadu_si512(xs[k].as_ptr().add(off) as *const __m512i);
6464                let d0 = _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(
6465                    zero,
6466                    wa0,
6467                    _mm512_mask_sub_epi8(xv, neg0, zero, xv),
6468                ));
6469                let d1 = _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(
6470                    zero,
6471                    wa1,
6472                    _mm512_mask_sub_epi8(xv, neg1, zero, xv),
6473                ));
6474                v0[k] = _mm512_fmadd_ps(d0, s0, v0[k]);
6475                v1[k] = _mm512_fmadd_ps(d1, s1, v1[k]);
6476            }
6477        }
6478        let mut acc = [[0f32; 8]; 2];
6479        for k in 0..8 {
6480            acc[0][k] = _mm512_reduce_add_ps(v0[k]);
6481            acc[1][k] = _mm512_reduce_add_ps(v1[k]);
6482        }
6483        if gpr % 2 == 1 {
6484            let off = (gpr - 1) * GROUP_SIZE;
6485            for j in off..off + GROUP_SIZE {
6486                let (w0, sa) = q4tp_outlier(nib, r0, gpr, j, sc0);
6487                let (w1, sb) = q4tp_outlier(nib, r0 + 1, gpr, j, sc1);
6488                for k in 0..8 {
6489                    let x = *xs[k].get_unchecked(j) as f32;
6490                    acc[0][k] += w0 * sa * x;
6491                    acc[1][k] += w1 * sb * x;
6492                }
6493            }
6494        }
6495        acc
6496    }
6497}
6498
6499/// The same, eight columns at a time. One unpack then feeds twice as many
6500/// activation streams, so a wide batch reads the weight tile half as
6501/// often; the price is eight accumulators live at once. Measured 9.0 ->
6502/// 8.3 ms at 9216x2304, b=296 on a 48-thread EPYC 9B45.
6503#[cfg(target_arch = "x86_64")]
6504#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
6505unsafe fn dot_q4tp_row_1x8_avx512(
6506    nib: &[u8],
6507    r: usize,
6508    gpr: usize,
6509    xs: [&[i8]; 8],
6510    scales: &[f32],
6511) -> [f32; 8] {
6512    // SAFETY: as dot_q4tp_row_1x4_avx2; caller guarantees the ISA.
6513    unsafe {
6514        use core::arch::x86_64::*;
6515        let lomask = _mm256_set1_epi8(0x0F);
6516        let eight = _mm256_set1_epi8(8);
6517        let zero = _mm512_setzero_si512();
6518        let (mut v0, mut v1, mut v2, mut v3) = (
6519            _mm512_setzero_ps(),
6520            _mm512_setzero_ps(),
6521            _mm512_setzero_ps(),
6522            _mm512_setzero_ps(),
6523        );
6524        let (mut v4, mut v5, mut v6, mut v7) = (
6525            _mm512_setzero_ps(),
6526            _mm512_setzero_ps(),
6527            _mm512_setzero_ps(),
6528            _mm512_setzero_ps(),
6529        );
6530        let pairs = gpr / 2;
6531        for gp in 0..pairs {
6532            let gi = gp * 2;
6533            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
6534            let bb = _mm256_loadu_si256(t as *const __m256i);
6535            let lo = _mm256_and_si256(bb, lomask);
6536            let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
6537            // `unpack` works per 128-bit lane, so the halves come out as
6538            // [A.lo, B.lo] and [A.hi, B.hi]; the shuffle reorders the four
6539            // 128-bit lanes into the weights' natural order, which is what
6540            // the straight activation load expects.
6541            let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
6542            let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
6543            let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
6544            let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
6545            let wabs = _mm512_abs_epi8(w);
6546            let neg = _mm512_movepi8_mask(w);
6547            let off = gi * GROUP_SIZE;
6548            let sv = _mm512_insertf32x8::<1>(
6549                _mm512_castps256_ps512(_mm256_set1_ps(*scales.get_unchecked(gi))),
6550                _mm256_set1_ps(*scales.get_unchecked(gi + 1)),
6551            );
6552            let dot = |x: &[i8]| -> __m512 {
6553                let xv = _mm512_loadu_si512(x.as_ptr().add(off) as *const __m512i);
6554                let sx = _mm512_mask_sub_epi8(xv, neg, zero, xv);
6555                _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(zero, wabs, sx))
6556            };
6557            v0 = _mm512_fmadd_ps(dot(xs[0]), sv, v0);
6558            v1 = _mm512_fmadd_ps(dot(xs[1]), sv, v1);
6559            v2 = _mm512_fmadd_ps(dot(xs[2]), sv, v2);
6560            v3 = _mm512_fmadd_ps(dot(xs[3]), sv, v3);
6561            v4 = _mm512_fmadd_ps(dot(xs[4]), sv, v4);
6562            v5 = _mm512_fmadd_ps(dot(xs[5]), sv, v5);
6563            v6 = _mm512_fmadd_ps(dot(xs[6]), sv, v6);
6564            v7 = _mm512_fmadd_ps(dot(xs[7]), sv, v7);
6565        }
6566        let mut acc = [
6567            _mm512_reduce_add_ps(v0),
6568            _mm512_reduce_add_ps(v1),
6569            _mm512_reduce_add_ps(v2),
6570            _mm512_reduce_add_ps(v3),
6571            _mm512_reduce_add_ps(v4),
6572            _mm512_reduce_add_ps(v5),
6573            _mm512_reduce_add_ps(v6),
6574            _mm512_reduce_add_ps(v7),
6575        ];
6576        // An odd group count leaves one group over; the narrow kernel
6577        // finishes it rather than the tail being a special case here.
6578        if gpr % 2 == 1 {
6579            let off = (gpr - 1) * GROUP_SIZE;
6580            for j in off..off + GROUP_SIZE {
6581                let (w, s) = q4tp_outlier(nib, r, gpr, j, scales);
6582                let ws = w * s;
6583                for k in 0..8 {
6584                    acc[k] += ws * *xs[k].get_unchecked(j) as f32;
6585                }
6586            }
6587        }
6588        acc
6589    }
6590}
6591
6592/// The same four columns, 512 bits wide. Two groups (64 weights) ride one
6593/// unpack and one `vpdpbusd`, where AVX2 needs two unpacks and four
6594/// `maddubs`/`madd` pairs — about 2.3x fewer instructions for the same
6595/// arithmetic. The two groups carry different scales, so the fma takes a
6596/// vector whose halves hold each group's scale rather than a broadcast.
6597///
6598/// There is no 512-bit `vpsignb`, so the activation's sign is applied by
6599/// negating under a mask taken from the weight's sign bits. That mask is
6600/// per-tile, so it is hoisted out of the column loop and the per-column
6601/// cost stays exactly one instruction, as with `sign_epi8`. Weights of
6602/// zero are not zeroed by the mask trick and do not need to be: their
6603/// magnitude is zero, so the product is.
6604#[cfg(target_arch = "x86_64")]
6605#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
6606unsafe fn dot_q4tp_row_1x4_avx512(
6607    nib: &[u8],
6608    r: usize,
6609    gpr: usize,
6610    xs: [&[i8]; 4],
6611    scales: &[f32],
6612) -> [f32; 4] {
6613    // SAFETY: as dot_q4tp_row_1x4_avx2; caller guarantees the ISA.
6614    unsafe {
6615        use core::arch::x86_64::*;
6616        let lomask = _mm256_set1_epi8(0x0F);
6617        let eight = _mm256_set1_epi8(8);
6618        let zero = _mm512_setzero_si512();
6619        let (mut v0, mut v1, mut v2, mut v3) = (
6620            _mm512_setzero_ps(),
6621            _mm512_setzero_ps(),
6622            _mm512_setzero_ps(),
6623            _mm512_setzero_ps(),
6624        );
6625        let pairs = gpr / 2;
6626        for gp in 0..pairs {
6627            let gi = gp * 2;
6628            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
6629            let bb = _mm256_loadu_si256(t as *const __m256i);
6630            let lo = _mm256_and_si256(bb, lomask);
6631            let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
6632            // `unpack` works per 128-bit lane, so the halves come out as
6633            // [A.lo, B.lo] and [A.hi, B.hi]; the shuffle reorders the four
6634            // 128-bit lanes into the weights' natural order, which is what
6635            // the straight activation load expects.
6636            let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
6637            let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
6638            let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
6639            let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
6640            let wabs = _mm512_abs_epi8(w);
6641            let neg = _mm512_movepi8_mask(w);
6642            let off = gi * GROUP_SIZE;
6643            let sv = _mm512_insertf32x8::<1>(
6644                _mm512_castps256_ps512(_mm256_set1_ps(*scales.get_unchecked(gi))),
6645                _mm256_set1_ps(*scales.get_unchecked(gi + 1)),
6646            );
6647            let dot = |x: &[i8]| -> __m512 {
6648                let xv = _mm512_loadu_si512(x.as_ptr().add(off) as *const __m512i);
6649                let sx = _mm512_mask_sub_epi8(xv, neg, zero, xv);
6650                _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(zero, wabs, sx))
6651            };
6652            v0 = _mm512_fmadd_ps(dot(xs[0]), sv, v0);
6653            v1 = _mm512_fmadd_ps(dot(xs[1]), sv, v1);
6654            v2 = _mm512_fmadd_ps(dot(xs[2]), sv, v2);
6655            v3 = _mm512_fmadd_ps(dot(xs[3]), sv, v3);
6656        }
6657        let mut acc = [
6658            _mm512_reduce_add_ps(v0),
6659            _mm512_reduce_add_ps(v1),
6660            _mm512_reduce_add_ps(v2),
6661            _mm512_reduce_add_ps(v3),
6662        ];
6663        // An odd group count leaves one group over; the narrow kernel
6664        // finishes it rather than the tail being a special case here.
6665        if gpr % 2 == 1 {
6666            let off = (gpr - 1) * GROUP_SIZE;
6667            for j in off..off + GROUP_SIZE {
6668                let (w, s) = q4tp_outlier(nib, r, gpr, j, scales);
6669                let ws = w * s;
6670                for k in 0..4 {
6671                    acc[k] += ws * *xs[k].get_unchecked(j) as f32;
6672                }
6673            }
6674        }
6675        acc
6676    }
6677}
6678
6679/// Four batch columns against one q4tp row: the tile is unpacked ONCE and
6680/// spent on four activation streams, which is where a prefill batch stops
6681/// being weight-bandwidth-bound. Twin of `dot_q4t_row_1x4_sdot`.
6682#[cfg(target_arch = "aarch64")]
6683#[target_feature(enable = "neon,dotprod")]
6684unsafe fn dot_q4tp_row_1x4_sdot(
6685    nib: &[u8],
6686    r: usize,
6687    gpr: usize,
6688    xs: [&[i8]; 4],
6689    scales: &[f32],
6690) -> [f32; 4] {
6691    // SAFETY: see dot_q4tp_row_sdot; every xs[k] is gpr·GROUP_SIZE long.
6692    unsafe {
6693        use core::arch::aarch64::*;
6694        use core::arch::asm;
6695        let lomask = vdupq_n_u8(0x0F);
6696        let eight = vdupq_n_s8(8);
6697        // Named accumulators, NOT an array indexed by a loop variable: the
6698        // latter does not stay in registers (the same defect cost 2x in the
6699        // AVX2 q4t kernel and again in WGSL).
6700        //
6701        // They are VECTORS, and the horizontal add happens once at the end
6702        // instead of once per group per column. `vaddvq` is a cross-lane
6703        // reduction — with 72 groups and four columns the old shape paid
6704        // 288 of them per row, each one a dependency stall the pipeline
6705        // cannot hide, to save four float adds. The group's scale now
6706        // rides an fma into the lane accumulators, so the arithmetic per
6707        // group is one convert and one fma. Summation order changes (the
6708        // lanes carry independent partial sums), which is the same
6709        // round-off class the SDOT path already lives in — the strict
6710        // kernel (`CMF_SDOT=0`, what `cortiq ppl` runs) is unchanged and
6711        // stays the reference.
6712        let (mut v0, mut v1, mut v2, mut v3) = (
6713            vdupq_n_f32(0.0),
6714            vdupq_n_f32(0.0),
6715            vdupq_n_f32(0.0),
6716            vdupq_n_f32(0.0),
6717        );
6718        for gi in 0..gpr {
6719            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
6720            let s = *scales.get_unchecked(gi);
6721            let bb = vld1q_u8(t);
6722            let lo = vandq_u8(bb, lomask);
6723            let hi = vshrq_n_u8::<4>(bb);
6724            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
6725            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
6726            let off = gi * GROUP_SIZE;
6727            let dot4 = |x: &[i8]| -> int32x4_t {
6728                let x0 = vld1q_s8(x.as_ptr().add(off));
6729                let x1 = vld1q_s8(x.as_ptr().add(off + 16));
6730                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6731                asm!(
6732                    "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
6733                    "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
6734                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6735                    e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
6736                    options(pure, nomem, nostack),
6737                );
6738                vaddq_s32(a0, a1)
6739            };
6740            v0 = vfmaq_n_f32(v0, vcvtq_f32_s32(dot4(xs[0])), s);
6741            v1 = vfmaq_n_f32(v1, vcvtq_f32_s32(dot4(xs[1])), s);
6742            v2 = vfmaq_n_f32(v2, vcvtq_f32_s32(dot4(xs[2])), s);
6743            v3 = vfmaq_n_f32(v3, vcvtq_f32_s32(dot4(xs[3])), s);
6744        }
6745        [
6746            vaddvq_f32(v0),
6747            vaddvq_f32(v1),
6748            vaddvq_f32(v2),
6749            vaddvq_f32(v3),
6750        ]
6751    }
6752}
6753
6754/// Fused q4tp matmat — the same three arms `q4t_matmat` has. Shipping only
6755/// the scalar one made Nanbeige-3B decode at 1.2 tok/s against q4t's 5.9:
6756/// the format was fine, the missing arms were the whole regression.
6757///
6758/// Under `row_exact()` every cell is summed in `q4tp_matvec`'s order on
6759/// every architecture; the mode is read once, so one call never mixes
6760/// the two contracts when a concurrent scope opens or closes mid-call.
6761fn q4tp_matmat(
6762    bytes: &[u8],
6763    xs_all: &[f32],
6764    b: usize,
6765    rows: usize,
6766    cols: usize,
6767    out: &mut [f32],
6768    pool: Option<&Pool>,
6769) {
6770    q4tp_matmat_with(bytes, xs_all, b, rows, cols, out, pool, row_exact())
6771}
6772
6773/// `q4tp_matmat` with the row-exact mode passed in rather than read from
6774/// the shared counter: `exact` = each cell equals its token's matvec bit
6775/// for bit, otherwise the fast blocked / AMX arms are free to reorder.
6776#[allow(clippy::too_many_arguments)]
6777fn q4tp_matmat_with(
6778    bytes: &[u8],
6779    xs_all: &[f32],
6780    b: usize,
6781    rows: usize,
6782    cols: usize,
6783    out: &mut [f32],
6784    pool: Option<&Pool>,
6785    exact: bool,
6786) {
6787    debug_assert_eq!(out.len(), b * rows);
6788    let gpr = cols / GROUP_SIZE;
6789    let v = Q4tpView::new(bytes, rows, cols);
6790
6791    // Wide batches ride the AMX through a dequant-tile sgemm, as in q4t.
6792    // An f32 GEMM over dequantized weights is not the matvec's int8 sum,
6793    // so a row-exact batch never takes it.
6794    #[cfg(target_os = "macos")]
6795    if !exact && b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
6796        dequant_matmat_accel(
6797            &|r, dst| {
6798                let mut sc = [0f32; 32];
6799                let mut scv;
6800                let s: &[f32] = if gpr <= 32 {
6801                    v.scales_into(r, gpr, &mut sc);
6802                    &sc[..gpr]
6803                } else {
6804                    scv = vec![0f32; gpr];
6805                    v.scales_into(r, gpr, &mut scv);
6806                    &scv
6807                };
6808                for gi in 0..gpr {
6809                    let tile = &v.nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
6810                    for (k, &bb) in tile.iter().enumerate() {
6811                        dst[gi * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s[gi];
6812                        dst[gi * GROUP_SIZE + k * 2 + 1] =
6813                            (((bb >> 4) & 0x0F) as f32 - 8.0) * s[gi];
6814                    }
6815                }
6816            },
6817            xs_all,
6818            b,
6819            rows,
6820            cols,
6821            out,
6822            pool,
6823        );
6824        return;
6825    }
6826
6827    let out_addr = SendMut(out.as_mut_ptr());
6828    if a8w8_enabled() {
6829        let acts: Vec<SplitAct> = (0..b)
6830            .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
6831            .collect();
6832        let acts = &acts;
6833        // ARM stays blocked under `exact` too: the 1x4 kernel then runs in
6834        // its v1 shape, whose per-group `int dot as f32 * scale` and scalar
6835        // running sum are exactly `dot_q4tp_row_sdot`'s, so the tile is
6836        // still unpacked once for four columns and each column equals its
6837        // matvec. The tuned shape (fma into lane partials, one horizontal
6838        // add per row) is 1-16 ulp off the matvec and is kept for `!exact`.
6839        #[cfg(target_arch = "aarch64")]
6840        let blocked_ok = sdot_enabled() && blocked_enabled();
6841        // x86 gets the same blocking: one tile unpack spent on four
6842        // columns. Without it every column re-decoded the row, which is
6843        // why a 48-core EPYC measured a sixth of an M4's per-core rate.
6844        // The gate is `avx2_enabled`, as in q4t — `sdot_enabled` answers
6845        // for ARM's dotprod and is hard-wired false everywhere else, so
6846        // asking it here left the whole blocked path unreachable on x86.
6847        #[cfg(target_arch = "x86_64")]
6848        let blocked_ok = q4tp_blocked_x86() && !exact;
6849        #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
6850        let blocked_ok = {
6851            let _ = exact;
6852            false
6853        };
6854        // Columns are swept in panels that fit L2. Without this a
6855        // row-pair walks every activation in the batch — 4.8 MB at
6856        // 512x512 — and does it again for the next pair, so the whole
6857        // batch streams out of the shared cache once per row. Measured
6858        // 800 GB/s of it, flat across batch sizes, which is the signature
6859        // of a loop bound by traffic rather than by arithmetic. A panel of
6860        // 256 columns is 590 KB beside 221 KB of this worker's weights:
6861        // both stay resident and the batch crosses L3 once instead of
6862        // once per row.
6863        let panel_cols: usize = std::env::var("CMF_Q4TP_PANEL")
6864            .ok()
6865            .and_then(|v| v.parse().ok())
6866            .filter(|v| *v > 0)
6867            .unwrap_or(256);
6868        let run = |start: usize, end: usize| {
6869            for abase in (0..acts.len()).step_by(panel_cols) {
6870                let alen = (acts.len() - abase).min(panel_cols);
6871                let mut sc = vec![0f32; gpr];
6872                #[cfg(target_arch = "x86_64")]
6873                let mut r_lo = start;
6874                #[cfg(target_arch = "x86_64")]
6875                if blocked_ok && alen >= 8 {
6876                    let mut sc1 = vec![0f32; gpr];
6877                    while r_lo + 2 <= end {
6878                        v.scales_into(r_lo, gpr, &mut sc);
6879                        v.scales_into(r_lo + 1, gpr, &mut sc1);
6880                        let mut bi = 0usize;
6881                        while bi + 8 <= alen {
6882                            let xs = [
6883                                acts[abase + bi].xq.as_slice(),
6884                                acts[abase + bi + 1].xq.as_slice(),
6885                                acts[abase + bi + 2].xq.as_slice(),
6886                                acts[abase + bi + 3].xq.as_slice(),
6887                                acts[abase + bi + 4].xq.as_slice(),
6888                                acts[abase + bi + 5].xq.as_slice(),
6889                                acts[abase + bi + 6].xq.as_slice(),
6890                                acts[abase + bi + 7].xq.as_slice(),
6891                            ];
6892                            let d = unsafe { dot_q4tp_2x8_avx512(v.nib, r_lo, gpr, xs, &sc, &sc1) };
6893                            for (row, dr, scr) in [(r_lo, &d[0], &sc), (r_lo + 1, &d[1], &sc1)] {
6894                                for k in 0..8 {
6895                                    let act = &acts[abase + bi + k];
6896                                    let mut acc = dr[k] * act.sx;
6897                                    for &(j, xv) in &act.outliers {
6898                                        let (w, s) = q4tp_outlier(v.nib, row, gpr, j, scr);
6899                                        acc += w * s * xv;
6900                                    }
6901                                    // SAFETY: disjoint (bi, r) cells per worker.
6902                                    unsafe { *out_addr.at((abase + bi + k) * rows + row) = acc };
6903                                }
6904                            }
6905                            bi += 8;
6906                        }
6907                        // Columns past the last group of eight, both rows —
6908                        // the same single-row kernel the tail below uses.
6909                        for row in [r_lo, r_lo + 1] {
6910                            let scr: &[f32] = if row == r_lo { &sc } else { &sc1 };
6911                            for b2 in bi..alen {
6912                                let act = &acts[abase + b2];
6913                                let xs4 = [
6914                                    act.xq.as_slice(),
6915                                    act.xq.as_slice(),
6916                                    act.xq.as_slice(),
6917                                    act.xq.as_slice(),
6918                                ];
6919                                let d =
6920                                    unsafe { dot_q4tp_row_1x4_avx512(v.nib, row, gpr, xs4, scr) };
6921                                let mut acc = d[0] * act.sx;
6922                                for &(j, xv) in &act.outliers {
6923                                    let (w, s) = q4tp_outlier(v.nib, row, gpr, j, scr);
6924                                    acc += w * s * xv;
6925                                }
6926                                // SAFETY: disjoint (bi, r) cells per worker.
6927                                unsafe { *out_addr.at((abase + b2) * rows + row) = acc };
6928                            }
6929                        }
6930                        r_lo += 2;
6931                    }
6932                }
6933                #[cfg(target_arch = "x86_64")]
6934                let row_start = r_lo;
6935                #[cfg(not(target_arch = "x86_64"))]
6936                let row_start = start;
6937                for r in row_start..end {
6938                    v.scales_into(r, gpr, &mut sc);
6939                    let mut bi = 0usize;
6940                    #[cfg(target_arch = "x86_64")]
6941                    if blocked_ok {
6942                        while bi + 8 <= alen {
6943                            let xs = [
6944                                acts[abase + bi].xq.as_slice(),
6945                                acts[abase + bi + 1].xq.as_slice(),
6946                                acts[abase + bi + 2].xq.as_slice(),
6947                                acts[abase + bi + 3].xq.as_slice(),
6948                                acts[abase + bi + 4].xq.as_slice(),
6949                                acts[abase + bi + 5].xq.as_slice(),
6950                                acts[abase + bi + 6].xq.as_slice(),
6951                                acts[abase + bi + 7].xq.as_slice(),
6952                            ];
6953                            let d = unsafe { dot_q4tp_row_1x8_avx512(v.nib, r, gpr, xs, &sc) };
6954                            for k in 0..8 {
6955                                let act = &acts[abase + bi + k];
6956                                let mut acc = d[k] * act.sx;
6957                                for &(j, xv) in &act.outliers {
6958                                    let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
6959                                    acc += w * s * xv;
6960                                }
6961                                // SAFETY: disjoint (bi, r) cells per worker.
6962                                unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
6963                            }
6964                            bi += 8;
6965                        }
6966                        while bi + 4 <= alen {
6967                            let xs = [
6968                                acts[abase + bi].xq.as_slice(),
6969                                acts[abase + bi + 1].xq.as_slice(),
6970                                acts[abase + bi + 2].xq.as_slice(),
6971                                acts[abase + bi + 3].xq.as_slice(),
6972                            ];
6973                            let d = unsafe { dot_q4tp_row_1x4_avx512(v.nib, r, gpr, xs, &sc) };
6974                            for k in 0..4 {
6975                                let act = &acts[abase + bi + k];
6976                                let mut acc = d[k] * act.sx;
6977                                for &(j, xv) in &act.outliers {
6978                                    let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
6979                                    acc += w * s * xv;
6980                                }
6981                                // SAFETY: disjoint (bi, r) cells per worker.
6982                                unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
6983                            }
6984                            bi += 4;
6985                        }
6986                    }
6987                    #[cfg(target_arch = "aarch64")]
6988                    if blocked_ok {
6989                        while bi + 4 <= alen {
6990                            let xs = [
6991                                acts[abase + bi].xq.as_slice(),
6992                                acts[abase + bi + 1].xq.as_slice(),
6993                                acts[abase + bi + 2].xq.as_slice(),
6994                                acts[abase + bi + 3].xq.as_slice(),
6995                            ];
6996                            let d = unsafe {
6997                                if exact || q4tp_v1() {
6998                                    dot_q4tp_row_1x4_sdot_v1(v.nib, r, gpr, xs, &sc)
6999                                } else {
7000                                    dot_q4tp_row_1x4_sdot(v.nib, r, gpr, xs, &sc)
7001                                }
7002                            };
7003                            for k in 0..4 {
7004                                let act = &acts[abase + bi + k];
7005                                let mut acc = d[k] * act.sx;
7006                                for &(j, xv) in &act.outliers {
7007                                    let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
7008                                    acc += w * s * xv;
7009                                }
7010                                // SAFETY: disjoint (bi, r) cells per worker.
7011                                unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
7012                            }
7013                            bi += 4;
7014                        }
7015                    }
7016                    let _ = blocked_ok;
7017                    while bi < alen {
7018                        let act = &acts[abase + bi];
7019                        let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, &sc) * act.sx;
7020                        for &(j, xv) in &act.outliers {
7021                            let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
7022                            acc += w * s * xv;
7023                        }
7024                        // SAFETY: disjoint (bi, r) cells per worker range.
7025                        unsafe { *out_addr.at((abase + bi) * rows + r) = acc };
7026                        bi += 1;
7027                    }
7028                }
7029            }
7030        };
7031        dispatch_rows(pool, rows, &run);
7032        return;
7033    }
7034
7035    let run = |start: usize, end: usize| {
7036        let mut sc = vec![0f32; gpr];
7037        for r in start..end {
7038            v.scales_into(r, gpr, &mut sc);
7039            for bi in 0..b {
7040                let x = &xs_all[bi * cols..(bi + 1) * cols];
7041                // SAFETY: disjoint (bi, r) cells per worker range.
7042                unsafe { *out_addr.at(bi * rows + r) = q4tp_row_exact(v.nib, r, gpr, x, &sc) };
7043            }
7044        }
7045    };
7046    dispatch_rows(pool, rows, &run);
7047}
7048
7049/// Fused q4_tiled matvec (dispatch mirrors `q4matvec`).
7050fn q4t_matvec(
7051    bytes: &[u8],
7052    x: &[f32],
7053    rows: usize,
7054    cols: usize,
7055    out: &mut [f32],
7056    pool: Option<&Pool>,
7057) {
7058    debug_assert_eq!(out.len(), rows);
7059    let gpr = cols / GROUP_SIZE;
7060    let out_addr = SendMut(out.as_mut_ptr());
7061    if a8w8_enabled() {
7062        let act = split_act(x);
7063        let run = move |start: usize, end: usize| {
7064            for r in start..end {
7065                let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
7066                for &(j, xv) in &act.outliers {
7067                    let (w, s) = q4t_outlier(bytes, r, gpr, j);
7068                    acc += w * s * xv;
7069                }
7070                // SAFETY: disjoint row ranges per worker.
7071                unsafe { *out_addr.at(r) = acc };
7072            }
7073        };
7074        dispatch_rows(pool, rows, &run);
7075        return;
7076    }
7077    let run = move |start: usize, end: usize| {
7078        for r in start..end {
7079            // SAFETY: disjoint row ranges per worker.
7080            unsafe { *out_addr.at(r) = q4t_row_exact(bytes, r, gpr, x) };
7081        }
7082    };
7083    dispatch_rows(pool, rows, &run);
7084}
7085
7086/// Fused two-input q4_tiled matvec (weights read once per pair).
7087#[allow(clippy::too_many_arguments)]
7088fn q4t_matvec2(
7089    bytes: &[u8],
7090    x1: &[f32],
7091    x2: &[f32],
7092    rows: usize,
7093    cols: usize,
7094    o1: &mut [f32],
7095    o2: &mut [f32],
7096    pool: Option<&Pool>,
7097) {
7098    let gpr = cols / GROUP_SIZE;
7099    let p1 = SendMut(o1.as_mut_ptr());
7100    let p2 = SendMut(o2.as_mut_ptr());
7101    if a8w8_enabled() {
7102        let a1 = split_act(x1);
7103        let a2 = split_act(x2);
7104        let run = move |start: usize, end: usize| {
7105            for r in start..end {
7106                let mut v1 = dot_q4t_row_i8(bytes, r, gpr, &a1.xq) * a1.sx;
7107                let mut v2 = dot_q4t_row_i8(bytes, r, gpr, &a2.xq) * a2.sx;
7108                for &(j, xv) in &a1.outliers {
7109                    let (w, s) = q4t_outlier(bytes, r, gpr, j);
7110                    v1 += w * s * xv;
7111                }
7112                for &(j, xv) in &a2.outliers {
7113                    let (w, s) = q4t_outlier(bytes, r, gpr, j);
7114                    v2 += w * s * xv;
7115                }
7116                // SAFETY: disjoint row ranges per worker.
7117                unsafe {
7118                    *p1.at(r) = v1;
7119                    *p2.at(r) = v2;
7120                }
7121            }
7122        };
7123        dispatch_rows(pool, rows, &run);
7124        return;
7125    }
7126    let run = move |start: usize, end: usize| {
7127        for r in start..end {
7128            // SAFETY: disjoint row ranges per worker.
7129            unsafe {
7130                *p1.at(r) = q4t_row_exact(bytes, r, gpr, x1);
7131                *p2.at(r) = q4t_row_exact(bytes, r, gpr, x2);
7132            }
7133        }
7134    };
7135    dispatch_rows(pool, rows, &run);
7136}
7137
7138/// Batched q4_tiled matmat: each row's tiles stream once per microbatch.
7139#[allow(clippy::too_many_arguments)]
7140/// Prefill GEMM through Accelerate for group-quantized codecs: a
7141/// caller-supplied row dequantizer fills f32 tiles (pool-parallel) and
7142/// each tile rides the AMX with one sgemm — the generic sibling of
7143/// `qmatmat_accel` (q8). Numerics are f32-GEMM (tolerance class);
7144/// decode (b=1) never takes this path.
7145#[cfg(target_os = "macos")]
7146fn dequant_matmat_accel(
7147    dequant_row: &(dyn Fn(usize, &mut [f32]) + Sync),
7148    xs_all: &[f32],
7149    b: usize,
7150    rows: usize,
7151    cols: usize,
7152    out: &mut [f32],
7153    pool: Option<&Pool>,
7154) {
7155    const TR: usize = 2048;
7156    thread_local! {
7157        static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
7158    }
7159    WTILE.with(|wt| {
7160        let mut wtile = wt.borrow_mut();
7161        wtile.resize(TR * cols, 0.0);
7162        let mut r0 = 0usize;
7163        while r0 < rows {
7164            let tr = TR.min(rows - r0);
7165            let wt_addr = SendMut(wtile.as_mut_ptr());
7166            let run = |start: usize, end: usize| {
7167                for r in start..end {
7168                    // SAFETY: workers cover disjoint r ranges.
7169                    let dst = unsafe { std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols) };
7170                    dequant_row(r0 + r, dst);
7171                }
7172            };
7173            dispatch_rows(pool, tr, &run);
7174            unsafe {
7175                accel_blas::cblas_sgemm(
7176                    101, // RowMajor
7177                    111, // NoTrans A
7178                    112, // Trans B
7179                    b as i32,
7180                    tr as i32,
7181                    cols as i32,
7182                    1.0,
7183                    xs_all.as_ptr(),
7184                    cols as i32,
7185                    wtile.as_ptr(),
7186                    cols as i32,
7187                    0.0,
7188                    out.as_mut_ptr().add(r0),
7189                    rows as i32,
7190                );
7191            }
7192            r0 += tr;
7193        }
7194    });
7195}
7196
7197fn q4t_matmat(
7198    bytes: &[u8],
7199    xs_all: &[f32],
7200    b: usize,
7201    rows: usize,
7202    cols: usize,
7203    out: &mut [f32],
7204    pool: Option<&Pool>,
7205) {
7206    debug_assert_eq!(out.len(), b * rows);
7207    let gpr = cols / GROUP_SIZE;
7208    // Wide batches ride the AMX like q8's qmatmat: on Apple silicon
7209    // the dequant-tile sgemm is an order above the SDOT row loop for
7210    // prefill shapes (imagegen DiT forwards are exactly this).
7211    #[cfg(target_os = "macos")]
7212    if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
7213        dequant_matmat_accel(
7214            &|r, dst| {
7215                for gi in 0..gpr {
7216                    let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
7217                    let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
7218                    for (k, &bb) in tile[2..].iter().enumerate() {
7219                        dst[gi * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s;
7220                        dst[gi * GROUP_SIZE + k * 2 + 1] = (((bb >> 4) & 0x0F) as f32 - 8.0) * s;
7221                    }
7222                }
7223            },
7224            xs_all,
7225            b,
7226            rows,
7227            cols,
7228            out,
7229            pool,
7230        );
7231        return;
7232    }
7233    let out_addr = SendMut(out.as_mut_ptr());
7234    if a8w8_enabled() {
7235        let acts: Vec<SplitAct> = (0..b)
7236            .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
7237            .collect();
7238        let acts = &acts;
7239        #[cfg(target_arch = "x86_64")]
7240        let blocked_ok = avx2_enabled() && blocked_enabled();
7241        #[cfg(target_arch = "aarch64")]
7242        let blocked_ok = sdot_enabled() && blocked_enabled();
7243        #[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))]
7244        let blocked_ok = false;
7245        let run = move |start: usize, end: usize| {
7246            for r in start..end {
7247                let mut bi = 0usize;
7248                #[cfg(target_arch = "aarch64")]
7249                if blocked_ok {
7250                    while bi + 4 <= acts.len() {
7251                        let xs = [
7252                            acts[bi].xq.as_slice(),
7253                            acts[bi + 1].xq.as_slice(),
7254                            acts[bi + 2].xq.as_slice(),
7255                            acts[bi + 3].xq.as_slice(),
7256                        ];
7257                        let d = unsafe { dot_q4t_row_1x4_sdot(bytes, r, gpr, xs) };
7258                        for k in 0..4 {
7259                            let act = &acts[bi + k];
7260                            let mut acc = d[k] * act.sx;
7261                            for &(j, xv) in &act.outliers {
7262                                let (w, sc) = q4t_outlier(bytes, r, gpr, j);
7263                                acc += w * sc * xv;
7264                            }
7265                            // SAFETY: disjoint (bi, r) cells per worker.
7266                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
7267                        }
7268                        bi += 4;
7269                    }
7270                }
7271                #[cfg(target_arch = "x86_64")]
7272                if blocked_ok {
7273                    while bi + 4 <= acts.len() {
7274                        let xs = [
7275                            acts[bi].xq.as_slice(),
7276                            acts[bi + 1].xq.as_slice(),
7277                            acts[bi + 2].xq.as_slice(),
7278                            acts[bi + 3].xq.as_slice(),
7279                        ];
7280                        let d = unsafe {
7281                            if vnni_tiles_enabled() {
7282                                dot_q4t_row_1x4_vnni(bytes, r, gpr, xs)
7283                            } else {
7284                                dot_q4t_row_1x4_avx2(bytes, r, gpr, xs)
7285                            }
7286                        };
7287                        for k in 0..4 {
7288                            let act = &acts[bi + k];
7289                            let mut acc = d[k] * act.sx;
7290                            for &(j, xv) in &act.outliers {
7291                                let (w, sc) = q4t_outlier(bytes, r, gpr, j);
7292                                acc += w * sc * xv;
7293                            }
7294                            // SAFETY: disjoint (bi, r) cells per worker.
7295                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
7296                        }
7297                        bi += 4;
7298                    }
7299                }
7300                let _ = blocked_ok;
7301                while bi < acts.len() {
7302                    let act = &acts[bi];
7303                    let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
7304                    for &(j, xv) in &act.outliers {
7305                        let (w, s) = q4t_outlier(bytes, r, gpr, j);
7306                        acc += w * s * xv;
7307                    }
7308                    // SAFETY: disjoint (bi, r) cells per worker range.
7309                    unsafe { *out_addr.at(bi * rows + r) = acc };
7310                    bi += 1;
7311                }
7312            }
7313        };
7314        dispatch_rows(pool, rows, &run);
7315        return;
7316    }
7317    let run = move |start: usize, end: usize| {
7318        for r in start..end {
7319            for bi in 0..b {
7320                let x = &xs_all[bi * cols..(bi + 1) * cols];
7321                // SAFETY: disjoint (bi, r) cells per worker range.
7322                unsafe { *out_addr.at(bi * rows + r) = q4t_row_exact(bytes, r, gpr, x) };
7323            }
7324        }
7325    };
7326    dispatch_rows(pool, rows, &run);
7327}
7328
7329// ── q1 (dtype 12): binary weights, [f16 scale][4B sign bits] per
7330// 32-group tile. The kernel family mirrors q4_tiled: one sequential
7331// stream of 6-byte tiles, per-tile integer dot × scale, exact outlier
7332// correction (A8W8 contract), exact scalar path under CMF_SDOT=0. ──
7333
7334/// Per-32-group sums of the quantized activation — the ±1 identity's
7335/// shared half: `dot = −2·sdot(mask, x) − gsum[g]`, computed ONCE per
7336/// matvec and reused by every row.
7337fn q1_group_sums(xq: &[i8], gpr: usize) -> Vec<i32> {
7338    (0..gpr)
7339        .map(|gi| {
7340            xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE]
7341                .iter()
7342                .map(|&v| v as i32)
7343                .sum()
7344        })
7345        .collect()
7346}
7347
7348/// One q1 row via the A8W8 int8 path — mask-SDOT on ARM (no ±1
7349/// expansion at all), scalar bit loop elsewhere (AVX2 queued with the
7350/// x86 pass).
7351#[inline]
7352#[allow(unreachable_code)]
7353/// AVX2 q1 row via the same ±1 identity as the ARM sdot kernel: the
7354/// sign bits expand to a {0, −1} byte mask through shuffle+cmpeq, the
7355/// masked activation sums through maddubs(1, x&mask), and
7356/// `dot = −(2·masked_sum + Σx_group)` — bit-identical integer math.
7357#[cfg(target_arch = "x86_64")]
7358#[target_feature(enable = "avx2")]
7359unsafe fn dot_q1_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
7360    // SAFETY: callers uphold the 6B-tile and xq/gsum length contracts.
7361    unsafe {
7362        use core::arch::x86_64::*;
7363        // Byte j of the mask must replicate bits-byte j/8.
7364        let expand = _mm256_setr_epi8(
7365            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,
7366            3, 3, 3,
7367        );
7368        let bitsel = _mm256_setr_epi8(
7369            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
7370            -128, 1, 2, 4, 8, 16, 32, 64, -128,
7371        );
7372        let ones8 = _mm256_set1_epi8(1);
7373        let ones16 = _mm256_set1_epi16(1);
7374        let mut acc = 0f32;
7375        for gi in 0..gpr {
7376            let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
7377            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
7378            let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
7379            let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
7380            let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
7381            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7382            let sel = _mm256_and_si256(x, mask);
7383            // Σ of selected i8 lanes: maddubs(1u8, sel_i8) pairs → madd.
7384            let p16 = _mm256_maddubs_epi16(ones8, sel);
7385            let d32 = _mm256_madd_epi16(p16, ones16);
7386            let hi128 = _mm256_extracti128_si256::<1>(d32);
7387            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
7388            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7389            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7390            let msum = _mm_cvtsi128_si32(s32);
7391            // The and-select keeps x UN-negated (unlike ARM's −1-mask
7392            // sdot): d = Σ_set − Σ_unset = 2·Σ_set − Σ_all.
7393            let d = 2 * msum - gsum[gi];
7394            acc += d as f32 * s;
7395        }
7396        acc
7397    }
7398}
7399
7400/// VNNI twin of `dot_q1_row_avx2`: the masked-select sum goes through
7401/// one `vpdpbusd(1u8, sel)` (see `dpbusd_hsum` — bit-identical).
7402#[cfg(target_arch = "x86_64")]
7403#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
7404unsafe fn dot_q1_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
7405    // SAFETY: callers uphold the 6B-tile and xq/gsum length contracts.
7406    unsafe {
7407        use core::arch::x86_64::*;
7408        let expand = _mm256_setr_epi8(
7409            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,
7410            3, 3, 3,
7411        );
7412        let bitsel = _mm256_setr_epi8(
7413            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
7414            -128, 1, 2, 4, 8, 16, 32, 64, -128,
7415        );
7416        let ones8 = _mm256_set1_epi8(1);
7417        let mut acc = 0f32;
7418        for gi in 0..gpr {
7419            let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
7420            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
7421            let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
7422            let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
7423            let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
7424            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7425            let msum = dpbusd_hsum(ones8, _mm256_and_si256(x, mask));
7426            let d = 2 * msum - gsum[gi];
7427            acc += d as f32 * s;
7428        }
7429        acc
7430    }
7431}
7432
7433/// VNNI twin of `dot_q1_row_1x4_avx2` (see `dpbusd_hsum`).
7434#[cfg(target_arch = "x86_64")]
7435#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
7436unsafe fn dot_q1_row_1x4_vnni(
7437    bytes: &[u8],
7438    r: usize,
7439    gpr: usize,
7440    xs: [&[i8]; 4],
7441    gsums: [&[i32]; 4],
7442) -> [f32; 4] {
7443    // SAFETY: callers uphold the 6B-tile and xq/gsum length contracts.
7444    unsafe {
7445        use core::arch::x86_64::*;
7446        let expand = _mm256_setr_epi8(
7447            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,
7448            3, 3, 3,
7449        );
7450        let bitsel = _mm256_setr_epi8(
7451            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
7452            -128, 1, 2, 4, 8, 16, 32, 64, -128,
7453        );
7454        let ones8 = _mm256_set1_epi8(1);
7455        let mut acc = [0f32; 4];
7456        for gi in 0..gpr {
7457            let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
7458            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
7459            let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
7460            let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
7461            let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
7462            for (k, xq) in xs.iter().enumerate() {
7463                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7464                let msum = dpbusd_hsum(ones8, _mm256_and_si256(x, mask));
7465                let d = 2 * msum - gsums[k][gi];
7466                acc[k] += d as f32 * s;
7467            }
7468        }
7469        acc
7470    }
7471}
7472
7473/// The blocked 1×4 flavor: the expanded bit mask serves four activation
7474/// streams per group (mask build once, four select+reduce chains).
7475#[cfg(target_arch = "x86_64")]
7476#[target_feature(enable = "avx2")]
7477unsafe fn dot_q1_row_1x4_avx2(
7478    bytes: &[u8],
7479    r: usize,
7480    gpr: usize,
7481    xs: [&[i8]; 4],
7482    gsums: [&[i32]; 4],
7483) -> [f32; 4] {
7484    // SAFETY: callers uphold the 6B-tile and xq/gsum length contracts.
7485    unsafe {
7486        use core::arch::x86_64::*;
7487        let expand = _mm256_setr_epi8(
7488            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,
7489            3, 3, 3,
7490        );
7491        let bitsel = _mm256_setr_epi8(
7492            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
7493            -128, 1, 2, 4, 8, 16, 32, 64, -128,
7494        );
7495        let ones8 = _mm256_set1_epi8(1);
7496        let ones16 = _mm256_set1_epi16(1);
7497        let mut acc = [0f32; 4];
7498        for gi in 0..gpr {
7499            let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
7500            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
7501            let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
7502            let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
7503            let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
7504            for (k, xq) in xs.iter().enumerate() {
7505                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7506                let sel = _mm256_and_si256(x, mask);
7507                let p16 = _mm256_maddubs_epi16(ones8, sel);
7508                let d32 = _mm256_madd_epi16(p16, ones16);
7509                let hi128 = _mm256_extracti128_si256::<1>(d32);
7510                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
7511                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7512                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7513                let msum = _mm_cvtsi128_si32(s32);
7514                let d = 2 * msum - gsums[k][gi];
7515                acc[k] += d as f32 * s;
7516            }
7517        }
7518        acc
7519    }
7520}
7521
7522#[allow(unreachable_code)]
7523fn dot_q1_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
7524    #[cfg(target_arch = "aarch64")]
7525    unsafe {
7526        return dot_q1_row_sdot(bytes, r, gpr, xq, gsum);
7527    }
7528    #[cfg(target_arch = "x86_64")]
7529    if avx2_enabled() {
7530        unsafe {
7531            if vnni_tiles_enabled() {
7532                return dot_q1_row_vnni(bytes, r, gpr, xq, gsum);
7533            }
7534            return dot_q1_row_avx2(bytes, r, gpr, xq, gsum);
7535        }
7536    }
7537    let _ = gsum;
7538    let mut acc = 0f32;
7539    for gi in 0..gpr {
7540        let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
7541        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
7542        let mut d = 0i32;
7543        for (j, &b) in tile[2..].iter().enumerate() {
7544            for k in 0..8 {
7545                let w = ((b >> k) & 1) as i32 * 2 - 1;
7546                d += w * xq[gi * GROUP_SIZE + j * 8 + k] as i32;
7547            }
7548        }
7549        acc += d as f32 * s;
7550    }
7551    acc
7552}
7553
7554/// SDOT q1 row via the ±1 identity: the vtst mask (0xFF where the bit
7555/// is set, i.e. −1 as i8) feeds `sdot` DIRECTLY — no expansion to ±1
7556/// lanes at all — and `dot = −(2·sdot(mask, x) + Σx_group)`, with the
7557/// per-group activation sums shared across every row of the matvec.
7558/// Four tiles (128 weights) per iteration: integer dots reduce through
7559/// a vpaddq tree into ONE i32x4 that meets its four scales in a single
7560/// fused f32 multiply-add. Integer math throughout — bit-identical to
7561/// the scalar ±1 reference.
7562#[cfg(target_arch = "aarch64")]
7563#[target_feature(enable = "neon,dotprod")]
7564unsafe fn dot_q1_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
7565    // SAFETY: callers uphold slice-length contracts (6B tile per group,
7566    // xq.len() == gpr·GROUP_SIZE, gsum.len() == gpr).
7567    unsafe {
7568        use core::arch::aarch64::*;
7569        use core::arch::asm;
7570        const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
7571        let m = vld1q_u8(MASKS.as_ptr());
7572        // One tile's −Σ_set(x) as an UNREDUCED i32x4 (two mask-sdots).
7573        macro_rules! tile_dot {
7574            ($t:expr, $x:expr) => {{
7575                let v0 = vcombine_u8(vdup_n_u8(*$t.add(2)), vdup_n_u8(*$t.add(3)));
7576                let v1 = vcombine_u8(vdup_n_u8(*$t.add(4)), vdup_n_u8(*$t.add(5)));
7577                let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
7578                let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
7579                let x0 = vld1q_s8($x);
7580                let x1 = vld1q_s8($x.add(16));
7581                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
7582                asm!(
7583                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
7584                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
7585                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
7586                    w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
7587                    options(pure, nomem, nostack),
7588                );
7589                vaddq_s32(a0, a1)
7590            }};
7591        }
7592        // TBL unpack over PAIR loads: one vld1q covers two 6B tiles
7593        // ([s s b b b b][s s b b b b] + 4B slack), TBL replicates each
7594        // bit-byte across 8 lanes for vtst, and the four scales gather
7595        // through tbl2 into one fcvtl — the 16 ld1r broadcast loads and
7596        // 4 branchy software f16 conversions per 128 weights (the
7597        // measured load-port wall of this kernel) become 2 vector
7598        // loads + 9 table lookups. Integer math order is unchanged —
7599        // bit-identical results (FCVTL is exact on every f16).
7600        const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
7601        const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
7602        const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
7603        const IW11: [u8; 16] = [
7604            10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
7605        ];
7606        const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
7607        let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
7608        let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
7609        let isc = vld1_u8(ISC.as_ptr());
7610        // One tile's −Σ_set(x) from a TBL-unpacked pair load.
7611        macro_rules! tile_dot_tbl {
7612            ($ld:expr, $i0:expr, $i1:expr, $x:expr) => {{
7613                let w0 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i0), m));
7614                let w1 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i1), m));
7615                let x0 = vld1q_s8($x);
7616                let x1 = vld1q_s8($x.add(16));
7617                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
7618                asm!(
7619                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
7620                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
7621                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
7622                    w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
7623                    options(pure, nomem, nostack),
7624                );
7625                vaddq_s32(a0, a1)
7626            }};
7627        }
7628        let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
7629        let row_base = r * gpr * Q1_TILE;
7630        let abs_end = bytes.len();
7631        let xp = xq.as_ptr();
7632        let gp = gsum.as_ptr();
7633        let mut accv = vdupq_n_f32(0.0);
7634        let mut gi = 0;
7635        // The second pair load reads 4B past tile gi+3 — stay inside
7636        // the payload slice (only the file's final tiles fall back).
7637        while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
7638            let t0 = base.add(gi * Q1_TILE);
7639            let ld_a = vld1q_u8(t0);
7640            let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
7641            let d0 = tile_dot_tbl!(ld_a, iw00, iw01, xp.add(gi * GROUP_SIZE));
7642            let d1 = tile_dot_tbl!(ld_a, iw10, iw11, xp.add((gi + 1) * GROUP_SIZE));
7643            let d2 = tile_dot_tbl!(ld_b, iw00, iw01, xp.add((gi + 2) * GROUP_SIZE));
7644            let d3 = tile_dot_tbl!(ld_b, iw10, iw11, xp.add((gi + 3) * GROUP_SIZE));
7645            // [−Σ0, −Σ1, −Σ2, −Σ3] → dots = −(2·Σset_neg + gsum)
7646            let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
7647            let g = vld1q_s32(gp.add(gi));
7648            let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
7649            let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
7650            let scf: float32x4_t;
7651            asm!(
7652                "fcvtl {o:v}.4s, {i:v}.4h",
7653                o = out(vreg) scf, i = in(vreg) sc16,
7654                options(pure, nomem, nostack),
7655            );
7656            accv = vfmaq_f32(accv, vcvtq_f32_s32(dots), scf);
7657            gi += 4;
7658        }
7659        let mut acc = vaddvq_f32(accv);
7660        while gi < gpr {
7661            let t = base.add(gi * Q1_TILE);
7662            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
7663            let d = vaddvq_s32(tile_dot!(t, xp.add(gi * GROUP_SIZE)));
7664            acc += (-(2 * d + *gp.add(gi))) as f32 * s;
7665            gi += 1;
7666        }
7667        acc
7668    }
7669}
7670
7671/// Blocked q1 1×4: one TBL unpack of the tile pair serves FOUR
7672/// activation streams (prefill amortization — the same idea as the
7673/// AVX2 twin; per stream the group order, fma order and tail match the
7674/// single-row kernel exactly, so batch == matvec bit-for-bit).
7675#[cfg(target_arch = "aarch64")]
7676#[target_feature(enable = "neon,dotprod")]
7677unsafe fn dot_q1_row_1x4_sdot(
7678    bytes: &[u8],
7679    r: usize,
7680    gpr: usize,
7681    xs: [&[i8]; 4],
7682    gs: [&[i32]; 4],
7683) -> [f32; 4] {
7684    // SAFETY: same slice-length contracts as `dot_q1_row_sdot`, ×4.
7685    unsafe {
7686        use core::arch::aarch64::*;
7687        use core::arch::asm;
7688        const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
7689        const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
7690        const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
7691        const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
7692        const IW11: [u8; 16] = [
7693            10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
7694        ];
7695        const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
7696        let m = vld1q_u8(MASKS.as_ptr());
7697        let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
7698        let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
7699        let isc = vld1_u8(ISC.as_ptr());
7700        macro_rules! sdot2 {
7701            ($w0:expr, $w1:expr, $x:expr) => {{
7702                let x0 = vld1q_s8($x);
7703                let x1 = vld1q_s8($x.add(16));
7704                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
7705                asm!(
7706                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
7707                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
7708                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
7709                    w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
7710                    options(pure, nomem, nostack),
7711                );
7712                vaddq_s32(a0, a1)
7713            }};
7714        }
7715        let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
7716        let row_base = r * gpr * Q1_TILE;
7717        let abs_end = bytes.len();
7718        let mut accv = [vdupq_n_f32(0.0); 4];
7719        let mut gi = 0;
7720        while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
7721            let t0 = base.add(gi * Q1_TILE);
7722            let ld_a = vld1q_u8(t0);
7723            let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
7724            // Unpack ONCE — eight ±mask vectors serve all four streams.
7725            let w00 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw00), m));
7726            let w01 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw01), m));
7727            let w10 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw10), m));
7728            let w11 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw11), m));
7729            let w20 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw00), m));
7730            let w21 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw01), m));
7731            let w30 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw10), m));
7732            let w31 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw11), m));
7733            let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
7734            let scf: float32x4_t;
7735            asm!(
7736                "fcvtl {o:v}.4s, {i:v}.4h",
7737                o = out(vreg) scf, i = in(vreg) sc16,
7738                options(pure, nomem, nostack),
7739            );
7740            for k in 0..4 {
7741                let xp = xs[k].as_ptr();
7742                let d0 = sdot2!(w00, w01, xp.add(gi * GROUP_SIZE));
7743                let d1 = sdot2!(w10, w11, xp.add((gi + 1) * GROUP_SIZE));
7744                let d2 = sdot2!(w20, w21, xp.add((gi + 2) * GROUP_SIZE));
7745                let d3 = sdot2!(w30, w31, xp.add((gi + 3) * GROUP_SIZE));
7746                let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
7747                let g = vld1q_s32(gs[k].as_ptr().add(gi));
7748                let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
7749                accv[k] = vfmaq_f32(accv[k], vcvtq_f32_s32(dots), scf);
7750            }
7751            gi += 4;
7752        }
7753        let mut acc = [
7754            vaddvq_f32(accv[0]),
7755            vaddvq_f32(accv[1]),
7756            vaddvq_f32(accv[2]),
7757            vaddvq_f32(accv[3]),
7758        ];
7759        while gi < gpr {
7760            let t = base.add(gi * Q1_TILE);
7761            let sc = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
7762            let v0 = vcombine_u8(vdup_n_u8(*t.add(2)), vdup_n_u8(*t.add(3)));
7763            let v1 = vcombine_u8(vdup_n_u8(*t.add(4)), vdup_n_u8(*t.add(5)));
7764            let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
7765            let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
7766            for k in 0..4 {
7767                let d = vaddvq_s32(sdot2!(w0, w1, xs[k].as_ptr().add(gi * GROUP_SIZE)));
7768                acc[k] += (-(2 * d + *gs[k].as_ptr().add(gi))) as f32 * sc;
7769            }
7770            gi += 1;
7771        }
7772        acc
7773    }
7774}
7775
7776/// (weight ±1, scale) of one q1 element — the exact outlier term.
7777#[inline]
7778fn q1_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
7779    let gi = j / GROUP_SIZE;
7780    let k = j % GROUP_SIZE;
7781    let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
7782    let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
7783    let bit = (tile[2 + k / 8] >> (k % 8)) & 1;
7784    ((bit as i32 * 2 - 1) as f32, s)
7785}
7786
7787/// Exact scalar q1 row (CMF_SDOT=0 contract).
7788#[inline]
7789fn q1_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
7790    let mut acc = 0f32;
7791    for gi in 0..gpr {
7792        let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
7793        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
7794        let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
7795        let mut ga = 0f32;
7796        for (j, &b) in tile[2..].iter().enumerate() {
7797            for k in 0..8 {
7798                ga += (((b >> k) & 1) as f32 * 2.0 - 1.0) * xg[j * 8 + k];
7799            }
7800        }
7801        acc += ga * s;
7802    }
7803    acc
7804}
7805
7806/// One binary-weight row range via A8W8 (the body of `bin_matvec`'s hot
7807/// loop, extracted so multi-matrix jobs drive the same kernel).
7808#[allow(clippy::too_many_arguments)]
7809fn bin_range_a8w8(
7810    w: BinW,
7811    gpr: usize,
7812    act: &SplitAct,
7813    gsum: &[i32],
7814    out: SendMut,
7815    start: usize,
7816    end: usize,
7817) {
7818    for r in start..end {
7819        let mut acc = w.dot_i8(r, gpr, &act.xq, gsum) * act.sx;
7820        for &(j, xv) in &act.outliers {
7821            let (wv, s) = w.outlier(r, gpr, j);
7822            acc += wv * s * xv;
7823        }
7824        // SAFETY: disjoint row ranges per worker.
7825        unsafe { *out.at(r) = acc };
7826    }
7827}
7828
7829/// Exact-scalar binary-weight row range (CMF_SDOT=0 contract).
7830fn bin_range_f32(w: BinW, gpr: usize, x: &[f32], out: SendMut, start: usize, end: usize) {
7831    for r in start..end {
7832        // SAFETY: disjoint row ranges per worker.
7833        unsafe { *out.at(r) = w.row_exact(r, gpr, x) };
7834    }
7835}
7836
7837/// A binary (±s) weight in either of its layouts: `q1` 6-byte tiles, or
7838/// `q1g` planes with one scale per 128-group. The row drivers — matvec,
7839/// matvec2, matmat, the multi-projection job and the fused gate/up pair —
7840/// are written once against this; only the per-row kernels differ, and
7841/// the match costs one predictable branch per row. `gpr` is always the
7842/// row's count of 32-tiles (cols / 32): the activation group sums and the
7843/// tile walk are shared by both layouts.
7844#[derive(Clone, Copy)]
7845pub(crate) enum BinW<'a> {
7846    Q1(&'a [u8]),
7847    Q1g {
7848        bytes: &'a [u8],
7849        /// Byte offset of the f16 scale plane (= rows · cols/8).
7850        sc_off: usize,
7851        /// Row count: the q1 kernels treat the payload's final row
7852        /// specially (see `dot_q1g_row_sdot`), and q1g mirrors that.
7853        rows: usize,
7854    },
7855}
7856
7857impl<'a> BinW<'a> {
7858    fn q1g(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
7859        Self::Q1g {
7860            bytes,
7861            sc_off: cortiq_core::quant::q1g_scale_offset(rows, cols),
7862            rows,
7863        }
7864    }
7865
7866    #[inline]
7867    fn dot_i8(self, r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
7868        match self {
7869            Self::Q1(b) => dot_q1_row_i8(b, r, gpr, xq, gsum),
7870            Self::Q1g {
7871                bytes,
7872                sc_off,
7873                rows,
7874            } => dot_q1g_row_i8(bytes, sc_off, rows, r, gpr, xq, gsum),
7875        }
7876    }
7877
7878    #[inline]
7879    fn row_exact(self, r: usize, gpr: usize, x: &[f32]) -> f32 {
7880        match self {
7881            Self::Q1(b) => q1_row_exact(b, r, gpr, x),
7882            Self::Q1g { bytes, sc_off, .. } => q1g_row_exact(bytes, sc_off, r, gpr, x),
7883        }
7884    }
7885
7886    #[inline]
7887    fn outlier(self, r: usize, gpr: usize, j: usize) -> (f32, f32) {
7888        match self {
7889            Self::Q1(b) => q1_outlier(b, r, gpr, j),
7890            Self::Q1g { bytes, sc_off, .. } => q1g_outlier(bytes, sc_off, r, gpr, j),
7891        }
7892    }
7893
7894    /// Blocked 1×4 row (one unpack, four activation streams).
7895    ///
7896    /// # Safety
7897    /// The SDOT/dotprod target feature must be present (`sdot_enabled`).
7898    #[cfg(target_arch = "aarch64")]
7899    #[inline]
7900    unsafe fn dot_1x4_sdot(
7901        self,
7902        r: usize,
7903        gpr: usize,
7904        xs: [&[i8]; 4],
7905        gs: [&[i32]; 4],
7906    ) -> [f32; 4] {
7907        // SAFETY: forwarded caller contract.
7908        unsafe {
7909            match self {
7910                Self::Q1(b) => dot_q1_row_1x4_sdot(b, r, gpr, xs, gs),
7911                Self::Q1g {
7912                    bytes,
7913                    sc_off,
7914                    rows,
7915                } => dot_q1g_row_1x4_sdot(bytes, sc_off, rows, r, gpr, xs, gs),
7916            }
7917        }
7918    }
7919
7920    /// Blocked 1×4 row on x86 (VNNI when the tile switch is on).
7921    ///
7922    /// # Safety
7923    /// AVX2 must be present (`avx2_enabled`).
7924    #[cfg(target_arch = "x86_64")]
7925    #[inline]
7926    unsafe fn dot_1x4_x86(self, r: usize, gpr: usize, xs: [&[i8]; 4], gs: [&[i32]; 4]) -> [f32; 4] {
7927        // SAFETY: forwarded caller contract.
7928        unsafe {
7929            match (self, vnni_tiles_enabled()) {
7930                (Self::Q1(b), true) => dot_q1_row_1x4_vnni(b, r, gpr, xs, gs),
7931                (Self::Q1(b), false) => dot_q1_row_1x4_avx2(b, r, gpr, xs, gs),
7932                (Self::Q1g { bytes, sc_off, .. }, true) => {
7933                    dot_q1g_row_1x4_vnni(bytes, sc_off, r, gpr, xs, gs)
7934                }
7935                (Self::Q1g { bytes, sc_off, .. }, false) => {
7936                    dot_q1g_row_1x4_avx2(bytes, sc_off, r, gpr, xs, gs)
7937                }
7938            }
7939        }
7940    }
7941}
7942
7943// ── q1g (dtype 20): binary weights, ONE f16 scale per 128-group, in sign
7944// and scale PLANES (`[bits: rows·cols/8][f16: rows·cols/128]`). Same
7945// kernel family as q1 and — deliberately — the same float accumulation
7946// order (one `acc += d·s` per 32-tile, the group's scale reused four
7947// times), so a q1g repack of a q1 tensor whose tile scales agree gives
7948// the same f32 outputs bit for bit on every path. The win is the stream:
7949// 18 bytes per 128 weights instead of 24, the bits 16-byte aligned. ──
7950
7951/// Sign bits of q1g tile `t` (flat 32-tile index `r·gpr + gi`).
7952#[inline]
7953fn q1g_bits(bytes: &[u8], t: usize) -> u32 {
7954    u32::from_le_bytes([
7955        bytes[t * 4],
7956        bytes[t * 4 + 1],
7957        bytes[t * 4 + 2],
7958        bytes[t * 4 + 3],
7959    ])
7960}
7961
7962/// Raw f16 bits of the scale of the 128-group holding q1g tile `t`.
7963#[inline]
7964fn q1g_scale_bits(bytes: &[u8], sc_off: usize, t: usize) -> u16 {
7965    let o = sc_off + (t / 4) * 2;
7966    u16::from_le_bytes([bytes[o], bytes[o + 1]])
7967}
7968
7969/// One q1g row via the A8W8 int8 path (dispatch twin of `dot_q1_row_i8`).
7970#[allow(unreachable_code)]
7971#[allow(clippy::too_many_arguments)]
7972fn dot_q1g_row_i8(
7973    bytes: &[u8],
7974    sc_off: usize,
7975    rows: usize,
7976    r: usize,
7977    gpr: usize,
7978    xq: &[i8],
7979    gsum: &[i32],
7980) -> f32 {
7981    #[cfg(target_arch = "aarch64")]
7982    unsafe {
7983        return dot_q1g_row_sdot(bytes, sc_off, rows, r, gpr, xq, gsum);
7984    }
7985    #[cfg(target_arch = "x86_64")]
7986    if avx2_enabled() {
7987        unsafe {
7988            if vnni_tiles_enabled() {
7989                return dot_q1g_row_vnni(bytes, sc_off, r, gpr, xq, gsum);
7990            }
7991            return dot_q1g_row_avx2(bytes, sc_off, r, gpr, xq, gsum);
7992        }
7993    }
7994    let _ = (gsum, rows);
7995    let mut acc = 0f32;
7996    for gi in 0..gpr {
7997        let t = r * gpr + gi;
7998        let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, t));
7999        let bits = q1g_bits(bytes, t);
8000        let mut d = 0i32;
8001        for k in 0..GROUP_SIZE {
8002            let w = ((bits >> k) & 1) as i32 * 2 - 1;
8003            d += w * xq[gi * GROUP_SIZE + k] as i32;
8004        }
8005        acc += d as f32 * s;
8006    }
8007    acc
8008}
8009
8010/// Four tiles' unreduced `i32×8` partial sums → `[Σ0, Σ1, Σ2, Σ3]` in one
8011/// hadd tree: 3 hadds + 1 add for four tiles where the per-tile reduce of
8012/// the q1 kernels costs 7 ops each. Integer math, so any order gives the
8013/// q1 kernel's per-tile sums exactly.
8014#[cfg(target_arch = "x86_64")]
8015#[target_feature(enable = "avx2")]
8016#[inline]
8017unsafe fn q1g_hsum4(
8018    d0: core::arch::x86_64::__m256i,
8019    d1: core::arch::x86_64::__m256i,
8020    d2: core::arch::x86_64::__m256i,
8021    d3: core::arch::x86_64::__m256i,
8022) -> core::arch::x86_64::__m128i {
8023    // SAFETY: pure register math.
8024    unsafe {
8025        use core::arch::x86_64::*;
8026        let h = _mm256_hadd_epi32(_mm256_hadd_epi32(d0, d1), _mm256_hadd_epi32(d2, d3));
8027        _mm_add_epi32(_mm256_castsi256_si128(h), _mm256_extracti128_si256::<1>(h))
8028    }
8029}
8030
8031/// One 128-group's epilogue, in q1's float order: `d = 2·Σset − Σx` per
8032/// tile (exact i32 → exact f32), then `acc += d·s` tile by tile — four
8033/// separately rounded multiply/add pairs, the q1 kernels' sequence.
8034#[cfg(target_arch = "x86_64")]
8035#[target_feature(enable = "avx2")]
8036#[inline]
8037unsafe fn q1g_group_acc(
8038    acc: f32,
8039    msum4: core::arch::x86_64::__m128i,
8040    gsum4: &[i32],
8041    s: f32,
8042) -> f32 {
8043    // SAFETY: gsum4 holds the group's four tile sums (callers slice it).
8044    unsafe {
8045        use core::arch::x86_64::*;
8046        let g = _mm_loadu_si128(gsum4.as_ptr() as *const __m128i);
8047        let d = _mm_cvtepi32_ps(_mm_sub_epi32(_mm_slli_epi32::<1>(msum4), g));
8048        let mut f = [0f32; 4];
8049        _mm_storeu_ps(f.as_mut_ptr(), d);
8050        let mut acc = acc;
8051        for v in f {
8052            acc += v * s;
8053        }
8054        acc
8055    }
8056}
8057
8058/// AVX2 q1g row: `dot_q1_row_avx2`'s mask/select/maddubs per tile, but the
8059/// four tiles of a 128-group reduce together (`q1g_hsum4`) and the scale
8060/// converts once per group — bit-identical results (integer sums, then
8061/// q1's per-tile float sequence), fewer horizontal ops and f16 decodes.
8062#[cfg(target_arch = "x86_64")]
8063#[target_feature(enable = "avx2")]
8064unsafe fn dot_q1g_row_avx2(
8065    bytes: &[u8],
8066    sc_off: usize,
8067    r: usize,
8068    gpr: usize,
8069    xq: &[i8],
8070    gsum: &[i32],
8071) -> f32 {
8072    // SAFETY: callers uphold the plane-length and xq/gsum contracts
8073    // (gpr % 4 == 0, xq.len() == gpr·32, gsum.len() == gpr).
8074    unsafe {
8075        use core::arch::x86_64::*;
8076        let expand = _mm256_setr_epi8(
8077            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,
8078            3, 3, 3,
8079        );
8080        let bitsel = _mm256_setr_epi8(
8081            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
8082            -128, 1, 2, 4, 8, 16, 32, 64, -128,
8083        );
8084        let ones8 = _mm256_set1_epi8(1);
8085        let ones16 = _mm256_set1_epi16(1);
8086        let bp = bytes.as_ptr().add(r * gpr * 4);
8087        let mut acc = 0f32;
8088        let mut gi = 0;
8089        while gi < gpr {
8090            let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
8091            let mut d = [_mm256_setzero_si256(); 4];
8092            for (k, dk) in d.iter_mut().enumerate() {
8093                let t = gi + k;
8094                let bits = (bp.add(t * 4) as *const u32).read_unaligned();
8095                let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
8096                let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
8097                let x = _mm256_loadu_si256(xq.as_ptr().add(t * GROUP_SIZE) as *const __m256i);
8098                let sel = _mm256_and_si256(x, mask);
8099                *dk = _mm256_madd_epi16(_mm256_maddubs_epi16(ones8, sel), ones16);
8100            }
8101            let m4 = q1g_hsum4(d[0], d[1], d[2], d[3]);
8102            acc = q1g_group_acc(acc, m4, &gsum[gi..gi + 4], s);
8103            gi += 4;
8104        }
8105        acc
8106    }
8107}
8108
8109/// VNNI twin of `dot_q1g_row_avx2`: `vpdpbusd(1u8, sel)` per tile (the
8110/// same i32 lanes as maddubs+madd — bounded products), same group reduce.
8111#[cfg(target_arch = "x86_64")]
8112#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
8113unsafe fn dot_q1g_row_vnni(
8114    bytes: &[u8],
8115    sc_off: usize,
8116    r: usize,
8117    gpr: usize,
8118    xq: &[i8],
8119    gsum: &[i32],
8120) -> f32 {
8121    // SAFETY: callers uphold the plane-length and xq/gsum contracts.
8122    unsafe {
8123        use core::arch::x86_64::*;
8124        let expand = _mm256_setr_epi8(
8125            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,
8126            3, 3, 3,
8127        );
8128        let bitsel = _mm256_setr_epi8(
8129            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
8130            -128, 1, 2, 4, 8, 16, 32, 64, -128,
8131        );
8132        let ones8 = _mm256_set1_epi8(1);
8133        let bp = bytes.as_ptr().add(r * gpr * 4);
8134        let mut acc = 0f32;
8135        let mut gi = 0;
8136        while gi < gpr {
8137            let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
8138            let mut d = [_mm256_setzero_si256(); 4];
8139            for (k, dk) in d.iter_mut().enumerate() {
8140                let t = gi + k;
8141                let bits = (bp.add(t * 4) as *const u32).read_unaligned();
8142                let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
8143                let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
8144                let x = _mm256_loadu_si256(xq.as_ptr().add(t * GROUP_SIZE) as *const __m256i);
8145                *dk = _mm256_dpbusd_epi32(_mm256_setzero_si256(), ones8, _mm256_and_si256(x, mask));
8146            }
8147            let m4 = q1g_hsum4(d[0], d[1], d[2], d[3]);
8148            acc = q1g_group_acc(acc, m4, &gsum[gi..gi + 4], s);
8149            gi += 4;
8150        }
8151        acc
8152    }
8153}
8154
8155/// Blocked 1×4 AVX2 q1g row: each tile's mask is built once and serves
8156/// four activation streams; per stream, the group reduce and float order
8157/// of `dot_q1g_row_avx2` (so batch == matvec bit for bit).
8158#[cfg(target_arch = "x86_64")]
8159#[target_feature(enable = "avx2")]
8160unsafe fn dot_q1g_row_1x4_avx2(
8161    bytes: &[u8],
8162    sc_off: usize,
8163    r: usize,
8164    gpr: usize,
8165    xs: [&[i8]; 4],
8166    gsums: [&[i32]; 4],
8167) -> [f32; 4] {
8168    // SAFETY: callers uphold the plane-length and xq/gsum contracts.
8169    unsafe {
8170        use core::arch::x86_64::*;
8171        let expand = _mm256_setr_epi8(
8172            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,
8173            3, 3, 3,
8174        );
8175        let bitsel = _mm256_setr_epi8(
8176            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
8177            -128, 1, 2, 4, 8, 16, 32, 64, -128,
8178        );
8179        let ones8 = _mm256_set1_epi8(1);
8180        let ones16 = _mm256_set1_epi16(1);
8181        let bp = bytes.as_ptr().add(r * gpr * 4);
8182        let mut acc = [0f32; 4];
8183        let mut gi = 0;
8184        while gi < gpr {
8185            let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
8186            // d[stream][tile]
8187            let mut d = [[_mm256_setzero_si256(); 4]; 4];
8188            for t in 0..4 {
8189                let bits = (bp.add((gi + t) * 4) as *const u32).read_unaligned();
8190                let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
8191                let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
8192                for (k, xq) in xs.iter().enumerate() {
8193                    let x = _mm256_loadu_si256(
8194                        xq.as_ptr().add((gi + t) * GROUP_SIZE) as *const __m256i
8195                    );
8196                    let sel = _mm256_and_si256(x, mask);
8197                    d[k][t] = _mm256_madd_epi16(_mm256_maddubs_epi16(ones8, sel), ones16);
8198                }
8199            }
8200            for k in 0..4 {
8201                let m4 = q1g_hsum4(d[k][0], d[k][1], d[k][2], d[k][3]);
8202                acc[k] = q1g_group_acc(acc[k], m4, &gsums[k][gi..gi + 4], s);
8203            }
8204            gi += 4;
8205        }
8206        acc
8207    }
8208}
8209
8210/// VNNI twin of `dot_q1g_row_1x4_avx2`.
8211#[cfg(target_arch = "x86_64")]
8212#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
8213unsafe fn dot_q1g_row_1x4_vnni(
8214    bytes: &[u8],
8215    sc_off: usize,
8216    r: usize,
8217    gpr: usize,
8218    xs: [&[i8]; 4],
8219    gsums: [&[i32]; 4],
8220) -> [f32; 4] {
8221    // SAFETY: callers uphold the plane-length and xq/gsum contracts.
8222    unsafe {
8223        use core::arch::x86_64::*;
8224        let expand = _mm256_setr_epi8(
8225            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,
8226            3, 3, 3,
8227        );
8228        let bitsel = _mm256_setr_epi8(
8229            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
8230            -128, 1, 2, 4, 8, 16, 32, 64, -128,
8231        );
8232        let ones8 = _mm256_set1_epi8(1);
8233        let bp = bytes.as_ptr().add(r * gpr * 4);
8234        let mut acc = [0f32; 4];
8235        let mut gi = 0;
8236        while gi < gpr {
8237            let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
8238            let mut d = [[_mm256_setzero_si256(); 4]; 4];
8239            for t in 0..4 {
8240                let bits = (bp.add((gi + t) * 4) as *const u32).read_unaligned();
8241                let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
8242                let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
8243                for (k, xq) in xs.iter().enumerate() {
8244                    let x = _mm256_loadu_si256(
8245                        xq.as_ptr().add((gi + t) * GROUP_SIZE) as *const __m256i
8246                    );
8247                    d[k][t] = _mm256_dpbusd_epi32(
8248                        _mm256_setzero_si256(),
8249                        ones8,
8250                        _mm256_and_si256(x, mask),
8251                    );
8252                }
8253            }
8254            for k in 0..4 {
8255                let m4 = q1g_hsum4(d[k][0], d[k][1], d[k][2], d[k][3]);
8256                acc[k] = q1g_group_acc(acc[k], m4, &gsums[k][gi..gi + 4], s);
8257            }
8258            gi += 4;
8259        }
8260        acc
8261    }
8262}
8263
8264/// SDOT q1g row: the `dot_q1_row_sdot` math on the plane layout. One
8265/// aligned 16-byte load carries a whole 128-group's signs (q1 needs two
8266/// pair loads and a TBL gather of four scales for the same weights), TBL
8267/// replicates each sign byte across 8 lanes for vtst, and the group's ONE
8268/// scale converts once (FCVTL, exact) and broadcasts. Lane order, the
8269/// vpaddq tree, the fma into the four lane accumulators and the tail are
8270/// the q1 kernel's — including its quirk: q1's pair load reads 4 bytes
8271/// past the 4th tile, so the payload's final four tiles (last row) take
8272/// the scalar tail there. q1g has no overread, but takes the same tail on
8273/// the last row to stay bit-identical to q1.
8274#[cfg(target_arch = "aarch64")]
8275#[target_feature(enable = "neon,dotprod")]
8276#[allow(clippy::too_many_arguments)]
8277unsafe fn dot_q1g_row_sdot(
8278    bytes: &[u8],
8279    sc_off: usize,
8280    rows: usize,
8281    r: usize,
8282    gpr: usize,
8283    xq: &[i8],
8284    gsum: &[i32],
8285) -> f32 {
8286    // SAFETY: callers uphold the plane-length and xq/gsum contracts
8287    // (gpr % 4 == 0, xq.len() == gpr·32, gsum.len() == gpr).
8288    unsafe {
8289        use core::arch::aarch64::*;
8290        use core::arch::asm;
8291        const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
8292        // Tile t of the group = sign bytes 4t..4t+4 of the 16-byte load.
8293        const IW: [[u8; 16]; 8] = [
8294            [0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1],
8295            [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3],
8296            [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5],
8297            [6, 6, 6, 6, 6, 6, 6, 6, 7, 7, 7, 7, 7, 7, 7, 7],
8298            [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9],
8299            [
8300                10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
8301            ],
8302            [
8303                12, 12, 12, 12, 12, 12, 12, 12, 13, 13, 13, 13, 13, 13, 13, 13,
8304            ],
8305            [
8306                14, 14, 14, 14, 14, 14, 14, 14, 15, 15, 15, 15, 15, 15, 15, 15,
8307            ],
8308        ];
8309        let m = vld1q_u8(MASKS.as_ptr());
8310        let iw: [uint8x16_t; 8] = core::array::from_fn(|i| vld1q_u8(IW[i].as_ptr()));
8311        macro_rules! tile_dot {
8312            ($w0:expr, $w1:expr, $x:expr) => {{
8313                let x0 = vld1q_s8($x);
8314                let x1 = vld1q_s8($x.add(16));
8315                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
8316                asm!(
8317                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
8318                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
8319                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
8320                    w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
8321                    options(pure, nomem, nostack),
8322                );
8323                vaddq_s32(a0, a1)
8324            }};
8325        }
8326        macro_rules! tbl_mask {
8327            ($ld:expr, $i:expr) => {
8328                vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, iw[$i]), m))
8329            };
8330        }
8331        let bp = bytes.as_ptr().add(r * gpr * 4);
8332        let xp = xq.as_ptr();
8333        let gp = gsum.as_ptr();
8334        let vec_end = if r + 1 == rows { gpr - 4 } else { gpr };
8335        let mut accv = vdupq_n_f32(0.0);
8336        let mut gi = 0;
8337        while gi + 4 <= vec_end {
8338            let ld = vld1q_u8(bp.add(gi * 4));
8339            let d0 = tile_dot!(tbl_mask!(ld, 0), tbl_mask!(ld, 1), xp.add(gi * GROUP_SIZE));
8340            let d1 = tile_dot!(
8341                tbl_mask!(ld, 2),
8342                tbl_mask!(ld, 3),
8343                xp.add((gi + 1) * GROUP_SIZE)
8344            );
8345            let d2 = tile_dot!(
8346                tbl_mask!(ld, 4),
8347                tbl_mask!(ld, 5),
8348                xp.add((gi + 2) * GROUP_SIZE)
8349            );
8350            let d3 = tile_dot!(
8351                tbl_mask!(ld, 6),
8352                tbl_mask!(ld, 7),
8353                xp.add((gi + 3) * GROUP_SIZE)
8354            );
8355            let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
8356            let g = vld1q_s32(gp.add(gi));
8357            let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
8358            let h = vdup_n_u16(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
8359            let scf: float32x4_t;
8360            asm!(
8361                "fcvtl {o:v}.4s, {i:v}.4h",
8362                o = out(vreg) scf, i = in(vreg) h,
8363                options(pure, nomem, nostack),
8364            );
8365            accv = vfmaq_f32(accv, vcvtq_f32_s32(dots), scf);
8366            gi += 4;
8367        }
8368        let mut acc = vaddvq_f32(accv);
8369        while gi < gpr {
8370            let t = bp.add(gi * 4);
8371            let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
8372            let v0 = vcombine_u8(vdup_n_u8(*t), vdup_n_u8(*t.add(1)));
8373            let v1 = vcombine_u8(vdup_n_u8(*t.add(2)), vdup_n_u8(*t.add(3)));
8374            let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
8375            let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
8376            let d = vaddvq_s32(tile_dot!(w0, w1, xp.add(gi * GROUP_SIZE)));
8377            acc += (-(2 * d + *gp.add(gi))) as f32 * s;
8378            gi += 1;
8379        }
8380        acc
8381    }
8382}
8383
8384/// Blocked q1g 1×4: one unpack of the group serves four activation
8385/// streams; per stream bit-identical to `dot_q1g_row_sdot` (and so to the
8386/// q1 1×4 kernel, same last-row tail).
8387#[cfg(target_arch = "aarch64")]
8388#[target_feature(enable = "neon,dotprod")]
8389#[allow(clippy::too_many_arguments)]
8390unsafe fn dot_q1g_row_1x4_sdot(
8391    bytes: &[u8],
8392    sc_off: usize,
8393    rows: usize,
8394    r: usize,
8395    gpr: usize,
8396    xs: [&[i8]; 4],
8397    gs: [&[i32]; 4],
8398) -> [f32; 4] {
8399    // SAFETY: same contracts as `dot_q1g_row_sdot`, ×4.
8400    unsafe {
8401        use core::arch::aarch64::*;
8402        use core::arch::asm;
8403        const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
8404        const IW: [[u8; 16]; 8] = [
8405            [0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1],
8406            [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3],
8407            [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5],
8408            [6, 6, 6, 6, 6, 6, 6, 6, 7, 7, 7, 7, 7, 7, 7, 7],
8409            [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9],
8410            [
8411                10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
8412            ],
8413            [
8414                12, 12, 12, 12, 12, 12, 12, 12, 13, 13, 13, 13, 13, 13, 13, 13,
8415            ],
8416            [
8417                14, 14, 14, 14, 14, 14, 14, 14, 15, 15, 15, 15, 15, 15, 15, 15,
8418            ],
8419        ];
8420        let m = vld1q_u8(MASKS.as_ptr());
8421        let iw: [uint8x16_t; 8] = core::array::from_fn(|i| vld1q_u8(IW[i].as_ptr()));
8422        macro_rules! sdot2 {
8423            ($w0:expr, $w1:expr, $x:expr) => {{
8424                let x0 = vld1q_s8($x);
8425                let x1 = vld1q_s8($x.add(16));
8426                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
8427                asm!(
8428                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
8429                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
8430                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
8431                    w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
8432                    options(pure, nomem, nostack),
8433                );
8434                vaddq_s32(a0, a1)
8435            }};
8436        }
8437        let bp = bytes.as_ptr().add(r * gpr * 4);
8438        let vec_end = if r + 1 == rows { gpr - 4 } else { gpr };
8439        let mut accv = [vdupq_n_f32(0.0); 4];
8440        let mut gi = 0;
8441        while gi + 4 <= vec_end {
8442            let ld = vld1q_u8(bp.add(gi * 4));
8443            // Unpack ONCE — eight ±mask vectors serve all four streams.
8444            let w: [int8x16_t; 8] =
8445                core::array::from_fn(|i| vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld, iw[i]), m)));
8446            let h = vdup_n_u16(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
8447            let scf: float32x4_t;
8448            asm!(
8449                "fcvtl {o:v}.4s, {i:v}.4h",
8450                o = out(vreg) scf, i = in(vreg) h,
8451                options(pure, nomem, nostack),
8452            );
8453            for k in 0..4 {
8454                let xp = xs[k].as_ptr();
8455                let d0 = sdot2!(w[0], w[1], xp.add(gi * GROUP_SIZE));
8456                let d1 = sdot2!(w[2], w[3], xp.add((gi + 1) * GROUP_SIZE));
8457                let d2 = sdot2!(w[4], w[5], xp.add((gi + 2) * GROUP_SIZE));
8458                let d3 = sdot2!(w[6], w[7], xp.add((gi + 3) * GROUP_SIZE));
8459                let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
8460                let g = vld1q_s32(gs[k].as_ptr().add(gi));
8461                let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
8462                accv[k] = vfmaq_f32(accv[k], vcvtq_f32_s32(dots), scf);
8463            }
8464            gi += 4;
8465        }
8466        let mut acc = [
8467            vaddvq_f32(accv[0]),
8468            vaddvq_f32(accv[1]),
8469            vaddvq_f32(accv[2]),
8470            vaddvq_f32(accv[3]),
8471        ];
8472        while gi < gpr {
8473            let t = bp.add(gi * 4);
8474            let sc = f16_to_f32(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
8475            let v0 = vcombine_u8(vdup_n_u8(*t), vdup_n_u8(*t.add(1)));
8476            let v1 = vcombine_u8(vdup_n_u8(*t.add(2)), vdup_n_u8(*t.add(3)));
8477            let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
8478            let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
8479            for k in 0..4 {
8480                let d = vaddvq_s32(sdot2!(w0, w1, xs[k].as_ptr().add(gi * GROUP_SIZE)));
8481                acc[k] += (-(2 * d + *gs[k].as_ptr().add(gi))) as f32 * sc;
8482            }
8483            gi += 1;
8484        }
8485        acc
8486    }
8487}
8488
8489/// (weight ±1, scale) of one q1g element — the exact outlier term.
8490#[inline]
8491fn q1g_outlier(bytes: &[u8], sc_off: usize, r: usize, gpr: usize, j: usize) -> (f32, f32) {
8492    let t = r * gpr + j / GROUP_SIZE;
8493    let k = j % GROUP_SIZE;
8494    let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, t));
8495    let bit = (bytes[t * 4 + k / 8] >> (k % 8)) & 1;
8496    ((bit as i32 * 2 - 1) as f32, s)
8497}
8498
8499/// Exact scalar q1g row (CMF_SDOT=0 contract) — `q1_row_exact`'s order.
8500#[inline]
8501fn q1g_row_exact(bytes: &[u8], sc_off: usize, r: usize, gpr: usize, x: &[f32]) -> f32 {
8502    let mut acc = 0f32;
8503    for gi in 0..gpr {
8504        let t = r * gpr + gi;
8505        let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, t));
8506        let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
8507        let mut ga = 0f32;
8508        for (j, &b) in bytes[t * 4..t * 4 + 4].iter().enumerate() {
8509            for k in 0..8 {
8510                ga += (((b >> k) & 1) as f32 * 2.0 - 1.0) * xg[j * 8 + k];
8511            }
8512        }
8513        acc += ga * s;
8514    }
8515    acc
8516}
8517
8518/// q1t per-row overlay locator. After the base (`base_len`) come
8519/// `[u32 row_ptr[rows+1]]` then `[(u16 col, f16 val)]` grouped by row (row
8520/// `r`'s entries are `[row_ptr[r], row_ptr[r+1])`). Returns
8521/// `(row_ptr offset, entries offset, present)`.
8522fn q1t_overlay(bytes: &[u8], base_len: usize, rows: usize) -> (usize, usize, bool) {
8523    let entries = base_len + (rows + 1) * 4;
8524    (base_len, entries, entries <= bytes.len())
8525}
8526
8527/// Read `row_ptr[r]` from the overlay's prefix-sum table.
8528#[inline]
8529fn q1t_rowptr(bytes: &[u8], rp_off: usize, r: usize) -> usize {
8530    let o = rp_off + r * 4;
8531    u32::from_le_bytes([bytes[o], bytes[o + 1], bytes[o + 2], bytes[o + 3]]) as usize
8532}
8533
8534/// Byte → the 5 ternary signs it packs `{−1,0,+1}` as f32, precomputed so
8535/// decoding a q1t code is a table load, not the base-3 divide/modulo per
8536/// weight (division is ~20–40× the cost of a load). Built at compile time.
8537const SIGN5: [[f32; 5]; 256] = {
8538    let mut lut = [[0.0f32; 5]; 256];
8539    let pow3 = [1u16, 3, 9, 27, 81];
8540    let mut byte = 0usize;
8541    while byte < 256 {
8542        let mut i = 0usize;
8543        while i < 5 {
8544            let code = (byte as u16 / pow3[i]) % 3;
8545            lut[byte][i] = if code == 1 {
8546                1.0
8547            } else if code == 2 {
8548                -1.0
8549            } else {
8550                0.0
8551            };
8552            i += 1;
8553        }
8554        byte += 1;
8555    }
8556    lut
8557};
8558
8559/// Same table, as i8 signs — the operand for the int8 SDOT base kernel.
8560const SIGN5_I8: [[i8; 5]; 256] = {
8561    let mut lut = [[0i8; 5]; 256];
8562    let pow3 = [1u16, 3, 9, 27, 81];
8563    let mut byte = 0usize;
8564    while byte < 256 {
8565        let mut i = 0usize;
8566        while i < 5 {
8567            let code = (byte as u16 / pow3[i]) % 3;
8568            lut[byte][i] = if code == 1 {
8569                1
8570            } else if code == 2 {
8571                -1
8572            } else {
8573                0
8574            };
8575            i += 1;
8576        }
8577        byte += 1;
8578    }
8579    lut
8580};
8581
8582/// The same 5 i8 signs packed into a u64 (`[s0 s1 s2 s3 s4 0 0 0]`, LE) so the
8583/// group unpack is 7 unaligned u64 stores at offsets 0,5,10,…,30 instead of
8584/// six 5-byte copies + LUT indexing — each store's trailing zeros are fixed by
8585/// the next store, and the last one runs 6 B past the 32nd weight (the unpack
8586/// buffer is padded to 40). This is the decode/prefill hot inner op.
8587const SIGN5_U64: [u64; 256] = {
8588    let mut lut = [0u64; 256];
8589    let pow3 = [1u16, 3, 9, 27, 81];
8590    let mut byte = 0usize;
8591    while byte < 256 {
8592        let mut v = 0u64;
8593        let mut i = 0usize;
8594        while i < 5 {
8595            let code = (byte as u16 / pow3[i]) % 3;
8596            let s: u8 = if code == 1 {
8597                1
8598            } else if code == 2 {
8599                0xFF
8600            } else {
8601                0
8602            };
8603            v |= (s as u64) << (i * 8);
8604            i += 1;
8605        }
8606        lut[byte] = v;
8607        byte += 1;
8608    }
8609    lut
8610};
8611
8612/// Ternary base weight at `(row r, col j)` = `sign(code)·s_group`. Used to add
8613/// back activation-outlier columns, whose `x` was zeroed for the int8 bulk dot
8614/// (`split_act`). At a weight-outlier position the code is 0, so this is 0 and
8615/// the overlay correction owns that column — no double counting.
8616#[inline]
8617fn q1t_base_weight(bytes: &[u8], r: usize, gpr: usize, j: usize) -> f32 {
8618    const TILE: usize = cortiq_core::quant::Q1T_TILE;
8619    let off = (r * gpr + j / GROUP_SIZE) * TILE;
8620    let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
8621    let within = j % GROUP_SIZE;
8622    SIGN5[bytes[off + 2 + within / 5] as usize][within % 5] * s
8623}
8624
8625/// One 32-group int8 dot via two SDOTs. Bit-exact vs the scalar i8 sum
8626/// (integer accumulation is order-independent).
8627#[cfg(target_arch = "aarch64")]
8628#[target_feature(enable = "neon,dotprod")]
8629#[inline]
8630unsafe fn sdot32_i8(w: *const i8, x: *const i8) -> i32 {
8631    // SAFETY: caller guarantees 32 readable i8 at each pointer.
8632    unsafe {
8633        use core::arch::aarch64::*;
8634        use core::arch::asm;
8635        let w0 = vld1q_s8(w);
8636        let w1 = vld1q_s8(w.add(16));
8637        let x0 = vld1q_s8(x);
8638        let x1 = vld1q_s8(x.add(16));
8639        let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
8640        asm!(
8641            "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
8642            "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
8643            a0 = inout(vreg) a0, a1 = inout(vreg) a1,
8644            w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
8645            options(pure, nomem, nostack),
8646        );
8647        vaddvq_s32(vaddq_s32(a0, a1))
8648    }
8649}
8650
8651/// One 32-group int8 dot via AVX2: signed·signed as `maddubs(|w|, sign(x,w))`
8652/// then `madd` and a horizontal reduce (the same idiom as `dot_q4t_row_avx2`).
8653#[cfg(target_arch = "x86_64")]
8654#[target_feature(enable = "avx2")]
8655#[inline]
8656unsafe fn i8dot32_avx2(w: *const i8, x: *const i8) -> i32 {
8657    // SAFETY: caller guarantees 32 readable i8 at each pointer.
8658    unsafe {
8659        use core::arch::x86_64::*;
8660        let wv = _mm256_loadu_si256(w as *const __m256i);
8661        let xv = _mm256_loadu_si256(x as *const __m256i);
8662        let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
8663        let d = _mm256_madd_epi16(p16, _mm256_set1_epi16(1));
8664        let hi128 = _mm256_extracti128_si256::<1>(d);
8665        let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
8666        let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8667        let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8668        _mm_cvtsi128_si32(s32)
8669    }
8670}
8671
8672/// Unpack one q1t group's base-3 codes into 32 i8 signs via 7 unaligned u64
8673/// stores (see `SIGN5_U64`). `dst` MUST have ≥ 40 bytes: the 7th store writes
8674/// `dst[30..38]`. Stores go in order so each one's trailing zeros are
8675/// overwritten by the next; the final 6 padding bytes are unused by the dot.
8676#[inline]
8677fn q1t_unpack_group_i8(codes: *const u8, dst: &mut [i8]) {
8678    debug_assert!(dst.len() >= 40);
8679    // SAFETY: codes points at 7 readable bytes; dst has ≥ 40 bytes so every
8680    // 8-byte store at offset bi*5 (bi ≤ 6 → ≤ 30) stays in bounds.
8681    unsafe {
8682        let p = dst.as_mut_ptr();
8683        for bi in 0..7 {
8684            core::ptr::write_unaligned(
8685                p.add(bi * 5) as *mut u64,
8686                SIGN5_U64[*codes.add(bi) as usize],
8687            );
8688        }
8689    }
8690}
8691
8692/// One 32-group int8 dot, arch-dispatched (the matmat inner loop, where the
8693/// row's signs are unpacked once and dotted against every batch input).
8694/// Callers are gated by `a8w8_enabled()`, so the target-feature arms are
8695/// reachable; the scalar arm is a non-SIMD-arch fallback.
8696#[inline]
8697fn q1t_i8dot32(w: *const i8, x: *const i8) -> i32 {
8698    #[cfg(target_arch = "aarch64")]
8699    unsafe {
8700        return sdot32_i8(w, x);
8701    }
8702    #[cfg(target_arch = "x86_64")]
8703    unsafe {
8704        return i8dot32_avx2(w, x);
8705    }
8706    #[allow(unreachable_code)]
8707    unsafe {
8708        let mut s = 0i32;
8709        for k in 0..GROUP_SIZE {
8710            s += *w.add(k) as i32 * *x.add(k) as i32;
8711        }
8712        s
8713    }
8714}
8715
8716#[inline]
8717unsafe fn q1t_unpack_reg_u64s(codes: *const u8) -> (u64, u64, u64, u64) {
8718    let (s0, s1, s2, s3, s4, s5, s6) = unsafe {
8719        (
8720            SIGN5_U64[*codes as usize],
8721            SIGN5_U64[*codes.add(1) as usize],
8722            SIGN5_U64[*codes.add(2) as usize],
8723            SIGN5_U64[*codes.add(3) as usize],
8724            SIGN5_U64[*codes.add(4) as usize],
8725            SIGN5_U64[*codes.add(5) as usize],
8726            SIGN5_U64[*codes.add(6) as usize],
8727        )
8728    };
8729
8730    let u0 = s0 | (s1 << 40);
8731    let u1 = (s1 >> 24) | (s2 << 16) | (s3 << 56);
8732    let u2 = (s3 >> 8) | (s4 << 32);
8733    let u3 = (s4 >> 32) | (s5 << 8) | (s6 << 48);
8734
8735    (u0, u1, u2, u3)
8736}
8737
8738/// One q1t row's int8 base dot: `Σ_group s·dot(signs, xq)` (before the shared
8739/// `sx`). Direct register unpacking (zero stack stores/loads, no STLF stalls).
8740/// ARM SDOT.
8741#[cfg(target_arch = "aarch64")]
8742#[target_feature(enable = "neon,dotprod")]
8743unsafe fn q1t_dot_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
8744    use core::arch::aarch64::*;
8745    use core::arch::asm;
8746    unsafe {
8747        const TILE: usize = cortiq_core::quant::Q1T_TILE;
8748        let mut acc = 0f32;
8749        let bytes_ptr = bytes.as_ptr();
8750        let xq_ptr = xq.as_ptr();
8751        let row_off = r * gpr * TILE;
8752
8753        let gpr2 = gpr & !1;
8754        let mut gi = 0;
8755        while gi < gpr2 {
8756            let off0 = row_off + gi * TILE;
8757            let off1 = off0 + TILE;
8758            let s0 = f16_to_f32(u16::from_le_bytes([
8759                *bytes_ptr.add(off0),
8760                *bytes_ptr.add(off0 + 1),
8761            ]));
8762            let s1 = f16_to_f32(u16::from_le_bytes([
8763                *bytes_ptr.add(off1),
8764                *bytes_ptr.add(off1 + 1),
8765            ]));
8766
8767            let (u0_0, u1_0, u2_0, u3_0) = q1t_unpack_reg_u64s(bytes_ptr.add(off0 + 2));
8768            let (u0_1, u1_1, u2_1, u3_1) = q1t_unpack_reg_u64s(bytes_ptr.add(off1 + 2));
8769
8770            let w0_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_0), vcreate_u64(u1_0)));
8771            let w1_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_0), vcreate_u64(u3_0)));
8772            let w0_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_1), vcreate_u64(u1_1)));
8773            let w1_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_1), vcreate_u64(u3_1)));
8774
8775            let x0_0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE));
8776            let x1_0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE + 16));
8777            let x0_1 = vld1q_s8(xq_ptr.add((gi + 1) * GROUP_SIZE));
8778            let x1_1 = vld1q_s8(xq_ptr.add((gi + 1) * GROUP_SIZE + 16));
8779
8780            let (mut a0_0, mut a1_0) = (vdupq_n_s32(0), vdupq_n_s32(0));
8781            let (mut a0_1, mut a1_1) = (vdupq_n_s32(0), vdupq_n_s32(0));
8782            asm!(
8783                "sdot {a0_0:v}.4s, {w0_0:v}.16b, {x0_0:v}.16b",
8784                "sdot {a1_0:v}.4s, {w1_0:v}.16b, {x1_0:v}.16b",
8785                "sdot {a0_1:v}.4s, {w0_1:v}.16b, {x0_1:v}.16b",
8786                "sdot {a1_1:v}.4s, {w1_1:v}.16b, {x1_1:v}.16b",
8787                a0_0 = inout(vreg) a0_0, a1_0 = inout(vreg) a1_0,
8788                a0_1 = inout(vreg) a0_1, a1_1 = inout(vreg) a1_1,
8789                w0_0 = in(vreg) w0_0, x0_0 = in(vreg) x0_0, w1_0 = in(vreg) w1_0, x1_0 = in(vreg) x1_0,
8790                w0_1 = in(vreg) w0_1, x0_1 = in(vreg) x0_1, w1_1 = in(vreg) w1_1, x1_1 = in(vreg) x1_1,
8791                options(pure, nomem, nostack),
8792            );
8793            let d0 = vaddvq_s32(vaddq_s32(a0_0, a1_0));
8794            let d1 = vaddvq_s32(vaddq_s32(a0_1, a1_1));
8795            acc += d0 as f32 * s0 + d1 as f32 * s1;
8796            gi += 2;
8797        }
8798
8799        if gi < gpr {
8800            let off = row_off + gi * TILE;
8801            let s = f16_to_f32(u16::from_le_bytes([
8802                *bytes_ptr.add(off),
8803                *bytes_ptr.add(off + 1),
8804            ]));
8805            let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
8806            let w0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0), vcreate_u64(u1)));
8807            let w1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2), vcreate_u64(u3)));
8808            let x0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE));
8809            let x1 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE + 16));
8810            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
8811            asm!(
8812                "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
8813                "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
8814                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
8815                w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
8816                options(pure, nomem, nostack),
8817            );
8818            let d = vaddvq_s32(vaddq_s32(a0, a1));
8819            acc += d as f32 * s;
8820        }
8821        acc
8822    }
8823}
8824
8825/// x86 AVX2 mirror of `q1t_dot_row_sdot` (maddubs int8 dot per group).
8826#[cfg(target_arch = "x86_64")]
8827#[target_feature(enable = "avx2")]
8828unsafe fn q1t_dot_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
8829    use core::arch::x86_64::*;
8830    unsafe {
8831        const TILE: usize = cortiq_core::quant::Q1T_TILE;
8832        let mut acc = 0f32;
8833        let bytes_ptr = bytes.as_ptr();
8834        let xq_ptr = xq.as_ptr();
8835        let row_off = r * gpr * TILE;
8836
8837        let ones = _mm256_set1_epi16(1);
8838        for gi in 0..gpr {
8839            let off = row_off + gi * TILE;
8840            let s = f16_to_f32(u16::from_le_bytes([
8841                *bytes_ptr.add(off),
8842                *bytes_ptr.add(off + 1),
8843            ]));
8844            let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
8845            let wv = _mm256_set_epi64x(u3 as i64, u2 as i64, u1 as i64, u0 as i64);
8846            let xv = _mm256_loadu_si256(xq_ptr.add(gi * GROUP_SIZE) as *const __m256i);
8847            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
8848            let d256 = _mm256_madd_epi16(p16, ones);
8849            let d128 = _mm_add_epi32(
8850                _mm256_castsi256_si128(d256),
8851                _mm256_extracti128_si256(d256, 1),
8852            );
8853            let d64 = _mm_add_epi32(d128, _mm_shuffle_epi32(d128, 0xee));
8854            let d32 = _mm_cvtsi128_si32(_mm_add_epi32(d64, _mm_shuffle_epi32(d64, 0x55)));
8855            acc += d32 as f32 * s;
8856        }
8857        acc
8858    }
8859}
8860
8861/// VNNI twin of `q1t_dot_row_avx2` (see `dpbusd_hsum`).
8862#[cfg(target_arch = "x86_64")]
8863#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
8864unsafe fn q1t_dot_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
8865    use core::arch::x86_64::*;
8866    // SAFETY: same tile/xq contracts as `q1t_dot_row_avx2`.
8867    unsafe {
8868        const TILE: usize = cortiq_core::quant::Q1T_TILE;
8869        let mut acc = 0f32;
8870        let bytes_ptr = bytes.as_ptr();
8871        let xq_ptr = xq.as_ptr();
8872        let row_off = r * gpr * TILE;
8873        for gi in 0..gpr {
8874            let off = row_off + gi * TILE;
8875            let s = f16_to_f32(u16::from_le_bytes([
8876                *bytes_ptr.add(off),
8877                *bytes_ptr.add(off + 1),
8878            ]));
8879            let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
8880            let wv = _mm256_set_epi64x(u3 as i64, u2 as i64, u1 as i64, u0 as i64);
8881            let xv = _mm256_loadu_si256(xq_ptr.add(gi * GROUP_SIZE) as *const __m256i);
8882            let d = dpbusd_hsum(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
8883            acc += d as f32 * s;
8884        }
8885        acc
8886    }
8887}
8888
8889/// Per-row int8 base dot, dispatched once per row (matvec decode hot path).
8890/// Callers are gated by `a8w8_enabled()`, so the target-feature kernels are
8891/// reachable.
8892#[inline]
8893fn q1t_dot_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
8894    #[cfg(target_arch = "aarch64")]
8895    unsafe {
8896        return q1t_dot_row_sdot(bytes, r, gpr, xq);
8897    }
8898    #[cfg(target_arch = "x86_64")]
8899    unsafe {
8900        if vnni_tiles_enabled() {
8901            return q1t_dot_row_vnni(bytes, r, gpr, xq);
8902        }
8903        return q1t_dot_row_avx2(bytes, r, gpr, xq);
8904    }
8905    #[allow(unreachable_code)]
8906    {
8907        const TILE: usize = cortiq_core::quant::Q1T_TILE;
8908        let mut acc = 0f32;
8909        let mut sg = [0i8; GROUP_SIZE + 8]; // +8 slack for the u64-store unpack
8910        for gi in 0..gpr {
8911            let off = (r * gpr + gi) * TILE;
8912            let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
8913            q1t_unpack_group_i8(bytes.as_ptr().wrapping_add(off + 2), &mut sg);
8914            let mut d = 0i32;
8915            for k in 0..GROUP_SIZE {
8916                d += sg[k] as i32 * xq[gi * GROUP_SIZE + k] as i32;
8917            }
8918            acc += d as f32 * s;
8919        }
8920        acc
8921    }
8922}
8923
8924/// Σ over a row's outliers of `value·x[col]` — the correction that adds the
8925/// overlay's exact weights on top of the base dot. INVARIANT: the encoder
8926/// writes ternary code 0 at every outlier position (`quantize_q1t`), so the
8927/// base contributes nothing there and this is a plain `value·x`, not
8928/// `(value − base)·x` — no scattered per-outlier scale read. Row `r`'s entries
8929/// are the contiguous slice `[row_ptr[r], row_ptr[r+1])`, so no binary search.
8930fn q1t_row_outlier_correction(
8931    bytes: &[u8],
8932    r: usize,
8933    rp_off: usize,
8934    entries_off: usize,
8935    has_ov: bool,
8936    x: &[f32],
8937) -> f32 {
8938    if !has_ov {
8939        return 0.0;
8940    }
8941    let (c0, c1) = (
8942        q1t_rowptr(bytes, rp_off, r),
8943        q1t_rowptr(bytes, rp_off, r + 1),
8944    );
8945    let mut corr = 0f32;
8946    for p in c0..c1 {
8947        let e = entries_off + p * 4;
8948        let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
8949        let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
8950        corr += val * x[col];
8951    }
8952    corr
8953}
8954
8955/// Dequantize one q1t row into `buf[..cols]` via the sign LUT (no division),
8956/// then apply the row's outliers (its `[row_ptr[r], row_ptr[r+1])` slice).
8957/// Used by the batched (prefill) path where the decode amortizes over the batch.
8958fn q1t_dequant_row(
8959    bytes: &[u8],
8960    r: usize,
8961    gpr: usize,
8962    rp_off: usize,
8963    entries_off: usize,
8964    has_ov: bool,
8965    buf: &mut [f32],
8966) {
8967    const TILE: usize = cortiq_core::quant::Q1T_TILE;
8968    for g in 0..gpr {
8969        let off = (r * gpr + g) * TILE;
8970        let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
8971        let codes = &bytes[off + 2..off + TILE];
8972        let bc = g * GROUP_SIZE;
8973        // 6 full bytes (30 codes) + a 7th byte holding the last 2.
8974        for bi in 0..6 {
8975            let lut = &SIGN5[codes[bi] as usize];
8976            let d = &mut buf[bc + bi * 5..bc + bi * 5 + 5];
8977            for i in 0..5 {
8978                d[i] = lut[i] * s;
8979            }
8980        }
8981        let lut = &SIGN5[codes[6] as usize];
8982        buf[bc + 30] = lut[0] * s;
8983        buf[bc + 31] = lut[1] * s;
8984    }
8985    if !has_ov {
8986        return;
8987    }
8988    let (c0, c1) = (
8989        q1t_rowptr(bytes, rp_off, r),
8990        q1t_rowptr(bytes, rp_off, r + 1),
8991    );
8992    for p in c0..c1 {
8993        let e = entries_off + p * 4;
8994        let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
8995        buf[col] = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
8996    }
8997}
8998
8999/// Add the sparse outlier overlay onto a base dot already in `out` (the GPU
9000/// computes the ternary base; the overlay stays on the CPU — its entries are
9001/// few and its per-row gather doesn't vectorize on the GPU). Row-parallel.
9002fn q1t_add_overlay(
9003    bytes: &[u8],
9004    x: &[f32],
9005    rows: usize,
9006    cols: usize,
9007    out: &mut [f32],
9008    pool: Option<&Pool>,
9009) {
9010    const TILE: usize = cortiq_core::quant::Q1T_TILE;
9011    let gpr = cols / GROUP_SIZE;
9012    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
9013    if !has_ov {
9014        return;
9015    }
9016    let out_addr = SendMut(out.as_mut_ptr());
9017    let run = move |start: usize, end: usize| {
9018        for r in start..end {
9019            let corr = q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
9020            // SAFETY: disjoint rows; add onto the base the GPU already wrote.
9021            unsafe { *out_addr.at(r) += corr };
9022        }
9023    };
9024    dispatch_rows(pool, rows, &run);
9025}
9026
9027/// Q1T row range via the A8W8 int8 path — shared activation split,
9028/// per-row: base SDOT dot + outlier correction + overlay.
9029#[allow(clippy::too_many_arguments)]
9030fn q1t_range_a8w8(
9031    bytes: &[u8],
9032    gpr: usize,
9033    rp_off: usize,
9034    ent_off: usize,
9035    has_ov: bool,
9036    act: &SplitAct,
9037    x: &[f32],
9038    out: SendMut,
9039    start: usize,
9040    end: usize,
9041) {
9042    for r in start..end {
9043        let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
9044        for &(j, xv) in &act.outliers {
9045            acc += q1t_base_weight(bytes, r, gpr, j) * xv;
9046        }
9047        acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
9048        // SAFETY: disjoint row ranges per worker.
9049        unsafe { *out.at(r) = acc };
9050    }
9051}
9052
9053/// Q1T row range via the f32 path (no SDOT) — for matvec_many batched
9054/// dispatch when a8w8 is unavailable.
9055#[allow(clippy::too_many_arguments)]
9056fn q1t_range_f32_batch(
9057    bytes: &[u8],
9058    gpr: usize,
9059    rp_off: usize,
9060    ent_off: usize,
9061    has_ov: bool,
9062    x: &[f32],
9063    out: SendMut,
9064    start: usize,
9065    end: usize,
9066) {
9067    const TILE: usize = cortiq_core::quant::Q1T_TILE;
9068    let mut sg = [0f32; GROUP_SIZE];
9069    for r in start..end {
9070        let mut acc = 0f32;
9071        for g in 0..gpr {
9072            let off = (r * gpr + g) * TILE;
9073            let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
9074            let codes = &bytes[off + 2..off + TILE];
9075            let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
9076            for bi in 0..6 {
9077                sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
9078            }
9079            let lut = &SIGN5[codes[6] as usize];
9080            sg[30] = lut[0];
9081            sg[31] = lut[1];
9082            let mut gsum = 0f32;
9083            for k in 0..GROUP_SIZE {
9084                gsum += sg[k] * xg[k];
9085            }
9086            acc += s * gsum;
9087        }
9088        acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
9089        // SAFETY: disjoint row ranges per worker.
9090        unsafe { *out.at(r) = acc };
9091    }
9092}
9093
9094/// Ternary (q1t) matvec — decode+dot straight from mmap, one group at a time:
9095/// no per-ROW buffer, no division (the sign LUT), and a tiny per-group sign
9096/// buffer so the 32-wide dot vectorizes. This is the decode hot path.
9097fn q1t_matvec(
9098    bytes: &[u8],
9099    x: &[f32],
9100    rows: usize,
9101    cols: usize,
9102    out: &mut [f32],
9103    pool: Option<&Pool>,
9104) {
9105    debug_assert_eq!(out.len(), rows);
9106    const TILE: usize = cortiq_core::quant::Q1T_TILE;
9107    let gpr = cols / GROUP_SIZE;
9108    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
9109    let out_addr = SendMut(out.as_mut_ptr());
9110    // int8 SDOT base dot (ARM dotprod): ~4× the f32 arithmetic. x → i8 once
9111    // (`split_act`), activation outliers added back exactly in f32, weight
9112    // overlay on top. ARM SDOT / x86 AVX2; CMF_SDOT=0 keeps the exact f32 path.
9113    if a8w8_enabled() {
9114        let act = split_act(x);
9115        let act = &act;
9116        let run = move |start: usize, end: usize| {
9117            for r in start..end {
9118                let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
9119                for &(j, xv) in &act.outliers {
9120                    acc += q1t_base_weight(bytes, r, gpr, j) * xv;
9121                }
9122                acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
9123                // SAFETY: disjoint row ranges per worker.
9124                unsafe { *out_addr.at(r) = acc };
9125            }
9126        };
9127        dispatch_rows(pool, rows, &run);
9128        return;
9129    }
9130    let run = move |start: usize, end: usize| {
9131        // Per-group signs, unpacked contiguously so the dot below is a clean
9132        // 32-wide reduction the autovectorizer turns into f32x4 FMAs — the
9133        // 5-values-per-byte base-3 layout won't SIMD in place.
9134        let mut sg = [0f32; GROUP_SIZE];
9135        for r in start..end {
9136            let mut acc = 0f32;
9137            for g in 0..gpr {
9138                let off = (r * gpr + g) * TILE;
9139                let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
9140                let codes = &bytes[off + 2..off + TILE];
9141                let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
9142                for bi in 0..6 {
9143                    sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
9144                }
9145                let lut = &SIGN5[codes[6] as usize];
9146                sg[30] = lut[0];
9147                sg[31] = lut[1];
9148                let mut gsum = 0f32;
9149                for k in 0..GROUP_SIZE {
9150                    gsum += sg[k] * xg[k];
9151                }
9152                acc += s * gsum;
9153            }
9154            acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
9155            unsafe { *out_addr.at(r) = acc };
9156        }
9157    };
9158    dispatch_rows(pool, rows, &run);
9159}
9160
9161/// Fused-pair twin of `q1t_dot_row_sdot`: ONE register unpack of the
9162/// ternary codes serves BOTH activation streams (the unpack chain is
9163/// the dominant per-row cost — MTP verify pairs paid it twice). Per
9164/// stream the group order and f32 accumulation match the single-row
9165/// kernel exactly, so pair == 2×matvec bit-for-bit.
9166#[cfg(target_arch = "aarch64")]
9167#[target_feature(enable = "neon,dotprod")]
9168unsafe fn q1t_dot_row_sdot2(bytes: &[u8], r: usize, gpr: usize, xa: &[i8], xb: &[i8]) -> [f32; 2] {
9169    use core::arch::aarch64::*;
9170    use core::arch::asm;
9171    // SAFETY: same slice-length contracts as `q1t_dot_row_sdot`, ×2.
9172    unsafe {
9173        const TILE: usize = cortiq_core::quant::Q1T_TILE;
9174        let bytes_ptr = bytes.as_ptr();
9175        let row_off = r * gpr * TILE;
9176        let xp = [xa.as_ptr(), xb.as_ptr()];
9177        let mut acc = [0f32; 2];
9178        macro_rules! sdot2 {
9179            ($w0:expr, $w1:expr, $x:expr) => {{
9180                let x0 = vld1q_s8($x);
9181                let x1 = vld1q_s8($x.add(16));
9182                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
9183                asm!(
9184                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
9185                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
9186                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
9187                    w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
9188                    options(pure, nomem, nostack),
9189                );
9190                vaddvq_s32(vaddq_s32(a0, a1))
9191            }};
9192        }
9193        let gpr2 = gpr & !1;
9194        let mut gi = 0;
9195        while gi < gpr2 {
9196            let off0 = row_off + gi * TILE;
9197            let off1 = off0 + TILE;
9198            let s0 = f16_to_f32(u16::from_le_bytes([
9199                *bytes_ptr.add(off0),
9200                *bytes_ptr.add(off0 + 1),
9201            ]));
9202            let s1 = f16_to_f32(u16::from_le_bytes([
9203                *bytes_ptr.add(off1),
9204                *bytes_ptr.add(off1 + 1),
9205            ]));
9206            let (u0_0, u1_0, u2_0, u3_0) = q1t_unpack_reg_u64s(bytes_ptr.add(off0 + 2));
9207            let (u0_1, u1_1, u2_1, u3_1) = q1t_unpack_reg_u64s(bytes_ptr.add(off1 + 2));
9208            let w0_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_0), vcreate_u64(u1_0)));
9209            let w1_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_0), vcreate_u64(u3_0)));
9210            let w0_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_1), vcreate_u64(u1_1)));
9211            let w1_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_1), vcreate_u64(u3_1)));
9212            for k in 0..2 {
9213                let d0 = sdot2!(w0_0, w1_0, xp[k].add(gi * GROUP_SIZE));
9214                let d1 = sdot2!(w0_1, w1_1, xp[k].add((gi + 1) * GROUP_SIZE));
9215                acc[k] += d0 as f32 * s0 + d1 as f32 * s1;
9216            }
9217            gi += 2;
9218        }
9219        if gi < gpr {
9220            let off = row_off + gi * TILE;
9221            let s = f16_to_f32(u16::from_le_bytes([
9222                *bytes_ptr.add(off),
9223                *bytes_ptr.add(off + 1),
9224            ]));
9225            let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
9226            let w0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0), vcreate_u64(u1)));
9227            let w1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2), vcreate_u64(u3)));
9228            for k in 0..2 {
9229                let d = sdot2!(w0, w1, xp[k].add(gi * GROUP_SIZE));
9230                acc[k] += d as f32 * s;
9231            }
9232        }
9233        acc
9234    }
9235}
9236
9237/// Fused Q1T pair matvec: ONE pass over the rows serves both
9238/// activation streams — on ARM the ternary register unpack happens
9239/// once per tile pair (`q1t_dot_row_sdot2`); elsewhere the second dot
9240/// rides the row's L1-warm tile bytes. Per stream the math matches
9241/// `q1t_matvec` exactly.
9242fn q1t_matvec2(
9243    bytes: &[u8],
9244    x1: &[f32],
9245    x2: &[f32],
9246    rows: usize,
9247    cols: usize,
9248    o1: &mut [f32],
9249    o2: &mut [f32],
9250    pool: Option<&Pool>,
9251) {
9252    debug_assert_eq!(o1.len(), rows);
9253    debug_assert_eq!(o2.len(), rows);
9254    const TILE: usize = cortiq_core::quant::Q1T_TILE;
9255    let gpr = cols / GROUP_SIZE;
9256    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
9257    let out1 = SendMut(o1.as_mut_ptr());
9258    let out2 = SendMut(o2.as_mut_ptr());
9259    if a8w8_enabled() {
9260        let a1 = split_act(x1);
9261        let a2 = split_act(x2);
9262        let (a1, a2) = (&a1, &a2);
9263        let run = move |start: usize, end: usize| {
9264            for r in start..end {
9265                #[cfg(target_arch = "aarch64")]
9266                // a8w8 on aarch64 ⇔ sdot_enabled(), so the kernel's
9267                // target features are present.
9268                let ds = unsafe { q1t_dot_row_sdot2(bytes, r, gpr, &a1.xq, &a2.xq) };
9269                #[cfg(not(target_arch = "aarch64"))]
9270                let ds = [
9271                    q1t_dot_row_i8(bytes, r, gpr, &a1.xq),
9272                    q1t_dot_row_i8(bytes, r, gpr, &a2.xq),
9273                ];
9274                let mut acc1 = ds[0] * a1.sx;
9275                for &(j, xv) in &a1.outliers {
9276                    acc1 += q1t_base_weight(bytes, r, gpr, j) * xv;
9277                }
9278                acc1 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x1);
9279                let mut acc2 = ds[1] * a2.sx;
9280                for &(j, xv) in &a2.outliers {
9281                    acc2 += q1t_base_weight(bytes, r, gpr, j) * xv;
9282                }
9283                acc2 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x2);
9284                // SAFETY: disjoint row ranges per worker.
9285                unsafe {
9286                    *out1.at(r) = acc1;
9287                    *out2.at(r) = acc2;
9288                }
9289            }
9290        };
9291        dispatch_rows(pool, rows, &run);
9292        return;
9293    }
9294    let run = move |start: usize, end: usize| {
9295        // Exact path (CMF_SDOT=0): unpack the sign LUT once per group,
9296        // dot both streams — same op order per stream as `q1t_matvec`.
9297        let mut sg = [0f32; GROUP_SIZE];
9298        for r in start..end {
9299            let mut acc1 = 0f32;
9300            let mut acc2 = 0f32;
9301            for g in 0..gpr {
9302                let off = (r * gpr + g) * TILE;
9303                let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
9304                let codes = &bytes[off + 2..off + TILE];
9305                for bi in 0..6 {
9306                    sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
9307                }
9308                let lut = &SIGN5[codes[6] as usize];
9309                sg[30] = lut[0];
9310                sg[31] = lut[1];
9311                let xg1 = &x1[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
9312                let xg2 = &x2[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
9313                let mut gsum1 = 0f32;
9314                for k in 0..GROUP_SIZE {
9315                    gsum1 += sg[k] * xg1[k];
9316                }
9317                acc1 += s * gsum1;
9318                let mut gsum2 = 0f32;
9319                for k in 0..GROUP_SIZE {
9320                    gsum2 += sg[k] * xg2[k];
9321                }
9322                acc2 += s * gsum2;
9323            }
9324            acc1 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x1);
9325            acc2 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x2);
9326            // SAFETY: disjoint row ranges per worker.
9327            unsafe {
9328                *out1.at(r) = acc1;
9329                *out2.at(r) = acc2;
9330            }
9331        }
9332    };
9333    dispatch_rows(pool, rows, &run);
9334}
9335
9336/// Ternary (q1t) matmat (prefill) — dequant each row once, dot the whole
9337/// batch against it (amortizes the per-row decode).
9338fn q1t_matmat(
9339    bytes: &[u8],
9340    xs: &[f32],
9341    b: usize,
9342    rows: usize,
9343    cols: usize,
9344    out: &mut [f32],
9345    pool: Option<&Pool>,
9346) {
9347    debug_assert_eq!(out.len(), b * rows);
9348    const TILE: usize = cortiq_core::quant::Q1T_TILE;
9349    let gpr = cols / GROUP_SIZE;
9350    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
9351    let out_addr = SendMut(out.as_mut_ptr());
9352    // int8 prefill (ARM SDOT / x86 AVX2): quantize the B inputs once, unpack
9353    // each weight row's signs to i8 ONCE, then int8-dot against every input —
9354    // the row sign-decode amortizes over the whole batch. CMF_SDOT=0 → f32.
9355    if a8w8_enabled() {
9356        let acts: Vec<SplitAct> = (0..b)
9357            .map(|bi| split_act(&xs[bi * cols..(bi + 1) * cols]))
9358            .collect();
9359        let acts = &acts;
9360        let run = move |start: usize, end: usize| {
9361            let mut sg = vec![0i8; cols + 8]; // row signs, i8 (+8 unpack slack)
9362            let mut sc = vec![0f32; gpr]; // per-group scales
9363            let mut accs = vec![0f32; b]; // per-batch accumulators, reused per row
9364            for r in start..end {
9365                for g in 0..gpr {
9366                    let off = (r * gpr + g) * TILE;
9367                    sc[g] = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
9368                    q1t_unpack_group_i8(
9369                        bytes.as_ptr().wrapping_add(off + 2),
9370                        &mut sg[g * GROUP_SIZE..],
9371                    );
9372                }
9373                for bi in 0..b {
9374                    let act = &acts[bi];
9375                    let mut isum = 0f32;
9376                    for g in 0..gpr {
9377                        let d = q1t_i8dot32(
9378                            sg.as_ptr().wrapping_add(g * GROUP_SIZE),
9379                            act.xq.as_ptr().wrapping_add(g * GROUP_SIZE),
9380                        );
9381                        isum += d as f32 * sc[g];
9382                    }
9383                    let mut acc = isum * act.sx;
9384                    for &(j, xv) in &act.outliers {
9385                        acc += q1t_base_weight(bytes, r, gpr, j) * xv;
9386                    }
9387                    accs[bi] = acc;
9388                }
9389                // Overlay ONCE per row for the whole batch: read each (col, val)
9390                // from mmap a single time (was b× — the re-read dominated prefill)
9391                // and fan it out over the batch via the cached inputs.
9392                if has_ov {
9393                    let (c0, c1) = (
9394                        q1t_rowptr(bytes, rp_off, r),
9395                        q1t_rowptr(bytes, rp_off, r + 1),
9396                    );
9397                    for p in c0..c1 {
9398                        let e = ent_off + p * 4;
9399                        let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
9400                        let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
9401                        for bi in 0..b {
9402                            accs[bi] += val * xs[bi * cols + col];
9403                        }
9404                    }
9405                }
9406                for bi in 0..b {
9407                    unsafe { *out_addr.at(bi * rows + r) = accs[bi] };
9408                }
9409            }
9410        };
9411        dispatch_rows(pool, rows, &run);
9412        return;
9413    }
9414    let run = move |start: usize, end: usize| {
9415        let mut buf = vec![0f32; cols];
9416        for r in start..end {
9417            q1t_dequant_row(bytes, r, gpr, rp_off, ent_off, has_ov, &mut buf);
9418            for bi in 0..b {
9419                let xr = &xs[bi * cols..(bi + 1) * cols];
9420                let mut acc = 0f32;
9421                for j in 0..cols {
9422                    acc += buf[j] * xr[j];
9423                }
9424                unsafe { *out_addr.at(bi * rows + r) = acc };
9425            }
9426        }
9427    };
9428    dispatch_rows(pool, rows, &run);
9429}
9430
9431/// q1-tile entry points of the binary drivers (the tests' handles; the
9432/// runtime goes through `QTensor::bin_w`).
9433#[cfg(test)]
9434fn q1_matvec(
9435    bytes: &[u8],
9436    x: &[f32],
9437    rows: usize,
9438    cols: usize,
9439    out: &mut [f32],
9440    pool: Option<&Pool>,
9441) {
9442    bin_matvec(BinW::Q1(bytes), x, rows, cols, out, pool)
9443}
9444
9445/// Binary-weight matvec (q1 or q1g): A8W8 rows with the shared group
9446/// sums, or the exact scalar walk under CMF_SDOT=0.
9447fn bin_matvec(w: BinW, x: &[f32], rows: usize, cols: usize, out: &mut [f32], pool: Option<&Pool>) {
9448    debug_assert_eq!(out.len(), rows);
9449    let gpr = cols / GROUP_SIZE;
9450    let out_addr = SendMut(out.as_mut_ptr());
9451    if a8w8_enabled() {
9452        let act = split_act(x);
9453        let gsum = q1_group_sums(&act.xq, gpr);
9454        let (act, gsum) = (&act, &gsum);
9455        let run =
9456            move |start: usize, end: usize| bin_range_a8w8(w, gpr, act, gsum, out_addr, start, end);
9457        dispatch_rows(pool, rows, &run);
9458        return;
9459    }
9460    let run = move |start: usize, end: usize| bin_range_f32(w, gpr, x, out_addr, start, end);
9461    dispatch_rows(pool, rows, &run);
9462}
9463
9464/// Fused two-input q1 matvec (weights read once per pair).
9465#[cfg(test)]
9466#[allow(clippy::too_many_arguments)]
9467fn q1_matvec2(
9468    bytes: &[u8],
9469    x1: &[f32],
9470    x2: &[f32],
9471    rows: usize,
9472    cols: usize,
9473    o1: &mut [f32],
9474    o2: &mut [f32],
9475    pool: Option<&Pool>,
9476) {
9477    bin_matvec2(BinW::Q1(bytes), x1, x2, rows, cols, o1, o2, pool)
9478}
9479
9480/// Fused two-input binary-weight matvec (q1 or q1g).
9481#[allow(clippy::too_many_arguments)]
9482fn bin_matvec2(
9483    w: BinW,
9484    x1: &[f32],
9485    x2: &[f32],
9486    rows: usize,
9487    cols: usize,
9488    o1: &mut [f32],
9489    o2: &mut [f32],
9490    pool: Option<&Pool>,
9491) {
9492    let gpr = cols / GROUP_SIZE;
9493    let p1 = SendMut(o1.as_mut_ptr());
9494    let p2 = SendMut(o2.as_mut_ptr());
9495    if a8w8_enabled() {
9496        let a1 = split_act(x1);
9497        let a2 = split_act(x2);
9498        let g1 = q1_group_sums(&a1.xq, gpr);
9499        let g2 = q1_group_sums(&a2.xq, gpr);
9500        let (a1, a2, g1, g2) = (&a1, &a2, &g1, &g2);
9501        let run = move |start: usize, end: usize| {
9502            for r in start..end {
9503                let mut v1 = w.dot_i8(r, gpr, &a1.xq, g1) * a1.sx;
9504                let mut v2 = w.dot_i8(r, gpr, &a2.xq, g2) * a2.sx;
9505                for &(j, xv) in &a1.outliers {
9506                    let (wv, s) = w.outlier(r, gpr, j);
9507                    v1 += wv * s * xv;
9508                }
9509                for &(j, xv) in &a2.outliers {
9510                    let (wv, s) = w.outlier(r, gpr, j);
9511                    v2 += wv * s * xv;
9512                }
9513                // SAFETY: disjoint row ranges per worker.
9514                unsafe {
9515                    *p1.at(r) = v1;
9516                    *p2.at(r) = v2;
9517                }
9518            }
9519        };
9520        dispatch_rows(pool, rows, &run);
9521        return;
9522    }
9523    let run = move |start: usize, end: usize| {
9524        for r in start..end {
9525            // SAFETY: disjoint row ranges per worker.
9526            unsafe {
9527                *p1.at(r) = w.row_exact(r, gpr, x1);
9528                *p2.at(r) = w.row_exact(r, gpr, x2);
9529            }
9530        }
9531    };
9532    dispatch_rows(pool, rows, &run);
9533}
9534
9535/// Batched q1 matmat: each row's tiles stream once per microbatch.
9536#[cfg(test)]
9537#[allow(clippy::too_many_arguments)]
9538fn q1_matmat(
9539    bytes: &[u8],
9540    xs_all: &[f32],
9541    b: usize,
9542    rows: usize,
9543    cols: usize,
9544    out: &mut [f32],
9545    pool: Option<&Pool>,
9546) {
9547    bin_matmat(BinW::Q1(bytes), xs_all, b, rows, cols, out, pool)
9548}
9549
9550/// Batched binary-weight matmat (q1 or q1g).
9551#[allow(clippy::too_many_arguments)]
9552fn bin_matmat(
9553    w: BinW,
9554    xs_all: &[f32],
9555    b: usize,
9556    rows: usize,
9557    cols: usize,
9558    out: &mut [f32],
9559    pool: Option<&Pool>,
9560) {
9561    debug_assert_eq!(out.len(), b * rows);
9562    let gpr = cols / GROUP_SIZE;
9563    let out_addr = SendMut(out.as_mut_ptr());
9564    if a8w8_enabled() {
9565        let acts: Vec<(SplitAct, Vec<i32>)> = (0..b)
9566            .map(|bi| {
9567                let act = split_act(&xs_all[bi * cols..(bi + 1) * cols]);
9568                let gsum = q1_group_sums(&act.xq, gpr);
9569                (act, gsum)
9570            })
9571            .collect();
9572        let acts = &acts;
9573        #[cfg(target_arch = "x86_64")]
9574        let blocked_ok = avx2_enabled() && blocked_enabled();
9575        #[cfg(target_arch = "aarch64")]
9576        let blocked_ok = sdot_enabled() && blocked_enabled();
9577        let run = move |start: usize, end: usize| {
9578            for r in start..end {
9579                let mut bi = 0usize;
9580                // Blocked 1×4: the unpacked bit mask serves four
9581                // activation streams per group.
9582                #[cfg(target_arch = "aarch64")]
9583                if blocked_ok {
9584                    while bi + 4 <= acts.len() {
9585                        let xs = [
9586                            acts[bi].0.xq.as_slice(),
9587                            acts[bi + 1].0.xq.as_slice(),
9588                            acts[bi + 2].0.xq.as_slice(),
9589                            acts[bi + 3].0.xq.as_slice(),
9590                        ];
9591                        let gs = [
9592                            acts[bi].1.as_slice(),
9593                            acts[bi + 1].1.as_slice(),
9594                            acts[bi + 2].1.as_slice(),
9595                            acts[bi + 3].1.as_slice(),
9596                        ];
9597                        // SAFETY: blocked_ok checked SDOT.
9598                        let d = unsafe { w.dot_1x4_sdot(r, gpr, xs, gs) };
9599                        for k in 0..4 {
9600                            let (act, _) = &acts[bi + k];
9601                            let mut acc = d[k] * act.sx;
9602                            for &(j, xv) in &act.outliers {
9603                                let (wv, sc) = w.outlier(r, gpr, j);
9604                                acc += wv * sc * xv;
9605                            }
9606                            // SAFETY: disjoint (bi, r) cells per worker.
9607                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
9608                        }
9609                        bi += 4;
9610                    }
9611                }
9612                #[cfg(target_arch = "x86_64")]
9613                if blocked_ok {
9614                    while bi + 4 <= acts.len() {
9615                        let xs = [
9616                            acts[bi].0.xq.as_slice(),
9617                            acts[bi + 1].0.xq.as_slice(),
9618                            acts[bi + 2].0.xq.as_slice(),
9619                            acts[bi + 3].0.xq.as_slice(),
9620                        ];
9621                        let gs = [
9622                            acts[bi].1.as_slice(),
9623                            acts[bi + 1].1.as_slice(),
9624                            acts[bi + 2].1.as_slice(),
9625                            acts[bi + 3].1.as_slice(),
9626                        ];
9627                        // SAFETY: blocked_ok checked AVX2.
9628                        let d = unsafe { w.dot_1x4_x86(r, gpr, xs, gs) };
9629                        for k in 0..4 {
9630                            let (act, _) = &acts[bi + k];
9631                            let mut acc = d[k] * act.sx;
9632                            for &(j, xv) in &act.outliers {
9633                                let (wv, sc) = w.outlier(r, gpr, j);
9634                                acc += wv * sc * xv;
9635                            }
9636                            // SAFETY: disjoint (bi, r) cells per worker.
9637                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
9638                        }
9639                        bi += 4;
9640                    }
9641                }
9642                while bi < acts.len() {
9643                    let (act, gsum) = &acts[bi];
9644                    let mut acc = w.dot_i8(r, gpr, &act.xq, gsum) * act.sx;
9645                    for &(j, xv) in &act.outliers {
9646                        let (wv, s) = w.outlier(r, gpr, j);
9647                        acc += wv * s * xv;
9648                    }
9649                    // SAFETY: disjoint (bi, r) cells per worker range.
9650                    unsafe { *out_addr.at(bi * rows + r) = acc };
9651                    bi += 1;
9652                }
9653            }
9654        };
9655        dispatch_rows(pool, rows, &run);
9656        return;
9657    }
9658    let run = move |start: usize, end: usize| {
9659        for r in start..end {
9660            for bi in 0..b {
9661                let x = &xs_all[bi * cols..(bi + 1) * cols];
9662                // SAFETY: disjoint (bi, r) cells per worker range.
9663                unsafe { *out_addr.at(bi * rows + r) = w.row_exact(r, gpr, x) };
9664            }
9665        }
9666    };
9667    dispatch_rows(pool, rows, &run);
9668}
9669
9670/// Fused q4_block matvec straight from the mapped bytes. SDOT path when
9671/// dotprod is available (port of vmfcore `dot_q4_block_sdot`, measured
9672/// +23% on q4 decode): nibbles → centered i8, int8×int8 `sdot` per
9673/// 32-group, exact outlier correction — the same A8W8 contract as q8.
9674/// `CMF_SDOT=0` keeps the exact scalar path.
9675fn q4matvec(
9676    bytes: &[u8],
9677    x: &[f32],
9678    rows: usize,
9679    cols: usize,
9680    out: &mut [f32],
9681    pool: Option<&Pool>,
9682) {
9683    debug_assert_eq!(out.len(), rows);
9684    let (packed, scales) = q4_split(bytes, rows, cols);
9685    let gpr = cols / GROUP_SIZE;
9686    let out_addr = SendMut(out.as_mut_ptr());
9687
9688    if a8w8_enabled() {
9689        let act = split_act(x);
9690        let run = move |start: usize, end: usize| {
9691            q4_range_a8w8(packed, scales, gpr, cols, &act, out_addr, start, end)
9692        };
9693        dispatch_rows(pool, rows, &run);
9694        return;
9695    }
9696
9697    let run =
9698        move |start: usize, end: usize| q4_range_f32(packed, scales, gpr, x, out_addr, start, end);
9699    dispatch_rows(pool, rows, &run);
9700}
9701
9702/// One q4 row via the A8W8 int8 path — SDOT on ARM, AVX2 maddubs on
9703/// x86 (scalar fallback is unreachable: callers gate on a8w8_enabled).
9704#[inline]
9705#[allow(unreachable_code)]
9706/// One UNPACKED q4 row (centered i8 in `buf`) against four activation
9707/// streams: the 32-byte weight chunk and its abs() load once per group,
9708/// the per-group f16 scale decodes once — four maddubs+reduce chains
9709/// instead of four full (load, abs, dot) rounds.
9710#[cfg(target_arch = "x86_64")]
9711#[target_feature(enable = "avx2")]
9712unsafe fn dot_q4b_row_1x4_avx2(
9713    buf: &[u8],
9714    scales: &[u8],
9715    g0: usize,
9716    gpr: usize,
9717    xs: [&[i8]; 4],
9718) -> [f32; 4] {
9719    // SAFETY: callers uphold buffer contracts (buf.len() == gpr·32).
9720    unsafe {
9721        use core::arch::x86_64::*;
9722        let ones = _mm256_set1_epi16(1);
9723        let mut acc = [0f32; 4];
9724        for gi in 0..gpr {
9725            let s = f16_to_f32(u16::from_le_bytes([
9726                scales[(g0 + gi) * 2],
9727                scales[(g0 + gi) * 2 + 1],
9728            ]));
9729            let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9730            let aw = _mm256_abs_epi8(w);
9731            for (k, xq) in xs.iter().enumerate() {
9732                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9733                let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
9734                let d = _mm256_madd_epi16(p16, ones);
9735                let hi128 = _mm256_extracti128_si256::<1>(d);
9736                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
9737                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
9738                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
9739                acc[k] += _mm_cvtsi128_si32(s32) as f32 * s;
9740            }
9741        }
9742        acc
9743    }
9744}
9745
9746/// VNNI twin of `dot_q4b_row_1x4_avx2` (see `dpbusd_hsum`).
9747#[cfg(target_arch = "x86_64")]
9748#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
9749unsafe fn dot_q4b_row_1x4_vnni(
9750    buf: &[u8],
9751    scales: &[u8],
9752    g0: usize,
9753    gpr: usize,
9754    xs: [&[i8]; 4],
9755) -> [f32; 4] {
9756    // SAFETY: callers uphold buffer contracts (buf.len() == gpr·32).
9757    unsafe {
9758        use core::arch::x86_64::*;
9759        let mut acc = [0f32; 4];
9760        for gi in 0..gpr {
9761            let s = f16_to_f32(u16::from_le_bytes([
9762                scales[(g0 + gi) * 2],
9763                scales[(g0 + gi) * 2 + 1],
9764            ]));
9765            let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9766            let aw = _mm256_abs_epi8(w);
9767            for (k, xq) in xs.iter().enumerate() {
9768                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9769                let d = dpbusd_hsum(aw, _mm256_sign_epi8(x, w));
9770                acc[k] += d as f32 * s;
9771            }
9772        }
9773        acc
9774    }
9775}
9776
9777/// The vbit flavor of the blocked 1×4: the per-activation A8W8 scale
9778/// folds in PER GROUP as `(d·sx)·s` — bit-matching the single-matvec
9779/// accumulation order (the q4_block flavor applies sx once at the end,
9780/// matching ITS single path; the two conventions are historical and
9781/// each blocked leg must mirror its own).
9782#[cfg(target_arch = "x86_64")]
9783#[target_feature(enable = "avx2")]
9784unsafe fn dot_q4b_row_1x4_sx_avx2(
9785    buf: &[u8],
9786    scales: &[u8],
9787    g0: usize,
9788    gpr: usize,
9789    xs: [&[i8]; 4],
9790    sxs: [f32; 4],
9791) -> [f32; 4] {
9792    // SAFETY: callers uphold buffer contracts (buf.len() == gpr·32).
9793    unsafe {
9794        use core::arch::x86_64::*;
9795        let ones = _mm256_set1_epi16(1);
9796        let mut acc = [0f32; 4];
9797        for gi in 0..gpr {
9798            let s = f16_to_f32(u16::from_le_bytes([
9799                scales[(g0 + gi) * 2],
9800                scales[(g0 + gi) * 2 + 1],
9801            ]));
9802            let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9803            let aw = _mm256_abs_epi8(w);
9804            for (k, xq) in xs.iter().enumerate() {
9805                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9806                let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
9807                let d = _mm256_madd_epi16(p16, ones);
9808                let hi128 = _mm256_extracti128_si256::<1>(d);
9809                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
9810                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
9811                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
9812                acc[k] += (_mm_cvtsi128_si32(s32) as f32 * sxs[k]) * s;
9813            }
9814        }
9815        acc
9816    }
9817}
9818
9819/// VNNI twin of `dot_q4b_row_1x4_sx_avx2` (see `dpbusd_hsum`; the
9820/// per-group `(d·sx)·s` fold mirrors the vbit single path).
9821#[cfg(target_arch = "x86_64")]
9822#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
9823unsafe fn dot_q4b_row_1x4_sx_vnni(
9824    buf: &[u8],
9825    scales: &[u8],
9826    g0: usize,
9827    gpr: usize,
9828    xs: [&[i8]; 4],
9829    sxs: [f32; 4],
9830) -> [f32; 4] {
9831    // SAFETY: callers uphold buffer contracts (buf.len() == gpr·32).
9832    unsafe {
9833        use core::arch::x86_64::*;
9834        let mut acc = [0f32; 4];
9835        for gi in 0..gpr {
9836            let s = f16_to_f32(u16::from_le_bytes([
9837                scales[(g0 + gi) * 2],
9838                scales[(g0 + gi) * 2 + 1],
9839            ]));
9840            let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9841            let aw = _mm256_abs_epi8(w);
9842            for (k, xq) in xs.iter().enumerate() {
9843                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9844                let d = dpbusd_hsum(aw, _mm256_sign_epi8(x, w));
9845                acc[k] += (d as f32 * sxs[k]) * s;
9846            }
9847        }
9848        acc
9849    }
9850}
9851
9852#[allow(unreachable_code)]
9853fn dot_q4_row_i8(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
9854    #[cfg(target_arch = "aarch64")]
9855    unsafe {
9856        return dot_q4_row_sdot(packed, scales, g0, gpr, xq);
9857    }
9858    #[cfg(target_arch = "x86_64")]
9859    unsafe {
9860        return dot_q4_row_avx2(packed, scales, g0, gpr, xq);
9861    }
9862    let mut acc = 0f32;
9863    for gi in 0..gpr {
9864        let g = g0 + gi;
9865        let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
9866        let mut d = 0i32;
9867        for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
9868            d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
9869                + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
9870        }
9871        acc += d as f32 * s;
9872    }
9873    acc
9874}
9875
9876/// Two-activation q4 row via the A8W8 int8 path (see `dot_q4_row_i8`).
9877#[inline]
9878#[allow(unreachable_code)]
9879fn dot_q4_row_i8_2(
9880    packed: &[u8],
9881    scales: &[u8],
9882    g0: usize,
9883    gpr: usize,
9884    xq1: &[i8],
9885    xq2: &[i8],
9886) -> (f32, f32) {
9887    #[cfg(target_arch = "aarch64")]
9888    unsafe {
9889        return dot_q4_row_sdot2(packed, scales, g0, gpr, xq1, xq2);
9890    }
9891    #[cfg(target_arch = "x86_64")]
9892    unsafe {
9893        return dot_q4_row_avx2_2(packed, scales, g0, gpr, xq1, xq2);
9894    }
9895    (
9896        dot_q4_row_i8(packed, scales, g0, gpr, xq1),
9897        dot_q4_row_i8(packed, scales, g0, gpr, xq2),
9898    )
9899}
9900
9901/// One q4 row range via SDOT (kernel body of `q4matvec`, extracted so
9902/// multi-matrix jobs can drive it for several tensors in one dispatch).
9903#[allow(clippy::too_many_arguments)]
9904fn q4_range_a8w8(
9905    packed: &[u8],
9906    scales: &[u8],
9907    gpr: usize,
9908    cols: usize,
9909    act: &SplitAct,
9910    out: SendMut,
9911    start: usize,
9912    end: usize,
9913) {
9914    for r in start..end {
9915        let mut acc = dot_q4_row_i8(packed, scales, r * gpr, gpr, &act.xq) * act.sx;
9916        // xq is zeroed at outlier slots — add the exact terms.
9917        for &(j, xv) in &act.outliers {
9918            let flat = r * cols + j;
9919            let byte = packed[flat / 2];
9920            let nib = if flat & 1 == 0 {
9921                byte & 0x0F
9922            } else {
9923                byte >> 4
9924            };
9925            let s = f16_to_f32(u16::from_le_bytes([
9926                scales[(flat / GROUP_SIZE) * 2],
9927                scales[(flat / GROUP_SIZE) * 2 + 1],
9928            ]));
9929            acc += ((nib as i32 - 8) as f32) * s * xv;
9930        }
9931        // SAFETY: disjoint row ranges per worker.
9932        unsafe { *out.at(r) = acc };
9933    }
9934}
9935
9936/// Two-input q4 row range via the A8W8 int8 path — kernel body of
9937/// `q4matvec2`, extracted for pair multi-matrix jobs.
9938#[allow(clippy::too_many_arguments)]
9939fn q4_range2_a8w8(
9940    packed: &[u8],
9941    scales: &[u8],
9942    gpr: usize,
9943    cols: usize,
9944    a1: &SplitAct,
9945    a2: &SplitAct,
9946    p1: SendMut,
9947    p2: SendMut,
9948    start: usize,
9949    end: usize,
9950) {
9951    for r in start..end {
9952        let (s1, s2) = dot_q4_row_i8_2(packed, scales, r * gpr, gpr, &a1.xq, &a2.xq);
9953        let mut acc1 = s1 * a1.sx;
9954        let mut acc2 = s2 * a2.sx;
9955        // xq is zeroed at outlier slots — add the exact terms.
9956        let fix = |outliers: &[(usize, f32)], acc: &mut f32| {
9957            for &(j, xv) in outliers {
9958                let flat = r * cols + j;
9959                let byte = packed[flat / 2];
9960                let nib = if flat & 1 == 0 {
9961                    byte & 0x0F
9962                } else {
9963                    byte >> 4
9964                };
9965                let s = f16_to_f32(u16::from_le_bytes([
9966                    scales[(flat / GROUP_SIZE) * 2],
9967                    scales[(flat / GROUP_SIZE) * 2 + 1],
9968                ]));
9969                *acc += ((nib as i32 - 8) as f32) * s * xv;
9970            }
9971        };
9972        fix(&a1.outliers, &mut acc1);
9973        fix(&a2.outliers, &mut acc2);
9974        // SAFETY: disjoint row ranges per worker.
9975        unsafe {
9976            *p1.at(r) = acc1;
9977            *p2.at(r) = acc2;
9978        }
9979    }
9980}
9981
9982/// Exact scalar q4 row range (same extraction, non-SDOT path).
9983fn q4_range_f32(
9984    packed: &[u8],
9985    scales: &[u8],
9986    gpr: usize,
9987    x: &[f32],
9988    out: SendMut,
9989    start: usize,
9990    end: usize,
9991) {
9992    for r in start..end {
9993        let mut acc = 0f32;
9994        for gi in 0..gpr {
9995            let g = r * gpr + gi;
9996            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
9997            let pk = &packed[g * 16..(g + 1) * 16];
9998            let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
9999            let mut ga = 0f32;
10000            for (k, &b) in pk.iter().enumerate() {
10001                ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
10002                    + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
10003            }
10004            acc += ga * s;
10005        }
10006        // SAFETY: disjoint row ranges per worker.
10007        unsafe { *out.at(r) = acc };
10008    }
10009}
10010
10011/// Fused two-input q4 matvec: nibbles are unpacked ONCE per group and
10012/// dotted against both activations (was: two full matvecs — double
10013/// weight traffic). Per-lane math matches `q4matvec` exactly.
10014#[allow(clippy::too_many_arguments)]
10015fn q4matvec2(
10016    bytes: &[u8],
10017    x1: &[f32],
10018    x2: &[f32],
10019    rows: usize,
10020    cols: usize,
10021    o1: &mut [f32],
10022    o2: &mut [f32],
10023    pool: Option<&Pool>,
10024) {
10025    debug_assert_eq!(o1.len(), rows);
10026    debug_assert_eq!(o2.len(), rows);
10027    let (packed, scales) = q4_split(bytes, rows, cols);
10028    let gpr = cols / GROUP_SIZE;
10029
10030    if a8w8_enabled() {
10031        let a1 = split_act(x1);
10032        let a2 = split_act(x2);
10033        let p1 = SendMut(o1.as_mut_ptr());
10034        let p2 = SendMut(o2.as_mut_ptr());
10035        let run = move |start: usize, end: usize| {
10036            q4_range2_a8w8(packed, scales, gpr, cols, &a1, &a2, p1, p2, start, end)
10037        };
10038        dispatch_rows(pool, rows, &run);
10039        return;
10040    }
10041
10042    let p1 = SendMut(o1.as_mut_ptr());
10043    let p2 = SendMut(o2.as_mut_ptr());
10044    let run = move |start: usize, end: usize| {
10045        q4_range2_f32(packed, scales, gpr, x1, x2, p1, p2, start, end)
10046    };
10047    dispatch_rows(pool, rows, &run);
10048}
10049
10050/// Two-input exact scalar q4 row range (same extraction).
10051#[allow(clippy::too_many_arguments)]
10052fn q4_range2_f32(
10053    packed: &[u8],
10054    scales: &[u8],
10055    gpr: usize,
10056    x1: &[f32],
10057    x2: &[f32],
10058    p1: SendMut,
10059    p2: SendMut,
10060    start: usize,
10061    end: usize,
10062) {
10063    for r in start..end {
10064        let (mut acc1, mut acc2) = (0f32, 0f32);
10065        for gi in 0..gpr {
10066            let g = r * gpr + gi;
10067            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
10068            let pk = &packed[g * 16..(g + 1) * 16];
10069            let x1g = &x1[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
10070            let x2g = &x2[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
10071            let (mut g1, mut g2) = (0f32, 0f32);
10072            for (k, &b) in pk.iter().enumerate() {
10073                let wl = (b & 0x0F) as f32 - 8.0;
10074                let wh = ((b >> 4) & 0x0F) as f32 - 8.0;
10075                g1 += wl * x1g[k * 2] + wh * x1g[k * 2 + 1];
10076                g2 += wl * x2g[k * 2] + wh * x2g[k * 2 + 1];
10077            }
10078            acc1 += g1 * s;
10079            acc2 += g2 * s;
10080        }
10081        // SAFETY: disjoint row ranges per worker.
10082        unsafe {
10083            *p1.at(r) = acc1;
10084            *p2.at(r) = acc2;
10085        }
10086    }
10087}
10088
10089thread_local! {
10090    /// Per-worker decoded-row scratch for the batched q4/vbit kernels
10091    /// (centered i8 for SDOT, f32 for the exact/scalar paths).
10092    static ROW_I8: std::cell::RefCell<Vec<u8>> = const { std::cell::RefCell::new(Vec::new()) };
10093    static ROW_F32: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
10094}
10095
10096/// Batched q4 matmat: each weight row is unpacked from the mmap ONCE
10097/// and dotted against ALL b activations (prefill used to fall back to b
10098/// full matvecs — b× weight traffic and b× nibble decode). Per-position
10099/// math matches `q4matvec` exactly: same group order, same accumulation.
10100/// `out` is row-major [b, rows] like `qmatmat`.
10101#[allow(clippy::too_many_arguments)]
10102fn q4matmat(
10103    bytes: &[u8],
10104    xs_all: &[f32],
10105    b: usize,
10106    rows: usize,
10107    cols: usize,
10108    out: &mut [f32],
10109    pool: Option<&Pool>,
10110) {
10111    debug_assert_eq!(xs_all.len(), b * cols);
10112    debug_assert_eq!(out.len(), b * rows);
10113    let (packed, scales) = q4_split(bytes, rows, cols);
10114    let gpr = cols / GROUP_SIZE;
10115    let gscale = |g: usize| f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
10116
10117    if a8w8_enabled() {
10118        let acts: Vec<SplitAct> = (0..b)
10119            .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
10120            .collect();
10121        let acts = &acts;
10122        let out_addr = SendMut(out.as_mut_ptr());
10123        let run = move |start: usize, end: usize| {
10124            ROW_I8.with(|rb| {
10125                let mut buf = rb.borrow_mut();
10126                buf.resize(cols, 0);
10127                for r in start..end {
10128                    // Unpack the row's nibbles to centered i8 once
10129                    // (element 2k = low nibble, 2k+1 = high — flat order,
10130                    // same as dot_q4_row_sdot's zip).
10131                    for gi in 0..gpr {
10132                        let g = r * gpr + gi;
10133                        for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
10134                            buf[gi * GROUP_SIZE + k * 2] = ((bt & 0x0F) as i32 - 8) as i8 as u8;
10135                            buf[gi * GROUP_SIZE + k * 2 + 1] =
10136                                (((bt >> 4) & 0x0F) as i32 - 8) as i8 as u8;
10137                        }
10138                    }
10139                    let mut bi = 0usize;
10140                    #[cfg(target_arch = "x86_64")]
10141                    if avx2_enabled() && blocked_enabled() {
10142                        while bi + 4 <= acts.len() {
10143                            let xs = [
10144                                acts[bi].xq.as_slice(),
10145                                acts[bi + 1].xq.as_slice(),
10146                                acts[bi + 2].xq.as_slice(),
10147                                acts[bi + 3].xq.as_slice(),
10148                            ];
10149                            let d = unsafe {
10150                                if vnni_tiles_enabled() {
10151                                    dot_q4b_row_1x4_vnni(&buf, scales, r * gpr, gpr, xs)
10152                                } else {
10153                                    dot_q4b_row_1x4_avx2(&buf, scales, r * gpr, gpr, xs)
10154                                }
10155                            };
10156                            for k in 0..4 {
10157                                let act = &acts[bi + k];
10158                                let mut acc = d[k] * act.sx;
10159                                for &(j, xv) in &act.outliers {
10160                                    acc += (buf[j] as i8) as f32
10161                                        * gscale((r * cols + j) / GROUP_SIZE)
10162                                        * xv;
10163                                }
10164                                // SAFETY: disjoint (bi, r) cells per worker.
10165                                unsafe { *out_addr.at((bi + k) * rows + r) = acc };
10166                            }
10167                            bi += 4;
10168                        }
10169                    }
10170                    while bi < acts.len() {
10171                        let act = &acts[bi];
10172                        let mut acc = 0f32;
10173                        for gi in 0..gpr {
10174                            let d = dot_i8_i8(
10175                                &buf[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
10176                                &act.xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
10177                            );
10178                            acc += d as f32 * gscale(r * gpr + gi);
10179                        }
10180                        acc *= act.sx;
10181                        // xq is zeroed at outlier slots — exact terms.
10182                        for &(j, xv) in &act.outliers {
10183                            acc += (buf[j] as i8) as f32 * gscale((r * cols + j) / GROUP_SIZE) * xv;
10184                        }
10185                        // SAFETY: disjoint (bi, r) cells per worker row range.
10186                        unsafe { *out_addr.at(bi * rows + r) = acc };
10187                        bi += 1;
10188                    }
10189                }
10190            })
10191        };
10192        dispatch_rows(pool, rows, &run);
10193        return;
10194    }
10195
10196    let out_addr = SendMut(out.as_mut_ptr());
10197    let run = move |start: usize, end: usize| {
10198        ROW_F32.with(|rb| {
10199            let mut buf = rb.borrow_mut();
10200            buf.resize(cols, 0.0);
10201            for r in start..end {
10202                // Decode raw (nib − 8) values once; scales stay per-group
10203                // so the accumulation order matches q4matvec bit-for-bit.
10204                for gi in 0..gpr {
10205                    let g = r * gpr + gi;
10206                    for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
10207                        buf[gi * GROUP_SIZE + k * 2] = (bt & 0x0F) as f32 - 8.0;
10208                        buf[gi * GROUP_SIZE + k * 2 + 1] = ((bt >> 4) & 0x0F) as f32 - 8.0;
10209                    }
10210                }
10211                for bi in 0..b {
10212                    let x = &xs_all[bi * cols..(bi + 1) * cols];
10213                    let mut acc = 0f32;
10214                    for gi in 0..gpr {
10215                        let mut ga = 0f32;
10216                        // Pairwise (lo + hi) addition, matching
10217                        // q4matvec's `ga += lo·x + hi·x` shape exactly —
10218                        // a flat one-per-element loop rounds differently
10219                        // and broke bit-parity on the scalar (x86) path.
10220                        for k in 0..GROUP_SIZE / 2 {
10221                            let e = gi * GROUP_SIZE + k * 2;
10222                            ga += buf[e] * x[e] + buf[e + 1] * x[e + 1];
10223                        }
10224                        acc += ga * gscale(r * gpr + gi);
10225                    }
10226                    // SAFETY: disjoint (bi, r) cells per worker row range.
10227                    unsafe { *out_addr.at(bi * rows + r) = acc };
10228                }
10229            }
10230        })
10231    };
10232    dispatch_rows(pool, rows, &run);
10233}
10234
10235/// Batched vbit matmat: each variable-bit row is decoded from the mmap
10236/// ONCE for the whole microbatch. Same per-position math as
10237/// `vbitmatvec` (SDOT A8W8 with exact outliers / exact f32 for b=8 rows
10238/// and the scalar path).
10239#[allow(clippy::too_many_arguments)]
10240fn vbitmatmat(
10241    bytes: &[u8],
10242    offsets: &[usize],
10243    xs_all: &[f32],
10244    b: usize,
10245    rows: usize,
10246    cols: usize,
10247    out: &mut [f32],
10248    pool: Option<&Pool>,
10249) {
10250    debug_assert_eq!(xs_all.len(), b * cols);
10251    debug_assert_eq!(out.len(), b * rows);
10252    debug_assert_eq!(offsets.len(), rows + 1);
10253    let ng = cols / GROUP_SIZE;
10254    let bits = &bytes[..rows];
10255    let sc_off = rows;
10256    let gscale = |r: usize, g: usize| {
10257        let so = (r * ng + g) * 2;
10258        f16_to_f32(u16::from_le_bytes([
10259            bytes[sc_off + so],
10260            bytes[sc_off + so + 1],
10261        ]))
10262    };
10263
10264    // Decode row r's raw (u − L) values into `dst` (f32, unscaled).
10265    let decode_f32 = |r: usize, dst: &mut [f32]| {
10266        let bw = bits[r] as usize;
10267        let l = ((1i32 << (bw - 1)) - 1) as f32;
10268        let data = &bytes[offsets[r]..offsets[r + 1]];
10269        let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
10270        for d in dst.iter_mut() {
10271            while nbits < bw {
10272                acc = (acc << 8) | data[idx] as u64;
10273                idx += 1;
10274                nbits += 8;
10275            }
10276            let u = ((acc >> (nbits - bw)) & ((1u64 << bw) - 1)) as f32;
10277            nbits -= bw;
10278            *d = u - l;
10279        }
10280    };
10281
10282    if a8w8_enabled() {
10283        let acts: Vec<SplitAct> = (0..b)
10284            .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
10285            .collect();
10286        let acts = &acts;
10287        let out_addr = SendMut(out.as_mut_ptr());
10288        let run = move |start: usize, end: usize| {
10289            for r in start..end {
10290                let bw = bits[r] as usize;
10291                if bw == 8 {
10292                    // u−L reaches 128 → no i8 path; decode once, exact
10293                    // f32 dots for every position (same as vbitmatvec).
10294                    ROW_F32.with(|rb| {
10295                        let mut buf = rb.borrow_mut();
10296                        buf.resize(cols, 0.0);
10297                        decode_f32(r, &mut buf);
10298                        for bi in 0..b {
10299                            let x = &xs_all[bi * cols..(bi + 1) * cols];
10300                            let mut dot = 0f32;
10301                            for g in 0..ng {
10302                                let mut gd = 0f32;
10303                                for k in 0..GROUP_SIZE {
10304                                    gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
10305                                }
10306                                dot += gd * gscale(r, g);
10307                            }
10308                            // SAFETY: disjoint (bi, r) cells per worker range.
10309                            unsafe { *out_addr.at(bi * rows + r) = dot };
10310                        }
10311                    });
10312                    continue;
10313                }
10314                let l = (1i32 << (bw - 1)) - 1;
10315                let data = &bytes[offsets[r]..offsets[r + 1]];
10316                ROW_I8.with(|rb| {
10317                    let mut buf = rb.borrow_mut();
10318                    buf.resize(cols, 0);
10319                    #[inline(always)]
10320                    fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
10321                        for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
10322                            let u = unpack8::<B>(&data[blk * B..]);
10323                            for k in 0..8 {
10324                                chunk[k] = (u[k] - l) as i8 as u8;
10325                            }
10326                        }
10327                    }
10328                    match bw {
10329                        3 => fill::<3>(data, l, &mut buf),
10330                        4 => vbit_fill4(data, &mut buf),
10331                        5 => fill::<5>(data, l, &mut buf),
10332                        6 => fill::<6>(data, l, &mut buf),
10333                        _ => unreachable!("vbit bit-width {bw} (validated at load)"),
10334                    }
10335                    let mut bi = 0usize;
10336                    // The vbit scale table shares q4_block's layout
10337                    // (contiguous f16 per (row·ng + g)), so the same
10338                    // blocked 1×4 kernel serves the decoded row.
10339                    #[cfg(target_arch = "x86_64")]
10340                    if avx2_enabled() && blocked_enabled() {
10341                        while bi + 4 <= acts.len() {
10342                            let xs = [
10343                                acts[bi].xq.as_slice(),
10344                                acts[bi + 1].xq.as_slice(),
10345                                acts[bi + 2].xq.as_slice(),
10346                                acts[bi + 3].xq.as_slice(),
10347                            ];
10348                            let sxs = [
10349                                acts[bi].sx,
10350                                acts[bi + 1].sx,
10351                                acts[bi + 2].sx,
10352                                acts[bi + 3].sx,
10353                            ];
10354                            let d = unsafe {
10355                                if vnni_tiles_enabled() {
10356                                    dot_q4b_row_1x4_sx_vnni(
10357                                        &buf,
10358                                        &bytes[sc_off..],
10359                                        r * ng,
10360                                        ng,
10361                                        xs,
10362                                        sxs,
10363                                    )
10364                                } else {
10365                                    dot_q4b_row_1x4_sx_avx2(
10366                                        &buf,
10367                                        &bytes[sc_off..],
10368                                        r * ng,
10369                                        ng,
10370                                        xs,
10371                                        sxs,
10372                                    )
10373                                }
10374                            };
10375                            for k in 0..4 {
10376                                let act = &acts[bi + k];
10377                                let mut dot = d[k];
10378                                for &(j, xv) in &act.outliers {
10379                                    dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
10380                                }
10381                                // SAFETY: disjoint (bi, r) cells per worker.
10382                                unsafe { *out_addr.at((bi + k) * rows + r) = dot };
10383                            }
10384                            bi += 4;
10385                        }
10386                    }
10387                    while bi < acts.len() {
10388                        let act = &acts[bi];
10389                        let mut dot = 0f32;
10390                        for g in 0..ng {
10391                            let d = dot_i8_i8(
10392                                &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
10393                                &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
10394                            ) as f32
10395                                * act.sx;
10396                            dot += d * gscale(r, g);
10397                        }
10398                        for &(j, xv) in &act.outliers {
10399                            dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
10400                        }
10401                        // SAFETY: disjoint (bi, r) cells per worker range.
10402                        unsafe { *out_addr.at(bi * rows + r) = dot };
10403                        bi += 1;
10404                    }
10405                });
10406            }
10407        };
10408        dispatch_rows(pool, rows, &run);
10409        return;
10410    }
10411
10412    let out_addr = SendMut(out.as_mut_ptr());
10413    let run = move |start: usize, end: usize| {
10414        ROW_F32.with(|rb| {
10415            let mut buf = rb.borrow_mut();
10416            buf.resize(cols, 0.0);
10417            for r in start..end {
10418                decode_f32(r, &mut buf);
10419                for bi in 0..b {
10420                    let x = &xs_all[bi * cols..(bi + 1) * cols];
10421                    let mut dot = 0f32;
10422                    for g in 0..ng {
10423                        let mut gd = 0f32;
10424                        for k in 0..GROUP_SIZE {
10425                            gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
10426                        }
10427                        dot += gd * gscale(r, g);
10428                    }
10429                    // SAFETY: disjoint (bi, r) cells per worker range.
10430                    unsafe { *out_addr.at(bi * rows + r) = dot };
10431                }
10432            }
10433        })
10434    };
10435    dispatch_rows(pool, rows, &run);
10436}
10437
10438/// Build a GPU batch job for a q8-family mapped tensor (primary
10439/// shard): prescaled input + directory coordinates. None → not
10440/// GPU-eligible, caller stays on the CPU.
10441pub(crate) fn gpu_batch_job<'a>(
10442    t: &'a QTensor,
10443    x: &[f32],
10444) -> Option<(std::sync::Arc<CmfModel>, crate::gpu::BatchJob<'a>)> {
10445    match t {
10446        QTensor::Mapped {
10447            model,
10448            idx,
10449            dtype: dt @ (TensorDtype::Q8Row | TensorDtype::Q8_2f),
10450            rows,
10451            cols,
10452            row_scale,
10453            col_field,
10454            ..
10455        } => Some((
10456            model.clone(),
10457            crate::gpu::BatchJob {
10458                idx: *idx,
10459                rows: *rows,
10460                cols: *cols,
10461                row_scale,
10462                xs: prescale(x, col_field, *dt).into_owned(),
10463                layout: crate::gpu::BatchLayout::Q8,
10464            },
10465        )),
10466        // q1: raw f32 activations, tile-embedded scales.
10467        QTensor::Mapped {
10468            model,
10469            idx,
10470            dtype: TensorDtype::Q1,
10471            rows,
10472            cols,
10473            ..
10474        } => Some((
10475            model.clone(),
10476            crate::gpu::BatchJob {
10477                idx: *idx,
10478                rows: *rows,
10479                cols: *cols,
10480                row_scale: &[],
10481                xs: x.to_vec(),
10482                layout: crate::gpu::BatchLayout::Q1,
10483            },
10484        )),
10485        // q1g: same raw-f32 contract, sign + scale planes.
10486        QTensor::Mapped {
10487            model,
10488            idx,
10489            dtype: TensorDtype::Q1G,
10490            rows,
10491            cols,
10492            ..
10493        } => Some((
10494            model.clone(),
10495            crate::gpu::BatchJob {
10496                idx: *idx,
10497                rows: *rows,
10498                cols: *cols,
10499                row_scale: &[],
10500                xs: x.to_vec(),
10501                layout: crate::gpu::BatchLayout::Q1g,
10502            },
10503        )),
10504        // q4_tiled / q4tp: raw f32 activations; the scales live in the
10505        // payload (inline tiles / row ladder), so row_scale stays empty.
10506        // The GDN projection batch already runs these layouts on Metal —
10507        // this arm lets the attention QKV batch reach the same kernels.
10508        QTensor::Mapped {
10509            model,
10510            idx,
10511            dtype: dt @ (TensorDtype::Q4Tiled | TensorDtype::Q4TiledP),
10512            rows,
10513            cols,
10514            ..
10515        } => Some((
10516            model.clone(),
10517            crate::gpu::BatchJob {
10518                idx: *idx,
10519                rows: *rows,
10520                cols: *cols,
10521                row_scale: &[],
10522                xs: x.to_vec(),
10523                layout: if *dt == TensorDtype::Q4Tiled {
10524                    crate::gpu::BatchLayout::Q4t
10525                } else {
10526                    crate::gpu::BatchLayout::Q4tp
10527                },
10528            },
10529        )),
10530        _ => None,
10531    }
10532}
10533
10534thread_local! {
10535    static PRESCALE_BUF1: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
10536    static PRESCALE_BUF2: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
10537}
10538
10539pub(crate) fn prescale<'a>(
10540    x: &'a [f32],
10541    col_field: &[f32],
10542    dtype: TensorDtype,
10543) -> std::borrow::Cow<'a, [f32]> {
10544    if dtype == TensorDtype::Q8_2f {
10545        x.iter().zip(col_field).map(|(a, c)| a * c).collect()
10546    } else {
10547        std::borrow::Cow::Borrowed(x)
10548    }
10549}
10550
10551/// θ col-field fold for q8_2f activations. Borrowed pass-through for
10552/// every other dtype, using thread-local buffers to eliminate per-matvec allocations.
10553pub(crate) fn prescale_with<R, F: FnOnce(&[f32]) -> R>(
10554    x: &[f32],
10555    col_field: &[f32],
10556    dtype: TensorDtype,
10557    buf_id: u8,
10558    f: F,
10559) -> R {
10560    if dtype == TensorDtype::Q8_2f {
10561        if buf_id == 1 {
10562            PRESCALE_BUF1.with(|b| {
10563                let mut buf = b.borrow_mut();
10564                buf.clear();
10565                buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
10566                f(&buf)
10567            })
10568        } else {
10569            PRESCALE_BUF2.with(|b| {
10570                let mut buf = b.borrow_mut();
10571                buf.clear();
10572                buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
10573                f(&buf)
10574            })
10575        }
10576    } else {
10577        f(x)
10578    }
10579}
10580
10581// ───────────────────── x86-64 AVX2 kernels (roadmap этап 2) ─────────────────────
10582
10583/// AVX2+FMA available? Default ON when the CPU supports both;
10584/// `CMF_AVX2=0` disables (falls back to the autovectorized loops).
10585#[cfg(target_arch = "x86_64")]
10586pub(crate) fn avx2_enabled() -> bool {
10587    use std::sync::OnceLock;
10588    static ON: OnceLock<bool> = OnceLock::new();
10589    *ON.get_or_init(|| {
10590        std::env::var("CMF_AVX2").map(|v| v != "0").unwrap_or(true)
10591            && std::arch::is_x86_feature_detected!("avx2")
10592            && std::arch::is_x86_feature_detected!("fma")
10593    })
10594}
10595
10596/// AVX2 A8W8 allowed? The quantized-activation contract is switched by
10597/// the SAME env as the ARM SDOT path: `CMF_SDOT=0` keeps exact kernels
10598/// (the golden-parity exact gate relies on it) — AVX2 f32 kernels stay
10599/// active either way, they are exact (regrouped sums only).
10600#[cfg(target_arch = "x86_64")]
10601fn avx2_a8w8_enabled() -> bool {
10602    if FLOAT_ACTIVATIONS.get() {
10603        return false;
10604    }
10605    use std::sync::OnceLock;
10606    static ON: OnceLock<bool> = OnceLock::new();
10607    *ON.get_or_init(|| {
10608        avx2_enabled() && std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true)
10609    })
10610}
10611
10612thread_local! {
10613    static FULL_GPU_Q8: std::cell::Cell<bool> = const { std::cell::Cell::new(false) };
10614}
10615
10616/// Match the graph's full-device q8 projection precision on MiMo's host
10617/// tail. Does not enable the GPU or bypass a CPU-only/device-refusal gate.
10618pub(crate) fn enter_full_gpu_q8_scope() -> impl Drop {
10619    struct Restore(bool, std::marker::PhantomData<std::rc::Rc<()>>);
10620    impl Drop for Restore {
10621        fn drop(&mut self) {
10622            FULL_GPU_Q8.set(self.0);
10623        }
10624    }
10625    Restore(FULL_GPU_Q8.replace(true), std::marker::PhantomData)
10626}
10627
10628// Dynamic MiMo experts must not change activation precision when a cache
10629// fill moves them from CPU to GPU. Thread-local: only the cold-expert
10630// dispatch selects float kernels; concurrent pipelines keep their policy.
10631thread_local! {
10632    static FLOAT_ACTIVATIONS: std::cell::Cell<bool> = const { std::cell::Cell::new(false) };
10633}
10634
10635pub(crate) fn float_activations_scope<R>(f: impl FnOnce() -> R) -> R {
10636    struct Restore(bool);
10637    impl Drop for Restore {
10638        fn drop(&mut self) {
10639            FLOAT_ACTIVATIONS.set(self.0);
10640        }
10641    }
10642    let _restore = Restore(FLOAT_ACTIVATIONS.replace(true));
10643    f()
10644}
10645
10646/// Row-exact batching: while set, the x86 batched kernels (`qmatmat`,
10647/// `q4tp_matmat`) compute every (weight row, token) cell with the
10648/// single-token kernel instead of the blocked 2×4 / 1×4 / 1×8 tiles, so a
10649/// token's result does not depend on the batch it rides in and equals its
10650/// matvec. On ARM `q4tp_matmat` keeps its 1×4 tile but in the matvec's
10651/// reduction order, and no q4tp batch takes the AMX or device GEMM. The
10652/// MiMo speculative verify holds it (`row_exact_scope`) — its accepted
10653/// rows must be the rows plain decode would have produced.
10654// Shared with pool workers, so overlapping requests must keep the mode
10655// enabled until the LAST scope leaves. Saving/restoring a global bool is
10656// incorrect when two threads enter and leave in a non-LIFO order.
10657static ROW_EXACT: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
10658
10659pub(crate) fn row_exact() -> bool {
10660    ROW_EXACT.load(std::sync::atomic::Ordering::Acquire) != 0
10661}
10662
10663fn counted_row_exact_scope<R>(active: &std::sync::atomic::AtomicUsize, f: impl FnOnce() -> R) -> R {
10664    struct Restore<'a>(&'a std::sync::atomic::AtomicUsize);
10665    impl Drop for Restore<'_> {
10666        fn drop(&mut self) {
10667            self.0.fetch_sub(1, std::sync::atomic::Ordering::AcqRel);
10668        }
10669    }
10670    active.fetch_add(1, std::sync::atomic::Ordering::AcqRel);
10671    let _restore = Restore(active);
10672    f()
10673}
10674
10675/// Run `f` with row-exact batching on (also released on unwind).
10676pub(crate) fn row_exact_scope<R>(f: impl FnOnce() -> R) -> R {
10677    counted_row_exact_scope(&ROW_EXACT, f)
10678}
10679
10680/// A8W8 quantized-activation path available on THIS machine? One
10681/// switch across architectures: ARM dotprod (CMF_SDOT) or x86 AVX2
10682/// (CMF_AVX2 + the same CMF_SDOT exact-contract override).
10683#[inline]
10684pub(crate) fn a8w8_enabled() -> bool {
10685    #[cfg(target_arch = "aarch64")]
10686    {
10687        sdot_enabled()
10688    }
10689    #[cfg(target_arch = "x86_64")]
10690    {
10691        avx2_a8w8_enabled()
10692    }
10693    #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
10694    {
10695        false
10696    }
10697}
10698
10699/// int8·int8 dot dispatch: SDOT on ARM; AVX-512 VNNI (vpdpbusd) or AVX2
10700/// maddubs on x86. Callers are gated by `a8w8_enabled()`.
10701#[inline]
10702#[allow(unreachable_code)]
10703fn dot_i8_i8(w: &[u8], xq: &[i8]) -> i32 {
10704    #[cfg(target_arch = "aarch64")]
10705    unsafe {
10706        return dot_i8_sdot(w, xq);
10707    }
10708    #[cfg(target_arch = "x86_64")]
10709    unsafe {
10710        if avx512vnni_enabled() {
10711            return dot_i8_i8_vnni(w, xq);
10712        }
10713        return dot_i8_i8_avx2(w, xq);
10714    }
10715    w.iter()
10716        .zip(xq)
10717        .map(|(&a, &b)| (a as i8) as i32 * b as i32)
10718        .sum()
10719}
10720
10721/// AVX-512 VNNI available? (F+BW+VL+VNNI; `CMF_AVX512=0` falls back to
10722/// AVX2.) VL matters: short 32-byte groups (q4/vbit) ride the 256-bit
10723/// `vpdpbusd` encoding.
10724#[cfg(target_arch = "x86_64")]
10725fn avx512vnni_enabled() -> bool {
10726    use std::sync::OnceLock;
10727    static ON: OnceLock<bool> = OnceLock::new();
10728    *ON.get_or_init(|| {
10729        std::env::var("CMF_AVX512")
10730            .map(|v| v != "0")
10731            .unwrap_or(true)
10732            && std::arch::is_x86_feature_detected!("avx512f")
10733            && std::arch::is_x86_feature_detected!("avx512bw")
10734            && std::arch::is_x86_feature_detected!("avx512vl")
10735            && std::arch::is_x86_feature_detected!("avx512vnni")
10736    })
10737}
10738
10739/// Grouped-codec VNNI arms (the q4t/q4b/q1/q1t tile kernels): default
10740/// ON where AVX-512 VNNI exists (`CMF_VNNI_TILES=0` opt-out). Measured
10741/// on Ryzen 7950X (Zen4, 3 alternating process pairs, blocked GEMM
10742/// 4864×896 b=256): q4t 63→68 GF/s (+8%), q1 53→56 (+6%), q4b 72→75
10743/// (+4%) — consistent, no leg regressed. The tile kernels keep a
10744/// horizontal reduce per 32-weight group, so the `vpdpbusd` saving is
10745/// smaller than the long-dot q8 win (+13%), but it is real and free.
10746#[cfg(target_arch = "x86_64")]
10747fn vnni_tiles_enabled() -> bool {
10748    use std::sync::OnceLock;
10749    static ON: OnceLock<bool> = OnceLock::new();
10750    *ON.get_or_init(|| {
10751        std::env::var("CMF_VNNI_TILES")
10752            .map(|v| v != "0")
10753            .unwrap_or(true)
10754            && avx512vnni_enabled()
10755    })
10756}
10757
10758/// One 256-bit u8×i8 dot → i32 via `vpdpbusd` into a fresh accumulator
10759/// plus the same horizontal reduce the AVX2 kernels use. Products are
10760/// bounded (|w| ≤ 8 or ≤ 1), so maddubs never saturated — the i32 sum
10761/// is bit-identical to the maddubs+madd pair it replaces.
10762#[cfg(target_arch = "x86_64")]
10763#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
10764#[inline]
10765unsafe fn dpbusd_hsum(aw: core::arch::x86_64::__m256i, xs: core::arch::x86_64::__m256i) -> i32 {
10766    // SAFETY: pure register math.
10767    unsafe {
10768        use core::arch::x86_64::*;
10769        let d = _mm256_dpbusd_epi32(_mm256_setzero_si256(), aw, xs);
10770        let hi128 = _mm256_extracti128_si256::<1>(d);
10771        let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
10772        let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
10773        let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
10774        _mm_cvtsi128_si32(s32)
10775    }
10776}
10777
10778/// int8·int8 via AVX-512 VNNI: `vpdpbusd` fuses the maddubs+madd+add
10779/// triple into one u8×i8 dot-accumulate. AVX-512 has no vpsignb, so the
10780/// |w|·sign(x,w) trick becomes |w| × (x negated where w<0) via a mask
10781/// subtract — w==0 lanes contribute 0 through |w|=0 either way.
10782#[cfg(target_arch = "x86_64")]
10783#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
10784unsafe fn dot_i8_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
10785    // SAFETY: callers uphold slice-length contracts (see call sites).
10786    unsafe {
10787        use core::arch::x86_64::*;
10788        let n = w.len();
10789        let mut j = 0usize;
10790        let mut total: i32;
10791        // 4 independent accumulators: vpdpbusd is its own loop-carried
10792        // dependency (~5-cycle latency) — a single-acc loop runs
10793        // latency-bound and LOSES to the AVX2 maddubs kernel, measured
10794        // on Granite Rapids.
10795        {
10796            #[inline(always)]
10797            unsafe fn step(
10798                w: *const u8,
10799                x: *const i8,
10800                acc: core::arch::x86_64::__m512i,
10801            ) -> core::arch::x86_64::__m512i {
10802                unsafe {
10803                    use core::arch::x86_64::*;
10804                    let wv = _mm512_loadu_si512(w as *const _);
10805                    let xv = _mm512_loadu_si512(x as *const _);
10806                    let aw = _mm512_abs_epi8(wv);
10807                    let neg = _mm512_movepi8_mask(wv);
10808                    let sx = _mm512_mask_sub_epi8(xv, neg, _mm512_setzero_si512(), xv);
10809                    _mm512_dpbusd_epi32(acc, aw, sx)
10810                }
10811            }
10812            let (mut a0, mut a1, mut a2, mut a3) = (
10813                _mm512_setzero_si512(),
10814                _mm512_setzero_si512(),
10815                _mm512_setzero_si512(),
10816                _mm512_setzero_si512(),
10817            );
10818            while j + 256 <= n {
10819                a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
10820                a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), a1);
10821                a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), a2);
10822                a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), a3);
10823                j += 256;
10824            }
10825            while j + 64 <= n {
10826                a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
10827                j += 64;
10828            }
10829            let s01 = _mm512_add_epi32(a0, a1);
10830            let s23 = _mm512_add_epi32(a2, a3);
10831            total = _mm512_reduce_add_epi32(_mm512_add_epi32(s01, s23));
10832        }
10833        // 32-wide (q4/vbit groups are exactly 32 bytes).
10834        if j + 32 <= n {
10835            let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
10836            let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
10837            let d = _mm256_dpbusd_epi32(
10838                _mm256_setzero_si256(),
10839                _mm256_abs_epi8(wv),
10840                _mm256_sign_epi8(xv, wv),
10841            );
10842            let hi128 = _mm256_extracti128_si256::<1>(d);
10843            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
10844            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
10845            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
10846            total += _mm_cvtsi128_si32(s32);
10847            j += 32;
10848        }
10849        while j < n {
10850            total += (w[j] as i8) as i32 * xq[j] as i32;
10851            j += 1;
10852        }
10853        total
10854    }
10855}
10856
10857/// i8 row · f32 x via AVX2/FMA (x86 mirror of `dot_i8_f32_neon`).
10858#[cfg(target_arch = "x86_64")]
10859#[target_feature(enable = "avx2,fma")]
10860unsafe fn dot_i8_f32_avx2(w: &[u8], x: &[f32]) -> f32 {
10861    // SAFETY: callers uphold slice-length contracts (see call sites).
10862    unsafe {
10863        use core::arch::x86_64::*;
10864        let n = x.len();
10865        let wp = w.as_ptr();
10866        let xp = x.as_ptr();
10867        let (mut a0, mut a1) = (_mm256_setzero_ps(), _mm256_setzero_ps());
10868        let mut j = 0usize;
10869        while j + 16 <= n {
10870            let wb = _mm_loadu_si128(wp.add(j) as *const __m128i);
10871            let lo = _mm256_cvtepi8_epi32(wb);
10872            let hi = _mm256_cvtepi8_epi32(_mm_srli_si128::<8>(wb));
10873            a0 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(lo), _mm256_loadu_ps(xp.add(j)), a0);
10874            a1 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(hi), _mm256_loadu_ps(xp.add(j + 8)), a1);
10875            j += 16;
10876        }
10877        let acc = _mm256_add_ps(a0, a1);
10878        let hi128 = _mm256_extractf128_ps::<1>(acc);
10879        let s128 = _mm_add_ps(_mm256_castps256_ps128(acc), hi128);
10880        let s64 = _mm_add_ps(s128, _mm_movehl_ps(s128, s128));
10881        let s32 = _mm_add_ss(s64, _mm_shuffle_ps::<1>(s64, s64));
10882        let mut sum = _mm_cvtss_f32(s32);
10883        while j < n {
10884            sum += (*wp.add(j) as i8) as f32 * *xp.add(j);
10885            j += 1;
10886        }
10887        sum
10888    }
10889}
10890
10891/// int8(weight)·int8(activation) → i32 via AVX2 maddubs — the x86
10892/// analogue of the SDOT A8W8 path. `maddubs` takes u8×i8, so the
10893/// standard sign trick applies: |w| × sign(x, w) ≡ w × x per lane.
10894/// Pair saturation is safe: |w|≤128, |x|≤127 → 2·128·127 < 32767.
10895#[cfg(target_arch = "x86_64")]
10896#[target_feature(enable = "avx2")]
10897unsafe fn dot_i8_i8_avx2(w: &[u8], xq: &[i8]) -> i32 {
10898    // SAFETY: callers uphold slice-length contracts (see call sites).
10899    unsafe {
10900        use core::arch::x86_64::*;
10901        let n = w.len();
10902        let ones = _mm256_set1_epi16(1);
10903        let mut acc = _mm256_setzero_si256();
10904        let mut j = 0usize;
10905        while j + 32 <= n {
10906            let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
10907            let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
10908            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
10909            acc = _mm256_add_epi32(acc, _mm256_madd_epi16(p16, ones));
10910            j += 32;
10911        }
10912        let hi128 = _mm256_extracti128_si256::<1>(acc);
10913        let s128 = _mm_add_epi32(_mm256_castsi256_si128(acc), hi128);
10914        let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
10915        let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
10916        let mut s = _mm_cvtsi128_si32(s32);
10917        while j < n {
10918            s += (w[j] as i8) as i32 * xq[j] as i32;
10919            j += 1;
10920        }
10921        s
10922    }
10923}
10924
10925/// smmla 2×4: one instruction covers a 2-row × 2-activation × 8-deep
10926/// tile (32 MACs vs sdot's 16) — the weight pair loads once per 8-k
10927/// slice as a combined 2×8 register and meets two activation pairs.
10928#[cfg(target_arch = "aarch64")]
10929#[target_feature(enable = "neon,i8mm")]
10930unsafe fn dot_i8_smmla_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
10931    // SAFETY: callers uphold slice-length contracts.
10932    unsafe {
10933        use core::arch::aarch64::*;
10934        use core::arch::asm;
10935        let n = w0.len();
10936        let w0p = w0.as_ptr() as *const i8;
10937        let w1p = w1.as_ptr() as *const i8;
10938        // acc01 holds [c(r0,x0) c(r0,x1) c(r1,x0) c(r1,x1)]; acc23 the
10939        // same for x2/x3.
10940        let mut acc01 = vdupq_n_s32(0);
10941        let mut acc23 = vdupq_n_s32(0);
10942        let mut i = 0usize;
10943        while i + 8 <= n {
10944            let wa = vcombine_s8(vld1_s8(w0p.add(i)), vld1_s8(w1p.add(i)));
10945            let xb01 = vcombine_s8(
10946                vld1_s8(xs[0].as_ptr().add(i)),
10947                vld1_s8(xs[1].as_ptr().add(i)),
10948            );
10949            let xb23 = vcombine_s8(
10950                vld1_s8(xs[2].as_ptr().add(i)),
10951                vld1_s8(xs[3].as_ptr().add(i)),
10952            );
10953            asm!(
10954                "smmla {a01:v}.4s, {w:v}.16b, {x01:v}.16b",
10955                "smmla {a23:v}.4s, {w:v}.16b, {x23:v}.16b",
10956                a01 = inout(vreg) acc01, a23 = inout(vreg) acc23,
10957                w = in(vreg) wa, x01 = in(vreg) xb01, x23 = in(vreg) xb23,
10958                options(pure, nomem, nostack),
10959            );
10960            i += 8;
10961        }
10962        let mut out = [[0i32; 4]; 2];
10963        let a01: [i32; 4] = core::mem::transmute(acc01);
10964        let a23: [i32; 4] = core::mem::transmute(acc23);
10965        out[0][0] = a01[0];
10966        out[0][1] = a01[1];
10967        out[1][0] = a01[2];
10968        out[1][1] = a01[3];
10969        out[0][2] = a23[0];
10970        out[0][3] = a23[1];
10971        out[1][2] = a23[2];
10972        out[1][3] = a23[3];
10973        if i < n {
10974            for (k, x) in xs.iter().enumerate() {
10975                for j in i..n {
10976                    out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
10977                    out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
10978                }
10979            }
10980        }
10981        out
10982    }
10983}
10984
10985/// ARM twin of the x86 blocked prefill GEMM: two weight rows stay in
10986/// registers across four activation streams, eight sdot accumulators.
10987/// (The per-row form re-read each W row once per activation.)
10988#[cfg(target_arch = "aarch64")]
10989#[target_feature(enable = "neon,dotprod")]
10990unsafe fn dot_i8_sdot_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
10991    // SAFETY: callers uphold slice-length contracts.
10992    unsafe {
10993        use core::arch::aarch64::*;
10994        use core::arch::asm;
10995        let n = w0.len();
10996        let w0p = w0.as_ptr() as *const i8;
10997        let w1p = w1.as_ptr() as *const i8;
10998        let mut acc = [[vdupq_n_s32(0); 4]; 2];
10999        let mut i = 0usize;
11000        while i + 16 <= n {
11001            let wv0 = vld1q_s8(w0p.add(i));
11002            let wv1 = vld1q_s8(w1p.add(i));
11003            for (k, x) in xs.iter().enumerate() {
11004                let xv = vld1q_s8(x.as_ptr().add(i));
11005                let (mut a0, mut a1) = (acc[0][k], acc[1][k]);
11006                asm!(
11007                    "sdot {a0:v}.4s, {w0:v}.16b, {x:v}.16b",
11008                    "sdot {a1:v}.4s, {w1:v}.16b, {x:v}.16b",
11009                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
11010                    w0 = in(vreg) wv0, w1 = in(vreg) wv1, x = in(vreg) xv,
11011                    options(pure, nomem, nostack),
11012                );
11013                acc[0][k] = a0;
11014                acc[1][k] = a1;
11015            }
11016            i += 16;
11017        }
11018        let mut out = [[0i32; 4]; 2];
11019        for r in 0..2 {
11020            for k in 0..4 {
11021                out[r][k] = vaddvq_s32(acc[r][k]);
11022            }
11023        }
11024        if i < n {
11025            for (k, x) in xs.iter().enumerate() {
11026                for j in i..n {
11027                    out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
11028                    out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
11029                }
11030            }
11031        }
11032        out
11033    }
11034}
11035
11036/// Blocked 2 weight rows × 4 activations for the prefill GEMM
11037/// (roadmap P0: packed panels + multi-row accumulators). The two rows'
11038/// abs() live in registers across all four activation streams; the
11039/// sign-fixup is recomputed per pair (the price of the maddubs trick).
11040/// Returns raw i8·i8 dots; the caller applies scales and outliers.
11041#[cfg(target_arch = "x86_64")]
11042#[target_feature(enable = "avx2")]
11043unsafe fn dot_i8_i8_avx2_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
11044    // SAFETY: callers uphold slice-length contracts.
11045    unsafe {
11046        use core::arch::x86_64::*;
11047        let n = w0.len();
11048        let ones = _mm256_set1_epi16(1);
11049        let mut acc = [[_mm256_setzero_si256(); 4]; 2];
11050        let mut j = 0usize;
11051        while j + 32 <= n {
11052            let wv0 = _mm256_loadu_si256(w0.as_ptr().add(j) as *const __m256i);
11053            let wv1 = _mm256_loadu_si256(w1.as_ptr().add(j) as *const __m256i);
11054            let aw0 = _mm256_abs_epi8(wv0);
11055            let aw1 = _mm256_abs_epi8(wv1);
11056            for (k, x) in xs.iter().enumerate() {
11057                let xv = _mm256_loadu_si256(x.as_ptr().add(j) as *const __m256i);
11058                let p0 = _mm256_maddubs_epi16(aw0, _mm256_sign_epi8(xv, wv0));
11059                acc[0][k] = _mm256_add_epi32(acc[0][k], _mm256_madd_epi16(p0, ones));
11060                let p1 = _mm256_maddubs_epi16(aw1, _mm256_sign_epi8(xv, wv1));
11061                acc[1][k] = _mm256_add_epi32(acc[1][k], _mm256_madd_epi16(p1, ones));
11062            }
11063            j += 32;
11064        }
11065        let mut out = [[0i32; 4]; 2];
11066        for r in 0..2 {
11067            for k in 0..4 {
11068                let a = acc[r][k];
11069                let hi128 = _mm256_extracti128_si256::<1>(a);
11070                let s128 = _mm_add_epi32(_mm256_castsi256_si128(a), hi128);
11071                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
11072                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
11073                out[r][k] = _mm_cvtsi128_si32(s32);
11074            }
11075        }
11076        if j < n {
11077            for (k, x) in xs.iter().enumerate() {
11078                for i in j..n {
11079                    out[0][k] += (w0[i] as i8) as i32 * x[i] as i32;
11080                    out[1][k] += (w1[i] as i8) as i32 * x[i] as i32;
11081                }
11082            }
11083        }
11084        out
11085    }
11086}
11087
11088/// AVX2/VNNI q8 row dot with exact outlier correction (x86 mirror of
11089/// `row_dot_sdot` — same A8W8 contract). With AVX-512 VNNI the row goes
11090/// through the bias trick: Σ(w+128)·x via pure `vpdpbusd` (no per-lane
11091/// sign fixups), corrected by −128·Σx with Σx precomputed per split.
11092#[cfg(target_arch = "x86_64")]
11093#[inline]
11094fn row_dot_avx2(row: &[u8], act: &SplitAct) -> f32 {
11095    let dot = if avx512vnni_enabled() && row.len() >= 64 {
11096        (unsafe { dot_u8p128_i8_vnni(row, &act.xq) }) - 128 * act.xsum
11097    } else {
11098        unsafe { dot_i8_i8_avx2(row, &act.xq) }
11099    };
11100    let mut acc = dot as f32 * act.sx;
11101    for &(j, xv) in &act.outliers {
11102        acc += (row[j] as i8) as f32 * xv;
11103    }
11104    acc
11105}
11106
11107/// Σ (w[i]+128)·x[i] via pure `vpdpbusd` — the caller subtracts
11108/// 128·Σx. Four independent accumulators (dpbusd is ~5-cycle latency;
11109/// a single-acc loop runs latency-bound, measured on Granite Rapids).
11110#[cfg(target_arch = "x86_64")]
11111#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
11112unsafe fn dot_u8p128_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
11113    // SAFETY: callers uphold slice-length contracts (see call sites).
11114    unsafe {
11115        use core::arch::x86_64::*;
11116        let n = w.len();
11117        let flip = _mm512_set1_epi8(-128); // XOR 0x80: i8 w → u8 (w+128)
11118        #[inline(always)]
11119        unsafe fn step(
11120            w: *const u8,
11121            x: *const i8,
11122            flip: core::arch::x86_64::__m512i,
11123            acc: core::arch::x86_64::__m512i,
11124        ) -> core::arch::x86_64::__m512i {
11125            unsafe {
11126                use core::arch::x86_64::*;
11127                let wv = _mm512_xor_si512(_mm512_loadu_si512(w as *const _), flip);
11128                _mm512_dpbusd_epi32(acc, wv, _mm512_loadu_si512(x as *const _))
11129            }
11130        }
11131        let (mut a0, mut a1, mut a2, mut a3) = (
11132            _mm512_setzero_si512(),
11133            _mm512_setzero_si512(),
11134            _mm512_setzero_si512(),
11135            _mm512_setzero_si512(),
11136        );
11137        let mut j = 0usize;
11138        while j + 256 <= n {
11139            a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
11140            a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), flip, a1);
11141            a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), flip, a2);
11142            a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), flip, a3);
11143            j += 256;
11144        }
11145        while j + 64 <= n {
11146            a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
11147            j += 64;
11148        }
11149        let mut total = _mm512_reduce_add_epi32(_mm512_add_epi32(
11150            _mm512_add_epi32(a0, a1),
11151            _mm512_add_epi32(a2, a3),
11152        ));
11153        // Scalar tail: (w as i8) + 128 ≡ (w as u8) ^ 0x80.
11154        while j < n {
11155            total += ((w[j] ^ 0x80) as i32) * xq[j] as i32;
11156            j += 1;
11157        }
11158        total
11159    }
11160}
11161
11162/// One q4 row via AVX2: nibbles → centered i8 (unpacklo/hi restores the
11163/// writer's flat order, same as the NEON vzip pair), maddubs against
11164/// the pre-quantized activation group, × the group's f16 scale. Pair
11165/// saturation safe: |w|≤8, |x|≤127 → 2·8·127 ≪ 32767. Mirror of
11166/// `dot_q4_row_sdot`.
11167#[cfg(target_arch = "x86_64")]
11168#[target_feature(enable = "avx2")]
11169unsafe fn dot_q4_row_avx2(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
11170    // SAFETY: callers uphold slice-length contracts (16 packed bytes and
11171    // 2 scale bytes per group; xq.len() == gpr·GROUP_SIZE).
11172    unsafe {
11173        use core::arch::x86_64::*;
11174        let lomask = _mm_set1_epi8(0x0F);
11175        let eight = _mm256_set1_epi8(8);
11176        let ones = _mm256_set1_epi16(1);
11177        let mut acc = 0f32;
11178        for gi in 0..gpr {
11179            let g = g0 + gi;
11180            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
11181            let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
11182            let lo = _mm_and_si128(b, lomask);
11183            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
11184            let w = _mm256_sub_epi8(
11185                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
11186                eight,
11187            );
11188            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
11189            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
11190            let d = _mm256_madd_epi16(p16, ones);
11191            let hi128 = _mm256_extracti128_si256::<1>(d);
11192            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
11193            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
11194            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
11195            acc += _mm_cvtsi128_si32(s32) as f32 * s;
11196        }
11197        acc
11198    }
11199}
11200
11201/// Two-activation q4 row via AVX2: nibbles unpacked ONCE per group,
11202/// both activations dotted against the same centered i8 register.
11203#[cfg(target_arch = "x86_64")]
11204#[target_feature(enable = "avx2")]
11205unsafe fn dot_q4_row_avx2_2(
11206    packed: &[u8],
11207    scales: &[u8],
11208    g0: usize,
11209    gpr: usize,
11210    xq1: &[i8],
11211    xq2: &[i8],
11212) -> (f32, f32) {
11213    // SAFETY: callers uphold slice-length contracts (see dot_q4_row_avx2).
11214    unsafe {
11215        use core::arch::x86_64::*;
11216        let lomask = _mm_set1_epi8(0x0F);
11217        let eight = _mm256_set1_epi8(8);
11218        let ones = _mm256_set1_epi16(1);
11219        let (mut acc1, mut acc2) = (0f32, 0f32);
11220        #[inline(always)]
11221        unsafe fn hsum(d: core::arch::x86_64::__m256i) -> i32 {
11222            unsafe {
11223                use core::arch::x86_64::*;
11224                let hi128 = _mm256_extracti128_si256::<1>(d);
11225                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
11226                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
11227                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
11228                _mm_cvtsi128_si32(s32)
11229            }
11230        }
11231        for gi in 0..gpr {
11232            let g = g0 + gi;
11233            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
11234            let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
11235            let lo = _mm_and_si128(b, lomask);
11236            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
11237            let w = _mm256_sub_epi8(
11238                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
11239                eight,
11240            );
11241            let aw = _mm256_abs_epi8(w);
11242            let x1 = _mm256_loadu_si256(xq1.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
11243            let x2 = _mm256_loadu_si256(xq2.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
11244            let d1 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x1, w)), ones);
11245            let d2 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x2, w)), ones);
11246            acc1 += hsum(d1) as f32 * s;
11247            acc2 += hsum(d2) as f32 * s;
11248        }
11249        (acc1, acc2)
11250    }
11251}
11252
11253/// One q8 row range via AVX2 (x86 mirror of `q8_range_sdot`).
11254#[cfg(target_arch = "x86_64")]
11255fn q8_range_avx2(
11256    q: &[u8],
11257    row_scale: &[f32],
11258    act: &SplitAct,
11259    cols: usize,
11260    out_addr: SendMut,
11261    start: usize,
11262    end: usize,
11263) {
11264    for o in start..end {
11265        let v = row_dot_avx2(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
11266        // SAFETY: disjoint row ranges per worker.
11267        unsafe { *out_addr.at(o) = v };
11268    }
11269}
11270
11271/// Two-input q8 row range via AVX2 (x86 mirror of `q8_range2_sdot`).
11272#[cfg(target_arch = "x86_64")]
11273#[allow(clippy::too_many_arguments)]
11274fn q8_range2_avx2(
11275    q: &[u8],
11276    row_scale: &[f32],
11277    a1: &SplitAct,
11278    a2: &SplitAct,
11279    cols: usize,
11280    p1: SendMut,
11281    p2: SendMut,
11282    start: usize,
11283    end: usize,
11284) {
11285    for o in start..end {
11286        let row = &q[o * cols..(o + 1) * cols];
11287        // SAFETY: disjoint row ranges per worker.
11288        unsafe {
11289            *p1.at(o) = row_dot_avx2(row, a1) * row_scale[o];
11290            *p2.at(o) = row_dot_avx2(row, a2) * row_scale[o];
11291        }
11292    }
11293}
11294
11295// ───────────────────── A8W8 SDOT path (port of vmfcore, ×1.78 decode) ─────────────────────
11296
11297/// ARMv8.6 i8mm (smmla): 32 int8 MACs per instruction vs sdot's 16 —
11298/// yet MEASURED 2.4× SLOWER than the blocked sdot on Apple silicon
11299/// (108 vs 264 GF/s): the on-the-fly vcombine packing and the two-
11300/// accumulator dependency chain swamp the MAC advantage, and Apple's
11301/// four SIMD pipes already keep sdot fed. OPT-IN (CMF_I8MM=1) for
11302/// field trials on Cortex-A710/X-class parts with two pipes, where the
11303/// balance may differ; a pre-interleaved weight layout (repack infra)
11304/// is the known path if it ever earns its keep.
11305#[cfg(target_arch = "aarch64")]
11306fn i8mm_enabled() -> bool {
11307    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
11308    *ON.get_or_init(|| {
11309        std::env::var("CMF_I8MM").map(|v| v == "1").unwrap_or(false)
11310            && std::arch::is_aarch64_feature_detected!("i8mm")
11311    })
11312}
11313
11314/// SDOT enabled? Default ON when the CPU has ARMv8.2 dotprod;
11315/// `CMF_SDOT=0` disables (falls back to i8×f32 NEON).
11316/// (On non-ARM release builds only the test tolerance switch calls it.)
11317#[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
11318fn sdot_enabled() -> bool {
11319    if FLOAT_ACTIVATIONS.get() {
11320        return false;
11321    }
11322    use std::sync::OnceLock;
11323    static ON: OnceLock<bool> = OnceLock::new();
11324    *ON.get_or_init(|| {
11325        let want = std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true);
11326        if !want {
11327            return false;
11328        }
11329
11330        #[cfg(target_arch = "aarch64")]
11331        {
11332            if std::arch::is_aarch64_feature_detected!("dotprod") {
11333                return true;
11334            }
11335            #[cfg(target_os = "android")]
11336            {
11337                if let Ok(cpuinfo) = std::fs::read_to_string("/proc/cpuinfo") {
11338                    if cpuinfo.lines().any(|l| {
11339                        (l.starts_with("Features") || l.starts_with("features"))
11340                            && l.contains("asimddp")
11341                    }) {
11342                        return true;
11343                    }
11344                }
11345            }
11346            false
11347        }
11348        #[cfg(not(target_arch = "aarch64"))]
11349        {
11350            false
11351        }
11352    })
11353}
11354
11355/// Two-field activation split (≡ vmfcore `q8_split_prep`): outlier
11356/// channels (>8·rms) are computed exactly in f32; the bulk (outliers
11357/// zeroed → clean absmax) goes through int8 SDOT. Computed ONCE per
11358/// matvec, shared by all rows/workers.
11359struct SplitAct {
11360    xq: Vec<i8>,
11361    sx: f32,
11362    outliers: Vec<(usize, f32)>,
11363    /// Σ xq — the VNNI bias-trick correction (`(w+128)·x` sums need
11364    /// `−128·Σx`); one i32 per split, computed once per matvec.
11365    #[cfg_attr(not(target_arch = "x86_64"), allow(dead_code))]
11366    xsum: i32,
11367}
11368
11369thread_local! {
11370    /// Recycled xq buffers: split_act runs for every matvec (~200/token)
11371    /// and its hidden-size allocation was steady-state heap churn.
11372    static XQ_FREE: std::cell::RefCell<Vec<Vec<i8>>> =
11373        const { std::cell::RefCell::new(Vec::new()) };
11374}
11375
11376impl Drop for SplitAct {
11377    fn drop(&mut self) {
11378        let buf = std::mem::take(&mut self.xq);
11379        if buf.capacity() > 0 {
11380            XQ_FREE.with(|f| {
11381                let mut f = f.borrow_mut();
11382                if f.len() < 16 {
11383                    f.push(buf);
11384                }
11385            });
11386        }
11387    }
11388}
11389
11390thread_local! {
11391    /// One scratch row per WORKER, kept for the life of the thread.
11392    ///
11393    /// The kernels take a row of group scales per dispatch, and a fresh
11394    /// `vec![0f32; gpr]` inside the closure is one allocation per worker per
11395    /// dispatch — on the release checkpoint about six thousand a token, a
11396    /// quarter of everything the benchmark counts.
11397    static KROW: UnsafeCell<Vec<f32>> = const { UnsafeCell::new(Vec::new()) };
11398    /// Two scratch rows for a gate/up pair.  Q2TP and Q4TP fused SwiGLU
11399    /// decode each projection's ladder independently, but both ladders have
11400    /// the same lifetime; retaining them per worker removes two heap trips
11401    /// from every dispatched expert pair.
11402    static KROWS: UnsafeCell<[Vec<f32>; 2]> =
11403        const { UnsafeCell::new([Vec::new(), Vec::new()]) };
11404}
11405
11406/// Borrow `n` floats of the calling worker's scratch. Nothing inside a
11407/// kernel body borrows it again, which is what keeps the RefCell honest.
11408#[inline]
11409fn with_krow<R>(n: usize, f: impl FnOnce(&mut [f32]) -> R) -> R {
11410    KROW.with(|s| {
11411        // SAFETY: `KROW` is thread-local and this function does not recurse;
11412        // each worker has exclusive access to its own scratch row.
11413        let b = unsafe { &mut *s.get() };
11414        if b.len() < n {
11415            b.resize(n, 0.0);
11416        }
11417        f(&mut b[..n])
11418    })
11419}
11420
11421/// `t.round().clamp(-127.0, 127.0) as i8`, bit for bit, without the libm
11422/// call. On baseline x86-64 (no SSE4.1 `roundps`) `f32::round` is a
11423/// function call per element, and split_act runs it over every hidden
11424/// state before every matvec: measured 27 us a call on a 2048-wide
11425/// activation on an EPYC 7763 — 5.4 ms of a 55 ms decode token, all of
11426/// it on the caller's thread while thirty workers wait. Clamping first is
11427/// equivalent (round is monotonic and ±127 are integers), and after the
11428/// clamp `t - trunc(t)` is exact, so the half-away-from-zero decision is
11429/// the one `round` makes. NaN clamps to NaN and converts to 0, as before.
11430/// The loop vectorizes (cvttps2dq + compare/select).
11431#[inline(always)]
11432fn q8_round(t: f32) -> i8 {
11433    let t = t.clamp(-127.0, 127.0);
11434    let i = t as i32;
11435    let f = t - i as f32;
11436    let r = if f >= 0.5 {
11437        i + 1
11438    } else if f <= -0.5 {
11439        i - 1
11440    } else {
11441        i
11442    };
11443    r as i8
11444}
11445
11446#[inline]
11447fn with_krows<R>(n: usize, f: impl FnOnce(&mut [f32], &mut [f32]) -> R) -> R {
11448    KROWS.with(|s| {
11449        // SAFETY: KROWS is thread-local and no kernel body recursively calls
11450        // this helper; each participant owns its two vectors exclusively.
11451        let b = unsafe { &mut *s.get() };
11452        for row in b.iter_mut() {
11453            if row.len() < n {
11454                row.resize(n, 0.0);
11455            }
11456        }
11457        let (a, b) = b.split_at_mut(1);
11458        f(&mut a[0][..n], &mut b[0][..n])
11459    })
11460}
11461
11462#[inline]
11463fn silu_mul_limited(mut gate: f32, mut up: f32, limit: f32) -> f32 {
11464    if limit > 0.0 {
11465        up = up.clamp(-limit, limit);
11466        gate = gate.min(limit);
11467    }
11468    gate / (1.0 + (-gate).exp()) * up
11469}
11470
11471fn split_act(x: &[f32]) -> SplitAct {
11472    let _prof = crate::cpuprof::time(crate::cpuprof::Slot::SplitAct);
11473    let n = x.len();
11474    let rms = (x.iter().map(|&v| (v * v) as f64).sum::<f64>() / n.max(1) as f64).sqrt() as f32;
11475    let thr = 8.0 * rms;
11476    // One pass: collect outliers and the bulk absmax (outliers excluded —
11477    // identical to the old zero-then-fold over a copied buffer, minus the
11478    // full-vector copy).
11479    let mut outliers: Vec<(usize, f32)> = Vec::new();
11480    let mut amax = 0f32;
11481    for (j, &v) in x.iter().enumerate() {
11482        let a = v.abs();
11483        if a > thr {
11484            outliers.push((j, v));
11485        } else if a > amax {
11486            amax = a;
11487        }
11488    }
11489    let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
11490    let inv = 1.0 / sx;
11491    let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
11492    xq.clear();
11493    xq.reserve(n);
11494    if outliers.is_empty() {
11495        xq.extend(
11496            x.iter()
11497                .map(|&v| q8_round(v * inv)),
11498        );
11499    } else {
11500        // Outlier slots quantize to 0 (their exact term is added later).
11501        xq.extend(x.iter().map(|&v| {
11502            if v.abs() > thr {
11503                0
11504            } else {
11505                q8_round(v * inv)
11506            }
11507        }));
11508    }
11509    let xsum = xq.iter().map(|&v| v as i32).sum();
11510    SplitAct {
11511        xq,
11512        sx,
11513        outliers,
11514        xsum,
11515    }
11516}
11517
11518fn split_act_q8_2f(x: &[f32], col: &[f32]) -> SplitAct {
11519    let _prof = crate::cpuprof::time(crate::cpuprof::Slot::SplitAct);
11520    let n = x.len();
11521    let rms = (x
11522        .iter()
11523        .zip(col)
11524        .map(|(&a, &c)| {
11525            let v = a * c;
11526            (v * v) as f64
11527        })
11528        .sum::<f64>()
11529        / n.max(1) as f64)
11530        .sqrt() as f32;
11531    let thr = 8.0 * rms;
11532
11533    let mut outliers = Vec::new();
11534    let mut amax = 0f32;
11535    for (j, (&a, &c)) in x.iter().zip(col).enumerate() {
11536        let v = a * c;
11537        let s = v.abs();
11538        if s > thr {
11539            outliers.push((j, v));
11540        } else if s > amax {
11541            amax = s;
11542        }
11543    }
11544
11545    let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
11546    let inv = 1.0 / sx;
11547    let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
11548    xq.clear();
11549    xq.reserve(n);
11550    if outliers.is_empty() {
11551        xq.extend(
11552            x.iter()
11553                .zip(col)
11554                .map(|(&a, &c)| q8_round((a * c) * inv)),
11555        );
11556    } else {
11557        xq.extend(x.iter().zip(col).map(|(&a, &c)| {
11558            let v = a * c;
11559            if v.abs() > thr {
11560                0
11561            } else {
11562                q8_round(v * inv)
11563            }
11564        }));
11565    }
11566    let xsum = xq.iter().map(|&v| v as i32).sum();
11567    SplitAct {
11568        xq,
11569        sx,
11570        outliers,
11571        xsum,
11572    }
11573}
11574
11575/// int8(weight)·int8(activation) → i32 via `sdot` (inline asm — the
11576/// vdotq intrinsic is unstable; port of vmfcore `dot_i8_sdot`).
11577#[cfg(target_arch = "aarch64")]
11578#[target_feature(enable = "neon,dotprod")]
11579unsafe fn dot_i8_sdot(w: &[u8], xq: &[i8]) -> i32 {
11580    // SAFETY: callers uphold slice-length contracts (see call sites).
11581    unsafe {
11582        use core::arch::aarch64::*;
11583        use core::arch::asm;
11584        let wp = w.as_ptr() as *const i8;
11585        let n = w.len();
11586        let (mut a0, mut a1, mut a2, mut a3) = (
11587            vdupq_n_s32(0),
11588            vdupq_n_s32(0),
11589            vdupq_n_s32(0),
11590            vdupq_n_s32(0),
11591        );
11592        let mut i = 0;
11593        while i + 64 <= n {
11594            let (w0, x0) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
11595            let (w1, x1) = (vld1q_s8(wp.add(i + 16)), vld1q_s8(xq.as_ptr().add(i + 16)));
11596            let (w2, x2) = (vld1q_s8(wp.add(i + 32)), vld1q_s8(xq.as_ptr().add(i + 32)));
11597            let (w3, x3) = (vld1q_s8(wp.add(i + 48)), vld1q_s8(xq.as_ptr().add(i + 48)));
11598            asm!(
11599                "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
11600                "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
11601                "sdot {a2:v}.4s, {w2:v}.16b, {x2:v}.16b",
11602                "sdot {a3:v}.4s, {w3:v}.16b, {x3:v}.16b",
11603                a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
11604                w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
11605                w2 = in(vreg) w2, x2 = in(vreg) x2, w3 = in(vreg) w3, x3 = in(vreg) x3,
11606                options(pure, nomem, nostack),
11607            );
11608            i += 64;
11609        }
11610        while i + 16 <= n {
11611            let (wv, xv) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
11612            asm!("sdot {a:v}.4s, {w:v}.16b, {x:v}.16b",
11613                 a = inout(vreg) a0, w = in(vreg) wv, x = in(vreg) xv, options(pure, nomem, nostack));
11614            i += 16;
11615        }
11616        let mut s = vaddvq_s32(vaddq_s32(vaddq_s32(a0, a1), vaddq_s32(a2, a3)));
11617        while i < n {
11618            s += (*wp.add(i)) as i32 * xq[i] as i32;
11619            i += 1;
11620        }
11621        s
11622    }
11623}
11624
11625/// Row-blocked SDOT: 4 output rows per pass — the activation chunk is
11626/// loaded once and reused, 4 independent accumulators hide sdot latency
11627/// (port of vmfcore `dot_i8_sdot_4rows`).
11628#[cfg(target_arch = "aarch64")]
11629#[target_feature(enable = "neon,dotprod")]
11630unsafe fn dot_i8_sdot_4rows(w0: &[u8], w1: &[u8], w2: &[u8], w3: &[u8], xq: &[i8]) -> [i32; 4] {
11631    // SAFETY: callers uphold slice-length contracts (see call sites).
11632    unsafe {
11633        use core::arch::aarch64::*;
11634        use core::arch::asm;
11635        let n = xq.len();
11636        let px = xq.as_ptr();
11637        let (p0, p1, p2, p3) = (
11638            w0.as_ptr() as *const i8,
11639            w1.as_ptr() as *const i8,
11640            w2.as_ptr() as *const i8,
11641            w3.as_ptr() as *const i8,
11642        );
11643        let (mut a0, mut a1, mut a2, mut a3) = (
11644            vdupq_n_s32(0),
11645            vdupq_n_s32(0),
11646            vdupq_n_s32(0),
11647            vdupq_n_s32(0),
11648        );
11649        let mut i = 0;
11650        while i + 16 <= n {
11651            let x = vld1q_s8(px.add(i));
11652            let v0 = vld1q_s8(p0.add(i));
11653            let v1 = vld1q_s8(p1.add(i));
11654            let v2 = vld1q_s8(p2.add(i));
11655            let v3 = vld1q_s8(p3.add(i));
11656            asm!(
11657                "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
11658                "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
11659                "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
11660                "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
11661                a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
11662                v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
11663                options(pure, nomem, nostack),
11664            );
11665            i += 16;
11666        }
11667        let mut r = [
11668            vaddvq_s32(a0),
11669            vaddvq_s32(a1),
11670            vaddvq_s32(a2),
11671            vaddvq_s32(a3),
11672        ];
11673        while i < n {
11674            let xi = *px.add(i) as i32;
11675            r[0] += (*p0.add(i)) as i32 * xi;
11676            r[1] += (*p1.add(i)) as i32 * xi;
11677            r[2] += (*p2.add(i)) as i32 * xi;
11678            r[3] += (*p3.add(i)) as i32 * xi;
11679            i += 1;
11680        }
11681        r
11682    }
11683}
11684
11685/// 4 interleaved rows in one pass: the repacked group is [r0[c], r1[c],
11686/// r2[c], r3[c]] per 16-byte chunk, so each iteration reads ONE 64-byte
11687/// line plus the shared activation chunk — a single sequential weight
11688/// stream per worker. Per-row accumulation is the same one-accumulator
11689/// scheme as `dot_i8_sdot_4rows`; integer sums are exact, so outputs
11690/// are bit-identical to the mmap-layout kernel.
11691#[cfg(target_arch = "aarch64")]
11692#[target_feature(enable = "neon,dotprod")]
11693unsafe fn dot_i8_sdot_4rows_il(g: &[u8], xq: &[i8]) -> [i32; 4] {
11694    // SAFETY: callers uphold slice-length contracts (g.len() == 4·n,
11695    // n % 16 == 0 — guaranteed by the repack gate).
11696    unsafe {
11697        use core::arch::aarch64::*;
11698        use core::arch::asm;
11699        let n = xq.len();
11700        let px = xq.as_ptr();
11701        let pg = g.as_ptr() as *const i8;
11702        let (mut a0, mut a1, mut a2, mut a3) = (
11703            vdupq_n_s32(0),
11704            vdupq_n_s32(0),
11705            vdupq_n_s32(0),
11706            vdupq_n_s32(0),
11707        );
11708        let mut i = 0;
11709        while i + 16 <= n {
11710            let x = vld1q_s8(px.add(i));
11711            let base = pg.add(4 * i);
11712            let v0 = vld1q_s8(base);
11713            let v1 = vld1q_s8(base.add(16));
11714            let v2 = vld1q_s8(base.add(32));
11715            let v3 = vld1q_s8(base.add(48));
11716            asm!(
11717                "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
11718                "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
11719                "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
11720                "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
11721                a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
11722                v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
11723                options(pure, nomem, nostack),
11724            );
11725            i += 16;
11726        }
11727        [
11728            vaddvq_s32(a0),
11729            vaddvq_s32(a1),
11730            vaddvq_s32(a2),
11731            vaddvq_s32(a3),
11732        ]
11733    }
11734}
11735
11736/// One q8 row range via SDOT (4-row blocks + tail) — the body of
11737/// `qmatvec`'s hot loop, extracted so multi-matrix jobs can drive the
11738/// SAME kernel for several tensors under one pool dispatch. `rep` — the
11739/// load-time interleaved repack (empty = mmap layout only); rows outside
11740/// full 4-row groups always come from the mmap layout.
11741#[cfg(target_arch = "aarch64")]
11742fn q8_range_sdot(
11743    q: &[u8],
11744    rep: &[u8],
11745    row_scale: &[f32],
11746    act: &SplitAct,
11747    cols: usize,
11748    out_addr: SendMut,
11749    start: usize,
11750    end: usize,
11751) {
11752    let mut o = start;
11753    // Leading rows to the group boundary (repack path only): the pool
11754    // splits row ranges arbitrarily, groups are absolute.
11755    if !rep.is_empty() {
11756        while o < end && o % 4 != 0 {
11757            let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
11758            unsafe { *out_addr.at(o) = v };
11759            o += 1;
11760        }
11761    }
11762    while o + 4 <= end {
11763        let r = if rep.is_empty() {
11764            unsafe {
11765                dot_i8_sdot_4rows(
11766                    &q[o * cols..(o + 1) * cols],
11767                    &q[(o + 1) * cols..(o + 2) * cols],
11768                    &q[(o + 2) * cols..(o + 3) * cols],
11769                    &q[(o + 3) * cols..(o + 4) * cols],
11770                    &act.xq,
11771                )
11772            }
11773        } else {
11774            unsafe { dot_i8_sdot_4rows_il(&rep[o * cols..(o + 4) * cols], &act.xq) }
11775        };
11776        for k in 0..4 {
11777            let mut acc = r[k] as f32 * act.sx;
11778            for &(j, xv) in &act.outliers {
11779                acc += (q[(o + k) * cols + j] as i8) as f32 * xv;
11780            }
11781            // SAFETY: disjoint row ranges per worker.
11782            unsafe { *out_addr.at(o + k) = acc * row_scale[o + k] };
11783        }
11784        o += 4;
11785    }
11786    while o < end {
11787        let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
11788        unsafe { *out_addr.at(o) = v };
11789        o += 1;
11790    }
11791}
11792
11793/// Two-input q8 row range via SDOT — `qmatvec2`'s hot loop, extracted
11794/// for the fused pair multi-matrix job (`matvec2_many`).
11795#[cfg(target_arch = "aarch64")]
11796#[allow(clippy::too_many_arguments)]
11797fn q8_range2_sdot(
11798    q: &[u8],
11799    row_scale: &[f32],
11800    a1: &SplitAct,
11801    a2: &SplitAct,
11802    cols: usize,
11803    p1: SendMut,
11804    p2: SendMut,
11805    start: usize,
11806    end: usize,
11807) {
11808    for o in start..end {
11809        let row = &q[o * cols..(o + 1) * cols];
11810        // SAFETY: disjoint row ranges per worker.
11811        unsafe {
11812            *p1.at(o) = row_dot_sdot(row, a1) * row_scale[o];
11813            *p2.at(o) = row_dot_sdot(row, a2) * row_scale[o];
11814        }
11815    }
11816}
11817
11818/// Two-input q8 row range, f32 kernel (non-SDOT) — same extraction.
11819#[allow(clippy::too_many_arguments)]
11820fn q8_range2_f32(
11821    q: &[u8],
11822    row_scale: &[f32],
11823    x1: &[f32],
11824    x2: &[f32],
11825    cols: usize,
11826    p1: SendMut,
11827    p2: SendMut,
11828    start: usize,
11829    end: usize,
11830) {
11831    for o in start..end {
11832        let row = &q[o * cols..(o + 1) * cols];
11833        // SAFETY: disjoint row ranges per worker.
11834        unsafe {
11835            *p1.at(o) = dot_i8_f32(row, x1) * row_scale[o];
11836            *p2.at(o) = dot_i8_f32(row, x2) * row_scale[o];
11837        }
11838    }
11839}
11840
11841/// Scalar/NEON-f32 q8 row range (non-SDOT platforms) — same extraction.
11842fn q8_range_f32(
11843    q: &[u8],
11844    row_scale: &[f32],
11845    xs: &[f32],
11846    cols: usize,
11847    out_addr: SendMut,
11848    start: usize,
11849    end: usize,
11850) {
11851    for o in start..end {
11852        let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
11853        // SAFETY: disjoint row ranges per worker.
11854        unsafe { *out_addr.at(o) = v };
11855    }
11856}
11857
11858/// One q8 row against a split activation, portable: the per-arch fast
11859/// dots where they exist, the exact scalar loop elsewhere. The scalar
11860/// arm is also the test oracle for both fast arms.
11861#[inline]
11862fn q8_row_dot(row: &[u8], act: &SplitAct) -> f32 {
11863    #[cfg(target_arch = "aarch64")]
11864    return row_dot_sdot(row, act);
11865    #[cfg(target_arch = "x86_64")]
11866    return row_dot_avx2(row, act);
11867    #[allow(unreachable_code)]
11868    q8_row_dot_scalar(row, act)
11869}
11870
11871#[allow(dead_code)]
11872fn q8_row_dot_scalar(row: &[u8], act: &SplitAct) -> f32 {
11873    let mut acc = 0i32;
11874    for (k, &b) in row.iter().enumerate() {
11875        acc += (b as i8) as i32 * act.xq[k] as i32;
11876    }
11877    let mut acc = acc as f32 * act.sx;
11878    for &(j, xv) in &act.outliers {
11879        acc += (row[j] as i8) as f32 * xv;
11880    }
11881    acc
11882}
11883
11884/// SDOT row dot with exact outlier correction:
11885/// `dot = sdot(w, xq)·sx + Σ_outl w[j]·x[j]` (then × row_scale by caller).
11886#[cfg(target_arch = "aarch64")]
11887#[inline]
11888fn row_dot_sdot(row: &[u8], act: &SplitAct) -> f32 {
11889    let mut acc = unsafe { dot_i8_sdot(row, &act.xq) } as f32 * act.sx;
11890    for &(j, xv) in &act.outliers {
11891        acc += (row[j] as i8) as f32 * xv;
11892    }
11893    acc
11894}
11895
11896/// One q4 row via SDOT: each 32-group's nibbles unpack to centered i8
11897/// (nib−8 ∈ [−8,7]), int8×int8 `sdot` against the pre-quantized
11898/// activation group, × the group's f16 scale. Returns Σ_g dot_g·s_g;
11899/// the caller multiplies by the activation scale and adds the exact
11900/// outlier terms (port of vmfcore `dot_q4_block_sdot`, +23% measured).
11901/// Nibble order matches the writer: element 2k = low nibble, 2k+1 = high
11902/// → zip(lo,hi) restores flat order.
11903#[cfg(target_arch = "aarch64")]
11904#[target_feature(enable = "neon,dotprod")]
11905unsafe fn dot_q4_row_sdot(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
11906    // SAFETY: callers uphold slice-length contracts (16 packed bytes and
11907    // 2 scale bytes per group; xq.len() == gpr·GROUP_SIZE).
11908    unsafe {
11909        use core::arch::aarch64::*;
11910        use core::arch::asm;
11911        let lomask = vdupq_n_u8(0x0F);
11912        let eight = vdupq_n_s8(8);
11913        let mut acc = 0f32;
11914        for gi in 0..gpr {
11915            let g = g0 + gi;
11916            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
11917            let b = vld1q_u8(packed.as_ptr().add(g * 16));
11918            let lo = vandq_u8(b, lomask);
11919            let hi = vshrq_n_u8::<4>(b);
11920            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
11921            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
11922            let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
11923            let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
11924            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
11925            asm!(
11926                "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
11927                "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
11928                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
11929                e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
11930                options(pure, nomem, nostack),
11931            );
11932            acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
11933        }
11934        acc
11935    }
11936}
11937
11938/// Two-activation q4 row via SDOT: the nibble unpack (the expensive
11939/// part) happens ONCE per group; both pre-quantized activations are
11940/// dotted against the same centered i8 registers. Per-lane math matches
11941/// `dot_q4_row_sdot` exactly.
11942#[cfg(target_arch = "aarch64")]
11943#[target_feature(enable = "neon,dotprod")]
11944unsafe fn dot_q4_row_sdot2(
11945    packed: &[u8],
11946    scales: &[u8],
11947    g0: usize,
11948    gpr: usize,
11949    xq1: &[i8],
11950    xq2: &[i8],
11951) -> (f32, f32) {
11952    // SAFETY: callers uphold slice-length contracts (16 packed bytes and
11953    // 2 scale bytes per group; xq*.len() == gpr·GROUP_SIZE).
11954    unsafe {
11955        use core::arch::aarch64::*;
11956        use core::arch::asm;
11957        let lomask = vdupq_n_u8(0x0F);
11958        let eight = vdupq_n_s8(8);
11959        let (mut acc1, mut acc2) = (0f32, 0f32);
11960        for gi in 0..gpr {
11961            let g = g0 + gi;
11962            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
11963            let b = vld1q_u8(packed.as_ptr().add(g * 16));
11964            let lo = vandq_u8(b, lomask);
11965            let hi = vshrq_n_u8::<4>(b);
11966            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
11967            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
11968            let x10 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE));
11969            let x11 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE + 16));
11970            let x20 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE));
11971            let x21 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE + 16));
11972            let (mut a0, mut a1, mut b0, mut b1) = (
11973                vdupq_n_s32(0),
11974                vdupq_n_s32(0),
11975                vdupq_n_s32(0),
11976                vdupq_n_s32(0),
11977            );
11978            asm!(
11979                "sdot {a0:v}.4s, {e0:v}.16b, {x10:v}.16b",
11980                "sdot {a1:v}.4s, {e1:v}.16b, {x11:v}.16b",
11981                "sdot {b0:v}.4s, {e0:v}.16b, {x20:v}.16b",
11982                "sdot {b1:v}.4s, {e1:v}.16b, {x21:v}.16b",
11983                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
11984                b0 = inout(vreg) b0, b1 = inout(vreg) b1,
11985                e0 = in(vreg) e0, e1 = in(vreg) e1,
11986                x10 = in(vreg) x10, x11 = in(vreg) x11,
11987                x20 = in(vreg) x20, x21 = in(vreg) x21,
11988                options(pure, nomem, nostack),
11989            );
11990            acc1 += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
11991            acc2 += vaddvq_s32(vaddq_s32(b0, b1)) as f32 * s;
11992        }
11993        (acc1, acc2)
11994    }
11995}
11996
11997// ───────────────────── fused int8 kernels ─────────────────────
11998
11999/// `acc += w · row` where the row is centered i8 — NEON widen+fma on
12000/// aarch64, scalar elsewhere. The KV-cache q8 value path rides on this.
12001#[inline]
12002pub(crate) fn axpy_i8_f32(acc: &mut [f32], row: &[i8], w: f32) {
12003    #[cfg(target_arch = "aarch64")]
12004    unsafe {
12005        return axpy_i8_f32_neon(acc, row, w);
12006    }
12007    #[cfg(target_arch = "x86_64")]
12008    if avx2_enabled() {
12009        return unsafe { axpy_i8_f32_avx2(acc, row, w) };
12010    }
12011    #[allow(unreachable_code)]
12012    {
12013        for (a, &b) in acc.iter_mut().zip(row) {
12014            *a += w * b as f32;
12015        }
12016    }
12017}
12018
12019/// i8→f32 axpy via AVX2/FMA (x86 mirror of `axpy_i8_f32_neon`).
12020#[cfg(target_arch = "x86_64")]
12021#[target_feature(enable = "avx2,fma")]
12022unsafe fn axpy_i8_f32_avx2(acc: &mut [f32], row: &[i8], w: f32) {
12023    // SAFETY: callers uphold slice-length contracts (see call sites).
12024    unsafe {
12025        use core::arch::x86_64::*;
12026        let n = acc.len().min(row.len());
12027        let ap = acc.as_mut_ptr();
12028        let rp = row.as_ptr();
12029        let wv = _mm256_set1_ps(w);
12030        let mut j = 0usize;
12031        while j + 16 <= n {
12032            let rb = _mm_loadu_si128(rp.add(j) as *const __m128i);
12033            let lo = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(rb));
12034            let hi = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(rb)));
12035            let v0 = _mm256_fmadd_ps(wv, lo, _mm256_loadu_ps(ap.add(j)));
12036            let v1 = _mm256_fmadd_ps(wv, hi, _mm256_loadu_ps(ap.add(j + 8)));
12037            _mm256_storeu_ps(ap.add(j), v0);
12038            _mm256_storeu_ps(ap.add(j + 8), v1);
12039            j += 16;
12040        }
12041        while j < n {
12042            *ap.add(j) += w * (*rp.add(j)) as f32;
12043            j += 1;
12044        }
12045    }
12046}
12047
12048#[cfg(target_arch = "aarch64")]
12049#[target_feature(enable = "neon")]
12050unsafe fn axpy_i8_f32_neon(acc: &mut [f32], row: &[i8], w: f32) {
12051    // SAFETY: callers uphold slice-length contracts (see call sites).
12052    unsafe {
12053        use core::arch::aarch64::*;
12054        let n = acc.len().min(row.len());
12055        let ap = acc.as_mut_ptr();
12056        let rp = row.as_ptr();
12057        let wv = vdupq_n_f32(w);
12058        let mut j = 0usize;
12059        while j + 16 <= n {
12060            let rb = vld1q_s8(rp.add(j));
12061            let lo = vmovl_s8(vget_low_s8(rb));
12062            let hi = vmovl_s8(vget_high_s8(rb));
12063            for (off, half) in [(0, lo), (8, hi)] {
12064                let f0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(half)));
12065                let f1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(half)));
12066                let o = j + off;
12067                vst1q_f32(ap.add(o), vfmaq_f32(vld1q_f32(ap.add(o)), wv, f0));
12068                vst1q_f32(ap.add(o + 4), vfmaq_f32(vld1q_f32(ap.add(o + 4)), wv, f1));
12069            }
12070            j += 16;
12071        }
12072        while j < n {
12073            *ap.add(j) += w * (*rp.add(j)) as f32;
12074            j += 1;
12075        }
12076    }
12077}
12078
12079/// i8 row · f32 x. NEON on aarch64 (ported from vmfcore `dot_i8_f32_neon`,
12080/// ≈9× scalar), scalar elsewhere.
12081#[inline]
12082pub(crate) fn dot_i8_f32(w: &[u8], x: &[f32]) -> f32 {
12083    #[cfg(target_arch = "aarch64")]
12084    unsafe {
12085        return dot_i8_f32_neon(w, x);
12086    }
12087    #[cfg(target_arch = "x86_64")]
12088    if avx2_enabled() {
12089        return unsafe { dot_i8_f32_avx2(w, x) };
12090    }
12091    #[allow(unreachable_code)]
12092    {
12093        let mut sum = 0.0f32;
12094        for (j, &b) in w.iter().enumerate() {
12095            sum += (b as i8) as f32 * x[j];
12096        }
12097        sum
12098    }
12099}
12100
12101/// i8 row · (x ⊙ col_field) — the q8_2f row dot with the θ col-field
12102/// folded into the product (no prescaled copy of x). NEON on aarch64,
12103/// scalar elsewhere. Used by the active-neuron path `row_dot`.
12104#[inline]
12105fn dot_i8_col_f32(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
12106    #[cfg(target_arch = "aarch64")]
12107    unsafe {
12108        return dot_i8_col_f32_neon(w, x, col);
12109    }
12110    #[allow(unreachable_code)]
12111    {
12112        let mut sum = 0.0f32;
12113        for (j, &b) in w.iter().enumerate() {
12114            sum += (b as i8) as f32 * x[j] * col[j];
12115        }
12116        sum
12117    }
12118}
12119
12120#[cfg(target_arch = "aarch64")]
12121#[target_feature(enable = "neon")]
12122unsafe fn dot_i8_col_f32_neon(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
12123    // SAFETY: callers uphold slice-length contracts (see call sites).
12124    unsafe {
12125        use core::arch::aarch64::*;
12126        let n = x.len();
12127        let wp = w.as_ptr() as *const i8;
12128        let xp = x.as_ptr();
12129        let cp = col.as_ptr();
12130        let (mut a0, mut a1, mut a2, mut a3) = (
12131            vdupq_n_f32(0.0),
12132            vdupq_n_f32(0.0),
12133            vdupq_n_f32(0.0),
12134            vdupq_n_f32(0.0),
12135        );
12136        let mut j = 0usize;
12137        while j + 16 <= n {
12138            let wb = vld1q_s8(wp.add(j));
12139            let lo = vmovl_s8(vget_low_s8(wb));
12140            let hi = vmovl_s8(vget_high_s8(wb));
12141            let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
12142            let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
12143            let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
12144            let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
12145            a0 = vfmaq_f32(
12146                a0,
12147                w0,
12148                vmulq_f32(vld1q_f32(xp.add(j)), vld1q_f32(cp.add(j))),
12149            );
12150            a1 = vfmaq_f32(
12151                a1,
12152                w1,
12153                vmulq_f32(vld1q_f32(xp.add(j + 4)), vld1q_f32(cp.add(j + 4))),
12154            );
12155            a2 = vfmaq_f32(
12156                a2,
12157                w2,
12158                vmulq_f32(vld1q_f32(xp.add(j + 8)), vld1q_f32(cp.add(j + 8))),
12159            );
12160            a3 = vfmaq_f32(
12161                a3,
12162                w3,
12163                vmulq_f32(vld1q_f32(xp.add(j + 12)), vld1q_f32(cp.add(j + 12))),
12164            );
12165            j += 16;
12166        }
12167        let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
12168        while j < n {
12169            sum += (*wp.add(j)) as f32 * *xp.add(j) * *cp.add(j);
12170            j += 1;
12171        }
12172        sum
12173    }
12174}
12175
12176#[cfg(target_arch = "aarch64")]
12177#[target_feature(enable = "neon")]
12178unsafe fn dot_i8_f32_neon(w: &[u8], x: &[f32]) -> f32 {
12179    // SAFETY: callers uphold slice-length contracts (see call sites).
12180    unsafe {
12181        use core::arch::aarch64::*;
12182        let n = x.len();
12183        let wp = w.as_ptr() as *const i8;
12184        let xp = x.as_ptr();
12185        let (mut a0, mut a1, mut a2, mut a3) = (
12186            vdupq_n_f32(0.0),
12187            vdupq_n_f32(0.0),
12188            vdupq_n_f32(0.0),
12189            vdupq_n_f32(0.0),
12190        );
12191        let mut j = 0usize;
12192        while j + 16 <= n {
12193            let wb = vld1q_s8(wp.add(j));
12194            let lo = vmovl_s8(vget_low_s8(wb));
12195            let hi = vmovl_s8(vget_high_s8(wb));
12196            let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
12197            let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
12198            let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
12199            let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
12200            a0 = vfmaq_f32(a0, w0, vld1q_f32(xp.add(j)));
12201            a1 = vfmaq_f32(a1, w1, vld1q_f32(xp.add(j + 4)));
12202            a2 = vfmaq_f32(a2, w2, vld1q_f32(xp.add(j + 8)));
12203            a3 = vfmaq_f32(a3, w3, vld1q_f32(xp.add(j + 12)));
12204            j += 16;
12205        }
12206        let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
12207        while j < n {
12208            sum += (*wp.add(j)) as f32 * *xp.add(j);
12209            j += 1;
12210        }
12211        sum
12212    }
12213}
12214
12215#[allow(clippy::too_many_arguments)]
12216fn qmatvec(
12217    q: &[u8],
12218    rep: &[u8],
12219    row_scale: &[f32],
12220    x: &[f32],
12221    col_field: &[f32],
12222    dtype: TensorDtype,
12223    rows: usize,
12224    cols: usize,
12225    out: &mut [f32],
12226    pool: Option<&Pool>,
12227) {
12228    debug_assert_eq!(out.len(), rows);
12229    #[cfg(not(target_arch = "aarch64"))]
12230    let _ = rep;
12231
12232    #[cfg(target_arch = "aarch64")]
12233    if sdot_enabled() {
12234        let act = if dtype == TensorDtype::Q8_2f {
12235            split_act_q8_2f(x, col_field)
12236        } else {
12237            split_act(x)
12238        };
12239        let out_addr = SendMut(out.as_mut_ptr());
12240        let run_range = |start: usize, end: usize| {
12241            q8_range_sdot(q, rep, row_scale, &act, cols, out_addr, start, end)
12242        };
12243        match pool {
12244            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
12245            _ => run_range(0, rows),
12246        }
12247        return;
12248    }
12249    // x86 A8W8 via AVX2 maddubs — same quantized-activation contract as
12250    // the SDOT path (CMF_AVX2=0 keeps the exact i8×f32 loop).
12251    #[cfg(target_arch = "x86_64")]
12252    if avx2_a8w8_enabled() {
12253        let act = if dtype == TensorDtype::Q8_2f {
12254            split_act_q8_2f(x, col_field)
12255        } else {
12256            split_act(x)
12257        };
12258        let out_addr = SendMut(out.as_mut_ptr());
12259        let run_range = |start: usize, end: usize| {
12260            q8_range_avx2(q, row_scale, &act, cols, out_addr, start, end)
12261        };
12262        match pool {
12263            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
12264            _ => run_range(0, rows),
12265        }
12266        return;
12267    }
12268
12269    prescale_with(x, col_field, dtype, 1, |xs| {
12270        let out_addr = SendMut(out.as_mut_ptr());
12271        let run_range = move |start: usize, end: usize| {
12272            for o in start..end {
12273                let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
12274                // SAFETY: disjoint row ranges per worker.
12275                unsafe { *out_addr.at(o) = v };
12276            }
12277        };
12278        match pool {
12279            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
12280            _ => run_range(0, rows),
12281        }
12282    });
12283}
12284
12285#[allow(clippy::too_many_arguments)]
12286fn qmatvec2(
12287    q: &[u8],
12288    row_scale: &[f32],
12289    x1: &[f32],
12290    x2: &[f32],
12291    col_field: &[f32],
12292    dtype: TensorDtype,
12293    rows: usize,
12294    cols: usize,
12295    o1: &mut [f32],
12296    o2: &mut [f32],
12297    pool: Option<&Pool>,
12298) {
12299    #[cfg(target_arch = "aarch64")]
12300    if sdot_enabled() {
12301        let a1s = if dtype == TensorDtype::Q8_2f {
12302            split_act_q8_2f(x1, col_field)
12303        } else {
12304            split_act(x1)
12305        };
12306        let a2s = if dtype == TensorDtype::Q8_2f {
12307            split_act_q8_2f(x2, col_field)
12308        } else {
12309            split_act(x2)
12310        };
12311        let p1 = SendMut(o1.as_mut_ptr());
12312        let p2 = SendMut(o2.as_mut_ptr());
12313        let run_range = |start: usize, end: usize| {
12314            q8_range2_sdot(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
12315        };
12316        match pool {
12317            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
12318            _ => run_range(0, rows),
12319        }
12320        return;
12321    }
12322    #[cfg(target_arch = "x86_64")]
12323    if avx2_a8w8_enabled() {
12324        let a1s = if dtype == TensorDtype::Q8_2f {
12325            split_act_q8_2f(x1, col_field)
12326        } else {
12327            split_act(x1)
12328        };
12329        let a2s = if dtype == TensorDtype::Q8_2f {
12330            split_act_q8_2f(x2, col_field)
12331        } else {
12332            split_act(x2)
12333        };
12334        let p1 = SendMut(o1.as_mut_ptr());
12335        let p2 = SendMut(o2.as_mut_ptr());
12336        let run_range = |start: usize, end: usize| {
12337            q8_range2_avx2(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
12338        };
12339        match pool {
12340            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
12341            _ => run_range(0, rows),
12342        }
12343        return;
12344    }
12345
12346    prescale_with(x1, col_field, dtype, 1, |x1s| {
12347        prescale_with(x2, col_field, dtype, 2, |x2s| {
12348            let p1 = SendMut(o1.as_mut_ptr());
12349            let p2 = SendMut(o2.as_mut_ptr());
12350            let run_range = move |start: usize, end: usize| {
12351                for o in start..end {
12352                    let row = &q[o * cols..(o + 1) * cols];
12353                    let s1 = dot_i8_f32(row, x1s) * row_scale[o];
12354                    let s2 = dot_i8_f32(row, x2s) * row_scale[o];
12355                    // SAFETY: disjoint row ranges per worker.
12356                    unsafe {
12357                        *p1.at(o) = s1;
12358                        *p2.at(o) = s2;
12359                    }
12360                }
12361            };
12362            match pool {
12363                Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
12364                _ => run_range(0, rows),
12365            }
12366        });
12367    });
12368}
12369
12370#[derive(Clone, Copy)]
12371struct SendMut(*mut f32);
12372unsafe impl Send for SendMut {}
12373unsafe impl Sync for SendMut {}
12374
12375impl SendMut {
12376    #[inline]
12377    fn at(self, i: usize) -> *mut f32 {
12378        unsafe { self.0.add(i) }
12379    }
12380}
12381
12382#[cfg(test)]
12383mod tests {
12384    /// `q8_round` must be `round().clamp(±127) as i8` bit for bit: every
12385    /// half-integer, their neighbours one ulp either side, the clamp
12386    /// boundary, huge values, infinities and NaN, plus a dense sweep.
12387    #[test]
12388    fn q8_round_is_round_clamp() {
12389        let reference = |t: f32| t.round().clamp(-127.0, 127.0) as i8;
12390        let mut probe = vec![
12391            0.0f32,
12392            -0.0,
12393            f32::NAN,
12394            f32::INFINITY,
12395            f32::NEG_INFINITY,
12396            f32::MAX,
12397            f32::MIN,
12398            1e30,
12399            -1e30,
12400            f32::MIN_POSITIVE,
12401            -f32::MIN_POSITIVE,
12402        ];
12403        for k in -300i32..=300 {
12404            let h = k as f32 * 0.5;
12405            let up = f32::from_bits(h.to_bits() + 1);
12406            let down = f32::from_bits(h.to_bits().wrapping_sub(1));
12407            for t in [h, up, down] {
12408                probe.push(t);
12409                probe.push(-t);
12410            }
12411        }
12412        let mut t = -140.0f32;
12413        while t < 140.0 {
12414            probe.push(t);
12415            t += 0.000_731;
12416        }
12417        for t in probe {
12418            assert_eq!(q8_round(t), reference(t), "t = {t:e} ({:#x})", t.to_bits());
12419        }
12420    }
12421
12422    use super::*;
12423
12424    /// The f32 `matmat` walks the batch four positions a row at a time; every
12425    /// output must still be the scalar-order sum `matvec` gives, bit for bit
12426    /// (a 7-row batch covers the four-wide block and the remainder).
12427    #[test]
12428    fn f32_matmat_equals_matvec_bitwise() {
12429        let (rows, cols, b) = (5usize, 37usize, 7usize);
12430        let w: Vec<f32> = (0..rows * cols).map(|i| ((i * 7919 % 113) as f32 - 56.0) / 37.0).collect();
12431        let xs: Vec<f32> = (0..b * cols).map(|i| ((i * 104_729 % 97) as f32 - 48.0) / 29.0).collect();
12432        let t = QTensor::from_f32(w, rows, cols);
12433        let mut out = vec![0f32; b * rows];
12434        t.matmat(&xs, b, &mut out, None);
12435        for bi in 0..b {
12436            let mut one = vec![0f32; rows];
12437            t.matvec(&xs[bi * cols..(bi + 1) * cols], &mut one, None);
12438            for o in 0..rows {
12439                assert_eq!(out[bi * rows + o].to_bits(), one[o].to_bits(), "bi {bi} row {o}");
12440            }
12441        }
12442    }
12443
12444    #[test]
12445    fn q2tp_i8_dot_matches_exact_on_grid() {
12446        // On-grid activations (±1 → sx=1/127, xq=±127 dequantizes
12447        // exactly, no outliers) must make the integer path agree with
12448        // the exact scalar walk to f32 rounding.
12449        let (rows, cols) = (5, 64);
12450        let gpr = cols / GROUP_SIZE;
12451        // Synthetic codes plane + a flat ladder: scales_into is not under
12452        // test here, so drive dot_q2tp_row_i8 / q2tp_row_exact directly
12453        // with hand-made scales.
12454        let chunks: Vec<u8> = (0..rows * gpr * Q2TP_CHUNK)
12455            .map(|i| (i as u32).wrapping_mul(2654435761) as u8)
12456            .collect();
12457        let scales: Vec<f32> = (0..gpr).map(|g| 0.5 + g as f32 * 0.25).collect();
12458        let x: Vec<f32> = (0..cols)
12459            .map(|i| if i % 3 == 0 { -1.0 } else { 1.0 })
12460            .collect();
12461        let act = split_act(&x);
12462        assert!(
12463            act.outliers.is_empty(),
12464            "on-grid input must have no outliers"
12465        );
12466        let gsum = q1_group_sums(&act.xq, gpr);
12467        for r in 0..rows {
12468            let exact = q2tp_row_exact(&chunks, r, gpr, &x, &scales);
12469            let fast = dot_q2tp_row_i8(&chunks, r, gpr, &act.xq, &gsum, &scales) * act.sx;
12470            assert!(
12471                (exact - fast).abs() <= exact.abs() * 1e-5 + 1e-5,
12472                "row {r}: exact {exact} vs i8 {fast}"
12473            );
12474        }
12475    }
12476
12477    #[cfg(target_arch = "x86_64")]
12478    #[test]
12479    fn q2tp_avx2_dot_matches_scalar_for_random_patterns() {
12480        // Compare the release AVX2 integer dot against the scalar oracle over
12481        // arbitrary packed bytes/activation signs.  This guards the exact
12482        // table-load path used after rejecting a faster-looking decoder whose
12483        // full-checkpoint greedy output drifted.
12484        if !std::arch::is_x86_feature_detected!("avx2") {
12485            return;
12486        }
12487        let mut seed = 0x9e3779b9u32;
12488        let mut next = || {
12489            seed = seed.wrapping_mul(1664525).wrapping_add(1013904223);
12490            seed
12491        };
12492        for _ in 0..20_000 {
12493            let mut ch = [0u8; Q2TP_CHUNK];
12494            let mut x = [0i8; GROUP_SIZE];
12495            for b in &mut ch {
12496                *b = next() as u8;
12497            }
12498            for v in &mut x {
12499                *v = (next() >> 24) as i8;
12500            }
12501            let mut reference = 0i32;
12502            for (k, &b) in ch.iter().enumerate() {
12503                reference += (b & 3) as i32 * x[k * 4] as i32;
12504                reference += ((b >> 2) & 3) as i32 * x[k * 4 + 1] as i32;
12505                reference += ((b >> 4) & 3) as i32 * x[k * 4 + 2] as i32;
12506                reference += ((b >> 6) & 3) as i32 * x[k * 4 + 3] as i32;
12507            }
12508            // SAFETY: guarded by the runtime AVX2 feature check and fixed
12509            // 8-byte/32-byte slice lengths above.
12510            let got = unsafe { q2tp_code_dot_avx2(&ch, &x) };
12511            assert_eq!(got, reference, "packed q2 lane mismatch");
12512        }
12513    }
12514
12515    #[test]
12516    fn q2tp_affine_fuses_half_scale_correction_without_changing_raw_decode() {
12517        let (rows, cols) = (1usize, GROUP_SIZE);
12518        let mut bytes = vec![0u8; Q2TP_CHUNK + 4 + 1];
12519        // Repeating symbols 0,1,2,0 at unit scale.  q2tp's raw B is
12520        // (c-1.5), while the affine Prism operator is (c-1.0).
12521        bytes[..Q2TP_CHUNK].fill(0x24); // codes 0,1,2,0 in LSB-first order
12522        bytes[Q2TP_CHUNK..Q2TP_CHUNK + 2].copy_from_slice(&0u16.to_le_bytes());
12523        bytes[Q2TP_CHUNK + 2..Q2TP_CHUNK + 4].copy_from_slice(&0u16.to_le_bytes());
12524        bytes[Q2TP_CHUNK + 4] = 1; // dtype16 rung 1 = 1.0
12525        let x = vec![1.0f32; cols];
12526        let mut raw = vec![0.0f32; rows];
12527        let mut affine = vec![0.0f32; rows];
12528        q2tp_matvec_for_test(&bytes, &x, rows, cols, &mut raw);
12529        q2tp_affine_matvec_for_test(&bytes, &x, rows, cols, &mut affine);
12530        assert_eq!(raw, vec![-24.0]);
12531        assert_eq!(affine, vec![-8.0]);
12532        assert!((affine[0] - (raw[0] + 0.5 * cols as f32)).abs() < 1e-6);
12533    }
12534
12535    #[test]
12536    fn q8_row_dot_fast_matches_scalar() {
12537        // The per-arch fast dot must agree with the exact scalar oracle
12538        // (same contract the fused q8 FFN arm rides on).
12539        let cols = 96;
12540        let row: Vec<u8> = (0..cols)
12541            .map(|i| ((i * 37 % 251) - 125) as i8 as u8)
12542            .collect();
12543        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.13).sin()).collect();
12544        let act = split_act(&x);
12545        let fast = q8_row_dot(&row, &act);
12546        let scalar = q8_row_dot_scalar(&row, &act);
12547        assert!(
12548            (fast - scalar).abs() <= scalar.abs() * 1e-5 + 1e-5,
12549            "fast {fast} vs scalar {scalar}"
12550        );
12551    }
12552
12553    #[test]
12554    fn f32_matvec_matches_matvec_rows_bitexact() {
12555        let (rows, cols) = (300, 40);
12556        let w: Vec<f32> = (0..rows * cols).map(|i| (i as f32 * 0.017).sin()).collect();
12557        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.05).cos()).collect();
12558        let qt = QTensor::from_f32(w.clone(), rows, cols);
12559
12560        let mut a = vec![0.0f32; rows];
12561        matvec_rows(None, &w, &x, &mut a);
12562        let mut b = vec![0.0f32; rows];
12563        qt.matvec(&x, &mut b, None);
12564        assert_eq!(a, b);
12565    }
12566
12567    #[test]
12568    fn sdot_kernel_exact_on_grid() {
12569        // Activations already on the i8 grid (±1 with amax=1 → sx=1/127,
12570        // xq=±127 dequantizes EXACTLY) → the SDOT path must match the
12571        // exact f32 dot to float rounding. This isolates kernel
12572        // correctness from quantization noise.
12573        eprintln!("sdot_enabled = {}", sdot_enabled());
12574        let (rows, cols) = (9, 80); // odd rows → exercises 4-row + tail
12575        let w: Vec<u8> = (0..rows * cols)
12576            .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
12577            .collect();
12578        let scales: Vec<f32> = (0..rows).map(|o| 0.005 + o as f32 * 0.001).collect();
12579        let x: Vec<f32> = (0..cols)
12580            .map(|i| match i % 3 {
12581                0 => 1.0,
12582                1 => -1.0,
12583                _ => 0.0,
12584            })
12585            .collect();
12586        let mut a = vec![0.0f32; rows];
12587        qmatvec(
12588            &w,
12589            &[],
12590            &scales,
12591            &x,
12592            &[],
12593            TensorDtype::Q8Row,
12594            rows,
12595            cols,
12596            &mut a,
12597            None,
12598        );
12599        for o in 0..rows {
12600            let mut acc = 0.0f32;
12601            for j in 0..cols {
12602                acc += (w[o * cols + j] as i8) as f32 * x[j];
12603            }
12604            let expect = acc * scales[o];
12605            assert!(
12606                (a[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
12607                "row {o}: {} vs {expect}",
12608                a[o]
12609            );
12610        }
12611    }
12612
12613    #[test]
12614    fn q1_tbl_fast_path_matches_reference() {
12615        // gpr = 8 exercises the TBL pair-load fast loop, and the LAST
12616        // row's final 4-tile window trips the 4B-overread guard (the
12617        // payload ends exactly at the last tile) — both paths must
12618        // agree with the dequant reference.
12619        let (rows, cols) = (5, 256);
12620        let gpr = cols / GROUP_SIZE;
12621        let mut bytes = Vec::new();
12622        for t in 0..rows * gpr {
12623            let s = 0.007 + (t % 11) as f32 * 0.004;
12624            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
12625            for j in 0..4 {
12626                bytes.push(((t * 53 + j * 89 + 7) % 249) as u8);
12627            }
12628        }
12629        let x: Vec<f32> = (0..cols)
12630            .map(|i| if (i * 5) % 7 < 3 { 1.0 } else { -1.0 })
12631            .collect();
12632        let mut w = vec![0.0f32; rows * cols];
12633        cortiq_core::quant::dequant_q1(&bytes, &mut w);
12634        let mut got = vec![0.0f32; rows];
12635        q1_matvec(&bytes, &x, rows, cols, &mut got, None);
12636        for o in 0..rows {
12637            let expect: f32 = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
12638            assert!(
12639                (got[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
12640                "row {o}: {} vs {expect}",
12641                got[o]
12642            );
12643        }
12644        // Blocked 1×4 batch (b=5: one quad + remainder) must equal the
12645        // single-matvec path bit-for-bit.
12646        let b = 5usize;
12647        let mut xs_all = Vec::new();
12648        for bi in 0..b {
12649            xs_all.extend(x.iter().map(|v| if bi % 2 == 0 { *v } else { -*v }));
12650        }
12651        let mut mm = vec![0.0f32; b * rows];
12652        q1_matmat(&bytes, &xs_all, b, rows, cols, &mut mm, None);
12653        for bi in 0..b {
12654            let mut single = vec![0.0f32; rows];
12655            q1_matvec(
12656                &bytes,
12657                &xs_all[bi * cols..(bi + 1) * cols],
12658                rows,
12659                cols,
12660                &mut single,
12661                None,
12662            );
12663            assert_eq!(&mm[bi * rows..(bi + 1) * rows], &single[..], "stream {bi}");
12664        }
12665    }
12666
12667    #[test]
12668    fn q1_kernels_match_exact_reference() {
12669        // Synthetic q1 payload: 6-byte tiles [f16 scale][4B bits].
12670        let (rows, cols) = (7, 96);
12671        let gpr = cols / GROUP_SIZE;
12672        let mut bytes = Vec::new();
12673        for t in 0..rows * gpr {
12674            let s = 0.01 + (t % 13) as f32 * 0.003;
12675            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
12676            for j in 0..4 {
12677                bytes.push(((t * 31 + j * 97) % 251) as u8);
12678            }
12679        }
12680        // On-grid activations (±1, amax 1) → the SDOT path is exact.
12681        let x: Vec<f32> = (0..cols)
12682            .map(|i| if i % 3 == 0 { 1.0 } else { -1.0 })
12683            .collect();
12684        // Reference through the core dequant.
12685        let mut w = vec![0.0f32; rows * cols];
12686        cortiq_core::quant::dequant_q1(&bytes, &mut w);
12687        let mut expect = vec![0.0f32; rows];
12688        for o in 0..rows {
12689            expect[o] = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
12690        }
12691        let mut got = vec![0.0f32; rows];
12692        q1_matvec(&bytes, &x, rows, cols, &mut got, None);
12693        for o in 0..rows {
12694            assert!(
12695                (got[o] - expect[o]).abs() < 1e-3 * expect[o].abs().max(1e-3),
12696                "row {o}: {} vs {}",
12697                got[o],
12698                expect[o]
12699            );
12700        }
12701        // Pair and batch paths agree with the single path.
12702        let x2: Vec<f32> = x.iter().map(|v| -v).collect();
12703        let (mut a1, mut a2) = (vec![0.0f32; rows], vec![0.0f32; rows]);
12704        q1_matvec2(&bytes, &x, &x2, rows, cols, &mut a1, &mut a2, None);
12705        assert_eq!(a1, got);
12706        let mut xs = x.clone();
12707        xs.extend_from_slice(&x2);
12708        let mut mm = vec![0.0f32; 2 * rows];
12709        q1_matmat(&bytes, &xs, 2, rows, cols, &mut mm, None);
12710        assert_eq!(&mm[..rows], got.as_slice());
12711        assert_eq!(&mm[rows..], a2.as_slice());
12712    }
12713
12714    /// A Bonsai-shaped q1 payload (one f16 scale per 128-group, stored in
12715    /// all four tiles; pseudo-random signs) and its q1g repack.
12716    fn q1_q1g_pair(rows: usize, cols: usize, salt: usize) -> (Vec<u8>, Vec<u8>) {
12717        let tiles = rows * cols / GROUP_SIZE;
12718        let mut q1 = Vec::with_capacity(tiles * Q1_TILE);
12719        for t in 0..tiles {
12720            let s = 0.004 + ((t / 4 * 37 + salt) % 29) as f32 * 0.0011;
12721            q1.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
12722            let bits = ((t * 2_654_435_761 + salt * 40_503) as u32).rotate_left((t % 31) as u32);
12723            q1.extend_from_slice(&bits.to_le_bytes());
12724        }
12725        let rp = cortiq_core::quant::repack_q1_to_q1g(&q1, rows, cols, 0.0);
12726        assert_eq!(rp.differing, 0, "uniform scales repack exactly");
12727        (q1, rp.data)
12728    }
12729
12730    /// Activations off the int8 grid with two outlier channels, so the
12731    /// A8W8 path quantizes AND takes the exact outlier correction.
12732    fn q1g_test_x(cols: usize, salt: usize) -> Vec<f32> {
12733        let mut x: Vec<f32> = (0..cols)
12734            .map(|i| ((i * 13 + salt * 7) as f32 * 0.071).sin() * 0.9)
12735            .collect();
12736        x[3 % cols] = 41.0;
12737        x[(cols / 2 + 5) % cols] = -37.5;
12738        x
12739    }
12740
12741    /// Every CPU row kernel of q1g returns the q1 kernel's f32 BIT FOR BIT
12742    /// on a uniform-scale repack: A8W8 rows (SIMD of this machine, and the
12743    /// x86 AVX2/VNNI twins directly), the exact scalar rows (CMF_SDOT=0
12744    /// contract), the outlier term, the drivers (matvec, matvec2, matmat
12745    /// with the blocked 1×4 quad + remainder) — including gpr = 4 and the
12746    /// last row, where q1's pair load falls back to its scalar tail.
12747    #[test]
12748    fn q1g_cpu_kernels_equal_q1_bit_for_bit() {
12749        for &(rows, cols) in &[(3usize, 128usize), (5, 256), (9, 2048), (2, 6144)] {
12750            let gpr = cols / GROUP_SIZE;
12751            let (q1, q1g) = q1_q1g_pair(rows, cols, rows + cols);
12752            let (wa, wb) = (BinW::Q1(&q1), BinW::q1g(&q1g, rows, cols));
12753            let x = q1g_test_x(cols, rows);
12754            let act = split_act(&x);
12755            assert!(!act.outliers.is_empty(), "test wants the outlier path");
12756            let gsum = q1_group_sums(&act.xq, gpr);
12757            for r in 0..rows {
12758                let (a, b) = (wa.row_exact(r, gpr, &x), wb.row_exact(r, gpr, &x));
12759                assert_eq!(a.to_bits(), b.to_bits(), "{rows}x{cols} exact row {r}");
12760                let (a, b) = (
12761                    wa.dot_i8(r, gpr, &act.xq, &gsum),
12762                    wb.dot_i8(r, gpr, &act.xq, &gsum),
12763                );
12764                assert_eq!(a.to_bits(), b.to_bits(), "{rows}x{cols} i8 row {r}");
12765                for j in [0, 31, 32, 127, cols - 1] {
12766                    let (a, b) = (wa.outlier(r, gpr, j), wb.outlier(r, gpr, j));
12767                    assert_eq!(
12768                        (a.0.to_bits(), a.1.to_bits()),
12769                        (b.0.to_bits(), b.1.to_bits())
12770                    );
12771                }
12772                #[cfg(target_arch = "x86_64")]
12773                if avx2_enabled() {
12774                    let sc = cortiq_core::quant::q1g_scale_offset(rows, cols);
12775                    // SAFETY: AVX2 checked; payloads sized for (rows, cols).
12776                    unsafe {
12777                        let a = dot_q1_row_avx2(&q1, r, gpr, &act.xq, &gsum);
12778                        let b = dot_q1g_row_avx2(&q1g, sc, r, gpr, &act.xq, &gsum);
12779                        assert_eq!(a.to_bits(), b.to_bits(), "avx2 row {r}");
12780                    }
12781                    if std::is_x86_feature_detected!("avx512vnni")
12782                        && std::is_x86_feature_detected!("avx512bw")
12783                        && std::is_x86_feature_detected!("avx512vl")
12784                    {
12785                        // SAFETY: VNNI family detected just above.
12786                        unsafe {
12787                            let a = dot_q1_row_vnni(&q1, r, gpr, &act.xq, &gsum);
12788                            let b = dot_q1g_row_vnni(&q1g, sc, r, gpr, &act.xq, &gsum);
12789                            assert_eq!(a.to_bits(), b.to_bits(), "vnni row {r}");
12790                        }
12791                    }
12792                }
12793            }
12794            // Drivers: single, pair, batch (b = 5: one blocked quad + one).
12795            let (mut ga, mut gb) = (vec![0f32; rows], vec![0f32; rows]);
12796            bin_matvec(wa, &x, rows, cols, &mut ga, None);
12797            bin_matvec(wb, &x, rows, cols, &mut gb, None);
12798            assert_eq!(ga, gb, "{rows}x{cols} matvec");
12799            let x2 = q1g_test_x(cols, rows + 1);
12800            let (mut a1, mut a2, mut b1, mut b2) = (
12801                vec![0f32; rows],
12802                vec![0f32; rows],
12803                vec![0f32; rows],
12804                vec![0f32; rows],
12805            );
12806            bin_matvec2(wa, &x, &x2, rows, cols, &mut a1, &mut a2, None);
12807            bin_matvec2(wb, &x, &x2, rows, cols, &mut b1, &mut b2, None);
12808            assert_eq!((&a1, &a2), (&b1, &b2), "{rows}x{cols} matvec2");
12809            let b = 5usize;
12810            let xs: Vec<f32> = (0..b).flat_map(|k| q1g_test_x(cols, k * 3 + 2)).collect();
12811            let (mut ma, mut mb) = (vec![0f32; b * rows], vec![0f32; b * rows]);
12812            bin_matmat(wa, &xs, b, rows, cols, &mut ma, None);
12813            bin_matmat(wb, &xs, b, rows, cols, &mut mb, None);
12814            assert_eq!(ma, mb, "{rows}x{cols} matmat");
12815            // The blocked 1×4 kernels directly (the driver only takes them
12816            // under CMF_X86_BLOCKED / SDOT, a process-wide switch).
12817            let acts: Vec<(SplitAct, Vec<i32>)> = (0..4)
12818                .map(|k| {
12819                    let a = split_act(&xs[k * cols..(k + 1) * cols]);
12820                    let g = q1_group_sums(&a.xq, gpr);
12821                    (a, g)
12822                })
12823                .collect();
12824            let xq4: [&[i8]; 4] = std::array::from_fn(|k| acts[k].0.xq.as_slice());
12825            let gs4: [&[i32]; 4] = std::array::from_fn(|k| acts[k].1.as_slice());
12826            for r in 0..rows {
12827                #[cfg(target_arch = "x86_64")]
12828                if avx2_enabled() {
12829                    // SAFETY: AVX2 checked; payloads sized for (rows, cols).
12830                    let (a, b) = unsafe {
12831                        (
12832                            wa.dot_1x4_x86(r, gpr, xq4, gs4),
12833                            wb.dot_1x4_x86(r, gpr, xq4, gs4),
12834                        )
12835                    };
12836                    assert_eq!(a.map(f32::to_bits), b.map(f32::to_bits), "1x4 row {r}");
12837                }
12838                #[cfg(target_arch = "aarch64")]
12839                if sdot_enabled() {
12840                    // SAFETY: SDOT checked; payloads sized for (rows, cols).
12841                    let (a, b) = unsafe {
12842                        (
12843                            wa.dot_1x4_sdot(r, gpr, xq4, gs4),
12844                            wb.dot_1x4_sdot(r, gpr, xq4, gs4),
12845                        )
12846                    };
12847                    assert_eq!(a.map(f32::to_bits), b.map(f32::to_bits), "1x4 row {r}");
12848                }
12849                let _ = (r, &xq4, &gs4);
12850            }
12851        }
12852    }
12853
12854    /// CPU kernel A/B, q1 vs q1g matvec (A8W8 path of this machine, the
12855    /// exact path too), Bonsai-1.7B decode shapes, arms alternated every
12856    /// round, medians printed. `CMF_THREADS` sets the pool (default 8).
12857    /// `cargo test -p cortiq-engine --release --lib q1g_cpu_matvec_speed --
12858    /// --ignored --nocapture`
12859    #[test]
12860    #[ignore]
12861    fn q1g_cpu_matvec_speed() {
12862        use crate::pool::Pool;
12863        let threads = std::env::var("CMF_THREADS")
12864            .ok()
12865            .and_then(|v| v.parse().ok())
12866            .unwrap_or(8usize);
12867        let pool = Pool::new(threads);
12868        for &(name, rows, cols, reps) in &[
12869            ("gate/up_proj", 6144usize, 2048usize, 41usize),
12870            ("down_proj", 2048, 6144, 41),
12871            ("lm_head", 151_669, 2048, 9),
12872        ] {
12873            let (q1, q1g) = q1_q1g_pair(rows, cols, 3);
12874            let (wa, wb) = (BinW::Q1(&q1), BinW::q1g(&q1g, rows, cols));
12875            let x = q1g_test_x(cols, 1);
12876            let mut out = vec![0f32; rows];
12877            let mut time = |w: BinW| {
12878                let t0 = std::time::Instant::now();
12879                bin_matvec(w, &x, rows, cols, &mut out, Some(&pool));
12880                t0.elapsed().as_secs_f64()
12881            };
12882            let _ = (time(wa), time(wb));
12883            let (mut ta, mut tb) = (Vec::new(), Vec::new());
12884            for r in 0..reps {
12885                if r % 2 == 0 {
12886                    ta.push(time(wa));
12887                    tb.push(time(wb));
12888                } else {
12889                    tb.push(time(wb));
12890                    ta.push(time(wa));
12891                }
12892            }
12893            ta.sort_by(|a, b| a.partial_cmp(b).unwrap());
12894            tb.sort_by(|a, b| a.partial_cmp(b).unwrap());
12895            let (ma, mb) = (ta[reps / 2], tb[reps / 2]);
12896            eprintln!(
12897                "cpu x{threads} a8w8={} {name:>12} {rows}x{cols}: q1 {:8.1} us | q1g {:8.1} us | q1g/q1 time {:.3}",
12898                a8w8_enabled(),
12899                ma * 1e6,
12900                mb * 1e6,
12901                mb / ma
12902            );
12903        }
12904    }
12905
12906    /// The mapped QTensor paths a real Bonsai q1g file takes — matvec, the
12907    /// multi-projection job (all-q1g and q1/q1g mixed), the fused gate/up
12908    /// with SiLU, matmat, matvec2 and the embedding row — equal the q1
12909    /// tensor they were repacked from, bit for bit.
12910    #[test]
12911    fn q1g_mapped_tensor_paths_equal_q1() {
12912        use crate::pool::Pool;
12913        use cortiq_core::{CmfHeader, CmfModel, TensorSpec};
12914        let (rows, cols) = (300usize, 256usize);
12915        let (gq1, gq1g) = q1_q1g_pair(rows, cols, 1);
12916        let (uq1, uq1g) = q1_q1g_pair(rows, cols, 2);
12917        let header: CmfHeader = serde_json::from_value(serde_json::json!({
12918            "format": "cmf", "version": cortiq_core::CMF_VERSION,
12919            "arch": {"arch_name": "tiny-q1g", "hidden_size": cols,
12920                "intermediate_size": rows, "num_layers": 1, "num_attention_heads": 2,
12921                "num_kv_heads": 1, "head_dim": 128, "vocab_size": rows,
12922                "layer_types": ["FullAttention"], "rms_norm_eps": 1e-6,
12923                "max_position_embeddings": 8},
12924            "quant_type": "VBIT"
12925        }))
12926        .unwrap();
12927        let spec = |name: &str, dtype, data: &Vec<u8>| TensorSpec {
12928            name: name.into(),
12929            dtype,
12930            shape: vec![rows, cols],
12931            data: data.clone(),
12932        };
12933        let specs = [
12934            spec("g1", TensorDtype::Q1, &gq1),
12935            spec("u1", TensorDtype::Q1, &uq1),
12936            spec("gg", TensorDtype::Q1G, &gq1g),
12937            spec("ug", TensorDtype::Q1G, &uq1g),
12938        ];
12939        let dir = std::env::temp_dir().join(format!("cmf-q1g-mapped-{}", std::process::id()));
12940        std::fs::create_dir_all(&dir).unwrap();
12941        let path = dir.join("m.cmf");
12942        CmfModel::write(&path, &header, &specs, None, None).unwrap();
12943        let model = Arc::new(CmfModel::open(&path).unwrap());
12944        assert!(model.verify().is_empty(), "{:?}", model.verify());
12945        let t = |n: &str| QTensor::from_model(&model, n).unwrap();
12946        let (g1, u1, gg, ug) = (t("g1"), t("u1"), t("gg"), t("ug"));
12947        assert_eq!(gg.model_dtype(), Some(TensorDtype::Q1G));
12948        assert!(gg.is_q1() && gg.graph_weight().map(|w| w.2) == Some(10));
12949        let pool = Pool::new(3);
12950        let x = q1g_test_x(cols, 4);
12951        let run = |a: &QTensor| {
12952            let mut o = vec![0f32; rows];
12953            a.matvec(&x, &mut o, Some(&pool));
12954            o
12955        };
12956        assert_eq!(run(&g1), run(&gg), "matvec");
12957        let many = |a: &QTensor, b: &QTensor| {
12958            let (mut oa, mut ob) = (vec![0f32; rows], vec![0f32; rows]);
12959            QTensor::matvec_many([a, b], &x, [&mut oa, &mut ob], Some(&pool));
12960            (oa, ob)
12961        };
12962        let want = many(&g1, &u1);
12963        assert_eq!(want, many(&gg, &ug), "matvec_many q1g");
12964        assert_eq!(want, many(&g1, &ug), "matvec_many mixed");
12965        assert_eq!(want, (run(&g1), run(&u1)), "matvec_many == singles");
12966        let fused = |g: &QTensor, u: &QTensor| {
12967            let mut o = vec![0f32; rows];
12968            let ok = QTensor::matvec_silu_mul(g, u, &x, &mut o, Some(&pool));
12969            (ok, o)
12970        };
12971        let (ok_a, fa) = fused(&g1, &u1);
12972        let (ok_b, fb) = fused(&gg, &ug);
12973        assert_eq!(ok_a, ok_b, "q1g takes the fused arm whenever q1 does");
12974        assert_eq!(fa, fb, "fused gate/up");
12975        let b = 6usize;
12976        let xs: Vec<f32> = (0..b).flat_map(|k| q1g_test_x(cols, k + 9)).collect();
12977        let mm = |a: &QTensor| {
12978            let mut o = vec![0f32; b * rows];
12979            a.matmat(&xs, b, &mut o, Some(&pool));
12980            o
12981        };
12982        assert_eq!(mm(&g1), mm(&gg), "matmat");
12983        let x2 = q1g_test_x(cols, 77);
12984        let pair = |a: &QTensor| {
12985            let (mut o1, mut o2) = (vec![0f32; rows], vec![0f32; rows]);
12986            a.matvec2(&x, &x2, &mut o1, &mut o2, Some(&pool));
12987            (o1, o2)
12988        };
12989        assert_eq!(pair(&g1), pair(&gg), "matvec2");
12990        for r in [0, 1, rows / 2, rows - 1] {
12991            let (mut ra, mut rb) = (vec![0f32; cols], vec![0f32; cols]);
12992            g1.row_f32(r, &mut ra);
12993            gg.row_f32(r, &mut rb);
12994            assert_eq!(ra, rb, "embedding row {r}");
12995        }
12996        drop((g1, u1, gg, ug, model));
12997        let _ = std::fs::remove_dir_all(&dir);
12998    }
12999
13000    /// A `blr` tensor read from a real CMF file: it maps to the factored
13001    /// arm (not a dense copy), declines every device-graph accessor, and
13002    /// every QTensor entry point a layer may take — matvec, matvec2,
13003    /// matmat (table loop and GEMM sizes), the multi-tensor job mixed with
13004    /// a q8_row sibling, the fused gate/up refusal, rows and columns — equals
13005    /// the dense reference built by the core dequant.
13006    #[test]
13007    fn blr_mapped_tensor_paths_equal_dense() {
13008        use crate::pool::Pool;
13009        use cortiq_core::{CmfHeader, CmfModel, TensorSpec};
13010        let (rows, cols, r) = (96usize, 128usize, 16usize);
13011        let mut rng = crate::blr::tests::Rng(41);
13012        let gate = crate::blr::tests::make(&mut rng, rows, cols, r, 2);
13013        let up = crate::blr::tests::make(&mut rng, rows, cols, r, 1);
13014        // q8_row sibling: [int8 rows·cols][f16 row scale · rows].
13015        let mut q8: Vec<u8> = (0..rows * cols)
13016            .map(|i| (((i * 89) % 253) as i32 - 126) as i8 as u8)
13017            .collect();
13018        for o in 0..rows {
13019            q8.extend_from_slice(&cortiq_core::quant::f32_to_f16(0.002 + o as f32 * 1e-4).to_le_bytes());
13020        }
13021        let header: CmfHeader = serde_json::from_value(serde_json::json!({
13022            "format": "cmf", "version": cortiq_core::CMF_VERSION,
13023            "arch": {"arch_name": "tiny-blr", "hidden_size": cols,
13024                "intermediate_size": rows, "num_layers": 1, "num_attention_heads": 2,
13025                "num_kv_heads": 1, "head_dim": 64, "vocab_size": rows,
13026                "layer_types": ["FullAttention"], "rms_norm_eps": 1e-6,
13027                "max_position_embeddings": 8},
13028            "quant_type": "Q8_2F"
13029        }))
13030        .unwrap();
13031        let specs = [
13032            TensorSpec {
13033                name: "g".into(),
13034                dtype: TensorDtype::BinaryLowRank,
13035                shape: vec![rows, cols],
13036                data: gate.clone(),
13037            },
13038            TensorSpec {
13039                name: "u".into(),
13040                dtype: TensorDtype::BinaryLowRank,
13041                shape: vec![rows, cols],
13042                data: up.clone(),
13043            },
13044            TensorSpec {
13045                name: "q".into(),
13046                dtype: TensorDtype::Q8Row,
13047                shape: vec![rows, cols],
13048                data: q8,
13049            },
13050        ];
13051        let dir = std::env::temp_dir().join(format!("cmf-blr-mapped-{}", std::process::id()));
13052        std::fs::create_dir_all(&dir).unwrap();
13053        let path = dir.join("m.cmf");
13054        CmfModel::write(&path, &header, &specs, None, None).unwrap();
13055        let model = Arc::new(CmfModel::open(&path).unwrap());
13056        assert!(model.verify().is_empty(), "{:?}", model.verify());
13057        let t = |n: &str| QTensor::from_model(&model, n).unwrap();
13058        let (g, u, q) = (t("g"), t("u"), t("q"));
13059        assert!(g.is_blr() && u.is_blr() && !q.is_blr());
13060        assert!(g.graph_weight().is_none() && g.graph_weight_descriptor().is_none());
13061        assert!(g.as_f32().is_none() && g.model_idx().is_none());
13062        assert!(g.mapped_device_gemm().is_none() && g.q1_parts().is_none());
13063        assert!(!g.sparse_col_ok());
13064        assert_eq!((g.rows(), g.cols()), (rows, cols));
13065        let deq = |bytes: &[u8]| {
13066            let mut w = vec![0f32; rows * cols];
13067            cortiq_core::blr::dequant_blr(bytes, rows, cols, &mut w).unwrap();
13068            w
13069        };
13070        let (wg, wu) = (deq(&gate), deq(&up));
13071        let mv = |w: &[f32], x: &[f32]| -> Vec<f32> {
13072            (0..rows)
13073                .map(|i| (0..cols).map(|j| w[i * cols + j] * x[j]).sum())
13074                .collect()
13075        };
13076        let close = |a: &[f32], b: &[f32], what: &str| {
13077            let s = b.iter().map(|v| v * v).sum::<f32>().sqrt().max(1e-30);
13078            let e = a.iter().zip(b).map(|(x, y)| (x - y) * (x - y)).sum::<f32>().sqrt();
13079            assert!(e / s < 1e-5, "{what}: rel {}", e / s);
13080        };
13081        let pool = Pool::new(3);
13082        let x: Vec<f32> = (0..cols).map(|_| rng.f()).collect();
13083        let x2: Vec<f32> = (0..cols).map(|_| rng.f()).collect();
13084        let mut o = vec![0f32; rows];
13085        g.matvec(&x, &mut o, Some(&pool));
13086        close(&o, &mv(&wg, &x), "matvec");
13087        let (mut o1, mut o2) = (vec![0f32; rows], vec![0f32; rows]);
13088        g.matvec2(&x, &x2, &mut o1, &mut o2, Some(&pool));
13089        close(&o1, &mv(&wg, &x), "matvec2 a");
13090        close(&o2, &mv(&wg, &x2), "matvec2 b");
13091        for b in [2usize, 5] {
13092            let xs: Vec<f32> = (0..b * cols).map(|_| rng.f()).collect();
13093            let mut om = vec![0f32; b * rows];
13094            g.matmat(&xs, b, &mut om, Some(&pool));
13095            for t in 0..b {
13096                close(
13097                    &om[t * rows..(t + 1) * rows],
13098                    &mv(&wg, &xs[t * cols..(t + 1) * cols]),
13099                    "matmat",
13100                );
13101            }
13102        }
13103        // Multi-tensor job mixed with a q8_row sibling: per-tensor fallback,
13104        // the q8_row output identical to its own matvec.
13105        let (mut a, mut bq, mut c) = (vec![0f32; rows], vec![0f32; rows], vec![0f32; rows]);
13106        QTensor::matvec_many([&g, &q, &u], &x, [&mut a, &mut bq, &mut c], Some(&pool));
13107        close(&a, &mv(&wg, &x), "matvec_many blr gate");
13108        close(&c, &mv(&wu, &x), "matvec_many blr up");
13109        let mut qo = vec![0f32; rows];
13110        q.matvec(&x, &mut qo, Some(&pool));
13111        assert_eq!(bq, qo, "matvec_many q8_row sibling");
13112        let mut f = vec![0f32; rows];
13113        assert!(!QTensor::matvec_silu_mul(&g, &u, &x, &mut f, Some(&pool)));
13114        let mut row = vec![0f32; cols];
13115        g.row_f32(rows / 3, &mut row);
13116        close(&row, &wg[(rows / 3) * cols..(rows / 3 + 1) * cols], "row_f32");
13117        let mut scratch = vec![0f32; cols];
13118        let d = g.row_dot(7, &x, &mut scratch);
13119        let want: f32 = (0..cols).map(|j| wg[7 * cols + j] * x[j]).sum();
13120        assert!((d - want).abs() < 1e-4 * (1.0 + want.abs()), "row_dot");
13121        let mut col = vec![0f32; rows];
13122        g.add_col_scaled(cols / 2, 1.5, &mut col);
13123        let want: Vec<f32> = (0..rows).map(|i| 1.5 * wg[i * cols + cols / 2]).collect();
13124        close(&col, &want, "add_col_scaled");
13125        drop((g, u, q, model));
13126        let _ = std::fs::remove_dir_all(&dir);
13127    }
13128
13129    #[test]
13130    fn repack_is_bit_identical() {
13131        // The interleaved-repack kernel must produce EXACTLY the same
13132        // bits as the mmap-layout kernel: integer accumulation is order-
13133        // exact, the f32 epilogue is identical. Odd rows exercise the
13134        // tail; direct range calls exercise unaligned pool splits.
13135        let (rows, cols) = (267, 96); // 66 groups + 3 tail rows, cols % 16 == 0
13136        let w: Vec<u8> = (0..rows * cols)
13137            .map(|i| (((i * 89) % 253) as i32 - 126) as i8 as u8)
13138            .collect();
13139        let scales: Vec<f32> = (0..rows).map(|o| 0.003 + o as f32 * 0.0007).collect();
13140        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.37).sin() * 2.0).collect();
13141        let rep = q8_repack_layout(&w, rows, cols);
13142        // Group interleave round-trips.
13143        for g in 0..rows / 4 {
13144            for c in 0..cols / 16 {
13145                for lane in 0..4 {
13146                    assert_eq!(
13147                        &rep[g * 4 * cols + c * 64 + lane * 16
13148                            ..g * 4 * cols + c * 64 + lane * 16 + 16],
13149                        &w[(g * 4 + lane) * cols + c * 16..(g * 4 + lane) * cols + c * 16 + 16],
13150                    );
13151                }
13152            }
13153        }
13154        let mut a = vec![0.0f32; rows];
13155        qmatvec(
13156            &w,
13157            &[],
13158            &scales,
13159            &x,
13160            &[],
13161            TensorDtype::Q8Row,
13162            rows,
13163            cols,
13164            &mut a,
13165            None,
13166        );
13167        let mut b = vec![0.0f32; rows];
13168        qmatvec(
13169            &w,
13170            &rep,
13171            &scales,
13172            &x,
13173            &[],
13174            TensorDtype::Q8Row,
13175            rows,
13176            cols,
13177            &mut b,
13178            None,
13179        );
13180        assert_eq!(a, b, "full-range repack output diverged");
13181
13182        #[cfg(target_arch = "aarch64")]
13183        if sdot_enabled() {
13184            // Unaligned range split (pool workers get arbitrary bounds).
13185            let act = split_act(&x);
13186            let mut c1 = vec![0.0f32; rows];
13187            let mut c2 = vec![0.0f32; rows];
13188            q8_range_sdot(
13189                &w,
13190                &[],
13191                &scales,
13192                &act,
13193                cols,
13194                SendMut(c1.as_mut_ptr()),
13195                3,
13196                rows - 2,
13197            );
13198            q8_range_sdot(
13199                &w,
13200                &rep,
13201                &scales,
13202                &act,
13203                cols,
13204                SendMut(c2.as_mut_ptr()),
13205                3,
13206                rows - 2,
13207            );
13208            assert_eq!(c1, c2, "unaligned-range repack output diverged");
13209        }
13210    }
13211
13212    #[test]
13213    fn sdot_a8w8_noise_is_bounded() {
13214        // Off-grid activations: A8 quantization noise must stay small in
13215        // relative L2 over the whole output (realistic accuracy contract;
13216        // vmfcore measured argmax-identical decode on real models).
13217        let (rows, cols) = (16, 512);
13218        let w: Vec<u8> = (0..rows * cols)
13219            .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
13220            .collect();
13221        let scales = vec![0.01f32; rows];
13222        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.21).sin()).collect();
13223        let mut a = vec![0.0f32; rows];
13224        qmatvec(
13225            &w,
13226            &[],
13227            &scales,
13228            &x,
13229            &[],
13230            TensorDtype::Q8Row,
13231            rows,
13232            cols,
13233            &mut a,
13234            None,
13235        );
13236        let (mut num, mut den) = (0f64, 0f64);
13237        for o in 0..rows {
13238            let mut acc = 0.0f32;
13239            for j in 0..cols {
13240                acc += (w[o * cols + j] as i8) as f32 * x[j];
13241            }
13242            let expect = acc * scales[o];
13243            num += ((a[o] - expect) as f64).powi(2);
13244            den += (expect as f64).powi(2);
13245        }
13246        let rel = (num / den.max(1e-12)).sqrt();
13247        assert!(rel < 0.05, "A8W8 relative L2 error too high: {rel}");
13248    }
13249
13250    #[test]
13251    fn i8_dot_neon_matches_scalar() {
13252        let n = 100;
13253        let w: Vec<u8> = (0..n).map(|i| ((i * 37 + 11) % 251) as u8).collect();
13254        let x: Vec<f32> = (0..n).map(|i| (i as f32 * 0.13).sin()).collect();
13255        let mut scalar = 0.0f32;
13256        for j in 0..n {
13257            scalar += (w[j] as i8) as f32 * x[j];
13258        }
13259        let fast = dot_i8_f32(&w, &x);
13260        assert!((scalar - fast).abs() < 1e-3 * scalar.abs().max(1.0));
13261    }
13262
13263    /// Fused vbit matvec must match full dequant_vbit + dense matvec.
13264    #[test]
13265    fn vbitmatvec_matches_full_dequant() {
13266        let (rows, cols) = (6, 64);
13267        let ng = cols / GROUP_SIZE;
13268        // Hand-craft: bits per row, f16 scales, packed rows.
13269        let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
13270        let mut bytes = bits.clone();
13271        for g in 0..rows * ng {
13272            let s = 0.02 + 0.001 * g as f32;
13273            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
13274        }
13275        for r in 0..rows {
13276            let b = bits[r] as usize;
13277            let (mut acc, mut nb) = (0u64, 0usize);
13278            let mut rowbytes = Vec::new();
13279            for i in 0..cols {
13280                let v = ((i * 7 + r * 13) % (1 << b)) as u64;
13281                acc = (acc << b) | v;
13282                nb += b;
13283                while nb >= 8 {
13284                    nb -= 8;
13285                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
13286                }
13287            }
13288            if nb > 0 {
13289                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
13290            }
13291            bytes.extend_from_slice(&rowbytes);
13292        }
13293        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
13294
13295        let mut reference = vec![0f32; rows * cols];
13296        cortiq_core::quant::dequant_vbit(&bytes, rows, cols, &mut reference).unwrap();
13297        let mut expect = vec![0f32; rows];
13298        for r in 0..rows {
13299            expect[r] = reference[r * cols..(r + 1) * cols]
13300                .iter()
13301                .zip(&x)
13302                .map(|(w, xv)| w * xv)
13303                .sum();
13304        }
13305        let mut got = vec![0f32; rows];
13306        let offsets = vbit_row_offsets(&bytes, rows, cols);
13307        vbitmatvec(&bytes, &offsets, &x, rows, cols, &mut got, None);
13308        // SDOT path quantizes activations to i8 (A8W8): bounded noise,
13309        // same contract as q8 (exact path is pinned by CMF_SDOT=0 in
13310        // the golden-parity gate).
13311        let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
13312        let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
13313        for r in 0..rows {
13314            assert!(
13315                (got[r] - expect[r]).abs() < tol * scale,
13316                "row {r}: {} vs {}",
13317                got[r],
13318                expect[r]
13319            );
13320        }
13321    }
13322
13323    /// Fused q4 matvec must match the reference full-dequant + dense
13324    /// matvec bit-for-bit in structure (same f32 math, group order).
13325    /// vbit matmat: the blocked 1×4 leg must match the per-row path
13326    /// (paired env toggle; larger shape so both code paths engage).
13327    #[test]
13328    #[cfg(target_arch = "x86_64")]
13329    fn vbit_matmat_blocked_matches_per_row() {
13330        let (rows, cols, b) = (64usize, 128usize, 9usize);
13331        let ng = cols / GROUP_SIZE;
13332        let bits: Vec<u8> = (0..rows).map(|r| [3u8, 4, 5, 6][r % 4]).collect();
13333        let mut bytes = bits.clone();
13334        for g in 0..rows * ng {
13335            let sc = 0.02 + 0.0005 * g as f32;
13336            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
13337        }
13338        for r in 0..rows {
13339            let bw = bits[r] as usize;
13340            let (mut acc, mut nb) = (0u64, 0usize);
13341            let mut rowbytes = Vec::new();
13342            for i in 0..cols {
13343                let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
13344                acc = (acc << bw) | v;
13345                nb += bw;
13346                while nb >= 8 {
13347                    nb -= 8;
13348                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
13349                }
13350            }
13351            if nb > 0 {
13352                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
13353            }
13354            bytes.extend_from_slice(&rowbytes);
13355        }
13356        let x: Vec<f32> = (0..b * cols)
13357            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
13358            .collect();
13359        let offsets = vbit_row_offsets(&bytes, rows, cols);
13360        let mut y_a = vec![0f32; b * rows];
13361        let mut y_b = vec![0f32; b * rows];
13362        unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
13363        vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_a, None);
13364        unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
13365        vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_b, None);
13366        unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
13367        let max_d = y_a
13368            .iter()
13369            .zip(&y_b)
13370            .map(|(p, q)| (p - q).abs())
13371            .fold(0.0f32, f32::max);
13372        assert!(max_d < 1e-4, "vbit blocked ≠ per-row: max|Δ| = {max_d}");
13373    }
13374
13375    /// q4t blocked 1×4 (SDOT on ARM, AVX2 on x86) must equal the
13376    /// per-row path exactly: same nibble unpack, same group order,
13377    /// same f32 accumulation — batch == matvec bit-for-bit. b=9 covers
13378    /// two full 1×4 blocks plus a remainder through the single-row
13379    /// kernel. (Both paths produce identical output, so the shared
13380    /// CMF_X86_BLOCKED env var racing with other tests cannot flip
13381    /// the verdict — worst case both sides take the same path.)
13382    #[test]
13383    fn q4t_matmat_blocked_matches_per_row() {
13384        let (rows, cols, b) = (16usize, 64usize, 9usize);
13385        let gpr = cols / GROUP_SIZE;
13386        let mut bytes = vec![0u8; rows * gpr * Q4_TILE];
13387        for r in 0..rows {
13388            for g in 0..gpr {
13389                let t = (r * gpr + g) * Q4_TILE;
13390                let sc = 0.02 + 0.001 * (r * gpr + g) as f32;
13391                bytes[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
13392                for k in 0..16 {
13393                    bytes[t + 2 + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
13394                }
13395            }
13396        }
13397        let x: Vec<f32> = (0..b * cols)
13398            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
13399            .collect();
13400        let mut y_blk = vec![0f32; b * rows];
13401        let mut y_row = vec![0f32; b * rows];
13402        unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
13403        q4t_matmat(&bytes, &x, b, rows, cols, &mut y_blk, None);
13404        unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
13405        q4t_matmat(&bytes, &x, b, rows, cols, &mut y_row, None);
13406        unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
13407        assert_eq!(y_blk, y_row, "q4t blocked 1x4 ≠ per-row");
13408    }
13409
13410    /// The wide-batch Accelerate arm of q4t_matmat vs a brute-force
13411    /// f32 dequant matmul: both are f32 GEMMs, so only reduction
13412    /// order differs — tight tolerance.
13413    /// A synthetic q4tp payload: random nibbles plus a per-row ladder whose
13414    /// span varies row to row, so the codes actually exercise the full 0..31
13415    /// range rather than clustering on one rung.
13416    fn synth_q4tp(rows: usize, cols: usize) -> Vec<u8> {
13417        use cortiq_core::quant::{f32_to_f16, q4tp_code_stride, q4tp_put_code};
13418        let gpr = cols / GROUP_SIZE;
13419        let stride = q4tp_code_stride(gpr);
13420        let (params_off, codes_off, _) = q4tp_sections(rows, cols);
13421        let mut b = vec![0u8; codes_off + rows * stride];
13422        for r in 0..rows {
13423            for g in 0..gpr {
13424                let t = (r * gpr + g) * Q4TP_NIB;
13425                for k in 0..16 {
13426                    b[t + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
13427                }
13428            }
13429            let lo = -6.0 - 0.03 * (r % 17) as f32;
13430            let step = 0.01 + 0.004 * (r % 11) as f32;
13431            let p = params_off + r * 4;
13432            b[p..p + 2].copy_from_slice(&f32_to_f16(lo).to_le_bytes());
13433            b[p + 2..p + 4].copy_from_slice(&f32_to_f16(step).to_le_bytes());
13434            let crow = &mut b[codes_off + r * stride..codes_off + (r + 1) * stride];
13435            for g in 0..gpr {
13436                q4tp_put_code(crow, g, (r * 5 + g * 3) % 32);
13437            }
13438        }
13439        b
13440    }
13441
13442    /// The same weights re-expressed as q4_tiled, so the proven kernel can
13443    /// be the reference: each tile stores the ladder scale its code selects.
13444    /// Only the f16 rounding of that scale separates the two payloads.
13445    fn q4tp_as_q4t(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
13446        let gpr = cols / GROUP_SIZE;
13447        let v = Q4tpView::new(bytes, rows, cols);
13448        let mut out = vec![0u8; rows * gpr * Q4_TILE];
13449        let mut sc = vec![0f32; gpr];
13450        for r in 0..rows {
13451            v.scales_into(r, gpr, &mut sc);
13452            for g in 0..gpr {
13453                let t = (r * gpr + g) * Q4_TILE;
13454                let s = sc[g];
13455                out[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
13456                let src = (r * gpr + g) * Q4TP_NIB;
13457                out[t + 2..t + Q4_TILE].copy_from_slice(&v.nib[src..src + Q4TP_NIB]);
13458            }
13459        }
13460        out
13461    }
13462
13463    /// The exact (`CMF_SDOT=0`) path must reproduce `dequant_q4tp` to f32
13464    /// rounding — that scalar routine is the format's definition, and the
13465    /// kernels re-derive the scale from the ladder independently. Call the
13466    /// row kernel directly: `matmat` picks the int8 arm when a8w8 is on,
13467    /// so routing through it would test the other path by accident.
13468    #[test]
13469    fn q4tp_exact_path_matches_dequant_reference() {
13470        let (rows, cols) = (256usize, 512usize);
13471        let gpr = cols / GROUP_SIZE;
13472        let bytes = synth_q4tp(rows, cols);
13473        let mut w = vec![0f32; rows * cols];
13474        cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
13475
13476        let x: Vec<f32> = (0..cols)
13477            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
13478            .collect();
13479        let v = Q4tpView::new(&bytes, rows, cols);
13480        let mut sc = vec![0f32; gpr];
13481        for r in 0..rows {
13482            v.scales_into(r, gpr, &mut sc);
13483            let got = q4tp_row_exact(v.nib, r, gpr, &x, &sc);
13484            let want: f32 = (0..cols).map(|c| w[r * cols + c] * x[c]).sum();
13485            // These dot products cancel down to ~1e-3 from terms of ~5e-2, so
13486            // the meaningful yardstick is the summed magnitude, not the result:
13487            // against the result any reordering of a 512-term f32 sum "fails".
13488            let mag: f32 = (0..cols).map(|c| (w[r * cols + c] * x[c]).abs()).sum();
13489            assert!(
13490                (got - want).abs() <= 1e-5 * mag,
13491                "row {r}: kernel {got} vs dequant {want}"
13492            );
13493        }
13494    }
13495
13496    /// The int8 (a8w8) path can't be checked against an f32 reference — the
13497    /// activation quantization dominates. Check it against the q4t kernel it
13498    /// was ported from instead, on payloads holding the same weights: that
13499    /// isolates exactly what the port could break (16 B stride, ladder
13500    /// lookup, nibble unpack) from what it deliberately shares.
13501    #[test]
13502    fn q4tp_matvec_matches_the_q4t_kernel_it_was_ported_from() {
13503        let (rows, cols) = (256usize, 512usize);
13504        let bytes = synth_q4tp(rows, cols);
13505        let twin = q4tp_as_q4t(&bytes, rows, cols);
13506        let x: Vec<f32> = (0..cols)
13507            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
13508            .collect();
13509
13510        let mut got = vec![0f32; rows];
13511        q4tp_matvec(&bytes, &x, rows, cols, &mut got, None);
13512        let mut want = vec![0f32; rows];
13513        q4t_matvec(&twin, &x, rows, cols, &mut want, None);
13514
13515        // Scale is f16 in the twin and f32 here, so allow that rounding on
13516        // top of the summed magnitude (same cancellation argument as above).
13517        let mut w = vec![0f32; rows * cols];
13518        cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
13519        for r in 0..rows {
13520            let mag: f32 = (0..cols).map(|c| (w[r * cols + c] * x[c]).abs()).sum();
13521            assert!(
13522                (got[r] - want[r]).abs() <= 1e-3 * mag,
13523                "row {r}: q4tp {} vs q4t {}",
13524                got[r],
13525                want[r]
13526            );
13527        }
13528    }
13529
13530    /// `matmat` carries three arms (Accelerate, blocked int8 1x4, scalar).
13531    /// Batch 5 crosses the blocked kernel's stride, so this exercises the
13532    /// 1x4 path AND its scalar tail in one run — the blocked kernel is new
13533    /// code and its four accumulators are exactly what tends to go wrong.
13534    #[test]
13535    fn q4tp_matmat_matches_the_q4t_kernel_it_was_ported_from() {
13536        let (rows, cols, b) = (256usize, 512usize, 5usize);
13537        let bytes = synth_q4tp(rows, cols);
13538        let twin = q4tp_as_q4t(&bytes, rows, cols);
13539        let xs: Vec<f32> = (0..b * cols)
13540            .map(|i| ((i * 29 + 11) % 89) as f32 / 89.0 - 0.5)
13541            .collect();
13542
13543        let mut got = vec![0f32; b * rows];
13544        q4tp_matmat(&bytes, &xs, b, rows, cols, &mut got, None);
13545        let mut want = vec![0f32; b * rows];
13546        q4t_matmat(&twin, &xs, b, rows, cols, &mut want, None);
13547
13548        let mut w = vec![0f32; rows * cols];
13549        cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
13550        for t in 0..b {
13551            for r in 0..rows {
13552                let mag: f32 = (0..cols)
13553                    .map(|c| (w[r * cols + c] * xs[t * cols + c]).abs())
13554                    .sum();
13555                let (g, wa) = (got[t * rows + r], want[t * rows + r]);
13556                assert!(
13557                    (g - wa).abs() <= 1e-3 * mag,
13558                    "batch {t} row {r}: q4tp {g} vs q4t {wa}"
13559                );
13560            }
13561        }
13562    }
13563
13564    #[test]
13565    fn q4tp_matvec2_matches_the_single_stream_kernel() {
13566        let (rows, cols) = (128usize, 256usize);
13567        let gpr = cols / GROUP_SIZE;
13568        let bytes = synth_q4tp(rows, cols);
13569        let xs: Vec<f32> = (0..2 * cols)
13570            .map(|i| ((i * 29 + 11) % 89) as f32 / 89.0 - 0.5)
13571            .collect();
13572
13573        let (mut o1, mut o2) = (vec![0f32; rows], vec![0f32; rows]);
13574        q4tp_matvec2(
13575            &bytes,
13576            &xs[..cols],
13577            &xs[cols..],
13578            rows,
13579            cols,
13580            &mut o1,
13581            &mut o2,
13582            None,
13583        );
13584
13585        // matvec2 takes the exact path for both streams, so the single-row
13586        // kernel is an exact reference — no tolerance for path differences.
13587        let v = Q4tpView::new(&bytes, rows, cols);
13588        let mut sc = vec![0f32; gpr];
13589        for r in 0..rows {
13590            v.scales_into(r, gpr, &mut sc);
13591            assert_eq!(o1[r], q4tp_row_exact(v.nib, r, gpr, &xs[..cols], &sc));
13592            assert_eq!(o2[r], q4tp_row_exact(v.nib, r, gpr, &xs[cols..], &sc));
13593        }
13594    }
13595
13596    /// q4tp must not COST speed — it exists to save bytes, and a format that
13597    /// trades 7% of a file for a slower model is a bad trade. This guard is
13598    /// here because correctness tests happily passed while `q4tp_matmat` was
13599    /// missing its int8 and Accelerate arms and the model ran 5x slower.
13600    /// Measured on M-series: 0.97-1.04x, i.e. parity (16 B tiles are better
13601    /// aligned than q4t's 18 B, which pays for the scale indirection).
13602    #[test]
13603    fn q4tp_matvec_keeps_pace_with_q4t() {
13604        let (rows, cols) = (4096usize, 3072usize);
13605        let bytes = synth_q4tp(rows, cols);
13606        let twin = q4tp_as_q4t(&bytes, rows, cols);
13607        let x: Vec<f32> = (0..cols).map(|i| (i % 97) as f32 / 97.0 - 0.5).collect();
13608        let mut o = vec![0f32; rows];
13609        let n = 12;
13610        let mut best = (f64::MAX, f64::MAX);
13611        // Interleaved A/B, minimum statistic: this machine throttles, and a
13612        // mean over a thermal ramp reliably indicts whichever ran second.
13613        for _ in 0..3 {
13614            let t0 = std::time::Instant::now();
13615            for _ in 0..n {
13616                q4t_matvec(&twin, &x, rows, cols, &mut o, None);
13617            }
13618            best.0 = best.0.min(t0.elapsed().as_secs_f64());
13619            let t0 = std::time::Instant::now();
13620            for _ in 0..n {
13621                q4tp_matvec(&bytes, &x, rows, cols, &mut o, None);
13622            }
13623            best.1 = best.1.min(t0.elapsed().as_secs_f64());
13624        }
13625        let ratio = best.1 / best.0;
13626        println!(
13627            "q4t {:.3} ms | q4tp {:.3} ms | {ratio:.2}x",
13628            best.0 * 1e3 / n as f64,
13629            best.1 * 1e3 / n as f64
13630        );
13631        assert!(ratio < 2.0, "q4tp matvec {ratio:.2}x slower than q4t");
13632    }
13633
13634    #[cfg(target_os = "macos")]
13635    #[test]
13636    fn q4t_matmat_accel_matches_dequant_reference() {
13637        if !accel_gemm_enabled() {
13638            return; // CMF_ACCEL=0
13639        }
13640        let (rows, cols, b) = (512usize, 1024usize, 8usize); // ≥500K → accel arm
13641        let gpr = cols / GROUP_SIZE;
13642        let mut bytes = vec![0u8; rows * gpr * Q4_TILE];
13643        for r in 0..rows {
13644            for g in 0..gpr {
13645                let t = (r * gpr + g) * Q4_TILE;
13646                let sc = 0.02 + 0.0005 * ((r * gpr + g) % 64) as f32;
13647                bytes[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
13648                for k in 0..16 {
13649                    bytes[t + 2 + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
13650                }
13651            }
13652        }
13653        let x: Vec<f32> = (0..b * cols)
13654            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
13655            .collect();
13656        let mut got = vec![0f32; b * rows];
13657        q4t_matmat(&bytes, &x, b, rows, cols, &mut got, None);
13658        // Brute-force reference off the same tiles.
13659        let mut w = vec![0f32; rows * cols];
13660        for r in 0..rows {
13661            for g in 0..gpr {
13662                let t = (r * gpr + g) * Q4_TILE;
13663                let s = f16_to_f32(u16::from_le_bytes([bytes[t], bytes[t + 1]]));
13664                for (k, &bb) in bytes[t + 2..t + Q4_TILE].iter().enumerate() {
13665                    w[r * cols + g * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s;
13666                    w[r * cols + g * GROUP_SIZE + k * 2 + 1] =
13667                        (((bb >> 4) & 0x0F) as f32 - 8.0) * s;
13668                }
13669            }
13670        }
13671        for bi in 0..b {
13672            for r in 0..rows {
13673                let want: f32 = (0..cols).map(|j| x[bi * cols + j] * w[r * cols + j]).sum();
13674                let d = (got[bi * rows + r] - want).abs();
13675                assert!(
13676                    d <= want.abs().max(1.0) * 1e-4,
13677                    "accel q4t GEMM diverged at ({bi},{r}): {} vs {want}",
13678                    got[bi * rows + r]
13679                );
13680            }
13681        }
13682    }
13683
13684    #[test]
13685    fn q4matvec_matches_full_dequant() {
13686        let (rows, cols) = (8, 64);
13687        let groups = rows * cols / GROUP_SIZE;
13688        // Hand-craft a q4_block blob: nibbles then f16 scales.
13689        let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
13690        for i in 0..groups * 16 {
13691            bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
13692        }
13693        for g in 0..groups {
13694            let s = 0.01 + 0.003 * g as f32;
13695            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
13696        }
13697        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
13698
13699        let mut reference = vec![0.0f32; rows * cols];
13700        cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
13701        let mut expect = vec![0.0f32; rows];
13702        for r in 0..rows {
13703            expect[r] = reference[r * cols..(r + 1) * cols]
13704                .iter()
13705                .zip(&x)
13706                .map(|(w, xv)| w * xv)
13707                .sum();
13708        }
13709
13710        let mut got = vec![0.0f32; rows];
13711        q4matvec(&bytes, &x, rows, cols, &mut got, None);
13712        // SDOT path quantizes activations to i8 (A8W8): bounded noise,
13713        // same contract as q8/vbit (exact path is pinned by CMF_SDOT=0
13714        // in the golden-parity gate).
13715        let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
13716        let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
13717        for r in 0..rows {
13718            assert!(
13719                (got[r] - expect[r]).abs() < tol * scale,
13720                "row {r}: {} vs {}",
13721                got[r],
13722                expect[r]
13723            );
13724        }
13725    }
13726
13727    /// Fused two-input vbit matvec must equal two single matvecs exactly
13728    /// (same per-lane accumulation order on both scalar and SDOT paths).
13729    #[test]
13730    fn vbitmatvec2_equals_two_singles() {
13731        let (rows, cols) = (6, 64);
13732        let ng = cols / GROUP_SIZE;
13733        let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
13734        let mut bytes = bits.clone();
13735        for g in 0..rows * ng {
13736            let s = 0.02 + 0.001 * g as f32;
13737            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
13738        }
13739        for r in 0..rows {
13740            let b = bits[r] as usize;
13741            let (mut acc, mut nb) = (0u64, 0usize);
13742            let mut rowbytes = Vec::new();
13743            for i in 0..cols {
13744                let v = ((i * 7 + r * 13) % (1 << b)) as u64;
13745                acc = (acc << b) | v;
13746                nb += b;
13747                while nb >= 8 {
13748                    nb -= 8;
13749                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
13750                }
13751            }
13752            if nb > 0 {
13753                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
13754            }
13755            bytes.extend_from_slice(&rowbytes);
13756        }
13757        let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
13758        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).cos()).collect();
13759        let offsets = vbit_row_offsets(&bytes, rows, cols);
13760
13761        let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
13762        vbitmatvec(&bytes, &offsets, &x1, rows, cols, &mut a1, None);
13763        vbitmatvec(&bytes, &offsets, &x2, rows, cols, &mut a2, None);
13764        let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
13765        vbitmatvec2(
13766            &bytes, &offsets, &x1, &x2, rows, cols, &mut b1, &mut b2, None,
13767        );
13768        assert_eq!(a1, b1, "fused vbit lane 1 must be bit-identical");
13769        assert_eq!(a2, b2, "fused vbit lane 2 must be bit-identical");
13770    }
13771
13772    /// Fused two-input q4 matvec must equal two single matvecs exactly.
13773    #[test]
13774    fn q4matvec2_equals_two_singles() {
13775        let (rows, cols) = (8, 128);
13776        let groups = rows * cols / GROUP_SIZE;
13777        let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
13778        for i in 0..groups * 16 {
13779            bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
13780        }
13781        for g in 0..groups {
13782            let s = 0.01 + 0.003 * g as f32;
13783            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
13784        }
13785        // Include an outlier channel so the SDOT correction path is
13786        // exercised in the pair kernel too.
13787        let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
13788        x1[9] = 250.0;
13789        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
13790
13791        let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
13792        q4matvec(&bytes, &x1, rows, cols, &mut a1, None);
13793        q4matvec(&bytes, &x2, rows, cols, &mut a2, None);
13794        let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
13795        q4matvec2(&bytes, &x1, &x2, rows, cols, &mut b1, &mut b2, None);
13796        assert_eq!(a1, b1, "fused q4 lane 1 must be bit-identical");
13797        assert_eq!(a2, b2, "fused q4 lane 2 must be bit-identical");
13798    }
13799
13800    /// Multi-matrix job must equal separate matvecs exactly — same
13801    /// kernels, only the dispatch is fused.
13802    #[test]
13803    fn matvec_many_equals_separate_matvecs() {
13804        use crate::pool::Pool;
13805        let (r1, r2, cols) = (300, 200, 64);
13806        let mk = |salt: usize, rows: usize| {
13807            QTensor::from_f32(
13808                (0..rows * cols)
13809                    .map(|i| ((i * 7 + salt) % 97) as f32 / 97.0 - 0.5)
13810                    .collect(),
13811                rows,
13812                cols,
13813            )
13814        };
13815        let (a, b) = (mk(1, r1), mk(5, r2));
13816        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).sin()).collect();
13817        let pool = Pool::new(3);
13818
13819        let (mut ea, mut eb) = (vec![0f32; r1], vec![0f32; r2]);
13820        a.matvec(&x, &mut ea, Some(&pool));
13821        b.matvec(&x, &mut eb, Some(&pool));
13822        let (mut ga, mut gb) = (vec![0f32; r1], vec![0f32; r2]);
13823        QTensor::matvec_many([&a, &b], &x, [&mut ga, &mut gb], Some(&pool));
13824        assert_eq!(ea, ga, "fused multi-matrix lane 1 must be bit-identical");
13825        assert_eq!(eb, gb, "fused multi-matrix lane 2 must be bit-identical");
13826    }
13827
13828    /// The public Q4TP operator must take the real mapped matvec_many arm,
13829    /// rather than the F32 fallback above.  Build a tiny valid CMF so both
13830    /// handles retain their mmap payloads, then compare the fused dispatch
13831    /// with two ordinary mapped matvec calls bit-for-bit.
13832    #[test]
13833    fn q4tp_matvec_many_equals_separate_matvecs() {
13834        use crate::pool::Pool;
13835        use cortiq_core::{CMF_VERSION, CmfHeader, CmfModel, QuantType, TensorSpec};
13836
13837        let (r1, r2, cols) = (300usize, 200usize, 64usize);
13838        let arch: cortiq_core::ModelArch = serde_json::from_value(serde_json::json!({
13839            "arch_name": "tiny-q4tp",
13840            "hidden_size": cols,
13841            "intermediate_size": cols * 2,
13842            "num_layers": 1,
13843            "num_attention_heads": 2,
13844            "num_kv_heads": 1,
13845            "head_dim": 32,
13846            "vocab_size": r1,
13847            "layer_types": ["FullAttention"],
13848            "rms_norm_eps": 1e-6,
13849            "max_position_embeddings": 8,
13850            "linear_conv_kernel_dim": 0,
13851            "linear_num_key_heads": 0,
13852            "linear_num_value_heads": 0
13853        }))
13854        .unwrap();
13855        let header = CmfHeader {
13856            format: "cmf".into(),
13857            version: CMF_VERSION,
13858            arch,
13859            quant_type: QuantType::Q4Block,
13860            provenance: None,
13861            tokenizer_config: None,
13862            section_hashes: None,
13863            skills: Vec::new(),
13864            shard: None,
13865            calibration: None,
13866            routing: None,
13867            genome: None,
13868            lineage: Vec::new(),
13869            router: None,
13870            segments: Vec::new(),
13871        };
13872        let specs = [
13873            TensorSpec {
13874                name: "q".into(),
13875                dtype: TensorDtype::Q4TiledP,
13876                shape: vec![r1, cols],
13877                data: synth_q4tp(r1, cols),
13878            },
13879            TensorSpec {
13880                name: "kv".into(),
13881                dtype: TensorDtype::Q4TiledP,
13882                shape: vec![r2, cols],
13883                data: synth_q4tp(r2, cols),
13884            },
13885        ];
13886        let dir = std::env::temp_dir().join(format!("cmf-q4tp-many-{}", std::process::id()));
13887        std::fs::create_dir_all(&dir).unwrap();
13888        let path = dir.join("m.cmf");
13889        CmfModel::write(&path, &header, &specs, None, None).unwrap();
13890        let model = Arc::new(CmfModel::open(&path).unwrap());
13891        let (a, b) = (
13892            QTensor::from_model(&model, "q").unwrap(),
13893            QTensor::from_model(&model, "kv").unwrap(),
13894        );
13895        assert_eq!(a.model_dtype(), Some(TensorDtype::Q4TiledP));
13896        assert_eq!(b.model_dtype(), Some(TensorDtype::Q4TiledP));
13897        let x: Vec<f32> = (0..cols)
13898            .map(|i| ((i * 17 + 3) % 97) as f32 / 97.0 - 0.5)
13899            .collect();
13900        let pool = Pool::new(3);
13901        let (mut ea, mut eb) = (vec![0.0f32; r1], vec![0.0f32; r2]);
13902        a.matvec(&x, &mut ea, Some(&pool));
13903        b.matvec(&x, &mut eb, Some(&pool));
13904        let (mut ga, mut gb) = (vec![0.0f32; r1], vec![0.0f32; r2]);
13905        QTensor::matvec_many([&a, &b], &x, [&mut ga, &mut gb], Some(&pool));
13906        assert_eq!(ea, ga, "Q4TP fused lane 1 must be bit-identical");
13907        assert_eq!(eb, gb, "Q4TP fused lane 2 must be bit-identical");
13908        let _ = std::fs::remove_dir_all(&dir);
13909    }
13910
13911    /// The MiMo speculative verify's kernels: several tokens' MoE through
13912    /// `moe_gate_up_rows` / `moe_down_rows` (+ the caller's route-order sum)
13913    /// is bit-identical to each token alone through `moe_gate_up_many` /
13914    /// `moe_down_many` (decode), and a row-exact `q4tp_matmat` of five
13915    /// tokens (wide enough for the blocked tiles) equals five matvecs.
13916    #[test]
13917    fn multi_token_moe_rows_equal_single_token_decode() {
13918        use crate::pool::Pool;
13919        use cortiq_core::{CMF_VERSION, CmfHeader, CmfModel, QuantType, TensorSpec};
13920
13921        let (h, inter, ne) = (64usize, 128usize, 3usize);
13922        let arch: cortiq_core::ModelArch = serde_json::from_value(serde_json::json!({
13923            "arch_name": "tiny-q4tp-moe",
13924            "hidden_size": h,
13925            "intermediate_size": inter,
13926            "num_layers": 1,
13927            "num_attention_heads": 2,
13928            "num_kv_heads": 1,
13929            "head_dim": 32,
13930            "vocab_size": 8,
13931            "layer_types": ["FullAttention"],
13932            "rms_norm_eps": 1e-6,
13933            "max_position_embeddings": 8,
13934            "linear_conv_kernel_dim": 0,
13935            "linear_num_key_heads": 0,
13936            "linear_num_value_heads": 0
13937        }))
13938        .unwrap();
13939        let header = CmfHeader {
13940            format: "cmf".into(),
13941            version: CMF_VERSION,
13942            arch,
13943            quant_type: QuantType::Q4Block,
13944            provenance: None,
13945            tokenizer_config: None,
13946            section_hashes: None,
13947            skills: Vec::new(),
13948            shard: None,
13949            calibration: None,
13950            routing: None,
13951            genome: None,
13952            lineage: Vec::new(),
13953            router: None,
13954            segments: Vec::new(),
13955        };
13956        let mut specs = Vec::new();
13957        for e in 0..ne {
13958            for (k, (n, r, c)) in [("g", inter, h), ("u", inter, h), ("d", h, inter)]
13959                .into_iter()
13960                .enumerate()
13961            {
13962                // Distinct experts: perturb only the nibble plane (any byte
13963                // is a valid pair of codes; the ladder stays intact).
13964                let mut data = synth_q4tp(r, c);
13965                for (i, byte) in data[..r * (c / GROUP_SIZE) * Q4TP_NIB]
13966                    .iter_mut()
13967                    .enumerate()
13968                {
13969                    *byte ^= ((i * (e * 3 + k + 1)) % 251) as u8;
13970                }
13971                specs.push(TensorSpec {
13972                    name: format!("{n}{e}"),
13973                    dtype: TensorDtype::Q4TiledP,
13974                    shape: vec![r, c],
13975                    data,
13976                });
13977            }
13978        }
13979        let dir = std::env::temp_dir().join(format!(
13980            "cmf-moe-rows-{}-{}",
13981            std::process::id(),
13982            FLOAT_ACTIVATIONS.get()
13983        ));
13984        std::fs::create_dir_all(&dir).unwrap();
13985        let path = dir.join("m.cmf");
13986        CmfModel::write(&path, &header, &specs, None, None).unwrap();
13987        let model = Arc::new(CmfModel::open(&path).unwrap());
13988        let t = |n: String| QTensor::from_model(&model, &n).unwrap();
13989        let g: Vec<QTensor> = (0..ne).map(|e| t(format!("g{e}"))).collect();
13990        let u: Vec<QTensor> = (0..ne).map(|e| t(format!("u{e}"))).collect();
13991        let d: Vec<QTensor> = (0..ne).map(|e| t(format!("d{e}"))).collect();
13992        let b = 4usize;
13993        let mut xs: Vec<f32> = (0..b * h)
13994            .map(|i| ((i * 31 + 7) % 89) as f32 / 89.0 - 0.5)
13995            .collect();
13996        xs[5] = 9.0; // an activation outlier on token 0
13997        // Token -> (experts in route order, weights).
13998        let routes: Vec<(Vec<usize>, Vec<f32>)> = vec![
13999            (vec![2, 0], vec![0.6, 0.4]),
14000            (vec![0, 1, 2], vec![0.2, 0.5, 0.3]),
14001            (vec![1], vec![1.0]),
14002            (vec![2, 1, 0], vec![0.25, 0.25, 0.5]),
14003        ];
14004        let pool = Pool::new(3);
14005        // Decode reference, token by token.
14006        let mut want = vec![0f32; b * h];
14007        for (tk, (idx, w)) in routes.iter().enumerate() {
14008            let x = &xs[tk * h..(tk + 1) * h];
14009            let pairs: Vec<(&QTensor, &QTensor)> = idx.iter().map(|&e| (&g[e], &u[e])).collect();
14010            let mut gs: Vec<Vec<f32>> = idx.iter().map(|_| vec![0f32; inter]).collect();
14011            assert!(QTensor::moe_gate_up_many(&pairs, x, &mut gs, Some(&pool)));
14012            if FLOAT_ACTIVATIONS.get() {
14013                for (slot, &e) in idx.iter().enumerate() {
14014                    let (mut gate, mut up) = (vec![0.0; inter], vec![0.0; inter]);
14015                    g[e].matvec(x, &mut gate, Some(&pool));
14016                    u[e].matvec(x, &mut up, Some(&pool));
14017                    for (v, u) in gate.iter_mut().zip(up) {
14018                        *v = (*v / (1.0 + (-*v).exp())) * u;
14019                    }
14020                    assert_eq!(gs[slot], gate, "float gate/up must equal ordinary matvecs");
14021                }
14022            }
14023            let downs: Vec<&QTensor> = idx.iter().map(|&e| &d[e]).collect();
14024            assert!(QTensor::moe_down_many(
14025                &downs,
14026                &gs,
14027                w,
14028                &mut want[tk * h..(tk + 1) * h],
14029                Some(&pool)
14030            ));
14031        }
14032        if FLOAT_ACTIVATIONS.get() {
14033            for (tk, (idx, w)) in routes.iter().enumerate() {
14034                let mut scalar = vec![0.0; h];
14035                for (&e, &weight) in idx.iter().zip(w) {
14036                    let (mut gate, mut up, mut down) =
14037                        (vec![0.0; inter], vec![0.0; inter], vec![0.0; h]);
14038                    g[e].matvec(&xs[tk * h..(tk + 1) * h], &mut gate, Some(&pool));
14039                    u[e].matvec(&xs[tk * h..(tk + 1) * h], &mut up, Some(&pool));
14040                    for (v, u) in gate.iter_mut().zip(up) {
14041                        *v = (*v / (1.0 + (-*v).exp())) * u;
14042                    }
14043                    d[e].matvec(&gate, &mut down, Some(&pool));
14044                    for (v, d) in scalar.iter_mut().zip(down) {
14045                        *v += weight * d;
14046                    }
14047                }
14048                assert_eq!(
14049                    &want[tk * h..(tk + 1) * h],
14050                    scalar,
14051                    "float many equals scalar experts"
14052                );
14053            }
14054        }
14055        // All four tokens at once, grouped by expert.
14056        let mut experts: Vec<usize> = Vec::new();
14057        let mut groups: Vec<Vec<usize>> = Vec::new();
14058        for (tk, (idx, _)) in routes.iter().enumerate() {
14059            for &e in idx {
14060                match experts.iter().position(|&x| x == e) {
14061                    Some(k) => groups[k].push(tk),
14062                    None => {
14063                        experts.push(e);
14064                        groups.push(vec![tk]);
14065                    }
14066                }
14067            }
14068        }
14069        let n_pairs: usize = groups.iter().map(|g| g.len()).sum();
14070        let pairs: Vec<(&QTensor, &QTensor)> = experts.iter().map(|&e| (&g[e], &u[e])).collect();
14071        let mut gs: Vec<Vec<f32>> = (0..n_pairs).map(|_| vec![0f32; inter]).collect();
14072        assert!(QTensor::moe_gate_up_rows(
14073            &pairs,
14074            &groups,
14075            &xs,
14076            &mut gs,
14077            Some(&pool)
14078        ));
14079        let downs: Vec<&QTensor> = experts.iter().map(|&e| &d[e]).collect();
14080        let lens: Vec<usize> = groups.iter().map(|g| g.len()).collect();
14081        let mut ds: Vec<Vec<f32>> = (0..n_pairs).map(|_| vec![0f32; h]).collect();
14082        assert!(QTensor::moe_down_rows(
14083            &downs,
14084            &lens,
14085            &gs,
14086            &mut ds,
14087            Some(&pool)
14088        ));
14089        let slot = |tk: usize, e: usize| {
14090            let k = experts.iter().position(|&x| x == e).unwrap();
14091            groups[..k].iter().map(|g| g.len()).sum::<usize>()
14092                + groups[k].iter().position(|&x| x == tk).unwrap()
14093        };
14094        let mut got = vec![0f32; b * h];
14095        for (tk, (idx, w)) in routes.iter().enumerate() {
14096            for i in 0..h {
14097                let mut acc = 0f32;
14098                for (&e, &we) in idx.iter().zip(w) {
14099                    acc += we * ds[slot(tk, e)][i];
14100                }
14101                got[tk * h + i] = acc;
14102            }
14103        }
14104        assert!(want.iter().any(|v| *v != 0.0));
14105        assert_eq!(
14106            want.iter().map(|v| v.to_bits()).collect::<Vec<_>>(),
14107            got.iter().map(|v| v.to_bits()).collect::<Vec<_>>(),
14108            "multi-token MoE must equal decode bit for bit"
14109        );
14110
14111        // Row-exact q4tp_matmat: five tokens (a blocked 1x4 tile + tail
14112        // otherwise) equal five matvecs.
14113        let b5 = 5usize;
14114        let x5: Vec<f32> = (0..b5 * h)
14115            .map(|i| ((i * 13 + 5) % 71) as f32 / 71.0 - 0.5)
14116            .collect();
14117        let mut mm = vec![0f32; b5 * inter];
14118        row_exact_scope(|| g[1].matmat(&x5, b5, &mut mm, Some(&pool)));
14119        for tk in 0..b5 {
14120            let mut mv = vec![0f32; inter];
14121            g[1].matvec(&x5[tk * h..(tk + 1) * h], &mut mv, Some(&pool));
14122            assert_eq!(
14123                mv.iter().map(|v| v.to_bits()).collect::<Vec<_>>(),
14124                mm[tk * inter..(tk + 1) * inter]
14125                    .iter()
14126                    .map(|v| v.to_bits())
14127                    .collect::<Vec<_>>(),
14128                "row-exact matmat token {tk}"
14129            );
14130        }
14131        // Other concurrent tests/requests may still hold the shared mode.
14132        // Nested, overlapping and unwind restoration is checked separately.
14133        let _ = std::fs::remove_dir_all(&dir);
14134    }
14135
14136    /// The row-exact fix must not touch the fast path. Outside the scope
14137    /// `q4tp_matmat` has to produce exactly what it did before, and on ARM
14138    /// "before" is spelled out below: the tuned 1x4 SDOT tile for every
14139    /// four columns and the single-row kernel for the tail. Inside the
14140    /// scope every column equals its token's matvec. `q4tp_matmat_with`
14141    /// takes the mode as an argument, so a concurrent test holding the
14142    /// shared scope cannot flip it under this one.
14143    #[test]
14144    fn q4tp_matmat_fast_path_unchanged_outside_row_exact() {
14145        use crate::pool::Pool;
14146        use std::sync::atomic::Ordering::Relaxed;
14147        let _alt = Q4TP_ALT_TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner());
14148        // The tuned ARM shape, which is also what an unset switch picks
14149        // unless CMF_Q4TP_V1 is exported.
14150        Q4TP_ALT.store(2, Relaxed);
14151        let pool = Pool::new(3);
14152        // Under 500k cells, so macOS keeps the matmat off the AMX; the
14153        // second shape runs across pool workers (rows >= 256) with 32
14154        // groups of accumulation and a tail after two 1x4 tiles.
14155        for &(rows, cols, b) in &[(64usize, 256usize, 7usize), (320, 1024, 9)] {
14156            let bytes = synth_q4tp(rows, cols);
14157            let mut xs: Vec<f32> = (0..b * cols)
14158                .map(|i| ((i * 29 + 11) % 83) as f32 / 83.0 - 0.5)
14159                .collect();
14160            xs[3] = 7.5; // an activation outlier on token 0
14161            let run = |exact: bool| {
14162                let mut out = vec![0f32; b * rows];
14163                q4tp_matmat_with(&bytes, &xs, b, rows, cols, &mut out, Some(&pool), exact);
14164                out
14165            };
14166            let (fast, exact) = (run(false), run(true));
14167            #[cfg(not(target_arch = "aarch64"))]
14168            let _ = fast;
14169            let bits = |v: &[f32]| v.iter().map(|x| x.to_bits()).collect::<Vec<_>>();
14170            let mut matvecs = vec![0f32; b * rows];
14171            for (bi, o) in matvecs.chunks_mut(rows).enumerate() {
14172                q4tp_matvec(
14173                    &bytes,
14174                    &xs[bi * cols..(bi + 1) * cols],
14175                    rows,
14176                    cols,
14177                    o,
14178                    Some(&pool),
14179                );
14180            }
14181            assert!(matvecs.iter().any(|v| *v != 0.0));
14182            assert_eq!(
14183                bits(&exact),
14184                bits(&matvecs),
14185                "{rows}x{cols} b={b}: row-exact matmat must equal per-token matvecs"
14186            );
14187            #[cfg(target_arch = "aarch64")]
14188            {
14189                // The pre-fix ARM loop, cell for cell.
14190                let gpr = cols / GROUP_SIZE;
14191                let v = Q4tpView::new(&bytes, rows, cols);
14192                let mut old = vec![0f32; b * rows];
14193                let mut sc = vec![0f32; gpr];
14194                let a8w8 = a8w8_enabled();
14195                let blocked = sdot_enabled() && blocked_enabled();
14196                let acts: Vec<SplitAct> = (0..b)
14197                    .map(|bi| split_act(&xs[bi * cols..(bi + 1) * cols]))
14198                    .collect();
14199                for r in 0..rows {
14200                    v.scales_into(r, gpr, &mut sc);
14201                    if !a8w8 {
14202                        for bi in 0..b {
14203                            let x = &xs[bi * cols..(bi + 1) * cols];
14204                            old[bi * rows + r] = q4tp_row_exact(v.nib, r, gpr, x, &sc);
14205                        }
14206                        continue;
14207                    }
14208                    let finish = |d: f32, act: &SplitAct| {
14209                        let mut acc = d * act.sx;
14210                        for &(j, xv) in &act.outliers {
14211                            let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
14212                            acc += w * s * xv;
14213                        }
14214                        acc
14215                    };
14216                    let mut bi = 0usize;
14217                    while blocked && bi + 4 <= b {
14218                        let xs4 = [
14219                            acts[bi].xq.as_slice(),
14220                            acts[bi + 1].xq.as_slice(),
14221                            acts[bi + 2].xq.as_slice(),
14222                            acts[bi + 3].xq.as_slice(),
14223                        ];
14224                        let d = unsafe { dot_q4tp_row_1x4_sdot(v.nib, r, gpr, xs4, &sc) };
14225                        for k in 0..4 {
14226                            old[(bi + k) * rows + r] = finish(d[k], &acts[bi + k]);
14227                        }
14228                        bi += 4;
14229                    }
14230                    for (bi, act) in acts.iter().enumerate().skip(bi) {
14231                        let d = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, &sc);
14232                        old[bi * rows + r] = finish(d, act);
14233                    }
14234                }
14235                assert_eq!(
14236                    bits(&fast),
14237                    bits(&old),
14238                    "{rows}x{cols} b={b}: the fast path changed outside row_exact"
14239                );
14240                // And it is still the fast tile that runs: its lane-parallel
14241                // fma sum rounds differently from the matvec somewhere.
14242                if blocked {
14243                    assert_ne!(
14244                        bits(&fast),
14245                        bits(&matvecs),
14246                        "{rows}x{cols} b={b}: the tuned tile no longer runs outside row_exact"
14247                    );
14248                }
14249            }
14250        }
14251        Q4TP_ALT.store(0, Relaxed);
14252    }
14253
14254    #[test]
14255    fn row_exact_scopes_survive_overlap_nesting_and_unwind() {
14256        use std::sync::{Barrier, atomic::{AtomicUsize, Ordering}};
14257        // A private counter makes this restoration test independent of
14258        // numerical tests concurrently using the production counter.
14259        let active = AtomicUsize::new(0);
14260        counted_row_exact_scope(&active, || {
14261            assert_eq!(active.load(Ordering::Acquire), 1);
14262            counted_row_exact_scope(&active, || {
14263                assert_eq!(active.load(Ordering::Acquire), 2);
14264            });
14265            assert_eq!(active.load(Ordering::Acquire), 1);
14266        });
14267        assert_eq!(active.load(Ordering::Acquire), 0);
14268
14269        let both_entered = Barrier::new(2);
14270        let release_last = Barrier::new(2);
14271        std::thread::scope(|s| {
14272            let first = s.spawn(|| counted_row_exact_scope(&active, || {
14273                both_entered.wait();
14274            }));
14275            let last = s.spawn(|| counted_row_exact_scope(&active, || {
14276                both_entered.wait();
14277                release_last.wait();
14278            }));
14279            first.join().unwrap();
14280            let after_first = active.load(Ordering::Acquire);
14281            release_last.wait();
14282            last.join().unwrap();
14283            assert_eq!(after_first, 1, "second request must remain exact");
14284        });
14285        assert_eq!(active.load(Ordering::Acquire), 0);
14286        let panic = std::panic::catch_unwind(|| {
14287            counted_row_exact_scope(&active, || panic!("scope unwind"));
14288        });
14289        assert!(panic.is_err());
14290        assert_eq!(active.load(Ordering::Acquire), 0);
14291    }
14292
14293    #[test]
14294    #[cfg(target_arch = "x86_64")]
14295    fn q4tp_float_avx2_is_bitwise_scalar() {
14296        if !avx2_enabled() {
14297            return;
14298        }
14299        for cols in [32, 64, 96, 2048, 4096] {
14300            let rows = 9;
14301            let bytes = synth_q4tp(rows, cols);
14302            let v = Q4tpView::new(&bytes, rows, cols);
14303            let gpr = cols / GROUP_SIZE;
14304            let mut sc = vec![0.0; gpr];
14305            for seed in 1..=5 {
14306                let xs: Vec<f32> = (0..cols)
14307                    .map(|i| (((i * 104729 + seed * 8191) % 100003) as f32 - 50001.0) / 7919.0)
14308                    .collect();
14309                for r in 0..rows {
14310                    v.scales_into(r, gpr, &mut sc);
14311                    let scalar = q4tp_row_float_scalar(v.nib, r, gpr, &xs, &sc);
14312                    let vector = unsafe { q4tp_row_float_avx2(v.nib, r, gpr, &xs, &sc) };
14313                    assert_eq!(
14314                        scalar.to_bits(),
14315                        vector.to_bits(),
14316                        "cols={cols} row={r} seed={seed}"
14317                    );
14318                }
14319            }
14320        }
14321    }
14322
14323    #[test]
14324    fn multi_token_moe_rows_float_equal_single_token_decode() {
14325        float_activations_scope(multi_token_moe_rows_equal_single_token_decode);
14326    }
14327
14328    #[test]
14329    fn full_gpu_q8_scope_is_nested_and_thread_local() {
14330        assert!(!FULL_GPU_Q8.get());
14331        let before = gpu_split_frac();
14332        {
14333            let _guard = enter_full_gpu_q8_scope();
14334            assert_eq!(gpu_split_frac(), 1.0);
14335            {
14336                let _nested = enter_full_gpu_q8_scope();
14337            }
14338            assert_eq!(gpu_split_frac(), 1.0);
14339            std::thread::spawn(|| assert!(!FULL_GPU_Q8.get())).join().unwrap();
14340        }
14341        assert!(!FULL_GPU_Q8.get());
14342        assert_eq!(gpu_split_frac(), before);
14343    }
14344
14345    #[test]
14346    fn float_activation_scope_is_nested_thread_local_and_unwind_safe() {
14347        assert!(!FLOAT_ACTIVATIONS.get());
14348        let before = a8w8_enabled();
14349        float_activations_scope(|| {
14350            assert!(!a8w8_enabled());
14351            float_activations_scope(|| assert!(!a8w8_enabled()));
14352            assert!(FLOAT_ACTIVATIONS.get());
14353            std::thread::spawn(|| assert!(!FLOAT_ACTIVATIONS.get()))
14354                .join()
14355                .unwrap();
14356        });
14357        assert!(!FLOAT_ACTIVATIONS.get());
14358        assert_eq!(a8w8_enabled(), before);
14359        let _ = std::panic::catch_unwind(|| float_activations_scope(|| panic!("test unwind")));
14360        assert!(!FLOAT_ACTIVATIONS.get());
14361    }
14362
14363    /// Batched q4/vbit matmat must equal per-position matvec calls
14364    /// exactly (the fallback it replaced) — same kernels, same order.
14365    #[test]
14366    fn batched_matmat_equals_per_position_matvec() {
14367        let (rows, cols, b) = (8, 64, 5);
14368        // q4 blob.
14369        let groups = rows * cols / GROUP_SIZE;
14370        let mut q4 = Vec::new();
14371        for i in 0..groups * 16 {
14372            q4.push((((i * 7 + 3) % 256) & 0xFF) as u8);
14373        }
14374        for g in 0..groups {
14375            q4.extend_from_slice(
14376                &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
14377            );
14378        }
14379        // vbit blob (mixed widths incl. 8).
14380        let ng = cols / GROUP_SIZE;
14381        let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4, 5, 3];
14382        let mut vb = bits.clone();
14383        for g in 0..rows * ng {
14384            vb.extend_from_slice(
14385                &cortiq_core::quant::f32_to_f16(0.02 + 0.001 * g as f32).to_le_bytes(),
14386            );
14387        }
14388        for r in 0..rows {
14389            let bw = bits[r] as usize;
14390            let (mut acc, mut nb) = (0u64, 0usize);
14391            let mut rowbytes = Vec::new();
14392            for i in 0..cols {
14393                let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
14394                acc = (acc << bw) | v;
14395                nb += bw;
14396                while nb >= 8 {
14397                    nb -= 8;
14398                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
14399                }
14400            }
14401            if nb > 0 {
14402                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
14403            }
14404            vb.extend_from_slice(&rowbytes);
14405        }
14406        let offsets = vbit_row_offsets(&vb, rows, cols);
14407
14408        let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
14409
14410        // q4: batch vs singles.
14411        let mut got = vec![0f32; b * rows];
14412        q4matmat(&q4, &xs, b, rows, cols, &mut got, None);
14413        for bi in 0..b {
14414            let mut expect = vec![0f32; rows];
14415            q4matvec(
14416                &q4,
14417                &xs[bi * cols..(bi + 1) * cols],
14418                rows,
14419                cols,
14420                &mut expect,
14421                None,
14422            );
14423            assert_eq!(
14424                &got[bi * rows..(bi + 1) * rows],
14425                &expect[..],
14426                "q4 batch pos {bi}"
14427            );
14428        }
14429
14430        // vbit: batch vs singles.
14431        let mut got = vec![0f32; b * rows];
14432        vbitmatmat(&vb, &offsets, &xs, b, rows, cols, &mut got, None);
14433        for bi in 0..b {
14434            let mut expect = vec![0f32; rows];
14435            vbitmatvec(
14436                &vb,
14437                &offsets,
14438                &xs[bi * cols..(bi + 1) * cols],
14439                rows,
14440                cols,
14441                &mut expect,
14442                None,
14443            );
14444            assert_eq!(
14445                &got[bi * rows..(bi + 1) * rows],
14446                &expect[..],
14447                "vbit batch pos {bi}"
14448            );
14449        }
14450    }
14451
14452    /// q4_tiled kernels must produce BIT-identical outputs to the q4
14453    /// split kernels on the same values (same ints, same order — only
14454    /// the byte placement differs).
14455    #[test]
14456    fn q4_tiled_matches_q4_block_bitexact() {
14457        let (rows, cols, b) = (8usize, 128usize, 3usize);
14458        let groups = rows * cols / GROUP_SIZE;
14459        let mut split = Vec::with_capacity(groups * 18);
14460        for i in 0..groups * 16 {
14461            split.push((((i * 7 + 3) % 256) & 0xFF) as u8);
14462        }
14463        for g in 0..groups {
14464            split.extend_from_slice(
14465                &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
14466            );
14467        }
14468        // Re-tile: [scale][nibbles] per group.
14469        let (packed, scales) = split.split_at(groups * 16);
14470        let mut tiled = Vec::with_capacity(groups * Q4_TILE);
14471        for g in 0..groups {
14472            tiled.extend_from_slice(&scales[g * 2..g * 2 + 2]);
14473            tiled.extend_from_slice(&packed[g * 16..(g + 1) * 16]);
14474        }
14475
14476        let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
14477        x1[9] = 250.0; // exercise the outlier path
14478        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
14479
14480        let (mut a, mut t) = (vec![0f32; rows], vec![0f32; rows]);
14481        q4matvec(&split, &x1, rows, cols, &mut a, None);
14482        q4t_matvec(&tiled, &x1, rows, cols, &mut t, None);
14483        assert_eq!(a, t, "q4t matvec must match q4 bit-for-bit");
14484
14485        let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
14486        let (mut t1, mut t2) = (vec![0f32; rows], vec![0f32; rows]);
14487        q4matvec2(&split, &x1, &x2, rows, cols, &mut a1, &mut a2, None);
14488        q4t_matvec2(&tiled, &x1, &x2, rows, cols, &mut t1, &mut t2, None);
14489        assert_eq!(a1, t1);
14490        assert_eq!(a2, t2);
14491
14492        let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
14493        let (mut am, mut tm) = (vec![0f32; b * rows], vec![0f32; b * rows]);
14494        q4matmat(&split, &xs, b, rows, cols, &mut am, None);
14495        q4t_matmat(&tiled, &xs, b, rows, cols, &mut tm, None);
14496        assert_eq!(am, tm, "q4t matmat must match q4 bit-for-bit");
14497    }
14498
14499    /// q4 SDOT outlier correction: a single huge activation channel
14500    /// (>8·rms → outlier, zeroed in xq) must still contribute its EXACT
14501    /// term. On-grid bulk (±1/0 → xq dequantizes exactly) isolates the
14502    /// correction from A8W8 noise. cols must exceed 64: at n=64 the
14503    /// 8·rms threshold equals sqrt(v²+rest) ≥ v, so a single outlier
14504    /// can never qualify (8² = n).
14505    #[test]
14506    fn q4matvec_sdot_outlier_exact() {
14507        let (rows, cols) = (4, 128);
14508        let groups = rows * cols / GROUP_SIZE;
14509        let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
14510        for i in 0..groups * 16 {
14511            bytes.push(((i * 11 + 5) % 256) as u8);
14512        }
14513        for g in 0..groups {
14514            let s = 0.02 + 0.002 * g as f32;
14515            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
14516        }
14517        let mut x: Vec<f32> = (0..cols)
14518            .map(|i| match i % 3 {
14519                0 => 1.0,
14520                1 => -1.0,
14521                _ => 0.0,
14522            })
14523            .collect();
14524        x[17] = 300.0; // ≫ 8·rms → outlier channel
14525
14526        let mut reference = vec![0.0f32; rows * cols];
14527        cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
14528        let mut expect = vec![0.0f32; rows];
14529        for r in 0..rows {
14530            expect[r] = reference[r * cols..(r + 1) * cols]
14531                .iter()
14532                .zip(&x)
14533                .map(|(w, xv)| w * xv)
14534                .sum();
14535        }
14536        let mut got = vec![0.0f32; rows];
14537        q4matvec(&bytes, &x, rows, cols, &mut got, None);
14538        let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
14539        for r in 0..rows {
14540            assert!(
14541                (got[r] - expect[r]).abs() < 2e-3 * scale,
14542                "row {r}: {} vs {} (outlier term must be exact)",
14543                got[r],
14544                expect[r]
14545            );
14546        }
14547    }
14548
14549    /// The fused q1t matvec must equal the reference (dequant_q1t → dot),
14550    /// including the ternary zero level and the binary-searched outlier
14551    /// overlay. Guards the mmap kernel that makes a 12B q1t runnable.
14552    #[test]
14553    fn q1t_matvec_matches_reference() {
14554        use cortiq_core::quant::{dequant_q1t, f32_to_f16};
14555        let (rows, cols) = (3usize, 64usize); // gpr = 2
14556        let gpr = cols / GROUP_SIZE;
14557        let scales = [0.5f32, 0.3, 0.7, 0.2, 0.6, 0.15];
14558        // Overlay (must be sorted by flat index): a few spikes across rows.
14559        let outliers: [(u32, f32); 3] = [(5, 9.0), (70, -4.5), (150, 3.25)];
14560        let is_out = |flat: usize| outliers.iter().any(|&(i, _)| i as usize == flat);
14561        let mut bytes = Vec::new();
14562        for r in 0..rows {
14563            for g in 0..gpr {
14564                bytes.extend_from_slice(&f32_to_f16(scales[r * gpr + g]).to_le_bytes());
14565                let mut c = [0u8; 7];
14566                for k in 0..GROUP_SIZE {
14567                    // Encoder invariant: code 0 at outlier positions.
14568                    let code = if is_out(r * cols + g * GROUP_SIZE + k) {
14569                        0
14570                    } else {
14571                        ((k + r * 3 + g) % 3) as u8 // 0,1,2
14572                    };
14573                    cortiq_core::quant::q1t_pack(&mut c, k, code);
14574                }
14575                bytes.extend_from_slice(&c);
14576            }
14577        }
14578        // Per-row overlay: [u32 row_ptr[rows+1]] then [(u16 col, f16 val)] by
14579        // row (outliers are sorted by flat index → already grouped by row).
14580        let mut row_ptr = vec![0u32; rows + 1];
14581        for &(idx, _) in &outliers {
14582            row_ptr[idx as usize / cols + 1] += 1;
14583        }
14584        for r in 0..rows {
14585            row_ptr[r + 1] += row_ptr[r];
14586        }
14587        for &p in &row_ptr {
14588            bytes.extend_from_slice(&p.to_le_bytes());
14589        }
14590        for &(idx, v) in &outliers {
14591            bytes.extend_from_slice(&((idx as usize % cols) as u16).to_le_bytes());
14592            bytes.extend_from_slice(&f32_to_f16(v).to_le_bytes());
14593        }
14594
14595        let mut refw = vec![0f32; rows * cols];
14596        dequant_q1t(&bytes, rows, cols, &mut refw);
14597        // On-grid activations (±1, amax 1) so the int8 SDOT path reconstructs
14598        // x exactly and matches the f32 reference (same trick as the q1 test).
14599        let x: Vec<f32> = (0..cols)
14600            .map(|j| if j % 3 == 0 { 1.0 } else { -1.0 })
14601            .collect();
14602        let mut expect = vec![0f32; rows];
14603        for r in 0..rows {
14604            let mut a = 0.0f32;
14605            for j in 0..cols {
14606                a += refw[r * cols + j] * x[j];
14607            }
14608            expect[r] = a;
14609        }
14610        let tol = |e: f32| 1e-3 * e.abs().max(1e-3);
14611        let mut got = vec![0f32; rows];
14612        q1t_matvec(&bytes, &x, rows, cols, &mut got, None);
14613        for r in 0..rows {
14614            assert!(
14615                (got[r] - expect[r]).abs() < tol(expect[r]),
14616                "row {r}: {} vs {}",
14617                got[r],
14618                expect[r]
14619            );
14620        }
14621        // matmat (b=2, f32 decode path) must agree too.
14622        let x2: Vec<f32> = x.iter().chain(x.iter()).copied().collect();
14623        let mut gm = vec![0f32; 2 * rows];
14624        q1t_matmat(&bytes, &x2, 2, rows, cols, &mut gm, None);
14625        for r in 0..rows {
14626            assert!((gm[r] - expect[r]).abs() < tol(expect[r]));
14627            assert!((gm[rows + r] - expect[r]).abs() < tol(expect[r]));
14628        }
14629        // Fused pair (q1t_matvec2) must equal two single matvecs
14630        // bit-for-bit: same unpack, same group order, same f32
14631        // accumulation per stream. Distinct x2 exercises both lanes.
14632        let xb: Vec<f32> = (0..cols)
14633            .map(|j| if j % 5 == 0 { -1.0 } else { 1.0 })
14634            .collect();
14635        let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
14636        q1t_matvec(&bytes, &x, rows, cols, &mut s1, None);
14637        q1t_matvec(&bytes, &xb, rows, cols, &mut s2, None);
14638        let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
14639        q1t_matvec2(&bytes, &x, &xb, rows, cols, &mut p1, &mut p2, None);
14640        assert_eq!(p1, s1, "q1t pair lane 1 ≠ single matvec");
14641        assert_eq!(p2, s2, "q1t pair lane 2 ≠ single matvec");
14642    }
14643
14644    /// Pair == 2×matvec with an ODD group count (the kernel's tail
14645    /// group) and no overlay section.
14646    #[test]
14647    fn q1t_matvec2_odd_gpr_matches_singles() {
14648        use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_pack};
14649        let (rows, cols) = (5usize, 96usize); // gpr = 3 → paired + tail
14650        let gpr = cols / GROUP_SIZE;
14651        let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE);
14652        for r in 0..rows {
14653            for g in 0..gpr {
14654                bytes.extend_from_slice(&f32_to_f16(0.1 + 0.05 * (r + g) as f32).to_le_bytes());
14655                let mut c = [0u8; 7];
14656                for k in 0..GROUP_SIZE {
14657                    q1t_pack(&mut c, k, ((k * 7 + r * 5 + g * 3) % 3) as u8);
14658                }
14659                bytes.extend_from_slice(&c);
14660            }
14661        }
14662        let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.31).sin()).collect();
14663        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).cos()).collect();
14664        let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
14665        q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
14666        q1t_matvec(&bytes, &x2, rows, cols, &mut s2, None);
14667        let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
14668        q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
14669        assert_eq!(p1, s1, "odd-gpr pair lane 1 ≠ single");
14670        assert_eq!(p2, s2, "odd-gpr pair lane 2 ≠ single");
14671    }
14672
14673    // Speed A/B: fused pair (one unpack, two streams) vs two single
14674    // matvecs. Single-threaded, FFN-sized, min-of paired in-process.
14675    //   cargo test -p cortiq-engine --release q1t_matvec2_speed -- --ignored --nocapture
14676    #[test]
14677    #[ignore]
14678    fn q1t_matvec2_speed() {
14679        use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_pack};
14680        use std::time::Instant;
14681        let (rows, cols) = (8192usize, 4096usize);
14682        let gpr = cols / GROUP_SIZE;
14683        let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE);
14684        for r in 0..rows {
14685            for g in 0..gpr {
14686                let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
14687                bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
14688                let mut c = [0u8; 7];
14689                for k in 0..GROUP_SIZE {
14690                    q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
14691                }
14692                bytes.extend_from_slice(&c);
14693            }
14694        }
14695        let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.31).sin()).collect();
14696        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).cos()).collect();
14697        let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
14698        let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
14699        // Warm both paths once.
14700        q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
14701        q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
14702        let (mut t_pair, mut t_two) = (f64::MAX, f64::MAX);
14703        for _ in 0..8 {
14704            let t0 = Instant::now();
14705            q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
14706            t_pair = t_pair.min(t0.elapsed().as_secs_f64() * 1000.0);
14707            let t1 = Instant::now();
14708            q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
14709            q1t_matvec(&bytes, &x2, rows, cols, &mut s2, None);
14710            t_two = t_two.min(t1.elapsed().as_secs_f64() * 1000.0);
14711        }
14712        assert_eq!(p1, s1);
14713        assert_eq!(p2, s2);
14714        println!("q1t pair {rows}x{cols}: fused {t_pair:.2} ms | two singles {t_two:.2} ms");
14715    }
14716
14717    // Speed A/B: the base-3-division decode (what the packing commit left in
14718    // place) vs the fused sign-LUT matvec. Both single-threaded, same bytes.
14719    //   cargo test -p cortiq-engine q1t_matvec_speed -- --ignored --nocapture
14720    #[test]
14721    #[ignore]
14722    fn q1t_matvec_speed() {
14723        use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_code, q1t_pack};
14724        use std::time::Instant;
14725        let (rows, cols) = (8192usize, 4096usize); // FFN-sized
14726        let gpr = cols / GROUP_SIZE;
14727        let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE + 16);
14728        for r in 0..rows {
14729            for g in 0..gpr {
14730                let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
14731                bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
14732                let mut c = [0u8; 7];
14733                for k in 0..GROUP_SIZE {
14734                    q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
14735                }
14736                bytes.extend_from_slice(&c);
14737            }
14738        }
14739        let (n, stride) = (rows * cols, 40usize); // ~2.5% outliers, per-row overlay
14740        let mut row_ptr = vec![0u32; rows + 1];
14741        let mut idx = 0usize;
14742        while idx < n {
14743            row_ptr[idx / cols + 1] += 1;
14744            idx += stride;
14745        }
14746        for r in 0..rows {
14747            row_ptr[r + 1] += row_ptr[r];
14748        }
14749        for &p in &row_ptr {
14750            bytes.extend_from_slice(&p.to_le_bytes());
14751        }
14752        let mut idx = 0usize;
14753        while idx < n {
14754            bytes.extend_from_slice(&((idx % cols) as u16).to_le_bytes());
14755            bytes.extend_from_slice(&f32_to_f16((idx % 13) as f32 * 0.1 - 0.6).to_le_bytes());
14756            idx += stride;
14757        }
14758        // On-grid ±1 so the fast path's int8 SDOT is exact vs the f32 "slow"
14759        // reference (the A/B is a timing check; values must still agree).
14760        let x: Vec<f32> = (0..cols)
14761            .map(|j| if j % 3 == 0 { 1.0 } else { -1.0 })
14762            .collect();
14763        let (rp_off, ent_off, has_ov) = q1t_overlay(&bytes, rows * gpr * Q1T_TILE, rows);
14764
14765        // "before": base-3 division decode into a buffer, then dot.
14766        let slow = |out: &mut [f32]| {
14767            let mut buf = vec![0f32; cols];
14768            for r in 0..rows {
14769                for g in 0..gpr {
14770                    let off = (r * gpr + g) * Q1T_TILE;
14771                    let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
14772                    let codes = &bytes[off + 2..off + Q1T_TILE];
14773                    for k in 0..GROUP_SIZE {
14774                        buf[g * GROUP_SIZE + k] = match q1t_code(codes, k) {
14775                            1 => s,
14776                            2 => -s,
14777                            _ => 0.0,
14778                        };
14779                    }
14780                }
14781                out[r] = q1t_row_outlier_correction(&bytes, r, rp_off, ent_off, has_ov, &x)
14782                    + (0..cols).map(|j| buf[j] * x[j]).sum::<f32>();
14783            }
14784        };
14785        let iters = 5;
14786        let mut a = vec![0f32; rows];
14787        slow(&mut a); // warm
14788        let t = Instant::now();
14789        for _ in 0..iters {
14790            slow(&mut a);
14791        }
14792        let slow_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
14793
14794        let mut b = vec![0f32; rows];
14795        q1t_matvec(&bytes, &x, rows, cols, &mut b, None); // warm
14796        let t = Instant::now();
14797        for _ in 0..iters {
14798            q1t_matvec(&bytes, &x, rows, cols, &mut b, None);
14799        }
14800        let fast_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
14801
14802        for r in 0..rows {
14803            assert!((a[r] - b[r]).abs() < 1e-2, "mismatch row {r}");
14804        }
14805        println!(
14806            "q1t matvec {rows}x{cols} (1 thread): div-decode {slow_ms:.2} ms  fused-LUT {fast_ms:.2} ms  => {:.2}x",
14807            slow_ms / fast_ms
14808        );
14809    }
14810}
14811
14812#[cfg(test)]
14813mod gemm_bench {
14814    /// `cargo test -p cortiq-engine --release q4tp_matmat_throughput -- --ignored --nocapture`
14815    /// Times the batched q4tp GEMM at the shapes the image DiT runs
14816    /// (b=296 tokens, 2304 -> 9216), on synthetic bytes: no model, no
14817    /// mmap, no thermal drift over minutes — a kernel change shows up
14818    /// here in seconds where a full render hides it in noise.
14819    ///
14820    /// On macOS add `CMF_ACCEL=0`: this shape is over the 500k-cell mark
14821    /// where the matmat hands off to Accelerate's dequant sgemm, and
14822    /// without the opt-out both rows below measure the AMX, not the
14823    /// kernel under test.
14824    #[test]
14825    #[ignore]
14826    fn q4tp_matmat_throughput() {
14827        let _alt = super::Q4TP_ALT_TEST_LOCK
14828            .lock()
14829            .unwrap_or_else(|e| e.into_inner());
14830        // 296 is a prompt-encode batch; the image DiT runs 2085 at
14831        // 512x512, where the activation panel stops fitting L2 and the
14832        // loop's shape starts to matter more than its instructions.
14833        let b: usize = std::env::var("CMF_BENCH_B")
14834            .ok()
14835            .and_then(|v| v.parse().ok())
14836            .unwrap_or(296);
14837        let (rows, cols) = (9216usize, 2304usize);
14838        let (_, _, _) = (rows, cols, b);
14839        let total =
14840            cortiq_core::quant::expected_nbytes(cortiq_core::TensorDtype::Q4TiledP, &[rows, cols])
14841                .unwrap();
14842        // Random nibbles are fine, but the row params are f16 (lo, step)
14843        // of a geometric ladder: garbage there gives exp2 of a huge
14844        // exponent, the scales come back inf, and the whole bench times
14845        // NaN arithmetic instead of the kernel.
14846        let (params_off, codes_off, _) = cortiq_core::quant::q4tp_sections(rows, cols);
14847        let mut bytes: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
14848        let lo = cortiq_core::quant::f32_to_f16(-4.0);
14849        let step = cortiq_core::quant::f32_to_f16(0.1);
14850        for r in 0..rows {
14851            let o = params_off + r * 4;
14852            bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
14853            bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
14854        }
14855        let _ = codes_off;
14856        let xs: Vec<f32> = (0..b * cols)
14857            .map(|i| ((i % 97) as f32 - 48.0) / 48.0)
14858            .collect();
14859        let mut out = vec![0f32; b * rows];
14860        let pool = crate::pool::Pool::from_env();
14861        // A shared 48-core stand drifts ±25% run to run, which is wider
14862        // than any kernel change worth making. So: alternate the two
14863        // kernels inside one process and keep the BEST time for
14864        // each. Interleaving makes both see the same interference, and a
14865        // minimum is the one statistic another tenant cannot inflate.
14866        super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
14867        let reps: usize = std::env::var("CMF_BENCH_REPS")
14868            .ok()
14869            .and_then(|v| v.parse().ok())
14870            .unwrap_or(10);
14871        let mut best = [f64::MAX; 2];
14872        let mut sums = [0f32; 2];
14873        for _ in 0..reps {
14874            for (k, w) in [(0usize, 1u8), (1usize, 2u8)] {
14875                super::Q4TP_ALT.store(w, std::sync::atomic::Ordering::Relaxed);
14876                let t = std::time::Instant::now();
14877                super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
14878                best[k] = best[k].min(t.elapsed().as_secs_f64());
14879                sums[k] = out.iter().take(64).sum::<f32>();
14880            }
14881        }
14882        let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
14883        for (k, name) in ["previous", "tuned   "].iter().enumerate() {
14884            println!(
14885                "q4tp matmat {rows}x{cols} b={b} {name}: {:.1} ms  {:.1} GFLOP/s  (checksum {:.3})",
14886                best[k] * 1e3,
14887                flops / best[k] / 1e9,
14888                sums[k]
14889            );
14890        }
14891        assert!(
14892            (sums[0] - sums[1]).abs() < 1e-2,
14893            "the tuned kernel changed the result: {} vs {}",
14894            sums[0],
14895            sums[1]
14896        );
14897    }
14898
14899    /// The blocked kernel must agree with the per-column path exactly —
14900    /// same weights, same activation split, only a different instruction
14901    /// mix. Shapes are chosen to hit the awkward cases: a column count
14902    /// that leaves an odd group (the 512-bit kernel does two at a time),
14903    /// and a batch that does not divide by four.
14904    #[test]
14905    fn q4tp_matmat_blocked_matches_scalar() {
14906        use std::sync::atomic::Ordering::Relaxed;
14907        let _alt = super::Q4TP_ALT_TEST_LOCK
14908            .lock()
14909            .unwrap_or_else(|e| e.into_inner());
14910        // The last shape carries the image DiT's column count — 2304, so
14911        // 72 groups of accumulation, which is where a reordered sum can
14912        // actually drift — and runs through the thread pool, since the
14913        // blocked path splits rows across workers. Its row count stays
14914        // under 500k cells on purpose: above that, macOS diverts the whole
14915        // matmat to the Accelerate/AMX dequant sgemm and neither kernel
14916        // here would run.
14917        for &(rows, cols, b) in &[
14918            (64usize, 128usize, 7usize),
14919            (33, 96, 4),
14920            (16, 256, 9),
14921            (192, 2304, 37),
14922        ] {
14923            let total = cortiq_core::quant::expected_nbytes(
14924                cortiq_core::TensorDtype::Q4TiledP,
14925                &[rows, cols],
14926            )
14927            .unwrap();
14928            let (params_off, _, _) = cortiq_core::quant::q4tp_sections(rows, cols);
14929            let mut bytes: Vec<u8> = (0..total).map(|i| (i * 61 % 251) as u8).collect();
14930            let lo = cortiq_core::quant::f32_to_f16(-4.0);
14931            let step = cortiq_core::quant::f32_to_f16(0.1);
14932            for r in 0..rows {
14933                let o = params_off + r * 4;
14934                bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
14935                bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
14936            }
14937            let xs: Vec<f32> = (0..b * cols)
14938                .map(|i| ((i % 89) as f32 - 44.0) / 44.0)
14939                .collect();
14940            let mut got = vec![0f32; b * rows];
14941            let mut want = vec![0f32; b * rows];
14942            let gpr = cols / 32;
14943            let view = super::Q4tpView::new(&bytes, rows, cols);
14944            let pool = crate::pool::Pool::from_env();
14945            super::Q4TP_ALT.store(2, Relaxed);
14946            super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut got, pool.as_deref());
14947            super::Q4TP_ALT.store(1, Relaxed);
14948            super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut want, pool.as_deref());
14949            super::Q4TP_ALT.store(0, Relaxed);
14950            // Measured against the output's scale, not cell by cell: a
14951            // dot product of 2304 terms lands near zero wherever the row
14952            // and the activation nearly cancel, and there a per-cell
14953            // ratio reports 1e-3 for an absolute error of 5e-6 — f32's
14954            // own rounding, reordered. What must stay small is the error
14955            // relative to what the layer actually outputs.
14956            let scale = want.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
14957            let (mut worst, mut at) = (0f32, 0usize);
14958            for (i, (g, w)) in got.iter().zip(&want).enumerate() {
14959                if (g - w).abs() > worst {
14960                    worst = (g - w).abs();
14961                    at = i;
14962                }
14963            }
14964            assert!(
14965                worst <= 1e-4 * scale,
14966                "{rows}x{cols} b={b}: blocked and scalar disagree by {worst:.3e} \
14967                 (scale {scale:.3e}) at cell {at}: {} vs {}",
14968                got[at],
14969                want[at]
14970            );
14971
14972            // "Same speed, no quality loss" is a claim about which answer
14973            // is RIGHT, not about which two agree. Both paths sum the same
14974            // 2304 products in different orders, so f64 decides: the
14975            // blocked kernel keeps sixteen partial sums and folds them at
14976            // the end, which is a shallower addition tree than the
14977            // per-column path's running scalar, and it must not be worse.
14978            let (mut e_blocked, mut e_scalar) = (0f64, 0f64);
14979            for bi in 0..b {
14980                let act = super::split_act(&xs[bi * cols..(bi + 1) * cols]);
14981                for r in 0..rows {
14982                    let mut sc = vec![0f32; gpr];
14983                    view.scales_into(r, gpr, &mut sc);
14984                    let mut exact = 0f64;
14985                    for j in 0..cols {
14986                        let (w, sq) = super::q4tp_outlier(view.nib, r, gpr, j, &sc);
14987                        exact += w as f64 * sq as f64 * act.xq[j] as f64;
14988                    }
14989                    exact *= act.sx as f64;
14990                    for &(j, xv) in &act.outliers {
14991                        let (w, sq) = super::q4tp_outlier(view.nib, r, gpr, j, &sc);
14992                        exact += w as f64 * sq as f64 * xv as f64;
14993                    }
14994                    let i = bi * rows + r;
14995                    e_blocked = e_blocked.max((got[i] as f64 - exact).abs());
14996                    e_scalar = e_scalar.max((want[i] as f64 - exact).abs());
14997                }
14998            }
14999            println!(
15000                "{rows}x{cols} b={b}: worst error vs f64 — blocked {e_blocked:.3e}, \
15001                 per-column {e_scalar:.3e}"
15002            );
15003            // An absolute bar, not a race between the two: at these
15004            // magnitudes both sit in f32's last bits, and on a small shape
15005            // whichever one happens to round the unluckiest cell "wins" by
15006            // a factor the next seed reverses.
15007            assert!(
15008                e_blocked <= 1e-5 * scale as f64 && e_scalar <= 1e-5 * scale as f64,
15009                "{rows}x{cols} b={b}: error against f64 too large — blocked \
15010                 {e_blocked:.3e}, per-column {e_scalar:.3e}, scale {scale:.3e}"
15011            );
15012        }
15013    }
15014
15015    /// What the row-exact contract costs a speculative-verify panel on the
15016    /// host: the same batch through the fast arms, through the row-exact
15017    /// arms, and as one matvec per token (the other way to be exact).
15018    /// Arms alternate inside one process and keep their best time.
15019    /// `CMF_GPU=0 cargo test -p cortiq-engine --release q4tp_matmat_row_exact_cost -- --ignored --nocapture`
15020    #[test]
15021    #[ignore]
15022    fn q4tp_matmat_row_exact_cost() {
15023        let _alt = super::Q4TP_ALT_TEST_LOCK
15024            .lock()
15025            .unwrap_or_else(|e| e.into_inner());
15026        let pool = crate::pool::Pool::from_env();
15027        let reps: usize = std::env::var("CMF_BENCH_REPS")
15028            .ok()
15029            .and_then(|v| v.parse().ok())
15030            .unwrap_or(30);
15031        for &(rows, cols) in &[(2048usize, 4096usize), (4096, 2048), (4096, 4096)] {
15032            let total = cortiq_core::quant::expected_nbytes(
15033                cortiq_core::TensorDtype::Q4TiledP,
15034                &[rows, cols],
15035            )
15036            .unwrap();
15037            let (params_off, _, _) = cortiq_core::quant::q4tp_sections(rows, cols);
15038            let mut bytes: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
15039            let lo = cortiq_core::quant::f32_to_f16(-4.0);
15040            let step = cortiq_core::quant::f32_to_f16(0.1);
15041            for r in 0..rows {
15042                let o = params_off + r * 4;
15043                bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
15044                bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
15045            }
15046            for &b in &[2usize, 4, 5, 8] {
15047                let xs: Vec<f32> = (0..b * cols)
15048                    .map(|i| ((i % 97) as f32 - 48.0) / 48.0)
15049                    .collect();
15050                let mut out = vec![0f32; b * rows];
15051                let mut best = [f64::MAX; 3];
15052                for _ in 0..reps {
15053                    for (k, best_k) in best.iter_mut().enumerate() {
15054                        let t = std::time::Instant::now();
15055                        match k {
15056                            0 | 1 => super::q4tp_matmat_with(
15057                                &bytes,
15058                                &xs,
15059                                b,
15060                                rows,
15061                                cols,
15062                                &mut out,
15063                                pool.as_deref(),
15064                                k == 1,
15065                            ),
15066                            _ => {
15067                                for (bi, o) in out.chunks_mut(rows).enumerate() {
15068                                    super::q4tp_matvec(
15069                                        &bytes,
15070                                        &xs[bi * cols..(bi + 1) * cols],
15071                                        rows,
15072                                        cols,
15073                                        o,
15074                                        pool.as_deref(),
15075                                    );
15076                                }
15077                            }
15078                        }
15079                        *best_k = best_k.min(t.elapsed().as_secs_f64());
15080                    }
15081                }
15082                println!(
15083                    "q4tp {rows}x{cols} b={b}: fast {:.3} ms, row-exact {:.3} ms, \
15084                     {b} matvecs {:.3} ms",
15085                    best[0] * 1e3,
15086                    best[1] * 1e3,
15087                    best[2] * 1e3
15088                );
15089            }
15090        }
15091    }
15092
15093    /// The q4t twin of the throughput bench, same shape and rules, so the
15094    /// two quantisations' batch kernels can be read against each other.
15095    /// `cargo test -p cortiq-engine --release q4t_matmat_throughput -- --ignored --nocapture`
15096    #[test]
15097    #[ignore]
15098    fn q4t_matmat_throughput() {
15099        let (rows, cols, b) = (9216usize, 2304usize, 296usize);
15100        let total =
15101            cortiq_core::quant::expected_nbytes(cortiq_core::TensorDtype::Q4Tiled, &[rows, cols])
15102                .unwrap();
15103        // q4t carries a per-group f16 scale in the tile's first two bytes;
15104        // random bytes there decode to inf and the bench would time NaNs.
15105        let mut bytes: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
15106        let sc = cortiq_core::quant::f32_to_f16(0.02);
15107        for t in bytes.chunks_mut(super::Q4_TILE) {
15108            t[..2].copy_from_slice(&sc.to_le_bytes());
15109        }
15110        let xs: Vec<f32> = (0..b * cols)
15111            .map(|i| ((i % 97) as f32 - 48.0) / 48.0)
15112            .collect();
15113        let mut out = vec![0f32; b * rows];
15114        let pool = crate::pool::Pool::from_env();
15115        super::q4t_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
15116        let reps: usize = std::env::var("CMF_BENCH_REPS")
15117            .ok()
15118            .and_then(|v| v.parse().ok())
15119            .unwrap_or(10);
15120        let mut best = f64::MAX;
15121        for _ in 0..reps {
15122            let t = std::time::Instant::now();
15123            super::q4t_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
15124            best = best.min(t.elapsed().as_secs_f64());
15125        }
15126        let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
15127        println!(
15128            "q4t matmat {rows}x{cols} b={b}: {:.1} ms  {:.1} GFLOP/s  (checksum {:.3})",
15129            best * 1e3,
15130            flops / best / 1e9,
15131            out.iter().take(64).sum::<f32>()
15132        );
15133    }
15134}