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
23/// `mxfp4` (OCP MX FP4) CPU kernels.
24mod mxfp4;
25
26pub enum QTensor {
27    F32 {
28        data: Vec<f32>,
29        rows: usize,
30        cols: usize,
31    },
32    Mapped {
33        model: Arc<CmfModel>,
34        /// Index into the model's tensor directory.
35        idx: usize,
36        dtype: TensorDtype,
37        rows: usize,
38        cols: usize,
39        /// Per-row scales, dequantized to f32 up front (tiny).
40        row_scale: Vec<f32>,
41        /// q8_2f column field (θ), dequantized up front; empty for q8_row.
42        col_field: Vec<f32>,
43        /// Vbit only: byte offset of each row's packed data within the
44        /// tensor blob (`[rows + 1]`, computed once at load — the per-
45        /// matvec prefix scan over row bit-widths was O(rows) each call).
46        vbit_offsets: Vec<usize>,
47        /// q8-family decode repack (load-time, optional): rows in groups
48        /// of 4, interleaved in 16-byte units — one 64-byte line per
49        /// iteration feeds all 4 sdot lanes, ONE sequential weight
50        /// stream per worker instead of four (this is where llama.cpp's
51        /// repacked Q8 kernels get their bandwidth). Empty = off
52        /// (CMF_REPACK=0, non-SDOT arch, or an ineligible shape). Trades
53        /// an anonymous copy of the quants for mmap pages that go cold.
54        repack: Vec<u8>,
55    },
56    /// Binary low-rank (`blr`): ±1 paths, signs from the mmap, scales as
57    /// f32 (docs/BLR_NVG.md §3). Every device-graph accessor answers `None`
58    /// for it (whole-token graphs decline to per-op/host paths); the CPU
59    /// kernels live in `crate::blr`, the per-op wgpu kernels in
60    /// `gpu_wgpu::blr`.
61    Blr(Box<crate::blr::BlrTensor>),
62}
63
64/// Load-time q8 repack gate (see `Mapped::repack`). OPT-IN
65/// (`CMF_REPACK=1`): the single-stream hypothesis LOST on Apple Silicon
66/// (M4, interleaved A/B: decode 101 vs 94 tok/s — four adjacent row
67/// streams per worker feed the prefetcher MORE memory-level parallelism
68/// than one); kept as an experiment flag for x86, where the tradeoff
69/// may land differently.
70fn repack_enabled() -> bool {
71    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
72    *ON.get_or_init(|| {
73        std::env::var("CMF_REPACK")
74            .map(|v| v == "1")
75            .unwrap_or(cfg!(target_os = "android"))
76    })
77}
78
79/// Interleave q8 rows for the decode kernel: group g holds rows
80/// 4g..4g+4 as [r0[c], r1[c], r2[c], r3[c]] per 16-byte chunk c. Only
81/// full groups are packed — tail rows keep reading the mmap layout.
82fn q8_repack(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
83    #[cfg(target_arch = "aarch64")]
84    let arch_ok = sdot_enabled();
85    #[cfg(not(target_arch = "aarch64"))]
86    let arch_ok = false;
87    if !arch_ok || !repack_enabled() || rows < 256 || cols % 16 != 0 {
88        return Vec::new();
89    }
90    q8_repack_layout(bytes, rows, cols)
91}
92
93/// The pure layout transform behind `q8_repack` (tested directly —
94/// the gate depends on arch and env).
95fn q8_repack_layout(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
96    let groups = rows / 4;
97    let mut rep = vec![0u8; groups * 4 * cols];
98    for g in 0..groups {
99        let dst = &mut rep[g * 4 * cols..(g + 1) * 4 * cols];
100        for c in 0..cols / 16 {
101            for lane in 0..4 {
102                let src = (g * 4 + lane) * cols + c * 16;
103                dst[c * 64 + lane * 16..c * 64 + lane * 16 + 16]
104                    .copy_from_slice(&bytes[src..src + 16]);
105            }
106        }
107    }
108    rep
109}
110
111/// Prefix-sum of vbit row payload offsets (absolute within the tensor
112/// bytes). `offsets[r]..offsets[r+1]` is row r's packed data.
113fn vbit_row_offsets(bytes: &[u8], rows: usize, cols: usize) -> Vec<usize> {
114    let ng = cols / GROUP_SIZE;
115    let bits = &bytes[..rows];
116    let mut offsets = Vec::with_capacity(rows + 1);
117    let mut off = rows + rows * ng * 2;
118    for r in 0..rows {
119        offsets.push(off);
120        off += (cols * bits[r] as usize).div_ceil(8);
121    }
122    offsets.push(off);
123    offsets
124}
125
126/// `CMF_X86_BLOCKED` / `CMF_GPU_LMHEAD` / `CMF_GPU_SPLIT`, read once. They
127/// used to be read from the environment on every large matvec and on every
128/// matmat in six places — microseconds each, but also a knob that could
129/// change under a running process, which is not a thing a kernel choice
130/// should be able to do mid-sequence.
131fn blocked_enabled() -> bool {
132    use std::sync::atomic::Ordering::Relaxed;
133    match BLOCKED_OVERRIDE.load(Relaxed) {
134        1 => false,
135        2 => true,
136        _ => {
137            static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
138            *ON.get_or_init(|| {
139                std::env::var("CMF_X86_BLOCKED")
140                    .map(|v| v != "0")
141                    .unwrap_or(true)
142            })
143        }
144    }
145}
146
147static BLOCKED_OVERRIDE: std::sync::atomic::AtomicU8 = std::sync::atomic::AtomicU8::new(0);
148
149/// Force the blocked GEMM on or off, ignoring the environment; `None`
150/// restores it. For tests that need to run BOTH paths and compare them:
151/// `blocked_enabled` caches its answer for the life of the process, which
152/// is right when the environment is the only input, but leaves a test that
153/// flips `CMF_X86_BLOCKED` between two calls comparing a path against
154/// itself — or against whatever a test running in parallel latched first.
155pub fn set_blocked_override(on: Option<bool>) {
156    let v = match on {
157        None => 0,
158        Some(false) => 1,
159        Some(true) => 2,
160    };
161    BLOCKED_OVERRIDE.store(v, std::sync::atomic::Ordering::Relaxed);
162}
163
164fn gpu_lmhead_enabled() -> bool {
165    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
166    *ON.get_or_init(|| {
167        std::env::var("CMF_GPU_LMHEAD")
168            .map(|v| v != "0")
169            .unwrap_or(true)
170    })
171}
172
173fn gpu_split_frac() -> f32 {
174    // MiMo's banked verification projects every row on the device. Plain
175    // decode must not quantize half of its O/head activation on the CPU.
176    if FULL_GPU_Q8.get() {
177        return 1.0;
178    }
179    static F: std::sync::OnceLock<f32> = std::sync::OnceLock::new();
180    *F.get_or_init(|| {
181        std::env::var("CMF_GPU_SPLIT")
182            .ok()
183            .and_then(|v| v.parse::<f32>().ok())
184            .unwrap_or(0.5)
185            .clamp(0.0, 1.0)
186    })
187}
188
189impl QTensor {
190    pub fn from_f32(data: Vec<f32>, rows: usize, cols: usize) -> Self {
191        debug_assert_eq!(data.len(), rows * cols);
192        Self::F32 { data, rows, cols }
193    }
194
195    /// Wrap a directory tensor without dequantizing the payload.
196    /// Falls back to dequantized f32 for dtypes without a fused kernel.
197    pub fn from_model(model: &Arc<CmfModel>, name: &str) -> Result<Self, String> {
198        // Indexed lookup: the linear directory scan made pipeline build
199        // O(N²) on MoE/skills files with thousands of tensors.
200        let idx = model
201            .tensor_index(name)
202            .ok_or_else(|| format!("tensor '{name}' not found in CMF directory"))?;
203        let entry = &model.tensors[idx];
204        if entry.shape.len() != 2 {
205            return Err(format!("QTensor::from_model needs 2-D, got '{name}'"));
206        }
207        let (rows, cols) = (entry.shape[0], entry.shape[1]);
208        let bytes = model.entry_bytes(entry);
209
210        match entry.dtype {
211            TensorDtype::Q8Row | TensorDtype::Q8_2f => {
212                let n = rows * cols;
213                let scales_off = n;
214                let row_scale: Vec<f32> = (0..rows)
215                    .map(|o| {
216                        f16_to_f32(u16::from_le_bytes([
217                            bytes[scales_off + o * 2],
218                            bytes[scales_off + o * 2 + 1],
219                        ]))
220                    })
221                    .collect();
222                let col_field: Vec<f32> = if entry.dtype == TensorDtype::Q8_2f {
223                    let col_off = n + rows * 2;
224                    (0..cols)
225                        .map(|i| {
226                            f16_to_f32(u16::from_le_bytes([
227                                bytes[col_off + i * 2],
228                                bytes[col_off + i * 2 + 1],
229                            ]))
230                        })
231                        .collect()
232                } else {
233                    Vec::new()
234                };
235                Ok(Self::Mapped {
236                    model: model.clone(),
237                    idx,
238                    dtype: entry.dtype,
239                    rows,
240                    cols,
241                    row_scale,
242                    col_field,
243                    vbit_offsets: Vec::new(),
244                    repack: q8_repack(bytes, rows, cols),
245                })
246            }
247            // vbit: fused kernel unpacks variable-bit rows from mmap.
248            TensorDtype::Vbit if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
249                model: model.clone(),
250                idx,
251                dtype: entry.dtype,
252                rows,
253                cols,
254                row_scale: Vec::new(),
255                col_field: Vec::new(),
256                vbit_offsets: vbit_row_offsets(bytes, rows, cols),
257                repack: Vec::new(),
258            }),
259            // vbit_ro (§4.2): the offset table comes straight from the
260            // file — no load-time prefix scan; kernels are shared with
261            // legacy vbit (they consume absolute offsets either way).
262            TensorDtype::VbitRo if cols % GROUP_SIZE == 0 => {
263                let (_, off_off, packed_off) = cortiq_core::quant::vbit_ro_sections(rows, cols);
264                let offsets: Vec<usize> = (0..=rows)
265                    .map(|r| packed_off + cortiq_core::quant::vbit_ro_offset(bytes, off_off, r))
266                    .collect();
267                Ok(Self::Mapped {
268                    model: model.clone(),
269                    idx,
270                    dtype: entry.dtype,
271                    rows,
272                    cols,
273                    row_scale: Vec::new(),
274                    col_field: Vec::new(),
275                    vbit_offsets: offsets,
276                    repack: Vec::new(),
277                })
278            }
279            // q4_block: fused kernel reads nibbles straight from mmap —
280            // a 14B q4 file no longer explodes into ×8 f32 RAM.
281            // q4_tiled (§4.3): interleaved [scale][nibbles] tiles — one
282            // sequential memory stream (measured ×1.66 ARM / ×1.13 AVX2
283            // at kernel level over the split layout).
284            TensorDtype::Q4Tiled if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
285                model: model.clone(),
286                idx,
287                dtype: entry.dtype,
288                rows,
289                cols,
290                row_scale: Vec::new(),
291                col_field: Vec::new(),
292                vbit_offsets: Vec::new(),
293                repack: Vec::new(),
294            }),
295            // q4tp (§4.10): nibbles from mmap, scale from the row ladder —
296            // 7.3% less file than q4t at the same 4-bit grid.
297            TensorDtype::Q4TiledP if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
298                model: model.clone(),
299                idx,
300                dtype: entry.dtype,
301                rows,
302                cols,
303                row_scale: Vec::new(),
304                col_field: Vec::new(),
305                vbit_offsets: Vec::new(),
306                repack: Vec::new(),
307            }),
308            // q2tp: 2-bit chunks from mmap, scale from the same row ladder.
309            TensorDtype::Q2TiledP if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
310                model: model.clone(),
311                idx,
312                dtype: entry.dtype,
313                rows,
314                cols,
315                row_scale: Vec::new(),
316                col_field: Vec::new(),
317                vbit_offsets: Vec::new(),
318                repack: Vec::new(),
319            }),
320            TensorDtype::Q4Block if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
321                model: model.clone(),
322                idx,
323                dtype: entry.dtype,
324                rows,
325                cols,
326                row_scale: Vec::new(),
327                col_field: Vec::new(),
328                vbit_offsets: Vec::new(),
329                repack: Vec::new(),
330            }),
331            // q1: binary sign-bit tiles from mmap (1-bit-trained models).
332            // q1g: the same signs as planes with one scale per 128-group.
333            TensorDtype::Q1 if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
334                model: model.clone(),
335                idx,
336                dtype: entry.dtype,
337                rows,
338                cols,
339                row_scale: Vec::new(),
340                col_field: Vec::new(),
341                vbit_offsets: Vec::new(),
342                repack: Vec::new(),
343            }),
344            TensorDtype::Q1G if cols % cortiq_core::quant::Q1G_GROUP == 0 => Ok(Self::Mapped {
345                model: model.clone(),
346                idx,
347                dtype: entry.dtype,
348                rows,
349                cols,
350                row_scale: Vec::new(),
351                col_field: Vec::new(),
352                vbit_offsets: Vec::new(),
353                repack: Vec::new(),
354            }),
355            // mxfp4: E2M1 codes + E8M0 scales straight from the mmap (the
356            // native FP4 expert codec; `qtensor::mxfp4`).
357            TensorDtype::Mxfp4 if cols % cortiq_core::mxfp4::GROUP == 0 => Ok(Self::Mapped {
358                model: model.clone(),
359                idx,
360                dtype: entry.dtype,
361                rows,
362                cols,
363                row_scale: Vec::new(),
364                col_field: Vec::new(),
365                vbit_offsets: Vec::new(),
366                repack: Vec::new(),
367            }),
368            // q1t (ternary + outlier overlay): fused per-row dequant kernel
369            // reads straight from mmap — a 12B q1t stays ~its file size in
370            // RAM instead of dequantizing to ~48 GB of f32.
371            TensorDtype::Q1T if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
372                model: model.clone(),
373                idx,
374                dtype: entry.dtype,
375                rows,
376                cols,
377                row_scale: Vec::new(),
378                col_field: Vec::new(),
379                vbit_offsets: Vec::new(),
380                repack: Vec::new(),
381            }),
382            // blr: factored kernels over the packed signs. CMF_BLR_DENSE=1
383            // dequantizes to f32 instead (A/B parity runs).
384            TensorDtype::BinaryLowRank if std::env::var("CMF_BLR_DENSE").as_deref() != Ok("1") => {
385                Ok(Self::Blr(Box::new(crate::blr::BlrTensor::from_model(
386                    model, idx,
387                )?)))
388            }
389            // No fused kernel yet → dequantize once (correct, more RAM).
390            _ => {
391                let mut data = vec![0.0f32; rows * cols];
392                cortiq_core::quant::dequant_tensor(entry, bytes, &mut data)?;
393                Ok(Self::from_f32(data, rows, cols))
394            }
395        }
396    }
397
398    /// Binary low-rank tensor? Fused multi-tensor kernels refuse it and fall
399    /// back to per-tensor calls.
400    pub fn is_blr(&self) -> bool {
401        matches!(self, Self::Blr(_))
402    }
403
404    /// Memory-mapped `mxfp4` tensor? Like `blr`, it has no row-range kernel
405    /// in the fused multi-tensor jobs: those run it per tensor.
406    pub fn is_mxfp4(&self) -> bool {
407        matches!(
408            self,
409            Self::Mapped {
410                dtype: TensorDtype::Mxfp4,
411                ..
412            }
413        )
414    }
415
416    /// Codecs the fused multi-tensor jobs (`matvec_many`, the SwiGLU and
417    /// MoE row kernels) do not carry: they fall back to per-tensor calls.
418    fn per_tensor_only(&self) -> bool {
419        self.is_blr() || self.is_mxfp4()
420    }
421
422    /// (model, idx) of a `blr` tensor, for the wgpu graphs only (their
423    /// `blr` stage kernels, kind [`crate::gpu::GRAPH_KIND_BLR`]). Kept apart
424    /// from [`graph_weight`](Self::graph_weight), which stays None for
425    /// `blr`: every other graph consumer keys its kernels on that answer.
426    /// A file with a Prism contract keeps the per-op path.
427    pub fn graph_blr(&self) -> Option<(&Arc<CmfModel>, usize)> {
428        match self {
429            Self::Blr(t) if !crate::prism::has_contract(&t.model) => Some((&t.model, t.idx)),
430            _ => None,
431        }
432    }
433
434    /// [`graph_weight`](Self::graph_weight), then the Prism descriptor
435    /// form, then a `blr` tensor as kind [`crate::gpu::GRAPH_KIND_BLR`] —
436    /// the wgpu token and batch graph builders' lookup.
437    pub(crate) fn graph_weight_or_blr(
438        &self,
439    ) -> Option<(&std::sync::Arc<CmfModel>, usize, u8, &[f32])> {
440        self.graph_weight()
441            .or_else(|| self.graph_weight_descriptor())
442            .or_else(|| {
443                self.graph_blr()
444                    .map(|(m, i)| (m, i, crate::gpu::GRAPH_KIND_BLR, &[][..]))
445            })
446    }
447
448    /// Binary-weight-mapped tensor, q1 or q1g? (GPU gates: the binary CPU
449    /// kernels are compute-bound, so offload pays at much smaller shapes
450    /// than q8.)
451    pub(crate) fn is_q1(&self) -> bool {
452        matches!(
453            self,
454            Self::Mapped {
455                dtype: TensorDtype::Q1 | TensorDtype::Q1G,
456                ..
457            }
458        )
459    }
460
461    /// The binary-weight view (q1 tiles or q1g planes) the shared CPU row
462    /// drivers take; None for every other storage.
463    fn bin_w(&self) -> Option<BinW<'_>> {
464        match self {
465            Self::Mapped {
466                dtype: TensorDtype::Q1,
467                ..
468            } => Some(BinW::Q1(self.quant_bytes())),
469            Self::Mapped {
470                dtype: TensorDtype::Q1G,
471                rows,
472                cols,
473                ..
474            } => Some(BinW::q1g(self.quant_bytes(), *rows, *cols)),
475            _ => None,
476        }
477    }
478
479    /// Owned-f32 view (data, rows, cols) — the GDN a/b gate projections
480    /// arrive dequantized (force-f16 in the converter → F32 in RAM).
481    pub(crate) fn f32_parts(&self) -> Option<(&[f32], usize, usize)> {
482        match self {
483            Self::F32 { data, rows, cols } => Some((data, *rows, *cols)),
484            _ => None,
485        }
486    }
487
488    /// (directory idx, rows, cols) of a q1-mapped tensor — the
489    /// whole-block GPU path resolves offsets itself.
490    /// (idx, rows, cols) of a mapped tensor the whole-token GPU graph can drive
491    /// — Q1, Q1T or Q4-block (it resolves the offset and picks the kernel by
492    /// dtype). Q4-block lets a precise down_proj/lm_head stay on-device.
493    /// Named `q1_parts` for historical reasons.
494    pub(crate) fn q1_parts(&self) -> Option<(usize, usize, usize)> {
495        if self.has_prism_contract() {
496            return None;
497        }
498        match self {
499            #[cfg(target_os = "macos")]
500            Self::Mapped {
501                dtype: TensorDtype::Q1T,
502                ..
503            } if !crate::gpu::metal_q1t_enabled() => None,
504            Self::Mapped {
505                idx,
506                dtype:
507                    TensorDtype::Q1
508                    | TensorDtype::Q1G
509                    | TensorDtype::Q1T
510                    | TensorDtype::Q4Block
511                    | TensorDtype::Q4Tiled
512                    // Q2TiledP deliberately absent: the Metal graph has no
513                    // q2tp kernel, and advertising it here made the block
514                    // plan truncate mid-run at the first q2tp layer.
515                    | TensorDtype::Q4TiledP
516                    | TensorDtype::Q8Row
517                    | TensorDtype::Q8_2f,
518                rows,
519                cols,
520                ..
521            } => Some((*idx, *rows, *cols)),
522            _ => None,
523        }
524    }
525
526    /// `(directory idx, rows, cols)` for the native Metal token graph.  The
527    /// historical q1 graph gate intentionally refuses every Prism tensor so
528    /// an untransformed q2 payload cannot slip into the resident path.  The
529    /// Metal2 graph is descriptor-aware and admits only the production
530    /// q2tp-affine forward targets; ordinary q1/q4 callers retain the old
531    /// `q1_parts` behaviour.
532    #[cfg(target_os = "macos")]
533    pub(crate) fn metal_graph_parts(&self) -> Option<(usize, usize, usize)> {
534        if let Some((model, idx, kind, _)) = self.graph_weight_descriptor() {
535            let name = &model.tensors[idx].name;
536            let forward = kind == 9 && crate::prism::is_forward_weight(model, name);
537            let affine = kind == 9 && crate::prism::is_affine_target(model, name);
538            if forward && affine {
539                let e = model.tensors.get(idx)?;
540                return Some((idx, *e.shape.first()?, *e.shape.get(1)?));
541            }
542        }
543        self.q1_parts()
544    }
545
546    /// (directory idx, rows, cols) of a q4_tiled mapped tensor. The
547    /// chunk-prefill graph takes it in the same 4-tuple slot as
548    /// `q8_row_parts` with an EMPTY row_scale — q4t carries its scales
549    /// inside the 18-byte tiles, and the empty slice is what tells the
550    /// encoder to reach for the q4t kernels.
551    pub(crate) fn q4t_parts(&self) -> Option<(usize, usize, usize)> {
552        if self.has_prism_contract() {
553            return None;
554        }
555        match self {
556            Self::Mapped {
557                idx,
558                dtype: TensorDtype::Q4Tiled,
559                rows,
560                cols,
561                ..
562            } => Some((*idx, *rows, *cols)),
563            _ => None,
564        }
565    }
566
567    /// (directory idx, rows, cols) of a q4tp mapped tensor. Same empty-scale
568    /// slot as `q4t_parts` in the chunk graph — the encoder tells the two
569    /// apart by the tensor's dtype, not by the slot.
570    pub(crate) fn q4tp_parts(&self) -> Option<(usize, usize, usize)> {
571        if self.has_prism_contract() {
572            return None;
573        }
574        match self {
575            Self::Mapped {
576                idx,
577                dtype: TensorDtype::Q4TiledP,
578                rows,
579                cols,
580                ..
581            } => Some((*idx, *rows, *cols)),
582            _ => None,
583        }
584    }
585
586    /// (directory idx, rows, cols, row_scale) of a plain q8_row mapped
587    /// tensor — the chunk-prefill GPU graph resolves offsets itself.
588    /// q8_2f is excluded on purpose: its column field would need a
589    /// prescale stage on the device.
590    pub(crate) fn q8_row_parts(&self) -> Option<(usize, usize, usize, &[f32])> {
591        if self.has_prism_contract() {
592            return None;
593        }
594        match self {
595            Self::Mapped {
596                idx,
597                dtype: TensorDtype::Q8Row,
598                rows,
599                cols,
600                row_scale,
601                col_field,
602                ..
603            } if col_field.is_empty() => Some((*idx, *rows, *cols, row_scale)),
604            _ => None,
605        }
606    }
607
608    /// (directory idx, rows, cols, row_scale, col_field) of a q8_2f mapped
609    /// tensor — the chunk-prefill graph stages x·col on the device first.
610    pub(crate) fn q8_2f_parts(&self) -> Option<(usize, usize, usize, &[f32], &[f32])> {
611        if self.has_prism_contract() {
612            return None;
613        }
614        match self {
615            Self::Mapped {
616                idx,
617                dtype: TensorDtype::Q8_2f,
618                rows,
619                cols,
620                row_scale,
621                col_field,
622                ..
623            } if col_field.len() == *cols => Some((*idx, *rows, *cols, row_scale, col_field)),
624            _ => None,
625        }
626    }
627
628    /// The layout this tensor is stored in, when it is mapped from a model.
629    /// The frames branch on it — a q2tp gate against a q4tp down is a real
630    /// combination in the 2-bit profile and needs a different kernel.
631    pub fn model_dtype(&self) -> Option<cortiq_core::TensorDtype> {
632        match self {
633            Self::Mapped { dtype, .. } => Some(*dtype),
634            _ => None,
635        }
636    }
637
638    /// The tensor's index in the model directory, when it is mapped from one.
639    /// The GPU frames bind by index rather than by name — a name lookup per
640    /// layer per token is not free, and the index is what the device cache is
641    /// keyed on anyway.
642    pub fn model_idx(&self) -> Option<usize> {
643        match self {
644            Self::Mapped { idx, .. } => Some(*idx),
645            _ => None,
646        }
647    }
648
649    /// The model this tensor is mapped from, when it is mapped at all. The
650    /// GPU frames need the container to reach the bytes; a QTensor already
651    /// holds it, and threading a second handle down every call site to say
652    /// the same thing invites the two to disagree.
653    pub fn model_arc(&self) -> Option<std::sync::Arc<cortiq_core::CmfModel>> {
654        match self {
655            Self::Mapped { model, .. } => Some(model.clone()),
656            _ => None,
657        }
658    }
659
660    /// Whether this mapped tensor belongs to the Prism/Bonsai transform
661    /// contract.  Device graphs do not carry the descriptor, so callers use
662    /// this conservative predicate to stay on the descriptor-aware CPU path
663    /// instead of silently executing an unrotated matrix.
664    pub(crate) fn has_prism_contract(&self) -> bool {
665        matches!(self, Self::Mapped { model, .. } if crate::prism::has_contract(model))
666    }
667
668    pub fn rows(&self) -> usize {
669        match self {
670            Self::F32 { rows, .. } | Self::Mapped { rows, .. } => *rows,
671            Self::Blr(t) => t.rows,
672        }
673    }
674
675    /// Mapped q4t handle (model + directory index) — the fused GPU FFN
676    /// needs the raw file coordinates of its three projections.
677    pub(crate) fn mapped_q4t(&self) -> Option<(&Arc<CmfModel>, usize)> {
678        if self.has_prism_contract() {
679            return None;
680        }
681        match self {
682            Self::Mapped {
683                model,
684                idx,
685                dtype: TensorDtype::Q4Tiled,
686                ..
687            } => Some((model, *idx)),
688            _ => None,
689        }
690    }
691
692    /// Same slot as `mapped_q4t` for a q4tp tensor — the fused DiT FFN picks
693    /// its kernels by which of the two answers.
694    pub fn mapped_q4tp(&self) -> Option<(&Arc<CmfModel>, usize)> {
695        if self.has_prism_contract() {
696            return None;
697        }
698        match self {
699            Self::Mapped {
700                model,
701                idx,
702                dtype: TensorDtype::Q4TiledP,
703                ..
704            } => Some((model, *idx)),
705            _ => None,
706        }
707    }
708
709    /// (model, tensor idx) for a mapped weight in ANY codec the fused device
710    /// paths can run — four-bit tiled or either int8 layout.
711    ///
712    /// The fused DiT chains asked for `mapped_q4tp` by name, so an eight-bit
713    /// container never reached them and rendered through per-op GEMMs even
714    /// after those kernels learned its codec. The gate is what the codec has
715    /// a device GEMM for, not which codec it is.
716    pub fn mapped_device_gemm(&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::Q4TiledP | TensorDtype::Q8Row | TensorDtype::Q8_2f,
725                ..
726            } => Some((model, *idx)),
727            _ => None,
728        }
729    }
730
731    /// (model, tensor idx) for a q8_2f mapped weight — the eight-bit slot of
732    /// the MoE graph ladder (int8 body, f16 row scales, f16 column field).
733    pub fn mapped_q8_2f(&self) -> Option<(&Arc<CmfModel>, usize)> {
734        if self.has_prism_contract() {
735            return None;
736        }
737        match self {
738            Self::Mapped {
739                model,
740                idx,
741                dtype: TensorDtype::Q8_2f,
742                ..
743            } => Some((model, *idx)),
744            _ => None,
745        }
746    }
747
748    /// (model, tensor idx) for a q2tp mapped weight — the 2-bit twin of
749    /// `mapped_q4tp`, used by the mixed MoE profile.
750    pub fn mapped_q2tp(&self) -> Option<(&Arc<CmfModel>, usize)> {
751        if self.has_prism_contract() {
752            return None;
753        }
754        match self {
755            Self::Mapped {
756                model,
757                idx,
758                dtype: TensorDtype::Q2TiledP,
759                ..
760            } => Some((model, *idx)),
761            _ => None,
762        }
763    }
764
765    pub fn cols(&self) -> usize {
766        match self {
767            Self::F32 { cols, .. } | Self::Mapped { cols, .. } => *cols,
768            Self::Blr(t) => t.cols,
769        }
770    }
771
772    /// (model, tensor idx) for a q1 mapped weight — the wgpu token graph
773    /// keys its resident VRAM cache by idx. None for any other dtype/kind.
774    pub fn mapped_q1(&self) -> Option<(&std::sync::Arc<CmfModel>, usize)> {
775        if self.has_prism_contract() {
776            return None;
777        }
778        match self {
779            Self::Mapped {
780                model,
781                idx,
782                dtype: TensorDtype::Q1,
783                ..
784            } => Some((model, *idx)),
785            _ => None,
786        }
787    }
788
789    /// (model, idx, kind, row_scale) for a graph-capable mapped weight.
790    /// kind: 0=q8_row (per-row scales), 1=q1, 2=q4_block, 3=q1t
791    /// (tile-embedded, no rs), 5=q4_tiled, 6=q4tp, 7=q8_2f (both scale
792    /// planes live inside the tensor), 9=q2tp, 10=q1g (sign + scale
793    /// planes). None only for `vbit`.
794    ///
795    /// The old comment here claimed q4_block was unhandled while the arm
796    /// right below mapped it, and it named q8_2f as unhandled after that
797    /// stopped being true — a stale comment on this function is how a
798    /// model silently loses the graph, so it is worth keeping honest.
799    pub fn graph_weight(&self) -> Option<(&std::sync::Arc<CmfModel>, usize, u8, &[f32])> {
800        if self.has_prism_contract() {
801            return None;
802        }
803        self.graph_weight_descriptor()
804    }
805
806    /// Descriptor-aware graph handle used only by the Prism token graph.
807    /// Ordinary graph callers continue to use [`graph_weight`] and therefore
808    /// remain fail-closed until they provide the same explicit transform
809    /// contract.
810    pub(crate) fn graph_weight_descriptor(
811        &self,
812    ) -> Option<(&std::sync::Arc<CmfModel>, usize, u8, &[f32])> {
813        match self {
814            Self::Mapped {
815                model,
816                idx,
817                dtype: TensorDtype::Q8Row,
818                row_scale,
819                ..
820            } => Some((model, *idx, 0, row_scale.as_slice())),
821            Self::Mapped {
822                model,
823                idx,
824                dtype: TensorDtype::Q1,
825                ..
826            } => Some((model, *idx, 1, &[])),
827            // Kind 10, never 1: q1g's sign and scale planes are a different
828            // walk from q1's 6-byte tiles (the kind-2/5 mixup taught that).
829            Self::Mapped {
830                model,
831                idx,
832                dtype: TensorDtype::Q1G,
833                ..
834            } => Some((model, *idx, 10, &[])),
835            // Q4Tiled is kind 5, NOT 2: both carried 2 historically, and
836            // the wgpu token graph fed 18B interleaved tiles to the
837            // split-layout q4b kernel — garbage output on q4t models
838            // (caught by an end-to-end answer check on real Vulkan).
839            Self::Mapped {
840                model,
841                idx,
842                dtype: TensorDtype::Q4Tiled,
843                ..
844            } => Some((model, *idx, 5, &[])),
845            // Kind 6, not 5: q4tp's nibble stride and scale planes differ,
846            // and feeding them to the q4t kernel is exactly the mistake that
847            // produced garbage when Q4Tiled shared kind 2 with Q4Block.
848            Self::Mapped {
849                model,
850                idx,
851                dtype: TensorDtype::Q4TiledP,
852                ..
853            } => Some((model, *idx, 6, &[])),
854            Self::Mapped {
855                model,
856                idx,
857                dtype: TensorDtype::Q4Block,
858                ..
859            } => Some((model, *idx, 2, &[])),
860            // q8_2f carries BOTH scale planes after the int8 body (rows
861            // f16, then cols f16), so the graph takes the whole tensor
862            // and the kernel reads them where they lie — no host-side
863            // prescale, which is what the per-op path does instead.
864            Self::Mapped {
865                model,
866                idx,
867                dtype: TensorDtype::Q8_2f,
868                ..
869            } => Some((model, *idx, 7, &[])),
870            Self::Mapped {
871                model,
872                idx,
873                dtype: TensorDtype::Q1T,
874                ..
875            } => Some((model, *idx, 3, &[])),
876            // Kind 9: the 2-bit plane on the q4tp ladder (dense FFN gate/up
877            // of the q2tp profile). Its own kernel — 8 bytes a group where
878            // q4tp has 16, and rung 0 is the exact zero.
879            Self::Mapped {
880                model,
881                idx,
882                dtype: TensorDtype::Q2TiledP,
883                ..
884            } => Some((model, *idx, 9, &[])),
885            _ => None,
886        }
887    }
888
889    /// Dense f32 view — only for owned tensors. Masked/sparse execution
890    /// paths require it; quantized weights don't support masks yet.
891    pub fn as_f32(&self) -> Option<&[f32]> {
892        match self {
893            Self::F32 { data, .. } => Some(data),
894            Self::Mapped { .. } | Self::Blr(_) => None,
895        }
896    }
897
898    fn quant_bytes(&self) -> &[u8] {
899        match self {
900            Self::Mapped { model, idx, .. } => model.entry_bytes(&model.tensors[*idx]),
901            Self::Blr(t) => t.bytes(),
902            Self::F32 { .. } => unreachable!("quant_bytes on F32"),
903        }
904    }
905
906    /// Dequantize one row into `dst` (embedding lookup).
907    pub fn row_f32(&self, r: usize, dst: &mut [f32]) {
908        let cols = self.cols();
909        debug_assert_eq!(dst.len(), cols);
910        match self {
911            Self::F32 { data, .. } => dst.copy_from_slice(&data[r * cols..(r + 1) * cols]),
912            Self::Blr(t) => t.row_f32(r, dst),
913            Self::Mapped {
914                model,
915                idx,
916                dtype,
917                row_scale,
918                col_field,
919                vbit_offsets,
920                ..
921            } => {
922                if *dtype == TensorDtype::Q4Tiled {
923                    let bytes = self.quant_bytes();
924                    let gpr = cols / GROUP_SIZE;
925                    for gi in 0..gpr {
926                        let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
927                        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
928                        for (k, &b) in tile[2..].iter().enumerate() {
929                            dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
930                            dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
931                        }
932                    }
933                    if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
934                        crate::prism::inverse_embedding(model, dst);
935                    }
936                    return;
937                }
938                if *dtype == TensorDtype::Q4TiledP {
939                    let bytes = self.quant_bytes();
940                    let gpr = cols / GROUP_SIZE;
941                    let v = Q4tpView::new(bytes, self.rows(), cols);
942                    let mut sc = vec![0f32; gpr];
943                    v.scales_into(r, gpr, &mut sc);
944                    for gi in 0..gpr {
945                        let tile = &v.nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
946                        let s = sc[gi];
947                        for (k, &b) in tile.iter().enumerate() {
948                            dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
949                            dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
950                        }
951                    }
952                    if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
953                        crate::prism::inverse_embedding(model, dst);
954                    }
955                    return;
956                }
957                if *dtype == TensorDtype::Q2TiledP {
958                    let bytes = self.quant_bytes();
959                    let gpr = cols / GROUP_SIZE;
960                    let v = Q4tpView::new_q2(bytes, self.rows(), cols);
961                    let mut sc = vec![0f32; gpr];
962                    v.scales_into(r, gpr, &mut sc);
963                    for gi in 0..gpr {
964                        let ch =
965                            &v.nib[(r * gpr + gi) * Q2TP_CHUNK..(r * gpr + gi + 1) * Q2TP_CHUNK];
966                        let s = sc[gi];
967                        for (k, &b) in ch.iter().enumerate() {
968                            for j in 0..4 {
969                                let center = if crate::prism::is_affine_target(
970                                    model,
971                                    &model.tensors[*idx].name,
972                                ) {
973                                    1.0
974                                } else {
975                                    1.5
976                                };
977                                dst[gi * GROUP_SIZE + k * 4 + j] =
978                                    (((b >> (2 * j)) & 3) as f32 - center) * 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::Q4Block {
988                    let (packed, scales) = q4_split(self.quant_bytes(), self.rows(), cols);
989                    let gpr = cols / GROUP_SIZE;
990                    for gi in 0..gpr {
991                        let g = r * gpr + gi;
992                        let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
993                        for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
994                            dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
995                            dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
996                        }
997                    }
998                    if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
999                        crate::prism::inverse_embedding(model, dst);
1000                    }
1001                    return;
1002                }
1003                if *dtype == TensorDtype::Mxfp4 {
1004                    mxfp4::Mx::new(self.quant_bytes(), self.rows(), cols).row_f32(r, dst);
1005                    if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
1006                        crate::prism::inverse_embedding(model, dst);
1007                    }
1008                    return;
1009                }
1010                if *dtype == TensorDtype::Q1 {
1011                    let bytes = self.quant_bytes();
1012                    let gpr = cols / GROUP_SIZE;
1013                    for gi in 0..gpr {
1014                        let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
1015                        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
1016                        for (j, &b) in tile[2..].iter().enumerate() {
1017                            for k in 0..8 {
1018                                dst[gi * GROUP_SIZE + j * 8 + k] =
1019                                    (((b >> k) & 1) as f32 * 2.0 - 1.0) * s;
1020                            }
1021                        }
1022                    }
1023                    if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
1024                        crate::prism::inverse_embedding(model, dst);
1025                    }
1026                    return;
1027                }
1028                if *dtype == TensorDtype::Q1G {
1029                    // The tied Bonsai embedding is q1g too: one row = cols/8
1030                    // sign bytes + cols/128 scales, same ±s values as q1.
1031                    let bytes = self.quant_bytes();
1032                    let sc_off = cortiq_core::quant::q1g_scale_offset(self.rows(), cols);
1033                    let gpr = cols / GROUP_SIZE;
1034                    for gi in 0..gpr {
1035                        let t = r * gpr + gi;
1036                        let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, t));
1037                        for (j, &b) in bytes[t * 4..t * 4 + 4].iter().enumerate() {
1038                            for k in 0..8 {
1039                                dst[gi * GROUP_SIZE + j * 8 + k] =
1040                                    (((b >> k) & 1) as f32 * 2.0 - 1.0) * s;
1041                            }
1042                        }
1043                    }
1044                    if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
1045                        crate::prism::inverse_embedding(model, dst);
1046                    }
1047                    return;
1048                }
1049                if *dtype == TensorDtype::Q1T {
1050                    let bytes = self.quant_bytes();
1051                    let gpr = cols / GROUP_SIZE;
1052                    let base_len = self.rows() * gpr * cortiq_core::quant::Q1T_TILE;
1053                    for gi in 0..gpr {
1054                        let off = (r * gpr + gi) * cortiq_core::quant::Q1T_TILE;
1055                        let s = cortiq_core::quant::f16_to_f32(u16::from_le_bytes([
1056                            bytes[off],
1057                            bytes[off + 1],
1058                        ]));
1059                        let codes = &bytes[off + 2..off + cortiq_core::quant::Q1T_TILE];
1060                        for k in 0..GROUP_SIZE {
1061                            dst[gi * GROUP_SIZE + k] = match cortiq_core::quant::q1t_code(codes, k)
1062                            {
1063                                1 => s,
1064                                2 => -s,
1065                                _ => 0.0,
1066                            };
1067                        }
1068                    }
1069                    // Overlay
1070                    let rows = self.rows();
1071                    let entries = base_len + (rows + 1) * 4;
1072                    if entries <= bytes.len() {
1073                        let ptrs = &bytes[base_len..base_len + (rows + 1) * 4];
1074                        let r0 = u32::from_le_bytes([
1075                            ptrs[r * 4],
1076                            ptrs[r * 4 + 1],
1077                            ptrs[r * 4 + 2],
1078                            ptrs[r * 4 + 3],
1079                        ]) as usize;
1080                        let r1 = u32::from_le_bytes([
1081                            ptrs[(r + 1) * 4],
1082                            ptrs[(r + 1) * 4 + 1],
1083                            ptrs[(r + 1) * 4 + 2],
1084                            ptrs[(r + 1) * 4 + 3],
1085                        ]) as usize;
1086                        let off = entries + r0 * 4;
1087                        for i in 0..r1 - r0 {
1088                            let item = &bytes[off + i * 4..off + i * 4 + 4];
1089                            let c = u16::from_le_bytes([item[0], item[1]]) as usize;
1090                            let v = cortiq_core::quant::f16_to_f32(u16::from_le_bytes([
1091                                item[2], item[3],
1092                            ]));
1093                            if c < cols {
1094                                dst[c] = v;
1095                            }
1096                        }
1097                    }
1098                    if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
1099                        crate::prism::inverse_embedding(model, dst);
1100                    }
1101                    return;
1102                }
1103                if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1104                    let bytes = self.quant_bytes();
1105                    let rows = self.rows();
1106                    let ng = cols / GROUP_SIZE;
1107                    let bits = &bytes[..rows];
1108                    let sc_off = rows;
1109                    // Precomputed at load — embedding lookup used to scan
1110                    // the bit-widths of every preceding row (O(token_id)).
1111                    let off = vbit_offsets[r];
1112                    let b = bits[r] as usize;
1113                    let l = ((1usize << (b - 1)) - 1) as f32;
1114                    let data = &bytes[off..];
1115                    let (mut acc, mut nbits, mut byte_idx) = (0u64, 0usize, 0usize);
1116                    for (i, d) in dst.iter_mut().enumerate() {
1117                        while nbits < b {
1118                            acc = (acc << 8) | data[byte_idx] as u64;
1119                            byte_idx += 1;
1120                            nbits += 8;
1121                        }
1122                        let u = ((acc >> (nbits - b)) & ((1u64 << b) - 1)) as f32;
1123                        nbits -= b;
1124                        let so = (r * ng + i / GROUP_SIZE) * 2;
1125                        let sv = f16_to_f32(u16::from_le_bytes([
1126                            bytes[sc_off + so],
1127                            bytes[sc_off + so + 1],
1128                        ]));
1129                        *d = (u - l) * sv;
1130                    }
1131                    if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
1132                        crate::prism::inverse_embedding(model, dst);
1133                    }
1134                    return;
1135                }
1136                let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
1137                let s = row_scale[r];
1138                match dtype {
1139                    TensorDtype::Q8Row => {
1140                        for (d, &b) in dst.iter_mut().zip(q) {
1141                            *d = (b as i8) as f32 * s;
1142                        }
1143                    }
1144                    TensorDtype::Q8_2f => {
1145                        for (i, (d, &b)) in dst.iter_mut().zip(q).enumerate() {
1146                            *d = (b as i8) as f32 * s * col_field[i];
1147                        }
1148                    }
1149                    _ => unreachable!(),
1150                }
1151                if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
1152                    crate::prism::inverse_embedding(model, dst);
1153                }
1154            }
1155        }
1156    }
1157
1158    /// Can this tensor's columns be read cheaply (for sparse down_proj)?
1159    /// True for F32/Q8Row/Q8_2f (per-row scale, direct strided access);
1160    /// false for group-packed q4/vbit (column access would unpack whole
1161    /// groups — sparse execution falls back to f32 for those).
1162    pub fn sparse_col_ok(&self) -> bool {
1163        match self {
1164            Self::F32 { .. } => true,
1165            Self::Blr(_) => false,
1166            Self::Mapped { dtype, .. } => {
1167                matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
1168            }
1169        }
1170    }
1171
1172    /// down_proj [hidden, inter]: accumulate `w · col(c)` into `out`
1173    /// [hidden] — reads ONLY column `c` (one neuron) from the mmap,
1174    /// no full-matrix dequant. `out[k] += w · down[k, c]`.
1175    pub fn add_col_scaled(&self, c: usize, w: f32, out: &mut [f32]) {
1176        let inter = self.cols();
1177        let hidden = self.rows();
1178        debug_assert_eq!(out.len(), hidden);
1179        match self {
1180            Self::F32 { data, .. } => {
1181                for (k, o) in out.iter_mut().enumerate() {
1182                    *o += w * data[k * inter + c];
1183                }
1184            }
1185            Self::Blr(t) => t.add_col_scaled(c, w, out),
1186            Self::Mapped {
1187                dtype,
1188                row_scale,
1189                col_field,
1190                ..
1191            } => {
1192                let q = self.quant_bytes();
1193                let colf = if *dtype == TensorDtype::Q8_2f {
1194                    col_field[c]
1195                } else {
1196                    1.0
1197                };
1198                let wc = w * colf;
1199                for (k, o) in out.iter_mut().enumerate() {
1200                    let b = q[k * inter + c] as i8 as f32;
1201                    *o += wc * b * row_scale[k];
1202                }
1203            }
1204        }
1205    }
1206
1207    /// Touch the head of row `r` so the DRAM latency of the next
1208    /// neuron's weights overlaps the current one's arithmetic.
1209    ///
1210    /// Scattered rows are what per-token sparsity reads, and a 2 KB
1211    /// stride is past what the hardware prefetcher follows: without this
1212    /// every row starts with a cold miss that nothing hides. One touch
1213    /// per 512 bytes is enough — the rest of the row is a sequential run
1214    /// the prefetcher does pick up.
1215    #[inline]
1216    pub fn prefetch_row(&self, r: usize) {
1217        let Self::Mapped { dtype, .. } = self else {
1218            return;
1219        };
1220        if !matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f) {
1221            return;
1222        }
1223        let cols = self.cols();
1224        let q = self.quant_bytes();
1225        let (a, b) = (r * cols, (r + 1) * cols);
1226        if b > q.len() {
1227            return;
1228        }
1229        let mut j = a;
1230        while j < b {
1231            unsafe { std::ptr::read_volatile(q.as_ptr().add(j)) };
1232            j += 512;
1233        }
1234    }
1235
1236    /// `out += w · row(r)` — the transposed twin of `add_col_scaled`.
1237    ///
1238    /// A neuron's `down` weights are a COLUMN of `[hidden, inter]`, and a
1239    /// column is strided: reading one costs a cache line per element, so
1240    /// per-neuron dynamic sparsity saves arithmetic and no bytes. Stored
1241    /// transposed (`down_proj.t.weight`, `[inter, hidden]`) the same
1242    /// weights are a contiguous ROW, and this accumulate reads exactly
1243    /// the neurons the token asked for.
1244    pub fn add_row_scaled(&self, r: usize, w: f32, out: &mut [f32], scratch: &mut [f32]) {
1245        let cols = self.cols();
1246        debug_assert_eq!(out.len(), cols);
1247        match self {
1248            Self::F32 { data, .. } => {
1249                let row = &data[r * cols..(r + 1) * cols];
1250                for (o, v) in out.iter_mut().zip(row) {
1251                    *o += w * v;
1252                }
1253            }
1254            Self::Blr(t) => {
1255                let tmp = &mut scratch[..cols];
1256                t.row_f32(r, tmp);
1257                for (o, v) in out.iter_mut().zip(tmp.iter()) {
1258                    *o += w * v;
1259                }
1260            }
1261            Self::Mapped {
1262                dtype,
1263                row_scale,
1264                col_field,
1265                ..
1266            } => match dtype {
1267                TensorDtype::Q8Row => {
1268                    let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
1269                    let ws = w * row_scale[r];
1270                    let row: &[i8] =
1271                        unsafe { std::slice::from_raw_parts(q.as_ptr() as *const i8, q.len()) };
1272                    axpy_i8_f32(out, row, ws);
1273                }
1274                TensorDtype::Q8_2f => {
1275                    let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
1276                    let ws = w * row_scale[r];
1277                    for ((o, b), c) in out.iter_mut().zip(q).zip(col_field) {
1278                        *o += ws * c * (*b as i8 as f32);
1279                    }
1280                }
1281                _ => {
1282                    self.row_f32(r, scratch);
1283                    for (o, v) in out.iter_mut().zip(scratch.iter()) {
1284                        *o += w * v;
1285                    }
1286                }
1287            },
1288        }
1289    }
1290
1291    /// Dot of row `r` with `x` (gate/up active-neuron path). Reads only
1292    /// row `r` from the mmap — no full dequant. q4/vbit dequant the row
1293    /// into `scratch` first (rare for active-FFN weights).
1294    pub fn row_dot(&self, r: usize, x: &[f32], scratch: &mut [f32]) -> f32 {
1295        let cols = self.cols();
1296        match self {
1297            Self::F32 { data, .. } => {
1298                let row = &data[r * cols..(r + 1) * cols];
1299                row.iter().zip(x).map(|(w, v)| w * v).sum()
1300            }
1301            Self::Blr(t) => {
1302                let tmp = &mut scratch[..cols];
1303                t.row_f32(r, tmp);
1304                tmp.iter().zip(x).map(|(w, v)| w * v).sum()
1305            }
1306            Self::Mapped {
1307                model,
1308                idx,
1309                dtype,
1310                row_scale,
1311                col_field,
1312                ..
1313            } => {
1314                let prism_forward =
1315                    crate::prism::is_forward_weight(model, &model.tensors[*idx].name);
1316                if prism_forward {
1317                    let transformed = crate::prism::forward(model, &x[..cols]);
1318                    let gpr = cols / GROUP_SIZE;
1319                    match dtype {
1320                        TensorDtype::Q2TiledP => {
1321                            let v = Q4tpView::new_q2(self.quant_bytes(), self.rows(), cols);
1322                            let mut sc = vec![0f32; gpr];
1323                            v.scales_into(r, gpr, &mut sc);
1324                            if crate::prism::is_affine_target(model, &model.tensors[*idx].name) {
1325                                return q2tp_affine_row_exact(v.nib, r, gpr, &transformed, &sc);
1326                            }
1327                            return q2tp_row_exact(v.nib, r, gpr, &transformed, &sc);
1328                        }
1329                        _ => {
1330                            self.row_f32(r, scratch);
1331                            return scratch.iter().zip(&transformed).map(|(w, v)| w * v).sum();
1332                        }
1333                    }
1334                }
1335                match dtype {
1336                    TensorDtype::Q8Row => {
1337                        let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
1338                        dot_i8_f32(q, x) * row_scale[r]
1339                    }
1340                    TensorDtype::Q8_2f => {
1341                        let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
1342                        dot_i8_col_f32(q, x, col_field) * row_scale[r]
1343                    }
1344                    _ => {
1345                        self.row_f32(r, scratch);
1346                        scratch.iter().zip(x).map(|(w, v)| w * v).sum()
1347                    }
1348                }
1349            }
1350        }
1351    }
1352
1353    /// `out = W · x` (row-major). F32 delegates to the historical
1354    /// bit-exact path; Mapped runs the fused int8 kernel.
1355    pub fn matvec(&self, x: &[f32], out: &mut [f32], pool: Option<&Pool>) {
1356        match self {
1357            // NOTE: `out.len()` DRIVES this arm — it computes that many rows,
1358            // and `x.len()` is the stride. A short `out` is legitimate here,
1359            // which is why the check below lives in the Mapped arm only.
1360            Self::F32 { data, .. } => {
1361                if !crate::f32_backend::matvec(data, x, out) {
1362                    matvec_rows(pool, data, x, out);
1363                }
1364            }
1365            Self::Blr(t) => {
1366                if crate::gpu::enabled_here()
1367                    && crate::blr::gpu_wanted()
1368                    && crate::gpu::blr_matmat(t, &x[..t.cols], 1, &mut out[..t.rows])
1369                {
1370                    return;
1371                }
1372                t.matvec(x, out, pool)
1373            }
1374            Self::Mapped {
1375                model,
1376                idx,
1377                dtype,
1378                rows,
1379                cols,
1380                row_scale,
1381                col_field,
1382                vbit_offsets,
1383                repack,
1384            } => {
1385                let _ = (model, idx);
1386                // Every kernel below writes `rows` entries through a raw
1387                // pointer, so a short `out` is an out-of-bounds WRITE, not a
1388                // wrong answer: it scribbles on the allocator's metadata and
1389                // the process aborts much later, somewhere innocent
1390                // (`double free or corruption`, `corrupted double-linked
1391                // list`). The debug_assert two of the kernels carried is
1392                // compiled out of the release — exactly the build where it
1393                // matters. Fail here instead, while the caller is still on
1394                // the stack to be named.
1395                assert!(
1396                    out.len() >= *rows && x.len() >= *cols,
1397                    "matvec {rows}x{cols}: out {} (need {rows}), x {} (need {cols})",
1398                    out.len(),
1399                    x.len(),
1400                );
1401                let prism_forward =
1402                    crate::prism::is_forward_weight(model, &model.tensors[*idx].name);
1403                if *dtype == TensorDtype::Q2TiledP
1404                    && std::env::var("CMF_Q2TP_TRACE").as_deref() == Ok("1")
1405                {
1406                    use std::sync::atomic::{AtomicUsize, Ordering};
1407                    static N: AtomicUsize = AtomicUsize::new(0);
1408                    let n = N.fetch_add(1, Ordering::Relaxed);
1409                    if n < 128 {
1410                        eprintln!(
1411                            "q2tp-dispatch #{n} name={} prism={} rows={} cols={} gpu={} optin={} layer={}",
1412                            model.tensors[*idx].name,
1413                            prism_forward,
1414                            rows,
1415                            cols,
1416                            crate::gpu::enabled_here(),
1417                            crate::gpu::q2tp_gpu_opt_in(),
1418                            crate::gpu::cur_layer(),
1419                        );
1420                    }
1421                }
1422                // Prism stores every manifest-listed forward matrix in the
1423                // signed-Hadamard basis.  The q2tp WGSL path receives that
1424                // transformed vector and an explicit affine bit; codecs
1425                // without a descriptor-aware kernel remain on CPU below.
1426                if prism_forward {
1427                    let transformed = crate::prism::forward(model, &x[..*cols]);
1428                    match dtype {
1429                        TensorDtype::Q4Block => {
1430                            q4matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1431                        }
1432                        TensorDtype::Q4Tiled => {
1433                            q4t_matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1434                        }
1435                        TensorDtype::Q4TiledP => {
1436                            q4tp_matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1437                        }
1438                        TensorDtype::Q2TiledP => {
1439                            let affine =
1440                                crate::prism::is_affine_target(model, &model.tensors[*idx].name);
1441                            if *rows * *cols >= 8_388_608
1442                                && crate::gpu::enabled_here()
1443                                && crate::gpu::q2tp_gpu_opt_in()
1444                            {
1445                                let gpu_ok = if affine {
1446                                    crate::gpu::q2tp_affine_matvec(
1447                                        model,
1448                                        *idx,
1449                                        &transformed,
1450                                        *rows,
1451                                        *cols,
1452                                        out,
1453                                    )
1454                                } else {
1455                                    crate::gpu::q2tp_matvec(
1456                                        model,
1457                                        *idx,
1458                                        &transformed,
1459                                        *rows,
1460                                        *cols,
1461                                        out,
1462                                    )
1463                                };
1464                                if gpu_ok {
1465                                    return;
1466                                }
1467                            }
1468                            if affine {
1469                                q2tp_affine_matvec(
1470                                    self.quant_bytes(),
1471                                    &transformed,
1472                                    *rows,
1473                                    *cols,
1474                                    out,
1475                                    pool,
1476                                )
1477                            } else {
1478                                q2tp_matvec(
1479                                    self.quant_bytes(),
1480                                    &transformed,
1481                                    *rows,
1482                                    *cols,
1483                                    out,
1484                                    pool,
1485                                )
1486                            }
1487                        }
1488                        TensorDtype::Q1 | TensorDtype::Q1G => bin_matvec(
1489                            self.bin_w().expect("binary dtype"),
1490                            &transformed,
1491                            *rows,
1492                            *cols,
1493                            out,
1494                            pool,
1495                        ),
1496                        TensorDtype::Q1T => {
1497                            q1t_matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1498                        }
1499                        TensorDtype::Mxfp4 => {
1500                            mxfp4::matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1501                        }
1502                        TensorDtype::Vbit | TensorDtype::VbitRo => vbitmatvec(
1503                            self.quant_bytes(),
1504                            vbit_offsets,
1505                            &transformed,
1506                            *rows,
1507                            *cols,
1508                            out,
1509                            pool,
1510                        ),
1511                        TensorDtype::Q8Row | TensorDtype::Q8_2f => qmatvec(
1512                            self.quant_bytes(),
1513                            repack,
1514                            row_scale,
1515                            &transformed,
1516                            col_field,
1517                            *dtype,
1518                            *rows,
1519                            *cols,
1520                            out,
1521                            pool,
1522                        ),
1523                        _ => unreachable!("unsupported mapped Prism dtype {dtype:?}"),
1524                    }
1525                    return;
1526                }
1527                if *dtype == TensorDtype::Q4Block {
1528                    // GPU route (wgpu q4b kernel) for large q4_block matvecs —
1529                    // gives NVIDIA/AMD/Intel q4 models a GPU path. Probe keeps
1530                    // the winner; Metal returns false → the CPU kernel below.
1531                    if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1532                        let t0 = std::time::Instant::now();
1533                        match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
1534                            crate::gpu::ProbeArm::Gpu => {
1535                                if crate::gpu::q4b_matvec(model, *idx, x, *rows, *cols, out) {
1536                                    crate::gpu::probe_record(
1537                                        crate::gpu::OpClass::Matvec,
1538                                        true,
1539                                        t0.elapsed(),
1540                                    );
1541                                    return;
1542                                }
1543                            }
1544                            crate::gpu::ProbeArm::CpuTimed => {
1545                                q4matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1546                                crate::gpu::probe_record(
1547                                    crate::gpu::OpClass::Matvec,
1548                                    false,
1549                                    t0.elapsed(),
1550                                );
1551                                return;
1552                            }
1553                            crate::gpu::ProbeArm::Cpu => {}
1554                        }
1555                    }
1556                    q4matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1557                    return;
1558                }
1559                if *dtype == TensorDtype::Q4Tiled {
1560                    // GPU route for large q4t matvecs — the lm_head class,
1561                    // same shape as the q4tp arm below. The probe keeps the
1562                    // winner; a backend without the kernel refuses and the
1563                    // CPU path stays.
1564                    if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1565                        let t0 = std::time::Instant::now();
1566                        let cls = crate::gpu::matvec_class(*rows, *cols);
1567                        match crate::gpu::probe_arm(cls) {
1568                            crate::gpu::ProbeArm::Gpu => {
1569                                if crate::gpu::q4t_matvec(model, *idx, x, *rows, *cols, out) {
1570                                    crate::gpu::probe_record(cls, true, t0.elapsed());
1571                                    return;
1572                                }
1573                            }
1574                            crate::gpu::ProbeArm::CpuTimed => {
1575                                q4t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1576                                crate::gpu::probe_record(cls, false, t0.elapsed());
1577                                return;
1578                            }
1579                            crate::gpu::ProbeArm::Cpu => {}
1580                        }
1581                    }
1582                    q4t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1583                    return;
1584                }
1585                if *dtype == TensorDtype::Q4TiledP {
1586                    // GPU route for large q4tp matvecs — the lm_head class.
1587                    // On a q4tp checkpoint the head is the biggest single
1588                    // host matvec left in the decode step, and the batched
1589                    // kernel at b=1 already exists on both backends. Probe
1590                    // keeps the winner, same as q4_block above.
1591                    if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1592                        let t0 = std::time::Instant::now();
1593                        let cls = crate::gpu::matvec_class(*rows, *cols);
1594                        match crate::gpu::probe_arm(cls) {
1595                            crate::gpu::ProbeArm::Gpu => {
1596                                if crate::gpu::q4tp_matvec(model, *idx, x, *rows, *cols, out) {
1597                                    crate::gpu::probe_record(cls, true, t0.elapsed());
1598                                    return;
1599                                }
1600                            }
1601                            crate::gpu::ProbeArm::CpuTimed => {
1602                                q4tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1603                                crate::gpu::probe_record(cls, false, t0.elapsed());
1604                                return;
1605                            }
1606                            crate::gpu::ProbeArm::Cpu => {}
1607                        }
1608                    }
1609                    q4tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1610                    return;
1611                }
1612                if *dtype == TensorDtype::Q2TiledP {
1613                    q2tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1614                    return;
1615                }
1616                if *dtype == TensorDtype::Mxfp4 {
1617                    self.mxfp4_matvec(x, out, pool);
1618                    return;
1619                }
1620                if *dtype == TensorDtype::Q1 || *dtype == TensorDtype::Q1G {
1621                    // GPU route for large q1 matvecs (out_proj / lm_head
1622                    // class): the CPU q1 kernel is load-port-bound at
1623                    // ~4 GB/s/core, the GPU one is bandwidth-bound — the
1624                    // probe measures both arms and keeps the winner.
1625                    let w = self.bin_w().expect("binary dtype");
1626                    if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1627                        let t0 = std::time::Instant::now();
1628                        let arm = if crate::gpu::q1_force() {
1629                            crate::gpu::ProbeArm::Gpu
1630                        } else {
1631                            crate::gpu::probe_arm(crate::gpu::OpClass::Matvec)
1632                        };
1633                        match arm {
1634                            crate::gpu::ProbeArm::Gpu => {
1635                                let ok = if *dtype == TensorDtype::Q1G {
1636                                    crate::gpu::q1g_matvec(model, *idx, x, *rows, *cols, out)
1637                                } else {
1638                                    crate::gpu::q1_matvec(model, *idx, x, *rows, *cols, out)
1639                                };
1640                                if ok {
1641                                    crate::gpu::probe_record(
1642                                        crate::gpu::OpClass::Matvec,
1643                                        true,
1644                                        t0.elapsed(),
1645                                    );
1646                                    return;
1647                                }
1648                            }
1649                            crate::gpu::ProbeArm::CpuTimed => {
1650                                bin_matvec(w, x, *rows, *cols, out, pool);
1651                                crate::gpu::probe_record(
1652                                    crate::gpu::OpClass::Matvec,
1653                                    false,
1654                                    t0.elapsed(),
1655                                );
1656                                return;
1657                            }
1658                            crate::gpu::ProbeArm::Cpu => {}
1659                        }
1660                    }
1661                    bin_matvec(w, x, *rows, *cols, out, pool);
1662                    return;
1663                }
1664                if *dtype == TensorDtype::Q1T {
1665                    // GPU route for large q1t matvecs: the ternary BASE dot runs
1666                    // on the GPU (load-port-bound on CPU, like q1), then the
1667                    // sparse overlay is added on the CPU. Probe keeps the winner.
1668                    if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1669                        let t0 = std::time::Instant::now();
1670                        match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
1671                            crate::gpu::ProbeArm::Gpu => {
1672                                if crate::gpu::q1t_matvec(model, *idx, x, *rows, *cols, out) {
1673                                    q1t_add_overlay(self.quant_bytes(), x, *rows, *cols, out, pool);
1674                                    crate::gpu::probe_record(
1675                                        crate::gpu::OpClass::Matvec,
1676                                        true,
1677                                        t0.elapsed(),
1678                                    );
1679                                    return;
1680                                }
1681                            }
1682                            crate::gpu::ProbeArm::CpuTimed => {
1683                                q1t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1684                                crate::gpu::probe_record(
1685                                    crate::gpu::OpClass::Matvec,
1686                                    false,
1687                                    t0.elapsed(),
1688                                );
1689                                return;
1690                            }
1691                            crate::gpu::ProbeArm::Cpu => {}
1692                        }
1693                    }
1694                    q1t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1695                    return;
1696                }
1697                if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1698                    vbitmatvec(self.quant_bytes(), vbit_offsets, x, *rows, *cols, out, pool);
1699                    return;
1700                }
1701                let xs = prescale(x, col_field, *dtype);
1702                // D5: large q8 matrices (lm_head-class) — hybrid
1703                // CPU∥GPU: split the rows, both sides compute
1704                // SIMULTANEOUSLY (same math, shared prescale).
1705                // GPU share: CMF_GPU_SPLIT (0..1, default 0.5).
1706                if *rows >= crate::gpu::min_rows()
1707                    && matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
1708                    && gpu_lmhead_enabled()
1709                    && crate::gpu::enabled_here()
1710                {
1711                    // Runtime probe: alternate the hybrid against the
1712                    // pure-CPU matvec, keep whichever is faster HERE.
1713                    let t0 = std::time::Instant::now();
1714                    match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
1715                        crate::gpu::ProbeArm::Gpu => {}
1716                        crate::gpu::ProbeArm::CpuTimed => {
1717                            qmatvec(
1718                                self.quant_bytes(),
1719                                repack,
1720                                row_scale,
1721                                x,
1722                                col_field,
1723                                *dtype,
1724                                *rows,
1725                                *cols,
1726                                out,
1727                                pool,
1728                            );
1729                            crate::gpu::probe_record(
1730                                crate::gpu::OpClass::Matvec,
1731                                false,
1732                                t0.elapsed(),
1733                            );
1734                            return;
1735                        }
1736                        crate::gpu::ProbeArm::Cpu => {
1737                            qmatvec(
1738                                self.quant_bytes(),
1739                                repack,
1740                                row_scale,
1741                                x,
1742                                col_field,
1743                                *dtype,
1744                                *rows,
1745                                *cols,
1746                                out,
1747                                pool,
1748                            );
1749                            return;
1750                        }
1751                    }
1752                    let frac = gpu_split_frac();
1753                    let cpu_rows = ((*rows as f32) * (1.0 - frac)) as usize;
1754                    let (out_cpu, out_gpu) = out.split_at_mut(cpu_rows);
1755                    let bytes = self.quant_bytes();
1756                    let ok = std::thread::scope(|sc| {
1757                        let g = sc.spawn(|| {
1758                            crate::gpu::q8_matvec_range(
1759                                model,
1760                                *idx,
1761                                cpu_rows,
1762                                &row_scale[cpu_rows..],
1763                                &xs,
1764                                *rows - cpu_rows,
1765                                *cols,
1766                                out_gpu,
1767                            )
1768                        });
1769                        if cpu_rows > 0 {
1770                            // Repack prefix covers the full groups of the
1771                            // CPU half (the split starts at row 0).
1772                            let rep_cpu = if repack.is_empty() {
1773                                &[][..]
1774                            } else {
1775                                &repack[..(cpu_rows / 4) * 4 * *cols]
1776                            };
1777                            qmatvec(
1778                                &bytes[..cpu_rows * *cols],
1779                                rep_cpu,
1780                                &row_scale[..cpu_rows],
1781                                x,
1782                                col_field,
1783                                *dtype,
1784                                cpu_rows,
1785                                *cols,
1786                                out_cpu,
1787                                pool,
1788                            );
1789                        }
1790                        g.join().unwrap_or(false)
1791                    });
1792                    if ok {
1793                        crate::gpu::probe_record(crate::gpu::OpClass::Matvec, true, t0.elapsed());
1794                        return;
1795                    }
1796                    // GPU failed — CPU finishes its half (rows rebased —
1797                    // group offsets don't line up, mmap layout only).
1798                    qmatvec(
1799                        &bytes[cpu_rows * *cols..(*rows) * *cols],
1800                        &[],
1801                        &row_scale[cpu_rows..],
1802                        x,
1803                        col_field,
1804                        *dtype,
1805                        *rows - cpu_rows,
1806                        *cols,
1807                        out_gpu,
1808                        pool,
1809                    );
1810                    return;
1811                }
1812                qmatvec(
1813                    self.quant_bytes(),
1814                    repack,
1815                    row_scale,
1816                    x,
1817                    col_field,
1818                    *dtype,
1819                    *rows,
1820                    *cols,
1821                    out,
1822                    pool,
1823                );
1824            }
1825        }
1826    }
1827
1828    /// Fused two-input matvec (MTP verify pair): weights streamed once.
1829    pub fn matvec2(
1830        &self,
1831        x1: &[f32],
1832        x2: &[f32],
1833        o1: &mut [f32],
1834        o2: &mut [f32],
1835        pool: Option<&Pool>,
1836    ) {
1837        match self {
1838            Self::F32 { data, .. } => matvec_rows2(pool, data, x1, x2, o1, o2),
1839            Self::Blr(_) => {
1840                self.matvec(x1, o1, pool);
1841                self.matvec(x2, o2, pool);
1842            }
1843            Self::Mapped {
1844                model,
1845                idx,
1846                dtype,
1847                rows,
1848                cols,
1849                row_scale,
1850                col_field,
1851                vbit_offsets,
1852                ..
1853            } => {
1854                if crate::prism::is_forward_weight(model, &model.tensors[*idx].name) {
1855                    let tx1 = crate::prism::forward(model, &x1[..*cols]);
1856                    let tx2 = crate::prism::forward(model, &x2[..*cols]);
1857                    match dtype {
1858                        TensorDtype::Q4Block => {
1859                            q4matvec2(self.quant_bytes(), &tx1, &tx2, *rows, *cols, o1, o2, pool)
1860                        }
1861                        TensorDtype::Q4Tiled => {
1862                            q4t_matvec2(self.quant_bytes(), &tx1, &tx2, *rows, *cols, o1, o2, pool)
1863                        }
1864                        TensorDtype::Q4TiledP => {
1865                            q4tp_matvec2(self.quant_bytes(), &tx1, &tx2, *rows, *cols, o1, o2, pool)
1866                        }
1867                        TensorDtype::Q2TiledP => {
1868                            if crate::prism::is_affine_target(model, &model.tensors[*idx].name) {
1869                                q2tp_affine_matvec2(
1870                                    self.quant_bytes(),
1871                                    &tx1,
1872                                    &tx2,
1873                                    *rows,
1874                                    *cols,
1875                                    o1,
1876                                    o2,
1877                                    pool,
1878                                )
1879                            } else {
1880                                q2tp_matvec2(
1881                                    self.quant_bytes(),
1882                                    &tx1,
1883                                    &tx2,
1884                                    *rows,
1885                                    *cols,
1886                                    o1,
1887                                    o2,
1888                                    pool,
1889                                )
1890                            }
1891                        }
1892                        TensorDtype::Q1 | TensorDtype::Q1G => bin_matvec2(
1893                            self.bin_w().expect("binary dtype"),
1894                            &tx1,
1895                            &tx2,
1896                            *rows,
1897                            *cols,
1898                            o1,
1899                            o2,
1900                            pool,
1901                        ),
1902                        TensorDtype::Q1T => {
1903                            q1t_matvec2(self.quant_bytes(), &tx1, &tx2, *rows, *cols, o1, o2, pool)
1904                        }
1905                        TensorDtype::Mxfp4 => mxfp4::matvec2(
1906                            self.quant_bytes(),
1907                            &tx1,
1908                            &tx2,
1909                            *rows,
1910                            *cols,
1911                            o1,
1912                            o2,
1913                            pool,
1914                        ),
1915                        TensorDtype::Vbit | TensorDtype::VbitRo => vbitmatvec2(
1916                            self.quant_bytes(),
1917                            vbit_offsets,
1918                            &tx1,
1919                            &tx2,
1920                            *rows,
1921                            *cols,
1922                            o1,
1923                            o2,
1924                            pool,
1925                        ),
1926                        TensorDtype::Q8Row | TensorDtype::Q8_2f => qmatvec2(
1927                            self.quant_bytes(),
1928                            row_scale,
1929                            &tx1,
1930                            &tx2,
1931                            col_field,
1932                            *dtype,
1933                            *rows,
1934                            *cols,
1935                            o1,
1936                            o2,
1937                            pool,
1938                        ),
1939                        _ => unreachable!("unsupported mapped Prism dtype {dtype:?}"),
1940                    }
1941                    return;
1942                }
1943                if *dtype == TensorDtype::Q4Block {
1944                    q4matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1945                    return;
1946                }
1947                if *dtype == TensorDtype::Q4Tiled {
1948                    q4t_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1949                    return;
1950                }
1951                if *dtype == TensorDtype::Q4TiledP {
1952                    q4tp_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1953                    return;
1954                }
1955                if *dtype == TensorDtype::Q2TiledP {
1956                    q2tp_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1957                    return;
1958                }
1959                if *dtype == TensorDtype::Mxfp4 {
1960                    mxfp4::matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1961                    return;
1962                }
1963                if let Some(w) = self.bin_w() {
1964                    bin_matvec2(w, x1, x2, *rows, *cols, o1, o2, pool);
1965                    return;
1966                }
1967                if *dtype == TensorDtype::Q1T {
1968                    // Fused ternary pair: one row pass, the register
1969                    // unpack shared across both streams on ARM. (Q1T
1970                    // lacks a row_scale array — scales live inline in
1971                    // the tiles — so it must not fall through to the
1972                    // q8 qmatvec2 below.)
1973                    q1t_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1974                    return;
1975                }
1976                if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1977                    vbitmatvec2(
1978                        self.quant_bytes(),
1979                        vbit_offsets,
1980                        x1,
1981                        x2,
1982                        *rows,
1983                        *cols,
1984                        o1,
1985                        o2,
1986                        pool,
1987                    );
1988                    return;
1989                }
1990                qmatvec2(
1991                    self.quant_bytes(),
1992                    row_scale,
1993                    x1,
1994                    x2,
1995                    col_field,
1996                    *dtype,
1997                    *rows,
1998                    *cols,
1999                    o1,
2000                    o2,
2001                    pool,
2002                );
2003            }
2004        }
2005    }
2006}
2007
2008impl QTensor {
2009    /// Batched matvec (prefill-GEMM): xs — row-major [b, cols],
2010    /// out — row-major [b, rows]. Element-wise semantics are IDENTICAL
2011    /// to b matvec calls (same dot kernels in the same order); the win —
2012    /// the weight row streams from DRAM once per batch, not b times.
2013    /// `(model, index)` when this is a memory-mapped q4tp tensor — the
2014    /// identity a device-resident chain needs to hand `tp_matmat` the
2015    /// weight without going through this struct's own dispatch.
2016    pub fn q4tp_mapped(&self) -> Option<(&std::sync::Arc<CmfModel>, usize)> {
2017        if self.has_prism_contract() {
2018            return None;
2019        }
2020        match self {
2021            Self::Mapped {
2022                model, idx, dtype, ..
2023            } if *dtype == TensorDtype::Q4TiledP => Some((model, *idx)),
2024            _ => None,
2025        }
2026    }
2027
2028    pub fn matmat(&self, xs_all: &[f32], b: usize, out: &mut [f32], pool: Option<&Pool>) {
2029        let cols = self.cols();
2030        let rows = self.rows();
2031        debug_assert_eq!(xs_all.len(), b * cols);
2032        debug_assert_eq!(out.len(), b * rows);
2033        let _prof = crate::cpuprof::time(crate::cpuprof::Slot::Matmat);
2034        // GPTQ calibration: fold this layer's inputs into its Hessian. Only
2035        // Mapped tensors carry a directory name; the check is a relaxed
2036        // atomic load, free when not calibrating.
2037        if crate::gptq_capture::capturing() {
2038            if let Self::Mapped { model, idx, .. } = self {
2039                crate::gptq_capture::accumulate(&model.tensors[*idx].name, xs_all, b, cols);
2040            }
2041            if let Self::Blr(t) = self {
2042                crate::gptq_capture::accumulate(t.name(), xs_all, b, cols);
2043            }
2044        }
2045        match self {
2046            Self::Blr(t) => {
2047                if crate::gpu::enabled_here()
2048                    && crate::blr::gpu_wanted()
2049                    && crate::gpu::blr_matmat(t, &xs_all[..b * cols], b, &mut out[..b * rows])
2050                {
2051                    return;
2052                }
2053                t.matmat(xs_all, b, out, pool)
2054            }
2055            Self::F32 { data, .. } => {
2056                if crate::f32_backend::matmat(data, xs_all, b, rows, cols, out) {
2057                    return;
2058                }
2059                let out_addr = SendMut(out.as_mut_ptr());
2060                // Over the batch, four positions a row at a time: four
2061                // independent chains, each still summing j = 0..cols in the
2062                // scalar order (bit-identical to the one-at-a-time loop and
2063                // to `matvec_rows`). The batch is the wide axis of a prefill;
2064                // the rows may be few (Spark-X2.5's 16-row g_proj spent
2065                // 17 ms a layer here split over its rows).
2066                let run = |start: usize, end: usize| {
2067                    let mut bi = start;
2068                    while bi < end {
2069                        let n = (end - bi).min(4);
2070                        for o in 0..rows {
2071                            let row = &data[o * cols..(o + 1) * cols];
2072                            if n == 4 {
2073                                let x = |k: usize| &xs_all[(bi + k) * cols..(bi + k + 1) * cols];
2074                                let (x0, x1, x2, x3) = (x(0), x(1), x(2), x(3));
2075                                let (mut a0, mut a1, mut a2, mut a3) = (0f32, 0f32, 0f32, 0f32);
2076                                for j in 0..cols {
2077                                    let w = row[j];
2078                                    a0 += w * x0[j];
2079                                    a1 += w * x1[j];
2080                                    a2 += w * x2[j];
2081                                    a3 += w * x3[j];
2082                                }
2083                                for (k, acc) in [a0, a1, a2, a3].into_iter().enumerate() {
2084                                    unsafe { *out_addr.at((bi + k) * rows + o) = acc };
2085                                }
2086                            } else {
2087                                for k in 0..n {
2088                                    let x = &xs_all[(bi + k) * cols..(bi + k + 1) * cols];
2089                                    let mut acc = 0f32;
2090                                    for j in 0..cols {
2091                                        acc += row[j] * x[j];
2092                                    }
2093                                    unsafe { *out_addr.at((bi + k) * rows + o) = acc };
2094                                }
2095                            }
2096                        }
2097                        bi += n;
2098                    }
2099                };
2100                dispatch_rows(pool, b, &run);
2101            }
2102            Self::Mapped {
2103                model,
2104                idx,
2105                dtype,
2106                row_scale,
2107                col_field,
2108                vbit_offsets,
2109                ..
2110            } => {
2111                if crate::prism::is_forward_weight(model, &model.tensors[*idx].name) {
2112                    let mut transformed = Vec::with_capacity(xs_all.len());
2113                    for bi in 0..b {
2114                        transformed.extend_from_slice(&crate::prism::forward(
2115                            model,
2116                            &xs_all[bi * cols..(bi + 1) * cols],
2117                        ));
2118                    }
2119                    match dtype {
2120                        TensorDtype::Q4Block => {
2121                            q4matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
2122                        }
2123                        TensorDtype::Q4Tiled => {
2124                            q4t_matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
2125                        }
2126                        TensorDtype::Q4TiledP => {
2127                            q4tp_matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
2128                        }
2129                        TensorDtype::Q2TiledP => {
2130                            let affine =
2131                                crate::prism::is_affine_target(model, &model.tensors[*idx].name);
2132                            // Affine Prism Q2TP has a descriptor-aware GPU
2133                            // kernel for short/tail batches too.  Unlike the
2134                            // ordinary Q2TP path, don't force b<32 back to a
2135                            // scalar CPU matmat: prefill chunks and the final
2136                            // tail both need to stay on the tested GPU arm.
2137                            let gpu_batch_ok = if affine {
2138                                b >= 2
2139                            } else {
2140                                b >= 32 && b * rows * cols >= 128_000_000
2141                            };
2142                            if gpu_batch_ok
2143                                && cols % 32 == 0
2144                                && crate::gpu::enabled_here()
2145                                && crate::gpu::q2tp_gpu_opt_in()
2146                            {
2147                                let gpu_ok = if affine {
2148                                    crate::gpu::q2tp_affine_matmat(
2149                                        model,
2150                                        *idx,
2151                                        &transformed,
2152                                        b,
2153                                        rows,
2154                                        cols,
2155                                        out,
2156                                    )
2157                                } else {
2158                                    crate::gpu::q2tp_matmat(
2159                                        model,
2160                                        *idx,
2161                                        &transformed,
2162                                        b,
2163                                        rows,
2164                                        cols,
2165                                        out,
2166                                    )
2167                                };
2168                                if gpu_ok {
2169                                    return;
2170                                }
2171                            }
2172                            // A one-token Prism decode is the other short
2173                            // case.  Use the descriptor-aware matvec kernel
2174                            // before falling back to the exact CPU path.
2175                            if affine
2176                                && b == 1
2177                                && cols % 32 == 0
2178                                && crate::gpu::enabled_here()
2179                                && crate::gpu::q2tp_gpu_opt_in()
2180                                && crate::gpu::q2tp_affine_matvec(
2181                                    model,
2182                                    *idx,
2183                                    &transformed[..cols],
2184                                    rows,
2185                                    cols,
2186                                    &mut out[..rows],
2187                                )
2188                            {
2189                                return;
2190                            }
2191                            if affine {
2192                                q2tp_affine_matmat(
2193                                    self.quant_bytes(),
2194                                    &transformed,
2195                                    b,
2196                                    rows,
2197                                    cols,
2198                                    out,
2199                                    pool,
2200                                )
2201                            } else {
2202                                q2tp_matmat(
2203                                    self.quant_bytes(),
2204                                    &transformed,
2205                                    b,
2206                                    rows,
2207                                    cols,
2208                                    out,
2209                                    pool,
2210                                )
2211                            }
2212                        }
2213                        TensorDtype::Q1 | TensorDtype::Q1G => bin_matmat(
2214                            self.bin_w().expect("binary dtype"),
2215                            &transformed,
2216                            b,
2217                            rows,
2218                            cols,
2219                            out,
2220                            pool,
2221                        ),
2222                        TensorDtype::Q1T => {
2223                            q1t_matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
2224                        }
2225                        TensorDtype::Mxfp4 => {
2226                            mxfp4::matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
2227                        }
2228                        TensorDtype::Vbit | TensorDtype::VbitRo => vbitmatmat(
2229                            self.quant_bytes(),
2230                            vbit_offsets,
2231                            &transformed,
2232                            b,
2233                            rows,
2234                            cols,
2235                            out,
2236                            pool,
2237                        ),
2238                        TensorDtype::Q8Row | TensorDtype::Q8_2f => {
2239                            let pre: Vec<std::borrow::Cow<'_, [f32]>> = (0..b)
2240                                .map(|bi| {
2241                                    prescale(
2242                                        &transformed[bi * cols..(bi + 1) * cols],
2243                                        col_field,
2244                                        *dtype,
2245                                    )
2246                                })
2247                                .collect();
2248                            qmatmat(self.quant_bytes(), row_scale, &pre, rows, cols, out, pool)
2249                        }
2250                        _ => unreachable!("unsupported mapped Prism dtype {dtype:?}"),
2251                    }
2252                    return;
2253                }
2254                if *dtype == TensorDtype::Q4Block {
2255                    q4matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2256                    return;
2257                }
2258                if *dtype == TensorDtype::Mxfp4 {
2259                    self.mxfp4_matmat(xs_all, b, out, pool);
2260                    return;
2261                }
2262                if *dtype == TensorDtype::Q4TiledP {
2263                    // GPU batched q4tp GEMM (dequant + f32nt mul_mm on the
2264                    // device); the probe keeps whichever beats the CPU arm.
2265                    // Narrow (prompt-encode) and wide (DiT) batches probe
2266                    // as separate classes — the regimes have opposite
2267                    // winners and one shared verdict locked the wrong arm.
2268                    // Kill switch (gpu::mm_kill): one grossly slow GPU op
2269                    // (a fair-condition op is ≤~100 ms even at 1024px)
2270                    // means the device is contended by another process
2271                    // (e.g. a simulator) — verdicts are per-process, so
2272                    // without the bail the whole render crawls behind
2273                    // someone else's queue.
2274                    // Row-exact batches stay on the host: the device GEMM
2275                    // is an f32 dequant-sgemm, not the host matvec's sum.
2276                    if b >= 32
2277                        && b * rows * cols >= 128_000_000
2278                        && cols % 32 == 0
2279                        && !row_exact()
2280                        && !crate::gpu::mm_killed()
2281                        && crate::gpu::enabled_here()
2282                    {
2283                        let class = if b >= 128 {
2284                            crate::gpu::OpClass::MatmatWide
2285                        } else {
2286                            crate::gpu::OpClass::Matmat
2287                        };
2288                        if let Self::Mapped { model, idx, .. } = self {
2289                            // In-process A/B (`CMF_MM_AB=1`). Three
2290                            // wall-clock A/Bs on a shared stand disagreed
2291                            // with each other by 25% on the same change,
2292                            // because the machine drifts between processes
2293                            // and interleaving whole renders does not fix
2294                            // that. Here both arms run back to back on the
2295                            // SAME data inside one call, so whatever the
2296                            // machine is doing, it does to both — and the
2297                            // disagreement between their outputs falls out
2298                            // for free. Doubles the work; a diagnostic,
2299                            // not a mode.
2300                            if crate::mm_ab::on() {
2301                                let mut g = vec![0f32; b * rows];
2302                                let t = std::time::Instant::now();
2303                                let took = crate::gpu::q4tp_matmat(
2304                                    model, *idx, xs_all, b, rows, cols, &mut g,
2305                                );
2306                                let dg = t.elapsed();
2307                                let t = std::time::Instant::now();
2308                                q4tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2309                                let dc = t.elapsed();
2310                                crate::mm_ab::record(b, rows, cols, took, dg, dc, &g, out);
2311                                return;
2312                            }
2313                            let t0 = std::time::Instant::now();
2314                            // A cold call takes the device arm: its sample
2315                            // is discarded either way, and the upload is
2316                            // what the next step needs.
2317                            let resident = crate::gpu::weight_is_resident(model, *idx);
2318                            match crate::gpu::probe_arm_cold_prefers_gpu(class, resident) {
2319                                crate::gpu::ProbeArm::Gpu => {
2320                                    if crate::gpu::q4tp_matmat(
2321                                        model, *idx, xs_all, b, rows, cols, out,
2322                                    ) {
2323                                        let el = t0.elapsed();
2324                                        // Work-proportional budget: ~8× the
2325                                        // fair-device estimate (+20 ms slack).
2326                                        // An absolute cap missed the worst
2327                                        // case — contended ops sit at
2328                                        // 100–240 ms each and still bury a
2329                                        // render whose fair op is 3–9 ms.
2330                                        // Cold ops (first PSO build, buffer
2331                                        // alloc) are exempt: a one-off
2332                                        // ~50 ms compile is not contention.
2333                                        let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
2334                                        let budget = std::time::Duration::from_secs_f64(
2335                                            flops / 1.5e12 * 8.0 + 0.020,
2336                                        );
2337                                        crate::gpu::mm_budget_check(
2338                                            "q4tp matmat",
2339                                            el,
2340                                            budget,
2341                                            crate::gpu::probe_was_cold() || !resident,
2342                                        );
2343                                        crate::gpu::probe_record(class, true, el);
2344                                        return;
2345                                    }
2346                                }
2347                                crate::gpu::ProbeArm::CpuTimed => {
2348                                    q4tp_matmat(
2349                                        self.quant_bytes(),
2350                                        xs_all,
2351                                        b,
2352                                        rows,
2353                                        cols,
2354                                        out,
2355                                        pool,
2356                                    );
2357                                    crate::gpu::probe_record(class, false, t0.elapsed());
2358                                    return;
2359                                }
2360                                crate::gpu::ProbeArm::Cpu => {}
2361                            }
2362                        }
2363                    }
2364                    q4tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2365                    return;
2366                }
2367                if *dtype == TensorDtype::Q2TiledP {
2368                    // Same device arm as q4tp, behind the same probe:
2369                    // the planes differ, the dispatch does not. Without
2370                    // this a q2tp file ran its widest projections on the
2371                    // host while the 4-bit one had the card, which is a
2372                    // codec paying for its size twice.
2373                    if b >= 32
2374                        && b * rows * cols >= 128_000_000
2375                        && cols % 32 == 0
2376                        && !crate::gpu::mm_killed()
2377                        && crate::gpu::enabled_here()
2378                    {
2379                        let class = if b >= 128 {
2380                            crate::gpu::OpClass::MatmatWide
2381                        } else {
2382                            crate::gpu::OpClass::Matmat
2383                        };
2384                        if let Self::Mapped { model, idx, .. } = self {
2385                            let t0 = std::time::Instant::now();
2386                            match crate::gpu::probe_arm(class) {
2387                                crate::gpu::ProbeArm::Gpu => {
2388                                    if crate::gpu::q2tp_matmat(
2389                                        model, *idx, xs_all, b, rows, cols, out,
2390                                    ) {
2391                                        crate::gpu::probe_record(class, true, t0.elapsed());
2392                                        return;
2393                                    }
2394                                }
2395                                crate::gpu::ProbeArm::CpuTimed => {
2396                                    q2tp_matmat(
2397                                        self.quant_bytes(),
2398                                        xs_all,
2399                                        b,
2400                                        rows,
2401                                        cols,
2402                                        out,
2403                                        pool,
2404                                    );
2405                                    crate::gpu::probe_record(class, false, t0.elapsed());
2406                                    return;
2407                                }
2408                                crate::gpu::ProbeArm::Cpu => {}
2409                            }
2410                        }
2411                    }
2412                    // Without a host arm a q2tp tensor falls through to
2413                    // the q8 fallback, which reads it at one BYTE per
2414                    // weight — a 2x overrun that killed pool workers
2415                    // mid-prefill while the dispatcher waited forever.
2416                    q2tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2417                    return;
2418                }
2419                if *dtype == TensorDtype::Q4Tiled {
2420                    // GPU batched q4t GEMM (dequant + f32nt mul_mm on the
2421                    // device); the probe keeps whichever beats the CPU arm.
2422                    // Narrow (prompt-encode) and wide (DiT) batches probe
2423                    // as separate classes — the regimes have opposite
2424                    // winners and one shared verdict locked the wrong arm.
2425                    // Kill switch (gpu::mm_kill): one grossly slow GPU op
2426                    // (a fair-condition op is ≤~100 ms even at 1024px)
2427                    // means the device is contended by another process
2428                    // (e.g. a simulator) — verdicts are per-process, so
2429                    // without the bail the whole render crawls behind
2430                    // someone else's queue.
2431                    if b >= 32
2432                        && b * rows * cols >= 128_000_000
2433                        && cols % 32 == 0
2434                        && !crate::gpu::mm_killed()
2435                        && crate::gpu::enabled_here()
2436                    {
2437                        let class = if b >= 128 {
2438                            crate::gpu::OpClass::MatmatWide
2439                        } else {
2440                            crate::gpu::OpClass::Matmat
2441                        };
2442                        if let Self::Mapped { model, idx, .. } = self {
2443                            let t0 = std::time::Instant::now();
2444                            match crate::gpu::probe_arm(class) {
2445                                crate::gpu::ProbeArm::Gpu => {
2446                                    if crate::gpu::q4t_matmat(
2447                                        model, *idx, xs_all, b, rows, cols, out,
2448                                    ) {
2449                                        let el = t0.elapsed();
2450                                        // Work-proportional budget: ~8× the
2451                                        // fair-device estimate (+20 ms slack).
2452                                        // An absolute cap missed the worst
2453                                        // case — contended ops sit at
2454                                        // 100–240 ms each and still bury a
2455                                        // render whose fair op is 3–9 ms.
2456                                        // Cold ops (first PSO build, buffer
2457                                        // alloc) are exempt: a one-off
2458                                        // ~50 ms compile is not contention.
2459                                        let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
2460                                        let budget = std::time::Duration::from_secs_f64(
2461                                            flops / 1.5e12 * 8.0 + 0.020,
2462                                        );
2463                                        crate::gpu::mm_budget_check(
2464                                            "q4t matmat",
2465                                            el,
2466                                            budget,
2467                                            crate::gpu::probe_was_cold(),
2468                                        );
2469                                        crate::gpu::probe_record(class, true, el);
2470                                        return;
2471                                    }
2472                                }
2473                                crate::gpu::ProbeArm::CpuTimed => {
2474                                    q4t_matmat(
2475                                        self.quant_bytes(),
2476                                        xs_all,
2477                                        b,
2478                                        rows,
2479                                        cols,
2480                                        out,
2481                                        pool,
2482                                    );
2483                                    crate::gpu::probe_record(class, false, t0.elapsed());
2484                                    return;
2485                                }
2486                                crate::gpu::ProbeArm::Cpu => {}
2487                            }
2488                        }
2489                    }
2490                    q4t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2491                    return;
2492                }
2493                if *dtype == TensorDtype::Q1 || *dtype == TensorDtype::Q1G {
2494                    // GPU batched q1 GEMM for wide prefill (q1_mul_mm on the
2495                    // device); the probe keeps whichever beats the CPU matmat.
2496                    // q1g has the twin `q1g_mul_mm` (cols % 128 by format).
2497                    let w = self.bin_w().expect("binary dtype");
2498                    if b >= 32
2499                        && b * rows * cols >= 128_000_000
2500                        && cols % 64 == 0
2501                        && crate::gpu::enabled_here()
2502                    {
2503                        if let Self::Mapped { model, idx, .. } = self {
2504                            let t0 = std::time::Instant::now();
2505                            match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
2506                                crate::gpu::ProbeArm::Gpu => {
2507                                    let ok = if *dtype == TensorDtype::Q1G {
2508                                        crate::gpu::q1g_matmat(
2509                                            model, *idx, xs_all, b, rows, cols, out,
2510                                        )
2511                                    } else {
2512                                        crate::gpu::q1_matmat(
2513                                            model, *idx, xs_all, b, rows, cols, out,
2514                                        )
2515                                    };
2516                                    if ok {
2517                                        crate::gpu::probe_record(
2518                                            crate::gpu::OpClass::Matmat,
2519                                            true,
2520                                            t0.elapsed(),
2521                                        );
2522                                        return;
2523                                    }
2524                                }
2525                                crate::gpu::ProbeArm::CpuTimed => {
2526                                    bin_matmat(w, xs_all, b, rows, cols, out, pool);
2527                                    crate::gpu::probe_record(
2528                                        crate::gpu::OpClass::Matmat,
2529                                        false,
2530                                        t0.elapsed(),
2531                                    );
2532                                    return;
2533                                }
2534                                crate::gpu::ProbeArm::Cpu => {}
2535                            }
2536                        }
2537                    }
2538                    bin_matmat(w, xs_all, b, rows, cols, out, pool);
2539                    return;
2540                }
2541                if *dtype == TensorDtype::Q1T {
2542                    // GPU batched GEMM for wide prefill (base + overlay on the
2543                    // device); probe keeps the winner vs the CPU matmat.
2544                    if b >= 32 && b * rows * cols >= 128_000_000 && crate::gpu::enabled_here() {
2545                        if let Self::Mapped { model, idx, .. } = self {
2546                            let t0 = std::time::Instant::now();
2547                            match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
2548                                crate::gpu::ProbeArm::Gpu => {
2549                                    if crate::gpu::q1t_matmat(
2550                                        model, *idx, xs_all, b, rows, cols, out,
2551                                    ) {
2552                                        crate::gpu::probe_record(
2553                                            crate::gpu::OpClass::Matmat,
2554                                            true,
2555                                            t0.elapsed(),
2556                                        );
2557                                        return;
2558                                    }
2559                                }
2560                                crate::gpu::ProbeArm::CpuTimed => {
2561                                    q1t_matmat(
2562                                        self.quant_bytes(),
2563                                        xs_all,
2564                                        b,
2565                                        rows,
2566                                        cols,
2567                                        out,
2568                                        pool,
2569                                    );
2570                                    crate::gpu::probe_record(
2571                                        crate::gpu::OpClass::Matmat,
2572                                        false,
2573                                        t0.elapsed(),
2574                                    );
2575                                    return;
2576                                }
2577                                crate::gpu::ProbeArm::Cpu => {}
2578                            }
2579                        }
2580                    }
2581                    q1t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2582                    return;
2583                }
2584                if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
2585                    vbitmatmat(
2586                        self.quant_bytes(),
2587                        vbit_offsets,
2588                        xs_all,
2589                        b,
2590                        rows,
2591                        cols,
2592                        out,
2593                        pool,
2594                    );
2595                    return;
2596                }
2597                let pre: Vec<std::borrow::Cow<'_, [f32]>> = (0..b)
2598                    .map(|bi| prescale(&xs_all[bi * cols..(bi + 1) * cols], col_field, *dtype))
2599                    .collect();
2600                // MiMo verification is a 2–4 row decode panel, not a wide
2601                // prompt GEMM. Keep q8 projections on the same device as
2602                // decode; the generic b>=8 gate otherwise silently moves
2603                // every projection back to CPU. The short wgpu matmat uses
2604                // the same 64-lane reduction as its single-token matvec.
2605                if row_exact()
2606                    && (1..=4).contains(&b)
2607                    && matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
2608                    && crate::gpu::enabled_here()
2609                    && crate::gpu::wgpu_active()
2610                {
2611                    let flat: Vec<f32> = pre.iter().flat_map(|v| v.iter().copied()).collect();
2612                    if crate::gpu::q8_matmat(model, *idx, row_scale, &flat, b, rows, cols, out) {
2613                        return;
2614                    }
2615                }
2616                // D5: large prefill-batch GEMMs — on the GPU (threshold by
2617                // work volume: submission carries b×rows×cols MACs).
2618                // Runtime probe: the naive GEMM shader + sync readback
2619                // lose to the CPU GEMM on slow driver stacks — alternate
2620                // both arms and keep the winner.
2621                if b >= 8 && b * rows * cols >= 128_000_000 && crate::gpu::enabled_here() {
2622                    if let Self::Mapped { model, idx, .. } = self {
2623                        let t0 = std::time::Instant::now();
2624                        match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
2625                            crate::gpu::ProbeArm::Gpu
2626                                if crate::gpu::probe_deciding(crate::gpu::OpClass::Matmat)
2627                                    && !crate::gpu::q8_resident_or_upload(model, *idx) =>
2628                            {
2629                                // Cold weights during probing: the upload
2630                                // has started, the count runs on the CPU —
2631                                // the GPU arm samples on the next touch.
2632                                let q = self.quant_bytes();
2633                                qmatmat(q, row_scale, &pre, rows, cols, out, pool);
2634                                return;
2635                            }
2636                            crate::gpu::ProbeArm::Gpu => {
2637                                // The two-field codec's own entry first, as
2638                                // `device_matmat` takes it: the column field
2639                                // folds into the weight plane and the GEMM
2640                                // runs on the matrix units. The plain int8
2641                                // entry below is the scalar f32 GEMM — on a
2642                                // Spark-X2.5 4B prefill it ran q|k|v|o at
2643                                // host speed.
2644                                if *dtype == TensorDtype::Q8_2f
2645                                    && std::env::var("CMF_Q8_2F_DEV").as_deref() != Ok("0")
2646                                    && crate::gpu::q8_matmat_2f(
2647                                        model, *idx, row_scale, col_field, xs_all, b, rows,
2648                                        cols, out,
2649                                    )
2650                                {
2651                                    crate::gpu::probe_record(
2652                                        crate::gpu::OpClass::Matmat,
2653                                        true,
2654                                        t0.elapsed(),
2655                                    );
2656                                    return;
2657                                }
2658                                let flat: Vec<f32> =
2659                                    pre.iter().flat_map(|v| v.iter().copied()).collect();
2660                                if crate::gpu::q8_matmat(
2661                                    model, *idx, row_scale, &flat, b, rows, cols, out,
2662                                ) {
2663                                    crate::gpu::probe_record(
2664                                        crate::gpu::OpClass::Matmat,
2665                                        true,
2666                                        t0.elapsed(),
2667                                    );
2668                                    return;
2669                                }
2670                            }
2671                            crate::gpu::ProbeArm::CpuTimed => {
2672                                let q = self.quant_bytes();
2673                                qmatmat(q, row_scale, &pre, rows, cols, out, pool);
2674                                crate::gpu::probe_record(
2675                                    crate::gpu::OpClass::Matmat,
2676                                    false,
2677                                    t0.elapsed(),
2678                                );
2679                                return;
2680                            }
2681                            crate::gpu::ProbeArm::Cpu => {}
2682                        }
2683                    }
2684                }
2685                let q = self.quant_bytes();
2686                qmatmat(q, row_scale, &pre, rows, cols, out, pool);
2687            }
2688        }
2689    }
2690}
2691
2692impl QTensor {
2693    /// `mxfp4` matvec: the wgpu kernel for lm_head/expert-class shapes when
2694    /// the probe prefers the device, else the CPU kernels of
2695    /// `qtensor::mxfp4` (always on native Metal, which has no mxfp4 kernel).
2696    fn mxfp4_matvec(&self, x: &[f32], out: &mut [f32], pool: Option<&Pool>) {
2697        let (rows, cols) = (self.rows(), self.cols());
2698        if let Self::Mapped { model, idx, .. } = self {
2699            if rows * cols >= 8_388_608 && crate::gpu::enabled_here() {
2700                let t0 = std::time::Instant::now();
2701                let cls = crate::gpu::matvec_class(rows, cols);
2702                match crate::gpu::probe_arm(cls) {
2703                    crate::gpu::ProbeArm::Gpu => {
2704                        if crate::gpu::mxfp4_matvec(model, *idx, x, rows, cols, out) {
2705                            crate::gpu::probe_record(cls, true, t0.elapsed());
2706                            return;
2707                        }
2708                    }
2709                    crate::gpu::ProbeArm::CpuTimed => {
2710                        mxfp4::matvec(self.quant_bytes(), x, rows, cols, out, pool);
2711                        crate::gpu::probe_record(cls, false, t0.elapsed());
2712                        return;
2713                    }
2714                    crate::gpu::ProbeArm::Cpu => {}
2715                }
2716            }
2717        }
2718        mxfp4::matvec(self.quant_bytes(), x, rows, cols, out, pool);
2719    }
2720
2721    /// `mxfp4` batched matvec (`xs` `[b, cols]` → `out` `[b, rows]`): the
2722    /// device's batched kernel behind the same probe as the other codecs'
2723    /// GEMMs (never for row-exact batches — the device sums in another
2724    /// order than the host matvec), else the CPU kernels.
2725    fn mxfp4_matmat(&self, xs: &[f32], b: usize, out: &mut [f32], pool: Option<&Pool>) {
2726        let (rows, cols) = (self.rows(), self.cols());
2727        if let Self::Mapped { model, idx, .. } = self {
2728            if b >= 2
2729                && b * rows * cols >= 33_554_432
2730                && !row_exact()
2731                && !crate::gpu::mm_killed()
2732                && crate::gpu::enabled_here()
2733            {
2734                let class = if b >= 128 {
2735                    crate::gpu::OpClass::MatmatWide
2736                } else {
2737                    crate::gpu::OpClass::Matmat
2738                };
2739                let t0 = std::time::Instant::now();
2740                match crate::gpu::probe_arm(class) {
2741                    crate::gpu::ProbeArm::Gpu => {
2742                        if crate::gpu::mxfp4_matmat(model, *idx, xs, b, rows, cols, out) {
2743                            crate::gpu::probe_record(class, true, t0.elapsed());
2744                            return;
2745                        }
2746                    }
2747                    crate::gpu::ProbeArm::CpuTimed => {
2748                        mxfp4::matmat(self.quant_bytes(), xs, b, rows, cols, out, pool);
2749                        crate::gpu::probe_record(class, false, t0.elapsed());
2750                        return;
2751                    }
2752                    crate::gpu::ProbeArm::Cpu => {}
2753                }
2754            }
2755        }
2756        mxfp4::matmat(self.quant_bytes(), xs, b, rows, cols, out, pool);
2757    }
2758}
2759
2760impl QTensor {
2761    /// The device GEMM this tensor would take, run once on the caller's
2762    /// data — the startup parity probe's arm, and the one place that knows
2763    /// which entry point each codec has.
2764    ///
2765    /// It exists because the probe used to look for a `q4tp` weight by
2766    /// name AND dtype, and a container packed any other way was declared
2767    /// "host path" for the whole render even though its codec had a device
2768    /// GEMM of its own. A gate that only recognizes one codec is a gate
2769    /// that silently downgrades every other one.
2770    pub fn device_matmat(&self, xs: &[f32], b: usize, out: &mut [f32]) -> bool {
2771        let (rows, cols) = (self.rows(), self.cols());
2772        let Self::Mapped {
2773            model,
2774            idx,
2775            dtype,
2776            row_scale,
2777            col_field,
2778            ..
2779        } = self
2780        else {
2781            return false;
2782        };
2783        if crate::prism::has_contract(model) {
2784            return false;
2785        }
2786        match *dtype {
2787            TensorDtype::Q4TiledP => crate::gpu::q4tp_matmat(model, *idx, xs, b, rows, cols, out),
2788            // The two-field codec folds its column field into the
2789            // activation, which leaves a plain per-row int8 GEMM — the
2790            // same kernel `q8_row` uses, on both backends.
2791            TensorDtype::Q8Row | TensorDtype::Q8_2f => {
2792                // The field belongs to the weight; only a backend that cannot
2793                // apply it there makes a scaled copy of the activation.
2794                if *dtype == TensorDtype::Q8_2f
2795                    && std::env::var("CMF_Q8_2F_DEV").as_deref() != Ok("0")
2796                    && crate::gpu::q8_matmat_2f(
2797                        model, *idx, row_scale, col_field, xs, b, rows, cols, out,
2798                    )
2799                {
2800                    return true;
2801                }
2802                let flat: Vec<f32> = (0..b)
2803                    .flat_map(|bi| {
2804                        prescale(&xs[bi * cols..(bi + 1) * cols], col_field, *dtype).into_owned()
2805                    })
2806                    .collect();
2807                crate::gpu::q8_matmat(model, *idx, row_scale, &flat, b, rows, cols, out)
2808            }
2809            _ => false,
2810        }
2811    }
2812
2813    /// Multi-matrix job (roadmap §3 P0): N tensors sharing one input
2814    /// run under a SINGLE pool dispatch — QKV or gate+up cost one
2815    /// barrier instead of N. Per-row math is the exact same kernel as
2816    /// `matvec` (bit-identical outputs); only the dispatch is fused.
2817    /// Falls back to N sequential matvecs when the set is not a uniform
2818    /// q8-family/F32 group or there is no pool.
2819    pub fn matvec_many<const N: usize>(
2820        ts: [&QTensor; N],
2821        x: &[f32],
2822        mut outs: [&mut [f32]; N],
2823        pool: Option<&Pool>,
2824    ) {
2825        let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
2826        if ts.iter().any(|t| t.has_prism_contract() || t.per_tensor_only()) {
2827            // The fused range kernels have no transform descriptor.  Let
2828            // each tensor's ordinary matvec dispatch perform the explicit
2829            // signed FWHT (and retain CPU fallback for mixed q2tp/q4tp).
2830            // blr has no row-range kernel either: per-tensor calls.
2831            for (t, o) in ts.iter().zip(outs.iter_mut()) {
2832                t.matvec(x, o, pool);
2833            }
2834            return;
2835        }
2836        let uniform_q8 = ts.iter().all(|t| {
2837            matches!(
2838                t,
2839                Self::Mapped {
2840                    dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f,
2841                    ..
2842                }
2843            )
2844        });
2845        let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
2846        if uniform_f32 && crate::f32_backend::active() {
2847            for (t, o) in ts.iter().zip(outs.iter_mut()) {
2848                t.matvec(x, o, pool);
2849            }
2850            return;
2851        }
2852        let uniform_q4 = ts.iter().all(|t| {
2853            matches!(
2854                t,
2855                Self::Mapped {
2856                    dtype: TensorDtype::Q4Block,
2857                    ..
2858                }
2859            )
2860        });
2861        let uniform_vbit = ts.iter().all(|t| {
2862            matches!(
2863                t,
2864                Self::Mapped {
2865                    dtype: TensorDtype::Vbit | TensorDtype::VbitRo,
2866                    ..
2867                }
2868            )
2869        });
2870        // q1 and q1g may mix freely: the shared activation split and group
2871        // sums are layout-independent, only the row kernel differs.
2872        let uniform_q1 =
2873            ts.iter().all(|t| t.bin_w().is_some()) && ts.iter().all(|t| t.cols() == ts[0].cols());
2874        let uniform_q1t = ts.iter().all(|t| {
2875            matches!(
2876                t,
2877                Self::Mapped {
2878                    dtype: TensorDtype::Q1T,
2879                    ..
2880                }
2881            )
2882        });
2883        // q4tp is the skeleton dtype of the big MoE files, and without an arm
2884        // here every projection that shares an input paid its own pool
2885        // barrier: DeepSeek-V4's attention step alone hands this function
2886        // wq_a, wkv and both compressors' pairs off the same hidden state.
2887        let uniform_q4tp = ts.iter().all(|t| {
2888            matches!(
2889                t,
2890                Self::Mapped {
2891                    dtype: TensorDtype::Q4TiledP,
2892                    ..
2893                }
2894            )
2895        }) && ts
2896            .iter()
2897            .all(|t| t.cols() == ts[0].cols() && t.cols() % GROUP_SIZE == 0);
2898        let Some(pool) = pool else {
2899            for (t, o) in ts.iter().zip(outs.iter_mut()) {
2900                t.matvec(x, o, None);
2901            }
2902            return;
2903        };
2904        if total_rows < 256
2905            || !(uniform_q8
2906                || uniform_f32
2907                || uniform_q4
2908                || uniform_vbit
2909                || uniform_q1
2910                || uniform_q1t
2911                || uniform_q4tp)
2912        {
2913            for (t, o) in ts.iter().zip(outs.iter_mut()) {
2914                t.matvec(x, o, Some(pool));
2915            }
2916            return;
2917        }
2918
2919        if uniform_q4tp {
2920            // Every tensor's rows laid end to end in one virtual row space,
2921            // so the whole set is ONE dispatch. The per-row body is the
2922            // `q4tp_matvec` arm verbatim — same activation split, same
2923            // accumulation order — so the outputs are bit-identical to the
2924            // sequential calls this replaces.
2925            let cols = ts[0].cols();
2926            let gpr = cols / GROUP_SIZE;
2927            let views: [Q4tpView; N] =
2928                std::array::from_fn(|i| Q4tpView::new(ts[i].quant_bytes(), ts[i].rows(), cols));
2929            let rows_of: [usize; N] = std::array::from_fn(|i| ts[i].rows());
2930            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2931            // flat index -> (which tensor, which of its rows)
2932            let locate = |flat: usize| -> (usize, usize) {
2933                let mut acc = 0;
2934                for (i, &r) in rows_of.iter().enumerate() {
2935                    if flat < acc + r {
2936                        return (i, flat - acc);
2937                    }
2938                    acc += r;
2939                }
2940                (rows_of.len() - 1, 0)
2941            };
2942            let (views, outs_addr) = (&views, &outs_addr);
2943            if a8w8_enabled() {
2944                let act = split_act(x);
2945                let act = &act;
2946                let run = |start: usize, end: usize| {
2947                    with_krow(gpr, |sc| {
2948                        for flat in start..end {
2949                            let (t, r) = locate(flat);
2950                            let v = &views[t];
2951                            v.scales_into(r, gpr, sc);
2952                            let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, sc) * act.sx;
2953                            for &(j, xv) in &act.outliers {
2954                                let (w, s) = q4tp_outlier(v.nib, r, gpr, j, sc);
2955                                acc += w * s * xv;
2956                            }
2957                            // SAFETY: one worker owns each (tensor, row) pair.
2958                            unsafe { *outs_addr[t].at(r) = acc };
2959                        }
2960                    });
2961                };
2962                pool.run_rows(total_rows, &run);
2963            } else {
2964                let run = |start: usize, end: usize| {
2965                    with_krow(gpr, |sc| {
2966                        for flat in start..end {
2967                            let (t, r) = locate(flat);
2968                            let v = &views[t];
2969                            v.scales_into(r, gpr, sc);
2970                            // SAFETY: one worker owns each (tensor, row) pair.
2971                            unsafe { *outs_addr[t].at(r) = q4tp_row_exact(v.nib, r, gpr, x, sc) };
2972                        }
2973                    });
2974                };
2975                pool.run_rows(total_rows, &run);
2976            }
2977            return;
2978        }
2979
2980        if uniform_q1 {
2981            // One shared activation split + group sums (q1 has no col
2982            // field; the same input feeds every tensor).
2983            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2984            if a8w8_enabled() {
2985                let act = split_act(x);
2986                let gsum = q1_group_sums(&act.xq, ts[0].cols() / GROUP_SIZE);
2987                let (act, gsum) = (&act, &gsum);
2988                let closures: [_; N] = std::array::from_fn(|i| {
2989                    let (w, gpr, out) = (
2990                        ts[i].bin_w().expect("uniform_q1 checked"),
2991                        ts[i].cols() / GROUP_SIZE,
2992                        outs_addr[i],
2993                    );
2994                    move |s: usize, e: usize| bin_range_a8w8(w, gpr, act, gsum, out, s, e)
2995                });
2996                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2997                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2998                pool.run_many(&parts);
2999            } else {
3000                let closures: [_; N] = std::array::from_fn(|i| {
3001                    let (w, gpr, out) = (
3002                        ts[i].bin_w().expect("uniform_q1 checked"),
3003                        ts[i].cols() / GROUP_SIZE,
3004                        outs_addr[i],
3005                    );
3006                    move |s: usize, e: usize| bin_range_f32(w, gpr, x, out, s, e)
3007                });
3008                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3009                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3010                pool.run_many(&parts);
3011            }
3012            return;
3013        }
3014
3015        if uniform_q1t {
3016            // Q1T batched: one shared activation split + overlay decode,
3017            // all tensors' rows in ONE pool dispatch (saves N−1 dispatches
3018            // and N−1 redundant split_act calls per layer).
3019            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
3020            const TILE: usize = cortiq_core::quant::Q1T_TILE;
3021            if a8w8_enabled() {
3022                let act = split_act(x);
3023                let act = &act;
3024                let x_ref = x;
3025                let closures: [_; N] = std::array::from_fn(|i| {
3026                    let bytes = ts[i].quant_bytes();
3027                    let (rows, cols) = (ts[i].rows(), ts[i].cols());
3028                    let gpr = cols / GROUP_SIZE;
3029                    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
3030                    let out = outs_addr[i];
3031                    move |s: usize, e: usize| {
3032                        q1t_range_a8w8(bytes, gpr, rp_off, ent_off, has_ov, act, x_ref, out, s, e)
3033                    }
3034                });
3035                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3036                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3037                pool.run_many(&parts);
3038            } else {
3039                let x_ref = x;
3040                let closures: [_; N] = std::array::from_fn(|i| {
3041                    let bytes = ts[i].quant_bytes();
3042                    let (rows, cols) = (ts[i].rows(), ts[i].cols());
3043                    let gpr = cols / GROUP_SIZE;
3044                    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
3045                    let out = outs_addr[i];
3046                    move |s: usize, e: usize| {
3047                        q1t_range_f32_batch(bytes, gpr, rp_off, ent_off, has_ov, x_ref, out, s, e)
3048                    }
3049                });
3050                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3051                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3052                pool.run_many(&parts);
3053            }
3054            return;
3055        }
3056
3057        if uniform_q4 || uniform_vbit {
3058            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
3059            // q4/vbit share one activation split — no per-tensor col field.
3060            if a8w8_enabled() {
3061                let act = split_act(x);
3062                let act = &act;
3063                if uniform_q4 {
3064                    let closures: [_; N] = std::array::from_fn(|i| {
3065                        let (packed, scales) =
3066                            q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
3067                        let (gpr, cols, out) =
3068                            (ts[i].cols() / GROUP_SIZE, ts[i].cols(), outs_addr[i]);
3069                        move |s: usize, e: usize| {
3070                            q4_range_a8w8(packed, scales, gpr, cols, act, out, s, e)
3071                        }
3072                    });
3073                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3074                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3075                    pool.run_many(&parts);
3076                } else {
3077                    let closures: [_; N] = std::array::from_fn(|i| {
3078                        let Self::Mapped { vbit_offsets, .. } = ts[i] else {
3079                            unreachable!()
3080                        };
3081                        let (bytes, rows, cols, out) = (
3082                            ts[i].quant_bytes(),
3083                            ts[i].rows(),
3084                            ts[i].cols(),
3085                            outs_addr[i],
3086                        );
3087                        move |s: usize, e: usize| {
3088                            vbit_range_a8w8(bytes, vbit_offsets, x, act, rows, cols, out, s, e)
3089                        }
3090                    });
3091                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3092                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3093                    pool.run_many(&parts);
3094                }
3095                return;
3096            }
3097            if uniform_q4 {
3098                let closures: [_; N] = std::array::from_fn(|i| {
3099                    let (packed, scales) =
3100                        q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
3101                    let (gpr, out) = (ts[i].cols() / GROUP_SIZE, outs_addr[i]);
3102                    move |s: usize, e: usize| q4_range_f32(packed, scales, gpr, x, out, s, e)
3103                });
3104                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3105                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3106                pool.run_many(&parts);
3107            } else {
3108                let closures: [_; N] = std::array::from_fn(|i| {
3109                    let Self::Mapped { vbit_offsets, .. } = ts[i] else {
3110                        unreachable!()
3111                    };
3112                    let (bytes, rows, cols, out) = (
3113                        ts[i].quant_bytes(),
3114                        ts[i].rows(),
3115                        ts[i].cols(),
3116                        outs_addr[i],
3117                    );
3118                    move |s: usize, e: usize| {
3119                        vbit_range_f32(bytes, vbit_offsets, x, rows, cols, out, s, e)
3120                    }
3121                });
3122                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3123                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3124                pool.run_many(&parts);
3125            }
3126            return;
3127        }
3128
3129        if uniform_f32 {
3130            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
3131            let closures: [_; N] = std::array::from_fn(|i| {
3132                let Self::F32 { data, cols, .. } = ts[i] else {
3133                    unreachable!()
3134                };
3135                let out = outs_addr[i];
3136                move |start: usize, end: usize| {
3137                    for o in start..end {
3138                        let row = &data[o * cols..(o + 1) * cols];
3139                        let mut sum = 0.0f32;
3140                        for j in 0..*cols {
3141                            sum += row[j] * x[j];
3142                        }
3143                        // SAFETY: disjoint (tensor, row) cells per worker.
3144                        unsafe { *out.at(o) = sum };
3145                    }
3146                }
3147            });
3148            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3149                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3150            pool.run_many(&parts);
3151            return;
3152        }
3153
3154        // Uniform q8-family: per-tensor prescale (q8_2f col fields
3155        // differ per tensor) + the shared range kernels.
3156        struct Ctx<'a> {
3157            bytes: &'a [u8],
3158            #[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
3159            rep: &'a [u8],
3160            row_scale: &'a [f32],
3161            cols: usize,
3162            xs: std::borrow::Cow<'a, [f32]>,
3163        }
3164        let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
3165            let Self::Mapped {
3166                dtype,
3167                cols,
3168                row_scale,
3169                col_field,
3170                repack,
3171                ..
3172            } = ts[i]
3173            else {
3174                unreachable!()
3175            };
3176            Ctx {
3177                bytes: ts[i].quant_bytes(),
3178                rep: repack,
3179                row_scale,
3180                cols: *cols,
3181                xs: prescale(x, col_field, *dtype),
3182            }
3183        });
3184        let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
3185        #[cfg(target_arch = "aarch64")]
3186        if sdot_enabled() {
3187            let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
3188            let closures: [_; N] = std::array::from_fn(|i| {
3189                let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
3190                move |start: usize, end: usize| {
3191                    q8_range_sdot(c.bytes, c.rep, c.row_scale, act, c.cols, out, start, end)
3192                }
3193            });
3194            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3195                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3196            pool.run_many(&parts);
3197            return;
3198        }
3199        #[cfg(target_arch = "x86_64")]
3200        if avx2_a8w8_enabled() {
3201            let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
3202            let closures: [_; N] = std::array::from_fn(|i| {
3203                let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
3204                move |start: usize, end: usize| {
3205                    q8_range_avx2(c.bytes, c.row_scale, act, c.cols, out, start, end)
3206                }
3207            });
3208            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3209                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3210            pool.run_many(&parts);
3211            return;
3212        }
3213        let closures: [_; N] = std::array::from_fn(|i| {
3214            let (c, out) = (&ctxs[i], outs_addr[i]);
3215            move |start: usize, end: usize| {
3216                q8_range_f32(c.bytes, c.row_scale, &c.xs, c.cols, out, start, end)
3217            }
3218        });
3219        let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3220            std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3221        pool.run_many(&parts);
3222    }
3223}
3224
3225impl QTensor {
3226    /// Pair-input multi-matrix job: N tensors × 2 shared inputs under a
3227    /// single pool dispatch — the MTP/pair decode path publishes one job
3228    /// for Q/K/V (and one for gate+up) instead of one per tensor.
3229    /// Per-row math is exactly `matvec2`'s kernels; bit-identical.
3230    #[allow(clippy::needless_range_loop)]
3231    pub fn matvec2_many<const N: usize>(
3232        ts: [&QTensor; N],
3233        x1: &[f32],
3234        x2: &[f32],
3235        mut o1s: [&mut [f32]; N],
3236        mut o2s: [&mut [f32]; N],
3237        pool: Option<&Pool>,
3238    ) {
3239        let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
3240        if ts.iter().any(|t| t.has_prism_contract() || t.per_tensor_only()) {
3241            for i in 0..N {
3242                ts[i].matvec2(x1, x2, o1s[i], o2s[i], pool);
3243            }
3244            return;
3245        }
3246        let uniform_q8 = ts.iter().all(|t| {
3247            matches!(
3248                t,
3249                Self::Mapped {
3250                    dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f,
3251                    ..
3252                }
3253            )
3254        });
3255        let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
3256        let uniform_q4 = ts.iter().all(|t| {
3257            matches!(
3258                t,
3259                Self::Mapped {
3260                    dtype: TensorDtype::Q4Block,
3261                    ..
3262                }
3263            )
3264        });
3265        let uniform_vbit = ts.iter().all(|t| {
3266            matches!(
3267                t,
3268                Self::Mapped {
3269                    dtype: TensorDtype::Vbit | TensorDtype::VbitRo,
3270                    ..
3271                }
3272            )
3273        });
3274        let fusable = pool.is_some()
3275            && total_rows >= 256
3276            && (uniform_q8 || uniform_f32 || uniform_q4 || uniform_vbit);
3277        if !fusable {
3278            for i in 0..N {
3279                ts[i].matvec2(x1, x2, o1s[i], o2s[i], pool);
3280            }
3281            return;
3282        }
3283        let pool = pool.unwrap();
3284
3285        if uniform_q4 || uniform_vbit {
3286            let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
3287            let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
3288            // q4/vbit share activation splits — no per-tensor col field.
3289            if a8w8_enabled() {
3290                let a1 = split_act(x1);
3291                let a2 = split_act(x2);
3292                let (a1, a2) = (&a1, &a2);
3293                if uniform_q4 {
3294                    let closures: [_; N] = std::array::from_fn(|i| {
3295                        let (packed, scales) =
3296                            q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
3297                        let (gpr, cols, o1, o2) =
3298                            (ts[i].cols() / GROUP_SIZE, ts[i].cols(), p1[i], p2[i]);
3299                        move |s: usize, e: usize| {
3300                            q4_range2_a8w8(packed, scales, gpr, cols, a1, a2, o1, o2, s, e)
3301                        }
3302                    });
3303                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3304                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3305                    pool.run_many(&parts);
3306                } else {
3307                    let closures: [_; N] = std::array::from_fn(|i| {
3308                        let Self::Mapped { vbit_offsets, .. } = ts[i] else {
3309                            unreachable!()
3310                        };
3311                        let (bytes, rows, cols, o1, o2) = (
3312                            ts[i].quant_bytes(),
3313                            ts[i].rows(),
3314                            ts[i].cols(),
3315                            p1[i],
3316                            p2[i],
3317                        );
3318                        move |s: usize, e: usize| {
3319                            vbit_range2_a8w8(
3320                                bytes,
3321                                vbit_offsets,
3322                                x1,
3323                                x2,
3324                                a1,
3325                                a2,
3326                                rows,
3327                                cols,
3328                                o1,
3329                                o2,
3330                                s,
3331                                e,
3332                            )
3333                        }
3334                    });
3335                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3336                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3337                    pool.run_many(&parts);
3338                }
3339                return;
3340            }
3341            if uniform_q4 {
3342                let closures: [_; N] = std::array::from_fn(|i| {
3343                    let (packed, scales) =
3344                        q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
3345                    let (gpr, o1, o2) = (ts[i].cols() / GROUP_SIZE, p1[i], p2[i]);
3346                    move |s: usize, e: usize| {
3347                        q4_range2_f32(packed, scales, gpr, x1, x2, o1, o2, s, e)
3348                    }
3349                });
3350                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3351                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3352                pool.run_many(&parts);
3353            } else {
3354                let closures: [_; N] = std::array::from_fn(|i| {
3355                    let Self::Mapped { vbit_offsets, .. } = ts[i] else {
3356                        unreachable!()
3357                    };
3358                    let (bytes, rows, cols, o1, o2) = (
3359                        ts[i].quant_bytes(),
3360                        ts[i].rows(),
3361                        ts[i].cols(),
3362                        p1[i],
3363                        p2[i],
3364                    );
3365                    move |s: usize, e: usize| {
3366                        vbit_range2_f32(bytes, vbit_offsets, x1, x2, rows, cols, o1, o2, s, e)
3367                    }
3368                });
3369                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3370                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3371                pool.run_many(&parts);
3372            }
3373            return;
3374        }
3375
3376        if uniform_f32 {
3377            let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
3378            let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
3379            let closures: [_; N] = std::array::from_fn(|i| {
3380                let Self::F32 { data, cols, .. } = ts[i] else {
3381                    unreachable!()
3382                };
3383                let (o1, o2) = (p1[i], p2[i]);
3384                move |start: usize, end: usize| {
3385                    for o in start..end {
3386                        let row = &data[o * cols..(o + 1) * cols];
3387                        let (mut s1, mut s2) = (0.0f32, 0.0f32);
3388                        for j in 0..*cols {
3389                            s1 += row[j] * x1[j];
3390                            s2 += row[j] * x2[j];
3391                        }
3392                        // SAFETY: disjoint (tensor, row) cells per worker.
3393                        unsafe {
3394                            *o1.at(o) = s1;
3395                            *o2.at(o) = s2;
3396                        }
3397                    }
3398                }
3399            });
3400            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3401                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3402            pool.run_many(&parts);
3403            return;
3404        }
3405
3406        struct Ctx<'a> {
3407            bytes: &'a [u8],
3408            row_scale: &'a [f32],
3409            cols: usize,
3410            xs1: std::borrow::Cow<'a, [f32]>,
3411            xs2: std::borrow::Cow<'a, [f32]>,
3412        }
3413        let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
3414            let Self::Mapped {
3415                dtype,
3416                cols,
3417                row_scale,
3418                col_field,
3419                ..
3420            } = ts[i]
3421            else {
3422                unreachable!()
3423            };
3424            Ctx {
3425                bytes: ts[i].quant_bytes(),
3426                row_scale,
3427                cols: *cols,
3428                xs1: prescale(x1, col_field, *dtype),
3429                xs2: prescale(x2, col_field, *dtype),
3430            }
3431        });
3432        let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
3433        let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
3434        #[cfg(target_arch = "aarch64")]
3435        if sdot_enabled() {
3436            let acts: [(SplitAct, SplitAct); N] =
3437                std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
3438            let closures: [_; N] = std::array::from_fn(|i| {
3439                let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
3440                move |start: usize, end: usize| {
3441                    q8_range2_sdot(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
3442                }
3443            });
3444            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3445                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3446            pool.run_many(&parts);
3447            return;
3448        }
3449        #[cfg(target_arch = "x86_64")]
3450        if avx2_a8w8_enabled() {
3451            let acts: [(SplitAct, SplitAct); N] =
3452                std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
3453            let closures: [_; N] = std::array::from_fn(|i| {
3454                let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
3455                move |start: usize, end: usize| {
3456                    q8_range2_avx2(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
3457                }
3458            });
3459            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3460                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3461            pool.run_many(&parts);
3462            return;
3463        }
3464        let closures: [_; N] = std::array::from_fn(|i| {
3465            let (c, o1, o2) = (&ctxs[i], p1[i], p2[i]);
3466            move |start: usize, end: usize| {
3467                q8_range2_f32(
3468                    c.bytes,
3469                    c.row_scale,
3470                    &c.xs1,
3471                    &c.xs2,
3472                    c.cols,
3473                    o1,
3474                    o2,
3475                    start,
3476                    end,
3477                )
3478            }
3479        });
3480        let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3481            std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3482        pool.run_many(&parts);
3483    }
3484
3485    /// Fused gate+up matvec with SiLU·mul: for each row r, computes
3486    /// `silu(gate·x) * (up·x)` and writes to `out[r]`. ONE pool dispatch,
3487    /// no intermediate g/u buffers, no separate silu pass. Falls back
3488    /// (returns false) for unsupported dtype combos.
3489    pub fn matvec_silu_mul(
3490        gate: &QTensor,
3491        up: &QTensor,
3492        x: &[f32],
3493        out: &mut [f32],
3494        pool: Option<&Pool>,
3495    ) -> bool {
3496        Self::matvec_silu_mul_limited(gate, up, x, out, 0.0, pool)
3497    }
3498
3499    /// Fused gate+up+SiLU with the GLM asymmetrical clamp.  `limit == 0`
3500    /// preserves the historical unclamped helper; a positive limit clamps
3501    /// `up` to both sides and `gate` only from above, matching the GLM
3502    /// SwiGLU reference.  Keeping the limit in the row kernel avoids the two
3503    /// intermediate vectors and the extra combine pass on the Q2TP experts.
3504    pub fn matvec_silu_mul_limited(
3505        gate: &QTensor,
3506        up: &QTensor,
3507        x: &[f32],
3508        out: &mut [f32],
3509        limit: f32,
3510        pool: Option<&Pool>,
3511    ) -> bool {
3512        if gate.has_prism_contract()
3513            || up.has_prism_contract()
3514            || gate.per_tensor_only()
3515            || up.per_tensor_only()
3516        {
3517            // The fused gate/up kernels consume x directly.  Prism requires
3518            // a per-matrix signed FWHT, so the caller must use two ordinary
3519            // descriptor-aware matvecs instead of an unrotated fast path.
3520            return false;
3521        }
3522        let inter = gate.rows();
3523        debug_assert_eq!(up.rows(), inter);
3524        debug_assert_eq!(out.len(), inter);
3525        debug_assert_eq!(gate.cols(), up.cols());
3526        if !a8w8_enabled() {
3527            return false;
3528        }
3529        let act = split_act(x);
3530        let act = &act;
3531        let x_ref = x;
3532        let out_addr = SendMut(out.as_mut_ptr());
3533
3534        match (gate, up) {
3535            // Q4Block gate + Q4Block up (most common mobile q4 models)
3536            (
3537                Self::Mapped {
3538                    dtype: TensorDtype::Q4Block,
3539                    ..
3540                },
3541                Self::Mapped {
3542                    dtype: TensorDtype::Q4Block,
3543                    ..
3544                },
3545            ) => {
3546                let (gp, gs) = q4_split(gate.quant_bytes(), gate.rows(), gate.cols());
3547                let (up_p, up_s) = q4_split(up.quant_bytes(), up.rows(), up.cols());
3548                let gpr = gate.cols() / GROUP_SIZE;
3549                let cols = gate.cols();
3550                let run = move |start: usize, end: usize| {
3551                    for r in start..end {
3552                        let mut gv = dot_q4_row_i8(gp, gs, r * gpr, gpr, &act.xq) * act.sx;
3553                        let mut uv = dot_q4_row_i8(up_p, up_s, r * gpr, gpr, &act.xq) * act.sx;
3554                        for &(j, xv) in &act.outliers {
3555                            let flat = r * cols + j;
3556                            let gb = gp[flat / 2];
3557                            let gn = if flat & 1 == 0 { gb & 0x0F } else { gb >> 4 };
3558                            let gsc = f16_to_f32(u16::from_le_bytes([
3559                                gs[(flat / GROUP_SIZE) * 2],
3560                                gs[(flat / GROUP_SIZE) * 2 + 1],
3561                            ]));
3562                            gv += ((gn as i32 - 8) as f32) * gsc * xv;
3563                            let ub = up_p[flat / 2];
3564                            let un = if flat & 1 == 0 { ub & 0x0F } else { ub >> 4 };
3565                            let usc = f16_to_f32(u16::from_le_bytes([
3566                                up_s[(flat / GROUP_SIZE) * 2],
3567                                up_s[(flat / GROUP_SIZE) * 2 + 1],
3568                            ]));
3569                            uv += ((un as i32 - 8) as f32) * usc * xv;
3570                        }
3571                        // SAFETY: disjoint row ranges per worker.
3572                        unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3573                    }
3574                };
3575                dispatch_rows(pool, inter, &run);
3576                true
3577            }
3578            // Q4Tiled gate + Q4Tiled up — one row pass, both tile
3579            // streams sequential, silu·mul fused (same per-row math as
3580            // `q4t_matvec`).
3581            (
3582                Self::Mapped {
3583                    dtype: TensorDtype::Q4Tiled,
3584                    ..
3585                },
3586                Self::Mapped {
3587                    dtype: TensorDtype::Q4Tiled,
3588                    ..
3589                },
3590            ) => {
3591                let g_bytes = gate.quant_bytes();
3592                let u_bytes = up.quant_bytes();
3593                let gpr = gate.cols() / GROUP_SIZE;
3594                let run = move |start: usize, end: usize| {
3595                    for r in start..end {
3596                        let mut gv = dot_q4t_row_i8(g_bytes, r, gpr, &act.xq) * act.sx;
3597                        let mut uv = dot_q4t_row_i8(u_bytes, r, gpr, &act.xq) * act.sx;
3598                        for &(j, xv) in &act.outliers {
3599                            let (w, s) = q4t_outlier(g_bytes, r, gpr, j);
3600                            gv += w * s * xv;
3601                            let (w, s) = q4t_outlier(u_bytes, r, gpr, j);
3602                            uv += w * s * xv;
3603                        }
3604                        // SAFETY: disjoint row ranges per worker.
3605                        unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3606                    }
3607                };
3608                dispatch_rows(pool, inter, &run);
3609                true
3610            }
3611            // Q4TiledP gate + Q4TiledP up — the same fused row pass, with
3612            // each row's two ladders built once and spent on both streams.
3613            (
3614                Self::Mapped {
3615                    dtype: TensorDtype::Q4TiledP,
3616                    ..
3617                },
3618                Self::Mapped {
3619                    dtype: TensorDtype::Q4TiledP,
3620                    ..
3621                },
3622            ) => {
3623                let cols = gate.cols();
3624                let gpr = cols / GROUP_SIZE;
3625                let gv_view = Q4tpView::new(gate.quant_bytes(), inter, cols);
3626                let uv_view = Q4tpView::new(up.quant_bytes(), inter, cols);
3627                let run = |start: usize, end: usize| {
3628                    with_krows(gpr, |gsc, usc| {
3629                        for r in start..end {
3630                            gv_view.scales_into(r, gpr, gsc);
3631                            uv_view.scales_into(r, gpr, usc);
3632                            let mut gv =
3633                                dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsc) * act.sx;
3634                            let mut uv =
3635                                dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, usc) * act.sx;
3636                            for &(j, xv) in &act.outliers {
3637                                let (w, s) = q4tp_outlier(gv_view.nib, r, gpr, j, gsc);
3638                                gv += w * s * xv;
3639                                let (w, s) = q4tp_outlier(uv_view.nib, r, gpr, j, usc);
3640                                uv += w * s * xv;
3641                            }
3642                            // SAFETY: disjoint row ranges per worker.
3643                            unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3644                        }
3645                    });
3646                };
3647                dispatch_rows(pool, inter, &run);
3648                true
3649            }
3650            // Q1 gate + Q1 up — one row pass over both sign streams,
3651            // silu·mul fused (the per-row math of `q1_range_a8w8`); the
3652            // activation group sums are shared by both streams. Without
3653            // this arm a q1 dense FFN paid two dispatches + a combine
3654            // loop — the exact barrier this function exists to remove.
3655            // q1g (and a q1/q1g mix) rides the same arm through `BinW`.
3656            (
3657                Self::Mapped {
3658                    dtype: TensorDtype::Q1 | TensorDtype::Q1G,
3659                    ..
3660                },
3661                Self::Mapped {
3662                    dtype: TensorDtype::Q1 | TensorDtype::Q1G,
3663                    ..
3664                },
3665            ) => {
3666                let gw = gate.bin_w().expect("binary dtype matched");
3667                let uw = up.bin_w().expect("binary dtype matched");
3668                let gpr = gate.cols() / GROUP_SIZE;
3669                let gsum = q1_group_sums(&act.xq, gpr);
3670                let gsum = &gsum;
3671                let run = move |start: usize, end: usize| {
3672                    for r in start..end {
3673                        let mut gv = gw.dot_i8(r, gpr, &act.xq, gsum) * act.sx;
3674                        let mut uv = uw.dot_i8(r, gpr, &act.xq, gsum) * act.sx;
3675                        for &(j, xv) in &act.outliers {
3676                            let (w, s) = gw.outlier(r, gpr, j);
3677                            gv += w * s * xv;
3678                            let (w, s) = uw.outlier(r, gpr, j);
3679                            uv += w * s * xv;
3680                        }
3681                        // SAFETY: disjoint row ranges per worker.
3682                        unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3683                    }
3684                };
3685                dispatch_rows(pool, inter, &run);
3686                true
3687            }
3688            // Q2TiledP gate + Q2TiledP up — the 2-bit expert pair (MoE
3689            // FFNs of the W2 class): one row pass, both ladders built
3690            // once, integer code dots with shared group sums.
3691            (
3692                Self::Mapped {
3693                    dtype: TensorDtype::Q2TiledP,
3694                    ..
3695                },
3696                Self::Mapped {
3697                    dtype: TensorDtype::Q2TiledP,
3698                    ..
3699                },
3700            ) => {
3701                let cols = gate.cols();
3702                let gpr = cols / GROUP_SIZE;
3703                let gv_view = Q4tpView::new_q2(gate.quant_bytes(), inter, cols);
3704                let uv_view = Q4tpView::new_q2(up.quant_bytes(), inter, cols);
3705                let gsum = q1_group_sums(&act.xq, gpr);
3706                let gsum = &gsum;
3707                let run = move |start: usize, end: usize| {
3708                    with_krows(gpr, |gsc, usc| {
3709                        for r in start..end {
3710                            gv_view.scales_into(r, gpr, gsc);
3711                            uv_view.scales_into(r, gpr, usc);
3712                            let mut gv =
3713                                dot_q2tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsum, gsc) * act.sx;
3714                            let mut uv =
3715                                dot_q2tp_row_i8(uv_view.nib, r, gpr, &act.xq, gsum, usc) * act.sx;
3716                            for &(j, xv) in &act.outliers {
3717                                let (w, s) = q2tp_outlier(gv_view.nib, r, gpr, j, gsc);
3718                                gv += w * s * xv;
3719                                let (w, s) = q2tp_outlier(uv_view.nib, r, gpr, j, usc);
3720                                uv += w * s * xv;
3721                            }
3722                            // SAFETY: disjoint row ranges per worker.
3723                            unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3724                        }
3725                    });
3726                };
3727                dispatch_rows(pool, inter, &run);
3728                true
3729            }
3730            // Q8Row gate + Q8Row up — one row pass over both i8 streams.
3731            // Q8_2f stays out on purpose: its column field prescales the
3732            // activations PER TENSOR, which breaks this fn's shared
3733            // split_act contract — it keeps the two-dispatch path.
3734            (
3735                Self::Mapped {
3736                    dtype: TensorDtype::Q8Row,
3737                    row_scale: g_rs,
3738                    ..
3739                },
3740                Self::Mapped {
3741                    dtype: TensorDtype::Q8Row,
3742                    row_scale: u_rs,
3743                    ..
3744                },
3745            ) => {
3746                let g_bytes = gate.quant_bytes();
3747                let u_bytes = up.quant_bytes();
3748                let cols = gate.cols();
3749                let run = move |start: usize, end: usize| {
3750                    for r in start..end {
3751                        let gv = q8_row_dot(&g_bytes[r * cols..(r + 1) * cols], act) * g_rs[r];
3752                        let uv = q8_row_dot(&u_bytes[r * cols..(r + 1) * cols], act) * u_rs[r];
3753                        // SAFETY: disjoint row ranges per worker.
3754                        unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3755                    }
3756                };
3757                dispatch_rows(pool, inter, &run);
3758                true
3759            }
3760            // Q1T gate + Q1T up
3761            (
3762                Self::Mapped {
3763                    dtype: TensorDtype::Q1T,
3764                    ..
3765                },
3766                Self::Mapped {
3767                    dtype: TensorDtype::Q1T,
3768                    ..
3769                },
3770            ) => {
3771                const TILE: usize = cortiq_core::quant::Q1T_TILE;
3772                let g_bytes = gate.quant_bytes();
3773                let u_bytes = up.quant_bytes();
3774                let gpr = gate.cols() / GROUP_SIZE;
3775                let (g_rp, g_ent, g_ov) = q1t_overlay(g_bytes, inter * gpr * TILE, inter);
3776                let (u_rp, u_ent, u_ov) = q1t_overlay(u_bytes, inter * gpr * TILE, inter);
3777                let run = move |start: usize, end: usize| {
3778                    for r in start..end {
3779                        let mut gv = q1t_dot_row_i8(g_bytes, r, gpr, &act.xq) * act.sx;
3780                        let mut uv = q1t_dot_row_i8(u_bytes, r, gpr, &act.xq) * act.sx;
3781                        for &(j, xv) in &act.outliers {
3782                            gv += q1t_base_weight(g_bytes, r, gpr, j) * xv;
3783                            uv += q1t_base_weight(u_bytes, r, gpr, j) * xv;
3784                        }
3785                        gv += q1t_row_outlier_correction(g_bytes, r, g_rp, g_ent, g_ov, x_ref);
3786                        uv += q1t_row_outlier_correction(u_bytes, r, u_rp, u_ent, u_ov, x_ref);
3787                        // SAFETY: disjoint row ranges per worker.
3788                        unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3789                    }
3790                };
3791                dispatch_rows(pool, inter, &run);
3792                true
3793            }
3794            _ => false,
3795        }
3796    }
3797
3798    /// Every routed expert's fused gate/up/SiLU under ONE pool dispatch.
3799    ///
3800    /// The per-expert path pays a pool barrier per expert per stage: at 9
3801    /// experts over 40 layers that is ~720 barriers a token, and a decode
3802    /// profile of Qwen3.6-35B-A3B showed the pool parked in
3803    /// `psynch_cvwait` about twice as long as it spent computing. Laying
3804    /// every expert's rows end-to-end in one virtual row space collapses
3805    /// the stage to a single dispatch. The per-row body is the
3806    /// single-expert q4tp arm verbatim, so outputs are bit-identical.
3807    ///
3808    /// `false` = something is outside the fused q4tp kernel (dtype, shape,
3809    /// or a transformed tensor); the caller walks the ordinary per-expert
3810    /// path. Float activations use the same exact scalar rows, still fused
3811    /// under one pool dispatch.
3812    pub fn moe_gate_up_many(
3813        pairs: &[(&QTensor, &QTensor)],
3814        x: &[f32],
3815        outs: &mut [Vec<f32>],
3816        pool: Option<&Pool>,
3817    ) -> bool {
3818        Self::moe_gate_up_many_limited(pairs, x, outs, 0.0, pool)
3819    }
3820
3821    /// Batched gate/up/SiLU with the optional GLM clamp.  The public legacy
3822    /// helper above keeps its historical unclamped semantics; callers that
3823    /// implement a reference with a positive SwiGLU limit use this variant.
3824    pub fn moe_gate_up_many_limited(
3825        pairs: &[(&QTensor, &QTensor)],
3826        x: &[f32],
3827        outs: &mut [Vec<f32>],
3828        limit: f32,
3829        pool: Option<&Pool>,
3830    ) -> bool {
3831        if pairs.is_empty() || pairs.len() != outs.len() {
3832            return false;
3833        }
3834        if pairs.iter().any(|(g, u)| g.per_tensor_only() || u.per_tensor_only()) {
3835            return false;
3836        }
3837        if !a8w8_enabled() {
3838            if limit > 0.0 {
3839                // The exact-row fallback applies no SwiGLU clamp; the
3840                // caller's per-expert path carries it instead.
3841                return false;
3842            }
3843            let groups = vec![vec![0]; pairs.len()];
3844            return Self::moe_gate_up_rows(pairs, &groups, x, outs, pool);
3845        }
3846        let inter = pairs[0].0.rows();
3847        let cols = pairs[0].0.cols();
3848        if cols % GROUP_SIZE != 0 {
3849            return false;
3850        }
3851        let gpr = cols / GROUP_SIZE;
3852        // Uniform layout across every routed pair: q4tp, or the 2-bit
3853        // profile's q2tp gate/up (the W2 class). Mixed sets refuse.
3854        let q2 = matches!(
3855            pairs[0].0,
3856            Self::Mapped {
3857                dtype: TensorDtype::Q2TiledP,
3858                ..
3859            }
3860        );
3861        let want = if q2 {
3862            TensorDtype::Q2TiledP
3863        } else {
3864            TensorDtype::Q4TiledP
3865        };
3866        let mut views = Vec::with_capacity(pairs.len() * 2);
3867        for ((g, u), o) in pairs.iter().zip(outs.iter()) {
3868            let both = matches!(g, Self::Mapped { dtype, .. } if *dtype == want)
3869                && matches!(u, Self::Mapped { dtype, .. } if *dtype == want);
3870            if !both
3871                || g.rows() != inter
3872                || u.rows() != inter
3873                || g.cols() != cols
3874                || u.cols() != cols
3875                || o.len() != inter
3876            {
3877                return false;
3878            }
3879            let mk = if q2 { Q4tpView::new_q2 } else { Q4tpView::new };
3880            views.push(mk(g.quant_bytes(), inter, cols));
3881            views.push(mk(u.quant_bytes(), inter, cols));
3882        }
3883        let act = split_act(x);
3884        let gsum = if q2 {
3885            q1_group_sums(&act.xq, gpr)
3886        } else {
3887            Vec::new()
3888        };
3889        let (act, gsum) = (&act, &gsum);
3890        let ptrs: Vec<SendMut> = outs.iter_mut().map(|o| SendMut(o.as_mut_ptr())).collect();
3891        let (views, ptrs) = (&views, &ptrs);
3892        let run = |start: usize, end: usize| {
3893            with_krows(gpr, |gsc, usc| {
3894                for flat in start..end {
3895                    let (e, r) = (flat / inter, flat % inter);
3896                    let gv_view = &views[e * 2];
3897                    let uv_view = &views[e * 2 + 1];
3898                    gv_view.scales_into(r, gpr, gsc);
3899                    uv_view.scales_into(r, gpr, usc);
3900                    let (mut gv, mut uv) = if q2 {
3901                        (
3902                            dot_q2tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsum, gsc) * act.sx,
3903                            dot_q2tp_row_i8(uv_view.nib, r, gpr, &act.xq, gsum, usc) * act.sx,
3904                        )
3905                    } else {
3906                        (
3907                            dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsc) * act.sx,
3908                            dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, usc) * act.sx,
3909                        )
3910                    };
3911                    for &(j, xv) in &act.outliers {
3912                        let (og, ou) = if q2 {
3913                            (
3914                                q2tp_outlier(gv_view.nib, r, gpr, j, gsc),
3915                                q2tp_outlier(uv_view.nib, r, gpr, j, usc),
3916                            )
3917                        } else {
3918                            (
3919                                q4tp_outlier(gv_view.nib, r, gpr, j, gsc),
3920                                q4tp_outlier(uv_view.nib, r, gpr, j, usc),
3921                            )
3922                        };
3923                        gv += og.0 * og.1 * xv;
3924                        uv += ou.0 * ou.1 * xv;
3925                    }
3926                    // SAFETY: one worker owns each (expert, row) pair.
3927                    unsafe { *ptrs[e].at(r) = silu_mul_limited(gv, uv, limit) };
3928                }
3929            });
3930        };
3931        dispatch_rows(pool, pairs.len() * inter, &run);
3932        true
3933    }
3934
3935    /// Every routed expert's down projection, weighted and summed into
3936    /// `out`, under ONE pool dispatch.
3937    ///
3938    /// Partitioned by OUTPUT row rather than by expert: each row is owned
3939    /// by a single worker, so the experts are summed in the caller's order
3940    /// — the same sequence of f32 adds the serial `out[i] += w·eo[i]` loop
3941    /// performs, hence bit-identical. Partitioning by expert instead would
3942    /// race on the shared accumulator.
3943    pub fn moe_down_many(
3944        downs: &[&QTensor],
3945        gs: &[Vec<f32>],
3946        weights: &[f32],
3947        out: &mut [f32],
3948        pool: Option<&Pool>,
3949    ) -> bool {
3950        if downs.is_empty() || downs.len() != gs.len() || downs.len() != weights.len() {
3951            return false;
3952        }
3953        if downs.iter().any(|d| d.per_tensor_only()) {
3954            return false;
3955        }
3956        if !a8w8_enabled() {
3957            let mut terms = vec![vec![0.0; out.len()]; downs.len()];
3958            if !Self::moe_down_rows(downs, &vec![1; downs.len()], gs, &mut terms, pool) {
3959                return false;
3960            }
3961            out.fill(0.0);
3962            for (row, &w) in terms.iter().zip(weights) {
3963                for (o, &v) in out.iter_mut().zip(row) {
3964                    *o += w * v;
3965                }
3966            }
3967            return true;
3968        }
3969        let rows = out.len();
3970        let cols = downs[0].cols();
3971        if cols % GROUP_SIZE != 0 {
3972            return false;
3973        }
3974        let gpr = cols / GROUP_SIZE;
3975        let mut views = Vec::with_capacity(downs.len());
3976        for (d, g) in downs.iter().zip(gs.iter()) {
3977            if !matches!(
3978                d,
3979                Self::Mapped {
3980                    dtype: TensorDtype::Q4TiledP,
3981                    ..
3982                }
3983            ) || d.rows() != rows
3984                || d.cols() != cols
3985                || g.len() != cols
3986            {
3987                return false;
3988            }
3989            views.push(Q4tpView::new(d.quant_bytes(), rows, cols));
3990        }
3991        // One int8 split per expert — the activation vectors differ.
3992        let acts: Vec<SplitAct> = gs.iter().map(|g| split_act(g)).collect();
3993        // Partitioned by OUTPUT row, with the experts folded inside: each
3994        // row is owned by one worker, so they are summed in the caller's
3995        // order — the same f32 sequence the serial `out[i] += w·eo[i]`
3996        // loop produces. Partitioning by expert instead would either race
3997        // on the accumulator or need a scratch plane and a second pass;
3998        // measured, that variant was a wash, so this keeps the simpler
3999        // shape.
4000        let out_addr = SendMut(out.as_mut_ptr());
4001        let (views, acts, weights) = (&views, &acts, &weights);
4002        let run = |start: usize, end: usize| {
4003            with_krow(gpr, |sc| {
4004                for r in start..end {
4005                    let mut acc = 0f32;
4006                    for (e, v) in views.iter().enumerate() {
4007                        v.scales_into(r, gpr, sc);
4008                        let a = &acts[e];
4009                        let mut d = dot_q4tp_row_i8(v.nib, r, gpr, &a.xq, sc) * a.sx;
4010                        for &(j, xv) in &a.outliers {
4011                            let (w, s) = q4tp_outlier(v.nib, r, gpr, j, sc);
4012                            d += w * s * xv;
4013                        }
4014                        acc += weights[e] * d;
4015                    }
4016                    // SAFETY: disjoint row ranges per worker.
4017                    unsafe { *out_addr.at(r) = acc };
4018                }
4019            });
4020        };
4021        dispatch_rows(pool, rows, &run);
4022        true
4023    }
4024
4025    /// `moe_gate_up_many` for SEVERAL tokens at once, decode-exact: expert
4026    /// `e` (`pairs[e]`, q4tp) serves the tokens `groups[e]` (row indices
4027    /// into `xs`, each `cols` wide). Every (expert, token) output is
4028    /// bit-identical to `moe_gate_up_many` run on that token alone — the
4029    /// same int8 activation split, VNNI dots, outlier terms and inline
4030    /// SiLU — while each weight row is read once for all the tokens routed
4031    /// to its expert (the speculative verify's expert sharing). `outs` is
4032    /// flat in (expert, token-of-group) order. False = not covered (not
4033    /// q4tp): the caller takes the per-token path. With float activations,
4034    /// the exact scalar row kernel replaces the int8 dot without changing
4035    /// the shared dispatch or route-order reduction.
4036    pub fn moe_gate_up_rows(
4037        pairs: &[(&QTensor, &QTensor)],
4038        groups: &[Vec<usize>],
4039        xs: &[f32],
4040        outs: &mut [Vec<f32>],
4041        pool: Option<&Pool>,
4042    ) -> bool {
4043        if pairs.is_empty() || pairs.len() != groups.len() {
4044            return false;
4045        }
4046        if pairs.iter().any(|(g, u)| g.per_tensor_only() || u.per_tensor_only()) {
4047            return false;
4048        }
4049        let inter = pairs[0].0.rows();
4050        let cols = pairs[0].0.cols();
4051        let n_pairs: usize = groups.iter().map(|g| g.len()).sum();
4052        if cols == 0 || cols % GROUP_SIZE != 0 || outs.len() != n_pairs || xs.len() % cols != 0 {
4053            return false;
4054        }
4055        let b = xs.len() / cols;
4056        let gpr = cols / GROUP_SIZE;
4057        let mut views = Vec::with_capacity(pairs.len() * 2);
4058        for (g, u) in pairs {
4059            let q4tp = |t: &QTensor| {
4060                matches!(
4061                    t,
4062                    Self::Mapped {
4063                        dtype: TensorDtype::Q4TiledP,
4064                        ..
4065                    }
4066                )
4067            };
4068            if g.has_prism_contract()
4069                || u.has_prism_contract()
4070                || !q4tp(g)
4071                || !q4tp(u)
4072                || g.rows() != inter
4073                || u.rows() != inter
4074                || g.cols() != cols
4075                || u.cols() != cols
4076            {
4077                return false;
4078            }
4079            views.push(Q4tpView::new(g.quant_bytes(), inter, cols));
4080            views.push(Q4tpView::new(u.quant_bytes(), inter, cols));
4081        }
4082        if outs.iter().any(|o| o.len() != inter) || groups.iter().flatten().any(|&t| t >= b) {
4083            return false;
4084        }
4085        let quantized = a8w8_enabled();
4086        let acts: Vec<SplitAct> = if quantized {
4087            (0..b)
4088                .map(|t| split_act(&xs[t * cols..(t + 1) * cols]))
4089                .collect()
4090        } else {
4091            Vec::new()
4092        };
4093        let mut offs = Vec::with_capacity(groups.len());
4094        let mut o = 0usize;
4095        for g in groups {
4096            offs.push(o);
4097            o += g.len();
4098        }
4099        let ptrs: Vec<SendMut> = outs.iter_mut().map(|o| SendMut(o.as_mut_ptr())).collect();
4100        let (views, ptrs, acts, offs) = (&views, &ptrs, &acts, &offs);
4101        let run = |start: usize, end: usize| {
4102            let (mut gsc, mut usc) = (vec![0f32; gpr], vec![0f32; gpr]);
4103            for flat in start..end {
4104                let (e, r) = (flat / inter, flat % inter);
4105                let (gv_view, uv_view) = (&views[e * 2], &views[e * 2 + 1]);
4106                gv_view.scales_into(r, gpr, &mut gsc);
4107                uv_view.scales_into(r, gpr, &mut usc);
4108                for (k, &t) in groups[e].iter().enumerate() {
4109                    if !quantized {
4110                        let x = &xs[t * cols..(t + 1) * cols];
4111                        let gv = q4tp_row_exact(gv_view.nib, r, gpr, x, &gsc);
4112                        let uv = q4tp_row_exact(uv_view.nib, r, gpr, x, &usc);
4113                        unsafe { *ptrs[offs[e] + k].at(r) = (gv / (1.0 + (-gv).exp())) * uv };
4114                        continue;
4115                    }
4116                    let act = &acts[t];
4117                    let mut gv = dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, &gsc) * act.sx;
4118                    let mut uv = dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, &usc) * act.sx;
4119                    for &(j, xv) in &act.outliers {
4120                        let og = q4tp_outlier(gv_view.nib, r, gpr, j, &gsc);
4121                        let ou = q4tp_outlier(uv_view.nib, r, gpr, j, &usc);
4122                        gv += og.0 * og.1 * xv;
4123                        uv += ou.0 * ou.1 * xv;
4124                    }
4125                    let silu_g = gv / (1.0 + (-gv).exp());
4126                    // SAFETY: one worker owns each (expert, row) cell of
4127                    // every output of the expert's group.
4128                    unsafe { *ptrs[offs[e] + k].at(r) = silu_g * uv };
4129                }
4130            }
4131        };
4132        dispatch_rows(pool, pairs.len() * inter, &run);
4133        true
4134    }
4135
4136    /// The per-(expert, token) down terms `moe_down_many` weights and sums,
4137    /// for SEVERAL tokens: `outs[p][o] = down_e[o] · gs[p]` (int8 split of
4138    /// `gs[p]`, VNNI dot, outlier terms — bit-identical to that kernel's
4139    /// `d`), each down row read once for its expert's whole group. The
4140    /// caller sums `w·d` per token in its route order, which reproduces
4141    /// `moe_down_many`'s f32 sequence exactly. Layout as `moe_gate_up_rows`.
4142    pub fn moe_down_rows(
4143        downs: &[&QTensor],
4144        group_lens: &[usize],
4145        gs: &[Vec<f32>],
4146        outs: &mut [Vec<f32>],
4147        pool: Option<&Pool>,
4148    ) -> bool {
4149        if downs.is_empty() || downs.len() != group_lens.len() {
4150            return false;
4151        }
4152        if downs.iter().any(|d| d.per_tensor_only()) {
4153            return false;
4154        }
4155        let rows = downs[0].rows();
4156        let cols = downs[0].cols();
4157        let n_pairs: usize = group_lens.iter().sum();
4158        if cols == 0 || cols % GROUP_SIZE != 0 || gs.len() != n_pairs || outs.len() != n_pairs {
4159            return false;
4160        }
4161        let gpr = cols / GROUP_SIZE;
4162        let mut views = Vec::with_capacity(downs.len());
4163        for d in downs {
4164            if d.has_prism_contract()
4165                || !matches!(
4166                    d,
4167                    Self::Mapped {
4168                        dtype: TensorDtype::Q4TiledP,
4169                        ..
4170                    }
4171                )
4172                || d.rows() != rows
4173                || d.cols() != cols
4174            {
4175                return false;
4176            }
4177            views.push(Q4tpView::new(d.quant_bytes(), rows, cols));
4178        }
4179        if gs.iter().any(|g| g.len() != cols) || outs.iter().any(|o| o.len() != rows) {
4180            return false;
4181        }
4182        let quantized = a8w8_enabled();
4183        let acts: Vec<SplitAct> = if quantized {
4184            gs.iter().map(|g| split_act(g)).collect()
4185        } else {
4186            Vec::new()
4187        };
4188        let mut offs = Vec::with_capacity(group_lens.len());
4189        let mut o = 0usize;
4190        for &l in group_lens {
4191            offs.push(o);
4192            o += l;
4193        }
4194        let ptrs: Vec<SendMut> = outs.iter_mut().map(|o| SendMut(o.as_mut_ptr())).collect();
4195        let (views, ptrs, acts, offs) = (&views, &ptrs, &acts, &offs);
4196        let run = |start: usize, end: usize| {
4197            let mut sc = vec![0f32; gpr];
4198            for flat in start..end {
4199                let (e, r) = (flat / rows, flat % rows);
4200                let v = &views[e];
4201                v.scales_into(r, gpr, &mut sc);
4202                for k in 0..group_lens[e] {
4203                    if !quantized {
4204                        let d = q4tp_row_exact(v.nib, r, gpr, &gs[offs[e] + k], &sc);
4205                        unsafe { *ptrs[offs[e] + k].at(r) = d };
4206                        continue;
4207                    }
4208                    let a = &acts[offs[e] + k];
4209                    let mut d = dot_q4tp_row_i8(v.nib, r, gpr, &a.xq, &sc) * a.sx;
4210                    for &(j, xv) in &a.outliers {
4211                        let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
4212                        d += w * s * xv;
4213                    }
4214                    // SAFETY: one worker owns each (expert, row) cell.
4215                    unsafe { *ptrs[offs[e] + k].at(r) = d };
4216                }
4217            }
4218        };
4219        dispatch_rows(pool, downs.len() * rows, &run);
4220        true
4221    }
4222}
4223
4224/// Batched q8 kernel: same math as qmatvec, the row makes a single
4225/// pass from memory for the whole batch.
4226/// Accelerate CBLAS — the Apple AMX matrix units, the same engine
4227/// llama.cpp's `-ngl 0` prefill rides via ggml-blas.
4228#[cfg(target_os = "macos")]
4229mod accel_blas {
4230    #[link(name = "Accelerate", kind = "framework")]
4231    unsafe extern "C" {
4232        pub fn cblas_sgemm(
4233            order: i32,
4234            trans_a: i32,
4235            trans_b: i32,
4236            m: i32,
4237            n: i32,
4238            k: i32,
4239            alpha: f32,
4240            a: *const f32,
4241            lda: i32,
4242            b: *const f32,
4243            ldb: i32,
4244            beta: f32,
4245            c: *mut f32,
4246            ldc: i32,
4247        );
4248    }
4249}
4250
4251#[cfg(target_os = "macos")]
4252pub(crate) fn accel_gemm_enabled() -> bool {
4253    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4254    *ON.get_or_init(|| std::env::var("CMF_ACCEL").map(|v| v != "0").unwrap_or(true))
4255}
4256
4257/// Off macOS the "accel" GEMM is the portable NEON micro-kernel below —
4258/// same entry point, so the batched-attention path opens on mobile.
4259#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
4260pub(crate) fn accel_gemm_enabled() -> bool {
4261    true
4262}
4263
4264/// Portable NEON f32 GEMM (row-major, optional Bᵀ): a 4×8 fmla
4265/// micro-kernel with A broadcast against B panels — the mobile stand-in
4266/// for Accelerate in the batched causal attention (QKᵀ and P·V). Not a
4267/// BLAS: shapes here are the attention panels (m ≤ heads·chunk,
4268/// k = head_dim or context), and the goal is removing the per-position
4269/// quadratic wall, not peak GEMM.
4270#[cfg(target_arch = "aarch64")]
4271#[allow(clippy::too_many_arguments)]
4272pub(crate) fn neon_gemm_rm(
4273    m: usize,
4274    n: usize,
4275    k: usize,
4276    alpha: f32,
4277    a: &[f32],
4278    lda: usize,
4279    b_mat: &[f32],
4280    ldb: usize,
4281    b_rows_are_n: bool,
4282    c: &mut [f32],
4283    ldc: usize,
4284) {
4285    debug_assert!(a.len() >= (m - 1) * lda + k);
4286    debug_assert!(c.len() >= (m - 1) * ldc + n);
4287    // SAFETY: bounds asserted above; NEON is baseline on aarch64.
4288    unsafe {
4289        use core::arch::aarch64::*;
4290        let mut i = 0usize;
4291        while i < m {
4292            let mi = (m - i).min(4);
4293            let mut j = 0usize;
4294            while j < n {
4295                let nj = (n - j).min(8);
4296                if mi == 4 && nj == 8 {
4297                    let (mut c0a, mut c0b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
4298                    let (mut c1a, mut c1b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
4299                    let (mut c2a, mut c2b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
4300                    let (mut c3a, mut c3b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
4301                    for p in 0..k {
4302                        let (b0, b1) = if b_rows_are_n {
4303                            // B is [n, k]: column p of Bᵀ = element p of
4304                            // eight consecutive B rows — gathered.
4305                            let base = b_mat.as_ptr().add(j * ldb + p);
4306                            let g = |o: usize| *base.add(o * ldb);
4307                            ([g(0), g(1), g(2), g(3)], [g(4), g(5), g(6), g(7)])
4308                        } else {
4309                            let base = b_mat.as_ptr().add(p * ldb + j);
4310                            (
4311                                [*base, *base.add(1), *base.add(2), *base.add(3)],
4312                                [*base.add(4), *base.add(5), *base.add(6), *base.add(7)],
4313                            )
4314                        };
4315                        let bv0 = vld1q_f32(b0.as_ptr());
4316                        let bv1 = vld1q_f32(b1.as_ptr());
4317                        let a0 = vdupq_n_f32(*a.as_ptr().add(i * lda + p));
4318                        let a1 = vdupq_n_f32(*a.as_ptr().add((i + 1) * lda + p));
4319                        let a2 = vdupq_n_f32(*a.as_ptr().add((i + 2) * lda + p));
4320                        let a3 = vdupq_n_f32(*a.as_ptr().add((i + 3) * lda + p));
4321                        c0a = vfmaq_f32(c0a, a0, bv0);
4322                        c0b = vfmaq_f32(c0b, a0, bv1);
4323                        c1a = vfmaq_f32(c1a, a1, bv0);
4324                        c1b = vfmaq_f32(c1b, a1, bv1);
4325                        c2a = vfmaq_f32(c2a, a2, bv0);
4326                        c2b = vfmaq_f32(c2b, a2, bv1);
4327                        c3a = vfmaq_f32(c3a, a3, bv0);
4328                        c3b = vfmaq_f32(c3b, a3, bv1);
4329                    }
4330                    let al = vdupq_n_f32(alpha);
4331                    for (r, (ca, cb)) in [(c0a, c0b), (c1a, c1b), (c2a, c2b), (c3a, c3b)]
4332                        .iter()
4333                        .enumerate()
4334                    {
4335                        let dst = c.as_mut_ptr().add((i + r) * ldc + j);
4336                        vst1q_f32(dst, vmulq_f32(*ca, al));
4337                        vst1q_f32(dst.add(4), vmulq_f32(*cb, al));
4338                    }
4339                } else {
4340                    for r in 0..mi {
4341                        for q in 0..nj {
4342                            let mut acc = 0f32;
4343                            for p in 0..k {
4344                                let bv = if b_rows_are_n {
4345                                    b_mat[(j + q) * ldb + p]
4346                                } else {
4347                                    b_mat[p * ldb + j + q]
4348                                };
4349                                acc += a[(i + r) * lda + p] * bv;
4350                            }
4351                            c[(i + r) * ldc + j + q] = acc * alpha;
4352                        }
4353                    }
4354                }
4355                j += nj;
4356            }
4357            i += mi;
4358        }
4359    }
4360}
4361
4362/// Off-macOS aarch64: the batched attention rides the NEON micro-GEMM.
4363#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
4364#[allow(clippy::too_many_arguments)]
4365pub(crate) fn sgemm_rm(
4366    m: usize,
4367    n: usize,
4368    k: usize,
4369    alpha: f32,
4370    a: &[f32],
4371    lda: usize,
4372    b_mat: &[f32],
4373    ldb: usize,
4374    b_rows_are_n: bool,
4375    c: &mut [f32],
4376    ldc: usize,
4377) {
4378    neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
4379}
4380
4381/// Row-major f32 GEMM, exposed for offline tools (the AWNP pass builds a
4382/// per-layer projection and applies it to every expert; a naive triple loop
4383/// would turn a two-minute job into half an hour).
4384#[allow(clippy::too_many_arguments)]
4385pub fn sgemm_public(
4386    m: usize,
4387    n: usize,
4388    k: usize,
4389    alpha: f32,
4390    a: &[f32],
4391    lda: usize,
4392    b_mat: &[f32],
4393    ldb: usize,
4394    b_rows_are_n: bool,
4395    c: &mut [f32],
4396    ldc: usize,
4397) {
4398    #[cfg(any(target_os = "macos", target_arch = "aarch64"))]
4399    {
4400        sgemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
4401    }
4402    // x86 without Accelerate has no sgemm_rm (the specialized paths there are
4403    // quantized kernels): the packed AVX/FMA/AVX-512 kernels of
4404    // matrixmultiply, single-threaded — offline callers split rows across
4405    // their own threads.
4406    #[cfg(not(any(target_os = "macos", target_arch = "aarch64")))]
4407    {
4408        if m == 0 || n == 0 {
4409            return;
4410        }
4411        assert!(a.len() >= (m - 1) * lda + k && c.len() >= (m - 1) * ldc + n);
4412        let (rsb, csb) = if b_rows_are_n { (1, ldb as isize) } else { (ldb as isize, 1) };
4413        assert!(k == 0 || b_mat.len() > if b_rows_are_n { (n - 1) * ldb + k - 1 } else { (k - 1) * ldb + n - 1 });
4414        // SAFETY: the asserts above bound every element the strides reach.
4415        unsafe {
4416            matrixmultiply::sgemm(
4417                m,
4418                k,
4419                n,
4420                alpha,
4421                a.as_ptr(),
4422                lda as isize,
4423                1,
4424                b_mat.as_ptr(),
4425                rsb,
4426                csb,
4427                0.0,
4428                c.as_mut_ptr(),
4429                ldc as isize,
4430                1,
4431            );
4432        }
4433    }
4434}
4435
4436/// Row-major f32 GEMM on Accelerate: C[m,n] = alpha·A[m,k] × B(ᵀ).
4437/// `b_rows_are_n` = true multiplies by Bᵀ where B is stored [n, k].
4438#[cfg(target_os = "macos")]
4439#[allow(clippy::too_many_arguments)]
4440pub(crate) fn sgemm_rm(
4441    m: usize,
4442    n: usize,
4443    k: usize,
4444    alpha: f32,
4445    a: &[f32],
4446    lda: usize,
4447    b_mat: &[f32],
4448    ldb: usize,
4449    b_rows_are_n: bool,
4450    c: &mut [f32],
4451    ldc: usize,
4452) {
4453    debug_assert!(a.len() >= (m - 1) * lda + k);
4454    debug_assert!(c.len() >= (m - 1) * ldc + n);
4455    // Test hook: route the attention GEMMs through the portable NEON
4456    // micro-kernel ON APPLE SILICON — how the mobile batched attend is
4457    // measured without a phone in the loop. (Intel macOS has no NEON —
4458    // the hook is a no-op there, Accelerate continues below.)
4459    #[cfg(target_arch = "aarch64")]
4460    if std::env::var("CMF_FORCE_NEON_GEMM")
4461        .map(|v| v == "1")
4462        .unwrap_or(false)
4463    {
4464        return neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
4465    }
4466    unsafe {
4467        accel_blas::cblas_sgemm(
4468            101, // RowMajor
4469            111, // NoTrans A
4470            if b_rows_are_n { 112 } else { 111 },
4471            m as i32,
4472            n as i32,
4473            k as i32,
4474            alpha,
4475            a.as_ptr(),
4476            lda as i32,
4477            b_mat.as_ptr(),
4478            ldb as i32,
4479            0.0,
4480            c.as_mut_ptr(),
4481            ldc as i32,
4482        );
4483    }
4484}
4485
4486/// Prefill GEMM through Accelerate (macOS): dequantize q8 rows into
4487/// f32 tiles (scale folded in, pool-parallel) and multiply each tile
4488/// on the AMX with one row-major sgemm. Tiles live in cache, weights
4489/// stream once. Numerics are f32-GEMM (not the int8 dot): prefill
4490/// logits shift within f32 rounding — tolerance-class, like every
4491/// reduction-order change; decode (M=1) never takes this path.
4492#[cfg(target_os = "macos")]
4493fn qmatmat_accel(
4494    q: &[u8],
4495    row_scale: &[f32],
4496    pre: &[std::borrow::Cow<'_, [f32]>],
4497    rows: usize,
4498    cols: usize,
4499    out: &mut [f32],
4500    pool: Option<&Pool>,
4501) {
4502    // NOTE: double-buffering the dequant against the sgemm (a scoped
4503    // thread driving the pool on tile k+1 while the caller multiplies
4504    // tile k) was tried and LOST ~6%: Accelerate's sgemm is itself
4505    // multithreaded, and the dequant workers just steal its cores.
4506    const TR: usize = 2048;
4507    let b = pre.len();
4508    thread_local! {
4509        static XPANEL: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
4510        static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
4511    }
4512    XPANEL.with(|xp| {
4513        WTILE.with(|wt| {
4514            let mut xpanel = xp.borrow_mut();
4515            xpanel.clear();
4516            for x in pre {
4517                xpanel.extend_from_slice(x);
4518            }
4519            let mut wtile = wt.borrow_mut();
4520            wtile.resize(TR * cols, 0.0);
4521            let mut r0 = 0usize;
4522            while r0 < rows {
4523                let tr = TR.min(rows - r0);
4524                // Dequant the tile (scale folded) — pool-parallel.
4525                let wt_addr = SendMut(wtile.as_mut_ptr());
4526                let run = |start: usize, end: usize| {
4527                    for r in start..end {
4528                        let row = &q[(r0 + r) * cols..(r0 + r + 1) * cols];
4529                        let s = row_scale[r0 + r];
4530                        // SAFETY: workers cover disjoint r ranges.
4531                        let dst =
4532                            unsafe { std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols) };
4533                        for (d, &v) in dst.iter_mut().zip(row) {
4534                            *d = (v as i8) as f32 * s;
4535                        }
4536                    }
4537                };
4538                dispatch_rows(pool, tr, &run);
4539                // C[b, tr] (at column r0 of out[b, rows]) = X · Wtileᵀ
4540                unsafe {
4541                    accel_blas::cblas_sgemm(
4542                        101, // RowMajor
4543                        111, // NoTrans A
4544                        112, // Trans B
4545                        b as i32,
4546                        tr as i32,
4547                        cols as i32,
4548                        1.0,
4549                        xpanel.as_ptr(),
4550                        cols as i32,
4551                        wtile.as_ptr(),
4552                        cols as i32,
4553                        0.0,
4554                        out.as_mut_ptr().add(r0),
4555                        rows as i32,
4556                    );
4557                }
4558                r0 += tr;
4559            }
4560        })
4561    });
4562}
4563
4564fn qmatmat(
4565    q: &[u8],
4566    row_scale: &[f32],
4567    pre: &[std::borrow::Cow<'_, [f32]>],
4568    rows: usize,
4569    cols: usize,
4570    out: &mut [f32],
4571    pool: Option<&Pool>,
4572) {
4573    let b = pre.len();
4574    debug_assert_eq!(out.len(), b * rows);
4575    // Big prefill batches ride the AMX (roadmap PR3): the row×batch
4576    // SDOT loop below peaks near the CPU's dot throughput, an order
4577    // below the matrix units. Small tensors and tiny test models stay
4578    // on the exact integer path.
4579    #[cfg(target_os = "macos")]
4580    if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
4581        qmatmat_accel(q, row_scale, pre, rows, cols, out, pool);
4582        return;
4583    }
4584    #[cfg(target_arch = "aarch64")]
4585    if sdot_enabled() {
4586        let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
4587        let out_addr = SendMut(out.as_mut_ptr());
4588        // Blocked 2×4 (mobile prefill: no AMX to fall back on — this
4589        // path IS the ARM prefill GEMM off Apple silicon).
4590        let blocked_ok = blocked_enabled();
4591        let use_i8mm = i8mm_enabled();
4592        if blocked_ok {
4593            let run = |start: usize, end: usize| {
4594                let mut o = start;
4595                while o < end {
4596                    if o + 2 <= end {
4597                        let r0 = &q[o * cols..(o + 1) * cols];
4598                        let r1 = &q[(o + 1) * cols..(o + 2) * cols];
4599                        let mut bi = 0usize;
4600                        while bi + 4 <= acts.len() {
4601                            let xs = [
4602                                acts[bi].xq.as_slice(),
4603                                acts[bi + 1].xq.as_slice(),
4604                                acts[bi + 2].xq.as_slice(),
4605                                acts[bi + 3].xq.as_slice(),
4606                            ];
4607                            let d = if use_i8mm {
4608                                unsafe { dot_i8_smmla_2x4(r0, r1, xs) }
4609                            } else {
4610                                unsafe { dot_i8_sdot_2x4(r0, r1, xs) }
4611                            };
4612                            for (r, row) in [r0, r1].into_iter().enumerate() {
4613                                for k in 0..4 {
4614                                    let act = &acts[bi + k];
4615                                    let mut v = d[r][k] as f32 * act.sx;
4616                                    for &(j, xv) in &act.outliers {
4617                                        v += (row[j] as i8) as f32 * xv;
4618                                    }
4619                                    unsafe {
4620                                        *out_addr.at((bi + k) * rows + o + r) = v * row_scale[o + r]
4621                                    };
4622                                }
4623                            }
4624                            bi += 4;
4625                        }
4626                        while bi < acts.len() {
4627                            for (r, row) in [r0, r1].into_iter().enumerate() {
4628                                let v = row_dot_sdot(row, &acts[bi]) * row_scale[o + r];
4629                                unsafe { *out_addr.at(bi * rows + o + r) = v };
4630                            }
4631                            bi += 1;
4632                        }
4633                        o += 2;
4634                    } else {
4635                        let row = &q[o * cols..(o + 1) * cols];
4636                        for (bi, act) in acts.iter().enumerate() {
4637                            let v = row_dot_sdot(row, act) * row_scale[o];
4638                            unsafe { *out_addr.at(bi * rows + o) = v };
4639                        }
4640                        o += 1;
4641                    }
4642                }
4643            };
4644            dispatch_rows(pool, rows, &run);
4645            return;
4646        }
4647        let run = |start: usize, end: usize| {
4648            for o in start..end {
4649                let row = &q[o * cols..(o + 1) * cols];
4650                for (bi, act) in acts.iter().enumerate() {
4651                    let v = row_dot_sdot(row, act) * row_scale[o];
4652                    unsafe { *out_addr.at(bi * rows + o) = v };
4653                }
4654            }
4655        };
4656        dispatch_rows(pool, rows, &run);
4657        return;
4658    }
4659    // x86 A8W8 batch. Non-VNNI parts take the BLOCKED 2×4 kernel
4660    // (roadmap P0: two weight rows' abs() stay in registers across four
4661    // activation streams); VNNI machines keep the per-row bias-trick
4662    // dot, which is already throughput-bound there.
4663    #[cfg(target_arch = "x86_64")]
4664    if avx2_a8w8_enabled() {
4665        let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
4666        let out_addr = SendMut(out.as_mut_ptr());
4667        // CMF_X86_BLOCKED=0 forces the per-row path (paired in-process
4668        // A/B on noisy shared-vCPU hosts).
4669        let blocked_ok = blocked_enabled();
4670        if !avx512vnni_enabled() && blocked_ok && !row_exact() {
4671            let run = |start: usize, end: usize| {
4672                let mut o = start;
4673                while o < end {
4674                    if o + 2 <= end {
4675                        let r0 = &q[o * cols..(o + 1) * cols];
4676                        let r1 = &q[(o + 1) * cols..(o + 2) * cols];
4677                        let mut bi = 0usize;
4678                        while bi + 4 <= acts.len() {
4679                            let xs = [
4680                                acts[bi].xq.as_slice(),
4681                                acts[bi + 1].xq.as_slice(),
4682                                acts[bi + 2].xq.as_slice(),
4683                                acts[bi + 3].xq.as_slice(),
4684                            ];
4685                            let d = unsafe { dot_i8_i8_avx2_2x4(r0, r1, xs) };
4686                            for (r, row) in [r0, r1].into_iter().enumerate() {
4687                                for k in 0..4 {
4688                                    let act = &acts[bi + k];
4689                                    let mut v = d[r][k] as f32 * act.sx;
4690                                    for &(j, xv) in &act.outliers {
4691                                        v += (row[j] as i8) as f32 * xv;
4692                                    }
4693                                    unsafe {
4694                                        *out_addr.at((bi + k) * rows + o + r) = v * row_scale[o + r]
4695                                    };
4696                                }
4697                            }
4698                            bi += 4;
4699                        }
4700                        while bi < acts.len() {
4701                            for (r, row) in [r0, r1].into_iter().enumerate() {
4702                                let v = row_dot_avx2(row, &acts[bi]) * row_scale[o + r];
4703                                unsafe { *out_addr.at(bi * rows + o + r) = v };
4704                            }
4705                            bi += 1;
4706                        }
4707                        o += 2;
4708                    } else {
4709                        let row = &q[o * cols..(o + 1) * cols];
4710                        for (bi, act) in acts.iter().enumerate() {
4711                            let v = row_dot_avx2(row, act) * row_scale[o];
4712                            unsafe { *out_addr.at(bi * rows + o) = v };
4713                        }
4714                        o += 1;
4715                    }
4716                }
4717            };
4718            dispatch_rows(pool, rows, &run);
4719            return;
4720        }
4721        let run = |start: usize, end: usize| {
4722            for o in start..end {
4723                let row = &q[o * cols..(o + 1) * cols];
4724                for (bi, act) in acts.iter().enumerate() {
4725                    let v = row_dot_avx2(row, act) * row_scale[o];
4726                    unsafe { *out_addr.at(bi * rows + o) = v };
4727                }
4728            }
4729        };
4730        dispatch_rows(pool, rows, &run);
4731        return;
4732    }
4733    let out_addr = SendMut(out.as_mut_ptr());
4734    let run = |start: usize, end: usize| {
4735        for o in start..end {
4736            let row = &q[o * cols..(o + 1) * cols];
4737            for (bi, x) in pre.iter().enumerate() {
4738                let mut acc = 0f32;
4739                for j in 0..cols {
4740                    acc += (row[j] as i8) as f32 * x[j];
4741                }
4742                unsafe { *out_addr.at(bi * rows + o) = acc * row_scale[o] };
4743            }
4744        }
4745    };
4746    dispatch_rows(pool, rows, &run);
4747}
4748
4749/// Split rows across pool workers (shared qmatvec pattern). Self-balancing
4750/// — see `Pool::run_rows` for why a static 1/n split is wrong here.
4751fn dispatch_rows(pool: Option<&Pool>, rows: usize, run: &(dyn Fn(usize, usize) + Sync)) {
4752    match pool {
4753        Some(pool) if rows >= 256 => pool.run_rows(rows, run),
4754        _ => run(0, rows),
4755    }
4756}
4757
4758/// Split a q4_block blob into (packed nibbles, f16 group scales).
4759fn q4_split(bytes: &[u8], rows: usize, cols: usize) -> (&[u8], &[u8]) {
4760    let groups = rows * cols / GROUP_SIZE;
4761    bytes.split_at(groups * 16)
4762}
4763
4764/// SIMD unpack for the dominant vbit width B=4 (94% of rows on the
4765/// log2-shape calibration): 16 packed bytes -> 32 centered i8 values.
4766/// vbit packs MSB-first, so the HIGH nibble is the even element
4767/// (opposite of q4_block's lo-first interleave). Centering is u-7.
4768#[inline]
4769fn vbit_fill4(data: &[u8], buf: &mut [u8]) {
4770    #[cfg(target_arch = "aarch64")]
4771    unsafe {
4772        return vbit_fill4_neon(data, buf);
4773    }
4774    #[cfg(target_arch = "x86_64")]
4775    if avx2_enabled() {
4776        return unsafe { vbit_fill4_avx2(data, buf) };
4777    }
4778    #[allow(unreachable_code)]
4779    for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
4780        let u = unpack8::<4>(&data[blk * 4..]);
4781        for k in 0..8 {
4782            chunk[k] = (u[k] - 7) as i8 as u8;
4783        }
4784    }
4785}
4786
4787#[cfg(target_arch = "aarch64")]
4788#[target_feature(enable = "neon")]
4789unsafe fn vbit_fill4_neon(data: &[u8], buf: &mut [u8]) {
4790    // SAFETY: buf.len() is a multiple of GROUP_SIZE=32; data holds
4791    // buf.len()/2 packed bytes (validated at load).
4792    unsafe {
4793        use core::arch::aarch64::*;
4794        let n = buf.len();
4795        let mask = vdupq_n_u8(0x0F);
4796        let seven = vdupq_n_s8(7);
4797        let mut g = 0usize;
4798        while g * 32 + 32 <= n {
4799            let b = vld1q_u8(data.as_ptr().add(g * 16));
4800            let hi = vshrq_n_u8::<4>(b);
4801            let lo = vandq_u8(b, mask);
4802            let z0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(hi, lo)), seven);
4803            let z1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(hi, lo)), seven);
4804            vst1q_u8(buf.as_mut_ptr().add(g * 32), vreinterpretq_u8_s8(z0));
4805            vst1q_u8(buf.as_mut_ptr().add(g * 32 + 16), vreinterpretq_u8_s8(z1));
4806            g += 1;
4807        }
4808    }
4809}
4810
4811#[cfg(target_arch = "x86_64")]
4812#[target_feature(enable = "avx2")]
4813unsafe fn vbit_fill4_avx2(data: &[u8], buf: &mut [u8]) {
4814    // SAFETY: see vbit_fill4_neon.
4815    unsafe {
4816        use core::arch::x86_64::*;
4817        let n = buf.len();
4818        let mask = _mm_set1_epi8(0x0F);
4819        let seven = _mm256_set1_epi8(7);
4820        let mut g = 0usize;
4821        while g * 32 + 32 <= n {
4822            let b = _mm_loadu_si128(data.as_ptr().add(g * 16) as *const __m128i);
4823            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), mask);
4824            let lo = _mm_and_si128(b, mask);
4825            let z = _mm256_sub_epi8(
4826                _mm256_set_m128i(_mm_unpackhi_epi8(hi, lo), _mm_unpacklo_epi8(hi, lo)),
4827                seven,
4828            );
4829            _mm256_storeu_si256(buf.as_mut_ptr().add(g * 32) as *mut __m256i, z);
4830            g += 1;
4831        }
4832    }
4833}
4834
4835/// Unpack 8 MSB-first B-bit values from exactly B bytes (fixed shifts —
4836/// no serial bit-buffer, auto-vectorizable). Every 32-value group starts
4837/// byte-aligned (32·B/8 is integral for B∈3..8), so groups decompose
4838/// into 4 such blocks.
4839#[inline(always)]
4840fn unpack8<const B: usize>(data: &[u8]) -> [i32; 8] {
4841    let mut acc = 0u64;
4842    for i in 0..B {
4843        acc = (acc << 8) | data[i] as u64;
4844    }
4845    let mask = (1u64 << B) - 1;
4846    let mut out = [0i32; 8];
4847    for (k, o) in out.iter_mut().enumerate() {
4848        *o = ((acc >> ((7 - k) * B)) & mask) as i32;
4849    }
4850    out
4851}
4852
4853/// Fused vbit matvec straight from the mapped bytes (spec §3, P13
4854/// FIG.3): [u8 bits: rows][f16 scales: rows·cols/32][bit-packed rows,
4855/// MSB-first, byte-padded]. Row data offsets are precomputed at load
4856/// (`vbit_row_offsets`) — the per-call prefix scan was O(rows) pure
4857/// overhead on every matvec.
4858#[allow(clippy::too_many_arguments)]
4859fn vbitmatvec(
4860    bytes: &[u8],
4861    offsets: &[usize],
4862    x: &[f32],
4863    rows: usize,
4864    cols: usize,
4865    out: &mut [f32],
4866    pool: Option<&Pool>,
4867) {
4868    debug_assert_eq!(out.len(), rows);
4869    debug_assert_eq!(offsets.len(), rows + 1);
4870
4871    // SDOT path: unpack the row to centered i8 once, then per-group
4872    // int8 dot against the quantized activations — same A8W8 contract
4873    // as q8 (bounded noise; CMF_SDOT=0 keeps the exact scalar path).
4874    if a8w8_enabled() {
4875        let act = split_act(x);
4876        let out_addr = SendMut(out.as_mut_ptr());
4877        let run = move |start: usize, end: usize| {
4878            vbit_range_a8w8(bytes, offsets, x, &act, rows, cols, out_addr, start, end)
4879        };
4880        dispatch_rows(pool, rows, &run);
4881        return;
4882    }
4883
4884    let out_addr = SendMut(out.as_mut_ptr());
4885    let run = move |start: usize, end: usize| {
4886        vbit_range_f32(bytes, offsets, x, rows, cols, out_addr, start, end)
4887    };
4888    dispatch_rows(pool, rows, &run);
4889}
4890
4891/// One vbit row range via the A8W8 int8 path — kernel body of
4892/// `vbitmatvec`, extracted so multi-matrix jobs can drive it for
4893/// several tensors in one dispatch (b=8 rows go exact f32).
4894#[allow(clippy::too_many_arguments)]
4895fn vbit_range_a8w8(
4896    bytes: &[u8],
4897    offsets: &[usize],
4898    x: &[f32],
4899    act: &SplitAct,
4900    rows: usize,
4901    cols: usize,
4902    out: SendMut,
4903    start: usize,
4904    end: usize,
4905) {
4906    let ng = cols / GROUP_SIZE;
4907    let bits = &bytes[..rows];
4908    let sc_off = rows;
4909    let row_dot = |r: usize| -> f32 {
4910        let b = bits[r] as usize;
4911        let l = (1i32 << (b - 1)) - 1;
4912        let mask = (1u64 << b) - 1;
4913        let data = &bytes[offsets[r]..offsets[r + 1]];
4914        if b == 8 {
4915            // u−L reaches 128 → does not fit i8; exact f32 path.
4916            let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
4917            let mut dot = 0f32;
4918            for g in 0..ng {
4919                let so = (r * ng + g) * 2;
4920                let sgf = f16_to_f32(u16::from_le_bytes([
4921                    bytes[sc_off + so],
4922                    bytes[sc_off + so + 1],
4923                ]));
4924                let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
4925                let mut gd = 0f32;
4926                for &xv in xg.iter() {
4927                    if nbits < 8 {
4928                        acc = (acc << 8) | data[idx] as u64;
4929                        idx += 1;
4930                        nbits += 8;
4931                    }
4932                    let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
4933                    nbits -= 8;
4934                    gd += (u - l) as f32 * xv;
4935                }
4936                dot += gd * sgf;
4937            }
4938            return dot;
4939        }
4940        // Per-worker scratch: this closure runs for every row of the
4941        // tensor (lm_head ≈ 150k rows/token) — a heap allocation per
4942        // row was measurable pure overhead.
4943        thread_local! {
4944            static VBIT_SCRATCH: std::cell::RefCell<Vec<u8>> =
4945                const { std::cell::RefCell::new(Vec::new()) };
4946        }
4947        #[inline(always)]
4948        fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
4949            for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
4950                let u = unpack8::<B>(&data[blk * B..]);
4951                for k in 0..8 {
4952                    chunk[k] = (u[k] - l) as i8 as u8;
4953                }
4954            }
4955        }
4956        let _ = mask;
4957        VBIT_SCRATCH.with(|scratch| {
4958            let mut buf = scratch.borrow_mut();
4959            buf.resize(cols, 0);
4960            match b {
4961                3 => fill::<3>(data, l, &mut buf),
4962                4 => vbit_fill4(data, &mut buf),
4963                5 => fill::<5>(data, l, &mut buf),
4964                6 => fill::<6>(data, l, &mut buf),
4965                _ => unreachable!(),
4966            }
4967            let mut dot = 0f32;
4968            for g in 0..ng {
4969                let so = (r * ng + g) * 2;
4970                let s = f16_to_f32(u16::from_le_bytes([
4971                    bytes[sc_off + so],
4972                    bytes[sc_off + so + 1],
4973                ]));
4974                let d = dot_i8_i8(
4975                    &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
4976                    &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
4977                ) as f32
4978                    * act.sx;
4979                dot += d * s;
4980            }
4981            for &(j, xv) in &act.outliers {
4982                let so = (r * ng + j / GROUP_SIZE) * 2;
4983                let s = f16_to_f32(u16::from_le_bytes([
4984                    bytes[sc_off + so],
4985                    bytes[sc_off + so + 1],
4986                ]));
4987                // xq is zeroed at outlier slots — add the exact term.
4988                dot += (buf[j] as i8) as f32 * s * xv;
4989            }
4990            dot
4991        })
4992    };
4993    for r in start..end {
4994        // SAFETY: disjoint row ranges per worker.
4995        unsafe { *out.at(r) = row_dot(r) };
4996    }
4997}
4998
4999/// Exact scalar vbit row range (same extraction, non-SDOT path).
5000#[allow(clippy::too_many_arguments)]
5001fn vbit_range_f32(
5002    bytes: &[u8],
5003    offsets: &[usize],
5004    x: &[f32],
5005    rows: usize,
5006    cols: usize,
5007    out: SendMut,
5008    start: usize,
5009    end: usize,
5010) {
5011    let ng = cols / GROUP_SIZE;
5012    let bits = &bytes[..rows];
5013    let sc_off = rows;
5014    // Per-bit-width specialized inner loops: the compiler unrolls the
5015    // constant shifts (the generic bit-buffer loop was branch-bound —
5016    // 5.6 vs 13.2 tok/s q4 on the 0.8B).
5017    #[inline(always)]
5018    fn dot_row<const B: usize>(
5019        data: &[u8],
5020        bytes: &[u8],
5021        sc_off: usize,
5022        r: usize,
5023        ng: usize,
5024        x: &[f32],
5025    ) -> f32 {
5026        let l = ((1i32 << (B - 1)) - 1) as f32;
5027        let gbytes = GROUP_SIZE * B / 8;
5028        let mut dot = 0f32;
5029        for g in 0..ng {
5030            let so = (r * ng + g) * 2;
5031            let s = f16_to_f32(u16::from_le_bytes([
5032                bytes[sc_off + so],
5033                bytes[sc_off + so + 1],
5034            ]));
5035            let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
5036            let gd0 = &data[g * gbytes..(g + 1) * gbytes];
5037            let mut gd = 0f32;
5038            for blk in 0..GROUP_SIZE / 8 {
5039                let u = unpack8::<B>(&gd0[blk * B..]);
5040                let xb = &xg[blk * 8..blk * 8 + 8];
5041                for k in 0..8 {
5042                    gd += (u[k] as f32 - l) * xb[k];
5043                }
5044            }
5045            dot += gd * s;
5046        }
5047        dot
5048    }
5049    for r in start..end {
5050        let data = &bytes[offsets[r]..offsets[r + 1]];
5051        let v = match bits[r] {
5052            3 => dot_row::<3>(data, bytes, sc_off, r, ng, x),
5053            4 => dot_row::<4>(data, bytes, sc_off, r, ng, x),
5054            5 => dot_row::<5>(data, bytes, sc_off, r, ng, x),
5055            6 => dot_row::<6>(data, bytes, sc_off, r, ng, x),
5056            8 => dot_row::<8>(data, bytes, sc_off, r, ng, x),
5057            b => unreachable!("vbit bit-width {b} (validated at load)"),
5058        };
5059        // SAFETY: disjoint row ranges per worker.
5060        unsafe { *out.at(r) = v };
5061    }
5062}
5063
5064/// Fused two-input vbit matvec: each row is unpacked from the mmap ONCE
5065/// and dotted against BOTH activations (MTP verify / pair prefill used
5066/// to run two full matvecs — double weight traffic and double unpack).
5067/// Per-input math is identical to `vbitmatvec` → same accuracy contract.
5068#[allow(clippy::too_many_arguments)]
5069fn vbitmatvec2(
5070    bytes: &[u8],
5071    offsets: &[usize],
5072    x1: &[f32],
5073    x2: &[f32],
5074    rows: usize,
5075    cols: usize,
5076    o1: &mut [f32],
5077    o2: &mut [f32],
5078    pool: Option<&Pool>,
5079) {
5080    debug_assert_eq!(o1.len(), rows);
5081    debug_assert_eq!(o2.len(), rows);
5082
5083    if a8w8_enabled() {
5084        let a1 = split_act(x1);
5085        let a2 = split_act(x2);
5086        let p1 = SendMut(o1.as_mut_ptr());
5087        let p2 = SendMut(o2.as_mut_ptr());
5088        let run = move |start: usize, end: usize| {
5089            vbit_range2_a8w8(
5090                bytes, offsets, x1, x2, &a1, &a2, rows, cols, p1, p2, start, end,
5091            )
5092        };
5093        dispatch_rows(pool, rows, &run);
5094        return;
5095    }
5096
5097    let p1 = SendMut(o1.as_mut_ptr());
5098    let p2 = SendMut(o2.as_mut_ptr());
5099    let run = move |start: usize, end: usize| {
5100        vbit_range2_f32(bytes, offsets, x1, x2, rows, cols, p1, p2, start, end)
5101    };
5102    dispatch_rows(pool, rows, &run);
5103}
5104
5105/// Two-input vbit row range via the A8W8 int8 path — kernel body of
5106/// `vbitmatvec2`, extracted for pair multi-matrix jobs (b=8 rows go
5107/// exact f32 for both lanes, bits streamed once).
5108#[allow(clippy::too_many_arguments)]
5109fn vbit_range2_a8w8(
5110    bytes: &[u8],
5111    offsets: &[usize],
5112    x1: &[f32],
5113    x2: &[f32],
5114    a1: &SplitAct,
5115    a2: &SplitAct,
5116    rows: usize,
5117    cols: usize,
5118    p1: SendMut,
5119    p2: SendMut,
5120    start: usize,
5121    end: usize,
5122) {
5123    let ng = cols / GROUP_SIZE;
5124    let bits = &bytes[..rows];
5125    let sc_off = rows;
5126    let row_dots = |r: usize| -> (f32, f32) {
5127        let b = bits[r] as usize;
5128        let l = (1i32 << (b - 1)) - 1;
5129        let data = &bytes[offsets[r]..offsets[r + 1]];
5130        if b == 8 {
5131            // u−L reaches 128 → does not fit i8; exact f32 path,
5132            // bits still streamed once for both lanes.
5133            let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
5134            let (mut d1, mut d2) = (0f32, 0f32);
5135            for g in 0..ng {
5136                let so = (r * ng + g) * 2;
5137                let sgf = f16_to_f32(u16::from_le_bytes([
5138                    bytes[sc_off + so],
5139                    bytes[sc_off + so + 1],
5140                ]));
5141                let (mut g1, mut g2) = (0f32, 0f32);
5142                for k in 0..GROUP_SIZE {
5143                    if nbits < 8 {
5144                        acc = (acc << 8) | data[idx] as u64;
5145                        idx += 1;
5146                        nbits += 8;
5147                    }
5148                    let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
5149                    nbits -= 8;
5150                    let w = (u - l) as f32;
5151                    g1 += w * x1[g * GROUP_SIZE + k];
5152                    g2 += w * x2[g * GROUP_SIZE + k];
5153                }
5154                d1 += g1 * sgf;
5155                d2 += g2 * sgf;
5156            }
5157            return (d1, d2);
5158        }
5159        thread_local! {
5160            static VBIT_SCRATCH2: std::cell::RefCell<Vec<u8>> =
5161                const { std::cell::RefCell::new(Vec::new()) };
5162        }
5163        #[inline(always)]
5164        fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
5165            for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
5166                let u = unpack8::<B>(&data[blk * B..]);
5167                for k in 0..8 {
5168                    chunk[k] = (u[k] - l) as i8 as u8;
5169                }
5170            }
5171        }
5172        VBIT_SCRATCH2.with(|scratch| {
5173            let mut buf = scratch.borrow_mut();
5174            buf.resize(cols, 0);
5175            match b {
5176                3 => fill::<3>(data, l, &mut buf),
5177                4 => vbit_fill4(data, &mut buf),
5178                5 => fill::<5>(data, l, &mut buf),
5179                6 => fill::<6>(data, l, &mut buf),
5180                _ => unreachable!(),
5181            }
5182            let (mut d1, mut d2) = (0f32, 0f32);
5183            for g in 0..ng {
5184                let so = (r * ng + g) * 2;
5185                let s = f16_to_f32(u16::from_le_bytes([
5186                    bytes[sc_off + so],
5187                    bytes[sc_off + so + 1],
5188                ]));
5189                let wg = &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
5190                let v1 = dot_i8_i8(wg, &a1.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a1.sx;
5191                let v2 = dot_i8_i8(wg, &a2.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a2.sx;
5192                d1 += v1 * s;
5193                d2 += v2 * s;
5194            }
5195            for &(j, xv) in &a1.outliers {
5196                let so = (r * ng + j / GROUP_SIZE) * 2;
5197                let s = f16_to_f32(u16::from_le_bytes([
5198                    bytes[sc_off + so],
5199                    bytes[sc_off + so + 1],
5200                ]));
5201                d1 += (buf[j] as i8) as f32 * s * xv;
5202            }
5203            for &(j, xv) in &a2.outliers {
5204                let so = (r * ng + j / GROUP_SIZE) * 2;
5205                let s = f16_to_f32(u16::from_le_bytes([
5206                    bytes[sc_off + so],
5207                    bytes[sc_off + so + 1],
5208                ]));
5209                d2 += (buf[j] as i8) as f32 * s * xv;
5210            }
5211            (d1, d2)
5212        })
5213    };
5214    for r in start..end {
5215        let (v1, v2) = row_dots(r);
5216        // SAFETY: disjoint row ranges per worker.
5217        unsafe {
5218            *p1.at(r) = v1;
5219            *p2.at(r) = v2;
5220        }
5221    }
5222}
5223
5224/// Two-input exact scalar vbit row range (same extraction) —
5225/// per-bit-width specialized, two accumulators per row; per-lane
5226/// accumulation order matches `vbitmatvec` exactly.
5227#[allow(clippy::too_many_arguments)]
5228fn vbit_range2_f32(
5229    bytes: &[u8],
5230    offsets: &[usize],
5231    x1: &[f32],
5232    x2: &[f32],
5233    rows: usize,
5234    cols: usize,
5235    p1: SendMut,
5236    p2: SendMut,
5237    start: usize,
5238    end: usize,
5239) {
5240    let ng = cols / GROUP_SIZE;
5241    let bits = &bytes[..rows];
5242    let sc_off = rows;
5243    #[inline(always)]
5244    #[allow(clippy::too_many_arguments)]
5245    fn dot_row2<const B: usize>(
5246        data: &[u8],
5247        bytes: &[u8],
5248        sc_off: usize,
5249        r: usize,
5250        ng: usize,
5251        x1: &[f32],
5252        x2: &[f32],
5253    ) -> (f32, f32) {
5254        let l = ((1i32 << (B - 1)) - 1) as f32;
5255        let gbytes = GROUP_SIZE * B / 8;
5256        let (mut d1, mut d2) = (0f32, 0f32);
5257        for g in 0..ng {
5258            let so = (r * ng + g) * 2;
5259            let s = f16_to_f32(u16::from_le_bytes([
5260                bytes[sc_off + so],
5261                bytes[sc_off + so + 1],
5262            ]));
5263            let x1g = &x1[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
5264            let x2g = &x2[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
5265            let gd0 = &data[g * gbytes..(g + 1) * gbytes];
5266            let (mut g1, mut g2) = (0f32, 0f32);
5267            for blk in 0..GROUP_SIZE / 8 {
5268                let u = unpack8::<B>(&gd0[blk * B..]);
5269                for k in 0..8 {
5270                    let w = u[k] as f32 - l;
5271                    g1 += w * x1g[blk * 8 + k];
5272                    g2 += w * x2g[blk * 8 + k];
5273                }
5274            }
5275            d1 += g1 * s;
5276            d2 += g2 * s;
5277        }
5278        (d1, d2)
5279    }
5280    for r in start..end {
5281        let data = &bytes[offsets[r]..offsets[r + 1]];
5282        let (v1, v2) = match bits[r] {
5283            3 => dot_row2::<3>(data, bytes, sc_off, r, ng, x1, x2),
5284            4 => dot_row2::<4>(data, bytes, sc_off, r, ng, x1, x2),
5285            5 => dot_row2::<5>(data, bytes, sc_off, r, ng, x1, x2),
5286            6 => dot_row2::<6>(data, bytes, sc_off, r, ng, x1, x2),
5287            8 => dot_row2::<8>(data, bytes, sc_off, r, ng, x1, x2),
5288            b => unreachable!("vbit bit-width {b} (validated at load)"),
5289        };
5290        // SAFETY: disjoint row ranges per worker.
5291        unsafe {
5292            *p1.at(r) = v1;
5293            *p2.at(r) = v2;
5294        }
5295    }
5296}
5297
5298// ───────────────────── q4_tiled kernels (§4.3) ─────────────────────
5299
5300/// One q4_tiled row dot on the A8W8 int8 path: per 32-group the tile
5301/// is ONE sequential read — [f16 scale][16B nibbles] — versus the two
5302/// distant streams of the split layout. Values/order identical to the
5303/// split kernels.
5304#[inline]
5305#[allow(unreachable_code)]
5306fn dot_q4t_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
5307    #[cfg(target_arch = "aarch64")]
5308    unsafe {
5309        return dot_q4t_row_sdot(bytes, r, gpr, xq);
5310    }
5311    #[cfg(target_arch = "x86_64")]
5312    unsafe {
5313        if vnni_tiles_enabled() {
5314            return dot_q4t_row_vnni(bytes, r, gpr, xq);
5315        }
5316        return dot_q4t_row_avx2(bytes, r, gpr, xq);
5317    }
5318    let mut acc = 0f32;
5319    for gi in 0..gpr {
5320        let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
5321        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
5322        let mut d = 0i32;
5323        for (k, &b) in tile[2..].iter().enumerate() {
5324            d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
5325                + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
5326        }
5327        acc += d as f32 * s;
5328    }
5329    acc
5330}
5331
5332#[cfg(target_arch = "aarch64")]
5333#[target_feature(enable = "neon,dotprod")]
5334unsafe fn dot_q4t_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
5335    // SAFETY: callers uphold slice-length contracts (18B tile per group,
5336    // xq.len() == gpr·GROUP_SIZE).
5337    unsafe {
5338        use core::arch::aarch64::*;
5339        use core::arch::asm;
5340        let lomask = vdupq_n_u8(0x0F);
5341        let eight = vdupq_n_s8(8);
5342        let mut acc = 0f32;
5343        for gi in 0..gpr {
5344            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
5345            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5346            let b = vld1q_u8(t.add(2));
5347            let lo = vandq_u8(b, lomask);
5348            let hi = vshrq_n_u8::<4>(b);
5349            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
5350            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
5351            let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
5352            let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
5353            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5354            asm!(
5355                "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
5356                "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
5357                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5358                e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
5359                options(pure, nomem, nostack),
5360            );
5361            acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
5362        }
5363        acc
5364    }
5365}
5366
5367#[cfg(target_arch = "x86_64")]
5368#[target_feature(enable = "avx2")]
5369unsafe fn dot_q4t_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
5370    // SAFETY: see dot_q4t_row_sdot.
5371    unsafe {
5372        use core::arch::x86_64::*;
5373        let lomask = _mm_set1_epi8(0x0F);
5374        let eight = _mm256_set1_epi8(8);
5375        let ones = _mm256_set1_epi16(1);
5376        let mut acc = 0f32;
5377        for gi in 0..gpr {
5378            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
5379            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5380            let b = _mm_loadu_si128(t.add(2) as *const __m128i);
5381            let lo = _mm_and_si128(b, lomask);
5382            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
5383            let w = _mm256_sub_epi8(
5384                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5385                eight,
5386            );
5387            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5388            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
5389            let d = _mm256_madd_epi16(p16, ones);
5390            let hi128 = _mm256_extracti128_si256::<1>(d);
5391            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
5392            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
5393            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
5394            acc += _mm_cvtsi128_si32(s32) as f32 * s;
5395        }
5396        acc
5397    }
5398}
5399
5400/// VNNI twin of `dot_q4t_row_avx2`: same unpack, `vpdpbusd` replaces
5401/// the maddubs+madd pair (see `dpbusd_hsum` — sums are bit-identical).
5402/// 256-bit VL encoding, so the VEX `vpsignb` stays usable.
5403#[cfg(target_arch = "x86_64")]
5404#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
5405unsafe fn dot_q4t_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
5406    // SAFETY: see dot_q4t_row_sdot.
5407    unsafe {
5408        use core::arch::x86_64::*;
5409        let lomask = _mm_set1_epi8(0x0F);
5410        let eight = _mm256_set1_epi8(8);
5411        let mut acc = 0f32;
5412        for gi in 0..gpr {
5413            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
5414            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5415            let b = _mm_loadu_si128(t.add(2) as *const __m128i);
5416            let lo = _mm_and_si128(b, lomask);
5417            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
5418            let w = _mm256_sub_epi8(
5419                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5420                eight,
5421            );
5422            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5423            let d = dpbusd_hsum(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
5424            acc += d as f32 * s;
5425        }
5426        acc
5427    }
5428}
5429
5430/// One q4_tiled row against FOUR activation streams: the nibble unpack
5431/// and abs() happen once per group instead of once per (group,
5432/// activation) — the unpack is the dominant per-element cost of the
5433/// tiled format (roadmap P0 portable blocking, q4t leg).
5434#[cfg(target_arch = "x86_64")]
5435// `fma` is NOT implied by `avx2`: without it LLVM lowers _mm256_fmadd_ps
5436// to a libm call per lane — measured 2x slower than the reduction this
5437// kernel replaces. The runtime gate (`avx2_enabled`) already requires
5438// both features, so declaring it here is safe.
5439#[target_feature(enable = "avx2,fma")]
5440unsafe fn dot_q4t_row_1x4_avx2(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
5441    // SAFETY: callers uphold the 18B-tile and xq-length contracts.
5442    unsafe {
5443        use core::arch::x86_64::*;
5444        let lomask = _mm_set1_epi8(0x0F);
5445        let eight = _mm256_set1_epi8(8);
5446        let ones = _mm256_set1_epi16(1);
5447        // One f32 accumulator VECTOR per activation, reduced once at the
5448        // end. Folding each group's i32 lanes to a scalar inside the loop
5449        // costs an extracti128 + three shift/add + a movd — a cross-lane
5450        // dependency chain per (group, activation), 288 of them per row at
5451        // cols=2304. The per-group scale is what forces a float
5452        // accumulator; it does not force a horizontal sum.
5453        //
5454        // The four accumulators are NAMED, not an array: as `[__m256; 4]`
5455        // indexed by a loop variable LLVM keeps them in memory and every
5456        // group pays four 32-byte loads and stores. That alone made this
5457        // kernel 2x SLOWER than the per-group reduction it replaces
5458        // (measured on the EPYC box: 150 s vs 71 s for two 256² steps).
5459        let mut f0 = _mm256_setzero_ps();
5460        let mut f1 = _mm256_setzero_ps();
5461        let mut f2 = _mm256_setzero_ps();
5462        let mut f3 = _mm256_setzero_ps();
5463        for gi in 0..gpr {
5464            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
5465            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5466            let sv = _mm256_set1_ps(s);
5467            let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
5468            let lo = _mm_and_si128(bb, lomask);
5469            let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
5470            let w = _mm256_sub_epi8(
5471                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5472                eight,
5473            );
5474            let aw = _mm256_abs_epi8(w);
5475            let off = gi * GROUP_SIZE;
5476            let dot = |xq: &[i8]| {
5477                let x = _mm256_loadu_si256(xq.as_ptr().add(off) as *const __m256i);
5478                let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
5479                _mm256_cvtepi32_ps(_mm256_madd_epi16(p16, ones))
5480            };
5481            f0 = _mm256_fmadd_ps(dot(xs[0]), sv, f0);
5482            f1 = _mm256_fmadd_ps(dot(xs[1]), sv, f1);
5483            f2 = _mm256_fmadd_ps(dot(xs[2]), sv, f2);
5484            f3 = _mm256_fmadd_ps(dot(xs[3]), sv, f3);
5485        }
5486        [
5487            hsum256_ps(f0),
5488            hsum256_ps(f1),
5489            hsum256_ps(f2),
5490            hsum256_ps(f3),
5491        ]
5492    }
5493}
5494
5495/// Horizontal sum of eight f32 lanes — the one cross-lane reduction the
5496/// blocked kernels pay, once per row instead of once per group.
5497#[cfg(target_arch = "x86_64")]
5498#[target_feature(enable = "avx2")]
5499#[inline]
5500unsafe fn hsum256_ps(v: core::arch::x86_64::__m256) -> f32 {
5501    // SAFETY: pure register arithmetic on the caller's vector.
5502    unsafe {
5503        use core::arch::x86_64::*;
5504        let hi = _mm256_extractf128_ps::<1>(v);
5505        let s = _mm_add_ps(_mm256_castps256_ps128(v), hi);
5506        let s = _mm_add_ps(s, _mm_movehl_ps(s, s));
5507        let s = _mm_add_ss(s, _mm_shuffle_ps::<0x55>(s, s));
5508        _mm_cvtss_f32(s)
5509    }
5510}
5511
5512/// VNNI twin of `dot_q4t_row_1x4_avx2` (see `dpbusd_hsum`).
5513#[cfg(target_arch = "x86_64")]
5514#[target_feature(enable = "avx2,fma,avx512f,avx512bw,avx512vl,avx512vnni")]
5515unsafe fn dot_q4t_row_1x4_vnni(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
5516    // SAFETY: callers uphold the 18B-tile and xq-length contracts.
5517    unsafe {
5518        use core::arch::x86_64::*;
5519        let lomask = _mm_set1_epi8(0x0F);
5520        let eight = _mm256_set1_epi8(8);
5521        // Same shape as the AVX2 twin: accumulate in f32 vectors and pay
5522        // one cross-lane reduction per row, not per (group, activation).
5523        let mut f0 = _mm256_setzero_ps();
5524        let mut f1 = _mm256_setzero_ps();
5525        let mut f2 = _mm256_setzero_ps();
5526        let mut f3 = _mm256_setzero_ps();
5527        for gi in 0..gpr {
5528            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
5529            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5530            let sv = _mm256_set1_ps(s);
5531            let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
5532            let lo = _mm_and_si128(bb, lomask);
5533            let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
5534            let w = _mm256_sub_epi8(
5535                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5536                eight,
5537            );
5538            let aw = _mm256_abs_epi8(w);
5539            let off = gi * GROUP_SIZE;
5540            let dot = |xq: &[i8]| {
5541                let x = _mm256_loadu_si256(xq.as_ptr().add(off) as *const __m256i);
5542                _mm256_cvtepi32_ps(_mm256_dpbusd_epi32(
5543                    _mm256_setzero_si256(),
5544                    aw,
5545                    _mm256_sign_epi8(x, w),
5546                ))
5547            };
5548            f0 = _mm256_fmadd_ps(dot(xs[0]), sv, f0);
5549            f1 = _mm256_fmadd_ps(dot(xs[1]), sv, f1);
5550            f2 = _mm256_fmadd_ps(dot(xs[2]), sv, f2);
5551            f3 = _mm256_fmadd_ps(dot(xs[3]), sv, f3);
5552        }
5553        let acc = [
5554            hsum256_ps(f0),
5555            hsum256_ps(f1),
5556            hsum256_ps(f2),
5557            hsum256_ps(f3),
5558        ];
5559        acc
5560    }
5561}
5562
5563/// ARM twin of `dot_q4t_row_1x4_avx2`: one nibble unpack per group
5564/// serves FOUR activation streams. Per stream the group order and f32
5565/// accumulation match `dot_q4t_row_sdot` exactly — batch == matvec
5566/// bit-for-bit.
5567#[cfg(target_arch = "aarch64")]
5568#[target_feature(enable = "neon,dotprod")]
5569unsafe fn dot_q4t_row_1x4_sdot(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
5570    // SAFETY: callers uphold the 18B-tile and xq-length contracts.
5571    unsafe {
5572        use core::arch::aarch64::*;
5573        use core::arch::asm;
5574        let lomask = vdupq_n_u8(0x0F);
5575        let eight = vdupq_n_s8(8);
5576        let mut acc = [0f32; 4];
5577        for gi in 0..gpr {
5578            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
5579            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5580            let b = vld1q_u8(t.add(2));
5581            let lo = vandq_u8(b, lomask);
5582            let hi = vshrq_n_u8::<4>(b);
5583            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
5584            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
5585            for (k, xq) in xs.iter().enumerate() {
5586                let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
5587                let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
5588                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5589                asm!(
5590                    "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
5591                    "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
5592                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5593                    e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
5594                    options(pure, nomem, nostack),
5595                );
5596                acc[k] += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
5597            }
5598        }
5599        acc
5600    }
5601}
5602
5603/// Exact-term correction for A8W8 outliers on a tiled row.
5604#[inline]
5605fn q4t_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
5606    let gi = j / GROUP_SIZE;
5607    let k = j % GROUP_SIZE;
5608    let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
5609    let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
5610    let byte = tile[2 + k / 2];
5611    let nib = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
5612    ((nib as i32 - 8) as f32, s)
5613}
5614
5615/// Exact scalar q4_tiled row (CMF_SDOT=0 contract) — same pairwise
5616/// accumulation shape as `q4_range_f32`.
5617#[inline]
5618fn q4t_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
5619    let mut acc = 0f32;
5620    for gi in 0..gpr {
5621        let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
5622        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
5623        let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
5624        let mut ga = 0f32;
5625        for (k, &b) in tile[2..].iter().enumerate() {
5626            ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
5627                + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
5628        }
5629        acc += ga * s;
5630    }
5631    acc
5632}
5633
5634/// Split view of a `q4tp` payload. The three planes are resolved once per
5635/// matvec instead of per row — `q4tp_sections` is cheap, but doing it inside
5636/// the row loop would put a division on the hot path for nothing.
5637struct Q4tpView<'a> {
5638    nib: &'a [u8],
5639    params: &'a [u8],
5640    codes: &'a [u8],
5641    stride: usize,
5642    /// q2tp reads the ladder with rung 0 = exact zero.
5643    zero_rung: bool,
5644}
5645
5646impl<'a> Q4tpView<'a> {
5647    fn new(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
5648        let (params_off, codes_off, stride) = q4tp_sections(rows, cols);
5649        Self {
5650            nib: &bytes[..params_off],
5651            params: &bytes[params_off..codes_off],
5652            codes: &bytes[codes_off..],
5653            stride,
5654            zero_rung: false,
5655        }
5656    }
5657
5658    /// The q2tp view: identical params/codes planes, 8 B weight chunks.
5659    fn new_q2(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
5660        let (params_off, codes_off, stride) = q2tp_sections(rows, cols);
5661        Self {
5662            nib: &bytes[..params_off],
5663            params: &bytes[params_off..codes_off],
5664            codes: &bytes[codes_off..],
5665            stride,
5666            zero_rung: true,
5667        }
5668    }
5669
5670    /// Expand row `r`'s per-tile scales into `out` (length `gpr`).
5671    ///
5672    /// Doing this once per row — rather than decoding a 5-bit code inside the
5673    /// tile loop — is what makes the format free at runtime. Random access to
5674    /// a packed 5-bit field costs a division, two bounds checks and a branch;
5675    /// the tile's actual work is two `sdot`s, so per-tile decoding dominated
5676    /// the kernel and cost 5x (measured: 1.4 vs 6.9 tok/s on Nanbeige-3B).
5677    /// Walking the plane sequentially with a bit accumulator is ~3 ops.
5678    /// Eight 5-bit codes are exactly five bytes, so a whole group of
5679    /// eight decodes from one little-endian word at fixed shifts. The
5680    /// bit-accumulator this replaces carried a data-dependent `while
5681    /// have < 5` refill whose branch sat in the innermost loop of every
5682    /// q4tp row; a decode profile put this function above the dot
5683    /// products it feeds. Same bitstream, same codes — just no branch
5684    /// and eight independent extractions.
5685    #[inline]
5686    fn scales_into(&self, r: usize, gpr: usize, out: &mut [f32]) {
5687        let tab = if self.zero_rung {
5688            q2tp_ladder(self.params, r)
5689        } else {
5690            q4tp_ladder(self.params, r)
5691        };
5692        let codes = &self.codes[r * self.stride..(r + 1) * self.stride];
5693        let out = &mut out[..gpr];
5694        let mut chunks = out.chunks_exact_mut(8);
5695        let mut ci = 0usize;
5696        for c in &mut chunks {
5697            let w = u64::from(codes[ci])
5698                | u64::from(codes[ci + 1]) << 8
5699                | u64::from(codes[ci + 2]) << 16
5700                | u64::from(codes[ci + 3]) << 24
5701                | u64::from(codes[ci + 4]) << 32;
5702            for (k, o) in c.iter_mut().enumerate() {
5703                *o = tab[((w >> (5 * k)) & 31) as usize];
5704            }
5705            ci += 5;
5706        }
5707        // Fewer than eight codes left: the shared total accessor, which
5708        // tolerates a 5-bit field whose spill byte is past the stride.
5709        let tail = &codes[ci..];
5710        for (k, o) in chunks.into_remainder().iter_mut().enumerate() {
5711            *o = tab[q4tp_code(tail, k)];
5712        }
5713    }
5714}
5715
5716#[inline]
5717fn dot_q4tp_row_i8(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
5718    #[cfg(target_arch = "aarch64")]
5719    unsafe {
5720        return dot_q4tp_row_sdot(nib, r, gpr, xq, scales);
5721    }
5722    #[cfg(target_arch = "x86_64")]
5723    unsafe {
5724        if vnni_tiles_enabled() {
5725            return dot_q4tp_row_vnni(nib, r, gpr, xq, scales);
5726        }
5727        return dot_q4tp_row_avx2(nib, r, gpr, xq, scales);
5728    }
5729    #[allow(unreachable_code)]
5730    {
5731        let mut acc = 0f32;
5732        for gi in 0..gpr {
5733            let tile = &nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
5734            let s = scales[gi];
5735            let mut d = 0i32;
5736            for (k, &b) in tile.iter().enumerate() {
5737                d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
5738                    + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
5739            }
5740            acc += d as f32 * s;
5741        }
5742        acc
5743    }
5744}
5745
5746/// q4tp twin of `dot_q4t_row_sdot`: identical nibble math, but the tile
5747/// stride is 16 B (no inline scale) and the scale is a ladder lookup.
5748#[cfg(target_arch = "aarch64")]
5749#[target_feature(enable = "neon,dotprod")]
5750unsafe fn dot_q4tp_row_sdot(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
5751    // SAFETY: callers uphold slice-length contracts (16B tile per group,
5752    // xq.len() == gpr·GROUP_SIZE, codes covering gpr 5-bit fields).
5753    unsafe {
5754        use core::arch::aarch64::*;
5755        use core::arch::asm;
5756        let lomask = vdupq_n_u8(0x0F);
5757        let eight = vdupq_n_s8(8);
5758        let mut acc = 0f32;
5759        for gi in 0..gpr {
5760            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5761            let s = *scales.get_unchecked(gi);
5762            let b = vld1q_u8(t);
5763            let lo = vandq_u8(b, lomask);
5764            let hi = vshrq_n_u8::<4>(b);
5765            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
5766            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
5767            let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
5768            let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
5769            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5770            asm!(
5771                "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
5772                "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
5773                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5774                e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
5775                options(pure, nomem, nostack),
5776            );
5777            acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
5778        }
5779        acc
5780    }
5781}
5782
5783#[cfg(target_arch = "x86_64")]
5784#[target_feature(enable = "avx2")]
5785unsafe fn dot_q4tp_row_avx2(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
5786    // SAFETY: see dot_q4tp_row_sdot.
5787    unsafe {
5788        use core::arch::x86_64::*;
5789        let lomask = _mm_set1_epi8(0x0F);
5790        let eight = _mm256_set1_epi8(8);
5791        let ones = _mm256_set1_epi16(1);
5792        let mut acc = 0f32;
5793        for gi in 0..gpr {
5794            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5795            let s = *scales.get_unchecked(gi);
5796            let b = _mm_loadu_si128(t as *const __m128i);
5797            let lo = _mm_and_si128(b, lomask);
5798            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
5799            let w = _mm256_sub_epi8(
5800                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5801                eight,
5802            );
5803            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5804            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
5805            let d = _mm256_madd_epi16(p16, ones);
5806            let hi128 = _mm256_extracti128_si256::<1>(d);
5807            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
5808            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
5809            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
5810            acc += _mm_cvtsi128_si32(s32) as f32 * s;
5811        }
5812        acc
5813    }
5814}
5815
5816/// VNNI twin of `dot_q4tp_row_avx2` (see `dot_q4t_row_vnni` for why the
5817/// 256-bit VL encoding is the one to use here).
5818#[cfg(target_arch = "x86_64")]
5819#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
5820unsafe fn dot_q4tp_row_vnni(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
5821    // SAFETY: see dot_q4tp_row_sdot.
5822    unsafe {
5823        use core::arch::x86_64::*;
5824        let lomask = _mm_set1_epi8(0x0F);
5825        let eight = _mm256_set1_epi8(8);
5826        let mut acc = 0f32;
5827        for gi in 0..gpr {
5828            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5829            let s = *scales.get_unchecked(gi);
5830            let b = _mm_loadu_si128(t as *const __m128i);
5831            let lo = _mm_and_si128(b, lomask);
5832            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
5833            let w = _mm256_sub_epi8(
5834                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5835                eight,
5836            );
5837            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5838            acc += dpbusd_hsum(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w)) as f32 * s;
5839        }
5840        acc
5841    }
5842}
5843
5844/// Exact scalar q4tp row — the `CMF_SDOT=0` contract, same pairwise
5845/// accumulation shape as `q4t_row_exact`.
5846#[inline]
5847fn q4tp_row_exact(nib: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
5848    #[cfg(target_arch = "x86_64")]
5849    if avx2_enabled() {
5850        // Keep the scalar pair/group reduction order, not a vector sum.
5851        return unsafe { q4tp_row_float_avx2(nib, r, gpr, x, scales) };
5852    }
5853    q4tp_row_float_scalar(nib, r, gpr, x, scales)
5854}
5855
5856#[inline]
5857fn q4tp_row_float_scalar(nib: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
5858    let mut acc = 0f32;
5859    for gi in 0..gpr {
5860        let tile = &nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
5861        let s = scales[gi];
5862        let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
5863        let mut ga = 0f32;
5864        for (k, &b) in tile.iter().enumerate() {
5865            ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
5866                + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
5867        }
5868        acc += ga * s;
5869    }
5870    acc
5871}
5872
5873/// Vectorize unpack, conversion and multiplication, but preserve every
5874/// pair addition and the scalar accumulation order. No activation rounding
5875/// or FMA: bit-identical to the float scalar row, including its group scale.
5876#[cfg(target_arch = "x86_64")]
5877#[target_feature(enable = "avx2")]
5878unsafe fn q4tp_row_float_avx2(nib: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
5879    // SAFETY: caller checks AVX2 and provides the same complete 32-element
5880    // groups as the scalar row. Loads/stores are explicitly unaligned.
5881    unsafe {
5882        use core::arch::x86_64::*;
5883        let mask = _mm_set1_epi8(15);
5884        let eight = _mm_set1_epi8(8);
5885        let order = _mm256_setr_epi32(0, 1, 4, 5, 2, 3, 6, 7);
5886        let mut acc = 0.0f32;
5887        for gi in 0..gpr {
5888            let packed = _mm_loadu_si128(nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB).cast());
5889            let lo = _mm_and_si128(packed, mask);
5890            let hi = _mm_and_si128(_mm_srli_epi16::<4>(packed), mask);
5891            let w0 = _mm_sub_epi8(_mm_unpacklo_epi8(lo, hi), eight);
5892            let w1 = _mm_sub_epi8(_mm_unpackhi_epi8(lo, hi), eight);
5893            let xp = x.as_ptr().add(gi * GROUP_SIZE);
5894            let a = _mm256_mul_ps(
5895                _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(w0)),
5896                _mm256_loadu_ps(xp),
5897            );
5898            let b = _mm256_mul_ps(
5899                _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(w0))),
5900                _mm256_loadu_ps(xp.add(8)),
5901            );
5902            let c = _mm256_mul_ps(
5903                _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(w1)),
5904                _mm256_loadu_ps(xp.add(16)),
5905            );
5906            let d = _mm256_mul_ps(
5907                _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(w1))),
5908                _mm256_loadu_ps(xp.add(24)),
5909            );
5910            let mut pairs = [0.0f32; 16];
5911            _mm256_storeu_ps(
5912                pairs.as_mut_ptr(),
5913                _mm256_permutevar8x32_ps(_mm256_hadd_ps(a, b), order),
5914            );
5915            _mm256_storeu_ps(
5916                pairs.as_mut_ptr().add(8),
5917                _mm256_permutevar8x32_ps(_mm256_hadd_ps(c, d), order),
5918            );
5919            let mut ga = 0.0f32;
5920            for v in pairs {
5921                ga += v;
5922            }
5923            acc += ga * scales[gi];
5924        }
5925        acc
5926    }
5927}
5928
5929/// Single weight of a q4tp tensor — the a8w8 outlier path, which restores
5930/// activation outliers at full precision after the int8 pass.
5931#[inline]
5932fn q4tp_outlier(nib: &[u8], r: usize, gpr: usize, j: usize, scales: &[f32]) -> (f32, f32) {
5933    let (gi, k) = (j / GROUP_SIZE, j % GROUP_SIZE);
5934    let byte = nib[(r * gpr + gi) * Q4TP_NIB + k / 2];
5935    let n = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
5936    ((n as i32 - 8) as f32, scales[gi])
5937}
5938
5939/// Fused q4tp matvec (dispatch mirrors `q4t_matvec`).
5940fn q4tp_matvec(
5941    bytes: &[u8],
5942    x: &[f32],
5943    rows: usize,
5944    cols: usize,
5945    out: &mut [f32],
5946    pool: Option<&Pool>,
5947) {
5948    debug_assert_eq!(out.len(), rows);
5949    let gpr = cols / GROUP_SIZE;
5950    let v = Q4tpView::new(bytes, rows, cols);
5951    let out_addr = SendMut(out.as_mut_ptr());
5952    if a8w8_enabled() {
5953        let act = split_act(x);
5954        let run = |start: usize, end: usize| {
5955            // One scratch row of scales per worker — borrowed, not minted.
5956            with_krow(gpr, |sc| {
5957                for r in start..end {
5958                    v.scales_into(r, gpr, sc);
5959                    let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, sc) * act.sx;
5960                    for &(j, xv) in &act.outliers {
5961                        let (w, s) = q4tp_outlier(v.nib, r, gpr, j, sc);
5962                        acc += w * s * xv;
5963                    }
5964                    // SAFETY: disjoint row ranges per worker.
5965                    unsafe { *out_addr.at(r) = acc };
5966                }
5967            })
5968        };
5969        dispatch_rows(pool, rows, &run);
5970        return;
5971    }
5972    let run = |start: usize, end: usize| {
5973        with_krow(gpr, |sc| {
5974            for r in start..end {
5975                v.scales_into(r, gpr, sc);
5976                // SAFETY: disjoint row ranges per worker.
5977                unsafe { *out_addr.at(r) = q4tp_row_exact(v.nib, r, gpr, x, sc) };
5978            }
5979        })
5980    };
5981    dispatch_rows(pool, rows, &run);
5982}
5983
5984/// Fused two-input q4tp matvec — the SwiGLU gate/up pair. Weights and the
5985/// row ladder are read once and spent on both activation streams.
5986#[allow(clippy::too_many_arguments)]
5987fn q4tp_matvec2(
5988    bytes: &[u8],
5989    x1: &[f32],
5990    x2: &[f32],
5991    rows: usize,
5992    cols: usize,
5993    o1: &mut [f32],
5994    o2: &mut [f32],
5995    pool: Option<&Pool>,
5996) {
5997    let gpr = cols / GROUP_SIZE;
5998    let v = Q4tpView::new(bytes, rows, cols);
5999    let (p1, p2) = (SendMut(o1.as_mut_ptr()), SendMut(o2.as_mut_ptr()));
6000    let run = |start: usize, end: usize| {
6001        let mut sc = vec![0f32; gpr];
6002        for r in start..end {
6003            v.scales_into(r, gpr, &mut sc);
6004            // SAFETY: disjoint row ranges per worker.
6005            unsafe {
6006                *p1.at(r) = q4tp_row_exact(v.nib, r, gpr, x1, &sc);
6007                *p2.at(r) = q4tp_row_exact(v.nib, r, gpr, x2, &sc);
6008            }
6009        }
6010    };
6011    dispatch_rows(pool, rows, &run);
6012}
6013
6014/// One q2tp outlier weight at column `j` of row `r`: the 2-bit code and
6015/// its group scale, mirrored on `q4tp_outlier`.
6016#[inline]
6017fn q2tp_outlier(chunks: &[u8], r: usize, gpr: usize, j: usize, scales: &[f32]) -> (f32, f32) {
6018    let (gi, k) = (j / GROUP_SIZE, j % GROUP_SIZE);
6019    let byte = chunks[(r * gpr + gi) * Q2TP_CHUNK + k / 4];
6020    let c = (byte >> (2 * (k % 4))) & 3;
6021    (c as f32 - 1.5, scales[gi])
6022}
6023
6024#[cfg(target_arch = "x86_64")]
6025const Q2TP_DECODE_U32: [u32; 256] = {
6026    let mut tab = [0u32; 256];
6027    let mut b = 0usize;
6028    while b < 256 {
6029        tab[b] = ((b as u32) & 3)
6030            | ((((b as u32) >> 2) & 3) << 8)
6031            | ((((b as u32) >> 4) & 3) << 16)
6032            | ((((b as u32) >> 6) & 3) << 24);
6033        b += 1;
6034    }
6035    tab
6036};
6037
6038/// Eight packed q2tp bytes against 32 signed activation bytes. `maddubs`
6039/// exactly computes unsigned 2-bit code × signed i8; its pair sums cannot
6040/// saturate (2 × 3 × 127 < i16::MAX), and the second madd widens to i32.
6041#[cfg(target_arch = "x86_64")]
6042#[target_feature(enable = "avx2")]
6043unsafe fn q2tp_code_dot_avx2(ch: &[u8], x: &[i8]) -> i32 {
6044    use core::arch::x86_64::*;
6045    debug_assert!(ch.len() >= Q2TP_CHUNK && x.len() >= GROUP_SIZE);
6046    let codes = _mm256_setr_epi32(
6047        Q2TP_DECODE_U32[ch[0] as usize] as i32,
6048        Q2TP_DECODE_U32[ch[1] as usize] as i32,
6049        Q2TP_DECODE_U32[ch[2] as usize] as i32,
6050        Q2TP_DECODE_U32[ch[3] as usize] as i32,
6051        Q2TP_DECODE_U32[ch[4] as usize] as i32,
6052        Q2TP_DECODE_U32[ch[5] as usize] as i32,
6053        Q2TP_DECODE_U32[ch[6] as usize] as i32,
6054        Q2TP_DECODE_U32[ch[7] as usize] as i32,
6055    );
6056    let xv = unsafe { _mm256_loadu_si256(x.as_ptr().cast()) };
6057    let pair = _mm256_maddubs_epi16(codes, xv);
6058    let quad = _mm256_madd_epi16(pair, _mm256_set1_epi16(1));
6059    let sum128 = _mm_add_epi32(
6060        _mm256_castsi256_si128(quad),
6061        _mm256_extracti128_si256(quad, 1),
6062    );
6063    let sum64 = _mm_hadd_epi32(sum128, sum128);
6064    _mm_cvtsi128_si32(_mm_hadd_epi32(sum64, sum64))
6065}
6066
6067#[cfg(target_arch = "x86_64")]
6068#[inline]
6069fn q2tp_avx2_enabled() -> bool {
6070    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6071    *ON.get_or_init(|| std::arch::is_x86_feature_detected!("avx2"))
6072}
6073
6074/// Integer dot of one q2tp row against pre-quantized activations:
6075/// Σ_g s_g · (Σ c·xq − 1.5·Σ xq). The half-integer grid (c − 1.5)
6076/// becomes exact integer math through the group sums — the same trick
6077/// every a8w8 kernel in this file rides. The codes decode into a
6078/// 32-byte scratch in natural order and the dot itself is the shared
6079/// SDOT primitive; elsewhere a scalar integer loop.
6080#[inline]
6081fn dot_q2tp_row_i8(
6082    chunks: &[u8],
6083    r: usize,
6084    gpr: usize,
6085    xq: &[i8],
6086    gsum: &[i32],
6087    scales: &[f32],
6088) -> f32 {
6089    let mut acc = 0f32;
6090    let base = r * gpr * Q2TP_CHUNK;
6091    #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
6092    let mut codes = [0i8; GROUP_SIZE];
6093    #[cfg(target_arch = "x86_64")]
6094    // Cache the process-wide AVX2 decision; probing it for every 32-weight
6095    // group adds a branch to the hottest Q2TP row loop while retaining the
6096    // established table-load AVX2 arithmetic (and without coupling this path
6097    // to the separate FMA-gated A8W8 switch).
6098    let avx2 = q2tp_avx2_enabled();
6099    for gi in 0..gpr {
6100        let ch = &chunks[base + gi * Q2TP_CHUNK..base + (gi + 1) * Q2TP_CHUNK];
6101        let xg = &xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
6102        #[cfg(target_arch = "aarch64")]
6103        // NEON: the byte's four 2-bit fields land in four lane vectors
6104        // (shift+mask), vld4 de-interleaves xq to match (xj[k] =
6105        // xq[4k+j]), widening MACs accumulate exactly in i32. A scalar
6106        // decode here cost as much as the dot it fed — the profile put
6107        // it at the top of the whole W2 decode.
6108        let dot = unsafe {
6109            use core::arch::aarch64::*;
6110            let b = vld1_u8(ch.as_ptr());
6111            let three = vdup_n_u8(3);
6112            let c0 = vreinterpret_s8_u8(vand_u8(b, three));
6113            let c1 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 2), three));
6114            let c2 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 4), three));
6115            let c3 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 6), three));
6116            let x4 = vld4_s8(xg.as_ptr());
6117            let mut acc4 = vdupq_n_s32(0);
6118            acc4 = vpadalq_s16(acc4, vmull_s8(c0, x4.0));
6119            acc4 = vpadalq_s16(acc4, vmull_s8(c1, x4.1));
6120            acc4 = vpadalq_s16(acc4, vmull_s8(c2, x4.2));
6121            acc4 = vpadalq_s16(acc4, vmull_s8(c3, x4.3));
6122            vaddvq_s32(acc4)
6123        };
6124        #[cfg(target_arch = "x86_64")]
6125        let dot: i32 = if avx2 {
6126            // SAFETY: the runtime feature check gates the target-feature body;
6127            // the group slices above are exactly 8 and 32 bytes long.
6128            unsafe { q2tp_code_dot_avx2(ch, xg) }
6129        } else {
6130            ch.iter()
6131                .enumerate()
6132                .map(|(k, &b)| {
6133                    ((b & 3) as i32) * xg[k * 4] as i32
6134                        + (((b >> 2) & 3) as i32) * xg[k * 4 + 1] as i32
6135                        + (((b >> 4) & 3) as i32) * xg[k * 4 + 2] as i32
6136                        + (((b >> 6) & 3) as i32) * xg[k * 4 + 3] as i32
6137                })
6138                .sum()
6139        };
6140        #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
6141        let dot: i32 = {
6142            for (k, &b) in ch.iter().enumerate() {
6143                codes[k * 4] = (b & 3) as i8;
6144                codes[k * 4 + 1] = ((b >> 2) & 3) as i8;
6145                codes[k * 4 + 2] = ((b >> 4) & 3) as i8;
6146                codes[k * 4 + 3] = ((b >> 6) & 3) as i8;
6147            }
6148            codes
6149                .iter()
6150                .zip(xg)
6151                .map(|(&c, &x)| c as i32 * x as i32)
6152                .sum()
6153        };
6154        acc += scales[gi] * (dot as f32 - 1.5 * gsum[gi] as f32);
6155    }
6156    acc
6157}
6158
6159/// Exact f32 dot of one q2tp row: 2-bit fields LSB-first, (c − 1.5)·s.
6160/// Scalar on purpose — the 2-bit class targets the GPU graph; the CPU
6161/// path exists for parity gates and small-machine fallback.
6162fn q2tp_row_exact(chunks: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
6163    q2tp_row_exact_center(chunks, r, gpr, x, scales, 1.5)
6164}
6165
6166/// Fused Prism affine row: the derived correction is applied inside the
6167/// decoded code, avoiding a second accumulated dot and avoiding cancellation
6168/// between `B=(c-1.5)s` and `+.5s` for long 5120/17408 rows.
6169#[inline]
6170fn q2tp_affine_row_exact(chunks: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
6171    q2tp_row_exact_center(chunks, r, gpr, x, scales, 1.0)
6172}
6173
6174#[inline]
6175fn q2tp_row_exact_center(
6176    chunks: &[u8],
6177    r: usize,
6178    gpr: usize,
6179    x: &[f32],
6180    scales: &[f32],
6181    center: f32,
6182) -> f32 {
6183    let mut acc = 0f32;
6184    for gi in 0..gpr {
6185        let ch = &chunks[(r * gpr + gi) * Q2TP_CHUNK..(r * gpr + gi + 1) * Q2TP_CHUNK];
6186        let s = scales[gi];
6187        let xb = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
6188        let mut g = 0f32;
6189        for (k, &b) in ch.iter().enumerate() {
6190            g += ((b & 3) as f32 - center) * xb[k * 4]
6191                + (((b >> 2) & 3) as f32 - center) * xb[k * 4 + 1]
6192                + (((b >> 4) & 3) as f32 - center) * xb[k * 4 + 2]
6193                + (((b >> 6) & 3) as f32 - center) * xb[k * 4 + 3];
6194        }
6195        acc += s * g;
6196    }
6197    acc
6198}
6199
6200fn q2tp_matvec(
6201    bytes: &[u8],
6202    x: &[f32],
6203    rows: usize,
6204    cols: usize,
6205    out: &mut [f32],
6206    pool: Option<&Pool>,
6207) {
6208    q2tp_matvec_mode(bytes, x, rows, cols, out, pool, false);
6209}
6210
6211fn q2tp_affine_matvec(
6212    bytes: &[u8],
6213    x: &[f32],
6214    rows: usize,
6215    cols: usize,
6216    out: &mut [f32],
6217    pool: Option<&Pool>,
6218) {
6219    q2tp_matvec_mode(bytes, x, rows, cols, out, pool, true);
6220}
6221
6222fn q2tp_matvec_mode(
6223    bytes: &[u8],
6224    x: &[f32],
6225    rows: usize,
6226    cols: usize,
6227    out: &mut [f32],
6228    pool: Option<&Pool>,
6229    affine: bool,
6230) {
6231    debug_assert_eq!(out.len(), rows);
6232    let gpr = cols / GROUP_SIZE;
6233    let v = Q4tpView::new_q2(bytes, rows, cols);
6234    let out_addr = SendMut(out.as_mut_ptr());
6235    // a8w8 fast path (CMF_SDOT=0 keeps the exact scalar walk): integer
6236    // code dots + group sums, exact outlier correction — the same
6237    // contract as every sibling kernel; measured 2-bit rows were the
6238    // only scalar holdout in the family.
6239    if !affine && a8w8_enabled() {
6240        let act = split_act(x);
6241        let gsum = q1_group_sums(&act.xq, gpr);
6242        let (act, gsum) = (&act, &gsum);
6243        let run = move |start: usize, end: usize| {
6244            with_krow(gpr, |sc| {
6245                for r in start..end {
6246                    v.scales_into(r, gpr, sc);
6247                    let mut acc = dot_q2tp_row_i8(v.nib, r, gpr, &act.xq, gsum, sc) * act.sx;
6248                    for &(j, xv) in &act.outliers {
6249                        let (w, s) = q2tp_outlier(v.nib, r, gpr, j, sc);
6250                        acc += w * s * xv;
6251                    }
6252                    // SAFETY: disjoint row ranges per worker.
6253                    unsafe { *out_addr.at(r) = acc };
6254                }
6255            })
6256        };
6257        dispatch_rows(pool, rows, &run);
6258        return;
6259    }
6260    let run = |start: usize, end: usize| {
6261        with_krow(gpr, |sc| {
6262            for r in start..end {
6263                v.scales_into(r, gpr, sc);
6264                // SAFETY: disjoint row ranges per worker.
6265                unsafe {
6266                    *out_addr.at(r) = if affine {
6267                        q2tp_affine_row_exact(v.nib, r, gpr, x, sc)
6268                    } else {
6269                        q2tp_row_exact(v.nib, r, gpr, x, sc)
6270                    }
6271                };
6272            }
6273        })
6274    };
6275    dispatch_rows(pool, rows, &run);
6276}
6277
6278/// Fused two-input q2tp matvec — the SwiGLU gate/up pair.
6279#[allow(clippy::too_many_arguments)]
6280fn q2tp_matvec2(
6281    bytes: &[u8],
6282    x1: &[f32],
6283    x2: &[f32],
6284    rows: usize,
6285    cols: usize,
6286    o1: &mut [f32],
6287    o2: &mut [f32],
6288    pool: Option<&Pool>,
6289) {
6290    q2tp_matvec2_mode(bytes, x1, x2, rows, cols, o1, o2, pool, false);
6291}
6292
6293#[allow(clippy::too_many_arguments)]
6294fn q2tp_affine_matvec2(
6295    bytes: &[u8],
6296    x1: &[f32],
6297    x2: &[f32],
6298    rows: usize,
6299    cols: usize,
6300    o1: &mut [f32],
6301    o2: &mut [f32],
6302    pool: Option<&Pool>,
6303) {
6304    q2tp_matvec2_mode(bytes, x1, x2, rows, cols, o1, o2, pool, true);
6305}
6306
6307#[allow(clippy::too_many_arguments)]
6308fn q2tp_matvec2_mode(
6309    bytes: &[u8],
6310    x1: &[f32],
6311    x2: &[f32],
6312    rows: usize,
6313    cols: usize,
6314    o1: &mut [f32],
6315    o2: &mut [f32],
6316    pool: Option<&Pool>,
6317    affine: bool,
6318) {
6319    let gpr = cols / GROUP_SIZE;
6320    let v = Q4tpView::new_q2(bytes, rows, cols);
6321    let (p1, p2) = (SendMut(o1.as_mut_ptr()), SendMut(o2.as_mut_ptr()));
6322    let run = |start: usize, end: usize| {
6323        let mut sc = vec![0f32; gpr];
6324        for r in start..end {
6325            v.scales_into(r, gpr, &mut sc);
6326            // SAFETY: disjoint row ranges per worker.
6327            unsafe {
6328                *p1.at(r) = if affine {
6329                    q2tp_affine_row_exact(v.nib, r, gpr, x1, &sc)
6330                } else {
6331                    q2tp_row_exact(v.nib, r, gpr, x1, &sc)
6332                };
6333                *p2.at(r) = if affine {
6334                    q2tp_affine_row_exact(v.nib, r, gpr, x2, &sc)
6335                } else {
6336                    q2tp_row_exact(v.nib, r, gpr, x2, &sc)
6337                };
6338            }
6339        }
6340    };
6341    dispatch_rows(pool, rows, &run);
6342}
6343
6344/// Batched q2tp matmat: scalar row kernel over every batch column. CPU
6345/// prefill only — decode rides the graph, so plain and correct beats
6346/// clever here.
6347/// Test doors into the host 2-bit kernels: the stand's heap corruption
6348/// pointed at down-shaped tensors, and the private fns need a way to be
6349/// held to a reference without a model file around them.
6350pub fn q2tp_matvec_for_test(bytes: &[u8], x: &[f32], rows: usize, cols: usize, out: &mut [f32]) {
6351    // The facade IS the reference: encoder oracles hold requant output
6352    // to the exact scalar walk. The production dispatch may take the i8
6353    // fast path, whose error scale is the ACTIVATIONS' — a different
6354    // claim than the encoder correctness these tests pin.
6355    let gpr = cols / GROUP_SIZE;
6356    let v = Q4tpView::new_q2(bytes, rows, cols);
6357    with_krow(gpr, |sc| {
6358        for r in 0..rows {
6359            v.scales_into(r, gpr, sc);
6360            out[r] = q2tp_row_exact(v.nib, r, gpr, x, sc);
6361        }
6362    });
6363}
6364
6365/// Test door for the descriptor-specific fused affine decode.  Production
6366/// callers select this through a validated Prism header, never by dtype alone.
6367pub fn q2tp_affine_matvec_for_test(
6368    bytes: &[u8],
6369    x: &[f32],
6370    rows: usize,
6371    cols: usize,
6372    out: &mut [f32],
6373) {
6374    q2tp_affine_matvec(bytes, x, rows, cols, out, None);
6375}
6376
6377pub fn q2tp_matmat_for_test(
6378    bytes: &[u8],
6379    xs_all: &[f32],
6380    b: usize,
6381    rows: usize,
6382    cols: usize,
6383    out: &mut [f32],
6384) {
6385    q2tp_matmat(bytes, xs_all, b, rows, cols, out, None);
6386}
6387
6388fn q2tp_matmat(
6389    bytes: &[u8],
6390    xs_all: &[f32],
6391    b: usize,
6392    rows: usize,
6393    cols: usize,
6394    out: &mut [f32],
6395    pool: Option<&Pool>,
6396) {
6397    q2tp_matmat_mode(bytes, xs_all, b, rows, cols, out, pool, false);
6398}
6399
6400fn q2tp_affine_matmat(
6401    bytes: &[u8],
6402    xs_all: &[f32],
6403    b: usize,
6404    rows: usize,
6405    cols: usize,
6406    out: &mut [f32],
6407    pool: Option<&Pool>,
6408) {
6409    q2tp_matmat_mode(bytes, xs_all, b, rows, cols, out, pool, true);
6410}
6411
6412fn q2tp_matmat_mode(
6413    bytes: &[u8],
6414    xs_all: &[f32],
6415    b: usize,
6416    rows: usize,
6417    cols: usize,
6418    out: &mut [f32],
6419    pool: Option<&Pool>,
6420    affine: bool,
6421) {
6422    debug_assert_eq!(out.len(), b * rows);
6423    let gpr = cols / GROUP_SIZE;
6424    let v = Q4tpView::new_q2(bytes, rows, cols);
6425    let out_addr = SendMut(out.as_mut_ptr());
6426    let run = |start: usize, end: usize| {
6427        let mut sc = vec![0f32; gpr];
6428        for r in start..end {
6429            v.scales_into(r, gpr, &mut sc);
6430            for bi in 0..b {
6431                let x = &xs_all[bi * cols..(bi + 1) * cols];
6432                // SAFETY: disjoint row ranges per worker.
6433                unsafe {
6434                    *out_addr.at(bi * rows + r) = if affine {
6435                        q2tp_affine_row_exact(v.nib, r, gpr, x, &sc)
6436                    } else {
6437                        q2tp_row_exact(v.nib, r, gpr, x, &sc)
6438                    }
6439                };
6440            }
6441        }
6442    };
6443    dispatch_rows(pool, rows, &run);
6444}
6445
6446/// The pre-vectorised shape, kept for A/B (`CMF_Q4TP_V1=1`): the
6447/// horizontal add lands once per group per column instead of once per
6448/// row. Same weights, same activations — only the reduction differs.
6449/// It is also the row-exact batch kernel: per group and column it forms
6450/// `int dot as f32 * scale` and adds it to a scalar running sum, which is
6451/// `dot_q4tp_row_sdot` step for step (Rust never contracts to an fma),
6452/// so each column is bit-identical to that column's matvec.
6453#[cfg(target_arch = "aarch64")]
6454#[target_feature(enable = "neon,dotprod")]
6455unsafe fn dot_q4tp_row_1x4_sdot_v1(
6456    nib: &[u8],
6457    r: usize,
6458    gpr: usize,
6459    xs: [&[i8]; 4],
6460    scales: &[f32],
6461) -> [f32; 4] {
6462    unsafe {
6463        use core::arch::aarch64::*;
6464        use core::arch::asm;
6465        let lomask = vdupq_n_u8(0x0F);
6466        let eight = vdupq_n_s8(8);
6467        let (mut f0, mut f1, mut f2, mut f3) = (0f32, 0f32, 0f32, 0f32);
6468        for gi in 0..gpr {
6469            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
6470            let s = *scales.get_unchecked(gi);
6471            let bb = vld1q_u8(t);
6472            let lo = vandq_u8(bb, lomask);
6473            let hi = vshrq_n_u8::<4>(bb);
6474            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
6475            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
6476            let mut d = [0f32; 4];
6477            for (k, dk) in d.iter_mut().enumerate() {
6478                let x0 = vld1q_s8(xs[k].as_ptr().add(gi * GROUP_SIZE));
6479                let x1 = vld1q_s8(xs[k].as_ptr().add(gi * GROUP_SIZE + 16));
6480                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6481                asm!(
6482                    "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
6483                    "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
6484                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6485                    e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
6486                    options(pure, nomem, nostack),
6487                );
6488                *dk = vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
6489            }
6490            f0 += d[0];
6491            f1 += d[1];
6492            f2 += d[2];
6493            f3 += d[3];
6494        }
6495        [f0, f1, f2, f3]
6496    }
6497}
6498
6499/// Which q4tp batch kernel to run: 1 = the previous one, 2 = the tuned
6500/// one, 0 = decide from the CPU. An atomic rather than a `OnceLock` so a
6501/// benchmark can alternate the two inside one process, where the machine's
6502/// mood — a shared box drifts ±25% between runs — is the same for both.
6503/// What the two mean is per-architecture: on x86 the blocked AVX-512 path
6504/// against the per-column one, on ARM the two reduction shapes.
6505#[allow(dead_code)]
6506static Q4TP_ALT: std::sync::atomic::AtomicU8 = std::sync::atomic::AtomicU8::new(0);
6507
6508/// Tests that store `Q4TP_ALT` hold this, so one test's kernel pick does
6509/// not leak into another's bit-exact comparison running in parallel.
6510#[cfg(test)]
6511static Q4TP_ALT_TEST_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
6512
6513/// Blocking pays on x86 only with 512-bit VNNI. With AVX2 alone, four
6514/// columns sharing an unpack still measured slower than the per-column
6515/// path (23.2 ms against 19.4 on a 48-thread EPYC), because that path
6516/// already dequantizes the row once — so the blocked kernel bought a
6517/// second unpack-free pass at the price of half the vector width.
6518#[cfg(target_arch = "x86_64")]
6519fn q4tp_blocked_x86() -> bool {
6520    match Q4TP_ALT.load(std::sync::atomic::Ordering::Relaxed) {
6521        1 => false,
6522        // A forced ON still asks the CPU. The switch exists so a bench can
6523        // pick a kernel, not so it can promise instructions the machine
6524        // does not have — CI caught that as a SIGILL on a runner without
6525        // AVX-512, where the parity test had turned the path on by hand.
6526        2 => avx512vnni_enabled(),
6527        // Deliberately not cached back into the switch: both gates below
6528        // hold their own `OnceLock`, and latching their answer here would
6529        // make a test's override outlive the test that set it.
6530        _ => blocked_enabled() && avx512vnni_enabled(),
6531    }
6532}
6533
6534/// `CMF_Q4TP_V1=1` picks the old reduction shape (A/B only).
6535#[cfg(target_arch = "aarch64")]
6536#[allow(dead_code)]
6537fn q4tp_v1() -> bool {
6538    match Q4TP_ALT.load(std::sync::atomic::Ordering::Relaxed) {
6539        1 => true,
6540        2 => false,
6541        _ => {
6542            static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6543            *ON.get_or_init(|| std::env::var("CMF_Q4TP_V1").is_ok_and(|v| v != "0"))
6544        }
6545    }
6546}
6547
6548/// Two weight rows against eight columns. The activation load is the
6549/// same for both rows, so it is paid once for twice the arithmetic, and
6550/// sixteen accumulator chains run where eight did — which is what a kernel
6551/// retiring 0.29 instructions a cycle is short of. Register pressure is
6552/// the limit: sixteen `zmm` accumulators, two weight tiles, one
6553/// activation, of thirty-two.
6554///
6555/// Four rows by four columns spends the same sixteen accumulators the
6556/// other way and measured worse — 1488 GFLOP/s against 1644 — so the
6557/// unpack, which four rows pay twice as often, costs more than the extra
6558/// sharing of one activation load buys.
6559#[cfg(target_arch = "x86_64")]
6560#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
6561unsafe fn dot_q4tp_2x8_avx512(
6562    nib: &[u8],
6563    r0: usize,
6564    gpr: usize,
6565    xs: [&[i8]; 8],
6566    sc0: &[f32],
6567    sc1: &[f32],
6568) -> [[f32; 8]; 2] {
6569    // SAFETY: as dot_q4tp_row_1x8_avx512, two adjacent rows at once; the
6570    // caller guarantees r0 + 1 < rows and the ISA.
6571    unsafe {
6572        use core::arch::x86_64::*;
6573        let lomask = _mm256_set1_epi8(0x0F);
6574        let eight = _mm256_set1_epi8(8);
6575        let zero = _mm512_setzero_si512();
6576        let mut v0 = [_mm512_setzero_ps(); 8];
6577        let mut v1 = [_mm512_setzero_ps(); 8];
6578        let pairs = gpr / 2;
6579        let unpack = |r: usize, gi: usize| -> (__m512i, __mmask64) {
6580            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
6581            let bb = _mm256_loadu_si256(t as *const __m256i);
6582            let lo = _mm256_and_si256(bb, lomask);
6583            let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
6584            let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
6585            let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
6586            let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
6587            let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
6588            (_mm512_abs_epi8(w), _mm512_movepi8_mask(w))
6589        };
6590        for gp in 0..pairs {
6591            let gi = gp * 2;
6592            let (wa0, neg0) = unpack(r0, gi);
6593            let (wa1, neg1) = unpack(r0 + 1, gi);
6594            let off = gi * GROUP_SIZE;
6595            let sv = |sc: &[f32]| {
6596                _mm512_insertf32x8::<1>(
6597                    _mm512_castps256_ps512(_mm256_set1_ps(*sc.get_unchecked(gi))),
6598                    _mm256_set1_ps(*sc.get_unchecked(gi + 1)),
6599                )
6600            };
6601            let s0 = sv(sc0);
6602            let s1 = sv(sc1);
6603            for k in 0..8 {
6604                let xv = _mm512_loadu_si512(xs[k].as_ptr().add(off) as *const __m512i);
6605                let d0 = _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(
6606                    zero,
6607                    wa0,
6608                    _mm512_mask_sub_epi8(xv, neg0, zero, xv),
6609                ));
6610                let d1 = _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(
6611                    zero,
6612                    wa1,
6613                    _mm512_mask_sub_epi8(xv, neg1, zero, xv),
6614                ));
6615                v0[k] = _mm512_fmadd_ps(d0, s0, v0[k]);
6616                v1[k] = _mm512_fmadd_ps(d1, s1, v1[k]);
6617            }
6618        }
6619        let mut acc = [[0f32; 8]; 2];
6620        for k in 0..8 {
6621            acc[0][k] = _mm512_reduce_add_ps(v0[k]);
6622            acc[1][k] = _mm512_reduce_add_ps(v1[k]);
6623        }
6624        if gpr % 2 == 1 {
6625            let off = (gpr - 1) * GROUP_SIZE;
6626            for j in off..off + GROUP_SIZE {
6627                let (w0, sa) = q4tp_outlier(nib, r0, gpr, j, sc0);
6628                let (w1, sb) = q4tp_outlier(nib, r0 + 1, gpr, j, sc1);
6629                for k in 0..8 {
6630                    let x = *xs[k].get_unchecked(j) as f32;
6631                    acc[0][k] += w0 * sa * x;
6632                    acc[1][k] += w1 * sb * x;
6633                }
6634            }
6635        }
6636        acc
6637    }
6638}
6639
6640/// The same, eight columns at a time. One unpack then feeds twice as many
6641/// activation streams, so a wide batch reads the weight tile half as
6642/// often; the price is eight accumulators live at once. Measured 9.0 ->
6643/// 8.3 ms at 9216x2304, b=296 on a 48-thread EPYC 9B45.
6644#[cfg(target_arch = "x86_64")]
6645#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
6646unsafe fn dot_q4tp_row_1x8_avx512(
6647    nib: &[u8],
6648    r: usize,
6649    gpr: usize,
6650    xs: [&[i8]; 8],
6651    scales: &[f32],
6652) -> [f32; 8] {
6653    // SAFETY: as dot_q4tp_row_1x4_avx2; caller guarantees the ISA.
6654    unsafe {
6655        use core::arch::x86_64::*;
6656        let lomask = _mm256_set1_epi8(0x0F);
6657        let eight = _mm256_set1_epi8(8);
6658        let zero = _mm512_setzero_si512();
6659        let (mut v0, mut v1, mut v2, mut v3) = (
6660            _mm512_setzero_ps(),
6661            _mm512_setzero_ps(),
6662            _mm512_setzero_ps(),
6663            _mm512_setzero_ps(),
6664        );
6665        let (mut v4, mut v5, mut v6, mut v7) = (
6666            _mm512_setzero_ps(),
6667            _mm512_setzero_ps(),
6668            _mm512_setzero_ps(),
6669            _mm512_setzero_ps(),
6670        );
6671        let pairs = gpr / 2;
6672        for gp in 0..pairs {
6673            let gi = gp * 2;
6674            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
6675            let bb = _mm256_loadu_si256(t as *const __m256i);
6676            let lo = _mm256_and_si256(bb, lomask);
6677            let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
6678            // `unpack` works per 128-bit lane, so the halves come out as
6679            // [A.lo, B.lo] and [A.hi, B.hi]; the shuffle reorders the four
6680            // 128-bit lanes into the weights' natural order, which is what
6681            // the straight activation load expects.
6682            let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
6683            let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
6684            let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
6685            let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
6686            let wabs = _mm512_abs_epi8(w);
6687            let neg = _mm512_movepi8_mask(w);
6688            let off = gi * GROUP_SIZE;
6689            let sv = _mm512_insertf32x8::<1>(
6690                _mm512_castps256_ps512(_mm256_set1_ps(*scales.get_unchecked(gi))),
6691                _mm256_set1_ps(*scales.get_unchecked(gi + 1)),
6692            );
6693            let dot = |x: &[i8]| -> __m512 {
6694                let xv = _mm512_loadu_si512(x.as_ptr().add(off) as *const __m512i);
6695                let sx = _mm512_mask_sub_epi8(xv, neg, zero, xv);
6696                _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(zero, wabs, sx))
6697            };
6698            v0 = _mm512_fmadd_ps(dot(xs[0]), sv, v0);
6699            v1 = _mm512_fmadd_ps(dot(xs[1]), sv, v1);
6700            v2 = _mm512_fmadd_ps(dot(xs[2]), sv, v2);
6701            v3 = _mm512_fmadd_ps(dot(xs[3]), sv, v3);
6702            v4 = _mm512_fmadd_ps(dot(xs[4]), sv, v4);
6703            v5 = _mm512_fmadd_ps(dot(xs[5]), sv, v5);
6704            v6 = _mm512_fmadd_ps(dot(xs[6]), sv, v6);
6705            v7 = _mm512_fmadd_ps(dot(xs[7]), sv, v7);
6706        }
6707        let mut acc = [
6708            _mm512_reduce_add_ps(v0),
6709            _mm512_reduce_add_ps(v1),
6710            _mm512_reduce_add_ps(v2),
6711            _mm512_reduce_add_ps(v3),
6712            _mm512_reduce_add_ps(v4),
6713            _mm512_reduce_add_ps(v5),
6714            _mm512_reduce_add_ps(v6),
6715            _mm512_reduce_add_ps(v7),
6716        ];
6717        // An odd group count leaves one group over; the narrow kernel
6718        // finishes it rather than the tail being a special case here.
6719        if gpr % 2 == 1 {
6720            let off = (gpr - 1) * GROUP_SIZE;
6721            for j in off..off + GROUP_SIZE {
6722                let (w, s) = q4tp_outlier(nib, r, gpr, j, scales);
6723                let ws = w * s;
6724                for k in 0..8 {
6725                    acc[k] += ws * *xs[k].get_unchecked(j) as f32;
6726                }
6727            }
6728        }
6729        acc
6730    }
6731}
6732
6733/// The same four columns, 512 bits wide. Two groups (64 weights) ride one
6734/// unpack and one `vpdpbusd`, where AVX2 needs two unpacks and four
6735/// `maddubs`/`madd` pairs — about 2.3x fewer instructions for the same
6736/// arithmetic. The two groups carry different scales, so the fma takes a
6737/// vector whose halves hold each group's scale rather than a broadcast.
6738///
6739/// There is no 512-bit `vpsignb`, so the activation's sign is applied by
6740/// negating under a mask taken from the weight's sign bits. That mask is
6741/// per-tile, so it is hoisted out of the column loop and the per-column
6742/// cost stays exactly one instruction, as with `sign_epi8`. Weights of
6743/// zero are not zeroed by the mask trick and do not need to be: their
6744/// magnitude is zero, so the product is.
6745#[cfg(target_arch = "x86_64")]
6746#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
6747unsafe fn dot_q4tp_row_1x4_avx512(
6748    nib: &[u8],
6749    r: usize,
6750    gpr: usize,
6751    xs: [&[i8]; 4],
6752    scales: &[f32],
6753) -> [f32; 4] {
6754    // SAFETY: as dot_q4tp_row_1x4_avx2; caller guarantees the ISA.
6755    unsafe {
6756        use core::arch::x86_64::*;
6757        let lomask = _mm256_set1_epi8(0x0F);
6758        let eight = _mm256_set1_epi8(8);
6759        let zero = _mm512_setzero_si512();
6760        let (mut v0, mut v1, mut v2, mut v3) = (
6761            _mm512_setzero_ps(),
6762            _mm512_setzero_ps(),
6763            _mm512_setzero_ps(),
6764            _mm512_setzero_ps(),
6765        );
6766        let pairs = gpr / 2;
6767        for gp in 0..pairs {
6768            let gi = gp * 2;
6769            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
6770            let bb = _mm256_loadu_si256(t as *const __m256i);
6771            let lo = _mm256_and_si256(bb, lomask);
6772            let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
6773            // `unpack` works per 128-bit lane, so the halves come out as
6774            // [A.lo, B.lo] and [A.hi, B.hi]; the shuffle reorders the four
6775            // 128-bit lanes into the weights' natural order, which is what
6776            // the straight activation load expects.
6777            let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
6778            let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
6779            let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
6780            let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
6781            let wabs = _mm512_abs_epi8(w);
6782            let neg = _mm512_movepi8_mask(w);
6783            let off = gi * GROUP_SIZE;
6784            let sv = _mm512_insertf32x8::<1>(
6785                _mm512_castps256_ps512(_mm256_set1_ps(*scales.get_unchecked(gi))),
6786                _mm256_set1_ps(*scales.get_unchecked(gi + 1)),
6787            );
6788            let dot = |x: &[i8]| -> __m512 {
6789                let xv = _mm512_loadu_si512(x.as_ptr().add(off) as *const __m512i);
6790                let sx = _mm512_mask_sub_epi8(xv, neg, zero, xv);
6791                _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(zero, wabs, sx))
6792            };
6793            v0 = _mm512_fmadd_ps(dot(xs[0]), sv, v0);
6794            v1 = _mm512_fmadd_ps(dot(xs[1]), sv, v1);
6795            v2 = _mm512_fmadd_ps(dot(xs[2]), sv, v2);
6796            v3 = _mm512_fmadd_ps(dot(xs[3]), sv, v3);
6797        }
6798        let mut acc = [
6799            _mm512_reduce_add_ps(v0),
6800            _mm512_reduce_add_ps(v1),
6801            _mm512_reduce_add_ps(v2),
6802            _mm512_reduce_add_ps(v3),
6803        ];
6804        // An odd group count leaves one group over; the narrow kernel
6805        // finishes it rather than the tail being a special case here.
6806        if gpr % 2 == 1 {
6807            let off = (gpr - 1) * GROUP_SIZE;
6808            for j in off..off + GROUP_SIZE {
6809                let (w, s) = q4tp_outlier(nib, r, gpr, j, scales);
6810                let ws = w * s;
6811                for k in 0..4 {
6812                    acc[k] += ws * *xs[k].get_unchecked(j) as f32;
6813                }
6814            }
6815        }
6816        acc
6817    }
6818}
6819
6820/// Four batch columns against one q4tp row: the tile is unpacked ONCE and
6821/// spent on four activation streams, which is where a prefill batch stops
6822/// being weight-bandwidth-bound. Twin of `dot_q4t_row_1x4_sdot`.
6823#[cfg(target_arch = "aarch64")]
6824#[target_feature(enable = "neon,dotprod")]
6825unsafe fn dot_q4tp_row_1x4_sdot(
6826    nib: &[u8],
6827    r: usize,
6828    gpr: usize,
6829    xs: [&[i8]; 4],
6830    scales: &[f32],
6831) -> [f32; 4] {
6832    // SAFETY: see dot_q4tp_row_sdot; every xs[k] is gpr·GROUP_SIZE long.
6833    unsafe {
6834        use core::arch::aarch64::*;
6835        use core::arch::asm;
6836        let lomask = vdupq_n_u8(0x0F);
6837        let eight = vdupq_n_s8(8);
6838        // Named accumulators, NOT an array indexed by a loop variable: the
6839        // latter does not stay in registers (the same defect cost 2x in the
6840        // AVX2 q4t kernel and again in WGSL).
6841        //
6842        // They are VECTORS, and the horizontal add happens once at the end
6843        // instead of once per group per column. `vaddvq` is a cross-lane
6844        // reduction — with 72 groups and four columns the old shape paid
6845        // 288 of them per row, each one a dependency stall the pipeline
6846        // cannot hide, to save four float adds. The group's scale now
6847        // rides an fma into the lane accumulators, so the arithmetic per
6848        // group is one convert and one fma. Summation order changes (the
6849        // lanes carry independent partial sums), which is the same
6850        // round-off class the SDOT path already lives in — the strict
6851        // kernel (`CMF_SDOT=0`, what `cortiq ppl` runs) is unchanged and
6852        // stays the reference.
6853        let (mut v0, mut v1, mut v2, mut v3) = (
6854            vdupq_n_f32(0.0),
6855            vdupq_n_f32(0.0),
6856            vdupq_n_f32(0.0),
6857            vdupq_n_f32(0.0),
6858        );
6859        for gi in 0..gpr {
6860            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
6861            let s = *scales.get_unchecked(gi);
6862            let bb = vld1q_u8(t);
6863            let lo = vandq_u8(bb, lomask);
6864            let hi = vshrq_n_u8::<4>(bb);
6865            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
6866            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
6867            let off = gi * GROUP_SIZE;
6868            let dot4 = |x: &[i8]| -> int32x4_t {
6869                let x0 = vld1q_s8(x.as_ptr().add(off));
6870                let x1 = vld1q_s8(x.as_ptr().add(off + 16));
6871                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6872                asm!(
6873                    "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
6874                    "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
6875                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6876                    e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
6877                    options(pure, nomem, nostack),
6878                );
6879                vaddq_s32(a0, a1)
6880            };
6881            v0 = vfmaq_n_f32(v0, vcvtq_f32_s32(dot4(xs[0])), s);
6882            v1 = vfmaq_n_f32(v1, vcvtq_f32_s32(dot4(xs[1])), s);
6883            v2 = vfmaq_n_f32(v2, vcvtq_f32_s32(dot4(xs[2])), s);
6884            v3 = vfmaq_n_f32(v3, vcvtq_f32_s32(dot4(xs[3])), s);
6885        }
6886        [
6887            vaddvq_f32(v0),
6888            vaddvq_f32(v1),
6889            vaddvq_f32(v2),
6890            vaddvq_f32(v3),
6891        ]
6892    }
6893}
6894
6895/// Fused q4tp matmat — the same three arms `q4t_matmat` has. Shipping only
6896/// the scalar one made Nanbeige-3B decode at 1.2 tok/s against q4t's 5.9:
6897/// the format was fine, the missing arms were the whole regression.
6898///
6899/// Under `row_exact()` every cell is summed in `q4tp_matvec`'s order on
6900/// every architecture; the mode is read once, so one call never mixes
6901/// the two contracts when a concurrent scope opens or closes mid-call.
6902fn q4tp_matmat(
6903    bytes: &[u8],
6904    xs_all: &[f32],
6905    b: usize,
6906    rows: usize,
6907    cols: usize,
6908    out: &mut [f32],
6909    pool: Option<&Pool>,
6910) {
6911    q4tp_matmat_with(bytes, xs_all, b, rows, cols, out, pool, row_exact())
6912}
6913
6914/// `q4tp_matmat` with the row-exact mode passed in rather than read from
6915/// the shared counter: `exact` = each cell equals its token's matvec bit
6916/// for bit, otherwise the fast blocked / AMX arms are free to reorder.
6917#[allow(clippy::too_many_arguments)]
6918fn q4tp_matmat_with(
6919    bytes: &[u8],
6920    xs_all: &[f32],
6921    b: usize,
6922    rows: usize,
6923    cols: usize,
6924    out: &mut [f32],
6925    pool: Option<&Pool>,
6926    exact: bool,
6927) {
6928    debug_assert_eq!(out.len(), b * rows);
6929    let gpr = cols / GROUP_SIZE;
6930    let v = Q4tpView::new(bytes, rows, cols);
6931
6932    // Wide batches ride the AMX through a dequant-tile sgemm, as in q4t.
6933    // An f32 GEMM over dequantized weights is not the matvec's int8 sum,
6934    // so a row-exact batch never takes it.
6935    #[cfg(target_os = "macos")]
6936    if !exact && b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
6937        dequant_matmat_accel(
6938            &|r, dst| {
6939                let mut sc = [0f32; 32];
6940                let mut scv;
6941                let s: &[f32] = if gpr <= 32 {
6942                    v.scales_into(r, gpr, &mut sc);
6943                    &sc[..gpr]
6944                } else {
6945                    scv = vec![0f32; gpr];
6946                    v.scales_into(r, gpr, &mut scv);
6947                    &scv
6948                };
6949                for gi in 0..gpr {
6950                    let tile = &v.nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
6951                    for (k, &bb) in tile.iter().enumerate() {
6952                        dst[gi * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s[gi];
6953                        dst[gi * GROUP_SIZE + k * 2 + 1] =
6954                            (((bb >> 4) & 0x0F) as f32 - 8.0) * s[gi];
6955                    }
6956                }
6957            },
6958            xs_all,
6959            b,
6960            rows,
6961            cols,
6962            out,
6963            pool,
6964        );
6965        return;
6966    }
6967
6968    let out_addr = SendMut(out.as_mut_ptr());
6969    if a8w8_enabled() {
6970        let acts: Vec<SplitAct> = (0..b)
6971            .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
6972            .collect();
6973        let acts = &acts;
6974        // ARM stays blocked under `exact` too: the 1x4 kernel then runs in
6975        // its v1 shape, whose per-group `int dot as f32 * scale` and scalar
6976        // running sum are exactly `dot_q4tp_row_sdot`'s, so the tile is
6977        // still unpacked once for four columns and each column equals its
6978        // matvec. The tuned shape (fma into lane partials, one horizontal
6979        // add per row) is 1-16 ulp off the matvec and is kept for `!exact`.
6980        #[cfg(target_arch = "aarch64")]
6981        let blocked_ok = sdot_enabled() && blocked_enabled();
6982        // x86 gets the same blocking: one tile unpack spent on four
6983        // columns. Without it every column re-decoded the row, which is
6984        // why a 48-core EPYC measured a sixth of an M4's per-core rate.
6985        // The gate is `avx2_enabled`, as in q4t — `sdot_enabled` answers
6986        // for ARM's dotprod and is hard-wired false everywhere else, so
6987        // asking it here left the whole blocked path unreachable on x86.
6988        #[cfg(target_arch = "x86_64")]
6989        let blocked_ok = q4tp_blocked_x86() && !exact;
6990        #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
6991        let blocked_ok = {
6992            let _ = exact;
6993            false
6994        };
6995        // Columns are swept in panels that fit L2. Without this a
6996        // row-pair walks every activation in the batch — 4.8 MB at
6997        // 512x512 — and does it again for the next pair, so the whole
6998        // batch streams out of the shared cache once per row. Measured
6999        // 800 GB/s of it, flat across batch sizes, which is the signature
7000        // of a loop bound by traffic rather than by arithmetic. A panel of
7001        // 256 columns is 590 KB beside 221 KB of this worker's weights:
7002        // both stay resident and the batch crosses L3 once instead of
7003        // once per row.
7004        let panel_cols: usize = std::env::var("CMF_Q4TP_PANEL")
7005            .ok()
7006            .and_then(|v| v.parse().ok())
7007            .filter(|v| *v > 0)
7008            .unwrap_or(256);
7009        let run = |start: usize, end: usize| {
7010            for abase in (0..acts.len()).step_by(panel_cols) {
7011                let alen = (acts.len() - abase).min(panel_cols);
7012                let mut sc = vec![0f32; gpr];
7013                #[cfg(target_arch = "x86_64")]
7014                let mut r_lo = start;
7015                #[cfg(target_arch = "x86_64")]
7016                if blocked_ok && alen >= 8 {
7017                    let mut sc1 = vec![0f32; gpr];
7018                    while r_lo + 2 <= end {
7019                        v.scales_into(r_lo, gpr, &mut sc);
7020                        v.scales_into(r_lo + 1, gpr, &mut sc1);
7021                        let mut bi = 0usize;
7022                        while bi + 8 <= alen {
7023                            let xs = [
7024                                acts[abase + bi].xq.as_slice(),
7025                                acts[abase + bi + 1].xq.as_slice(),
7026                                acts[abase + bi + 2].xq.as_slice(),
7027                                acts[abase + bi + 3].xq.as_slice(),
7028                                acts[abase + bi + 4].xq.as_slice(),
7029                                acts[abase + bi + 5].xq.as_slice(),
7030                                acts[abase + bi + 6].xq.as_slice(),
7031                                acts[abase + bi + 7].xq.as_slice(),
7032                            ];
7033                            let d = unsafe { dot_q4tp_2x8_avx512(v.nib, r_lo, gpr, xs, &sc, &sc1) };
7034                            for (row, dr, scr) in [(r_lo, &d[0], &sc), (r_lo + 1, &d[1], &sc1)] {
7035                                for k in 0..8 {
7036                                    let act = &acts[abase + bi + k];
7037                                    let mut acc = dr[k] * act.sx;
7038                                    for &(j, xv) in &act.outliers {
7039                                        let (w, s) = q4tp_outlier(v.nib, row, gpr, j, scr);
7040                                        acc += w * s * xv;
7041                                    }
7042                                    // SAFETY: disjoint (bi, r) cells per worker.
7043                                    unsafe { *out_addr.at((abase + bi + k) * rows + row) = acc };
7044                                }
7045                            }
7046                            bi += 8;
7047                        }
7048                        // Columns past the last group of eight, both rows —
7049                        // the same single-row kernel the tail below uses.
7050                        for row in [r_lo, r_lo + 1] {
7051                            let scr: &[f32] = if row == r_lo { &sc } else { &sc1 };
7052                            for b2 in bi..alen {
7053                                let act = &acts[abase + b2];
7054                                let xs4 = [
7055                                    act.xq.as_slice(),
7056                                    act.xq.as_slice(),
7057                                    act.xq.as_slice(),
7058                                    act.xq.as_slice(),
7059                                ];
7060                                let d =
7061                                    unsafe { dot_q4tp_row_1x4_avx512(v.nib, row, gpr, xs4, scr) };
7062                                let mut acc = d[0] * act.sx;
7063                                for &(j, xv) in &act.outliers {
7064                                    let (w, s) = q4tp_outlier(v.nib, row, gpr, j, scr);
7065                                    acc += w * s * xv;
7066                                }
7067                                // SAFETY: disjoint (bi, r) cells per worker.
7068                                unsafe { *out_addr.at((abase + b2) * rows + row) = acc };
7069                            }
7070                        }
7071                        r_lo += 2;
7072                    }
7073                }
7074                #[cfg(target_arch = "x86_64")]
7075                let row_start = r_lo;
7076                #[cfg(not(target_arch = "x86_64"))]
7077                let row_start = start;
7078                for r in row_start..end {
7079                    v.scales_into(r, gpr, &mut sc);
7080                    let mut bi = 0usize;
7081                    #[cfg(target_arch = "x86_64")]
7082                    if blocked_ok {
7083                        while bi + 8 <= alen {
7084                            let xs = [
7085                                acts[abase + bi].xq.as_slice(),
7086                                acts[abase + bi + 1].xq.as_slice(),
7087                                acts[abase + bi + 2].xq.as_slice(),
7088                                acts[abase + bi + 3].xq.as_slice(),
7089                                acts[abase + bi + 4].xq.as_slice(),
7090                                acts[abase + bi + 5].xq.as_slice(),
7091                                acts[abase + bi + 6].xq.as_slice(),
7092                                acts[abase + bi + 7].xq.as_slice(),
7093                            ];
7094                            let d = unsafe { dot_q4tp_row_1x8_avx512(v.nib, r, gpr, xs, &sc) };
7095                            for k in 0..8 {
7096                                let act = &acts[abase + bi + k];
7097                                let mut acc = d[k] * act.sx;
7098                                for &(j, xv) in &act.outliers {
7099                                    let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
7100                                    acc += w * s * xv;
7101                                }
7102                                // SAFETY: disjoint (bi, r) cells per worker.
7103                                unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
7104                            }
7105                            bi += 8;
7106                        }
7107                        while bi + 4 <= alen {
7108                            let xs = [
7109                                acts[abase + bi].xq.as_slice(),
7110                                acts[abase + bi + 1].xq.as_slice(),
7111                                acts[abase + bi + 2].xq.as_slice(),
7112                                acts[abase + bi + 3].xq.as_slice(),
7113                            ];
7114                            let d = unsafe { dot_q4tp_row_1x4_avx512(v.nib, r, gpr, xs, &sc) };
7115                            for k in 0..4 {
7116                                let act = &acts[abase + bi + k];
7117                                let mut acc = d[k] * act.sx;
7118                                for &(j, xv) in &act.outliers {
7119                                    let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
7120                                    acc += w * s * xv;
7121                                }
7122                                // SAFETY: disjoint (bi, r) cells per worker.
7123                                unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
7124                            }
7125                            bi += 4;
7126                        }
7127                    }
7128                    #[cfg(target_arch = "aarch64")]
7129                    if blocked_ok {
7130                        while bi + 4 <= alen {
7131                            let xs = [
7132                                acts[abase + bi].xq.as_slice(),
7133                                acts[abase + bi + 1].xq.as_slice(),
7134                                acts[abase + bi + 2].xq.as_slice(),
7135                                acts[abase + bi + 3].xq.as_slice(),
7136                            ];
7137                            let d = unsafe {
7138                                if exact || q4tp_v1() {
7139                                    dot_q4tp_row_1x4_sdot_v1(v.nib, r, gpr, xs, &sc)
7140                                } else {
7141                                    dot_q4tp_row_1x4_sdot(v.nib, r, gpr, xs, &sc)
7142                                }
7143                            };
7144                            for k in 0..4 {
7145                                let act = &acts[abase + bi + k];
7146                                let mut acc = d[k] * act.sx;
7147                                for &(j, xv) in &act.outliers {
7148                                    let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
7149                                    acc += w * s * xv;
7150                                }
7151                                // SAFETY: disjoint (bi, r) cells per worker.
7152                                unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
7153                            }
7154                            bi += 4;
7155                        }
7156                    }
7157                    let _ = blocked_ok;
7158                    while bi < alen {
7159                        let act = &acts[abase + bi];
7160                        let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, &sc) * act.sx;
7161                        for &(j, xv) in &act.outliers {
7162                            let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
7163                            acc += w * s * xv;
7164                        }
7165                        // SAFETY: disjoint (bi, r) cells per worker range.
7166                        unsafe { *out_addr.at((abase + bi) * rows + r) = acc };
7167                        bi += 1;
7168                    }
7169                }
7170            }
7171        };
7172        dispatch_rows(pool, rows, &run);
7173        return;
7174    }
7175
7176    let run = |start: usize, end: usize| {
7177        let mut sc = vec![0f32; gpr];
7178        for r in start..end {
7179            v.scales_into(r, gpr, &mut sc);
7180            for bi in 0..b {
7181                let x = &xs_all[bi * cols..(bi + 1) * cols];
7182                // SAFETY: disjoint (bi, r) cells per worker range.
7183                unsafe { *out_addr.at(bi * rows + r) = q4tp_row_exact(v.nib, r, gpr, x, &sc) };
7184            }
7185        }
7186    };
7187    dispatch_rows(pool, rows, &run);
7188}
7189
7190/// Fused q4_tiled matvec (dispatch mirrors `q4matvec`).
7191fn q4t_matvec(
7192    bytes: &[u8],
7193    x: &[f32],
7194    rows: usize,
7195    cols: usize,
7196    out: &mut [f32],
7197    pool: Option<&Pool>,
7198) {
7199    debug_assert_eq!(out.len(), rows);
7200    let gpr = cols / GROUP_SIZE;
7201    let out_addr = SendMut(out.as_mut_ptr());
7202    if a8w8_enabled() {
7203        let act = split_act(x);
7204        let run = move |start: usize, end: usize| {
7205            for r in start..end {
7206                let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
7207                for &(j, xv) in &act.outliers {
7208                    let (w, s) = q4t_outlier(bytes, r, gpr, j);
7209                    acc += w * s * xv;
7210                }
7211                // SAFETY: disjoint row ranges per worker.
7212                unsafe { *out_addr.at(r) = acc };
7213            }
7214        };
7215        dispatch_rows(pool, rows, &run);
7216        return;
7217    }
7218    let run = move |start: usize, end: usize| {
7219        for r in start..end {
7220            // SAFETY: disjoint row ranges per worker.
7221            unsafe { *out_addr.at(r) = q4t_row_exact(bytes, r, gpr, x) };
7222        }
7223    };
7224    dispatch_rows(pool, rows, &run);
7225}
7226
7227/// Fused two-input q4_tiled matvec (weights read once per pair).
7228#[allow(clippy::too_many_arguments)]
7229fn q4t_matvec2(
7230    bytes: &[u8],
7231    x1: &[f32],
7232    x2: &[f32],
7233    rows: usize,
7234    cols: usize,
7235    o1: &mut [f32],
7236    o2: &mut [f32],
7237    pool: Option<&Pool>,
7238) {
7239    let gpr = cols / GROUP_SIZE;
7240    let p1 = SendMut(o1.as_mut_ptr());
7241    let p2 = SendMut(o2.as_mut_ptr());
7242    if a8w8_enabled() {
7243        let a1 = split_act(x1);
7244        let a2 = split_act(x2);
7245        let run = move |start: usize, end: usize| {
7246            for r in start..end {
7247                let mut v1 = dot_q4t_row_i8(bytes, r, gpr, &a1.xq) * a1.sx;
7248                let mut v2 = dot_q4t_row_i8(bytes, r, gpr, &a2.xq) * a2.sx;
7249                for &(j, xv) in &a1.outliers {
7250                    let (w, s) = q4t_outlier(bytes, r, gpr, j);
7251                    v1 += w * s * xv;
7252                }
7253                for &(j, xv) in &a2.outliers {
7254                    let (w, s) = q4t_outlier(bytes, r, gpr, j);
7255                    v2 += w * s * xv;
7256                }
7257                // SAFETY: disjoint row ranges per worker.
7258                unsafe {
7259                    *p1.at(r) = v1;
7260                    *p2.at(r) = v2;
7261                }
7262            }
7263        };
7264        dispatch_rows(pool, rows, &run);
7265        return;
7266    }
7267    let run = move |start: usize, end: usize| {
7268        for r in start..end {
7269            // SAFETY: disjoint row ranges per worker.
7270            unsafe {
7271                *p1.at(r) = q4t_row_exact(bytes, r, gpr, x1);
7272                *p2.at(r) = q4t_row_exact(bytes, r, gpr, x2);
7273            }
7274        }
7275    };
7276    dispatch_rows(pool, rows, &run);
7277}
7278
7279/// Batched q4_tiled matmat: each row's tiles stream once per microbatch.
7280#[allow(clippy::too_many_arguments)]
7281/// Prefill GEMM through Accelerate for group-quantized codecs: a
7282/// caller-supplied row dequantizer fills f32 tiles (pool-parallel) and
7283/// each tile rides the AMX with one sgemm — the generic sibling of
7284/// `qmatmat_accel` (q8). Numerics are f32-GEMM (tolerance class);
7285/// decode (b=1) never takes this path.
7286#[cfg(target_os = "macos")]
7287fn dequant_matmat_accel(
7288    dequant_row: &(dyn Fn(usize, &mut [f32]) + Sync),
7289    xs_all: &[f32],
7290    b: usize,
7291    rows: usize,
7292    cols: usize,
7293    out: &mut [f32],
7294    pool: Option<&Pool>,
7295) {
7296    const TR: usize = 2048;
7297    thread_local! {
7298        static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
7299    }
7300    WTILE.with(|wt| {
7301        let mut wtile = wt.borrow_mut();
7302        wtile.resize(TR * cols, 0.0);
7303        let mut r0 = 0usize;
7304        while r0 < rows {
7305            let tr = TR.min(rows - r0);
7306            let wt_addr = SendMut(wtile.as_mut_ptr());
7307            let run = |start: usize, end: usize| {
7308                for r in start..end {
7309                    // SAFETY: workers cover disjoint r ranges.
7310                    let dst = unsafe { std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols) };
7311                    dequant_row(r0 + r, dst);
7312                }
7313            };
7314            dispatch_rows(pool, tr, &run);
7315            unsafe {
7316                accel_blas::cblas_sgemm(
7317                    101, // RowMajor
7318                    111, // NoTrans A
7319                    112, // Trans B
7320                    b as i32,
7321                    tr as i32,
7322                    cols as i32,
7323                    1.0,
7324                    xs_all.as_ptr(),
7325                    cols as i32,
7326                    wtile.as_ptr(),
7327                    cols as i32,
7328                    0.0,
7329                    out.as_mut_ptr().add(r0),
7330                    rows as i32,
7331                );
7332            }
7333            r0 += tr;
7334        }
7335    });
7336}
7337
7338fn q4t_matmat(
7339    bytes: &[u8],
7340    xs_all: &[f32],
7341    b: usize,
7342    rows: usize,
7343    cols: usize,
7344    out: &mut [f32],
7345    pool: Option<&Pool>,
7346) {
7347    debug_assert_eq!(out.len(), b * rows);
7348    let gpr = cols / GROUP_SIZE;
7349    // Wide batches ride the AMX like q8's qmatmat: on Apple silicon
7350    // the dequant-tile sgemm is an order above the SDOT row loop for
7351    // prefill shapes (imagegen DiT forwards are exactly this).
7352    #[cfg(target_os = "macos")]
7353    if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
7354        dequant_matmat_accel(
7355            &|r, dst| {
7356                for gi in 0..gpr {
7357                    let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
7358                    let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
7359                    for (k, &bb) in tile[2..].iter().enumerate() {
7360                        dst[gi * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s;
7361                        dst[gi * GROUP_SIZE + k * 2 + 1] = (((bb >> 4) & 0x0F) as f32 - 8.0) * s;
7362                    }
7363                }
7364            },
7365            xs_all,
7366            b,
7367            rows,
7368            cols,
7369            out,
7370            pool,
7371        );
7372        return;
7373    }
7374    let out_addr = SendMut(out.as_mut_ptr());
7375    if a8w8_enabled() {
7376        let acts: Vec<SplitAct> = (0..b)
7377            .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
7378            .collect();
7379        let acts = &acts;
7380        #[cfg(target_arch = "x86_64")]
7381        let blocked_ok = avx2_enabled() && blocked_enabled();
7382        #[cfg(target_arch = "aarch64")]
7383        let blocked_ok = sdot_enabled() && blocked_enabled();
7384        #[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))]
7385        let blocked_ok = false;
7386        let run = move |start: usize, end: usize| {
7387            for r in start..end {
7388                let mut bi = 0usize;
7389                #[cfg(target_arch = "aarch64")]
7390                if blocked_ok {
7391                    while bi + 4 <= acts.len() {
7392                        let xs = [
7393                            acts[bi].xq.as_slice(),
7394                            acts[bi + 1].xq.as_slice(),
7395                            acts[bi + 2].xq.as_slice(),
7396                            acts[bi + 3].xq.as_slice(),
7397                        ];
7398                        let d = unsafe { dot_q4t_row_1x4_sdot(bytes, r, gpr, xs) };
7399                        for k in 0..4 {
7400                            let act = &acts[bi + k];
7401                            let mut acc = d[k] * act.sx;
7402                            for &(j, xv) in &act.outliers {
7403                                let (w, sc) = q4t_outlier(bytes, r, gpr, j);
7404                                acc += w * sc * xv;
7405                            }
7406                            // SAFETY: disjoint (bi, r) cells per worker.
7407                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
7408                        }
7409                        bi += 4;
7410                    }
7411                }
7412                #[cfg(target_arch = "x86_64")]
7413                if blocked_ok {
7414                    while bi + 4 <= acts.len() {
7415                        let xs = [
7416                            acts[bi].xq.as_slice(),
7417                            acts[bi + 1].xq.as_slice(),
7418                            acts[bi + 2].xq.as_slice(),
7419                            acts[bi + 3].xq.as_slice(),
7420                        ];
7421                        let d = unsafe {
7422                            if vnni_tiles_enabled() {
7423                                dot_q4t_row_1x4_vnni(bytes, r, gpr, xs)
7424                            } else {
7425                                dot_q4t_row_1x4_avx2(bytes, r, gpr, xs)
7426                            }
7427                        };
7428                        for k in 0..4 {
7429                            let act = &acts[bi + k];
7430                            let mut acc = d[k] * act.sx;
7431                            for &(j, xv) in &act.outliers {
7432                                let (w, sc) = q4t_outlier(bytes, r, gpr, j);
7433                                acc += w * sc * xv;
7434                            }
7435                            // SAFETY: disjoint (bi, r) cells per worker.
7436                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
7437                        }
7438                        bi += 4;
7439                    }
7440                }
7441                let _ = blocked_ok;
7442                while bi < acts.len() {
7443                    let act = &acts[bi];
7444                    let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
7445                    for &(j, xv) in &act.outliers {
7446                        let (w, s) = q4t_outlier(bytes, r, gpr, j);
7447                        acc += w * s * xv;
7448                    }
7449                    // SAFETY: disjoint (bi, r) cells per worker range.
7450                    unsafe { *out_addr.at(bi * rows + r) = acc };
7451                    bi += 1;
7452                }
7453            }
7454        };
7455        dispatch_rows(pool, rows, &run);
7456        return;
7457    }
7458    let run = move |start: usize, end: usize| {
7459        for r in start..end {
7460            for bi in 0..b {
7461                let x = &xs_all[bi * cols..(bi + 1) * cols];
7462                // SAFETY: disjoint (bi, r) cells per worker range.
7463                unsafe { *out_addr.at(bi * rows + r) = q4t_row_exact(bytes, r, gpr, x) };
7464            }
7465        }
7466    };
7467    dispatch_rows(pool, rows, &run);
7468}
7469
7470// ── q1 (dtype 12): binary weights, [f16 scale][4B sign bits] per
7471// 32-group tile. The kernel family mirrors q4_tiled: one sequential
7472// stream of 6-byte tiles, per-tile integer dot × scale, exact outlier
7473// correction (A8W8 contract), exact scalar path under CMF_SDOT=0. ──
7474
7475/// Per-32-group sums of the quantized activation — the ±1 identity's
7476/// shared half: `dot = −2·sdot(mask, x) − gsum[g]`, computed ONCE per
7477/// matvec and reused by every row.
7478fn q1_group_sums(xq: &[i8], gpr: usize) -> Vec<i32> {
7479    (0..gpr)
7480        .map(|gi| {
7481            xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE]
7482                .iter()
7483                .map(|&v| v as i32)
7484                .sum()
7485        })
7486        .collect()
7487}
7488
7489/// One q1 row via the A8W8 int8 path — mask-SDOT on ARM (no ±1
7490/// expansion at all), scalar bit loop elsewhere (AVX2 queued with the
7491/// x86 pass).
7492#[inline]
7493#[allow(unreachable_code)]
7494/// AVX2 q1 row via the same ±1 identity as the ARM sdot kernel: the
7495/// sign bits expand to a {0, −1} byte mask through shuffle+cmpeq, the
7496/// masked activation sums through maddubs(1, x&mask), and
7497/// `dot = −(2·masked_sum + Σx_group)` — bit-identical integer math.
7498#[cfg(target_arch = "x86_64")]
7499#[target_feature(enable = "avx2")]
7500unsafe fn dot_q1_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
7501    // SAFETY: callers uphold the 6B-tile and xq/gsum length contracts.
7502    unsafe {
7503        use core::arch::x86_64::*;
7504        // Byte j of the mask must replicate bits-byte j/8.
7505        let expand = _mm256_setr_epi8(
7506            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,
7507            3, 3, 3,
7508        );
7509        let bitsel = _mm256_setr_epi8(
7510            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
7511            -128, 1, 2, 4, 8, 16, 32, 64, -128,
7512        );
7513        let ones8 = _mm256_set1_epi8(1);
7514        let ones16 = _mm256_set1_epi16(1);
7515        let mut acc = 0f32;
7516        for gi in 0..gpr {
7517            let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
7518            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
7519            let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
7520            let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
7521            let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
7522            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7523            let sel = _mm256_and_si256(x, mask);
7524            // Σ of selected i8 lanes: maddubs(1u8, sel_i8) pairs → madd.
7525            let p16 = _mm256_maddubs_epi16(ones8, sel);
7526            let d32 = _mm256_madd_epi16(p16, ones16);
7527            let hi128 = _mm256_extracti128_si256::<1>(d32);
7528            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
7529            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7530            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7531            let msum = _mm_cvtsi128_si32(s32);
7532            // The and-select keeps x UN-negated (unlike ARM's −1-mask
7533            // sdot): d = Σ_set − Σ_unset = 2·Σ_set − Σ_all.
7534            let d = 2 * msum - gsum[gi];
7535            acc += d as f32 * s;
7536        }
7537        acc
7538    }
7539}
7540
7541/// VNNI twin of `dot_q1_row_avx2`: the masked-select sum goes through
7542/// one `vpdpbusd(1u8, sel)` (see `dpbusd_hsum` — bit-identical).
7543#[cfg(target_arch = "x86_64")]
7544#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
7545unsafe fn dot_q1_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
7546    // SAFETY: callers uphold the 6B-tile and xq/gsum length contracts.
7547    unsafe {
7548        use core::arch::x86_64::*;
7549        let expand = _mm256_setr_epi8(
7550            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,
7551            3, 3, 3,
7552        );
7553        let bitsel = _mm256_setr_epi8(
7554            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
7555            -128, 1, 2, 4, 8, 16, 32, 64, -128,
7556        );
7557        let ones8 = _mm256_set1_epi8(1);
7558        let mut acc = 0f32;
7559        for gi in 0..gpr {
7560            let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
7561            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
7562            let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
7563            let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
7564            let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
7565            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7566            let msum = dpbusd_hsum(ones8, _mm256_and_si256(x, mask));
7567            let d = 2 * msum - gsum[gi];
7568            acc += d as f32 * s;
7569        }
7570        acc
7571    }
7572}
7573
7574/// VNNI twin of `dot_q1_row_1x4_avx2` (see `dpbusd_hsum`).
7575#[cfg(target_arch = "x86_64")]
7576#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
7577unsafe fn dot_q1_row_1x4_vnni(
7578    bytes: &[u8],
7579    r: usize,
7580    gpr: usize,
7581    xs: [&[i8]; 4],
7582    gsums: [&[i32]; 4],
7583) -> [f32; 4] {
7584    // SAFETY: callers uphold the 6B-tile and xq/gsum length contracts.
7585    unsafe {
7586        use core::arch::x86_64::*;
7587        let expand = _mm256_setr_epi8(
7588            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,
7589            3, 3, 3,
7590        );
7591        let bitsel = _mm256_setr_epi8(
7592            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
7593            -128, 1, 2, 4, 8, 16, 32, 64, -128,
7594        );
7595        let ones8 = _mm256_set1_epi8(1);
7596        let mut acc = [0f32; 4];
7597        for gi in 0..gpr {
7598            let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
7599            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
7600            let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
7601            let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
7602            let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
7603            for (k, xq) in xs.iter().enumerate() {
7604                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7605                let msum = dpbusd_hsum(ones8, _mm256_and_si256(x, mask));
7606                let d = 2 * msum - gsums[k][gi];
7607                acc[k] += d as f32 * s;
7608            }
7609        }
7610        acc
7611    }
7612}
7613
7614/// The blocked 1×4 flavor: the expanded bit mask serves four activation
7615/// streams per group (mask build once, four select+reduce chains).
7616#[cfg(target_arch = "x86_64")]
7617#[target_feature(enable = "avx2")]
7618unsafe fn dot_q1_row_1x4_avx2(
7619    bytes: &[u8],
7620    r: usize,
7621    gpr: usize,
7622    xs: [&[i8]; 4],
7623    gsums: [&[i32]; 4],
7624) -> [f32; 4] {
7625    // SAFETY: callers uphold the 6B-tile and xq/gsum length contracts.
7626    unsafe {
7627        use core::arch::x86_64::*;
7628        let expand = _mm256_setr_epi8(
7629            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,
7630            3, 3, 3,
7631        );
7632        let bitsel = _mm256_setr_epi8(
7633            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
7634            -128, 1, 2, 4, 8, 16, 32, 64, -128,
7635        );
7636        let ones8 = _mm256_set1_epi8(1);
7637        let ones16 = _mm256_set1_epi16(1);
7638        let mut acc = [0f32; 4];
7639        for gi in 0..gpr {
7640            let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
7641            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
7642            let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
7643            let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
7644            let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
7645            for (k, xq) in xs.iter().enumerate() {
7646                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7647                let sel = _mm256_and_si256(x, mask);
7648                let p16 = _mm256_maddubs_epi16(ones8, sel);
7649                let d32 = _mm256_madd_epi16(p16, ones16);
7650                let hi128 = _mm256_extracti128_si256::<1>(d32);
7651                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
7652                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7653                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7654                let msum = _mm_cvtsi128_si32(s32);
7655                let d = 2 * msum - gsums[k][gi];
7656                acc[k] += d as f32 * s;
7657            }
7658        }
7659        acc
7660    }
7661}
7662
7663#[allow(unreachable_code)]
7664fn dot_q1_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
7665    #[cfg(target_arch = "aarch64")]
7666    unsafe {
7667        return dot_q1_row_sdot(bytes, r, gpr, xq, gsum);
7668    }
7669    #[cfg(target_arch = "x86_64")]
7670    if avx2_enabled() {
7671        unsafe {
7672            if vnni_tiles_enabled() {
7673                return dot_q1_row_vnni(bytes, r, gpr, xq, gsum);
7674            }
7675            return dot_q1_row_avx2(bytes, r, gpr, xq, gsum);
7676        }
7677    }
7678    let _ = gsum;
7679    let mut acc = 0f32;
7680    for gi in 0..gpr {
7681        let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
7682        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
7683        let mut d = 0i32;
7684        for (j, &b) in tile[2..].iter().enumerate() {
7685            for k in 0..8 {
7686                let w = ((b >> k) & 1) as i32 * 2 - 1;
7687                d += w * xq[gi * GROUP_SIZE + j * 8 + k] as i32;
7688            }
7689        }
7690        acc += d as f32 * s;
7691    }
7692    acc
7693}
7694
7695/// SDOT q1 row via the ±1 identity: the vtst mask (0xFF where the bit
7696/// is set, i.e. −1 as i8) feeds `sdot` DIRECTLY — no expansion to ±1
7697/// lanes at all — and `dot = −(2·sdot(mask, x) + Σx_group)`, with the
7698/// per-group activation sums shared across every row of the matvec.
7699/// Four tiles (128 weights) per iteration: integer dots reduce through
7700/// a vpaddq tree into ONE i32x4 that meets its four scales in a single
7701/// fused f32 multiply-add. Integer math throughout — bit-identical to
7702/// the scalar ±1 reference.
7703#[cfg(target_arch = "aarch64")]
7704#[target_feature(enable = "neon,dotprod")]
7705unsafe fn dot_q1_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
7706    // SAFETY: callers uphold slice-length contracts (6B tile per group,
7707    // xq.len() == gpr·GROUP_SIZE, gsum.len() == gpr).
7708    unsafe {
7709        use core::arch::aarch64::*;
7710        use core::arch::asm;
7711        const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
7712        let m = vld1q_u8(MASKS.as_ptr());
7713        // One tile's −Σ_set(x) as an UNREDUCED i32x4 (two mask-sdots).
7714        macro_rules! tile_dot {
7715            ($t:expr, $x:expr) => {{
7716                let v0 = vcombine_u8(vdup_n_u8(*$t.add(2)), vdup_n_u8(*$t.add(3)));
7717                let v1 = vcombine_u8(vdup_n_u8(*$t.add(4)), vdup_n_u8(*$t.add(5)));
7718                let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
7719                let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
7720                let x0 = vld1q_s8($x);
7721                let x1 = vld1q_s8($x.add(16));
7722                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
7723                asm!(
7724                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
7725                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
7726                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
7727                    w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
7728                    options(pure, nomem, nostack),
7729                );
7730                vaddq_s32(a0, a1)
7731            }};
7732        }
7733        // TBL unpack over PAIR loads: one vld1q covers two 6B tiles
7734        // ([s s b b b b][s s b b b b] + 4B slack), TBL replicates each
7735        // bit-byte across 8 lanes for vtst, and the four scales gather
7736        // through tbl2 into one fcvtl — the 16 ld1r broadcast loads and
7737        // 4 branchy software f16 conversions per 128 weights (the
7738        // measured load-port wall of this kernel) become 2 vector
7739        // loads + 9 table lookups. Integer math order is unchanged —
7740        // bit-identical results (FCVTL is exact on every f16).
7741        const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
7742        const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
7743        const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
7744        const IW11: [u8; 16] = [
7745            10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
7746        ];
7747        const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
7748        let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
7749        let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
7750        let isc = vld1_u8(ISC.as_ptr());
7751        // One tile's −Σ_set(x) from a TBL-unpacked pair load.
7752        macro_rules! tile_dot_tbl {
7753            ($ld:expr, $i0:expr, $i1:expr, $x:expr) => {{
7754                let w0 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i0), m));
7755                let w1 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i1), m));
7756                let x0 = vld1q_s8($x);
7757                let x1 = vld1q_s8($x.add(16));
7758                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
7759                asm!(
7760                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
7761                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
7762                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
7763                    w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
7764                    options(pure, nomem, nostack),
7765                );
7766                vaddq_s32(a0, a1)
7767            }};
7768        }
7769        let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
7770        let row_base = r * gpr * Q1_TILE;
7771        let abs_end = bytes.len();
7772        let xp = xq.as_ptr();
7773        let gp = gsum.as_ptr();
7774        let mut accv = vdupq_n_f32(0.0);
7775        let mut gi = 0;
7776        // The second pair load reads 4B past tile gi+3 — stay inside
7777        // the payload slice (only the file's final tiles fall back).
7778        while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
7779            let t0 = base.add(gi * Q1_TILE);
7780            let ld_a = vld1q_u8(t0);
7781            let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
7782            let d0 = tile_dot_tbl!(ld_a, iw00, iw01, xp.add(gi * GROUP_SIZE));
7783            let d1 = tile_dot_tbl!(ld_a, iw10, iw11, xp.add((gi + 1) * GROUP_SIZE));
7784            let d2 = tile_dot_tbl!(ld_b, iw00, iw01, xp.add((gi + 2) * GROUP_SIZE));
7785            let d3 = tile_dot_tbl!(ld_b, iw10, iw11, xp.add((gi + 3) * GROUP_SIZE));
7786            // [−Σ0, −Σ1, −Σ2, −Σ3] → dots = −(2·Σset_neg + gsum)
7787            let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
7788            let g = vld1q_s32(gp.add(gi));
7789            let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
7790            let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
7791            let scf: float32x4_t;
7792            asm!(
7793                "fcvtl {o:v}.4s, {i:v}.4h",
7794                o = out(vreg) scf, i = in(vreg) sc16,
7795                options(pure, nomem, nostack),
7796            );
7797            accv = vfmaq_f32(accv, vcvtq_f32_s32(dots), scf);
7798            gi += 4;
7799        }
7800        let mut acc = vaddvq_f32(accv);
7801        while gi < gpr {
7802            let t = base.add(gi * Q1_TILE);
7803            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
7804            let d = vaddvq_s32(tile_dot!(t, xp.add(gi * GROUP_SIZE)));
7805            acc += (-(2 * d + *gp.add(gi))) as f32 * s;
7806            gi += 1;
7807        }
7808        acc
7809    }
7810}
7811
7812/// Blocked q1 1×4: one TBL unpack of the tile pair serves FOUR
7813/// activation streams (prefill amortization — the same idea as the
7814/// AVX2 twin; per stream the group order, fma order and tail match the
7815/// single-row kernel exactly, so batch == matvec bit-for-bit).
7816#[cfg(target_arch = "aarch64")]
7817#[target_feature(enable = "neon,dotprod")]
7818unsafe fn dot_q1_row_1x4_sdot(
7819    bytes: &[u8],
7820    r: usize,
7821    gpr: usize,
7822    xs: [&[i8]; 4],
7823    gs: [&[i32]; 4],
7824) -> [f32; 4] {
7825    // SAFETY: same slice-length contracts as `dot_q1_row_sdot`, ×4.
7826    unsafe {
7827        use core::arch::aarch64::*;
7828        use core::arch::asm;
7829        const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
7830        const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
7831        const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
7832        const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
7833        const IW11: [u8; 16] = [
7834            10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
7835        ];
7836        const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
7837        let m = vld1q_u8(MASKS.as_ptr());
7838        let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
7839        let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
7840        let isc = vld1_u8(ISC.as_ptr());
7841        macro_rules! sdot2 {
7842            ($w0:expr, $w1:expr, $x:expr) => {{
7843                let x0 = vld1q_s8($x);
7844                let x1 = vld1q_s8($x.add(16));
7845                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
7846                asm!(
7847                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
7848                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
7849                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
7850                    w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
7851                    options(pure, nomem, nostack),
7852                );
7853                vaddq_s32(a0, a1)
7854            }};
7855        }
7856        let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
7857        let row_base = r * gpr * Q1_TILE;
7858        let abs_end = bytes.len();
7859        let mut accv = [vdupq_n_f32(0.0); 4];
7860        let mut gi = 0;
7861        while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
7862            let t0 = base.add(gi * Q1_TILE);
7863            let ld_a = vld1q_u8(t0);
7864            let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
7865            // Unpack ONCE — eight ±mask vectors serve all four streams.
7866            let w00 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw00), m));
7867            let w01 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw01), m));
7868            let w10 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw10), m));
7869            let w11 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw11), m));
7870            let w20 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw00), m));
7871            let w21 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw01), m));
7872            let w30 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw10), m));
7873            let w31 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw11), m));
7874            let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
7875            let scf: float32x4_t;
7876            asm!(
7877                "fcvtl {o:v}.4s, {i:v}.4h",
7878                o = out(vreg) scf, i = in(vreg) sc16,
7879                options(pure, nomem, nostack),
7880            );
7881            for k in 0..4 {
7882                let xp = xs[k].as_ptr();
7883                let d0 = sdot2!(w00, w01, xp.add(gi * GROUP_SIZE));
7884                let d1 = sdot2!(w10, w11, xp.add((gi + 1) * GROUP_SIZE));
7885                let d2 = sdot2!(w20, w21, xp.add((gi + 2) * GROUP_SIZE));
7886                let d3 = sdot2!(w30, w31, xp.add((gi + 3) * GROUP_SIZE));
7887                let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
7888                let g = vld1q_s32(gs[k].as_ptr().add(gi));
7889                let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
7890                accv[k] = vfmaq_f32(accv[k], vcvtq_f32_s32(dots), scf);
7891            }
7892            gi += 4;
7893        }
7894        let mut acc = [
7895            vaddvq_f32(accv[0]),
7896            vaddvq_f32(accv[1]),
7897            vaddvq_f32(accv[2]),
7898            vaddvq_f32(accv[3]),
7899        ];
7900        while gi < gpr {
7901            let t = base.add(gi * Q1_TILE);
7902            let sc = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
7903            let v0 = vcombine_u8(vdup_n_u8(*t.add(2)), vdup_n_u8(*t.add(3)));
7904            let v1 = vcombine_u8(vdup_n_u8(*t.add(4)), vdup_n_u8(*t.add(5)));
7905            let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
7906            let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
7907            for k in 0..4 {
7908                let d = vaddvq_s32(sdot2!(w0, w1, xs[k].as_ptr().add(gi * GROUP_SIZE)));
7909                acc[k] += (-(2 * d + *gs[k].as_ptr().add(gi))) as f32 * sc;
7910            }
7911            gi += 1;
7912        }
7913        acc
7914    }
7915}
7916
7917/// (weight ±1, scale) of one q1 element — the exact outlier term.
7918#[inline]
7919fn q1_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
7920    let gi = j / GROUP_SIZE;
7921    let k = j % GROUP_SIZE;
7922    let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
7923    let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
7924    let bit = (tile[2 + k / 8] >> (k % 8)) & 1;
7925    ((bit as i32 * 2 - 1) as f32, s)
7926}
7927
7928/// Exact scalar q1 row (CMF_SDOT=0 contract).
7929#[inline]
7930fn q1_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
7931    let mut acc = 0f32;
7932    for gi in 0..gpr {
7933        let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
7934        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
7935        let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
7936        let mut ga = 0f32;
7937        for (j, &b) in tile[2..].iter().enumerate() {
7938            for k in 0..8 {
7939                ga += (((b >> k) & 1) as f32 * 2.0 - 1.0) * xg[j * 8 + k];
7940            }
7941        }
7942        acc += ga * s;
7943    }
7944    acc
7945}
7946
7947/// One binary-weight row range via A8W8 (the body of `bin_matvec`'s hot
7948/// loop, extracted so multi-matrix jobs drive the same kernel).
7949#[allow(clippy::too_many_arguments)]
7950fn bin_range_a8w8(
7951    w: BinW,
7952    gpr: usize,
7953    act: &SplitAct,
7954    gsum: &[i32],
7955    out: SendMut,
7956    start: usize,
7957    end: usize,
7958) {
7959    for r in start..end {
7960        let mut acc = w.dot_i8(r, gpr, &act.xq, gsum) * act.sx;
7961        for &(j, xv) in &act.outliers {
7962            let (wv, s) = w.outlier(r, gpr, j);
7963            acc += wv * s * xv;
7964        }
7965        // SAFETY: disjoint row ranges per worker.
7966        unsafe { *out.at(r) = acc };
7967    }
7968}
7969
7970/// Exact-scalar binary-weight row range (CMF_SDOT=0 contract).
7971fn bin_range_f32(w: BinW, gpr: usize, x: &[f32], out: SendMut, start: usize, end: usize) {
7972    for r in start..end {
7973        // SAFETY: disjoint row ranges per worker.
7974        unsafe { *out.at(r) = w.row_exact(r, gpr, x) };
7975    }
7976}
7977
7978/// A binary (±s) weight in either of its layouts: `q1` 6-byte tiles, or
7979/// `q1g` planes with one scale per 128-group. The row drivers — matvec,
7980/// matvec2, matmat, the multi-projection job and the fused gate/up pair —
7981/// are written once against this; only the per-row kernels differ, and
7982/// the match costs one predictable branch per row. `gpr` is always the
7983/// row's count of 32-tiles (cols / 32): the activation group sums and the
7984/// tile walk are shared by both layouts.
7985#[derive(Clone, Copy)]
7986pub(crate) enum BinW<'a> {
7987    Q1(&'a [u8]),
7988    Q1g {
7989        bytes: &'a [u8],
7990        /// Byte offset of the f16 scale plane (= rows · cols/8).
7991        sc_off: usize,
7992        /// Row count: the q1 kernels treat the payload's final row
7993        /// specially (see `dot_q1g_row_sdot`), and q1g mirrors that.
7994        rows: usize,
7995    },
7996}
7997
7998impl<'a> BinW<'a> {
7999    fn q1g(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
8000        Self::Q1g {
8001            bytes,
8002            sc_off: cortiq_core::quant::q1g_scale_offset(rows, cols),
8003            rows,
8004        }
8005    }
8006
8007    #[inline]
8008    fn dot_i8(self, r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
8009        match self {
8010            Self::Q1(b) => dot_q1_row_i8(b, r, gpr, xq, gsum),
8011            Self::Q1g {
8012                bytes,
8013                sc_off,
8014                rows,
8015            } => dot_q1g_row_i8(bytes, sc_off, rows, r, gpr, xq, gsum),
8016        }
8017    }
8018
8019    #[inline]
8020    fn row_exact(self, r: usize, gpr: usize, x: &[f32]) -> f32 {
8021        match self {
8022            Self::Q1(b) => q1_row_exact(b, r, gpr, x),
8023            Self::Q1g { bytes, sc_off, .. } => q1g_row_exact(bytes, sc_off, r, gpr, x),
8024        }
8025    }
8026
8027    #[inline]
8028    fn outlier(self, r: usize, gpr: usize, j: usize) -> (f32, f32) {
8029        match self {
8030            Self::Q1(b) => q1_outlier(b, r, gpr, j),
8031            Self::Q1g { bytes, sc_off, .. } => q1g_outlier(bytes, sc_off, r, gpr, j),
8032        }
8033    }
8034
8035    /// Blocked 1×4 row (one unpack, four activation streams).
8036    ///
8037    /// # Safety
8038    /// The SDOT/dotprod target feature must be present (`sdot_enabled`).
8039    #[cfg(target_arch = "aarch64")]
8040    #[inline]
8041    unsafe fn dot_1x4_sdot(
8042        self,
8043        r: usize,
8044        gpr: usize,
8045        xs: [&[i8]; 4],
8046        gs: [&[i32]; 4],
8047    ) -> [f32; 4] {
8048        // SAFETY: forwarded caller contract.
8049        unsafe {
8050            match self {
8051                Self::Q1(b) => dot_q1_row_1x4_sdot(b, r, gpr, xs, gs),
8052                Self::Q1g {
8053                    bytes,
8054                    sc_off,
8055                    rows,
8056                } => dot_q1g_row_1x4_sdot(bytes, sc_off, rows, r, gpr, xs, gs),
8057            }
8058        }
8059    }
8060
8061    /// Blocked 1×4 row on x86 (VNNI when the tile switch is on).
8062    ///
8063    /// # Safety
8064    /// AVX2 must be present (`avx2_enabled`).
8065    #[cfg(target_arch = "x86_64")]
8066    #[inline]
8067    unsafe fn dot_1x4_x86(self, r: usize, gpr: usize, xs: [&[i8]; 4], gs: [&[i32]; 4]) -> [f32; 4] {
8068        // SAFETY: forwarded caller contract.
8069        unsafe {
8070            match (self, vnni_tiles_enabled()) {
8071                (Self::Q1(b), true) => dot_q1_row_1x4_vnni(b, r, gpr, xs, gs),
8072                (Self::Q1(b), false) => dot_q1_row_1x4_avx2(b, r, gpr, xs, gs),
8073                (Self::Q1g { bytes, sc_off, .. }, true) => {
8074                    dot_q1g_row_1x4_vnni(bytes, sc_off, r, gpr, xs, gs)
8075                }
8076                (Self::Q1g { bytes, sc_off, .. }, false) => {
8077                    dot_q1g_row_1x4_avx2(bytes, sc_off, r, gpr, xs, gs)
8078                }
8079            }
8080        }
8081    }
8082}
8083
8084// ── q1g (dtype 20): binary weights, ONE f16 scale per 128-group, in sign
8085// and scale PLANES (`[bits: rows·cols/8][f16: rows·cols/128]`). Same
8086// kernel family as q1 and — deliberately — the same float accumulation
8087// order (one `acc += d·s` per 32-tile, the group's scale reused four
8088// times), so a q1g repack of a q1 tensor whose tile scales agree gives
8089// the same f32 outputs bit for bit on every path. The win is the stream:
8090// 18 bytes per 128 weights instead of 24, the bits 16-byte aligned. ──
8091
8092/// Sign bits of q1g tile `t` (flat 32-tile index `r·gpr + gi`).
8093#[inline]
8094fn q1g_bits(bytes: &[u8], t: usize) -> u32 {
8095    u32::from_le_bytes([
8096        bytes[t * 4],
8097        bytes[t * 4 + 1],
8098        bytes[t * 4 + 2],
8099        bytes[t * 4 + 3],
8100    ])
8101}
8102
8103/// Raw f16 bits of the scale of the 128-group holding q1g tile `t`.
8104#[inline]
8105fn q1g_scale_bits(bytes: &[u8], sc_off: usize, t: usize) -> u16 {
8106    let o = sc_off + (t / 4) * 2;
8107    u16::from_le_bytes([bytes[o], bytes[o + 1]])
8108}
8109
8110/// One q1g row via the A8W8 int8 path (dispatch twin of `dot_q1_row_i8`).
8111#[allow(unreachable_code)]
8112#[allow(clippy::too_many_arguments)]
8113fn dot_q1g_row_i8(
8114    bytes: &[u8],
8115    sc_off: usize,
8116    rows: usize,
8117    r: usize,
8118    gpr: usize,
8119    xq: &[i8],
8120    gsum: &[i32],
8121) -> f32 {
8122    #[cfg(target_arch = "aarch64")]
8123    unsafe {
8124        return dot_q1g_row_sdot(bytes, sc_off, rows, r, gpr, xq, gsum);
8125    }
8126    #[cfg(target_arch = "x86_64")]
8127    if avx2_enabled() {
8128        unsafe {
8129            if vnni_tiles_enabled() {
8130                return dot_q1g_row_vnni(bytes, sc_off, r, gpr, xq, gsum);
8131            }
8132            return dot_q1g_row_avx2(bytes, sc_off, r, gpr, xq, gsum);
8133        }
8134    }
8135    let _ = (gsum, rows);
8136    let mut acc = 0f32;
8137    for gi in 0..gpr {
8138        let t = r * gpr + gi;
8139        let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, t));
8140        let bits = q1g_bits(bytes, t);
8141        let mut d = 0i32;
8142        for k in 0..GROUP_SIZE {
8143            let w = ((bits >> k) & 1) as i32 * 2 - 1;
8144            d += w * xq[gi * GROUP_SIZE + k] as i32;
8145        }
8146        acc += d as f32 * s;
8147    }
8148    acc
8149}
8150
8151/// Four tiles' unreduced `i32×8` partial sums → `[Σ0, Σ1, Σ2, Σ3]` in one
8152/// hadd tree: 3 hadds + 1 add for four tiles where the per-tile reduce of
8153/// the q1 kernels costs 7 ops each. Integer math, so any order gives the
8154/// q1 kernel's per-tile sums exactly.
8155#[cfg(target_arch = "x86_64")]
8156#[target_feature(enable = "avx2")]
8157#[inline]
8158unsafe fn q1g_hsum4(
8159    d0: core::arch::x86_64::__m256i,
8160    d1: core::arch::x86_64::__m256i,
8161    d2: core::arch::x86_64::__m256i,
8162    d3: core::arch::x86_64::__m256i,
8163) -> core::arch::x86_64::__m128i {
8164    // SAFETY: pure register math.
8165    unsafe {
8166        use core::arch::x86_64::*;
8167        let h = _mm256_hadd_epi32(_mm256_hadd_epi32(d0, d1), _mm256_hadd_epi32(d2, d3));
8168        _mm_add_epi32(_mm256_castsi256_si128(h), _mm256_extracti128_si256::<1>(h))
8169    }
8170}
8171
8172/// One 128-group's epilogue, in q1's float order: `d = 2·Σset − Σx` per
8173/// tile (exact i32 → exact f32), then `acc += d·s` tile by tile — four
8174/// separately rounded multiply/add pairs, the q1 kernels' sequence.
8175#[cfg(target_arch = "x86_64")]
8176#[target_feature(enable = "avx2")]
8177#[inline]
8178unsafe fn q1g_group_acc(
8179    acc: f32,
8180    msum4: core::arch::x86_64::__m128i,
8181    gsum4: &[i32],
8182    s: f32,
8183) -> f32 {
8184    // SAFETY: gsum4 holds the group's four tile sums (callers slice it).
8185    unsafe {
8186        use core::arch::x86_64::*;
8187        let g = _mm_loadu_si128(gsum4.as_ptr() as *const __m128i);
8188        let d = _mm_cvtepi32_ps(_mm_sub_epi32(_mm_slli_epi32::<1>(msum4), g));
8189        let mut f = [0f32; 4];
8190        _mm_storeu_ps(f.as_mut_ptr(), d);
8191        let mut acc = acc;
8192        for v in f {
8193            acc += v * s;
8194        }
8195        acc
8196    }
8197}
8198
8199/// AVX2 q1g row: `dot_q1_row_avx2`'s mask/select/maddubs per tile, but the
8200/// four tiles of a 128-group reduce together (`q1g_hsum4`) and the scale
8201/// converts once per group — bit-identical results (integer sums, then
8202/// q1's per-tile float sequence), fewer horizontal ops and f16 decodes.
8203#[cfg(target_arch = "x86_64")]
8204#[target_feature(enable = "avx2")]
8205unsafe fn dot_q1g_row_avx2(
8206    bytes: &[u8],
8207    sc_off: usize,
8208    r: usize,
8209    gpr: usize,
8210    xq: &[i8],
8211    gsum: &[i32],
8212) -> f32 {
8213    // SAFETY: callers uphold the plane-length and xq/gsum contracts
8214    // (gpr % 4 == 0, xq.len() == gpr·32, gsum.len() == gpr).
8215    unsafe {
8216        use core::arch::x86_64::*;
8217        let expand = _mm256_setr_epi8(
8218            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,
8219            3, 3, 3,
8220        );
8221        let bitsel = _mm256_setr_epi8(
8222            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
8223            -128, 1, 2, 4, 8, 16, 32, 64, -128,
8224        );
8225        let ones8 = _mm256_set1_epi8(1);
8226        let ones16 = _mm256_set1_epi16(1);
8227        let bp = bytes.as_ptr().add(r * gpr * 4);
8228        let mut acc = 0f32;
8229        let mut gi = 0;
8230        while gi < gpr {
8231            let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
8232            let mut d = [_mm256_setzero_si256(); 4];
8233            for (k, dk) in d.iter_mut().enumerate() {
8234                let t = gi + k;
8235                let bits = (bp.add(t * 4) as *const u32).read_unaligned();
8236                let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
8237                let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
8238                let x = _mm256_loadu_si256(xq.as_ptr().add(t * GROUP_SIZE) as *const __m256i);
8239                let sel = _mm256_and_si256(x, mask);
8240                *dk = _mm256_madd_epi16(_mm256_maddubs_epi16(ones8, sel), ones16);
8241            }
8242            let m4 = q1g_hsum4(d[0], d[1], d[2], d[3]);
8243            acc = q1g_group_acc(acc, m4, &gsum[gi..gi + 4], s);
8244            gi += 4;
8245        }
8246        acc
8247    }
8248}
8249
8250/// VNNI twin of `dot_q1g_row_avx2`: `vpdpbusd(1u8, sel)` per tile (the
8251/// same i32 lanes as maddubs+madd — bounded products), same group reduce.
8252#[cfg(target_arch = "x86_64")]
8253#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
8254unsafe fn dot_q1g_row_vnni(
8255    bytes: &[u8],
8256    sc_off: usize,
8257    r: usize,
8258    gpr: usize,
8259    xq: &[i8],
8260    gsum: &[i32],
8261) -> f32 {
8262    // SAFETY: callers uphold the plane-length and xq/gsum contracts.
8263    unsafe {
8264        use core::arch::x86_64::*;
8265        let expand = _mm256_setr_epi8(
8266            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,
8267            3, 3, 3,
8268        );
8269        let bitsel = _mm256_setr_epi8(
8270            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
8271            -128, 1, 2, 4, 8, 16, 32, 64, -128,
8272        );
8273        let ones8 = _mm256_set1_epi8(1);
8274        let bp = bytes.as_ptr().add(r * gpr * 4);
8275        let mut acc = 0f32;
8276        let mut gi = 0;
8277        while gi < gpr {
8278            let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
8279            let mut d = [_mm256_setzero_si256(); 4];
8280            for (k, dk) in d.iter_mut().enumerate() {
8281                let t = gi + k;
8282                let bits = (bp.add(t * 4) as *const u32).read_unaligned();
8283                let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
8284                let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
8285                let x = _mm256_loadu_si256(xq.as_ptr().add(t * GROUP_SIZE) as *const __m256i);
8286                *dk = _mm256_dpbusd_epi32(_mm256_setzero_si256(), ones8, _mm256_and_si256(x, mask));
8287            }
8288            let m4 = q1g_hsum4(d[0], d[1], d[2], d[3]);
8289            acc = q1g_group_acc(acc, m4, &gsum[gi..gi + 4], s);
8290            gi += 4;
8291        }
8292        acc
8293    }
8294}
8295
8296/// Blocked 1×4 AVX2 q1g row: each tile's mask is built once and serves
8297/// four activation streams; per stream, the group reduce and float order
8298/// of `dot_q1g_row_avx2` (so batch == matvec bit for bit).
8299#[cfg(target_arch = "x86_64")]
8300#[target_feature(enable = "avx2")]
8301unsafe fn dot_q1g_row_1x4_avx2(
8302    bytes: &[u8],
8303    sc_off: usize,
8304    r: usize,
8305    gpr: usize,
8306    xs: [&[i8]; 4],
8307    gsums: [&[i32]; 4],
8308) -> [f32; 4] {
8309    // SAFETY: callers uphold the plane-length and xq/gsum contracts.
8310    unsafe {
8311        use core::arch::x86_64::*;
8312        let expand = _mm256_setr_epi8(
8313            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,
8314            3, 3, 3,
8315        );
8316        let bitsel = _mm256_setr_epi8(
8317            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
8318            -128, 1, 2, 4, 8, 16, 32, 64, -128,
8319        );
8320        let ones8 = _mm256_set1_epi8(1);
8321        let ones16 = _mm256_set1_epi16(1);
8322        let bp = bytes.as_ptr().add(r * gpr * 4);
8323        let mut acc = [0f32; 4];
8324        let mut gi = 0;
8325        while gi < gpr {
8326            let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
8327            // d[stream][tile]
8328            let mut d = [[_mm256_setzero_si256(); 4]; 4];
8329            for t in 0..4 {
8330                let bits = (bp.add((gi + t) * 4) as *const u32).read_unaligned();
8331                let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
8332                let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
8333                for (k, xq) in xs.iter().enumerate() {
8334                    let x = _mm256_loadu_si256(
8335                        xq.as_ptr().add((gi + t) * GROUP_SIZE) as *const __m256i
8336                    );
8337                    let sel = _mm256_and_si256(x, mask);
8338                    d[k][t] = _mm256_madd_epi16(_mm256_maddubs_epi16(ones8, sel), ones16);
8339                }
8340            }
8341            for k in 0..4 {
8342                let m4 = q1g_hsum4(d[k][0], d[k][1], d[k][2], d[k][3]);
8343                acc[k] = q1g_group_acc(acc[k], m4, &gsums[k][gi..gi + 4], s);
8344            }
8345            gi += 4;
8346        }
8347        acc
8348    }
8349}
8350
8351/// VNNI twin of `dot_q1g_row_1x4_avx2`.
8352#[cfg(target_arch = "x86_64")]
8353#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
8354unsafe fn dot_q1g_row_1x4_vnni(
8355    bytes: &[u8],
8356    sc_off: usize,
8357    r: usize,
8358    gpr: usize,
8359    xs: [&[i8]; 4],
8360    gsums: [&[i32]; 4],
8361) -> [f32; 4] {
8362    // SAFETY: callers uphold the plane-length and xq/gsum contracts.
8363    unsafe {
8364        use core::arch::x86_64::*;
8365        let expand = _mm256_setr_epi8(
8366            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,
8367            3, 3, 3,
8368        );
8369        let bitsel = _mm256_setr_epi8(
8370            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
8371            -128, 1, 2, 4, 8, 16, 32, 64, -128,
8372        );
8373        let ones8 = _mm256_set1_epi8(1);
8374        let bp = bytes.as_ptr().add(r * gpr * 4);
8375        let mut acc = [0f32; 4];
8376        let mut gi = 0;
8377        while gi < gpr {
8378            let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
8379            let mut d = [[_mm256_setzero_si256(); 4]; 4];
8380            for t in 0..4 {
8381                let bits = (bp.add((gi + t) * 4) as *const u32).read_unaligned();
8382                let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
8383                let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
8384                for (k, xq) in xs.iter().enumerate() {
8385                    let x = _mm256_loadu_si256(
8386                        xq.as_ptr().add((gi + t) * GROUP_SIZE) as *const __m256i
8387                    );
8388                    d[k][t] = _mm256_dpbusd_epi32(
8389                        _mm256_setzero_si256(),
8390                        ones8,
8391                        _mm256_and_si256(x, mask),
8392                    );
8393                }
8394            }
8395            for k in 0..4 {
8396                let m4 = q1g_hsum4(d[k][0], d[k][1], d[k][2], d[k][3]);
8397                acc[k] = q1g_group_acc(acc[k], m4, &gsums[k][gi..gi + 4], s);
8398            }
8399            gi += 4;
8400        }
8401        acc
8402    }
8403}
8404
8405/// SDOT q1g row: the `dot_q1_row_sdot` math on the plane layout. One
8406/// aligned 16-byte load carries a whole 128-group's signs (q1 needs two
8407/// pair loads and a TBL gather of four scales for the same weights), TBL
8408/// replicates each sign byte across 8 lanes for vtst, and the group's ONE
8409/// scale converts once (FCVTL, exact) and broadcasts. Lane order, the
8410/// vpaddq tree, the fma into the four lane accumulators and the tail are
8411/// the q1 kernel's — including its quirk: q1's pair load reads 4 bytes
8412/// past the 4th tile, so the payload's final four tiles (last row) take
8413/// the scalar tail there. q1g has no overread, but takes the same tail on
8414/// the last row to stay bit-identical to q1.
8415#[cfg(target_arch = "aarch64")]
8416#[target_feature(enable = "neon,dotprod")]
8417#[allow(clippy::too_many_arguments)]
8418unsafe fn dot_q1g_row_sdot(
8419    bytes: &[u8],
8420    sc_off: usize,
8421    rows: usize,
8422    r: usize,
8423    gpr: usize,
8424    xq: &[i8],
8425    gsum: &[i32],
8426) -> f32 {
8427    // SAFETY: callers uphold the plane-length and xq/gsum contracts
8428    // (gpr % 4 == 0, xq.len() == gpr·32, gsum.len() == gpr).
8429    unsafe {
8430        use core::arch::aarch64::*;
8431        use core::arch::asm;
8432        const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
8433        // Tile t of the group = sign bytes 4t..4t+4 of the 16-byte load.
8434        const IW: [[u8; 16]; 8] = [
8435            [0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1],
8436            [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3],
8437            [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5],
8438            [6, 6, 6, 6, 6, 6, 6, 6, 7, 7, 7, 7, 7, 7, 7, 7],
8439            [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9],
8440            [
8441                10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
8442            ],
8443            [
8444                12, 12, 12, 12, 12, 12, 12, 12, 13, 13, 13, 13, 13, 13, 13, 13,
8445            ],
8446            [
8447                14, 14, 14, 14, 14, 14, 14, 14, 15, 15, 15, 15, 15, 15, 15, 15,
8448            ],
8449        ];
8450        let m = vld1q_u8(MASKS.as_ptr());
8451        let iw: [uint8x16_t; 8] = core::array::from_fn(|i| vld1q_u8(IW[i].as_ptr()));
8452        macro_rules! tile_dot {
8453            ($w0:expr, $w1:expr, $x:expr) => {{
8454                let x0 = vld1q_s8($x);
8455                let x1 = vld1q_s8($x.add(16));
8456                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
8457                asm!(
8458                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
8459                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
8460                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
8461                    w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
8462                    options(pure, nomem, nostack),
8463                );
8464                vaddq_s32(a0, a1)
8465            }};
8466        }
8467        macro_rules! tbl_mask {
8468            ($ld:expr, $i:expr) => {
8469                vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, iw[$i]), m))
8470            };
8471        }
8472        let bp = bytes.as_ptr().add(r * gpr * 4);
8473        let xp = xq.as_ptr();
8474        let gp = gsum.as_ptr();
8475        let vec_end = if r + 1 == rows { gpr - 4 } else { gpr };
8476        let mut accv = vdupq_n_f32(0.0);
8477        let mut gi = 0;
8478        while gi + 4 <= vec_end {
8479            let ld = vld1q_u8(bp.add(gi * 4));
8480            let d0 = tile_dot!(tbl_mask!(ld, 0), tbl_mask!(ld, 1), xp.add(gi * GROUP_SIZE));
8481            let d1 = tile_dot!(
8482                tbl_mask!(ld, 2),
8483                tbl_mask!(ld, 3),
8484                xp.add((gi + 1) * GROUP_SIZE)
8485            );
8486            let d2 = tile_dot!(
8487                tbl_mask!(ld, 4),
8488                tbl_mask!(ld, 5),
8489                xp.add((gi + 2) * GROUP_SIZE)
8490            );
8491            let d3 = tile_dot!(
8492                tbl_mask!(ld, 6),
8493                tbl_mask!(ld, 7),
8494                xp.add((gi + 3) * GROUP_SIZE)
8495            );
8496            let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
8497            let g = vld1q_s32(gp.add(gi));
8498            let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
8499            let h = vdup_n_u16(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
8500            let scf: float32x4_t;
8501            asm!(
8502                "fcvtl {o:v}.4s, {i:v}.4h",
8503                o = out(vreg) scf, i = in(vreg) h,
8504                options(pure, nomem, nostack),
8505            );
8506            accv = vfmaq_f32(accv, vcvtq_f32_s32(dots), scf);
8507            gi += 4;
8508        }
8509        let mut acc = vaddvq_f32(accv);
8510        while gi < gpr {
8511            let t = bp.add(gi * 4);
8512            let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
8513            let v0 = vcombine_u8(vdup_n_u8(*t), vdup_n_u8(*t.add(1)));
8514            let v1 = vcombine_u8(vdup_n_u8(*t.add(2)), vdup_n_u8(*t.add(3)));
8515            let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
8516            let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
8517            let d = vaddvq_s32(tile_dot!(w0, w1, xp.add(gi * GROUP_SIZE)));
8518            acc += (-(2 * d + *gp.add(gi))) as f32 * s;
8519            gi += 1;
8520        }
8521        acc
8522    }
8523}
8524
8525/// Blocked q1g 1×4: one unpack of the group serves four activation
8526/// streams; per stream bit-identical to `dot_q1g_row_sdot` (and so to the
8527/// q1 1×4 kernel, same last-row tail).
8528#[cfg(target_arch = "aarch64")]
8529#[target_feature(enable = "neon,dotprod")]
8530#[allow(clippy::too_many_arguments)]
8531unsafe fn dot_q1g_row_1x4_sdot(
8532    bytes: &[u8],
8533    sc_off: usize,
8534    rows: usize,
8535    r: usize,
8536    gpr: usize,
8537    xs: [&[i8]; 4],
8538    gs: [&[i32]; 4],
8539) -> [f32; 4] {
8540    // SAFETY: same contracts as `dot_q1g_row_sdot`, ×4.
8541    unsafe {
8542        use core::arch::aarch64::*;
8543        use core::arch::asm;
8544        const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
8545        const IW: [[u8; 16]; 8] = [
8546            [0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1],
8547            [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3],
8548            [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5],
8549            [6, 6, 6, 6, 6, 6, 6, 6, 7, 7, 7, 7, 7, 7, 7, 7],
8550            [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9],
8551            [
8552                10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
8553            ],
8554            [
8555                12, 12, 12, 12, 12, 12, 12, 12, 13, 13, 13, 13, 13, 13, 13, 13,
8556            ],
8557            [
8558                14, 14, 14, 14, 14, 14, 14, 14, 15, 15, 15, 15, 15, 15, 15, 15,
8559            ],
8560        ];
8561        let m = vld1q_u8(MASKS.as_ptr());
8562        let iw: [uint8x16_t; 8] = core::array::from_fn(|i| vld1q_u8(IW[i].as_ptr()));
8563        macro_rules! sdot2 {
8564            ($w0:expr, $w1:expr, $x:expr) => {{
8565                let x0 = vld1q_s8($x);
8566                let x1 = vld1q_s8($x.add(16));
8567                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
8568                asm!(
8569                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
8570                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
8571                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
8572                    w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
8573                    options(pure, nomem, nostack),
8574                );
8575                vaddq_s32(a0, a1)
8576            }};
8577        }
8578        let bp = bytes.as_ptr().add(r * gpr * 4);
8579        let vec_end = if r + 1 == rows { gpr - 4 } else { gpr };
8580        let mut accv = [vdupq_n_f32(0.0); 4];
8581        let mut gi = 0;
8582        while gi + 4 <= vec_end {
8583            let ld = vld1q_u8(bp.add(gi * 4));
8584            // Unpack ONCE — eight ±mask vectors serve all four streams.
8585            let w: [int8x16_t; 8] =
8586                core::array::from_fn(|i| vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld, iw[i]), m)));
8587            let h = vdup_n_u16(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
8588            let scf: float32x4_t;
8589            asm!(
8590                "fcvtl {o:v}.4s, {i:v}.4h",
8591                o = out(vreg) scf, i = in(vreg) h,
8592                options(pure, nomem, nostack),
8593            );
8594            for k in 0..4 {
8595                let xp = xs[k].as_ptr();
8596                let d0 = sdot2!(w[0], w[1], xp.add(gi * GROUP_SIZE));
8597                let d1 = sdot2!(w[2], w[3], xp.add((gi + 1) * GROUP_SIZE));
8598                let d2 = sdot2!(w[4], w[5], xp.add((gi + 2) * GROUP_SIZE));
8599                let d3 = sdot2!(w[6], w[7], xp.add((gi + 3) * GROUP_SIZE));
8600                let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
8601                let g = vld1q_s32(gs[k].as_ptr().add(gi));
8602                let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
8603                accv[k] = vfmaq_f32(accv[k], vcvtq_f32_s32(dots), scf);
8604            }
8605            gi += 4;
8606        }
8607        let mut acc = [
8608            vaddvq_f32(accv[0]),
8609            vaddvq_f32(accv[1]),
8610            vaddvq_f32(accv[2]),
8611            vaddvq_f32(accv[3]),
8612        ];
8613        while gi < gpr {
8614            let t = bp.add(gi * 4);
8615            let sc = f16_to_f32(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
8616            let v0 = vcombine_u8(vdup_n_u8(*t), vdup_n_u8(*t.add(1)));
8617            let v1 = vcombine_u8(vdup_n_u8(*t.add(2)), vdup_n_u8(*t.add(3)));
8618            let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
8619            let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
8620            for k in 0..4 {
8621                let d = vaddvq_s32(sdot2!(w0, w1, xs[k].as_ptr().add(gi * GROUP_SIZE)));
8622                acc[k] += (-(2 * d + *gs[k].as_ptr().add(gi))) as f32 * sc;
8623            }
8624            gi += 1;
8625        }
8626        acc
8627    }
8628}
8629
8630/// (weight ±1, scale) of one q1g element — the exact outlier term.
8631#[inline]
8632fn q1g_outlier(bytes: &[u8], sc_off: usize, r: usize, gpr: usize, j: usize) -> (f32, f32) {
8633    let t = r * gpr + j / GROUP_SIZE;
8634    let k = j % GROUP_SIZE;
8635    let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, t));
8636    let bit = (bytes[t * 4 + k / 8] >> (k % 8)) & 1;
8637    ((bit as i32 * 2 - 1) as f32, s)
8638}
8639
8640/// Exact scalar q1g row (CMF_SDOT=0 contract) — `q1_row_exact`'s order.
8641#[inline]
8642fn q1g_row_exact(bytes: &[u8], sc_off: usize, r: usize, gpr: usize, x: &[f32]) -> f32 {
8643    let mut acc = 0f32;
8644    for gi in 0..gpr {
8645        let t = r * gpr + gi;
8646        let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, t));
8647        let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
8648        let mut ga = 0f32;
8649        for (j, &b) in bytes[t * 4..t * 4 + 4].iter().enumerate() {
8650            for k in 0..8 {
8651                ga += (((b >> k) & 1) as f32 * 2.0 - 1.0) * xg[j * 8 + k];
8652            }
8653        }
8654        acc += ga * s;
8655    }
8656    acc
8657}
8658
8659/// q1t per-row overlay locator. After the base (`base_len`) come
8660/// `[u32 row_ptr[rows+1]]` then `[(u16 col, f16 val)]` grouped by row (row
8661/// `r`'s entries are `[row_ptr[r], row_ptr[r+1])`). Returns
8662/// `(row_ptr offset, entries offset, present)`.
8663fn q1t_overlay(bytes: &[u8], base_len: usize, rows: usize) -> (usize, usize, bool) {
8664    let entries = base_len + (rows + 1) * 4;
8665    (base_len, entries, entries <= bytes.len())
8666}
8667
8668/// Read `row_ptr[r]` from the overlay's prefix-sum table.
8669#[inline]
8670fn q1t_rowptr(bytes: &[u8], rp_off: usize, r: usize) -> usize {
8671    let o = rp_off + r * 4;
8672    u32::from_le_bytes([bytes[o], bytes[o + 1], bytes[o + 2], bytes[o + 3]]) as usize
8673}
8674
8675/// Byte → the 5 ternary signs it packs `{−1,0,+1}` as f32, precomputed so
8676/// decoding a q1t code is a table load, not the base-3 divide/modulo per
8677/// weight (division is ~20–40× the cost of a load). Built at compile time.
8678const SIGN5: [[f32; 5]; 256] = {
8679    let mut lut = [[0.0f32; 5]; 256];
8680    let pow3 = [1u16, 3, 9, 27, 81];
8681    let mut byte = 0usize;
8682    while byte < 256 {
8683        let mut i = 0usize;
8684        while i < 5 {
8685            let code = (byte as u16 / pow3[i]) % 3;
8686            lut[byte][i] = if code == 1 {
8687                1.0
8688            } else if code == 2 {
8689                -1.0
8690            } else {
8691                0.0
8692            };
8693            i += 1;
8694        }
8695        byte += 1;
8696    }
8697    lut
8698};
8699
8700/// Same table, as i8 signs — the operand for the int8 SDOT base kernel.
8701const SIGN5_I8: [[i8; 5]; 256] = {
8702    let mut lut = [[0i8; 5]; 256];
8703    let pow3 = [1u16, 3, 9, 27, 81];
8704    let mut byte = 0usize;
8705    while byte < 256 {
8706        let mut i = 0usize;
8707        while i < 5 {
8708            let code = (byte as u16 / pow3[i]) % 3;
8709            lut[byte][i] = if code == 1 {
8710                1
8711            } else if code == 2 {
8712                -1
8713            } else {
8714                0
8715            };
8716            i += 1;
8717        }
8718        byte += 1;
8719    }
8720    lut
8721};
8722
8723/// The same 5 i8 signs packed into a u64 (`[s0 s1 s2 s3 s4 0 0 0]`, LE) so the
8724/// group unpack is 7 unaligned u64 stores at offsets 0,5,10,…,30 instead of
8725/// six 5-byte copies + LUT indexing — each store's trailing zeros are fixed by
8726/// the next store, and the last one runs 6 B past the 32nd weight (the unpack
8727/// buffer is padded to 40). This is the decode/prefill hot inner op.
8728const SIGN5_U64: [u64; 256] = {
8729    let mut lut = [0u64; 256];
8730    let pow3 = [1u16, 3, 9, 27, 81];
8731    let mut byte = 0usize;
8732    while byte < 256 {
8733        let mut v = 0u64;
8734        let mut i = 0usize;
8735        while i < 5 {
8736            let code = (byte as u16 / pow3[i]) % 3;
8737            let s: u8 = if code == 1 {
8738                1
8739            } else if code == 2 {
8740                0xFF
8741            } else {
8742                0
8743            };
8744            v |= (s as u64) << (i * 8);
8745            i += 1;
8746        }
8747        lut[byte] = v;
8748        byte += 1;
8749    }
8750    lut
8751};
8752
8753/// Ternary base weight at `(row r, col j)` = `sign(code)·s_group`. Used to add
8754/// back activation-outlier columns, whose `x` was zeroed for the int8 bulk dot
8755/// (`split_act`). At a weight-outlier position the code is 0, so this is 0 and
8756/// the overlay correction owns that column — no double counting.
8757#[inline]
8758fn q1t_base_weight(bytes: &[u8], r: usize, gpr: usize, j: usize) -> f32 {
8759    const TILE: usize = cortiq_core::quant::Q1T_TILE;
8760    let off = (r * gpr + j / GROUP_SIZE) * TILE;
8761    let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
8762    let within = j % GROUP_SIZE;
8763    SIGN5[bytes[off + 2 + within / 5] as usize][within % 5] * s
8764}
8765
8766/// One 32-group int8 dot via two SDOTs. Bit-exact vs the scalar i8 sum
8767/// (integer accumulation is order-independent).
8768#[cfg(target_arch = "aarch64")]
8769#[target_feature(enable = "neon,dotprod")]
8770#[inline]
8771unsafe fn sdot32_i8(w: *const i8, x: *const i8) -> i32 {
8772    // SAFETY: caller guarantees 32 readable i8 at each pointer.
8773    unsafe {
8774        use core::arch::aarch64::*;
8775        use core::arch::asm;
8776        let w0 = vld1q_s8(w);
8777        let w1 = vld1q_s8(w.add(16));
8778        let x0 = vld1q_s8(x);
8779        let x1 = vld1q_s8(x.add(16));
8780        let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
8781        asm!(
8782            "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
8783            "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
8784            a0 = inout(vreg) a0, a1 = inout(vreg) a1,
8785            w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
8786            options(pure, nomem, nostack),
8787        );
8788        vaddvq_s32(vaddq_s32(a0, a1))
8789    }
8790}
8791
8792/// One 32-group int8 dot via AVX2: signed·signed as `maddubs(|w|, sign(x,w))`
8793/// then `madd` and a horizontal reduce (the same idiom as `dot_q4t_row_avx2`).
8794#[cfg(target_arch = "x86_64")]
8795#[target_feature(enable = "avx2")]
8796#[inline]
8797unsafe fn i8dot32_avx2(w: *const i8, x: *const i8) -> i32 {
8798    // SAFETY: caller guarantees 32 readable i8 at each pointer.
8799    unsafe {
8800        use core::arch::x86_64::*;
8801        let wv = _mm256_loadu_si256(w as *const __m256i);
8802        let xv = _mm256_loadu_si256(x as *const __m256i);
8803        let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
8804        let d = _mm256_madd_epi16(p16, _mm256_set1_epi16(1));
8805        let hi128 = _mm256_extracti128_si256::<1>(d);
8806        let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
8807        let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8808        let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8809        _mm_cvtsi128_si32(s32)
8810    }
8811}
8812
8813/// Unpack one q1t group's base-3 codes into 32 i8 signs via 7 unaligned u64
8814/// stores (see `SIGN5_U64`). `dst` MUST have ≥ 40 bytes: the 7th store writes
8815/// `dst[30..38]`. Stores go in order so each one's trailing zeros are
8816/// overwritten by the next; the final 6 padding bytes are unused by the dot.
8817#[inline]
8818fn q1t_unpack_group_i8(codes: *const u8, dst: &mut [i8]) {
8819    debug_assert!(dst.len() >= 40);
8820    // SAFETY: codes points at 7 readable bytes; dst has ≥ 40 bytes so every
8821    // 8-byte store at offset bi*5 (bi ≤ 6 → ≤ 30) stays in bounds.
8822    unsafe {
8823        let p = dst.as_mut_ptr();
8824        for bi in 0..7 {
8825            core::ptr::write_unaligned(
8826                p.add(bi * 5) as *mut u64,
8827                SIGN5_U64[*codes.add(bi) as usize],
8828            );
8829        }
8830    }
8831}
8832
8833/// One 32-group int8 dot, arch-dispatched (the matmat inner loop, where the
8834/// row's signs are unpacked once and dotted against every batch input).
8835/// Callers are gated by `a8w8_enabled()`, so the target-feature arms are
8836/// reachable; the scalar arm is a non-SIMD-arch fallback.
8837#[inline]
8838fn q1t_i8dot32(w: *const i8, x: *const i8) -> i32 {
8839    #[cfg(target_arch = "aarch64")]
8840    unsafe {
8841        return sdot32_i8(w, x);
8842    }
8843    #[cfg(target_arch = "x86_64")]
8844    unsafe {
8845        return i8dot32_avx2(w, x);
8846    }
8847    #[allow(unreachable_code)]
8848    unsafe {
8849        let mut s = 0i32;
8850        for k in 0..GROUP_SIZE {
8851            s += *w.add(k) as i32 * *x.add(k) as i32;
8852        }
8853        s
8854    }
8855}
8856
8857#[inline]
8858unsafe fn q1t_unpack_reg_u64s(codes: *const u8) -> (u64, u64, u64, u64) {
8859    let (s0, s1, s2, s3, s4, s5, s6) = unsafe {
8860        (
8861            SIGN5_U64[*codes as usize],
8862            SIGN5_U64[*codes.add(1) as usize],
8863            SIGN5_U64[*codes.add(2) as usize],
8864            SIGN5_U64[*codes.add(3) as usize],
8865            SIGN5_U64[*codes.add(4) as usize],
8866            SIGN5_U64[*codes.add(5) as usize],
8867            SIGN5_U64[*codes.add(6) as usize],
8868        )
8869    };
8870
8871    let u0 = s0 | (s1 << 40);
8872    let u1 = (s1 >> 24) | (s2 << 16) | (s3 << 56);
8873    let u2 = (s3 >> 8) | (s4 << 32);
8874    let u3 = (s4 >> 32) | (s5 << 8) | (s6 << 48);
8875
8876    (u0, u1, u2, u3)
8877}
8878
8879/// One q1t row's int8 base dot: `Σ_group s·dot(signs, xq)` (before the shared
8880/// `sx`). Direct register unpacking (zero stack stores/loads, no STLF stalls).
8881/// ARM SDOT.
8882#[cfg(target_arch = "aarch64")]
8883#[target_feature(enable = "neon,dotprod")]
8884unsafe fn q1t_dot_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
8885    use core::arch::aarch64::*;
8886    use core::arch::asm;
8887    unsafe {
8888        const TILE: usize = cortiq_core::quant::Q1T_TILE;
8889        let mut acc = 0f32;
8890        let bytes_ptr = bytes.as_ptr();
8891        let xq_ptr = xq.as_ptr();
8892        let row_off = r * gpr * TILE;
8893
8894        let gpr2 = gpr & !1;
8895        let mut gi = 0;
8896        while gi < gpr2 {
8897            let off0 = row_off + gi * TILE;
8898            let off1 = off0 + TILE;
8899            let s0 = f16_to_f32(u16::from_le_bytes([
8900                *bytes_ptr.add(off0),
8901                *bytes_ptr.add(off0 + 1),
8902            ]));
8903            let s1 = f16_to_f32(u16::from_le_bytes([
8904                *bytes_ptr.add(off1),
8905                *bytes_ptr.add(off1 + 1),
8906            ]));
8907
8908            let (u0_0, u1_0, u2_0, u3_0) = q1t_unpack_reg_u64s(bytes_ptr.add(off0 + 2));
8909            let (u0_1, u1_1, u2_1, u3_1) = q1t_unpack_reg_u64s(bytes_ptr.add(off1 + 2));
8910
8911            let w0_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_0), vcreate_u64(u1_0)));
8912            let w1_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_0), vcreate_u64(u3_0)));
8913            let w0_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_1), vcreate_u64(u1_1)));
8914            let w1_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_1), vcreate_u64(u3_1)));
8915
8916            let x0_0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE));
8917            let x1_0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE + 16));
8918            let x0_1 = vld1q_s8(xq_ptr.add((gi + 1) * GROUP_SIZE));
8919            let x1_1 = vld1q_s8(xq_ptr.add((gi + 1) * GROUP_SIZE + 16));
8920
8921            let (mut a0_0, mut a1_0) = (vdupq_n_s32(0), vdupq_n_s32(0));
8922            let (mut a0_1, mut a1_1) = (vdupq_n_s32(0), vdupq_n_s32(0));
8923            asm!(
8924                "sdot {a0_0:v}.4s, {w0_0:v}.16b, {x0_0:v}.16b",
8925                "sdot {a1_0:v}.4s, {w1_0:v}.16b, {x1_0:v}.16b",
8926                "sdot {a0_1:v}.4s, {w0_1:v}.16b, {x0_1:v}.16b",
8927                "sdot {a1_1:v}.4s, {w1_1:v}.16b, {x1_1:v}.16b",
8928                a0_0 = inout(vreg) a0_0, a1_0 = inout(vreg) a1_0,
8929                a0_1 = inout(vreg) a0_1, a1_1 = inout(vreg) a1_1,
8930                w0_0 = in(vreg) w0_0, x0_0 = in(vreg) x0_0, w1_0 = in(vreg) w1_0, x1_0 = in(vreg) x1_0,
8931                w0_1 = in(vreg) w0_1, x0_1 = in(vreg) x0_1, w1_1 = in(vreg) w1_1, x1_1 = in(vreg) x1_1,
8932                options(pure, nomem, nostack),
8933            );
8934            let d0 = vaddvq_s32(vaddq_s32(a0_0, a1_0));
8935            let d1 = vaddvq_s32(vaddq_s32(a0_1, a1_1));
8936            acc += d0 as f32 * s0 + d1 as f32 * s1;
8937            gi += 2;
8938        }
8939
8940        if gi < gpr {
8941            let off = row_off + gi * TILE;
8942            let s = f16_to_f32(u16::from_le_bytes([
8943                *bytes_ptr.add(off),
8944                *bytes_ptr.add(off + 1),
8945            ]));
8946            let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
8947            let w0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0), vcreate_u64(u1)));
8948            let w1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2), vcreate_u64(u3)));
8949            let x0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE));
8950            let x1 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE + 16));
8951            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
8952            asm!(
8953                "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
8954                "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
8955                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
8956                w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
8957                options(pure, nomem, nostack),
8958            );
8959            let d = vaddvq_s32(vaddq_s32(a0, a1));
8960            acc += d as f32 * s;
8961        }
8962        acc
8963    }
8964}
8965
8966/// x86 AVX2 mirror of `q1t_dot_row_sdot` (maddubs int8 dot per group).
8967#[cfg(target_arch = "x86_64")]
8968#[target_feature(enable = "avx2")]
8969unsafe fn q1t_dot_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
8970    use core::arch::x86_64::*;
8971    unsafe {
8972        const TILE: usize = cortiq_core::quant::Q1T_TILE;
8973        let mut acc = 0f32;
8974        let bytes_ptr = bytes.as_ptr();
8975        let xq_ptr = xq.as_ptr();
8976        let row_off = r * gpr * TILE;
8977
8978        let ones = _mm256_set1_epi16(1);
8979        for gi in 0..gpr {
8980            let off = row_off + gi * TILE;
8981            let s = f16_to_f32(u16::from_le_bytes([
8982                *bytes_ptr.add(off),
8983                *bytes_ptr.add(off + 1),
8984            ]));
8985            let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
8986            let wv = _mm256_set_epi64x(u3 as i64, u2 as i64, u1 as i64, u0 as i64);
8987            let xv = _mm256_loadu_si256(xq_ptr.add(gi * GROUP_SIZE) as *const __m256i);
8988            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
8989            let d256 = _mm256_madd_epi16(p16, ones);
8990            let d128 = _mm_add_epi32(
8991                _mm256_castsi256_si128(d256),
8992                _mm256_extracti128_si256(d256, 1),
8993            );
8994            let d64 = _mm_add_epi32(d128, _mm_shuffle_epi32(d128, 0xee));
8995            let d32 = _mm_cvtsi128_si32(_mm_add_epi32(d64, _mm_shuffle_epi32(d64, 0x55)));
8996            acc += d32 as f32 * s;
8997        }
8998        acc
8999    }
9000}
9001
9002/// VNNI twin of `q1t_dot_row_avx2` (see `dpbusd_hsum`).
9003#[cfg(target_arch = "x86_64")]
9004#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
9005unsafe fn q1t_dot_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
9006    use core::arch::x86_64::*;
9007    // SAFETY: same tile/xq contracts as `q1t_dot_row_avx2`.
9008    unsafe {
9009        const TILE: usize = cortiq_core::quant::Q1T_TILE;
9010        let mut acc = 0f32;
9011        let bytes_ptr = bytes.as_ptr();
9012        let xq_ptr = xq.as_ptr();
9013        let row_off = r * gpr * TILE;
9014        for gi in 0..gpr {
9015            let off = row_off + gi * TILE;
9016            let s = f16_to_f32(u16::from_le_bytes([
9017                *bytes_ptr.add(off),
9018                *bytes_ptr.add(off + 1),
9019            ]));
9020            let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
9021            let wv = _mm256_set_epi64x(u3 as i64, u2 as i64, u1 as i64, u0 as i64);
9022            let xv = _mm256_loadu_si256(xq_ptr.add(gi * GROUP_SIZE) as *const __m256i);
9023            let d = dpbusd_hsum(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
9024            acc += d as f32 * s;
9025        }
9026        acc
9027    }
9028}
9029
9030/// Per-row int8 base dot, dispatched once per row (matvec decode hot path).
9031/// Callers are gated by `a8w8_enabled()`, so the target-feature kernels are
9032/// reachable.
9033#[inline]
9034fn q1t_dot_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
9035    #[cfg(target_arch = "aarch64")]
9036    unsafe {
9037        return q1t_dot_row_sdot(bytes, r, gpr, xq);
9038    }
9039    #[cfg(target_arch = "x86_64")]
9040    unsafe {
9041        if vnni_tiles_enabled() {
9042            return q1t_dot_row_vnni(bytes, r, gpr, xq);
9043        }
9044        return q1t_dot_row_avx2(bytes, r, gpr, xq);
9045    }
9046    #[allow(unreachable_code)]
9047    {
9048        const TILE: usize = cortiq_core::quant::Q1T_TILE;
9049        let mut acc = 0f32;
9050        let mut sg = [0i8; GROUP_SIZE + 8]; // +8 slack for the u64-store unpack
9051        for gi in 0..gpr {
9052            let off = (r * gpr + gi) * TILE;
9053            let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
9054            q1t_unpack_group_i8(bytes.as_ptr().wrapping_add(off + 2), &mut sg);
9055            let mut d = 0i32;
9056            for k in 0..GROUP_SIZE {
9057                d += sg[k] as i32 * xq[gi * GROUP_SIZE + k] as i32;
9058            }
9059            acc += d as f32 * s;
9060        }
9061        acc
9062    }
9063}
9064
9065/// Σ over a row's outliers of `value·x[col]` — the correction that adds the
9066/// overlay's exact weights on top of the base dot. INVARIANT: the encoder
9067/// writes ternary code 0 at every outlier position (`quantize_q1t`), so the
9068/// base contributes nothing there and this is a plain `value·x`, not
9069/// `(value − base)·x` — no scattered per-outlier scale read. Row `r`'s entries
9070/// are the contiguous slice `[row_ptr[r], row_ptr[r+1])`, so no binary search.
9071fn q1t_row_outlier_correction(
9072    bytes: &[u8],
9073    r: usize,
9074    rp_off: usize,
9075    entries_off: usize,
9076    has_ov: bool,
9077    x: &[f32],
9078) -> f32 {
9079    if !has_ov {
9080        return 0.0;
9081    }
9082    let (c0, c1) = (
9083        q1t_rowptr(bytes, rp_off, r),
9084        q1t_rowptr(bytes, rp_off, r + 1),
9085    );
9086    let mut corr = 0f32;
9087    for p in c0..c1 {
9088        let e = entries_off + p * 4;
9089        let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
9090        let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
9091        corr += val * x[col];
9092    }
9093    corr
9094}
9095
9096/// Dequantize one q1t row into `buf[..cols]` via the sign LUT (no division),
9097/// then apply the row's outliers (its `[row_ptr[r], row_ptr[r+1])` slice).
9098/// Used by the batched (prefill) path where the decode amortizes over the batch.
9099fn q1t_dequant_row(
9100    bytes: &[u8],
9101    r: usize,
9102    gpr: usize,
9103    rp_off: usize,
9104    entries_off: usize,
9105    has_ov: bool,
9106    buf: &mut [f32],
9107) {
9108    const TILE: usize = cortiq_core::quant::Q1T_TILE;
9109    for g in 0..gpr {
9110        let off = (r * gpr + g) * TILE;
9111        let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
9112        let codes = &bytes[off + 2..off + TILE];
9113        let bc = g * GROUP_SIZE;
9114        // 6 full bytes (30 codes) + a 7th byte holding the last 2.
9115        for bi in 0..6 {
9116            let lut = &SIGN5[codes[bi] as usize];
9117            let d = &mut buf[bc + bi * 5..bc + bi * 5 + 5];
9118            for i in 0..5 {
9119                d[i] = lut[i] * s;
9120            }
9121        }
9122        let lut = &SIGN5[codes[6] as usize];
9123        buf[bc + 30] = lut[0] * s;
9124        buf[bc + 31] = lut[1] * s;
9125    }
9126    if !has_ov {
9127        return;
9128    }
9129    let (c0, c1) = (
9130        q1t_rowptr(bytes, rp_off, r),
9131        q1t_rowptr(bytes, rp_off, r + 1),
9132    );
9133    for p in c0..c1 {
9134        let e = entries_off + p * 4;
9135        let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
9136        buf[col] = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
9137    }
9138}
9139
9140/// Add the sparse outlier overlay onto a base dot already in `out` (the GPU
9141/// computes the ternary base; the overlay stays on the CPU — its entries are
9142/// few and its per-row gather doesn't vectorize on the GPU). Row-parallel.
9143fn q1t_add_overlay(
9144    bytes: &[u8],
9145    x: &[f32],
9146    rows: usize,
9147    cols: usize,
9148    out: &mut [f32],
9149    pool: Option<&Pool>,
9150) {
9151    const TILE: usize = cortiq_core::quant::Q1T_TILE;
9152    let gpr = cols / GROUP_SIZE;
9153    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
9154    if !has_ov {
9155        return;
9156    }
9157    let out_addr = SendMut(out.as_mut_ptr());
9158    let run = move |start: usize, end: usize| {
9159        for r in start..end {
9160            let corr = q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
9161            // SAFETY: disjoint rows; add onto the base the GPU already wrote.
9162            unsafe { *out_addr.at(r) += corr };
9163        }
9164    };
9165    dispatch_rows(pool, rows, &run);
9166}
9167
9168/// Q1T row range via the A8W8 int8 path — shared activation split,
9169/// per-row: base SDOT dot + outlier correction + overlay.
9170#[allow(clippy::too_many_arguments)]
9171fn q1t_range_a8w8(
9172    bytes: &[u8],
9173    gpr: usize,
9174    rp_off: usize,
9175    ent_off: usize,
9176    has_ov: bool,
9177    act: &SplitAct,
9178    x: &[f32],
9179    out: SendMut,
9180    start: usize,
9181    end: usize,
9182) {
9183    for r in start..end {
9184        let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
9185        for &(j, xv) in &act.outliers {
9186            acc += q1t_base_weight(bytes, r, gpr, j) * xv;
9187        }
9188        acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
9189        // SAFETY: disjoint row ranges per worker.
9190        unsafe { *out.at(r) = acc };
9191    }
9192}
9193
9194/// Q1T row range via the f32 path (no SDOT) — for matvec_many batched
9195/// dispatch when a8w8 is unavailable.
9196#[allow(clippy::too_many_arguments)]
9197fn q1t_range_f32_batch(
9198    bytes: &[u8],
9199    gpr: usize,
9200    rp_off: usize,
9201    ent_off: usize,
9202    has_ov: bool,
9203    x: &[f32],
9204    out: SendMut,
9205    start: usize,
9206    end: usize,
9207) {
9208    const TILE: usize = cortiq_core::quant::Q1T_TILE;
9209    let mut sg = [0f32; GROUP_SIZE];
9210    for r in start..end {
9211        let mut acc = 0f32;
9212        for g in 0..gpr {
9213            let off = (r * gpr + g) * TILE;
9214            let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
9215            let codes = &bytes[off + 2..off + TILE];
9216            let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
9217            for bi in 0..6 {
9218                sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
9219            }
9220            let lut = &SIGN5[codes[6] as usize];
9221            sg[30] = lut[0];
9222            sg[31] = lut[1];
9223            let mut gsum = 0f32;
9224            for k in 0..GROUP_SIZE {
9225                gsum += sg[k] * xg[k];
9226            }
9227            acc += s * gsum;
9228        }
9229        acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
9230        // SAFETY: disjoint row ranges per worker.
9231        unsafe { *out.at(r) = acc };
9232    }
9233}
9234
9235/// Ternary (q1t) matvec — decode+dot straight from mmap, one group at a time:
9236/// no per-ROW buffer, no division (the sign LUT), and a tiny per-group sign
9237/// buffer so the 32-wide dot vectorizes. This is the decode hot path.
9238fn q1t_matvec(
9239    bytes: &[u8],
9240    x: &[f32],
9241    rows: usize,
9242    cols: usize,
9243    out: &mut [f32],
9244    pool: Option<&Pool>,
9245) {
9246    debug_assert_eq!(out.len(), rows);
9247    const TILE: usize = cortiq_core::quant::Q1T_TILE;
9248    let gpr = cols / GROUP_SIZE;
9249    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
9250    let out_addr = SendMut(out.as_mut_ptr());
9251    // int8 SDOT base dot (ARM dotprod): ~4× the f32 arithmetic. x → i8 once
9252    // (`split_act`), activation outliers added back exactly in f32, weight
9253    // overlay on top. ARM SDOT / x86 AVX2; CMF_SDOT=0 keeps the exact f32 path.
9254    if a8w8_enabled() {
9255        let act = split_act(x);
9256        let act = &act;
9257        let run = move |start: usize, end: usize| {
9258            for r in start..end {
9259                let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
9260                for &(j, xv) in &act.outliers {
9261                    acc += q1t_base_weight(bytes, r, gpr, j) * xv;
9262                }
9263                acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
9264                // SAFETY: disjoint row ranges per worker.
9265                unsafe { *out_addr.at(r) = acc };
9266            }
9267        };
9268        dispatch_rows(pool, rows, &run);
9269        return;
9270    }
9271    let run = move |start: usize, end: usize| {
9272        // Per-group signs, unpacked contiguously so the dot below is a clean
9273        // 32-wide reduction the autovectorizer turns into f32x4 FMAs — the
9274        // 5-values-per-byte base-3 layout won't SIMD in place.
9275        let mut sg = [0f32; GROUP_SIZE];
9276        for r in start..end {
9277            let mut acc = 0f32;
9278            for g in 0..gpr {
9279                let off = (r * gpr + g) * TILE;
9280                let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
9281                let codes = &bytes[off + 2..off + TILE];
9282                let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
9283                for bi in 0..6 {
9284                    sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
9285                }
9286                let lut = &SIGN5[codes[6] as usize];
9287                sg[30] = lut[0];
9288                sg[31] = lut[1];
9289                let mut gsum = 0f32;
9290                for k in 0..GROUP_SIZE {
9291                    gsum += sg[k] * xg[k];
9292                }
9293                acc += s * gsum;
9294            }
9295            acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
9296            unsafe { *out_addr.at(r) = acc };
9297        }
9298    };
9299    dispatch_rows(pool, rows, &run);
9300}
9301
9302/// Fused-pair twin of `q1t_dot_row_sdot`: ONE register unpack of the
9303/// ternary codes serves BOTH activation streams (the unpack chain is
9304/// the dominant per-row cost — MTP verify pairs paid it twice). Per
9305/// stream the group order and f32 accumulation match the single-row
9306/// kernel exactly, so pair == 2×matvec bit-for-bit.
9307#[cfg(target_arch = "aarch64")]
9308#[target_feature(enable = "neon,dotprod")]
9309unsafe fn q1t_dot_row_sdot2(bytes: &[u8], r: usize, gpr: usize, xa: &[i8], xb: &[i8]) -> [f32; 2] {
9310    use core::arch::aarch64::*;
9311    use core::arch::asm;
9312    // SAFETY: same slice-length contracts as `q1t_dot_row_sdot`, ×2.
9313    unsafe {
9314        const TILE: usize = cortiq_core::quant::Q1T_TILE;
9315        let bytes_ptr = bytes.as_ptr();
9316        let row_off = r * gpr * TILE;
9317        let xp = [xa.as_ptr(), xb.as_ptr()];
9318        let mut acc = [0f32; 2];
9319        macro_rules! sdot2 {
9320            ($w0:expr, $w1:expr, $x:expr) => {{
9321                let x0 = vld1q_s8($x);
9322                let x1 = vld1q_s8($x.add(16));
9323                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
9324                asm!(
9325                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
9326                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
9327                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
9328                    w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
9329                    options(pure, nomem, nostack),
9330                );
9331                vaddvq_s32(vaddq_s32(a0, a1))
9332            }};
9333        }
9334        let gpr2 = gpr & !1;
9335        let mut gi = 0;
9336        while gi < gpr2 {
9337            let off0 = row_off + gi * TILE;
9338            let off1 = off0 + TILE;
9339            let s0 = f16_to_f32(u16::from_le_bytes([
9340                *bytes_ptr.add(off0),
9341                *bytes_ptr.add(off0 + 1),
9342            ]));
9343            let s1 = f16_to_f32(u16::from_le_bytes([
9344                *bytes_ptr.add(off1),
9345                *bytes_ptr.add(off1 + 1),
9346            ]));
9347            let (u0_0, u1_0, u2_0, u3_0) = q1t_unpack_reg_u64s(bytes_ptr.add(off0 + 2));
9348            let (u0_1, u1_1, u2_1, u3_1) = q1t_unpack_reg_u64s(bytes_ptr.add(off1 + 2));
9349            let w0_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_0), vcreate_u64(u1_0)));
9350            let w1_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_0), vcreate_u64(u3_0)));
9351            let w0_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_1), vcreate_u64(u1_1)));
9352            let w1_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_1), vcreate_u64(u3_1)));
9353            for k in 0..2 {
9354                let d0 = sdot2!(w0_0, w1_0, xp[k].add(gi * GROUP_SIZE));
9355                let d1 = sdot2!(w0_1, w1_1, xp[k].add((gi + 1) * GROUP_SIZE));
9356                acc[k] += d0 as f32 * s0 + d1 as f32 * s1;
9357            }
9358            gi += 2;
9359        }
9360        if gi < gpr {
9361            let off = row_off + gi * TILE;
9362            let s = f16_to_f32(u16::from_le_bytes([
9363                *bytes_ptr.add(off),
9364                *bytes_ptr.add(off + 1),
9365            ]));
9366            let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
9367            let w0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0), vcreate_u64(u1)));
9368            let w1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2), vcreate_u64(u3)));
9369            for k in 0..2 {
9370                let d = sdot2!(w0, w1, xp[k].add(gi * GROUP_SIZE));
9371                acc[k] += d as f32 * s;
9372            }
9373        }
9374        acc
9375    }
9376}
9377
9378/// Fused Q1T pair matvec: ONE pass over the rows serves both
9379/// activation streams — on ARM the ternary register unpack happens
9380/// once per tile pair (`q1t_dot_row_sdot2`); elsewhere the second dot
9381/// rides the row's L1-warm tile bytes. Per stream the math matches
9382/// `q1t_matvec` exactly.
9383fn q1t_matvec2(
9384    bytes: &[u8],
9385    x1: &[f32],
9386    x2: &[f32],
9387    rows: usize,
9388    cols: usize,
9389    o1: &mut [f32],
9390    o2: &mut [f32],
9391    pool: Option<&Pool>,
9392) {
9393    debug_assert_eq!(o1.len(), rows);
9394    debug_assert_eq!(o2.len(), rows);
9395    const TILE: usize = cortiq_core::quant::Q1T_TILE;
9396    let gpr = cols / GROUP_SIZE;
9397    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
9398    let out1 = SendMut(o1.as_mut_ptr());
9399    let out2 = SendMut(o2.as_mut_ptr());
9400    if a8w8_enabled() {
9401        let a1 = split_act(x1);
9402        let a2 = split_act(x2);
9403        let (a1, a2) = (&a1, &a2);
9404        let run = move |start: usize, end: usize| {
9405            for r in start..end {
9406                #[cfg(target_arch = "aarch64")]
9407                // a8w8 on aarch64 ⇔ sdot_enabled(), so the kernel's
9408                // target features are present.
9409                let ds = unsafe { q1t_dot_row_sdot2(bytes, r, gpr, &a1.xq, &a2.xq) };
9410                #[cfg(not(target_arch = "aarch64"))]
9411                let ds = [
9412                    q1t_dot_row_i8(bytes, r, gpr, &a1.xq),
9413                    q1t_dot_row_i8(bytes, r, gpr, &a2.xq),
9414                ];
9415                let mut acc1 = ds[0] * a1.sx;
9416                for &(j, xv) in &a1.outliers {
9417                    acc1 += q1t_base_weight(bytes, r, gpr, j) * xv;
9418                }
9419                acc1 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x1);
9420                let mut acc2 = ds[1] * a2.sx;
9421                for &(j, xv) in &a2.outliers {
9422                    acc2 += q1t_base_weight(bytes, r, gpr, j) * xv;
9423                }
9424                acc2 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x2);
9425                // SAFETY: disjoint row ranges per worker.
9426                unsafe {
9427                    *out1.at(r) = acc1;
9428                    *out2.at(r) = acc2;
9429                }
9430            }
9431        };
9432        dispatch_rows(pool, rows, &run);
9433        return;
9434    }
9435    let run = move |start: usize, end: usize| {
9436        // Exact path (CMF_SDOT=0): unpack the sign LUT once per group,
9437        // dot both streams — same op order per stream as `q1t_matvec`.
9438        let mut sg = [0f32; GROUP_SIZE];
9439        for r in start..end {
9440            let mut acc1 = 0f32;
9441            let mut acc2 = 0f32;
9442            for g in 0..gpr {
9443                let off = (r * gpr + g) * TILE;
9444                let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
9445                let codes = &bytes[off + 2..off + TILE];
9446                for bi in 0..6 {
9447                    sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
9448                }
9449                let lut = &SIGN5[codes[6] as usize];
9450                sg[30] = lut[0];
9451                sg[31] = lut[1];
9452                let xg1 = &x1[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
9453                let xg2 = &x2[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
9454                let mut gsum1 = 0f32;
9455                for k in 0..GROUP_SIZE {
9456                    gsum1 += sg[k] * xg1[k];
9457                }
9458                acc1 += s * gsum1;
9459                let mut gsum2 = 0f32;
9460                for k in 0..GROUP_SIZE {
9461                    gsum2 += sg[k] * xg2[k];
9462                }
9463                acc2 += s * gsum2;
9464            }
9465            acc1 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x1);
9466            acc2 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x2);
9467            // SAFETY: disjoint row ranges per worker.
9468            unsafe {
9469                *out1.at(r) = acc1;
9470                *out2.at(r) = acc2;
9471            }
9472        }
9473    };
9474    dispatch_rows(pool, rows, &run);
9475}
9476
9477/// Ternary (q1t) matmat (prefill) — dequant each row once, dot the whole
9478/// batch against it (amortizes the per-row decode).
9479fn q1t_matmat(
9480    bytes: &[u8],
9481    xs: &[f32],
9482    b: usize,
9483    rows: usize,
9484    cols: usize,
9485    out: &mut [f32],
9486    pool: Option<&Pool>,
9487) {
9488    debug_assert_eq!(out.len(), b * rows);
9489    const TILE: usize = cortiq_core::quant::Q1T_TILE;
9490    let gpr = cols / GROUP_SIZE;
9491    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
9492    let out_addr = SendMut(out.as_mut_ptr());
9493    // int8 prefill (ARM SDOT / x86 AVX2): quantize the B inputs once, unpack
9494    // each weight row's signs to i8 ONCE, then int8-dot against every input —
9495    // the row sign-decode amortizes over the whole batch. CMF_SDOT=0 → f32.
9496    if a8w8_enabled() {
9497        let acts: Vec<SplitAct> = (0..b)
9498            .map(|bi| split_act(&xs[bi * cols..(bi + 1) * cols]))
9499            .collect();
9500        let acts = &acts;
9501        let run = move |start: usize, end: usize| {
9502            let mut sg = vec![0i8; cols + 8]; // row signs, i8 (+8 unpack slack)
9503            let mut sc = vec![0f32; gpr]; // per-group scales
9504            let mut accs = vec![0f32; b]; // per-batch accumulators, reused per row
9505            for r in start..end {
9506                for g in 0..gpr {
9507                    let off = (r * gpr + g) * TILE;
9508                    sc[g] = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
9509                    q1t_unpack_group_i8(
9510                        bytes.as_ptr().wrapping_add(off + 2),
9511                        &mut sg[g * GROUP_SIZE..],
9512                    );
9513                }
9514                for bi in 0..b {
9515                    let act = &acts[bi];
9516                    let mut isum = 0f32;
9517                    for g in 0..gpr {
9518                        let d = q1t_i8dot32(
9519                            sg.as_ptr().wrapping_add(g * GROUP_SIZE),
9520                            act.xq.as_ptr().wrapping_add(g * GROUP_SIZE),
9521                        );
9522                        isum += d as f32 * sc[g];
9523                    }
9524                    let mut acc = isum * act.sx;
9525                    for &(j, xv) in &act.outliers {
9526                        acc += q1t_base_weight(bytes, r, gpr, j) * xv;
9527                    }
9528                    accs[bi] = acc;
9529                }
9530                // Overlay ONCE per row for the whole batch: read each (col, val)
9531                // from mmap a single time (was b× — the re-read dominated prefill)
9532                // and fan it out over the batch via the cached inputs.
9533                if has_ov {
9534                    let (c0, c1) = (
9535                        q1t_rowptr(bytes, rp_off, r),
9536                        q1t_rowptr(bytes, rp_off, r + 1),
9537                    );
9538                    for p in c0..c1 {
9539                        let e = ent_off + p * 4;
9540                        let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
9541                        let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
9542                        for bi in 0..b {
9543                            accs[bi] += val * xs[bi * cols + col];
9544                        }
9545                    }
9546                }
9547                for bi in 0..b {
9548                    unsafe { *out_addr.at(bi * rows + r) = accs[bi] };
9549                }
9550            }
9551        };
9552        dispatch_rows(pool, rows, &run);
9553        return;
9554    }
9555    let run = move |start: usize, end: usize| {
9556        let mut buf = vec![0f32; cols];
9557        for r in start..end {
9558            q1t_dequant_row(bytes, r, gpr, rp_off, ent_off, has_ov, &mut buf);
9559            for bi in 0..b {
9560                let xr = &xs[bi * cols..(bi + 1) * cols];
9561                let mut acc = 0f32;
9562                for j in 0..cols {
9563                    acc += buf[j] * xr[j];
9564                }
9565                unsafe { *out_addr.at(bi * rows + r) = acc };
9566            }
9567        }
9568    };
9569    dispatch_rows(pool, rows, &run);
9570}
9571
9572/// q1-tile entry points of the binary drivers (the tests' handles; the
9573/// runtime goes through `QTensor::bin_w`).
9574#[cfg(test)]
9575fn q1_matvec(
9576    bytes: &[u8],
9577    x: &[f32],
9578    rows: usize,
9579    cols: usize,
9580    out: &mut [f32],
9581    pool: Option<&Pool>,
9582) {
9583    bin_matvec(BinW::Q1(bytes), x, rows, cols, out, pool)
9584}
9585
9586/// Binary-weight matvec (q1 or q1g): A8W8 rows with the shared group
9587/// sums, or the exact scalar walk under CMF_SDOT=0.
9588fn bin_matvec(w: BinW, x: &[f32], rows: usize, cols: usize, out: &mut [f32], pool: Option<&Pool>) {
9589    debug_assert_eq!(out.len(), rows);
9590    let gpr = cols / GROUP_SIZE;
9591    let out_addr = SendMut(out.as_mut_ptr());
9592    if a8w8_enabled() {
9593        let act = split_act(x);
9594        let gsum = q1_group_sums(&act.xq, gpr);
9595        let (act, gsum) = (&act, &gsum);
9596        let run =
9597            move |start: usize, end: usize| bin_range_a8w8(w, gpr, act, gsum, out_addr, start, end);
9598        dispatch_rows(pool, rows, &run);
9599        return;
9600    }
9601    let run = move |start: usize, end: usize| bin_range_f32(w, gpr, x, out_addr, start, end);
9602    dispatch_rows(pool, rows, &run);
9603}
9604
9605/// Fused two-input q1 matvec (weights read once per pair).
9606#[cfg(test)]
9607#[allow(clippy::too_many_arguments)]
9608fn q1_matvec2(
9609    bytes: &[u8],
9610    x1: &[f32],
9611    x2: &[f32],
9612    rows: usize,
9613    cols: usize,
9614    o1: &mut [f32],
9615    o2: &mut [f32],
9616    pool: Option<&Pool>,
9617) {
9618    bin_matvec2(BinW::Q1(bytes), x1, x2, rows, cols, o1, o2, pool)
9619}
9620
9621/// Fused two-input binary-weight matvec (q1 or q1g).
9622#[allow(clippy::too_many_arguments)]
9623fn bin_matvec2(
9624    w: BinW,
9625    x1: &[f32],
9626    x2: &[f32],
9627    rows: usize,
9628    cols: usize,
9629    o1: &mut [f32],
9630    o2: &mut [f32],
9631    pool: Option<&Pool>,
9632) {
9633    let gpr = cols / GROUP_SIZE;
9634    let p1 = SendMut(o1.as_mut_ptr());
9635    let p2 = SendMut(o2.as_mut_ptr());
9636    if a8w8_enabled() {
9637        let a1 = split_act(x1);
9638        let a2 = split_act(x2);
9639        let g1 = q1_group_sums(&a1.xq, gpr);
9640        let g2 = q1_group_sums(&a2.xq, gpr);
9641        let (a1, a2, g1, g2) = (&a1, &a2, &g1, &g2);
9642        let run = move |start: usize, end: usize| {
9643            for r in start..end {
9644                let mut v1 = w.dot_i8(r, gpr, &a1.xq, g1) * a1.sx;
9645                let mut v2 = w.dot_i8(r, gpr, &a2.xq, g2) * a2.sx;
9646                for &(j, xv) in &a1.outliers {
9647                    let (wv, s) = w.outlier(r, gpr, j);
9648                    v1 += wv * s * xv;
9649                }
9650                for &(j, xv) in &a2.outliers {
9651                    let (wv, s) = w.outlier(r, gpr, j);
9652                    v2 += wv * s * xv;
9653                }
9654                // SAFETY: disjoint row ranges per worker.
9655                unsafe {
9656                    *p1.at(r) = v1;
9657                    *p2.at(r) = v2;
9658                }
9659            }
9660        };
9661        dispatch_rows(pool, rows, &run);
9662        return;
9663    }
9664    let run = move |start: usize, end: usize| {
9665        for r in start..end {
9666            // SAFETY: disjoint row ranges per worker.
9667            unsafe {
9668                *p1.at(r) = w.row_exact(r, gpr, x1);
9669                *p2.at(r) = w.row_exact(r, gpr, x2);
9670            }
9671        }
9672    };
9673    dispatch_rows(pool, rows, &run);
9674}
9675
9676/// Batched q1 matmat: each row's tiles stream once per microbatch.
9677#[cfg(test)]
9678#[allow(clippy::too_many_arguments)]
9679fn q1_matmat(
9680    bytes: &[u8],
9681    xs_all: &[f32],
9682    b: usize,
9683    rows: usize,
9684    cols: usize,
9685    out: &mut [f32],
9686    pool: Option<&Pool>,
9687) {
9688    bin_matmat(BinW::Q1(bytes), xs_all, b, rows, cols, out, pool)
9689}
9690
9691/// Batched binary-weight matmat (q1 or q1g).
9692#[allow(clippy::too_many_arguments)]
9693fn bin_matmat(
9694    w: BinW,
9695    xs_all: &[f32],
9696    b: usize,
9697    rows: usize,
9698    cols: usize,
9699    out: &mut [f32],
9700    pool: Option<&Pool>,
9701) {
9702    debug_assert_eq!(out.len(), b * rows);
9703    let gpr = cols / GROUP_SIZE;
9704    let out_addr = SendMut(out.as_mut_ptr());
9705    if a8w8_enabled() {
9706        let acts: Vec<(SplitAct, Vec<i32>)> = (0..b)
9707            .map(|bi| {
9708                let act = split_act(&xs_all[bi * cols..(bi + 1) * cols]);
9709                let gsum = q1_group_sums(&act.xq, gpr);
9710                (act, gsum)
9711            })
9712            .collect();
9713        let acts = &acts;
9714        #[cfg(target_arch = "x86_64")]
9715        let blocked_ok = avx2_enabled() && blocked_enabled();
9716        #[cfg(target_arch = "aarch64")]
9717        let blocked_ok = sdot_enabled() && blocked_enabled();
9718        let run = move |start: usize, end: usize| {
9719            for r in start..end {
9720                let mut bi = 0usize;
9721                // Blocked 1×4: the unpacked bit mask serves four
9722                // activation streams per group.
9723                #[cfg(target_arch = "aarch64")]
9724                if blocked_ok {
9725                    while bi + 4 <= acts.len() {
9726                        let xs = [
9727                            acts[bi].0.xq.as_slice(),
9728                            acts[bi + 1].0.xq.as_slice(),
9729                            acts[bi + 2].0.xq.as_slice(),
9730                            acts[bi + 3].0.xq.as_slice(),
9731                        ];
9732                        let gs = [
9733                            acts[bi].1.as_slice(),
9734                            acts[bi + 1].1.as_slice(),
9735                            acts[bi + 2].1.as_slice(),
9736                            acts[bi + 3].1.as_slice(),
9737                        ];
9738                        // SAFETY: blocked_ok checked SDOT.
9739                        let d = unsafe { w.dot_1x4_sdot(r, gpr, xs, gs) };
9740                        for k in 0..4 {
9741                            let (act, _) = &acts[bi + k];
9742                            let mut acc = d[k] * act.sx;
9743                            for &(j, xv) in &act.outliers {
9744                                let (wv, sc) = w.outlier(r, gpr, j);
9745                                acc += wv * sc * xv;
9746                            }
9747                            // SAFETY: disjoint (bi, r) cells per worker.
9748                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
9749                        }
9750                        bi += 4;
9751                    }
9752                }
9753                #[cfg(target_arch = "x86_64")]
9754                if blocked_ok {
9755                    while bi + 4 <= acts.len() {
9756                        let xs = [
9757                            acts[bi].0.xq.as_slice(),
9758                            acts[bi + 1].0.xq.as_slice(),
9759                            acts[bi + 2].0.xq.as_slice(),
9760                            acts[bi + 3].0.xq.as_slice(),
9761                        ];
9762                        let gs = [
9763                            acts[bi].1.as_slice(),
9764                            acts[bi + 1].1.as_slice(),
9765                            acts[bi + 2].1.as_slice(),
9766                            acts[bi + 3].1.as_slice(),
9767                        ];
9768                        // SAFETY: blocked_ok checked AVX2.
9769                        let d = unsafe { w.dot_1x4_x86(r, gpr, xs, gs) };
9770                        for k in 0..4 {
9771                            let (act, _) = &acts[bi + k];
9772                            let mut acc = d[k] * act.sx;
9773                            for &(j, xv) in &act.outliers {
9774                                let (wv, sc) = w.outlier(r, gpr, j);
9775                                acc += wv * sc * xv;
9776                            }
9777                            // SAFETY: disjoint (bi, r) cells per worker.
9778                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
9779                        }
9780                        bi += 4;
9781                    }
9782                }
9783                while bi < acts.len() {
9784                    let (act, gsum) = &acts[bi];
9785                    let mut acc = w.dot_i8(r, gpr, &act.xq, gsum) * act.sx;
9786                    for &(j, xv) in &act.outliers {
9787                        let (wv, s) = w.outlier(r, gpr, j);
9788                        acc += wv * s * xv;
9789                    }
9790                    // SAFETY: disjoint (bi, r) cells per worker range.
9791                    unsafe { *out_addr.at(bi * rows + r) = acc };
9792                    bi += 1;
9793                }
9794            }
9795        };
9796        dispatch_rows(pool, rows, &run);
9797        return;
9798    }
9799    let run = move |start: usize, end: usize| {
9800        for r in start..end {
9801            for bi in 0..b {
9802                let x = &xs_all[bi * cols..(bi + 1) * cols];
9803                // SAFETY: disjoint (bi, r) cells per worker range.
9804                unsafe { *out_addr.at(bi * rows + r) = w.row_exact(r, gpr, x) };
9805            }
9806        }
9807    };
9808    dispatch_rows(pool, rows, &run);
9809}
9810
9811/// Fused q4_block matvec straight from the mapped bytes. SDOT path when
9812/// dotprod is available (port of vmfcore `dot_q4_block_sdot`, measured
9813/// +23% on q4 decode): nibbles → centered i8, int8×int8 `sdot` per
9814/// 32-group, exact outlier correction — the same A8W8 contract as q8.
9815/// `CMF_SDOT=0` keeps the exact scalar path.
9816fn q4matvec(
9817    bytes: &[u8],
9818    x: &[f32],
9819    rows: usize,
9820    cols: usize,
9821    out: &mut [f32],
9822    pool: Option<&Pool>,
9823) {
9824    debug_assert_eq!(out.len(), rows);
9825    let (packed, scales) = q4_split(bytes, rows, cols);
9826    let gpr = cols / GROUP_SIZE;
9827    let out_addr = SendMut(out.as_mut_ptr());
9828
9829    if a8w8_enabled() {
9830        let act = split_act(x);
9831        let run = move |start: usize, end: usize| {
9832            q4_range_a8w8(packed, scales, gpr, cols, &act, out_addr, start, end)
9833        };
9834        dispatch_rows(pool, rows, &run);
9835        return;
9836    }
9837
9838    let run =
9839        move |start: usize, end: usize| q4_range_f32(packed, scales, gpr, x, out_addr, start, end);
9840    dispatch_rows(pool, rows, &run);
9841}
9842
9843/// One q4 row via the A8W8 int8 path — SDOT on ARM, AVX2 maddubs on
9844/// x86 (scalar fallback is unreachable: callers gate on a8w8_enabled).
9845#[inline]
9846#[allow(unreachable_code)]
9847/// One UNPACKED q4 row (centered i8 in `buf`) against four activation
9848/// streams: the 32-byte weight chunk and its abs() load once per group,
9849/// the per-group f16 scale decodes once — four maddubs+reduce chains
9850/// instead of four full (load, abs, dot) rounds.
9851#[cfg(target_arch = "x86_64")]
9852#[target_feature(enable = "avx2")]
9853unsafe fn dot_q4b_row_1x4_avx2(
9854    buf: &[u8],
9855    scales: &[u8],
9856    g0: usize,
9857    gpr: usize,
9858    xs: [&[i8]; 4],
9859) -> [f32; 4] {
9860    // SAFETY: callers uphold buffer contracts (buf.len() == gpr·32).
9861    unsafe {
9862        use core::arch::x86_64::*;
9863        let ones = _mm256_set1_epi16(1);
9864        let mut acc = [0f32; 4];
9865        for gi in 0..gpr {
9866            let s = f16_to_f32(u16::from_le_bytes([
9867                scales[(g0 + gi) * 2],
9868                scales[(g0 + gi) * 2 + 1],
9869            ]));
9870            let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9871            let aw = _mm256_abs_epi8(w);
9872            for (k, xq) in xs.iter().enumerate() {
9873                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9874                let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
9875                let d = _mm256_madd_epi16(p16, ones);
9876                let hi128 = _mm256_extracti128_si256::<1>(d);
9877                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
9878                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
9879                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
9880                acc[k] += _mm_cvtsi128_si32(s32) as f32 * s;
9881            }
9882        }
9883        acc
9884    }
9885}
9886
9887/// VNNI twin of `dot_q4b_row_1x4_avx2` (see `dpbusd_hsum`).
9888#[cfg(target_arch = "x86_64")]
9889#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
9890unsafe fn dot_q4b_row_1x4_vnni(
9891    buf: &[u8],
9892    scales: &[u8],
9893    g0: usize,
9894    gpr: usize,
9895    xs: [&[i8]; 4],
9896) -> [f32; 4] {
9897    // SAFETY: callers uphold buffer contracts (buf.len() == gpr·32).
9898    unsafe {
9899        use core::arch::x86_64::*;
9900        let mut acc = [0f32; 4];
9901        for gi in 0..gpr {
9902            let s = f16_to_f32(u16::from_le_bytes([
9903                scales[(g0 + gi) * 2],
9904                scales[(g0 + gi) * 2 + 1],
9905            ]));
9906            let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9907            let aw = _mm256_abs_epi8(w);
9908            for (k, xq) in xs.iter().enumerate() {
9909                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9910                let d = dpbusd_hsum(aw, _mm256_sign_epi8(x, w));
9911                acc[k] += d as f32 * s;
9912            }
9913        }
9914        acc
9915    }
9916}
9917
9918/// The vbit flavor of the blocked 1×4: the per-activation A8W8 scale
9919/// folds in PER GROUP as `(d·sx)·s` — bit-matching the single-matvec
9920/// accumulation order (the q4_block flavor applies sx once at the end,
9921/// matching ITS single path; the two conventions are historical and
9922/// each blocked leg must mirror its own).
9923#[cfg(target_arch = "x86_64")]
9924#[target_feature(enable = "avx2")]
9925unsafe fn dot_q4b_row_1x4_sx_avx2(
9926    buf: &[u8],
9927    scales: &[u8],
9928    g0: usize,
9929    gpr: usize,
9930    xs: [&[i8]; 4],
9931    sxs: [f32; 4],
9932) -> [f32; 4] {
9933    // SAFETY: callers uphold buffer contracts (buf.len() == gpr·32).
9934    unsafe {
9935        use core::arch::x86_64::*;
9936        let ones = _mm256_set1_epi16(1);
9937        let mut acc = [0f32; 4];
9938        for gi in 0..gpr {
9939            let s = f16_to_f32(u16::from_le_bytes([
9940                scales[(g0 + gi) * 2],
9941                scales[(g0 + gi) * 2 + 1],
9942            ]));
9943            let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9944            let aw = _mm256_abs_epi8(w);
9945            for (k, xq) in xs.iter().enumerate() {
9946                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9947                let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
9948                let d = _mm256_madd_epi16(p16, ones);
9949                let hi128 = _mm256_extracti128_si256::<1>(d);
9950                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
9951                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
9952                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
9953                acc[k] += (_mm_cvtsi128_si32(s32) as f32 * sxs[k]) * s;
9954            }
9955        }
9956        acc
9957    }
9958}
9959
9960/// VNNI twin of `dot_q4b_row_1x4_sx_avx2` (see `dpbusd_hsum`; the
9961/// per-group `(d·sx)·s` fold mirrors the vbit single path).
9962#[cfg(target_arch = "x86_64")]
9963#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
9964unsafe fn dot_q4b_row_1x4_sx_vnni(
9965    buf: &[u8],
9966    scales: &[u8],
9967    g0: usize,
9968    gpr: usize,
9969    xs: [&[i8]; 4],
9970    sxs: [f32; 4],
9971) -> [f32; 4] {
9972    // SAFETY: callers uphold buffer contracts (buf.len() == gpr·32).
9973    unsafe {
9974        use core::arch::x86_64::*;
9975        let mut acc = [0f32; 4];
9976        for gi in 0..gpr {
9977            let s = f16_to_f32(u16::from_le_bytes([
9978                scales[(g0 + gi) * 2],
9979                scales[(g0 + gi) * 2 + 1],
9980            ]));
9981            let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9982            let aw = _mm256_abs_epi8(w);
9983            for (k, xq) in xs.iter().enumerate() {
9984                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9985                let d = dpbusd_hsum(aw, _mm256_sign_epi8(x, w));
9986                acc[k] += (d as f32 * sxs[k]) * s;
9987            }
9988        }
9989        acc
9990    }
9991}
9992
9993#[allow(unreachable_code)]
9994fn dot_q4_row_i8(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
9995    #[cfg(target_arch = "aarch64")]
9996    unsafe {
9997        return dot_q4_row_sdot(packed, scales, g0, gpr, xq);
9998    }
9999    #[cfg(target_arch = "x86_64")]
10000    unsafe {
10001        return dot_q4_row_avx2(packed, scales, g0, gpr, xq);
10002    }
10003    let mut acc = 0f32;
10004    for gi in 0..gpr {
10005        let g = g0 + gi;
10006        let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
10007        let mut d = 0i32;
10008        for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
10009            d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
10010                + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
10011        }
10012        acc += d as f32 * s;
10013    }
10014    acc
10015}
10016
10017/// Two-activation q4 row via the A8W8 int8 path (see `dot_q4_row_i8`).
10018#[inline]
10019#[allow(unreachable_code)]
10020fn dot_q4_row_i8_2(
10021    packed: &[u8],
10022    scales: &[u8],
10023    g0: usize,
10024    gpr: usize,
10025    xq1: &[i8],
10026    xq2: &[i8],
10027) -> (f32, f32) {
10028    #[cfg(target_arch = "aarch64")]
10029    unsafe {
10030        return dot_q4_row_sdot2(packed, scales, g0, gpr, xq1, xq2);
10031    }
10032    #[cfg(target_arch = "x86_64")]
10033    unsafe {
10034        return dot_q4_row_avx2_2(packed, scales, g0, gpr, xq1, xq2);
10035    }
10036    (
10037        dot_q4_row_i8(packed, scales, g0, gpr, xq1),
10038        dot_q4_row_i8(packed, scales, g0, gpr, xq2),
10039    )
10040}
10041
10042/// One q4 row range via SDOT (kernel body of `q4matvec`, extracted so
10043/// multi-matrix jobs can drive it for several tensors in one dispatch).
10044#[allow(clippy::too_many_arguments)]
10045fn q4_range_a8w8(
10046    packed: &[u8],
10047    scales: &[u8],
10048    gpr: usize,
10049    cols: usize,
10050    act: &SplitAct,
10051    out: SendMut,
10052    start: usize,
10053    end: usize,
10054) {
10055    for r in start..end {
10056        let mut acc = dot_q4_row_i8(packed, scales, r * gpr, gpr, &act.xq) * act.sx;
10057        // xq is zeroed at outlier slots — add the exact terms.
10058        for &(j, xv) in &act.outliers {
10059            let flat = r * cols + j;
10060            let byte = packed[flat / 2];
10061            let nib = if flat & 1 == 0 {
10062                byte & 0x0F
10063            } else {
10064                byte >> 4
10065            };
10066            let s = f16_to_f32(u16::from_le_bytes([
10067                scales[(flat / GROUP_SIZE) * 2],
10068                scales[(flat / GROUP_SIZE) * 2 + 1],
10069            ]));
10070            acc += ((nib as i32 - 8) as f32) * s * xv;
10071        }
10072        // SAFETY: disjoint row ranges per worker.
10073        unsafe { *out.at(r) = acc };
10074    }
10075}
10076
10077/// Two-input q4 row range via the A8W8 int8 path — kernel body of
10078/// `q4matvec2`, extracted for pair multi-matrix jobs.
10079#[allow(clippy::too_many_arguments)]
10080fn q4_range2_a8w8(
10081    packed: &[u8],
10082    scales: &[u8],
10083    gpr: usize,
10084    cols: usize,
10085    a1: &SplitAct,
10086    a2: &SplitAct,
10087    p1: SendMut,
10088    p2: SendMut,
10089    start: usize,
10090    end: usize,
10091) {
10092    for r in start..end {
10093        let (s1, s2) = dot_q4_row_i8_2(packed, scales, r * gpr, gpr, &a1.xq, &a2.xq);
10094        let mut acc1 = s1 * a1.sx;
10095        let mut acc2 = s2 * a2.sx;
10096        // xq is zeroed at outlier slots — add the exact terms.
10097        let fix = |outliers: &[(usize, f32)], acc: &mut f32| {
10098            for &(j, xv) in outliers {
10099                let flat = r * cols + j;
10100                let byte = packed[flat / 2];
10101                let nib = if flat & 1 == 0 {
10102                    byte & 0x0F
10103                } else {
10104                    byte >> 4
10105                };
10106                let s = f16_to_f32(u16::from_le_bytes([
10107                    scales[(flat / GROUP_SIZE) * 2],
10108                    scales[(flat / GROUP_SIZE) * 2 + 1],
10109                ]));
10110                *acc += ((nib as i32 - 8) as f32) * s * xv;
10111            }
10112        };
10113        fix(&a1.outliers, &mut acc1);
10114        fix(&a2.outliers, &mut acc2);
10115        // SAFETY: disjoint row ranges per worker.
10116        unsafe {
10117            *p1.at(r) = acc1;
10118            *p2.at(r) = acc2;
10119        }
10120    }
10121}
10122
10123/// Exact scalar q4 row range (same extraction, non-SDOT path).
10124fn q4_range_f32(
10125    packed: &[u8],
10126    scales: &[u8],
10127    gpr: usize,
10128    x: &[f32],
10129    out: SendMut,
10130    start: usize,
10131    end: usize,
10132) {
10133    for r in start..end {
10134        let mut acc = 0f32;
10135        for gi in 0..gpr {
10136            let g = r * gpr + gi;
10137            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
10138            let pk = &packed[g * 16..(g + 1) * 16];
10139            let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
10140            let mut ga = 0f32;
10141            for (k, &b) in pk.iter().enumerate() {
10142                ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
10143                    + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
10144            }
10145            acc += ga * s;
10146        }
10147        // SAFETY: disjoint row ranges per worker.
10148        unsafe { *out.at(r) = acc };
10149    }
10150}
10151
10152/// Fused two-input q4 matvec: nibbles are unpacked ONCE per group and
10153/// dotted against both activations (was: two full matvecs — double
10154/// weight traffic). Per-lane math matches `q4matvec` exactly.
10155#[allow(clippy::too_many_arguments)]
10156fn q4matvec2(
10157    bytes: &[u8],
10158    x1: &[f32],
10159    x2: &[f32],
10160    rows: usize,
10161    cols: usize,
10162    o1: &mut [f32],
10163    o2: &mut [f32],
10164    pool: Option<&Pool>,
10165) {
10166    debug_assert_eq!(o1.len(), rows);
10167    debug_assert_eq!(o2.len(), rows);
10168    let (packed, scales) = q4_split(bytes, rows, cols);
10169    let gpr = cols / GROUP_SIZE;
10170
10171    if a8w8_enabled() {
10172        let a1 = split_act(x1);
10173        let a2 = split_act(x2);
10174        let p1 = SendMut(o1.as_mut_ptr());
10175        let p2 = SendMut(o2.as_mut_ptr());
10176        let run = move |start: usize, end: usize| {
10177            q4_range2_a8w8(packed, scales, gpr, cols, &a1, &a2, p1, p2, start, end)
10178        };
10179        dispatch_rows(pool, rows, &run);
10180        return;
10181    }
10182
10183    let p1 = SendMut(o1.as_mut_ptr());
10184    let p2 = SendMut(o2.as_mut_ptr());
10185    let run = move |start: usize, end: usize| {
10186        q4_range2_f32(packed, scales, gpr, x1, x2, p1, p2, start, end)
10187    };
10188    dispatch_rows(pool, rows, &run);
10189}
10190
10191/// Two-input exact scalar q4 row range (same extraction).
10192#[allow(clippy::too_many_arguments)]
10193fn q4_range2_f32(
10194    packed: &[u8],
10195    scales: &[u8],
10196    gpr: usize,
10197    x1: &[f32],
10198    x2: &[f32],
10199    p1: SendMut,
10200    p2: SendMut,
10201    start: usize,
10202    end: usize,
10203) {
10204    for r in start..end {
10205        let (mut acc1, mut acc2) = (0f32, 0f32);
10206        for gi in 0..gpr {
10207            let g = r * gpr + gi;
10208            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
10209            let pk = &packed[g * 16..(g + 1) * 16];
10210            let x1g = &x1[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
10211            let x2g = &x2[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
10212            let (mut g1, mut g2) = (0f32, 0f32);
10213            for (k, &b) in pk.iter().enumerate() {
10214                let wl = (b & 0x0F) as f32 - 8.0;
10215                let wh = ((b >> 4) & 0x0F) as f32 - 8.0;
10216                g1 += wl * x1g[k * 2] + wh * x1g[k * 2 + 1];
10217                g2 += wl * x2g[k * 2] + wh * x2g[k * 2 + 1];
10218            }
10219            acc1 += g1 * s;
10220            acc2 += g2 * s;
10221        }
10222        // SAFETY: disjoint row ranges per worker.
10223        unsafe {
10224            *p1.at(r) = acc1;
10225            *p2.at(r) = acc2;
10226        }
10227    }
10228}
10229
10230thread_local! {
10231    /// Per-worker decoded-row scratch for the batched q4/vbit kernels
10232    /// (centered i8 for SDOT, f32 for the exact/scalar paths).
10233    static ROW_I8: std::cell::RefCell<Vec<u8>> = const { std::cell::RefCell::new(Vec::new()) };
10234    static ROW_F32: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
10235}
10236
10237/// Batched q4 matmat: each weight row is unpacked from the mmap ONCE
10238/// and dotted against ALL b activations (prefill used to fall back to b
10239/// full matvecs — b× weight traffic and b× nibble decode). Per-position
10240/// math matches `q4matvec` exactly: same group order, same accumulation.
10241/// `out` is row-major [b, rows] like `qmatmat`.
10242#[allow(clippy::too_many_arguments)]
10243fn q4matmat(
10244    bytes: &[u8],
10245    xs_all: &[f32],
10246    b: usize,
10247    rows: usize,
10248    cols: usize,
10249    out: &mut [f32],
10250    pool: Option<&Pool>,
10251) {
10252    debug_assert_eq!(xs_all.len(), b * cols);
10253    debug_assert_eq!(out.len(), b * rows);
10254    let (packed, scales) = q4_split(bytes, rows, cols);
10255    let gpr = cols / GROUP_SIZE;
10256    let gscale = |g: usize| f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
10257
10258    if a8w8_enabled() {
10259        let acts: Vec<SplitAct> = (0..b)
10260            .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
10261            .collect();
10262        let acts = &acts;
10263        let out_addr = SendMut(out.as_mut_ptr());
10264        let run = move |start: usize, end: usize| {
10265            ROW_I8.with(|rb| {
10266                let mut buf = rb.borrow_mut();
10267                buf.resize(cols, 0);
10268                for r in start..end {
10269                    // Unpack the row's nibbles to centered i8 once
10270                    // (element 2k = low nibble, 2k+1 = high — flat order,
10271                    // same as dot_q4_row_sdot's zip).
10272                    for gi in 0..gpr {
10273                        let g = r * gpr + gi;
10274                        for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
10275                            buf[gi * GROUP_SIZE + k * 2] = ((bt & 0x0F) as i32 - 8) as i8 as u8;
10276                            buf[gi * GROUP_SIZE + k * 2 + 1] =
10277                                (((bt >> 4) & 0x0F) as i32 - 8) as i8 as u8;
10278                        }
10279                    }
10280                    let mut bi = 0usize;
10281                    #[cfg(target_arch = "x86_64")]
10282                    if avx2_enabled() && blocked_enabled() {
10283                        while bi + 4 <= acts.len() {
10284                            let xs = [
10285                                acts[bi].xq.as_slice(),
10286                                acts[bi + 1].xq.as_slice(),
10287                                acts[bi + 2].xq.as_slice(),
10288                                acts[bi + 3].xq.as_slice(),
10289                            ];
10290                            let d = unsafe {
10291                                if vnni_tiles_enabled() {
10292                                    dot_q4b_row_1x4_vnni(&buf, scales, r * gpr, gpr, xs)
10293                                } else {
10294                                    dot_q4b_row_1x4_avx2(&buf, scales, r * gpr, gpr, xs)
10295                                }
10296                            };
10297                            for k in 0..4 {
10298                                let act = &acts[bi + k];
10299                                let mut acc = d[k] * act.sx;
10300                                for &(j, xv) in &act.outliers {
10301                                    acc += (buf[j] as i8) as f32
10302                                        * gscale((r * cols + j) / GROUP_SIZE)
10303                                        * xv;
10304                                }
10305                                // SAFETY: disjoint (bi, r) cells per worker.
10306                                unsafe { *out_addr.at((bi + k) * rows + r) = acc };
10307                            }
10308                            bi += 4;
10309                        }
10310                    }
10311                    while bi < acts.len() {
10312                        let act = &acts[bi];
10313                        let mut acc = 0f32;
10314                        for gi in 0..gpr {
10315                            let d = dot_i8_i8(
10316                                &buf[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
10317                                &act.xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
10318                            );
10319                            acc += d as f32 * gscale(r * gpr + gi);
10320                        }
10321                        acc *= act.sx;
10322                        // xq is zeroed at outlier slots — exact terms.
10323                        for &(j, xv) in &act.outliers {
10324                            acc += (buf[j] as i8) as f32 * gscale((r * cols + j) / GROUP_SIZE) * xv;
10325                        }
10326                        // SAFETY: disjoint (bi, r) cells per worker row range.
10327                        unsafe { *out_addr.at(bi * rows + r) = acc };
10328                        bi += 1;
10329                    }
10330                }
10331            })
10332        };
10333        dispatch_rows(pool, rows, &run);
10334        return;
10335    }
10336
10337    let out_addr = SendMut(out.as_mut_ptr());
10338    let run = move |start: usize, end: usize| {
10339        ROW_F32.with(|rb| {
10340            let mut buf = rb.borrow_mut();
10341            buf.resize(cols, 0.0);
10342            for r in start..end {
10343                // Decode raw (nib − 8) values once; scales stay per-group
10344                // so the accumulation order matches q4matvec bit-for-bit.
10345                for gi in 0..gpr {
10346                    let g = r * gpr + gi;
10347                    for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
10348                        buf[gi * GROUP_SIZE + k * 2] = (bt & 0x0F) as f32 - 8.0;
10349                        buf[gi * GROUP_SIZE + k * 2 + 1] = ((bt >> 4) & 0x0F) as f32 - 8.0;
10350                    }
10351                }
10352                for bi in 0..b {
10353                    let x = &xs_all[bi * cols..(bi + 1) * cols];
10354                    let mut acc = 0f32;
10355                    for gi in 0..gpr {
10356                        let mut ga = 0f32;
10357                        // Pairwise (lo + hi) addition, matching
10358                        // q4matvec's `ga += lo·x + hi·x` shape exactly —
10359                        // a flat one-per-element loop rounds differently
10360                        // and broke bit-parity on the scalar (x86) path.
10361                        for k in 0..GROUP_SIZE / 2 {
10362                            let e = gi * GROUP_SIZE + k * 2;
10363                            ga += buf[e] * x[e] + buf[e + 1] * x[e + 1];
10364                        }
10365                        acc += ga * gscale(r * gpr + gi);
10366                    }
10367                    // SAFETY: disjoint (bi, r) cells per worker row range.
10368                    unsafe { *out_addr.at(bi * rows + r) = acc };
10369                }
10370            }
10371        })
10372    };
10373    dispatch_rows(pool, rows, &run);
10374}
10375
10376/// Batched vbit matmat: each variable-bit row is decoded from the mmap
10377/// ONCE for the whole microbatch. Same per-position math as
10378/// `vbitmatvec` (SDOT A8W8 with exact outliers / exact f32 for b=8 rows
10379/// and the scalar path).
10380#[allow(clippy::too_many_arguments)]
10381fn vbitmatmat(
10382    bytes: &[u8],
10383    offsets: &[usize],
10384    xs_all: &[f32],
10385    b: usize,
10386    rows: usize,
10387    cols: usize,
10388    out: &mut [f32],
10389    pool: Option<&Pool>,
10390) {
10391    debug_assert_eq!(xs_all.len(), b * cols);
10392    debug_assert_eq!(out.len(), b * rows);
10393    debug_assert_eq!(offsets.len(), rows + 1);
10394    let ng = cols / GROUP_SIZE;
10395    let bits = &bytes[..rows];
10396    let sc_off = rows;
10397    let gscale = |r: usize, g: usize| {
10398        let so = (r * ng + g) * 2;
10399        f16_to_f32(u16::from_le_bytes([
10400            bytes[sc_off + so],
10401            bytes[sc_off + so + 1],
10402        ]))
10403    };
10404
10405    // Decode row r's raw (u − L) values into `dst` (f32, unscaled).
10406    let decode_f32 = |r: usize, dst: &mut [f32]| {
10407        let bw = bits[r] as usize;
10408        let l = ((1i32 << (bw - 1)) - 1) as f32;
10409        let data = &bytes[offsets[r]..offsets[r + 1]];
10410        let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
10411        for d in dst.iter_mut() {
10412            while nbits < bw {
10413                acc = (acc << 8) | data[idx] as u64;
10414                idx += 1;
10415                nbits += 8;
10416            }
10417            let u = ((acc >> (nbits - bw)) & ((1u64 << bw) - 1)) as f32;
10418            nbits -= bw;
10419            *d = u - l;
10420        }
10421    };
10422
10423    if a8w8_enabled() {
10424        let acts: Vec<SplitAct> = (0..b)
10425            .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
10426            .collect();
10427        let acts = &acts;
10428        let out_addr = SendMut(out.as_mut_ptr());
10429        let run = move |start: usize, end: usize| {
10430            for r in start..end {
10431                let bw = bits[r] as usize;
10432                if bw == 8 {
10433                    // u−L reaches 128 → no i8 path; decode once, exact
10434                    // f32 dots for every position (same as vbitmatvec).
10435                    ROW_F32.with(|rb| {
10436                        let mut buf = rb.borrow_mut();
10437                        buf.resize(cols, 0.0);
10438                        decode_f32(r, &mut buf);
10439                        for bi in 0..b {
10440                            let x = &xs_all[bi * cols..(bi + 1) * cols];
10441                            let mut dot = 0f32;
10442                            for g in 0..ng {
10443                                let mut gd = 0f32;
10444                                for k in 0..GROUP_SIZE {
10445                                    gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
10446                                }
10447                                dot += gd * gscale(r, g);
10448                            }
10449                            // SAFETY: disjoint (bi, r) cells per worker range.
10450                            unsafe { *out_addr.at(bi * rows + r) = dot };
10451                        }
10452                    });
10453                    continue;
10454                }
10455                let l = (1i32 << (bw - 1)) - 1;
10456                let data = &bytes[offsets[r]..offsets[r + 1]];
10457                ROW_I8.with(|rb| {
10458                    let mut buf = rb.borrow_mut();
10459                    buf.resize(cols, 0);
10460                    #[inline(always)]
10461                    fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
10462                        for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
10463                            let u = unpack8::<B>(&data[blk * B..]);
10464                            for k in 0..8 {
10465                                chunk[k] = (u[k] - l) as i8 as u8;
10466                            }
10467                        }
10468                    }
10469                    match bw {
10470                        3 => fill::<3>(data, l, &mut buf),
10471                        4 => vbit_fill4(data, &mut buf),
10472                        5 => fill::<5>(data, l, &mut buf),
10473                        6 => fill::<6>(data, l, &mut buf),
10474                        _ => unreachable!("vbit bit-width {bw} (validated at load)"),
10475                    }
10476                    let mut bi = 0usize;
10477                    // The vbit scale table shares q4_block's layout
10478                    // (contiguous f16 per (row·ng + g)), so the same
10479                    // blocked 1×4 kernel serves the decoded row.
10480                    #[cfg(target_arch = "x86_64")]
10481                    if avx2_enabled() && blocked_enabled() {
10482                        while bi + 4 <= acts.len() {
10483                            let xs = [
10484                                acts[bi].xq.as_slice(),
10485                                acts[bi + 1].xq.as_slice(),
10486                                acts[bi + 2].xq.as_slice(),
10487                                acts[bi + 3].xq.as_slice(),
10488                            ];
10489                            let sxs = [
10490                                acts[bi].sx,
10491                                acts[bi + 1].sx,
10492                                acts[bi + 2].sx,
10493                                acts[bi + 3].sx,
10494                            ];
10495                            let d = unsafe {
10496                                if vnni_tiles_enabled() {
10497                                    dot_q4b_row_1x4_sx_vnni(
10498                                        &buf,
10499                                        &bytes[sc_off..],
10500                                        r * ng,
10501                                        ng,
10502                                        xs,
10503                                        sxs,
10504                                    )
10505                                } else {
10506                                    dot_q4b_row_1x4_sx_avx2(
10507                                        &buf,
10508                                        &bytes[sc_off..],
10509                                        r * ng,
10510                                        ng,
10511                                        xs,
10512                                        sxs,
10513                                    )
10514                                }
10515                            };
10516                            for k in 0..4 {
10517                                let act = &acts[bi + k];
10518                                let mut dot = d[k];
10519                                for &(j, xv) in &act.outliers {
10520                                    dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
10521                                }
10522                                // SAFETY: disjoint (bi, r) cells per worker.
10523                                unsafe { *out_addr.at((bi + k) * rows + r) = dot };
10524                            }
10525                            bi += 4;
10526                        }
10527                    }
10528                    while bi < acts.len() {
10529                        let act = &acts[bi];
10530                        let mut dot = 0f32;
10531                        for g in 0..ng {
10532                            let d = dot_i8_i8(
10533                                &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
10534                                &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
10535                            ) as f32
10536                                * act.sx;
10537                            dot += d * gscale(r, g);
10538                        }
10539                        for &(j, xv) in &act.outliers {
10540                            dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
10541                        }
10542                        // SAFETY: disjoint (bi, r) cells per worker range.
10543                        unsafe { *out_addr.at(bi * rows + r) = dot };
10544                        bi += 1;
10545                    }
10546                });
10547            }
10548        };
10549        dispatch_rows(pool, rows, &run);
10550        return;
10551    }
10552
10553    let out_addr = SendMut(out.as_mut_ptr());
10554    let run = move |start: usize, end: usize| {
10555        ROW_F32.with(|rb| {
10556            let mut buf = rb.borrow_mut();
10557            buf.resize(cols, 0.0);
10558            for r in start..end {
10559                decode_f32(r, &mut buf);
10560                for bi in 0..b {
10561                    let x = &xs_all[bi * cols..(bi + 1) * cols];
10562                    let mut dot = 0f32;
10563                    for g in 0..ng {
10564                        let mut gd = 0f32;
10565                        for k in 0..GROUP_SIZE {
10566                            gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
10567                        }
10568                        dot += gd * gscale(r, g);
10569                    }
10570                    // SAFETY: disjoint (bi, r) cells per worker range.
10571                    unsafe { *out_addr.at(bi * rows + r) = dot };
10572                }
10573            }
10574        })
10575    };
10576    dispatch_rows(pool, rows, &run);
10577}
10578
10579/// Build a GPU batch job for a q8-family mapped tensor (primary
10580/// shard): prescaled input + directory coordinates. None → not
10581/// GPU-eligible, caller stays on the CPU.
10582pub(crate) fn gpu_batch_job<'a>(
10583    t: &'a QTensor,
10584    x: &[f32],
10585) -> Option<(std::sync::Arc<CmfModel>, crate::gpu::BatchJob<'a>)> {
10586    match t {
10587        QTensor::Mapped {
10588            model,
10589            idx,
10590            dtype: dt @ (TensorDtype::Q8Row | TensorDtype::Q8_2f),
10591            rows,
10592            cols,
10593            row_scale,
10594            col_field,
10595            ..
10596        } => Some((
10597            model.clone(),
10598            crate::gpu::BatchJob {
10599                idx: *idx,
10600                rows: *rows,
10601                cols: *cols,
10602                row_scale,
10603                xs: prescale(x, col_field, *dt).into_owned(),
10604                layout: crate::gpu::BatchLayout::Q8,
10605            },
10606        )),
10607        // q1: raw f32 activations, tile-embedded scales.
10608        QTensor::Mapped {
10609            model,
10610            idx,
10611            dtype: TensorDtype::Q1,
10612            rows,
10613            cols,
10614            ..
10615        } => Some((
10616            model.clone(),
10617            crate::gpu::BatchJob {
10618                idx: *idx,
10619                rows: *rows,
10620                cols: *cols,
10621                row_scale: &[],
10622                xs: x.to_vec(),
10623                layout: crate::gpu::BatchLayout::Q1,
10624            },
10625        )),
10626        // q1g: same raw-f32 contract, sign + scale planes.
10627        QTensor::Mapped {
10628            model,
10629            idx,
10630            dtype: TensorDtype::Q1G,
10631            rows,
10632            cols,
10633            ..
10634        } => Some((
10635            model.clone(),
10636            crate::gpu::BatchJob {
10637                idx: *idx,
10638                rows: *rows,
10639                cols: *cols,
10640                row_scale: &[],
10641                xs: x.to_vec(),
10642                layout: crate::gpu::BatchLayout::Q1g,
10643            },
10644        )),
10645        // q4_tiled / q4tp: raw f32 activations; the scales live in the
10646        // payload (inline tiles / row ladder), so row_scale stays empty.
10647        // The GDN projection batch already runs these layouts on Metal —
10648        // this arm lets the attention QKV batch reach the same kernels.
10649        QTensor::Mapped {
10650            model,
10651            idx,
10652            dtype: dt @ (TensorDtype::Q4Tiled | TensorDtype::Q4TiledP),
10653            rows,
10654            cols,
10655            ..
10656        } => Some((
10657            model.clone(),
10658            crate::gpu::BatchJob {
10659                idx: *idx,
10660                rows: *rows,
10661                cols: *cols,
10662                row_scale: &[],
10663                xs: x.to_vec(),
10664                layout: if *dt == TensorDtype::Q4Tiled {
10665                    crate::gpu::BatchLayout::Q4t
10666                } else {
10667                    crate::gpu::BatchLayout::Q4tp
10668                },
10669            },
10670        )),
10671        _ => None,
10672    }
10673}
10674
10675thread_local! {
10676    static PRESCALE_BUF1: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
10677    static PRESCALE_BUF2: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
10678}
10679
10680pub(crate) fn prescale<'a>(
10681    x: &'a [f32],
10682    col_field: &[f32],
10683    dtype: TensorDtype,
10684) -> std::borrow::Cow<'a, [f32]> {
10685    if dtype == TensorDtype::Q8_2f {
10686        x.iter().zip(col_field).map(|(a, c)| a * c).collect()
10687    } else {
10688        std::borrow::Cow::Borrowed(x)
10689    }
10690}
10691
10692/// θ col-field fold for q8_2f activations. Borrowed pass-through for
10693/// every other dtype, using thread-local buffers to eliminate per-matvec allocations.
10694pub(crate) fn prescale_with<R, F: FnOnce(&[f32]) -> R>(
10695    x: &[f32],
10696    col_field: &[f32],
10697    dtype: TensorDtype,
10698    buf_id: u8,
10699    f: F,
10700) -> R {
10701    if dtype == TensorDtype::Q8_2f {
10702        if buf_id == 1 {
10703            PRESCALE_BUF1.with(|b| {
10704                let mut buf = b.borrow_mut();
10705                buf.clear();
10706                buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
10707                f(&buf)
10708            })
10709        } else {
10710            PRESCALE_BUF2.with(|b| {
10711                let mut buf = b.borrow_mut();
10712                buf.clear();
10713                buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
10714                f(&buf)
10715            })
10716        }
10717    } else {
10718        f(x)
10719    }
10720}
10721
10722// ───────────────────── x86-64 AVX2 kernels (roadmap этап 2) ─────────────────────
10723
10724/// AVX2+FMA available? Default ON when the CPU supports both;
10725/// `CMF_AVX2=0` disables (falls back to the autovectorized loops).
10726#[cfg(target_arch = "x86_64")]
10727pub(crate) fn avx2_enabled() -> bool {
10728    use std::sync::OnceLock;
10729    static ON: OnceLock<bool> = OnceLock::new();
10730    *ON.get_or_init(|| {
10731        std::env::var("CMF_AVX2").map(|v| v != "0").unwrap_or(true)
10732            && std::arch::is_x86_feature_detected!("avx2")
10733            && std::arch::is_x86_feature_detected!("fma")
10734    })
10735}
10736
10737/// AVX2 A8W8 allowed? The quantized-activation contract is switched by
10738/// the SAME env as the ARM SDOT path: `CMF_SDOT=0` keeps exact kernels
10739/// (the golden-parity exact gate relies on it) — AVX2 f32 kernels stay
10740/// active either way, they are exact (regrouped sums only).
10741#[cfg(target_arch = "x86_64")]
10742fn avx2_a8w8_enabled() -> bool {
10743    if FLOAT_ACTIVATIONS.get() {
10744        return false;
10745    }
10746    use std::sync::OnceLock;
10747    static ON: OnceLock<bool> = OnceLock::new();
10748    *ON.get_or_init(|| {
10749        avx2_enabled() && std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true)
10750    })
10751}
10752
10753thread_local! {
10754    static FULL_GPU_Q8: std::cell::Cell<bool> = const { std::cell::Cell::new(false) };
10755}
10756
10757/// Match the graph's full-device q8 projection precision on MiMo's host
10758/// tail. Does not enable the GPU or bypass a CPU-only/device-refusal gate.
10759pub(crate) fn enter_full_gpu_q8_scope() -> impl Drop {
10760    struct Restore(bool, std::marker::PhantomData<std::rc::Rc<()>>);
10761    impl Drop for Restore {
10762        fn drop(&mut self) {
10763            FULL_GPU_Q8.set(self.0);
10764        }
10765    }
10766    Restore(FULL_GPU_Q8.replace(true), std::marker::PhantomData)
10767}
10768
10769// Dynamic MiMo experts must not change activation precision when a cache
10770// fill moves them from CPU to GPU. Thread-local: only the cold-expert
10771// dispatch selects float kernels; concurrent pipelines keep their policy.
10772thread_local! {
10773    static FLOAT_ACTIVATIONS: std::cell::Cell<bool> = const { std::cell::Cell::new(false) };
10774}
10775
10776/// Run `f` with the exact float-activation contract on this thread: the
10777/// 4-bit and 8-bit kernels skip the int8 activation split (`CMF_SDOT=0`
10778/// semantics, scoped). For parity tests and diagnostics — e.g. an `mxfp4`
10779/// model then computes bit for bit what its dequantized f32 twin computes.
10780pub fn with_float_activations<R>(f: impl FnOnce() -> R) -> R {
10781    float_activations_scope(f)
10782}
10783
10784/// Run `f` with exact float activations on this thread (no A8W8 int8
10785/// activation quantization in the quantized matvecs). Public for parity
10786/// tests that compare the default kernels with the strict ones in one
10787/// process; `CMF_SDOT=0` is the process-wide equivalent.
10788#[doc(hidden)]
10789pub fn float_activations_scope<R>(f: impl FnOnce() -> R) -> R {
10790    struct Restore(bool);
10791    impl Drop for Restore {
10792        fn drop(&mut self) {
10793            FLOAT_ACTIVATIONS.set(self.0);
10794        }
10795    }
10796    let _restore = Restore(FLOAT_ACTIVATIONS.replace(true));
10797    f()
10798}
10799
10800/// Row-exact batching: while set, the x86 batched kernels (`qmatmat`,
10801/// `q4tp_matmat`) compute every (weight row, token) cell with the
10802/// single-token kernel instead of the blocked 2×4 / 1×4 / 1×8 tiles, so a
10803/// token's result does not depend on the batch it rides in and equals its
10804/// matvec. On ARM `q4tp_matmat` keeps its 1×4 tile but in the matvec's
10805/// reduction order, and no q4tp batch takes the AMX or device GEMM. The
10806/// MiMo speculative verify holds it (`row_exact_scope`) — its accepted
10807/// rows must be the rows plain decode would have produced.
10808// Shared with pool workers, so overlapping requests must keep the mode
10809// enabled until the LAST scope leaves. Saving/restoring a global bool is
10810// incorrect when two threads enter and leave in a non-LIFO order.
10811static ROW_EXACT: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
10812
10813pub(crate) fn row_exact() -> bool {
10814    ROW_EXACT.load(std::sync::atomic::Ordering::Acquire) != 0
10815}
10816
10817fn counted_row_exact_scope<R>(active: &std::sync::atomic::AtomicUsize, f: impl FnOnce() -> R) -> R {
10818    struct Restore<'a>(&'a std::sync::atomic::AtomicUsize);
10819    impl Drop for Restore<'_> {
10820        fn drop(&mut self) {
10821            self.0.fetch_sub(1, std::sync::atomic::Ordering::AcqRel);
10822        }
10823    }
10824    active.fetch_add(1, std::sync::atomic::Ordering::AcqRel);
10825    let _restore = Restore(active);
10826    f()
10827}
10828
10829/// Run `f` with row-exact batching on (also released on unwind).
10830pub(crate) fn row_exact_scope<R>(f: impl FnOnce() -> R) -> R {
10831    counted_row_exact_scope(&ROW_EXACT, f)
10832}
10833
10834/// A8W8 quantized-activation path available on THIS machine? One
10835/// switch across architectures: ARM dotprod (CMF_SDOT) or x86 AVX2
10836/// (CMF_AVX2 + the same CMF_SDOT exact-contract override).
10837#[inline]
10838pub(crate) fn a8w8_enabled() -> bool {
10839    #[cfg(target_arch = "aarch64")]
10840    {
10841        sdot_enabled()
10842    }
10843    #[cfg(target_arch = "x86_64")]
10844    {
10845        avx2_a8w8_enabled()
10846    }
10847    #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
10848    {
10849        false
10850    }
10851}
10852
10853/// int8·int8 dot dispatch: SDOT on ARM; AVX-512 VNNI (vpdpbusd) or AVX2
10854/// maddubs on x86. Callers are gated by `a8w8_enabled()`.
10855#[inline]
10856#[allow(unreachable_code)]
10857fn dot_i8_i8(w: &[u8], xq: &[i8]) -> i32 {
10858    #[cfg(target_arch = "aarch64")]
10859    unsafe {
10860        return dot_i8_sdot(w, xq);
10861    }
10862    #[cfg(target_arch = "x86_64")]
10863    unsafe {
10864        if avx512vnni_enabled() {
10865            return dot_i8_i8_vnni(w, xq);
10866        }
10867        return dot_i8_i8_avx2(w, xq);
10868    }
10869    w.iter()
10870        .zip(xq)
10871        .map(|(&a, &b)| (a as i8) as i32 * b as i32)
10872        .sum()
10873}
10874
10875/// AVX-512 VNNI available? (F+BW+VL+VNNI; `CMF_AVX512=0` falls back to
10876/// AVX2.) VL matters: short 32-byte groups (q4/vbit) ride the 256-bit
10877/// `vpdpbusd` encoding.
10878#[cfg(target_arch = "x86_64")]
10879fn avx512vnni_enabled() -> bool {
10880    use std::sync::OnceLock;
10881    static ON: OnceLock<bool> = OnceLock::new();
10882    *ON.get_or_init(|| {
10883        std::env::var("CMF_AVX512")
10884            .map(|v| v != "0")
10885            .unwrap_or(true)
10886            && std::arch::is_x86_feature_detected!("avx512f")
10887            && std::arch::is_x86_feature_detected!("avx512bw")
10888            && std::arch::is_x86_feature_detected!("avx512vl")
10889            && std::arch::is_x86_feature_detected!("avx512vnni")
10890    })
10891}
10892
10893/// Grouped-codec VNNI arms (the q4t/q4b/q1/q1t tile kernels): default
10894/// ON where AVX-512 VNNI exists (`CMF_VNNI_TILES=0` opt-out). Measured
10895/// on Ryzen 7950X (Zen4, 3 alternating process pairs, blocked GEMM
10896/// 4864×896 b=256): q4t 63→68 GF/s (+8%), q1 53→56 (+6%), q4b 72→75
10897/// (+4%) — consistent, no leg regressed. The tile kernels keep a
10898/// horizontal reduce per 32-weight group, so the `vpdpbusd` saving is
10899/// smaller than the long-dot q8 win (+13%), but it is real and free.
10900#[cfg(target_arch = "x86_64")]
10901fn vnni_tiles_enabled() -> bool {
10902    use std::sync::OnceLock;
10903    static ON: OnceLock<bool> = OnceLock::new();
10904    *ON.get_or_init(|| {
10905        std::env::var("CMF_VNNI_TILES")
10906            .map(|v| v != "0")
10907            .unwrap_or(true)
10908            && avx512vnni_enabled()
10909    })
10910}
10911
10912/// One 256-bit u8×i8 dot → i32 via `vpdpbusd` into a fresh accumulator
10913/// plus the same horizontal reduce the AVX2 kernels use. Products are
10914/// bounded (|w| ≤ 8 or ≤ 1), so maddubs never saturated — the i32 sum
10915/// is bit-identical to the maddubs+madd pair it replaces.
10916#[cfg(target_arch = "x86_64")]
10917#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
10918#[inline]
10919unsafe fn dpbusd_hsum(aw: core::arch::x86_64::__m256i, xs: core::arch::x86_64::__m256i) -> i32 {
10920    // SAFETY: pure register math.
10921    unsafe {
10922        use core::arch::x86_64::*;
10923        let d = _mm256_dpbusd_epi32(_mm256_setzero_si256(), aw, xs);
10924        let hi128 = _mm256_extracti128_si256::<1>(d);
10925        let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
10926        let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
10927        let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
10928        _mm_cvtsi128_si32(s32)
10929    }
10930}
10931
10932/// int8·int8 via AVX-512 VNNI: `vpdpbusd` fuses the maddubs+madd+add
10933/// triple into one u8×i8 dot-accumulate. AVX-512 has no vpsignb, so the
10934/// |w|·sign(x,w) trick becomes |w| × (x negated where w<0) via a mask
10935/// subtract — w==0 lanes contribute 0 through |w|=0 either way.
10936#[cfg(target_arch = "x86_64")]
10937#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
10938unsafe fn dot_i8_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
10939    // SAFETY: callers uphold slice-length contracts (see call sites).
10940    unsafe {
10941        use core::arch::x86_64::*;
10942        let n = w.len();
10943        let mut j = 0usize;
10944        let mut total: i32;
10945        // 4 independent accumulators: vpdpbusd is its own loop-carried
10946        // dependency (~5-cycle latency) — a single-acc loop runs
10947        // latency-bound and LOSES to the AVX2 maddubs kernel, measured
10948        // on Granite Rapids.
10949        {
10950            #[inline(always)]
10951            unsafe fn step(
10952                w: *const u8,
10953                x: *const i8,
10954                acc: core::arch::x86_64::__m512i,
10955            ) -> core::arch::x86_64::__m512i {
10956                unsafe {
10957                    use core::arch::x86_64::*;
10958                    let wv = _mm512_loadu_si512(w as *const _);
10959                    let xv = _mm512_loadu_si512(x as *const _);
10960                    let aw = _mm512_abs_epi8(wv);
10961                    let neg = _mm512_movepi8_mask(wv);
10962                    let sx = _mm512_mask_sub_epi8(xv, neg, _mm512_setzero_si512(), xv);
10963                    _mm512_dpbusd_epi32(acc, aw, sx)
10964                }
10965            }
10966            let (mut a0, mut a1, mut a2, mut a3) = (
10967                _mm512_setzero_si512(),
10968                _mm512_setzero_si512(),
10969                _mm512_setzero_si512(),
10970                _mm512_setzero_si512(),
10971            );
10972            while j + 256 <= n {
10973                a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
10974                a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), a1);
10975                a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), a2);
10976                a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), a3);
10977                j += 256;
10978            }
10979            while j + 64 <= n {
10980                a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
10981                j += 64;
10982            }
10983            let s01 = _mm512_add_epi32(a0, a1);
10984            let s23 = _mm512_add_epi32(a2, a3);
10985            total = _mm512_reduce_add_epi32(_mm512_add_epi32(s01, s23));
10986        }
10987        // 32-wide (q4/vbit groups are exactly 32 bytes).
10988        if j + 32 <= n {
10989            let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
10990            let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
10991            let d = _mm256_dpbusd_epi32(
10992                _mm256_setzero_si256(),
10993                _mm256_abs_epi8(wv),
10994                _mm256_sign_epi8(xv, wv),
10995            );
10996            let hi128 = _mm256_extracti128_si256::<1>(d);
10997            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
10998            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
10999            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
11000            total += _mm_cvtsi128_si32(s32);
11001            j += 32;
11002        }
11003        while j < n {
11004            total += (w[j] as i8) as i32 * xq[j] as i32;
11005            j += 1;
11006        }
11007        total
11008    }
11009}
11010
11011/// i8 row · f32 x via AVX2/FMA (x86 mirror of `dot_i8_f32_neon`).
11012#[cfg(target_arch = "x86_64")]
11013#[target_feature(enable = "avx2,fma")]
11014unsafe fn dot_i8_f32_avx2(w: &[u8], x: &[f32]) -> f32 {
11015    // SAFETY: callers uphold slice-length contracts (see call sites).
11016    unsafe {
11017        use core::arch::x86_64::*;
11018        let n = x.len();
11019        let wp = w.as_ptr();
11020        let xp = x.as_ptr();
11021        let (mut a0, mut a1) = (_mm256_setzero_ps(), _mm256_setzero_ps());
11022        let mut j = 0usize;
11023        while j + 16 <= n {
11024            let wb = _mm_loadu_si128(wp.add(j) as *const __m128i);
11025            let lo = _mm256_cvtepi8_epi32(wb);
11026            let hi = _mm256_cvtepi8_epi32(_mm_srli_si128::<8>(wb));
11027            a0 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(lo), _mm256_loadu_ps(xp.add(j)), a0);
11028            a1 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(hi), _mm256_loadu_ps(xp.add(j + 8)), a1);
11029            j += 16;
11030        }
11031        let acc = _mm256_add_ps(a0, a1);
11032        let hi128 = _mm256_extractf128_ps::<1>(acc);
11033        let s128 = _mm_add_ps(_mm256_castps256_ps128(acc), hi128);
11034        let s64 = _mm_add_ps(s128, _mm_movehl_ps(s128, s128));
11035        let s32 = _mm_add_ss(s64, _mm_shuffle_ps::<1>(s64, s64));
11036        let mut sum = _mm_cvtss_f32(s32);
11037        while j < n {
11038            sum += (*wp.add(j) as i8) as f32 * *xp.add(j);
11039            j += 1;
11040        }
11041        sum
11042    }
11043}
11044
11045/// int8(weight)·int8(activation) → i32 via AVX2 maddubs — the x86
11046/// analogue of the SDOT A8W8 path. `maddubs` takes u8×i8, so the
11047/// standard sign trick applies: |w| × sign(x, w) ≡ w × x per lane.
11048/// Pair saturation is safe: |w|≤128, |x|≤127 → 2·128·127 < 32767.
11049#[cfg(target_arch = "x86_64")]
11050#[target_feature(enable = "avx2")]
11051unsafe fn dot_i8_i8_avx2(w: &[u8], xq: &[i8]) -> i32 {
11052    // SAFETY: callers uphold slice-length contracts (see call sites).
11053    unsafe {
11054        use core::arch::x86_64::*;
11055        let n = w.len();
11056        let ones = _mm256_set1_epi16(1);
11057        let mut acc = _mm256_setzero_si256();
11058        let mut j = 0usize;
11059        while j + 32 <= n {
11060            let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
11061            let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
11062            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
11063            acc = _mm256_add_epi32(acc, _mm256_madd_epi16(p16, ones));
11064            j += 32;
11065        }
11066        let hi128 = _mm256_extracti128_si256::<1>(acc);
11067        let s128 = _mm_add_epi32(_mm256_castsi256_si128(acc), hi128);
11068        let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
11069        let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
11070        let mut s = _mm_cvtsi128_si32(s32);
11071        while j < n {
11072            s += (w[j] as i8) as i32 * xq[j] as i32;
11073            j += 1;
11074        }
11075        s
11076    }
11077}
11078
11079/// smmla 2×4: one instruction covers a 2-row × 2-activation × 8-deep
11080/// tile (32 MACs vs sdot's 16) — the weight pair loads once per 8-k
11081/// slice as a combined 2×8 register and meets two activation pairs.
11082#[cfg(target_arch = "aarch64")]
11083#[target_feature(enable = "neon,i8mm")]
11084unsafe fn dot_i8_smmla_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
11085    // SAFETY: callers uphold slice-length contracts.
11086    unsafe {
11087        use core::arch::aarch64::*;
11088        use core::arch::asm;
11089        let n = w0.len();
11090        let w0p = w0.as_ptr() as *const i8;
11091        let w1p = w1.as_ptr() as *const i8;
11092        // acc01 holds [c(r0,x0) c(r0,x1) c(r1,x0) c(r1,x1)]; acc23 the
11093        // same for x2/x3.
11094        let mut acc01 = vdupq_n_s32(0);
11095        let mut acc23 = vdupq_n_s32(0);
11096        let mut i = 0usize;
11097        while i + 8 <= n {
11098            let wa = vcombine_s8(vld1_s8(w0p.add(i)), vld1_s8(w1p.add(i)));
11099            let xb01 = vcombine_s8(
11100                vld1_s8(xs[0].as_ptr().add(i)),
11101                vld1_s8(xs[1].as_ptr().add(i)),
11102            );
11103            let xb23 = vcombine_s8(
11104                vld1_s8(xs[2].as_ptr().add(i)),
11105                vld1_s8(xs[3].as_ptr().add(i)),
11106            );
11107            asm!(
11108                "smmla {a01:v}.4s, {w:v}.16b, {x01:v}.16b",
11109                "smmla {a23:v}.4s, {w:v}.16b, {x23:v}.16b",
11110                a01 = inout(vreg) acc01, a23 = inout(vreg) acc23,
11111                w = in(vreg) wa, x01 = in(vreg) xb01, x23 = in(vreg) xb23,
11112                options(pure, nomem, nostack),
11113            );
11114            i += 8;
11115        }
11116        let mut out = [[0i32; 4]; 2];
11117        let a01: [i32; 4] = core::mem::transmute(acc01);
11118        let a23: [i32; 4] = core::mem::transmute(acc23);
11119        out[0][0] = a01[0];
11120        out[0][1] = a01[1];
11121        out[1][0] = a01[2];
11122        out[1][1] = a01[3];
11123        out[0][2] = a23[0];
11124        out[0][3] = a23[1];
11125        out[1][2] = a23[2];
11126        out[1][3] = a23[3];
11127        if i < n {
11128            for (k, x) in xs.iter().enumerate() {
11129                for j in i..n {
11130                    out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
11131                    out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
11132                }
11133            }
11134        }
11135        out
11136    }
11137}
11138
11139/// ARM twin of the x86 blocked prefill GEMM: two weight rows stay in
11140/// registers across four activation streams, eight sdot accumulators.
11141/// (The per-row form re-read each W row once per activation.)
11142#[cfg(target_arch = "aarch64")]
11143#[target_feature(enable = "neon,dotprod")]
11144unsafe fn dot_i8_sdot_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
11145    // SAFETY: callers uphold slice-length contracts.
11146    unsafe {
11147        use core::arch::aarch64::*;
11148        use core::arch::asm;
11149        let n = w0.len();
11150        let w0p = w0.as_ptr() as *const i8;
11151        let w1p = w1.as_ptr() as *const i8;
11152        let mut acc = [[vdupq_n_s32(0); 4]; 2];
11153        let mut i = 0usize;
11154        while i + 16 <= n {
11155            let wv0 = vld1q_s8(w0p.add(i));
11156            let wv1 = vld1q_s8(w1p.add(i));
11157            for (k, x) in xs.iter().enumerate() {
11158                let xv = vld1q_s8(x.as_ptr().add(i));
11159                let (mut a0, mut a1) = (acc[0][k], acc[1][k]);
11160                asm!(
11161                    "sdot {a0:v}.4s, {w0:v}.16b, {x:v}.16b",
11162                    "sdot {a1:v}.4s, {w1:v}.16b, {x:v}.16b",
11163                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
11164                    w0 = in(vreg) wv0, w1 = in(vreg) wv1, x = in(vreg) xv,
11165                    options(pure, nomem, nostack),
11166                );
11167                acc[0][k] = a0;
11168                acc[1][k] = a1;
11169            }
11170            i += 16;
11171        }
11172        let mut out = [[0i32; 4]; 2];
11173        for r in 0..2 {
11174            for k in 0..4 {
11175                out[r][k] = vaddvq_s32(acc[r][k]);
11176            }
11177        }
11178        if i < n {
11179            for (k, x) in xs.iter().enumerate() {
11180                for j in i..n {
11181                    out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
11182                    out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
11183                }
11184            }
11185        }
11186        out
11187    }
11188}
11189
11190/// Blocked 2 weight rows × 4 activations for the prefill GEMM
11191/// (roadmap P0: packed panels + multi-row accumulators). The two rows'
11192/// abs() live in registers across all four activation streams; the
11193/// sign-fixup is recomputed per pair (the price of the maddubs trick).
11194/// Returns raw i8·i8 dots; the caller applies scales and outliers.
11195#[cfg(target_arch = "x86_64")]
11196#[target_feature(enable = "avx2")]
11197unsafe fn dot_i8_i8_avx2_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
11198    // SAFETY: callers uphold slice-length contracts.
11199    unsafe {
11200        use core::arch::x86_64::*;
11201        let n = w0.len();
11202        let ones = _mm256_set1_epi16(1);
11203        let mut acc = [[_mm256_setzero_si256(); 4]; 2];
11204        let mut j = 0usize;
11205        while j + 32 <= n {
11206            let wv0 = _mm256_loadu_si256(w0.as_ptr().add(j) as *const __m256i);
11207            let wv1 = _mm256_loadu_si256(w1.as_ptr().add(j) as *const __m256i);
11208            let aw0 = _mm256_abs_epi8(wv0);
11209            let aw1 = _mm256_abs_epi8(wv1);
11210            for (k, x) in xs.iter().enumerate() {
11211                let xv = _mm256_loadu_si256(x.as_ptr().add(j) as *const __m256i);
11212                let p0 = _mm256_maddubs_epi16(aw0, _mm256_sign_epi8(xv, wv0));
11213                acc[0][k] = _mm256_add_epi32(acc[0][k], _mm256_madd_epi16(p0, ones));
11214                let p1 = _mm256_maddubs_epi16(aw1, _mm256_sign_epi8(xv, wv1));
11215                acc[1][k] = _mm256_add_epi32(acc[1][k], _mm256_madd_epi16(p1, ones));
11216            }
11217            j += 32;
11218        }
11219        let mut out = [[0i32; 4]; 2];
11220        for r in 0..2 {
11221            for k in 0..4 {
11222                let a = acc[r][k];
11223                let hi128 = _mm256_extracti128_si256::<1>(a);
11224                let s128 = _mm_add_epi32(_mm256_castsi256_si128(a), hi128);
11225                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
11226                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
11227                out[r][k] = _mm_cvtsi128_si32(s32);
11228            }
11229        }
11230        if j < n {
11231            for (k, x) in xs.iter().enumerate() {
11232                for i in j..n {
11233                    out[0][k] += (w0[i] as i8) as i32 * x[i] as i32;
11234                    out[1][k] += (w1[i] as i8) as i32 * x[i] as i32;
11235                }
11236            }
11237        }
11238        out
11239    }
11240}
11241
11242/// AVX2/VNNI q8 row dot with exact outlier correction (x86 mirror of
11243/// `row_dot_sdot` — same A8W8 contract). With AVX-512 VNNI the row goes
11244/// through the bias trick: Σ(w+128)·x via pure `vpdpbusd` (no per-lane
11245/// sign fixups), corrected by −128·Σx with Σx precomputed per split.
11246#[cfg(target_arch = "x86_64")]
11247#[inline]
11248fn row_dot_avx2(row: &[u8], act: &SplitAct) -> f32 {
11249    let dot = if avx512vnni_enabled() && row.len() >= 64 {
11250        (unsafe { dot_u8p128_i8_vnni(row, &act.xq) }) - 128 * act.xsum
11251    } else {
11252        unsafe { dot_i8_i8_avx2(row, &act.xq) }
11253    };
11254    let mut acc = dot as f32 * act.sx;
11255    for &(j, xv) in &act.outliers {
11256        acc += (row[j] as i8) as f32 * xv;
11257    }
11258    acc
11259}
11260
11261/// Σ (w[i]+128)·x[i] via pure `vpdpbusd` — the caller subtracts
11262/// 128·Σx. Four independent accumulators (dpbusd is ~5-cycle latency;
11263/// a single-acc loop runs latency-bound, measured on Granite Rapids).
11264#[cfg(target_arch = "x86_64")]
11265#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
11266unsafe fn dot_u8p128_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
11267    // SAFETY: callers uphold slice-length contracts (see call sites).
11268    unsafe {
11269        use core::arch::x86_64::*;
11270        let n = w.len();
11271        let flip = _mm512_set1_epi8(-128); // XOR 0x80: i8 w → u8 (w+128)
11272        #[inline(always)]
11273        unsafe fn step(
11274            w: *const u8,
11275            x: *const i8,
11276            flip: core::arch::x86_64::__m512i,
11277            acc: core::arch::x86_64::__m512i,
11278        ) -> core::arch::x86_64::__m512i {
11279            unsafe {
11280                use core::arch::x86_64::*;
11281                let wv = _mm512_xor_si512(_mm512_loadu_si512(w as *const _), flip);
11282                _mm512_dpbusd_epi32(acc, wv, _mm512_loadu_si512(x as *const _))
11283            }
11284        }
11285        let (mut a0, mut a1, mut a2, mut a3) = (
11286            _mm512_setzero_si512(),
11287            _mm512_setzero_si512(),
11288            _mm512_setzero_si512(),
11289            _mm512_setzero_si512(),
11290        );
11291        let mut j = 0usize;
11292        while j + 256 <= n {
11293            a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
11294            a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), flip, a1);
11295            a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), flip, a2);
11296            a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), flip, a3);
11297            j += 256;
11298        }
11299        while j + 64 <= n {
11300            a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
11301            j += 64;
11302        }
11303        let mut total = _mm512_reduce_add_epi32(_mm512_add_epi32(
11304            _mm512_add_epi32(a0, a1),
11305            _mm512_add_epi32(a2, a3),
11306        ));
11307        // Scalar tail: (w as i8) + 128 ≡ (w as u8) ^ 0x80.
11308        while j < n {
11309            total += ((w[j] ^ 0x80) as i32) * xq[j] as i32;
11310            j += 1;
11311        }
11312        total
11313    }
11314}
11315
11316/// One q4 row via AVX2: nibbles → centered i8 (unpacklo/hi restores the
11317/// writer's flat order, same as the NEON vzip pair), maddubs against
11318/// the pre-quantized activation group, × the group's f16 scale. Pair
11319/// saturation safe: |w|≤8, |x|≤127 → 2·8·127 ≪ 32767. Mirror of
11320/// `dot_q4_row_sdot`.
11321#[cfg(target_arch = "x86_64")]
11322#[target_feature(enable = "avx2")]
11323unsafe fn dot_q4_row_avx2(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
11324    // SAFETY: callers uphold slice-length contracts (16 packed bytes and
11325    // 2 scale bytes per group; xq.len() == gpr·GROUP_SIZE).
11326    unsafe {
11327        use core::arch::x86_64::*;
11328        let lomask = _mm_set1_epi8(0x0F);
11329        let eight = _mm256_set1_epi8(8);
11330        let ones = _mm256_set1_epi16(1);
11331        let mut acc = 0f32;
11332        for gi in 0..gpr {
11333            let g = g0 + gi;
11334            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
11335            let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
11336            let lo = _mm_and_si128(b, lomask);
11337            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
11338            let w = _mm256_sub_epi8(
11339                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
11340                eight,
11341            );
11342            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
11343            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
11344            let d = _mm256_madd_epi16(p16, ones);
11345            let hi128 = _mm256_extracti128_si256::<1>(d);
11346            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
11347            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
11348            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
11349            acc += _mm_cvtsi128_si32(s32) as f32 * s;
11350        }
11351        acc
11352    }
11353}
11354
11355/// Two-activation q4 row via AVX2: nibbles unpacked ONCE per group,
11356/// both activations dotted against the same centered i8 register.
11357#[cfg(target_arch = "x86_64")]
11358#[target_feature(enable = "avx2")]
11359unsafe fn dot_q4_row_avx2_2(
11360    packed: &[u8],
11361    scales: &[u8],
11362    g0: usize,
11363    gpr: usize,
11364    xq1: &[i8],
11365    xq2: &[i8],
11366) -> (f32, f32) {
11367    // SAFETY: callers uphold slice-length contracts (see dot_q4_row_avx2).
11368    unsafe {
11369        use core::arch::x86_64::*;
11370        let lomask = _mm_set1_epi8(0x0F);
11371        let eight = _mm256_set1_epi8(8);
11372        let ones = _mm256_set1_epi16(1);
11373        let (mut acc1, mut acc2) = (0f32, 0f32);
11374        #[inline(always)]
11375        unsafe fn hsum(d: core::arch::x86_64::__m256i) -> i32 {
11376            unsafe {
11377                use core::arch::x86_64::*;
11378                let hi128 = _mm256_extracti128_si256::<1>(d);
11379                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
11380                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
11381                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
11382                _mm_cvtsi128_si32(s32)
11383            }
11384        }
11385        for gi in 0..gpr {
11386            let g = g0 + gi;
11387            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
11388            let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
11389            let lo = _mm_and_si128(b, lomask);
11390            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
11391            let w = _mm256_sub_epi8(
11392                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
11393                eight,
11394            );
11395            let aw = _mm256_abs_epi8(w);
11396            let x1 = _mm256_loadu_si256(xq1.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
11397            let x2 = _mm256_loadu_si256(xq2.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
11398            let d1 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x1, w)), ones);
11399            let d2 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x2, w)), ones);
11400            acc1 += hsum(d1) as f32 * s;
11401            acc2 += hsum(d2) as f32 * s;
11402        }
11403        (acc1, acc2)
11404    }
11405}
11406
11407/// One q8 row range via AVX2 (x86 mirror of `q8_range_sdot`).
11408#[cfg(target_arch = "x86_64")]
11409fn q8_range_avx2(
11410    q: &[u8],
11411    row_scale: &[f32],
11412    act: &SplitAct,
11413    cols: usize,
11414    out_addr: SendMut,
11415    start: usize,
11416    end: usize,
11417) {
11418    for o in start..end {
11419        let v = row_dot_avx2(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
11420        // SAFETY: disjoint row ranges per worker.
11421        unsafe { *out_addr.at(o) = v };
11422    }
11423}
11424
11425/// Two-input q8 row range via AVX2 (x86 mirror of `q8_range2_sdot`).
11426#[cfg(target_arch = "x86_64")]
11427#[allow(clippy::too_many_arguments)]
11428fn q8_range2_avx2(
11429    q: &[u8],
11430    row_scale: &[f32],
11431    a1: &SplitAct,
11432    a2: &SplitAct,
11433    cols: usize,
11434    p1: SendMut,
11435    p2: SendMut,
11436    start: usize,
11437    end: usize,
11438) {
11439    for o in start..end {
11440        let row = &q[o * cols..(o + 1) * cols];
11441        // SAFETY: disjoint row ranges per worker.
11442        unsafe {
11443            *p1.at(o) = row_dot_avx2(row, a1) * row_scale[o];
11444            *p2.at(o) = row_dot_avx2(row, a2) * row_scale[o];
11445        }
11446    }
11447}
11448
11449// ───────────────────── A8W8 SDOT path (port of vmfcore, ×1.78 decode) ─────────────────────
11450
11451/// ARMv8.6 i8mm (smmla): 32 int8 MACs per instruction vs sdot's 16 —
11452/// yet MEASURED 2.4× SLOWER than the blocked sdot on Apple silicon
11453/// (108 vs 264 GF/s): the on-the-fly vcombine packing and the two-
11454/// accumulator dependency chain swamp the MAC advantage, and Apple's
11455/// four SIMD pipes already keep sdot fed. OPT-IN (CMF_I8MM=1) for
11456/// field trials on Cortex-A710/X-class parts with two pipes, where the
11457/// balance may differ; a pre-interleaved weight layout (repack infra)
11458/// is the known path if it ever earns its keep.
11459#[cfg(target_arch = "aarch64")]
11460fn i8mm_enabled() -> bool {
11461    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
11462    *ON.get_or_init(|| {
11463        std::env::var("CMF_I8MM").map(|v| v == "1").unwrap_or(false)
11464            && std::arch::is_aarch64_feature_detected!("i8mm")
11465    })
11466}
11467
11468/// SDOT enabled? Default ON when the CPU has ARMv8.2 dotprod;
11469/// `CMF_SDOT=0` disables (falls back to i8×f32 NEON).
11470/// (On non-ARM release builds only the test tolerance switch calls it.)
11471#[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
11472fn sdot_enabled() -> bool {
11473    if FLOAT_ACTIVATIONS.get() {
11474        return false;
11475    }
11476    use std::sync::OnceLock;
11477    static ON: OnceLock<bool> = OnceLock::new();
11478    *ON.get_or_init(|| {
11479        let want = std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true);
11480        if !want {
11481            return false;
11482        }
11483
11484        #[cfg(target_arch = "aarch64")]
11485        {
11486            if std::arch::is_aarch64_feature_detected!("dotprod") {
11487                return true;
11488            }
11489            #[cfg(target_os = "android")]
11490            {
11491                if let Ok(cpuinfo) = std::fs::read_to_string("/proc/cpuinfo") {
11492                    if cpuinfo.lines().any(|l| {
11493                        (l.starts_with("Features") || l.starts_with("features"))
11494                            && l.contains("asimddp")
11495                    }) {
11496                        return true;
11497                    }
11498                }
11499            }
11500            false
11501        }
11502        #[cfg(not(target_arch = "aarch64"))]
11503        {
11504            false
11505        }
11506    })
11507}
11508
11509/// Two-field activation split (≡ vmfcore `q8_split_prep`): outlier
11510/// channels (>8·rms) are computed exactly in f32; the bulk (outliers
11511/// zeroed → clean absmax) goes through int8 SDOT. Computed ONCE per
11512/// matvec, shared by all rows/workers.
11513struct SplitAct {
11514    xq: Vec<i8>,
11515    sx: f32,
11516    outliers: Vec<(usize, f32)>,
11517    /// Σ xq — the VNNI bias-trick correction (`(w+128)·x` sums need
11518    /// `−128·Σx`); one i32 per split, computed once per matvec.
11519    #[cfg_attr(not(target_arch = "x86_64"), allow(dead_code))]
11520    xsum: i32,
11521}
11522
11523thread_local! {
11524    /// Recycled xq buffers: split_act runs for every matvec (~200/token)
11525    /// and its hidden-size allocation was steady-state heap churn.
11526    static XQ_FREE: std::cell::RefCell<Vec<Vec<i8>>> =
11527        const { std::cell::RefCell::new(Vec::new()) };
11528}
11529
11530impl Drop for SplitAct {
11531    fn drop(&mut self) {
11532        let buf = std::mem::take(&mut self.xq);
11533        if buf.capacity() > 0 {
11534            XQ_FREE.with(|f| {
11535                let mut f = f.borrow_mut();
11536                if f.len() < 16 {
11537                    f.push(buf);
11538                }
11539            });
11540        }
11541    }
11542}
11543
11544thread_local! {
11545    /// One scratch row per WORKER, kept for the life of the thread.
11546    ///
11547    /// The kernels take a row of group scales per dispatch, and a fresh
11548    /// `vec![0f32; gpr]` inside the closure is one allocation per worker per
11549    /// dispatch — on the release checkpoint about six thousand a token, a
11550    /// quarter of everything the benchmark counts.
11551    static KROW: UnsafeCell<Vec<f32>> = const { UnsafeCell::new(Vec::new()) };
11552    /// Two scratch rows for a gate/up pair.  Q2TP and Q4TP fused SwiGLU
11553    /// decode each projection's ladder independently, but both ladders have
11554    /// the same lifetime; retaining them per worker removes two heap trips
11555    /// from every dispatched expert pair.
11556    static KROWS: UnsafeCell<[Vec<f32>; 2]> =
11557        const { UnsafeCell::new([Vec::new(), Vec::new()]) };
11558}
11559
11560/// Borrow `n` floats of the calling worker's scratch. Nothing inside a
11561/// kernel body borrows it again, which is what keeps the RefCell honest.
11562#[inline]
11563fn with_krow<R>(n: usize, f: impl FnOnce(&mut [f32]) -> R) -> R {
11564    KROW.with(|s| {
11565        // SAFETY: `KROW` is thread-local and this function does not recurse;
11566        // each worker has exclusive access to its own scratch row.
11567        let b = unsafe { &mut *s.get() };
11568        if b.len() < n {
11569            b.resize(n, 0.0);
11570        }
11571        f(&mut b[..n])
11572    })
11573}
11574
11575/// `t.round().clamp(-127.0, 127.0) as i8`, bit for bit, without the libm
11576/// call. On baseline x86-64 (no SSE4.1 `roundps`) `f32::round` is a
11577/// function call per element, and split_act runs it over every hidden
11578/// state before every matvec: measured 27 us a call on a 2048-wide
11579/// activation on an EPYC 7763 — 5.4 ms of a 55 ms decode token, all of
11580/// it on the caller's thread while thirty workers wait. Clamping first is
11581/// equivalent (round is monotonic and ±127 are integers), and after the
11582/// clamp `t - trunc(t)` is exact, so the half-away-from-zero decision is
11583/// the one `round` makes. NaN clamps to NaN and converts to 0, as before.
11584/// The loop vectorizes (cvttps2dq + compare/select).
11585#[inline(always)]
11586fn q8_round(t: f32) -> i8 {
11587    let t = t.clamp(-127.0, 127.0);
11588    let i = t as i32;
11589    let f = t - i as f32;
11590    let r = if f >= 0.5 {
11591        i + 1
11592    } else if f <= -0.5 {
11593        i - 1
11594    } else {
11595        i
11596    };
11597    r as i8
11598}
11599
11600#[inline]
11601fn with_krows<R>(n: usize, f: impl FnOnce(&mut [f32], &mut [f32]) -> R) -> R {
11602    KROWS.with(|s| {
11603        // SAFETY: KROWS is thread-local and no kernel body recursively calls
11604        // this helper; each participant owns its two vectors exclusively.
11605        let b = unsafe { &mut *s.get() };
11606        for row in b.iter_mut() {
11607            if row.len() < n {
11608                row.resize(n, 0.0);
11609            }
11610        }
11611        let (a, b) = b.split_at_mut(1);
11612        f(&mut a[0][..n], &mut b[0][..n])
11613    })
11614}
11615
11616#[inline]
11617fn silu_mul_limited(mut gate: f32, mut up: f32, limit: f32) -> f32 {
11618    if limit > 0.0 {
11619        up = up.clamp(-limit, limit);
11620        gate = gate.min(limit);
11621    }
11622    gate / (1.0 + (-gate).exp()) * up
11623}
11624
11625fn split_act(x: &[f32]) -> SplitAct {
11626    let _prof = crate::cpuprof::time(crate::cpuprof::Slot::SplitAct);
11627    let n = x.len();
11628    let rms = (x.iter().map(|&v| (v * v) as f64).sum::<f64>() / n.max(1) as f64).sqrt() as f32;
11629    let thr = 8.0 * rms;
11630    // One pass: collect outliers and the bulk absmax (outliers excluded —
11631    // identical to the old zero-then-fold over a copied buffer, minus the
11632    // full-vector copy).
11633    let mut outliers: Vec<(usize, f32)> = Vec::new();
11634    let mut amax = 0f32;
11635    for (j, &v) in x.iter().enumerate() {
11636        let a = v.abs();
11637        if a > thr {
11638            outliers.push((j, v));
11639        } else if a > amax {
11640            amax = a;
11641        }
11642    }
11643    let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
11644    let inv = 1.0 / sx;
11645    let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
11646    xq.clear();
11647    xq.reserve(n);
11648    if outliers.is_empty() {
11649        xq.extend(
11650            x.iter()
11651                .map(|&v| q8_round(v * inv)),
11652        );
11653    } else {
11654        // Outlier slots quantize to 0 (their exact term is added later).
11655        xq.extend(x.iter().map(|&v| {
11656            if v.abs() > thr {
11657                0
11658            } else {
11659                q8_round(v * inv)
11660            }
11661        }));
11662    }
11663    let xsum = xq.iter().map(|&v| v as i32).sum();
11664    SplitAct {
11665        xq,
11666        sx,
11667        outliers,
11668        xsum,
11669    }
11670}
11671
11672fn split_act_q8_2f(x: &[f32], col: &[f32]) -> SplitAct {
11673    let _prof = crate::cpuprof::time(crate::cpuprof::Slot::SplitAct);
11674    let n = x.len();
11675    let rms = (x
11676        .iter()
11677        .zip(col)
11678        .map(|(&a, &c)| {
11679            let v = a * c;
11680            (v * v) as f64
11681        })
11682        .sum::<f64>()
11683        / n.max(1) as f64)
11684        .sqrt() as f32;
11685    let thr = 8.0 * rms;
11686
11687    let mut outliers = Vec::new();
11688    let mut amax = 0f32;
11689    for (j, (&a, &c)) in x.iter().zip(col).enumerate() {
11690        let v = a * c;
11691        let s = v.abs();
11692        if s > thr {
11693            outliers.push((j, v));
11694        } else if s > amax {
11695            amax = s;
11696        }
11697    }
11698
11699    let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
11700    let inv = 1.0 / sx;
11701    let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
11702    xq.clear();
11703    xq.reserve(n);
11704    if outliers.is_empty() {
11705        xq.extend(
11706            x.iter()
11707                .zip(col)
11708                .map(|(&a, &c)| q8_round((a * c) * inv)),
11709        );
11710    } else {
11711        xq.extend(x.iter().zip(col).map(|(&a, &c)| {
11712            let v = a * c;
11713            if v.abs() > thr {
11714                0
11715            } else {
11716                q8_round(v * inv)
11717            }
11718        }));
11719    }
11720    let xsum = xq.iter().map(|&v| v as i32).sum();
11721    SplitAct {
11722        xq,
11723        sx,
11724        outliers,
11725        xsum,
11726    }
11727}
11728
11729/// int8(weight)·int8(activation) → i32 via `sdot` (inline asm — the
11730/// vdotq intrinsic is unstable; port of vmfcore `dot_i8_sdot`).
11731#[cfg(target_arch = "aarch64")]
11732#[target_feature(enable = "neon,dotprod")]
11733unsafe fn dot_i8_sdot(w: &[u8], xq: &[i8]) -> i32 {
11734    // SAFETY: callers uphold slice-length contracts (see call sites).
11735    unsafe {
11736        use core::arch::aarch64::*;
11737        use core::arch::asm;
11738        let wp = w.as_ptr() as *const i8;
11739        let n = w.len();
11740        let (mut a0, mut a1, mut a2, mut a3) = (
11741            vdupq_n_s32(0),
11742            vdupq_n_s32(0),
11743            vdupq_n_s32(0),
11744            vdupq_n_s32(0),
11745        );
11746        let mut i = 0;
11747        while i + 64 <= n {
11748            let (w0, x0) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
11749            let (w1, x1) = (vld1q_s8(wp.add(i + 16)), vld1q_s8(xq.as_ptr().add(i + 16)));
11750            let (w2, x2) = (vld1q_s8(wp.add(i + 32)), vld1q_s8(xq.as_ptr().add(i + 32)));
11751            let (w3, x3) = (vld1q_s8(wp.add(i + 48)), vld1q_s8(xq.as_ptr().add(i + 48)));
11752            asm!(
11753                "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
11754                "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
11755                "sdot {a2:v}.4s, {w2:v}.16b, {x2:v}.16b",
11756                "sdot {a3:v}.4s, {w3:v}.16b, {x3:v}.16b",
11757                a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
11758                w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
11759                w2 = in(vreg) w2, x2 = in(vreg) x2, w3 = in(vreg) w3, x3 = in(vreg) x3,
11760                options(pure, nomem, nostack),
11761            );
11762            i += 64;
11763        }
11764        while i + 16 <= n {
11765            let (wv, xv) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
11766            asm!("sdot {a:v}.4s, {w:v}.16b, {x:v}.16b",
11767                 a = inout(vreg) a0, w = in(vreg) wv, x = in(vreg) xv, options(pure, nomem, nostack));
11768            i += 16;
11769        }
11770        let mut s = vaddvq_s32(vaddq_s32(vaddq_s32(a0, a1), vaddq_s32(a2, a3)));
11771        while i < n {
11772            s += (*wp.add(i)) as i32 * xq[i] as i32;
11773            i += 1;
11774        }
11775        s
11776    }
11777}
11778
11779/// Row-blocked SDOT: 4 output rows per pass — the activation chunk is
11780/// loaded once and reused, 4 independent accumulators hide sdot latency
11781/// (port of vmfcore `dot_i8_sdot_4rows`).
11782#[cfg(target_arch = "aarch64")]
11783#[target_feature(enable = "neon,dotprod")]
11784unsafe fn dot_i8_sdot_4rows(w0: &[u8], w1: &[u8], w2: &[u8], w3: &[u8], xq: &[i8]) -> [i32; 4] {
11785    // SAFETY: callers uphold slice-length contracts (see call sites).
11786    unsafe {
11787        use core::arch::aarch64::*;
11788        use core::arch::asm;
11789        let n = xq.len();
11790        let px = xq.as_ptr();
11791        let (p0, p1, p2, p3) = (
11792            w0.as_ptr() as *const i8,
11793            w1.as_ptr() as *const i8,
11794            w2.as_ptr() as *const i8,
11795            w3.as_ptr() as *const i8,
11796        );
11797        let (mut a0, mut a1, mut a2, mut a3) = (
11798            vdupq_n_s32(0),
11799            vdupq_n_s32(0),
11800            vdupq_n_s32(0),
11801            vdupq_n_s32(0),
11802        );
11803        let mut i = 0;
11804        while i + 16 <= n {
11805            let x = vld1q_s8(px.add(i));
11806            let v0 = vld1q_s8(p0.add(i));
11807            let v1 = vld1q_s8(p1.add(i));
11808            let v2 = vld1q_s8(p2.add(i));
11809            let v3 = vld1q_s8(p3.add(i));
11810            asm!(
11811                "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
11812                "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
11813                "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
11814                "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
11815                a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
11816                v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
11817                options(pure, nomem, nostack),
11818            );
11819            i += 16;
11820        }
11821        let mut r = [
11822            vaddvq_s32(a0),
11823            vaddvq_s32(a1),
11824            vaddvq_s32(a2),
11825            vaddvq_s32(a3),
11826        ];
11827        while i < n {
11828            let xi = *px.add(i) as i32;
11829            r[0] += (*p0.add(i)) as i32 * xi;
11830            r[1] += (*p1.add(i)) as i32 * xi;
11831            r[2] += (*p2.add(i)) as i32 * xi;
11832            r[3] += (*p3.add(i)) as i32 * xi;
11833            i += 1;
11834        }
11835        r
11836    }
11837}
11838
11839/// 4 interleaved rows in one pass: the repacked group is [r0[c], r1[c],
11840/// r2[c], r3[c]] per 16-byte chunk, so each iteration reads ONE 64-byte
11841/// line plus the shared activation chunk — a single sequential weight
11842/// stream per worker. Per-row accumulation is the same one-accumulator
11843/// scheme as `dot_i8_sdot_4rows`; integer sums are exact, so outputs
11844/// are bit-identical to the mmap-layout kernel.
11845#[cfg(target_arch = "aarch64")]
11846#[target_feature(enable = "neon,dotprod")]
11847unsafe fn dot_i8_sdot_4rows_il(g: &[u8], xq: &[i8]) -> [i32; 4] {
11848    // SAFETY: callers uphold slice-length contracts (g.len() == 4·n,
11849    // n % 16 == 0 — guaranteed by the repack gate).
11850    unsafe {
11851        use core::arch::aarch64::*;
11852        use core::arch::asm;
11853        let n = xq.len();
11854        let px = xq.as_ptr();
11855        let pg = g.as_ptr() as *const i8;
11856        let (mut a0, mut a1, mut a2, mut a3) = (
11857            vdupq_n_s32(0),
11858            vdupq_n_s32(0),
11859            vdupq_n_s32(0),
11860            vdupq_n_s32(0),
11861        );
11862        let mut i = 0;
11863        while i + 16 <= n {
11864            let x = vld1q_s8(px.add(i));
11865            let base = pg.add(4 * i);
11866            let v0 = vld1q_s8(base);
11867            let v1 = vld1q_s8(base.add(16));
11868            let v2 = vld1q_s8(base.add(32));
11869            let v3 = vld1q_s8(base.add(48));
11870            asm!(
11871                "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
11872                "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
11873                "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
11874                "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
11875                a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
11876                v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
11877                options(pure, nomem, nostack),
11878            );
11879            i += 16;
11880        }
11881        [
11882            vaddvq_s32(a0),
11883            vaddvq_s32(a1),
11884            vaddvq_s32(a2),
11885            vaddvq_s32(a3),
11886        ]
11887    }
11888}
11889
11890/// One q8 row range via SDOT (4-row blocks + tail) — the body of
11891/// `qmatvec`'s hot loop, extracted so multi-matrix jobs can drive the
11892/// SAME kernel for several tensors under one pool dispatch. `rep` — the
11893/// load-time interleaved repack (empty = mmap layout only); rows outside
11894/// full 4-row groups always come from the mmap layout.
11895#[cfg(target_arch = "aarch64")]
11896fn q8_range_sdot(
11897    q: &[u8],
11898    rep: &[u8],
11899    row_scale: &[f32],
11900    act: &SplitAct,
11901    cols: usize,
11902    out_addr: SendMut,
11903    start: usize,
11904    end: usize,
11905) {
11906    let mut o = start;
11907    // Leading rows to the group boundary (repack path only): the pool
11908    // splits row ranges arbitrarily, groups are absolute.
11909    if !rep.is_empty() {
11910        while o < end && o % 4 != 0 {
11911            let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
11912            unsafe { *out_addr.at(o) = v };
11913            o += 1;
11914        }
11915    }
11916    while o + 4 <= end {
11917        let r = if rep.is_empty() {
11918            unsafe {
11919                dot_i8_sdot_4rows(
11920                    &q[o * cols..(o + 1) * cols],
11921                    &q[(o + 1) * cols..(o + 2) * cols],
11922                    &q[(o + 2) * cols..(o + 3) * cols],
11923                    &q[(o + 3) * cols..(o + 4) * cols],
11924                    &act.xq,
11925                )
11926            }
11927        } else {
11928            unsafe { dot_i8_sdot_4rows_il(&rep[o * cols..(o + 4) * cols], &act.xq) }
11929        };
11930        for k in 0..4 {
11931            let mut acc = r[k] as f32 * act.sx;
11932            for &(j, xv) in &act.outliers {
11933                acc += (q[(o + k) * cols + j] as i8) as f32 * xv;
11934            }
11935            // SAFETY: disjoint row ranges per worker.
11936            unsafe { *out_addr.at(o + k) = acc * row_scale[o + k] };
11937        }
11938        o += 4;
11939    }
11940    while o < end {
11941        let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
11942        unsafe { *out_addr.at(o) = v };
11943        o += 1;
11944    }
11945}
11946
11947/// Two-input q8 row range via SDOT — `qmatvec2`'s hot loop, extracted
11948/// for the fused pair multi-matrix job (`matvec2_many`).
11949#[cfg(target_arch = "aarch64")]
11950#[allow(clippy::too_many_arguments)]
11951fn q8_range2_sdot(
11952    q: &[u8],
11953    row_scale: &[f32],
11954    a1: &SplitAct,
11955    a2: &SplitAct,
11956    cols: usize,
11957    p1: SendMut,
11958    p2: SendMut,
11959    start: usize,
11960    end: usize,
11961) {
11962    for o in start..end {
11963        let row = &q[o * cols..(o + 1) * cols];
11964        // SAFETY: disjoint row ranges per worker.
11965        unsafe {
11966            *p1.at(o) = row_dot_sdot(row, a1) * row_scale[o];
11967            *p2.at(o) = row_dot_sdot(row, a2) * row_scale[o];
11968        }
11969    }
11970}
11971
11972/// Two-input q8 row range, f32 kernel (non-SDOT) — same extraction.
11973#[allow(clippy::too_many_arguments)]
11974fn q8_range2_f32(
11975    q: &[u8],
11976    row_scale: &[f32],
11977    x1: &[f32],
11978    x2: &[f32],
11979    cols: usize,
11980    p1: SendMut,
11981    p2: SendMut,
11982    start: usize,
11983    end: usize,
11984) {
11985    for o in start..end {
11986        let row = &q[o * cols..(o + 1) * cols];
11987        // SAFETY: disjoint row ranges per worker.
11988        unsafe {
11989            *p1.at(o) = dot_i8_f32(row, x1) * row_scale[o];
11990            *p2.at(o) = dot_i8_f32(row, x2) * row_scale[o];
11991        }
11992    }
11993}
11994
11995/// Scalar/NEON-f32 q8 row range (non-SDOT platforms) — same extraction.
11996fn q8_range_f32(
11997    q: &[u8],
11998    row_scale: &[f32],
11999    xs: &[f32],
12000    cols: usize,
12001    out_addr: SendMut,
12002    start: usize,
12003    end: usize,
12004) {
12005    for o in start..end {
12006        let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
12007        // SAFETY: disjoint row ranges per worker.
12008        unsafe { *out_addr.at(o) = v };
12009    }
12010}
12011
12012/// One q8 row against a split activation, portable: the per-arch fast
12013/// dots where they exist, the exact scalar loop elsewhere. The scalar
12014/// arm is also the test oracle for both fast arms.
12015#[inline]
12016fn q8_row_dot(row: &[u8], act: &SplitAct) -> f32 {
12017    #[cfg(target_arch = "aarch64")]
12018    return row_dot_sdot(row, act);
12019    #[cfg(target_arch = "x86_64")]
12020    return row_dot_avx2(row, act);
12021    #[allow(unreachable_code)]
12022    q8_row_dot_scalar(row, act)
12023}
12024
12025#[allow(dead_code)]
12026fn q8_row_dot_scalar(row: &[u8], act: &SplitAct) -> f32 {
12027    let mut acc = 0i32;
12028    for (k, &b) in row.iter().enumerate() {
12029        acc += (b as i8) as i32 * act.xq[k] as i32;
12030    }
12031    let mut acc = acc as f32 * act.sx;
12032    for &(j, xv) in &act.outliers {
12033        acc += (row[j] as i8) as f32 * xv;
12034    }
12035    acc
12036}
12037
12038/// SDOT row dot with exact outlier correction:
12039/// `dot = sdot(w, xq)·sx + Σ_outl w[j]·x[j]` (then × row_scale by caller).
12040#[cfg(target_arch = "aarch64")]
12041#[inline]
12042fn row_dot_sdot(row: &[u8], act: &SplitAct) -> f32 {
12043    let mut acc = unsafe { dot_i8_sdot(row, &act.xq) } as f32 * act.sx;
12044    for &(j, xv) in &act.outliers {
12045        acc += (row[j] as i8) as f32 * xv;
12046    }
12047    acc
12048}
12049
12050/// One q4 row via SDOT: each 32-group's nibbles unpack to centered i8
12051/// (nib−8 ∈ [−8,7]), int8×int8 `sdot` against the pre-quantized
12052/// activation group, × the group's f16 scale. Returns Σ_g dot_g·s_g;
12053/// the caller multiplies by the activation scale and adds the exact
12054/// outlier terms (port of vmfcore `dot_q4_block_sdot`, +23% measured).
12055/// Nibble order matches the writer: element 2k = low nibble, 2k+1 = high
12056/// → zip(lo,hi) restores flat order.
12057#[cfg(target_arch = "aarch64")]
12058#[target_feature(enable = "neon,dotprod")]
12059unsafe fn dot_q4_row_sdot(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
12060    // SAFETY: callers uphold slice-length contracts (16 packed bytes and
12061    // 2 scale bytes per group; xq.len() == gpr·GROUP_SIZE).
12062    unsafe {
12063        use core::arch::aarch64::*;
12064        use core::arch::asm;
12065        let lomask = vdupq_n_u8(0x0F);
12066        let eight = vdupq_n_s8(8);
12067        let mut acc = 0f32;
12068        for gi in 0..gpr {
12069            let g = g0 + gi;
12070            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
12071            let b = vld1q_u8(packed.as_ptr().add(g * 16));
12072            let lo = vandq_u8(b, lomask);
12073            let hi = vshrq_n_u8::<4>(b);
12074            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
12075            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
12076            let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
12077            let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
12078            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
12079            asm!(
12080                "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
12081                "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
12082                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
12083                e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
12084                options(pure, nomem, nostack),
12085            );
12086            acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
12087        }
12088        acc
12089    }
12090}
12091
12092/// Two-activation q4 row via SDOT: the nibble unpack (the expensive
12093/// part) happens ONCE per group; both pre-quantized activations are
12094/// dotted against the same centered i8 registers. Per-lane math matches
12095/// `dot_q4_row_sdot` exactly.
12096#[cfg(target_arch = "aarch64")]
12097#[target_feature(enable = "neon,dotprod")]
12098unsafe fn dot_q4_row_sdot2(
12099    packed: &[u8],
12100    scales: &[u8],
12101    g0: usize,
12102    gpr: usize,
12103    xq1: &[i8],
12104    xq2: &[i8],
12105) -> (f32, f32) {
12106    // SAFETY: callers uphold slice-length contracts (16 packed bytes and
12107    // 2 scale bytes per group; xq*.len() == gpr·GROUP_SIZE).
12108    unsafe {
12109        use core::arch::aarch64::*;
12110        use core::arch::asm;
12111        let lomask = vdupq_n_u8(0x0F);
12112        let eight = vdupq_n_s8(8);
12113        let (mut acc1, mut acc2) = (0f32, 0f32);
12114        for gi in 0..gpr {
12115            let g = g0 + gi;
12116            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
12117            let b = vld1q_u8(packed.as_ptr().add(g * 16));
12118            let lo = vandq_u8(b, lomask);
12119            let hi = vshrq_n_u8::<4>(b);
12120            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
12121            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
12122            let x10 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE));
12123            let x11 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE + 16));
12124            let x20 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE));
12125            let x21 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE + 16));
12126            let (mut a0, mut a1, mut b0, mut b1) = (
12127                vdupq_n_s32(0),
12128                vdupq_n_s32(0),
12129                vdupq_n_s32(0),
12130                vdupq_n_s32(0),
12131            );
12132            asm!(
12133                "sdot {a0:v}.4s, {e0:v}.16b, {x10:v}.16b",
12134                "sdot {a1:v}.4s, {e1:v}.16b, {x11:v}.16b",
12135                "sdot {b0:v}.4s, {e0:v}.16b, {x20:v}.16b",
12136                "sdot {b1:v}.4s, {e1:v}.16b, {x21:v}.16b",
12137                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
12138                b0 = inout(vreg) b0, b1 = inout(vreg) b1,
12139                e0 = in(vreg) e0, e1 = in(vreg) e1,
12140                x10 = in(vreg) x10, x11 = in(vreg) x11,
12141                x20 = in(vreg) x20, x21 = in(vreg) x21,
12142                options(pure, nomem, nostack),
12143            );
12144            acc1 += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
12145            acc2 += vaddvq_s32(vaddq_s32(b0, b1)) as f32 * s;
12146        }
12147        (acc1, acc2)
12148    }
12149}
12150
12151// ───────────────────── fused int8 kernels ─────────────────────
12152
12153/// `acc += w · row` where the row is centered i8 — NEON widen+fma on
12154/// aarch64, scalar elsewhere. The KV-cache q8 value path rides on this.
12155#[inline]
12156pub(crate) fn axpy_i8_f32(acc: &mut [f32], row: &[i8], w: f32) {
12157    #[cfg(target_arch = "aarch64")]
12158    unsafe {
12159        return axpy_i8_f32_neon(acc, row, w);
12160    }
12161    #[cfg(target_arch = "x86_64")]
12162    if avx2_enabled() {
12163        return unsafe { axpy_i8_f32_avx2(acc, row, w) };
12164    }
12165    #[allow(unreachable_code)]
12166    {
12167        for (a, &b) in acc.iter_mut().zip(row) {
12168            *a += w * b as f32;
12169        }
12170    }
12171}
12172
12173/// i8→f32 axpy via AVX2/FMA (x86 mirror of `axpy_i8_f32_neon`).
12174#[cfg(target_arch = "x86_64")]
12175#[target_feature(enable = "avx2,fma")]
12176unsafe fn axpy_i8_f32_avx2(acc: &mut [f32], row: &[i8], w: f32) {
12177    // SAFETY: callers uphold slice-length contracts (see call sites).
12178    unsafe {
12179        use core::arch::x86_64::*;
12180        let n = acc.len().min(row.len());
12181        let ap = acc.as_mut_ptr();
12182        let rp = row.as_ptr();
12183        let wv = _mm256_set1_ps(w);
12184        let mut j = 0usize;
12185        while j + 16 <= n {
12186            let rb = _mm_loadu_si128(rp.add(j) as *const __m128i);
12187            let lo = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(rb));
12188            let hi = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(rb)));
12189            let v0 = _mm256_fmadd_ps(wv, lo, _mm256_loadu_ps(ap.add(j)));
12190            let v1 = _mm256_fmadd_ps(wv, hi, _mm256_loadu_ps(ap.add(j + 8)));
12191            _mm256_storeu_ps(ap.add(j), v0);
12192            _mm256_storeu_ps(ap.add(j + 8), v1);
12193            j += 16;
12194        }
12195        while j < n {
12196            *ap.add(j) += w * (*rp.add(j)) as f32;
12197            j += 1;
12198        }
12199    }
12200}
12201
12202#[cfg(target_arch = "aarch64")]
12203#[target_feature(enable = "neon")]
12204unsafe fn axpy_i8_f32_neon(acc: &mut [f32], row: &[i8], w: f32) {
12205    // SAFETY: callers uphold slice-length contracts (see call sites).
12206    unsafe {
12207        use core::arch::aarch64::*;
12208        let n = acc.len().min(row.len());
12209        let ap = acc.as_mut_ptr();
12210        let rp = row.as_ptr();
12211        let wv = vdupq_n_f32(w);
12212        let mut j = 0usize;
12213        while j + 16 <= n {
12214            let rb = vld1q_s8(rp.add(j));
12215            let lo = vmovl_s8(vget_low_s8(rb));
12216            let hi = vmovl_s8(vget_high_s8(rb));
12217            for (off, half) in [(0, lo), (8, hi)] {
12218                let f0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(half)));
12219                let f1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(half)));
12220                let o = j + off;
12221                vst1q_f32(ap.add(o), vfmaq_f32(vld1q_f32(ap.add(o)), wv, f0));
12222                vst1q_f32(ap.add(o + 4), vfmaq_f32(vld1q_f32(ap.add(o + 4)), wv, f1));
12223            }
12224            j += 16;
12225        }
12226        while j < n {
12227            *ap.add(j) += w * (*rp.add(j)) as f32;
12228            j += 1;
12229        }
12230    }
12231}
12232
12233/// i8 row · f32 x. NEON on aarch64 (ported from vmfcore `dot_i8_f32_neon`,
12234/// ≈9× scalar), scalar elsewhere.
12235#[inline]
12236pub(crate) fn dot_i8_f32(w: &[u8], x: &[f32]) -> f32 {
12237    #[cfg(target_arch = "aarch64")]
12238    unsafe {
12239        return dot_i8_f32_neon(w, x);
12240    }
12241    #[cfg(target_arch = "x86_64")]
12242    if avx2_enabled() {
12243        return unsafe { dot_i8_f32_avx2(w, x) };
12244    }
12245    #[allow(unreachable_code)]
12246    {
12247        let mut sum = 0.0f32;
12248        for (j, &b) in w.iter().enumerate() {
12249            sum += (b as i8) as f32 * x[j];
12250        }
12251        sum
12252    }
12253}
12254
12255/// i8 row · (x ⊙ col_field) — the q8_2f row dot with the θ col-field
12256/// folded into the product (no prescaled copy of x). NEON on aarch64,
12257/// scalar elsewhere. Used by the active-neuron path `row_dot`.
12258#[inline]
12259fn dot_i8_col_f32(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
12260    #[cfg(target_arch = "aarch64")]
12261    unsafe {
12262        return dot_i8_col_f32_neon(w, x, col);
12263    }
12264    #[allow(unreachable_code)]
12265    {
12266        let mut sum = 0.0f32;
12267        for (j, &b) in w.iter().enumerate() {
12268            sum += (b as i8) as f32 * x[j] * col[j];
12269        }
12270        sum
12271    }
12272}
12273
12274#[cfg(target_arch = "aarch64")]
12275#[target_feature(enable = "neon")]
12276unsafe fn dot_i8_col_f32_neon(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
12277    // SAFETY: callers uphold slice-length contracts (see call sites).
12278    unsafe {
12279        use core::arch::aarch64::*;
12280        let n = x.len();
12281        let wp = w.as_ptr() as *const i8;
12282        let xp = x.as_ptr();
12283        let cp = col.as_ptr();
12284        let (mut a0, mut a1, mut a2, mut a3) = (
12285            vdupq_n_f32(0.0),
12286            vdupq_n_f32(0.0),
12287            vdupq_n_f32(0.0),
12288            vdupq_n_f32(0.0),
12289        );
12290        let mut j = 0usize;
12291        while j + 16 <= n {
12292            let wb = vld1q_s8(wp.add(j));
12293            let lo = vmovl_s8(vget_low_s8(wb));
12294            let hi = vmovl_s8(vget_high_s8(wb));
12295            let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
12296            let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
12297            let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
12298            let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
12299            a0 = vfmaq_f32(
12300                a0,
12301                w0,
12302                vmulq_f32(vld1q_f32(xp.add(j)), vld1q_f32(cp.add(j))),
12303            );
12304            a1 = vfmaq_f32(
12305                a1,
12306                w1,
12307                vmulq_f32(vld1q_f32(xp.add(j + 4)), vld1q_f32(cp.add(j + 4))),
12308            );
12309            a2 = vfmaq_f32(
12310                a2,
12311                w2,
12312                vmulq_f32(vld1q_f32(xp.add(j + 8)), vld1q_f32(cp.add(j + 8))),
12313            );
12314            a3 = vfmaq_f32(
12315                a3,
12316                w3,
12317                vmulq_f32(vld1q_f32(xp.add(j + 12)), vld1q_f32(cp.add(j + 12))),
12318            );
12319            j += 16;
12320        }
12321        let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
12322        while j < n {
12323            sum += (*wp.add(j)) as f32 * *xp.add(j) * *cp.add(j);
12324            j += 1;
12325        }
12326        sum
12327    }
12328}
12329
12330#[cfg(target_arch = "aarch64")]
12331#[target_feature(enable = "neon")]
12332unsafe fn dot_i8_f32_neon(w: &[u8], x: &[f32]) -> f32 {
12333    // SAFETY: callers uphold slice-length contracts (see call sites).
12334    unsafe {
12335        use core::arch::aarch64::*;
12336        let n = x.len();
12337        let wp = w.as_ptr() as *const i8;
12338        let xp = x.as_ptr();
12339        let (mut a0, mut a1, mut a2, mut a3) = (
12340            vdupq_n_f32(0.0),
12341            vdupq_n_f32(0.0),
12342            vdupq_n_f32(0.0),
12343            vdupq_n_f32(0.0),
12344        );
12345        let mut j = 0usize;
12346        while j + 16 <= n {
12347            let wb = vld1q_s8(wp.add(j));
12348            let lo = vmovl_s8(vget_low_s8(wb));
12349            let hi = vmovl_s8(vget_high_s8(wb));
12350            let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
12351            let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
12352            let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
12353            let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
12354            a0 = vfmaq_f32(a0, w0, vld1q_f32(xp.add(j)));
12355            a1 = vfmaq_f32(a1, w1, vld1q_f32(xp.add(j + 4)));
12356            a2 = vfmaq_f32(a2, w2, vld1q_f32(xp.add(j + 8)));
12357            a3 = vfmaq_f32(a3, w3, vld1q_f32(xp.add(j + 12)));
12358            j += 16;
12359        }
12360        let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
12361        while j < n {
12362            sum += (*wp.add(j)) as f32 * *xp.add(j);
12363            j += 1;
12364        }
12365        sum
12366    }
12367}
12368
12369#[allow(clippy::too_many_arguments)]
12370fn qmatvec(
12371    q: &[u8],
12372    rep: &[u8],
12373    row_scale: &[f32],
12374    x: &[f32],
12375    col_field: &[f32],
12376    dtype: TensorDtype,
12377    rows: usize,
12378    cols: usize,
12379    out: &mut [f32],
12380    pool: Option<&Pool>,
12381) {
12382    debug_assert_eq!(out.len(), rows);
12383    #[cfg(not(target_arch = "aarch64"))]
12384    let _ = rep;
12385
12386    #[cfg(target_arch = "aarch64")]
12387    if sdot_enabled() {
12388        let act = if dtype == TensorDtype::Q8_2f {
12389            split_act_q8_2f(x, col_field)
12390        } else {
12391            split_act(x)
12392        };
12393        let out_addr = SendMut(out.as_mut_ptr());
12394        let run_range = |start: usize, end: usize| {
12395            q8_range_sdot(q, rep, row_scale, &act, cols, out_addr, start, end)
12396        };
12397        match pool {
12398            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
12399            _ => run_range(0, rows),
12400        }
12401        return;
12402    }
12403    // x86 A8W8 via AVX2 maddubs — same quantized-activation contract as
12404    // the SDOT path (CMF_AVX2=0 keeps the exact i8×f32 loop).
12405    #[cfg(target_arch = "x86_64")]
12406    if avx2_a8w8_enabled() {
12407        let act = if dtype == TensorDtype::Q8_2f {
12408            split_act_q8_2f(x, col_field)
12409        } else {
12410            split_act(x)
12411        };
12412        let out_addr = SendMut(out.as_mut_ptr());
12413        let run_range = |start: usize, end: usize| {
12414            q8_range_avx2(q, row_scale, &act, cols, out_addr, start, end)
12415        };
12416        match pool {
12417            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
12418            _ => run_range(0, rows),
12419        }
12420        return;
12421    }
12422
12423    prescale_with(x, col_field, dtype, 1, |xs| {
12424        let out_addr = SendMut(out.as_mut_ptr());
12425        let run_range = move |start: usize, end: usize| {
12426            for o in start..end {
12427                let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
12428                // SAFETY: disjoint row ranges per worker.
12429                unsafe { *out_addr.at(o) = v };
12430            }
12431        };
12432        match pool {
12433            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
12434            _ => run_range(0, rows),
12435        }
12436    });
12437}
12438
12439#[allow(clippy::too_many_arguments)]
12440fn qmatvec2(
12441    q: &[u8],
12442    row_scale: &[f32],
12443    x1: &[f32],
12444    x2: &[f32],
12445    col_field: &[f32],
12446    dtype: TensorDtype,
12447    rows: usize,
12448    cols: usize,
12449    o1: &mut [f32],
12450    o2: &mut [f32],
12451    pool: Option<&Pool>,
12452) {
12453    #[cfg(target_arch = "aarch64")]
12454    if sdot_enabled() {
12455        let a1s = if dtype == TensorDtype::Q8_2f {
12456            split_act_q8_2f(x1, col_field)
12457        } else {
12458            split_act(x1)
12459        };
12460        let a2s = if dtype == TensorDtype::Q8_2f {
12461            split_act_q8_2f(x2, col_field)
12462        } else {
12463            split_act(x2)
12464        };
12465        let p1 = SendMut(o1.as_mut_ptr());
12466        let p2 = SendMut(o2.as_mut_ptr());
12467        let run_range = |start: usize, end: usize| {
12468            q8_range2_sdot(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
12469        };
12470        match pool {
12471            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
12472            _ => run_range(0, rows),
12473        }
12474        return;
12475    }
12476    #[cfg(target_arch = "x86_64")]
12477    if avx2_a8w8_enabled() {
12478        let a1s = if dtype == TensorDtype::Q8_2f {
12479            split_act_q8_2f(x1, col_field)
12480        } else {
12481            split_act(x1)
12482        };
12483        let a2s = if dtype == TensorDtype::Q8_2f {
12484            split_act_q8_2f(x2, col_field)
12485        } else {
12486            split_act(x2)
12487        };
12488        let p1 = SendMut(o1.as_mut_ptr());
12489        let p2 = SendMut(o2.as_mut_ptr());
12490        let run_range = |start: usize, end: usize| {
12491            q8_range2_avx2(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
12492        };
12493        match pool {
12494            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
12495            _ => run_range(0, rows),
12496        }
12497        return;
12498    }
12499
12500    prescale_with(x1, col_field, dtype, 1, |x1s| {
12501        prescale_with(x2, col_field, dtype, 2, |x2s| {
12502            let p1 = SendMut(o1.as_mut_ptr());
12503            let p2 = SendMut(o2.as_mut_ptr());
12504            let run_range = move |start: usize, end: usize| {
12505                for o in start..end {
12506                    let row = &q[o * cols..(o + 1) * cols];
12507                    let s1 = dot_i8_f32(row, x1s) * row_scale[o];
12508                    let s2 = dot_i8_f32(row, x2s) * row_scale[o];
12509                    // SAFETY: disjoint row ranges per worker.
12510                    unsafe {
12511                        *p1.at(o) = s1;
12512                        *p2.at(o) = s2;
12513                    }
12514                }
12515            };
12516            match pool {
12517                Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
12518                _ => run_range(0, rows),
12519            }
12520        });
12521    });
12522}
12523
12524#[derive(Clone, Copy)]
12525struct SendMut(*mut f32);
12526unsafe impl Send for SendMut {}
12527unsafe impl Sync for SendMut {}
12528
12529impl SendMut {
12530    #[inline]
12531    fn at(self, i: usize) -> *mut f32 {
12532        unsafe { self.0.add(i) }
12533    }
12534}
12535
12536#[cfg(test)]
12537mod tests {
12538    /// `q8_round` must be `round().clamp(±127) as i8` bit for bit: every
12539    /// half-integer, their neighbours one ulp either side, the clamp
12540    /// boundary, huge values, infinities and NaN, plus a dense sweep.
12541    #[test]
12542    fn q8_round_is_round_clamp() {
12543        let reference = |t: f32| t.round().clamp(-127.0, 127.0) as i8;
12544        let mut probe = vec![
12545            0.0f32,
12546            -0.0,
12547            f32::NAN,
12548            f32::INFINITY,
12549            f32::NEG_INFINITY,
12550            f32::MAX,
12551            f32::MIN,
12552            1e30,
12553            -1e30,
12554            f32::MIN_POSITIVE,
12555            -f32::MIN_POSITIVE,
12556        ];
12557        for k in -300i32..=300 {
12558            let h = k as f32 * 0.5;
12559            let up = f32::from_bits(h.to_bits() + 1);
12560            let down = f32::from_bits(h.to_bits().wrapping_sub(1));
12561            for t in [h, up, down] {
12562                probe.push(t);
12563                probe.push(-t);
12564            }
12565        }
12566        let mut t = -140.0f32;
12567        while t < 140.0 {
12568            probe.push(t);
12569            t += 0.000_731;
12570        }
12571        for t in probe {
12572            assert_eq!(q8_round(t), reference(t), "t = {t:e} ({:#x})", t.to_bits());
12573        }
12574    }
12575
12576    use super::*;
12577
12578    /// The f32 `matmat` walks the batch four positions a row at a time; every
12579    /// output must still be the scalar-order sum `matvec` gives, bit for bit
12580    /// (a 7-row batch covers the four-wide block and the remainder).
12581    #[test]
12582    fn f32_matmat_equals_matvec_bitwise() {
12583        let (rows, cols, b) = (5usize, 37usize, 7usize);
12584        let w: Vec<f32> = (0..rows * cols).map(|i| ((i * 7919 % 113) as f32 - 56.0) / 37.0).collect();
12585        let xs: Vec<f32> = (0..b * cols).map(|i| ((i * 104_729 % 97) as f32 - 48.0) / 29.0).collect();
12586        let t = QTensor::from_f32(w, rows, cols);
12587        let mut out = vec![0f32; b * rows];
12588        t.matmat(&xs, b, &mut out, None);
12589        for bi in 0..b {
12590            let mut one = vec![0f32; rows];
12591            t.matvec(&xs[bi * cols..(bi + 1) * cols], &mut one, None);
12592            for o in 0..rows {
12593                assert_eq!(out[bi * rows + o].to_bits(), one[o].to_bits(), "bi {bi} row {o}");
12594            }
12595        }
12596    }
12597
12598    #[test]
12599    fn q2tp_i8_dot_matches_exact_on_grid() {
12600        // On-grid activations (±1 → sx=1/127, xq=±127 dequantizes
12601        // exactly, no outliers) must make the integer path agree with
12602        // the exact scalar walk to f32 rounding.
12603        let (rows, cols) = (5, 64);
12604        let gpr = cols / GROUP_SIZE;
12605        // Synthetic codes plane + a flat ladder: scales_into is not under
12606        // test here, so drive dot_q2tp_row_i8 / q2tp_row_exact directly
12607        // with hand-made scales.
12608        let chunks: Vec<u8> = (0..rows * gpr * Q2TP_CHUNK)
12609            .map(|i| (i as u32).wrapping_mul(2654435761) as u8)
12610            .collect();
12611        let scales: Vec<f32> = (0..gpr).map(|g| 0.5 + g as f32 * 0.25).collect();
12612        let x: Vec<f32> = (0..cols)
12613            .map(|i| if i % 3 == 0 { -1.0 } else { 1.0 })
12614            .collect();
12615        let act = split_act(&x);
12616        assert!(
12617            act.outliers.is_empty(),
12618            "on-grid input must have no outliers"
12619        );
12620        let gsum = q1_group_sums(&act.xq, gpr);
12621        for r in 0..rows {
12622            let exact = q2tp_row_exact(&chunks, r, gpr, &x, &scales);
12623            let fast = dot_q2tp_row_i8(&chunks, r, gpr, &act.xq, &gsum, &scales) * act.sx;
12624            assert!(
12625                (exact - fast).abs() <= exact.abs() * 1e-5 + 1e-5,
12626                "row {r}: exact {exact} vs i8 {fast}"
12627            );
12628        }
12629    }
12630
12631    #[cfg(target_arch = "x86_64")]
12632    #[test]
12633    fn q2tp_avx2_dot_matches_scalar_for_random_patterns() {
12634        // Compare the release AVX2 integer dot against the scalar oracle over
12635        // arbitrary packed bytes/activation signs.  This guards the exact
12636        // table-load path used after rejecting a faster-looking decoder whose
12637        // full-checkpoint greedy output drifted.
12638        if !std::arch::is_x86_feature_detected!("avx2") {
12639            return;
12640        }
12641        let mut seed = 0x9e3779b9u32;
12642        let mut next = || {
12643            seed = seed.wrapping_mul(1664525).wrapping_add(1013904223);
12644            seed
12645        };
12646        for _ in 0..20_000 {
12647            let mut ch = [0u8; Q2TP_CHUNK];
12648            let mut x = [0i8; GROUP_SIZE];
12649            for b in &mut ch {
12650                *b = next() as u8;
12651            }
12652            for v in &mut x {
12653                *v = (next() >> 24) as i8;
12654            }
12655            let mut reference = 0i32;
12656            for (k, &b) in ch.iter().enumerate() {
12657                reference += (b & 3) as i32 * x[k * 4] as i32;
12658                reference += ((b >> 2) & 3) as i32 * x[k * 4 + 1] as i32;
12659                reference += ((b >> 4) & 3) as i32 * x[k * 4 + 2] as i32;
12660                reference += ((b >> 6) & 3) as i32 * x[k * 4 + 3] as i32;
12661            }
12662            // SAFETY: guarded by the runtime AVX2 feature check and fixed
12663            // 8-byte/32-byte slice lengths above.
12664            let got = unsafe { q2tp_code_dot_avx2(&ch, &x) };
12665            assert_eq!(got, reference, "packed q2 lane mismatch");
12666        }
12667    }
12668
12669    #[test]
12670    fn q2tp_affine_fuses_half_scale_correction_without_changing_raw_decode() {
12671        let (rows, cols) = (1usize, GROUP_SIZE);
12672        let mut bytes = vec![0u8; Q2TP_CHUNK + 4 + 1];
12673        // Repeating symbols 0,1,2,0 at unit scale.  q2tp's raw B is
12674        // (c-1.5), while the affine Prism operator is (c-1.0).
12675        bytes[..Q2TP_CHUNK].fill(0x24); // codes 0,1,2,0 in LSB-first order
12676        bytes[Q2TP_CHUNK..Q2TP_CHUNK + 2].copy_from_slice(&0u16.to_le_bytes());
12677        bytes[Q2TP_CHUNK + 2..Q2TP_CHUNK + 4].copy_from_slice(&0u16.to_le_bytes());
12678        bytes[Q2TP_CHUNK + 4] = 1; // dtype16 rung 1 = 1.0
12679        let x = vec![1.0f32; cols];
12680        let mut raw = vec![0.0f32; rows];
12681        let mut affine = vec![0.0f32; rows];
12682        q2tp_matvec_for_test(&bytes, &x, rows, cols, &mut raw);
12683        q2tp_affine_matvec_for_test(&bytes, &x, rows, cols, &mut affine);
12684        assert_eq!(raw, vec![-24.0]);
12685        assert_eq!(affine, vec![-8.0]);
12686        assert!((affine[0] - (raw[0] + 0.5 * cols as f32)).abs() < 1e-6);
12687    }
12688
12689    #[test]
12690    fn q8_row_dot_fast_matches_scalar() {
12691        // The per-arch fast dot must agree with the exact scalar oracle
12692        // (same contract the fused q8 FFN arm rides on).
12693        let cols = 96;
12694        let row: Vec<u8> = (0..cols)
12695            .map(|i| ((i * 37 % 251) - 125) as i8 as u8)
12696            .collect();
12697        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.13).sin()).collect();
12698        let act = split_act(&x);
12699        let fast = q8_row_dot(&row, &act);
12700        let scalar = q8_row_dot_scalar(&row, &act);
12701        assert!(
12702            (fast - scalar).abs() <= scalar.abs() * 1e-5 + 1e-5,
12703            "fast {fast} vs scalar {scalar}"
12704        );
12705    }
12706
12707    #[test]
12708    fn f32_matvec_matches_matvec_rows_bitexact() {
12709        let (rows, cols) = (300, 40);
12710        let w: Vec<f32> = (0..rows * cols).map(|i| (i as f32 * 0.017).sin()).collect();
12711        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.05).cos()).collect();
12712        let qt = QTensor::from_f32(w.clone(), rows, cols);
12713
12714        let mut a = vec![0.0f32; rows];
12715        matvec_rows(None, &w, &x, &mut a);
12716        let mut b = vec![0.0f32; rows];
12717        qt.matvec(&x, &mut b, None);
12718        assert_eq!(a, b);
12719    }
12720
12721    #[test]
12722    fn sdot_kernel_exact_on_grid() {
12723        // Activations already on the i8 grid (±1 with amax=1 → sx=1/127,
12724        // xq=±127 dequantizes EXACTLY) → the SDOT path must match the
12725        // exact f32 dot to float rounding. This isolates kernel
12726        // correctness from quantization noise.
12727        eprintln!("sdot_enabled = {}", sdot_enabled());
12728        let (rows, cols) = (9, 80); // odd rows → exercises 4-row + tail
12729        let w: Vec<u8> = (0..rows * cols)
12730            .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
12731            .collect();
12732        let scales: Vec<f32> = (0..rows).map(|o| 0.005 + o as f32 * 0.001).collect();
12733        let x: Vec<f32> = (0..cols)
12734            .map(|i| match i % 3 {
12735                0 => 1.0,
12736                1 => -1.0,
12737                _ => 0.0,
12738            })
12739            .collect();
12740        let mut a = vec![0.0f32; rows];
12741        qmatvec(
12742            &w,
12743            &[],
12744            &scales,
12745            &x,
12746            &[],
12747            TensorDtype::Q8Row,
12748            rows,
12749            cols,
12750            &mut a,
12751            None,
12752        );
12753        for o in 0..rows {
12754            let mut acc = 0.0f32;
12755            for j in 0..cols {
12756                acc += (w[o * cols + j] as i8) as f32 * x[j];
12757            }
12758            let expect = acc * scales[o];
12759            assert!(
12760                (a[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
12761                "row {o}: {} vs {expect}",
12762                a[o]
12763            );
12764        }
12765    }
12766
12767    #[test]
12768    fn q1_tbl_fast_path_matches_reference() {
12769        // gpr = 8 exercises the TBL pair-load fast loop, and the LAST
12770        // row's final 4-tile window trips the 4B-overread guard (the
12771        // payload ends exactly at the last tile) — both paths must
12772        // agree with the dequant reference.
12773        let (rows, cols) = (5, 256);
12774        let gpr = cols / GROUP_SIZE;
12775        let mut bytes = Vec::new();
12776        for t in 0..rows * gpr {
12777            let s = 0.007 + (t % 11) as f32 * 0.004;
12778            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
12779            for j in 0..4 {
12780                bytes.push(((t * 53 + j * 89 + 7) % 249) as u8);
12781            }
12782        }
12783        let x: Vec<f32> = (0..cols)
12784            .map(|i| if (i * 5) % 7 < 3 { 1.0 } else { -1.0 })
12785            .collect();
12786        let mut w = vec![0.0f32; rows * cols];
12787        cortiq_core::quant::dequant_q1(&bytes, &mut w);
12788        let mut got = vec![0.0f32; rows];
12789        q1_matvec(&bytes, &x, rows, cols, &mut got, None);
12790        for o in 0..rows {
12791            let expect: f32 = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
12792            assert!(
12793                (got[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
12794                "row {o}: {} vs {expect}",
12795                got[o]
12796            );
12797        }
12798        // Blocked 1×4 batch (b=5: one quad + remainder) must equal the
12799        // single-matvec path bit-for-bit.
12800        let b = 5usize;
12801        let mut xs_all = Vec::new();
12802        for bi in 0..b {
12803            xs_all.extend(x.iter().map(|v| if bi % 2 == 0 { *v } else { -*v }));
12804        }
12805        let mut mm = vec![0.0f32; b * rows];
12806        q1_matmat(&bytes, &xs_all, b, rows, cols, &mut mm, None);
12807        for bi in 0..b {
12808            let mut single = vec![0.0f32; rows];
12809            q1_matvec(
12810                &bytes,
12811                &xs_all[bi * cols..(bi + 1) * cols],
12812                rows,
12813                cols,
12814                &mut single,
12815                None,
12816            );
12817            assert_eq!(&mm[bi * rows..(bi + 1) * rows], &single[..], "stream {bi}");
12818        }
12819    }
12820
12821    #[test]
12822    fn q1_kernels_match_exact_reference() {
12823        // Synthetic q1 payload: 6-byte tiles [f16 scale][4B bits].
12824        let (rows, cols) = (7, 96);
12825        let gpr = cols / GROUP_SIZE;
12826        let mut bytes = Vec::new();
12827        for t in 0..rows * gpr {
12828            let s = 0.01 + (t % 13) as f32 * 0.003;
12829            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
12830            for j in 0..4 {
12831                bytes.push(((t * 31 + j * 97) % 251) as u8);
12832            }
12833        }
12834        // On-grid activations (±1, amax 1) → the SDOT path is exact.
12835        let x: Vec<f32> = (0..cols)
12836            .map(|i| if i % 3 == 0 { 1.0 } else { -1.0 })
12837            .collect();
12838        // Reference through the core dequant.
12839        let mut w = vec![0.0f32; rows * cols];
12840        cortiq_core::quant::dequant_q1(&bytes, &mut w);
12841        let mut expect = vec![0.0f32; rows];
12842        for o in 0..rows {
12843            expect[o] = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
12844        }
12845        let mut got = vec![0.0f32; rows];
12846        q1_matvec(&bytes, &x, rows, cols, &mut got, None);
12847        for o in 0..rows {
12848            assert!(
12849                (got[o] - expect[o]).abs() < 1e-3 * expect[o].abs().max(1e-3),
12850                "row {o}: {} vs {}",
12851                got[o],
12852                expect[o]
12853            );
12854        }
12855        // Pair and batch paths agree with the single path.
12856        let x2: Vec<f32> = x.iter().map(|v| -v).collect();
12857        let (mut a1, mut a2) = (vec![0.0f32; rows], vec![0.0f32; rows]);
12858        q1_matvec2(&bytes, &x, &x2, rows, cols, &mut a1, &mut a2, None);
12859        assert_eq!(a1, got);
12860        let mut xs = x.clone();
12861        xs.extend_from_slice(&x2);
12862        let mut mm = vec![0.0f32; 2 * rows];
12863        q1_matmat(&bytes, &xs, 2, rows, cols, &mut mm, None);
12864        assert_eq!(&mm[..rows], got.as_slice());
12865        assert_eq!(&mm[rows..], a2.as_slice());
12866    }
12867
12868    /// A Bonsai-shaped q1 payload (one f16 scale per 128-group, stored in
12869    /// all four tiles; pseudo-random signs) and its q1g repack.
12870    fn q1_q1g_pair(rows: usize, cols: usize, salt: usize) -> (Vec<u8>, Vec<u8>) {
12871        let tiles = rows * cols / GROUP_SIZE;
12872        let mut q1 = Vec::with_capacity(tiles * Q1_TILE);
12873        for t in 0..tiles {
12874            let s = 0.004 + ((t / 4 * 37 + salt) % 29) as f32 * 0.0011;
12875            q1.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
12876            let bits = ((t * 2_654_435_761 + salt * 40_503) as u32).rotate_left((t % 31) as u32);
12877            q1.extend_from_slice(&bits.to_le_bytes());
12878        }
12879        let rp = cortiq_core::quant::repack_q1_to_q1g(&q1, rows, cols, 0.0);
12880        assert_eq!(rp.differing, 0, "uniform scales repack exactly");
12881        (q1, rp.data)
12882    }
12883
12884    /// Activations off the int8 grid with two outlier channels, so the
12885    /// A8W8 path quantizes AND takes the exact outlier correction.
12886    fn q1g_test_x(cols: usize, salt: usize) -> Vec<f32> {
12887        let mut x: Vec<f32> = (0..cols)
12888            .map(|i| ((i * 13 + salt * 7) as f32 * 0.071).sin() * 0.9)
12889            .collect();
12890        x[3 % cols] = 41.0;
12891        x[(cols / 2 + 5) % cols] = -37.5;
12892        x
12893    }
12894
12895    /// Every CPU row kernel of q1g returns the q1 kernel's f32 BIT FOR BIT
12896    /// on a uniform-scale repack: A8W8 rows (SIMD of this machine, and the
12897    /// x86 AVX2/VNNI twins directly), the exact scalar rows (CMF_SDOT=0
12898    /// contract), the outlier term, the drivers (matvec, matvec2, matmat
12899    /// with the blocked 1×4 quad + remainder) — including gpr = 4 and the
12900    /// last row, where q1's pair load falls back to its scalar tail.
12901    #[test]
12902    fn q1g_cpu_kernels_equal_q1_bit_for_bit() {
12903        for &(rows, cols) in &[(3usize, 128usize), (5, 256), (9, 2048), (2, 6144)] {
12904            let gpr = cols / GROUP_SIZE;
12905            let (q1, q1g) = q1_q1g_pair(rows, cols, rows + cols);
12906            let (wa, wb) = (BinW::Q1(&q1), BinW::q1g(&q1g, rows, cols));
12907            let x = q1g_test_x(cols, rows);
12908            let act = split_act(&x);
12909            assert!(!act.outliers.is_empty(), "test wants the outlier path");
12910            let gsum = q1_group_sums(&act.xq, gpr);
12911            for r in 0..rows {
12912                let (a, b) = (wa.row_exact(r, gpr, &x), wb.row_exact(r, gpr, &x));
12913                assert_eq!(a.to_bits(), b.to_bits(), "{rows}x{cols} exact row {r}");
12914                let (a, b) = (
12915                    wa.dot_i8(r, gpr, &act.xq, &gsum),
12916                    wb.dot_i8(r, gpr, &act.xq, &gsum),
12917                );
12918                assert_eq!(a.to_bits(), b.to_bits(), "{rows}x{cols} i8 row {r}");
12919                for j in [0, 31, 32, 127, cols - 1] {
12920                    let (a, b) = (wa.outlier(r, gpr, j), wb.outlier(r, gpr, j));
12921                    assert_eq!(
12922                        (a.0.to_bits(), a.1.to_bits()),
12923                        (b.0.to_bits(), b.1.to_bits())
12924                    );
12925                }
12926                #[cfg(target_arch = "x86_64")]
12927                if avx2_enabled() {
12928                    let sc = cortiq_core::quant::q1g_scale_offset(rows, cols);
12929                    // SAFETY: AVX2 checked; payloads sized for (rows, cols).
12930                    unsafe {
12931                        let a = dot_q1_row_avx2(&q1, r, gpr, &act.xq, &gsum);
12932                        let b = dot_q1g_row_avx2(&q1g, sc, r, gpr, &act.xq, &gsum);
12933                        assert_eq!(a.to_bits(), b.to_bits(), "avx2 row {r}");
12934                    }
12935                    if std::is_x86_feature_detected!("avx512vnni")
12936                        && std::is_x86_feature_detected!("avx512bw")
12937                        && std::is_x86_feature_detected!("avx512vl")
12938                    {
12939                        // SAFETY: VNNI family detected just above.
12940                        unsafe {
12941                            let a = dot_q1_row_vnni(&q1, r, gpr, &act.xq, &gsum);
12942                            let b = dot_q1g_row_vnni(&q1g, sc, r, gpr, &act.xq, &gsum);
12943                            assert_eq!(a.to_bits(), b.to_bits(), "vnni row {r}");
12944                        }
12945                    }
12946                }
12947            }
12948            // Drivers: single, pair, batch (b = 5: one blocked quad + one).
12949            let (mut ga, mut gb) = (vec![0f32; rows], vec![0f32; rows]);
12950            bin_matvec(wa, &x, rows, cols, &mut ga, None);
12951            bin_matvec(wb, &x, rows, cols, &mut gb, None);
12952            assert_eq!(ga, gb, "{rows}x{cols} matvec");
12953            let x2 = q1g_test_x(cols, rows + 1);
12954            let (mut a1, mut a2, mut b1, mut b2) = (
12955                vec![0f32; rows],
12956                vec![0f32; rows],
12957                vec![0f32; rows],
12958                vec![0f32; rows],
12959            );
12960            bin_matvec2(wa, &x, &x2, rows, cols, &mut a1, &mut a2, None);
12961            bin_matvec2(wb, &x, &x2, rows, cols, &mut b1, &mut b2, None);
12962            assert_eq!((&a1, &a2), (&b1, &b2), "{rows}x{cols} matvec2");
12963            let b = 5usize;
12964            let xs: Vec<f32> = (0..b).flat_map(|k| q1g_test_x(cols, k * 3 + 2)).collect();
12965            let (mut ma, mut mb) = (vec![0f32; b * rows], vec![0f32; b * rows]);
12966            bin_matmat(wa, &xs, b, rows, cols, &mut ma, None);
12967            bin_matmat(wb, &xs, b, rows, cols, &mut mb, None);
12968            assert_eq!(ma, mb, "{rows}x{cols} matmat");
12969            // The blocked 1×4 kernels directly (the driver only takes them
12970            // under CMF_X86_BLOCKED / SDOT, a process-wide switch).
12971            let acts: Vec<(SplitAct, Vec<i32>)> = (0..4)
12972                .map(|k| {
12973                    let a = split_act(&xs[k * cols..(k + 1) * cols]);
12974                    let g = q1_group_sums(&a.xq, gpr);
12975                    (a, g)
12976                })
12977                .collect();
12978            let xq4: [&[i8]; 4] = std::array::from_fn(|k| acts[k].0.xq.as_slice());
12979            let gs4: [&[i32]; 4] = std::array::from_fn(|k| acts[k].1.as_slice());
12980            for r in 0..rows {
12981                #[cfg(target_arch = "x86_64")]
12982                if avx2_enabled() {
12983                    // SAFETY: AVX2 checked; payloads sized for (rows, cols).
12984                    let (a, b) = unsafe {
12985                        (
12986                            wa.dot_1x4_x86(r, gpr, xq4, gs4),
12987                            wb.dot_1x4_x86(r, gpr, xq4, gs4),
12988                        )
12989                    };
12990                    assert_eq!(a.map(f32::to_bits), b.map(f32::to_bits), "1x4 row {r}");
12991                }
12992                #[cfg(target_arch = "aarch64")]
12993                if sdot_enabled() {
12994                    // SAFETY: SDOT checked; payloads sized for (rows, cols).
12995                    let (a, b) = unsafe {
12996                        (
12997                            wa.dot_1x4_sdot(r, gpr, xq4, gs4),
12998                            wb.dot_1x4_sdot(r, gpr, xq4, gs4),
12999                        )
13000                    };
13001                    assert_eq!(a.map(f32::to_bits), b.map(f32::to_bits), "1x4 row {r}");
13002                }
13003                let _ = (r, &xq4, &gs4);
13004            }
13005        }
13006    }
13007
13008    /// CPU kernel A/B, q1 vs q1g matvec (A8W8 path of this machine, the
13009    /// exact path too), Bonsai-1.7B decode shapes, arms alternated every
13010    /// round, medians printed. `CMF_THREADS` sets the pool (default 8).
13011    /// `cargo test -p cortiq-engine --release --lib q1g_cpu_matvec_speed --
13012    /// --ignored --nocapture`
13013    #[test]
13014    #[ignore]
13015    fn q1g_cpu_matvec_speed() {
13016        use crate::pool::Pool;
13017        let threads = std::env::var("CMF_THREADS")
13018            .ok()
13019            .and_then(|v| v.parse().ok())
13020            .unwrap_or(8usize);
13021        let pool = Pool::new(threads);
13022        for &(name, rows, cols, reps) in &[
13023            ("gate/up_proj", 6144usize, 2048usize, 41usize),
13024            ("down_proj", 2048, 6144, 41),
13025            ("lm_head", 151_669, 2048, 9),
13026        ] {
13027            let (q1, q1g) = q1_q1g_pair(rows, cols, 3);
13028            let (wa, wb) = (BinW::Q1(&q1), BinW::q1g(&q1g, rows, cols));
13029            let x = q1g_test_x(cols, 1);
13030            let mut out = vec![0f32; rows];
13031            let mut time = |w: BinW| {
13032                let t0 = std::time::Instant::now();
13033                bin_matvec(w, &x, rows, cols, &mut out, Some(&pool));
13034                t0.elapsed().as_secs_f64()
13035            };
13036            let _ = (time(wa), time(wb));
13037            let (mut ta, mut tb) = (Vec::new(), Vec::new());
13038            for r in 0..reps {
13039                if r % 2 == 0 {
13040                    ta.push(time(wa));
13041                    tb.push(time(wb));
13042                } else {
13043                    tb.push(time(wb));
13044                    ta.push(time(wa));
13045                }
13046            }
13047            ta.sort_by(|a, b| a.partial_cmp(b).unwrap());
13048            tb.sort_by(|a, b| a.partial_cmp(b).unwrap());
13049            let (ma, mb) = (ta[reps / 2], tb[reps / 2]);
13050            eprintln!(
13051                "cpu x{threads} a8w8={} {name:>12} {rows}x{cols}: q1 {:8.1} us | q1g {:8.1} us | q1g/q1 time {:.3}",
13052                a8w8_enabled(),
13053                ma * 1e6,
13054                mb * 1e6,
13055                mb / ma
13056            );
13057        }
13058    }
13059
13060    /// The mapped QTensor paths a real Bonsai q1g file takes — matvec, the
13061    /// multi-projection job (all-q1g and q1/q1g mixed), the fused gate/up
13062    /// with SiLU, matmat, matvec2 and the embedding row — equal the q1
13063    /// tensor they were repacked from, bit for bit.
13064    #[test]
13065    fn q1g_mapped_tensor_paths_equal_q1() {
13066        use crate::pool::Pool;
13067        use cortiq_core::{CmfHeader, CmfModel, TensorSpec};
13068        let (rows, cols) = (300usize, 256usize);
13069        let (gq1, gq1g) = q1_q1g_pair(rows, cols, 1);
13070        let (uq1, uq1g) = q1_q1g_pair(rows, cols, 2);
13071        let header: CmfHeader = serde_json::from_value(serde_json::json!({
13072            "format": "cmf", "version": cortiq_core::CMF_VERSION,
13073            "arch": {"arch_name": "tiny-q1g", "hidden_size": cols,
13074                "intermediate_size": rows, "num_layers": 1, "num_attention_heads": 2,
13075                "num_kv_heads": 1, "head_dim": 128, "vocab_size": rows,
13076                "layer_types": ["FullAttention"], "rms_norm_eps": 1e-6,
13077                "max_position_embeddings": 8},
13078            "quant_type": "VBIT"
13079        }))
13080        .unwrap();
13081        let spec = |name: &str, dtype, data: &Vec<u8>| TensorSpec {
13082            name: name.into(),
13083            dtype,
13084            shape: vec![rows, cols],
13085            data: data.clone(),
13086        };
13087        let specs = [
13088            spec("g1", TensorDtype::Q1, &gq1),
13089            spec("u1", TensorDtype::Q1, &uq1),
13090            spec("gg", TensorDtype::Q1G, &gq1g),
13091            spec("ug", TensorDtype::Q1G, &uq1g),
13092        ];
13093        let dir = std::env::temp_dir().join(format!("cmf-q1g-mapped-{}", std::process::id()));
13094        std::fs::create_dir_all(&dir).unwrap();
13095        let path = dir.join("m.cmf");
13096        CmfModel::write(&path, &header, &specs, None, None).unwrap();
13097        let model = Arc::new(CmfModel::open(&path).unwrap());
13098        assert!(model.verify().is_empty(), "{:?}", model.verify());
13099        let t = |n: &str| QTensor::from_model(&model, n).unwrap();
13100        let (g1, u1, gg, ug) = (t("g1"), t("u1"), t("gg"), t("ug"));
13101        assert_eq!(gg.model_dtype(), Some(TensorDtype::Q1G));
13102        assert!(gg.is_q1() && gg.graph_weight().map(|w| w.2) == Some(10));
13103        let pool = Pool::new(3);
13104        let x = q1g_test_x(cols, 4);
13105        let run = |a: &QTensor| {
13106            let mut o = vec![0f32; rows];
13107            a.matvec(&x, &mut o, Some(&pool));
13108            o
13109        };
13110        assert_eq!(run(&g1), run(&gg), "matvec");
13111        let many = |a: &QTensor, b: &QTensor| {
13112            let (mut oa, mut ob) = (vec![0f32; rows], vec![0f32; rows]);
13113            QTensor::matvec_many([a, b], &x, [&mut oa, &mut ob], Some(&pool));
13114            (oa, ob)
13115        };
13116        let want = many(&g1, &u1);
13117        assert_eq!(want, many(&gg, &ug), "matvec_many q1g");
13118        assert_eq!(want, many(&g1, &ug), "matvec_many mixed");
13119        assert_eq!(want, (run(&g1), run(&u1)), "matvec_many == singles");
13120        let fused = |g: &QTensor, u: &QTensor| {
13121            let mut o = vec![0f32; rows];
13122            let ok = QTensor::matvec_silu_mul(g, u, &x, &mut o, Some(&pool));
13123            (ok, o)
13124        };
13125        let (ok_a, fa) = fused(&g1, &u1);
13126        let (ok_b, fb) = fused(&gg, &ug);
13127        assert_eq!(ok_a, ok_b, "q1g takes the fused arm whenever q1 does");
13128        assert_eq!(fa, fb, "fused gate/up");
13129        let b = 6usize;
13130        let xs: Vec<f32> = (0..b).flat_map(|k| q1g_test_x(cols, k + 9)).collect();
13131        let mm = |a: &QTensor| {
13132            let mut o = vec![0f32; b * rows];
13133            a.matmat(&xs, b, &mut o, Some(&pool));
13134            o
13135        };
13136        assert_eq!(mm(&g1), mm(&gg), "matmat");
13137        let x2 = q1g_test_x(cols, 77);
13138        let pair = |a: &QTensor| {
13139            let (mut o1, mut o2) = (vec![0f32; rows], vec![0f32; rows]);
13140            a.matvec2(&x, &x2, &mut o1, &mut o2, Some(&pool));
13141            (o1, o2)
13142        };
13143        assert_eq!(pair(&g1), pair(&gg), "matvec2");
13144        for r in [0, 1, rows / 2, rows - 1] {
13145            let (mut ra, mut rb) = (vec![0f32; cols], vec![0f32; cols]);
13146            g1.row_f32(r, &mut ra);
13147            gg.row_f32(r, &mut rb);
13148            assert_eq!(ra, rb, "embedding row {r}");
13149        }
13150        drop((g1, u1, gg, ug, model));
13151        let _ = std::fs::remove_dir_all(&dir);
13152    }
13153
13154    /// A `blr` tensor read from a real CMF file: it maps to the factored
13155    /// arm (not a dense copy), declines every device-graph accessor, and
13156    /// every QTensor entry point a layer may take — matvec, matvec2,
13157    /// matmat (table loop and GEMM sizes), the multi-tensor job mixed with
13158    /// a q8_row sibling, the fused gate/up refusal, rows and columns — equals
13159    /// the dense reference built by the core dequant.
13160    #[test]
13161    fn blr_mapped_tensor_paths_equal_dense() {
13162        use crate::pool::Pool;
13163        use cortiq_core::{CmfHeader, CmfModel, TensorSpec};
13164        let (rows, cols, r) = (96usize, 128usize, 16usize);
13165        let mut rng = crate::blr::tests::Rng(41);
13166        let gate = crate::blr::tests::make(&mut rng, rows, cols, r, 2);
13167        let up = crate::blr::tests::make(&mut rng, rows, cols, r, 1);
13168        // q8_row sibling: [int8 rows·cols][f16 row scale · rows].
13169        let mut q8: Vec<u8> = (0..rows * cols)
13170            .map(|i| (((i * 89) % 253) as i32 - 126) as i8 as u8)
13171            .collect();
13172        for o in 0..rows {
13173            q8.extend_from_slice(&cortiq_core::quant::f32_to_f16(0.002 + o as f32 * 1e-4).to_le_bytes());
13174        }
13175        let header: CmfHeader = serde_json::from_value(serde_json::json!({
13176            "format": "cmf", "version": cortiq_core::CMF_VERSION,
13177            "arch": {"arch_name": "tiny-blr", "hidden_size": cols,
13178                "intermediate_size": rows, "num_layers": 1, "num_attention_heads": 2,
13179                "num_kv_heads": 1, "head_dim": 64, "vocab_size": rows,
13180                "layer_types": ["FullAttention"], "rms_norm_eps": 1e-6,
13181                "max_position_embeddings": 8},
13182            "quant_type": "Q8_2F"
13183        }))
13184        .unwrap();
13185        let specs = [
13186            TensorSpec {
13187                name: "g".into(),
13188                dtype: TensorDtype::BinaryLowRank,
13189                shape: vec![rows, cols],
13190                data: gate.clone(),
13191            },
13192            TensorSpec {
13193                name: "u".into(),
13194                dtype: TensorDtype::BinaryLowRank,
13195                shape: vec![rows, cols],
13196                data: up.clone(),
13197            },
13198            TensorSpec {
13199                name: "q".into(),
13200                dtype: TensorDtype::Q8Row,
13201                shape: vec![rows, cols],
13202                data: q8,
13203            },
13204        ];
13205        let dir = std::env::temp_dir().join(format!("cmf-blr-mapped-{}", std::process::id()));
13206        std::fs::create_dir_all(&dir).unwrap();
13207        let path = dir.join("m.cmf");
13208        CmfModel::write(&path, &header, &specs, None, None).unwrap();
13209        let model = Arc::new(CmfModel::open(&path).unwrap());
13210        assert!(model.verify().is_empty(), "{:?}", model.verify());
13211        let t = |n: &str| QTensor::from_model(&model, n).unwrap();
13212        let (g, u, q) = (t("g"), t("u"), t("q"));
13213        assert!(g.is_blr() && u.is_blr() && !q.is_blr());
13214        assert!(g.graph_weight().is_none() && g.graph_weight_descriptor().is_none());
13215        assert!(g.as_f32().is_none() && g.model_idx().is_none());
13216        assert!(g.mapped_device_gemm().is_none() && g.q1_parts().is_none());
13217        assert!(!g.sparse_col_ok());
13218        assert_eq!((g.rows(), g.cols()), (rows, cols));
13219        let deq = |bytes: &[u8]| {
13220            let mut w = vec![0f32; rows * cols];
13221            cortiq_core::blr::dequant_blr(bytes, rows, cols, &mut w).unwrap();
13222            w
13223        };
13224        let (wg, wu) = (deq(&gate), deq(&up));
13225        let mv = |w: &[f32], x: &[f32]| -> Vec<f32> {
13226            (0..rows)
13227                .map(|i| (0..cols).map(|j| w[i * cols + j] * x[j]).sum())
13228                .collect()
13229        };
13230        let close = |a: &[f32], b: &[f32], what: &str| {
13231            let s = b.iter().map(|v| v * v).sum::<f32>().sqrt().max(1e-30);
13232            let e = a.iter().zip(b).map(|(x, y)| (x - y) * (x - y)).sum::<f32>().sqrt();
13233            assert!(e / s < 1e-5, "{what}: rel {}", e / s);
13234        };
13235        let pool = Pool::new(3);
13236        let x: Vec<f32> = (0..cols).map(|_| rng.f()).collect();
13237        let x2: Vec<f32> = (0..cols).map(|_| rng.f()).collect();
13238        let mut o = vec![0f32; rows];
13239        g.matvec(&x, &mut o, Some(&pool));
13240        close(&o, &mv(&wg, &x), "matvec");
13241        let (mut o1, mut o2) = (vec![0f32; rows], vec![0f32; rows]);
13242        g.matvec2(&x, &x2, &mut o1, &mut o2, Some(&pool));
13243        close(&o1, &mv(&wg, &x), "matvec2 a");
13244        close(&o2, &mv(&wg, &x2), "matvec2 b");
13245        for b in [2usize, 5] {
13246            let xs: Vec<f32> = (0..b * cols).map(|_| rng.f()).collect();
13247            let mut om = vec![0f32; b * rows];
13248            g.matmat(&xs, b, &mut om, Some(&pool));
13249            for t in 0..b {
13250                close(
13251                    &om[t * rows..(t + 1) * rows],
13252                    &mv(&wg, &xs[t * cols..(t + 1) * cols]),
13253                    "matmat",
13254                );
13255            }
13256        }
13257        // Multi-tensor job mixed with a q8_row sibling: per-tensor fallback,
13258        // the q8_row output identical to its own matvec.
13259        let (mut a, mut bq, mut c) = (vec![0f32; rows], vec![0f32; rows], vec![0f32; rows]);
13260        QTensor::matvec_many([&g, &q, &u], &x, [&mut a, &mut bq, &mut c], Some(&pool));
13261        close(&a, &mv(&wg, &x), "matvec_many blr gate");
13262        close(&c, &mv(&wu, &x), "matvec_many blr up");
13263        let mut qo = vec![0f32; rows];
13264        q.matvec(&x, &mut qo, Some(&pool));
13265        assert_eq!(bq, qo, "matvec_many q8_row sibling");
13266        let mut f = vec![0f32; rows];
13267        assert!(!QTensor::matvec_silu_mul(&g, &u, &x, &mut f, Some(&pool)));
13268        let mut row = vec![0f32; cols];
13269        g.row_f32(rows / 3, &mut row);
13270        close(&row, &wg[(rows / 3) * cols..(rows / 3 + 1) * cols], "row_f32");
13271        let mut scratch = vec![0f32; cols];
13272        let d = g.row_dot(7, &x, &mut scratch);
13273        let want: f32 = (0..cols).map(|j| wg[7 * cols + j] * x[j]).sum();
13274        assert!((d - want).abs() < 1e-4 * (1.0 + want.abs()), "row_dot");
13275        let mut col = vec![0f32; rows];
13276        g.add_col_scaled(cols / 2, 1.5, &mut col);
13277        let want: Vec<f32> = (0..rows).map(|i| 1.5 * wg[i * cols + cols / 2]).collect();
13278        close(&col, &want, "add_col_scaled");
13279        drop((g, u, q, model));
13280        let _ = std::fs::remove_dir_all(&dir);
13281    }
13282
13283    #[test]
13284    fn repack_is_bit_identical() {
13285        // The interleaved-repack kernel must produce EXACTLY the same
13286        // bits as the mmap-layout kernel: integer accumulation is order-
13287        // exact, the f32 epilogue is identical. Odd rows exercise the
13288        // tail; direct range calls exercise unaligned pool splits.
13289        let (rows, cols) = (267, 96); // 66 groups + 3 tail rows, cols % 16 == 0
13290        let w: Vec<u8> = (0..rows * cols)
13291            .map(|i| (((i * 89) % 253) as i32 - 126) as i8 as u8)
13292            .collect();
13293        let scales: Vec<f32> = (0..rows).map(|o| 0.003 + o as f32 * 0.0007).collect();
13294        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.37).sin() * 2.0).collect();
13295        let rep = q8_repack_layout(&w, rows, cols);
13296        // Group interleave round-trips.
13297        for g in 0..rows / 4 {
13298            for c in 0..cols / 16 {
13299                for lane in 0..4 {
13300                    assert_eq!(
13301                        &rep[g * 4 * cols + c * 64 + lane * 16
13302                            ..g * 4 * cols + c * 64 + lane * 16 + 16],
13303                        &w[(g * 4 + lane) * cols + c * 16..(g * 4 + lane) * cols + c * 16 + 16],
13304                    );
13305                }
13306            }
13307        }
13308        let mut a = vec![0.0f32; rows];
13309        qmatvec(
13310            &w,
13311            &[],
13312            &scales,
13313            &x,
13314            &[],
13315            TensorDtype::Q8Row,
13316            rows,
13317            cols,
13318            &mut a,
13319            None,
13320        );
13321        let mut b = vec![0.0f32; rows];
13322        qmatvec(
13323            &w,
13324            &rep,
13325            &scales,
13326            &x,
13327            &[],
13328            TensorDtype::Q8Row,
13329            rows,
13330            cols,
13331            &mut b,
13332            None,
13333        );
13334        assert_eq!(a, b, "full-range repack output diverged");
13335
13336        #[cfg(target_arch = "aarch64")]
13337        if sdot_enabled() {
13338            // Unaligned range split (pool workers get arbitrary bounds).
13339            let act = split_act(&x);
13340            let mut c1 = vec![0.0f32; rows];
13341            let mut c2 = vec![0.0f32; rows];
13342            q8_range_sdot(
13343                &w,
13344                &[],
13345                &scales,
13346                &act,
13347                cols,
13348                SendMut(c1.as_mut_ptr()),
13349                3,
13350                rows - 2,
13351            );
13352            q8_range_sdot(
13353                &w,
13354                &rep,
13355                &scales,
13356                &act,
13357                cols,
13358                SendMut(c2.as_mut_ptr()),
13359                3,
13360                rows - 2,
13361            );
13362            assert_eq!(c1, c2, "unaligned-range repack output diverged");
13363        }
13364    }
13365
13366    #[test]
13367    fn sdot_a8w8_noise_is_bounded() {
13368        // Off-grid activations: A8 quantization noise must stay small in
13369        // relative L2 over the whole output (realistic accuracy contract;
13370        // vmfcore measured argmax-identical decode on real models).
13371        let (rows, cols) = (16, 512);
13372        let w: Vec<u8> = (0..rows * cols)
13373            .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
13374            .collect();
13375        let scales = vec![0.01f32; rows];
13376        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.21).sin()).collect();
13377        let mut a = vec![0.0f32; rows];
13378        qmatvec(
13379            &w,
13380            &[],
13381            &scales,
13382            &x,
13383            &[],
13384            TensorDtype::Q8Row,
13385            rows,
13386            cols,
13387            &mut a,
13388            None,
13389        );
13390        let (mut num, mut den) = (0f64, 0f64);
13391        for o in 0..rows {
13392            let mut acc = 0.0f32;
13393            for j in 0..cols {
13394                acc += (w[o * cols + j] as i8) as f32 * x[j];
13395            }
13396            let expect = acc * scales[o];
13397            num += ((a[o] - expect) as f64).powi(2);
13398            den += (expect as f64).powi(2);
13399        }
13400        let rel = (num / den.max(1e-12)).sqrt();
13401        assert!(rel < 0.05, "A8W8 relative L2 error too high: {rel}");
13402    }
13403
13404    #[test]
13405    fn i8_dot_neon_matches_scalar() {
13406        let n = 100;
13407        let w: Vec<u8> = (0..n).map(|i| ((i * 37 + 11) % 251) as u8).collect();
13408        let x: Vec<f32> = (0..n).map(|i| (i as f32 * 0.13).sin()).collect();
13409        let mut scalar = 0.0f32;
13410        for j in 0..n {
13411            scalar += (w[j] as i8) as f32 * x[j];
13412        }
13413        let fast = dot_i8_f32(&w, &x);
13414        assert!((scalar - fast).abs() < 1e-3 * scalar.abs().max(1.0));
13415    }
13416
13417    /// Fused vbit matvec must match full dequant_vbit + dense matvec.
13418    #[test]
13419    fn vbitmatvec_matches_full_dequant() {
13420        let (rows, cols) = (6, 64);
13421        let ng = cols / GROUP_SIZE;
13422        // Hand-craft: bits per row, f16 scales, packed rows.
13423        let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
13424        let mut bytes = bits.clone();
13425        for g in 0..rows * ng {
13426            let s = 0.02 + 0.001 * g as f32;
13427            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
13428        }
13429        for r in 0..rows {
13430            let b = bits[r] as usize;
13431            let (mut acc, mut nb) = (0u64, 0usize);
13432            let mut rowbytes = Vec::new();
13433            for i in 0..cols {
13434                let v = ((i * 7 + r * 13) % (1 << b)) as u64;
13435                acc = (acc << b) | v;
13436                nb += b;
13437                while nb >= 8 {
13438                    nb -= 8;
13439                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
13440                }
13441            }
13442            if nb > 0 {
13443                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
13444            }
13445            bytes.extend_from_slice(&rowbytes);
13446        }
13447        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
13448
13449        let mut reference = vec![0f32; rows * cols];
13450        cortiq_core::quant::dequant_vbit(&bytes, rows, cols, &mut reference).unwrap();
13451        let mut expect = vec![0f32; rows];
13452        for r in 0..rows {
13453            expect[r] = reference[r * cols..(r + 1) * cols]
13454                .iter()
13455                .zip(&x)
13456                .map(|(w, xv)| w * xv)
13457                .sum();
13458        }
13459        let mut got = vec![0f32; rows];
13460        let offsets = vbit_row_offsets(&bytes, rows, cols);
13461        vbitmatvec(&bytes, &offsets, &x, rows, cols, &mut got, None);
13462        // SDOT path quantizes activations to i8 (A8W8): bounded noise,
13463        // same contract as q8 (exact path is pinned by CMF_SDOT=0 in
13464        // the golden-parity gate).
13465        let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
13466        let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
13467        for r in 0..rows {
13468            assert!(
13469                (got[r] - expect[r]).abs() < tol * scale,
13470                "row {r}: {} vs {}",
13471                got[r],
13472                expect[r]
13473            );
13474        }
13475    }
13476
13477    /// Fused q4 matvec must match the reference full-dequant + dense
13478    /// matvec bit-for-bit in structure (same f32 math, group order).
13479    /// vbit matmat: the blocked 1×4 leg must match the per-row path
13480    /// (paired env toggle; larger shape so both code paths engage).
13481    #[test]
13482    #[cfg(target_arch = "x86_64")]
13483    fn vbit_matmat_blocked_matches_per_row() {
13484        let (rows, cols, b) = (64usize, 128usize, 9usize);
13485        let ng = cols / GROUP_SIZE;
13486        let bits: Vec<u8> = (0..rows).map(|r| [3u8, 4, 5, 6][r % 4]).collect();
13487        let mut bytes = bits.clone();
13488        for g in 0..rows * ng {
13489            let sc = 0.02 + 0.0005 * g as f32;
13490            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
13491        }
13492        for r in 0..rows {
13493            let bw = bits[r] as usize;
13494            let (mut acc, mut nb) = (0u64, 0usize);
13495            let mut rowbytes = Vec::new();
13496            for i in 0..cols {
13497                let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
13498                acc = (acc << bw) | v;
13499                nb += bw;
13500                while nb >= 8 {
13501                    nb -= 8;
13502                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
13503                }
13504            }
13505            if nb > 0 {
13506                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
13507            }
13508            bytes.extend_from_slice(&rowbytes);
13509        }
13510        let x: Vec<f32> = (0..b * cols)
13511            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
13512            .collect();
13513        let offsets = vbit_row_offsets(&bytes, rows, cols);
13514        let mut y_a = vec![0f32; b * rows];
13515        let mut y_b = vec![0f32; b * rows];
13516        unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
13517        vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_a, None);
13518        unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
13519        vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_b, None);
13520        unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
13521        let max_d = y_a
13522            .iter()
13523            .zip(&y_b)
13524            .map(|(p, q)| (p - q).abs())
13525            .fold(0.0f32, f32::max);
13526        assert!(max_d < 1e-4, "vbit blocked ≠ per-row: max|Δ| = {max_d}");
13527    }
13528
13529    /// q4t blocked 1×4 (SDOT on ARM, AVX2 on x86) must equal the
13530    /// per-row path exactly: same nibble unpack, same group order,
13531    /// same f32 accumulation — batch == matvec bit-for-bit. b=9 covers
13532    /// two full 1×4 blocks plus a remainder through the single-row
13533    /// kernel. (Both paths produce identical output, so the shared
13534    /// CMF_X86_BLOCKED env var racing with other tests cannot flip
13535    /// the verdict — worst case both sides take the same path.)
13536    #[test]
13537    fn q4t_matmat_blocked_matches_per_row() {
13538        let (rows, cols, b) = (16usize, 64usize, 9usize);
13539        let gpr = cols / GROUP_SIZE;
13540        let mut bytes = vec![0u8; rows * gpr * Q4_TILE];
13541        for r in 0..rows {
13542            for g in 0..gpr {
13543                let t = (r * gpr + g) * Q4_TILE;
13544                let sc = 0.02 + 0.001 * (r * gpr + g) as f32;
13545                bytes[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
13546                for k in 0..16 {
13547                    bytes[t + 2 + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
13548                }
13549            }
13550        }
13551        let x: Vec<f32> = (0..b * cols)
13552            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
13553            .collect();
13554        let mut y_blk = vec![0f32; b * rows];
13555        let mut y_row = vec![0f32; b * rows];
13556        unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
13557        q4t_matmat(&bytes, &x, b, rows, cols, &mut y_blk, None);
13558        unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
13559        q4t_matmat(&bytes, &x, b, rows, cols, &mut y_row, None);
13560        unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
13561        assert_eq!(y_blk, y_row, "q4t blocked 1x4 ≠ per-row");
13562    }
13563
13564    /// The wide-batch Accelerate arm of q4t_matmat vs a brute-force
13565    /// f32 dequant matmul: both are f32 GEMMs, so only reduction
13566    /// order differs — tight tolerance.
13567    /// A synthetic q4tp payload: random nibbles plus a per-row ladder whose
13568    /// span varies row to row, so the codes actually exercise the full 0..31
13569    /// range rather than clustering on one rung.
13570    fn synth_q4tp(rows: usize, cols: usize) -> Vec<u8> {
13571        use cortiq_core::quant::{f32_to_f16, q4tp_code_stride, q4tp_put_code};
13572        let gpr = cols / GROUP_SIZE;
13573        let stride = q4tp_code_stride(gpr);
13574        let (params_off, codes_off, _) = q4tp_sections(rows, cols);
13575        let mut b = vec![0u8; codes_off + rows * stride];
13576        for r in 0..rows {
13577            for g in 0..gpr {
13578                let t = (r * gpr + g) * Q4TP_NIB;
13579                for k in 0..16 {
13580                    b[t + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
13581                }
13582            }
13583            let lo = -6.0 - 0.03 * (r % 17) as f32;
13584            let step = 0.01 + 0.004 * (r % 11) as f32;
13585            let p = params_off + r * 4;
13586            b[p..p + 2].copy_from_slice(&f32_to_f16(lo).to_le_bytes());
13587            b[p + 2..p + 4].copy_from_slice(&f32_to_f16(step).to_le_bytes());
13588            let crow = &mut b[codes_off + r * stride..codes_off + (r + 1) * stride];
13589            for g in 0..gpr {
13590                q4tp_put_code(crow, g, (r * 5 + g * 3) % 32);
13591            }
13592        }
13593        b
13594    }
13595
13596    /// The same weights re-expressed as q4_tiled, so the proven kernel can
13597    /// be the reference: each tile stores the ladder scale its code selects.
13598    /// Only the f16 rounding of that scale separates the two payloads.
13599    fn q4tp_as_q4t(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
13600        let gpr = cols / GROUP_SIZE;
13601        let v = Q4tpView::new(bytes, rows, cols);
13602        let mut out = vec![0u8; rows * gpr * Q4_TILE];
13603        let mut sc = vec![0f32; gpr];
13604        for r in 0..rows {
13605            v.scales_into(r, gpr, &mut sc);
13606            for g in 0..gpr {
13607                let t = (r * gpr + g) * Q4_TILE;
13608                let s = sc[g];
13609                out[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
13610                let src = (r * gpr + g) * Q4TP_NIB;
13611                out[t + 2..t + Q4_TILE].copy_from_slice(&v.nib[src..src + Q4TP_NIB]);
13612            }
13613        }
13614        out
13615    }
13616
13617    /// The exact (`CMF_SDOT=0`) path must reproduce `dequant_q4tp` to f32
13618    /// rounding — that scalar routine is the format's definition, and the
13619    /// kernels re-derive the scale from the ladder independently. Call the
13620    /// row kernel directly: `matmat` picks the int8 arm when a8w8 is on,
13621    /// so routing through it would test the other path by accident.
13622    #[test]
13623    fn q4tp_exact_path_matches_dequant_reference() {
13624        let (rows, cols) = (256usize, 512usize);
13625        let gpr = cols / GROUP_SIZE;
13626        let bytes = synth_q4tp(rows, cols);
13627        let mut w = vec![0f32; rows * cols];
13628        cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
13629
13630        let x: Vec<f32> = (0..cols)
13631            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
13632            .collect();
13633        let v = Q4tpView::new(&bytes, rows, cols);
13634        let mut sc = vec![0f32; gpr];
13635        for r in 0..rows {
13636            v.scales_into(r, gpr, &mut sc);
13637            let got = q4tp_row_exact(v.nib, r, gpr, &x, &sc);
13638            let want: f32 = (0..cols).map(|c| w[r * cols + c] * x[c]).sum();
13639            // These dot products cancel down to ~1e-3 from terms of ~5e-2, so
13640            // the meaningful yardstick is the summed magnitude, not the result:
13641            // against the result any reordering of a 512-term f32 sum "fails".
13642            let mag: f32 = (0..cols).map(|c| (w[r * cols + c] * x[c]).abs()).sum();
13643            assert!(
13644                (got - want).abs() <= 1e-5 * mag,
13645                "row {r}: kernel {got} vs dequant {want}"
13646            );
13647        }
13648    }
13649
13650    /// The int8 (a8w8) path can't be checked against an f32 reference — the
13651    /// activation quantization dominates. Check it against the q4t kernel it
13652    /// was ported from instead, on payloads holding the same weights: that
13653    /// isolates exactly what the port could break (16 B stride, ladder
13654    /// lookup, nibble unpack) from what it deliberately shares.
13655    #[test]
13656    fn q4tp_matvec_matches_the_q4t_kernel_it_was_ported_from() {
13657        let (rows, cols) = (256usize, 512usize);
13658        let bytes = synth_q4tp(rows, cols);
13659        let twin = q4tp_as_q4t(&bytes, rows, cols);
13660        let x: Vec<f32> = (0..cols)
13661            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
13662            .collect();
13663
13664        let mut got = vec![0f32; rows];
13665        q4tp_matvec(&bytes, &x, rows, cols, &mut got, None);
13666        let mut want = vec![0f32; rows];
13667        q4t_matvec(&twin, &x, rows, cols, &mut want, None);
13668
13669        // Scale is f16 in the twin and f32 here, so allow that rounding on
13670        // top of the summed magnitude (same cancellation argument as above).
13671        let mut w = vec![0f32; rows * cols];
13672        cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
13673        for r in 0..rows {
13674            let mag: f32 = (0..cols).map(|c| (w[r * cols + c] * x[c]).abs()).sum();
13675            assert!(
13676                (got[r] - want[r]).abs() <= 1e-3 * mag,
13677                "row {r}: q4tp {} vs q4t {}",
13678                got[r],
13679                want[r]
13680            );
13681        }
13682    }
13683
13684    /// `matmat` carries three arms (Accelerate, blocked int8 1x4, scalar).
13685    /// Batch 5 crosses the blocked kernel's stride, so this exercises the
13686    /// 1x4 path AND its scalar tail in one run — the blocked kernel is new
13687    /// code and its four accumulators are exactly what tends to go wrong.
13688    #[test]
13689    fn q4tp_matmat_matches_the_q4t_kernel_it_was_ported_from() {
13690        let (rows, cols, b) = (256usize, 512usize, 5usize);
13691        let bytes = synth_q4tp(rows, cols);
13692        let twin = q4tp_as_q4t(&bytes, rows, cols);
13693        let xs: Vec<f32> = (0..b * cols)
13694            .map(|i| ((i * 29 + 11) % 89) as f32 / 89.0 - 0.5)
13695            .collect();
13696
13697        let mut got = vec![0f32; b * rows];
13698        q4tp_matmat(&bytes, &xs, b, rows, cols, &mut got, None);
13699        let mut want = vec![0f32; b * rows];
13700        q4t_matmat(&twin, &xs, b, rows, cols, &mut want, None);
13701
13702        let mut w = vec![0f32; rows * cols];
13703        cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
13704        for t in 0..b {
13705            for r in 0..rows {
13706                let mag: f32 = (0..cols)
13707                    .map(|c| (w[r * cols + c] * xs[t * cols + c]).abs())
13708                    .sum();
13709                let (g, wa) = (got[t * rows + r], want[t * rows + r]);
13710                assert!(
13711                    (g - wa).abs() <= 1e-3 * mag,
13712                    "batch {t} row {r}: q4tp {g} vs q4t {wa}"
13713                );
13714            }
13715        }
13716    }
13717
13718    #[test]
13719    fn q4tp_matvec2_matches_the_single_stream_kernel() {
13720        let (rows, cols) = (128usize, 256usize);
13721        let gpr = cols / GROUP_SIZE;
13722        let bytes = synth_q4tp(rows, cols);
13723        let xs: Vec<f32> = (0..2 * cols)
13724            .map(|i| ((i * 29 + 11) % 89) as f32 / 89.0 - 0.5)
13725            .collect();
13726
13727        let (mut o1, mut o2) = (vec![0f32; rows], vec![0f32; rows]);
13728        q4tp_matvec2(
13729            &bytes,
13730            &xs[..cols],
13731            &xs[cols..],
13732            rows,
13733            cols,
13734            &mut o1,
13735            &mut o2,
13736            None,
13737        );
13738
13739        // matvec2 takes the exact path for both streams, so the single-row
13740        // kernel is an exact reference — no tolerance for path differences.
13741        let v = Q4tpView::new(&bytes, rows, cols);
13742        let mut sc = vec![0f32; gpr];
13743        for r in 0..rows {
13744            v.scales_into(r, gpr, &mut sc);
13745            assert_eq!(o1[r], q4tp_row_exact(v.nib, r, gpr, &xs[..cols], &sc));
13746            assert_eq!(o2[r], q4tp_row_exact(v.nib, r, gpr, &xs[cols..], &sc));
13747        }
13748    }
13749
13750    /// q4tp must not COST speed — it exists to save bytes, and a format that
13751    /// trades 7% of a file for a slower model is a bad trade. This guard is
13752    /// here because correctness tests happily passed while `q4tp_matmat` was
13753    /// missing its int8 and Accelerate arms and the model ran 5x slower.
13754    /// Measured on M-series: 0.97-1.04x, i.e. parity (16 B tiles are better
13755    /// aligned than q4t's 18 B, which pays for the scale indirection).
13756    #[test]
13757    fn q4tp_matvec_keeps_pace_with_q4t() {
13758        let (rows, cols) = (4096usize, 3072usize);
13759        let bytes = synth_q4tp(rows, cols);
13760        let twin = q4tp_as_q4t(&bytes, rows, cols);
13761        let x: Vec<f32> = (0..cols).map(|i| (i % 97) as f32 / 97.0 - 0.5).collect();
13762        let mut o = vec![0f32; rows];
13763        let n = 12;
13764        let mut best = (f64::MAX, f64::MAX);
13765        // Interleaved A/B, minimum statistic: this machine throttles, and a
13766        // mean over a thermal ramp reliably indicts whichever ran second.
13767        for _ in 0..3 {
13768            let t0 = std::time::Instant::now();
13769            for _ in 0..n {
13770                q4t_matvec(&twin, &x, rows, cols, &mut o, None);
13771            }
13772            best.0 = best.0.min(t0.elapsed().as_secs_f64());
13773            let t0 = std::time::Instant::now();
13774            for _ in 0..n {
13775                q4tp_matvec(&bytes, &x, rows, cols, &mut o, None);
13776            }
13777            best.1 = best.1.min(t0.elapsed().as_secs_f64());
13778        }
13779        let ratio = best.1 / best.0;
13780        println!(
13781            "q4t {:.3} ms | q4tp {:.3} ms | {ratio:.2}x",
13782            best.0 * 1e3 / n as f64,
13783            best.1 * 1e3 / n as f64
13784        );
13785        assert!(ratio < 2.0, "q4tp matvec {ratio:.2}x slower than q4t");
13786    }
13787
13788    #[cfg(target_os = "macos")]
13789    #[test]
13790    fn q4t_matmat_accel_matches_dequant_reference() {
13791        if !accel_gemm_enabled() {
13792            return; // CMF_ACCEL=0
13793        }
13794        let (rows, cols, b) = (512usize, 1024usize, 8usize); // ≥500K → accel arm
13795        let gpr = cols / GROUP_SIZE;
13796        let mut bytes = vec![0u8; rows * gpr * Q4_TILE];
13797        for r in 0..rows {
13798            for g in 0..gpr {
13799                let t = (r * gpr + g) * Q4_TILE;
13800                let sc = 0.02 + 0.0005 * ((r * gpr + g) % 64) as f32;
13801                bytes[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
13802                for k in 0..16 {
13803                    bytes[t + 2 + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
13804                }
13805            }
13806        }
13807        let x: Vec<f32> = (0..b * cols)
13808            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
13809            .collect();
13810        let mut got = vec![0f32; b * rows];
13811        q4t_matmat(&bytes, &x, b, rows, cols, &mut got, None);
13812        // Brute-force reference off the same tiles.
13813        let mut w = vec![0f32; rows * cols];
13814        for r in 0..rows {
13815            for g in 0..gpr {
13816                let t = (r * gpr + g) * Q4_TILE;
13817                let s = f16_to_f32(u16::from_le_bytes([bytes[t], bytes[t + 1]]));
13818                for (k, &bb) in bytes[t + 2..t + Q4_TILE].iter().enumerate() {
13819                    w[r * cols + g * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s;
13820                    w[r * cols + g * GROUP_SIZE + k * 2 + 1] =
13821                        (((bb >> 4) & 0x0F) as f32 - 8.0) * s;
13822                }
13823            }
13824        }
13825        for bi in 0..b {
13826            for r in 0..rows {
13827                let want: f32 = (0..cols).map(|j| x[bi * cols + j] * w[r * cols + j]).sum();
13828                let d = (got[bi * rows + r] - want).abs();
13829                assert!(
13830                    d <= want.abs().max(1.0) * 1e-4,
13831                    "accel q4t GEMM diverged at ({bi},{r}): {} vs {want}",
13832                    got[bi * rows + r]
13833                );
13834            }
13835        }
13836    }
13837
13838    #[test]
13839    fn q4matvec_matches_full_dequant() {
13840        let (rows, cols) = (8, 64);
13841        let groups = rows * cols / GROUP_SIZE;
13842        // Hand-craft a q4_block blob: nibbles then f16 scales.
13843        let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
13844        for i in 0..groups * 16 {
13845            bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
13846        }
13847        for g in 0..groups {
13848            let s = 0.01 + 0.003 * g as f32;
13849            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
13850        }
13851        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
13852
13853        let mut reference = vec![0.0f32; rows * cols];
13854        cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
13855        let mut expect = vec![0.0f32; rows];
13856        for r in 0..rows {
13857            expect[r] = reference[r * cols..(r + 1) * cols]
13858                .iter()
13859                .zip(&x)
13860                .map(|(w, xv)| w * xv)
13861                .sum();
13862        }
13863
13864        let mut got = vec![0.0f32; rows];
13865        q4matvec(&bytes, &x, rows, cols, &mut got, None);
13866        // SDOT path quantizes activations to i8 (A8W8): bounded noise,
13867        // same contract as q8/vbit (exact path is pinned by CMF_SDOT=0
13868        // in the golden-parity gate).
13869        let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
13870        let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
13871        for r in 0..rows {
13872            assert!(
13873                (got[r] - expect[r]).abs() < tol * scale,
13874                "row {r}: {} vs {}",
13875                got[r],
13876                expect[r]
13877            );
13878        }
13879    }
13880
13881    /// Fused two-input vbit matvec must equal two single matvecs exactly
13882    /// (same per-lane accumulation order on both scalar and SDOT paths).
13883    #[test]
13884    fn vbitmatvec2_equals_two_singles() {
13885        let (rows, cols) = (6, 64);
13886        let ng = cols / GROUP_SIZE;
13887        let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
13888        let mut bytes = bits.clone();
13889        for g in 0..rows * ng {
13890            let s = 0.02 + 0.001 * g as f32;
13891            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
13892        }
13893        for r in 0..rows {
13894            let b = bits[r] as usize;
13895            let (mut acc, mut nb) = (0u64, 0usize);
13896            let mut rowbytes = Vec::new();
13897            for i in 0..cols {
13898                let v = ((i * 7 + r * 13) % (1 << b)) as u64;
13899                acc = (acc << b) | v;
13900                nb += b;
13901                while nb >= 8 {
13902                    nb -= 8;
13903                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
13904                }
13905            }
13906            if nb > 0 {
13907                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
13908            }
13909            bytes.extend_from_slice(&rowbytes);
13910        }
13911        let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
13912        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).cos()).collect();
13913        let offsets = vbit_row_offsets(&bytes, rows, cols);
13914
13915        let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
13916        vbitmatvec(&bytes, &offsets, &x1, rows, cols, &mut a1, None);
13917        vbitmatvec(&bytes, &offsets, &x2, rows, cols, &mut a2, None);
13918        let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
13919        vbitmatvec2(
13920            &bytes, &offsets, &x1, &x2, rows, cols, &mut b1, &mut b2, None,
13921        );
13922        assert_eq!(a1, b1, "fused vbit lane 1 must be bit-identical");
13923        assert_eq!(a2, b2, "fused vbit lane 2 must be bit-identical");
13924    }
13925
13926    /// Fused two-input q4 matvec must equal two single matvecs exactly.
13927    #[test]
13928    fn q4matvec2_equals_two_singles() {
13929        let (rows, cols) = (8, 128);
13930        let groups = rows * cols / GROUP_SIZE;
13931        let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
13932        for i in 0..groups * 16 {
13933            bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
13934        }
13935        for g in 0..groups {
13936            let s = 0.01 + 0.003 * g as f32;
13937            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
13938        }
13939        // Include an outlier channel so the SDOT correction path is
13940        // exercised in the pair kernel too.
13941        let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
13942        x1[9] = 250.0;
13943        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
13944
13945        let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
13946        q4matvec(&bytes, &x1, rows, cols, &mut a1, None);
13947        q4matvec(&bytes, &x2, rows, cols, &mut a2, None);
13948        let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
13949        q4matvec2(&bytes, &x1, &x2, rows, cols, &mut b1, &mut b2, None);
13950        assert_eq!(a1, b1, "fused q4 lane 1 must be bit-identical");
13951        assert_eq!(a2, b2, "fused q4 lane 2 must be bit-identical");
13952    }
13953
13954    /// Multi-matrix job must equal separate matvecs exactly — same
13955    /// kernels, only the dispatch is fused.
13956    #[test]
13957    fn matvec_many_equals_separate_matvecs() {
13958        use crate::pool::Pool;
13959        let (r1, r2, cols) = (300, 200, 64);
13960        let mk = |salt: usize, rows: usize| {
13961            QTensor::from_f32(
13962                (0..rows * cols)
13963                    .map(|i| ((i * 7 + salt) % 97) as f32 / 97.0 - 0.5)
13964                    .collect(),
13965                rows,
13966                cols,
13967            )
13968        };
13969        let (a, b) = (mk(1, r1), mk(5, r2));
13970        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).sin()).collect();
13971        let pool = Pool::new(3);
13972
13973        let (mut ea, mut eb) = (vec![0f32; r1], vec![0f32; r2]);
13974        a.matvec(&x, &mut ea, Some(&pool));
13975        b.matvec(&x, &mut eb, Some(&pool));
13976        let (mut ga, mut gb) = (vec![0f32; r1], vec![0f32; r2]);
13977        QTensor::matvec_many([&a, &b], &x, [&mut ga, &mut gb], Some(&pool));
13978        assert_eq!(ea, ga, "fused multi-matrix lane 1 must be bit-identical");
13979        assert_eq!(eb, gb, "fused multi-matrix lane 2 must be bit-identical");
13980    }
13981
13982    /// The public Q4TP operator must take the real mapped matvec_many arm,
13983    /// rather than the F32 fallback above.  Build a tiny valid CMF so both
13984    /// handles retain their mmap payloads, then compare the fused dispatch
13985    /// with two ordinary mapped matvec calls bit-for-bit.
13986    #[test]
13987    fn q4tp_matvec_many_equals_separate_matvecs() {
13988        use crate::pool::Pool;
13989        use cortiq_core::{CMF_VERSION, CmfHeader, CmfModel, QuantType, TensorSpec};
13990
13991        let (r1, r2, cols) = (300usize, 200usize, 64usize);
13992        let arch: cortiq_core::ModelArch = serde_json::from_value(serde_json::json!({
13993            "arch_name": "tiny-q4tp",
13994            "hidden_size": cols,
13995            "intermediate_size": cols * 2,
13996            "num_layers": 1,
13997            "num_attention_heads": 2,
13998            "num_kv_heads": 1,
13999            "head_dim": 32,
14000            "vocab_size": r1,
14001            "layer_types": ["FullAttention"],
14002            "rms_norm_eps": 1e-6,
14003            "max_position_embeddings": 8,
14004            "linear_conv_kernel_dim": 0,
14005            "linear_num_key_heads": 0,
14006            "linear_num_value_heads": 0
14007        }))
14008        .unwrap();
14009        let header = CmfHeader {
14010            format: "cmf".into(),
14011            version: CMF_VERSION,
14012            arch,
14013            quant_type: QuantType::Q4Block,
14014            provenance: None,
14015            tokenizer_config: None,
14016            section_hashes: None,
14017            skills: Vec::new(),
14018            shard: None,
14019            calibration: None,
14020            routing: None,
14021            genome: None,
14022            lineage: Vec::new(),
14023            router: None,
14024            segments: Vec::new(),
14025        };
14026        let specs = [
14027            TensorSpec {
14028                name: "q".into(),
14029                dtype: TensorDtype::Q4TiledP,
14030                shape: vec![r1, cols],
14031                data: synth_q4tp(r1, cols),
14032            },
14033            TensorSpec {
14034                name: "kv".into(),
14035                dtype: TensorDtype::Q4TiledP,
14036                shape: vec![r2, cols],
14037                data: synth_q4tp(r2, cols),
14038            },
14039        ];
14040        let dir = std::env::temp_dir().join(format!("cmf-q4tp-many-{}", std::process::id()));
14041        std::fs::create_dir_all(&dir).unwrap();
14042        let path = dir.join("m.cmf");
14043        CmfModel::write(&path, &header, &specs, None, None).unwrap();
14044        let model = Arc::new(CmfModel::open(&path).unwrap());
14045        let (a, b) = (
14046            QTensor::from_model(&model, "q").unwrap(),
14047            QTensor::from_model(&model, "kv").unwrap(),
14048        );
14049        assert_eq!(a.model_dtype(), Some(TensorDtype::Q4TiledP));
14050        assert_eq!(b.model_dtype(), Some(TensorDtype::Q4TiledP));
14051        let x: Vec<f32> = (0..cols)
14052            .map(|i| ((i * 17 + 3) % 97) as f32 / 97.0 - 0.5)
14053            .collect();
14054        let pool = Pool::new(3);
14055        let (mut ea, mut eb) = (vec![0.0f32; r1], vec![0.0f32; r2]);
14056        a.matvec(&x, &mut ea, Some(&pool));
14057        b.matvec(&x, &mut eb, Some(&pool));
14058        let (mut ga, mut gb) = (vec![0.0f32; r1], vec![0.0f32; r2]);
14059        QTensor::matvec_many([&a, &b], &x, [&mut ga, &mut gb], Some(&pool));
14060        assert_eq!(ea, ga, "Q4TP fused lane 1 must be bit-identical");
14061        assert_eq!(eb, gb, "Q4TP fused lane 2 must be bit-identical");
14062        let _ = std::fs::remove_dir_all(&dir);
14063    }
14064
14065    /// The MiMo speculative verify's kernels: several tokens' MoE through
14066    /// `moe_gate_up_rows` / `moe_down_rows` (+ the caller's route-order sum)
14067    /// is bit-identical to each token alone through `moe_gate_up_many` /
14068    /// `moe_down_many` (decode), and a row-exact `q4tp_matmat` of five
14069    /// tokens (wide enough for the blocked tiles) equals five matvecs.
14070    #[test]
14071    fn multi_token_moe_rows_equal_single_token_decode() {
14072        use crate::pool::Pool;
14073        use cortiq_core::{CMF_VERSION, CmfHeader, CmfModel, QuantType, TensorSpec};
14074
14075        let (h, inter, ne) = (64usize, 128usize, 3usize);
14076        let arch: cortiq_core::ModelArch = serde_json::from_value(serde_json::json!({
14077            "arch_name": "tiny-q4tp-moe",
14078            "hidden_size": h,
14079            "intermediate_size": inter,
14080            "num_layers": 1,
14081            "num_attention_heads": 2,
14082            "num_kv_heads": 1,
14083            "head_dim": 32,
14084            "vocab_size": 8,
14085            "layer_types": ["FullAttention"],
14086            "rms_norm_eps": 1e-6,
14087            "max_position_embeddings": 8,
14088            "linear_conv_kernel_dim": 0,
14089            "linear_num_key_heads": 0,
14090            "linear_num_value_heads": 0
14091        }))
14092        .unwrap();
14093        let header = CmfHeader {
14094            format: "cmf".into(),
14095            version: CMF_VERSION,
14096            arch,
14097            quant_type: QuantType::Q4Block,
14098            provenance: None,
14099            tokenizer_config: None,
14100            section_hashes: None,
14101            skills: Vec::new(),
14102            shard: None,
14103            calibration: None,
14104            routing: None,
14105            genome: None,
14106            lineage: Vec::new(),
14107            router: None,
14108            segments: Vec::new(),
14109        };
14110        let mut specs = Vec::new();
14111        for e in 0..ne {
14112            for (k, (n, r, c)) in [("g", inter, h), ("u", inter, h), ("d", h, inter)]
14113                .into_iter()
14114                .enumerate()
14115            {
14116                // Distinct experts: perturb only the nibble plane (any byte
14117                // is a valid pair of codes; the ladder stays intact).
14118                let mut data = synth_q4tp(r, c);
14119                for (i, byte) in data[..r * (c / GROUP_SIZE) * Q4TP_NIB]
14120                    .iter_mut()
14121                    .enumerate()
14122                {
14123                    *byte ^= ((i * (e * 3 + k + 1)) % 251) as u8;
14124                }
14125                specs.push(TensorSpec {
14126                    name: format!("{n}{e}"),
14127                    dtype: TensorDtype::Q4TiledP,
14128                    shape: vec![r, c],
14129                    data,
14130                });
14131            }
14132        }
14133        let dir = std::env::temp_dir().join(format!(
14134            "cmf-moe-rows-{}-{}",
14135            std::process::id(),
14136            FLOAT_ACTIVATIONS.get()
14137        ));
14138        std::fs::create_dir_all(&dir).unwrap();
14139        let path = dir.join("m.cmf");
14140        CmfModel::write(&path, &header, &specs, None, None).unwrap();
14141        let model = Arc::new(CmfModel::open(&path).unwrap());
14142        let t = |n: String| QTensor::from_model(&model, &n).unwrap();
14143        let g: Vec<QTensor> = (0..ne).map(|e| t(format!("g{e}"))).collect();
14144        let u: Vec<QTensor> = (0..ne).map(|e| t(format!("u{e}"))).collect();
14145        let d: Vec<QTensor> = (0..ne).map(|e| t(format!("d{e}"))).collect();
14146        let b = 4usize;
14147        let mut xs: Vec<f32> = (0..b * h)
14148            .map(|i| ((i * 31 + 7) % 89) as f32 / 89.0 - 0.5)
14149            .collect();
14150        xs[5] = 9.0; // an activation outlier on token 0
14151        // Token -> (experts in route order, weights).
14152        let routes: Vec<(Vec<usize>, Vec<f32>)> = vec![
14153            (vec![2, 0], vec![0.6, 0.4]),
14154            (vec![0, 1, 2], vec![0.2, 0.5, 0.3]),
14155            (vec![1], vec![1.0]),
14156            (vec![2, 1, 0], vec![0.25, 0.25, 0.5]),
14157        ];
14158        let pool = Pool::new(3);
14159        // Decode reference, token by token.
14160        let mut want = vec![0f32; b * h];
14161        for (tk, (idx, w)) in routes.iter().enumerate() {
14162            let x = &xs[tk * h..(tk + 1) * h];
14163            let pairs: Vec<(&QTensor, &QTensor)> = idx.iter().map(|&e| (&g[e], &u[e])).collect();
14164            let mut gs: Vec<Vec<f32>> = idx.iter().map(|_| vec![0f32; inter]).collect();
14165            assert!(QTensor::moe_gate_up_many(&pairs, x, &mut gs, Some(&pool)));
14166            if FLOAT_ACTIVATIONS.get() {
14167                for (slot, &e) in idx.iter().enumerate() {
14168                    let (mut gate, mut up) = (vec![0.0; inter], vec![0.0; inter]);
14169                    g[e].matvec(x, &mut gate, Some(&pool));
14170                    u[e].matvec(x, &mut up, Some(&pool));
14171                    for (v, u) in gate.iter_mut().zip(up) {
14172                        *v = (*v / (1.0 + (-*v).exp())) * u;
14173                    }
14174                    assert_eq!(gs[slot], gate, "float gate/up must equal ordinary matvecs");
14175                }
14176            }
14177            let downs: Vec<&QTensor> = idx.iter().map(|&e| &d[e]).collect();
14178            assert!(QTensor::moe_down_many(
14179                &downs,
14180                &gs,
14181                w,
14182                &mut want[tk * h..(tk + 1) * h],
14183                Some(&pool)
14184            ));
14185        }
14186        if FLOAT_ACTIVATIONS.get() {
14187            for (tk, (idx, w)) in routes.iter().enumerate() {
14188                let mut scalar = vec![0.0; h];
14189                for (&e, &weight) in idx.iter().zip(w) {
14190                    let (mut gate, mut up, mut down) =
14191                        (vec![0.0; inter], vec![0.0; inter], vec![0.0; h]);
14192                    g[e].matvec(&xs[tk * h..(tk + 1) * h], &mut gate, Some(&pool));
14193                    u[e].matvec(&xs[tk * h..(tk + 1) * h], &mut up, Some(&pool));
14194                    for (v, u) in gate.iter_mut().zip(up) {
14195                        *v = (*v / (1.0 + (-*v).exp())) * u;
14196                    }
14197                    d[e].matvec(&gate, &mut down, Some(&pool));
14198                    for (v, d) in scalar.iter_mut().zip(down) {
14199                        *v += weight * d;
14200                    }
14201                }
14202                assert_eq!(
14203                    &want[tk * h..(tk + 1) * h],
14204                    scalar,
14205                    "float many equals scalar experts"
14206                );
14207            }
14208        }
14209        // All four tokens at once, grouped by expert.
14210        let mut experts: Vec<usize> = Vec::new();
14211        let mut groups: Vec<Vec<usize>> = Vec::new();
14212        for (tk, (idx, _)) in routes.iter().enumerate() {
14213            for &e in idx {
14214                match experts.iter().position(|&x| x == e) {
14215                    Some(k) => groups[k].push(tk),
14216                    None => {
14217                        experts.push(e);
14218                        groups.push(vec![tk]);
14219                    }
14220                }
14221            }
14222        }
14223        let n_pairs: usize = groups.iter().map(|g| g.len()).sum();
14224        let pairs: Vec<(&QTensor, &QTensor)> = experts.iter().map(|&e| (&g[e], &u[e])).collect();
14225        let mut gs: Vec<Vec<f32>> = (0..n_pairs).map(|_| vec![0f32; inter]).collect();
14226        assert!(QTensor::moe_gate_up_rows(
14227            &pairs,
14228            &groups,
14229            &xs,
14230            &mut gs,
14231            Some(&pool)
14232        ));
14233        let downs: Vec<&QTensor> = experts.iter().map(|&e| &d[e]).collect();
14234        let lens: Vec<usize> = groups.iter().map(|g| g.len()).collect();
14235        let mut ds: Vec<Vec<f32>> = (0..n_pairs).map(|_| vec![0f32; h]).collect();
14236        assert!(QTensor::moe_down_rows(
14237            &downs,
14238            &lens,
14239            &gs,
14240            &mut ds,
14241            Some(&pool)
14242        ));
14243        let slot = |tk: usize, e: usize| {
14244            let k = experts.iter().position(|&x| x == e).unwrap();
14245            groups[..k].iter().map(|g| g.len()).sum::<usize>()
14246                + groups[k].iter().position(|&x| x == tk).unwrap()
14247        };
14248        let mut got = vec![0f32; b * h];
14249        for (tk, (idx, w)) in routes.iter().enumerate() {
14250            for i in 0..h {
14251                let mut acc = 0f32;
14252                for (&e, &we) in idx.iter().zip(w) {
14253                    acc += we * ds[slot(tk, e)][i];
14254                }
14255                got[tk * h + i] = acc;
14256            }
14257        }
14258        assert!(want.iter().any(|v| *v != 0.0));
14259        assert_eq!(
14260            want.iter().map(|v| v.to_bits()).collect::<Vec<_>>(),
14261            got.iter().map(|v| v.to_bits()).collect::<Vec<_>>(),
14262            "multi-token MoE must equal decode bit for bit"
14263        );
14264
14265        // Row-exact q4tp_matmat: five tokens (a blocked 1x4 tile + tail
14266        // otherwise) equal five matvecs.
14267        let b5 = 5usize;
14268        let x5: Vec<f32> = (0..b5 * h)
14269            .map(|i| ((i * 13 + 5) % 71) as f32 / 71.0 - 0.5)
14270            .collect();
14271        let mut mm = vec![0f32; b5 * inter];
14272        row_exact_scope(|| g[1].matmat(&x5, b5, &mut mm, Some(&pool)));
14273        for tk in 0..b5 {
14274            let mut mv = vec![0f32; inter];
14275            g[1].matvec(&x5[tk * h..(tk + 1) * h], &mut mv, Some(&pool));
14276            assert_eq!(
14277                mv.iter().map(|v| v.to_bits()).collect::<Vec<_>>(),
14278                mm[tk * inter..(tk + 1) * inter]
14279                    .iter()
14280                    .map(|v| v.to_bits())
14281                    .collect::<Vec<_>>(),
14282                "row-exact matmat token {tk}"
14283            );
14284        }
14285        // Other concurrent tests/requests may still hold the shared mode.
14286        // Nested, overlapping and unwind restoration is checked separately.
14287        let _ = std::fs::remove_dir_all(&dir);
14288    }
14289
14290    /// The row-exact fix must not touch the fast path. Outside the scope
14291    /// `q4tp_matmat` has to produce exactly what it did before, and on ARM
14292    /// "before" is spelled out below: the tuned 1x4 SDOT tile for every
14293    /// four columns and the single-row kernel for the tail. Inside the
14294    /// scope every column equals its token's matvec. `q4tp_matmat_with`
14295    /// takes the mode as an argument, so a concurrent test holding the
14296    /// shared scope cannot flip it under this one.
14297    #[test]
14298    fn q4tp_matmat_fast_path_unchanged_outside_row_exact() {
14299        use crate::pool::Pool;
14300        use std::sync::atomic::Ordering::Relaxed;
14301        let _alt = Q4TP_ALT_TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner());
14302        // The tuned ARM shape, which is also what an unset switch picks
14303        // unless CMF_Q4TP_V1 is exported.
14304        Q4TP_ALT.store(2, Relaxed);
14305        let pool = Pool::new(3);
14306        // Under 500k cells, so macOS keeps the matmat off the AMX; the
14307        // second shape runs across pool workers (rows >= 256) with 32
14308        // groups of accumulation and a tail after two 1x4 tiles.
14309        for &(rows, cols, b) in &[(64usize, 256usize, 7usize), (320, 1024, 9)] {
14310            let bytes = synth_q4tp(rows, cols);
14311            let mut xs: Vec<f32> = (0..b * cols)
14312                .map(|i| ((i * 29 + 11) % 83) as f32 / 83.0 - 0.5)
14313                .collect();
14314            xs[3] = 7.5; // an activation outlier on token 0
14315            let run = |exact: bool| {
14316                let mut out = vec![0f32; b * rows];
14317                q4tp_matmat_with(&bytes, &xs, b, rows, cols, &mut out, Some(&pool), exact);
14318                out
14319            };
14320            let (fast, exact) = (run(false), run(true));
14321            #[cfg(not(target_arch = "aarch64"))]
14322            let _ = fast;
14323            let bits = |v: &[f32]| v.iter().map(|x| x.to_bits()).collect::<Vec<_>>();
14324            let mut matvecs = vec![0f32; b * rows];
14325            for (bi, o) in matvecs.chunks_mut(rows).enumerate() {
14326                q4tp_matvec(
14327                    &bytes,
14328                    &xs[bi * cols..(bi + 1) * cols],
14329                    rows,
14330                    cols,
14331                    o,
14332                    Some(&pool),
14333                );
14334            }
14335            assert!(matvecs.iter().any(|v| *v != 0.0));
14336            assert_eq!(
14337                bits(&exact),
14338                bits(&matvecs),
14339                "{rows}x{cols} b={b}: row-exact matmat must equal per-token matvecs"
14340            );
14341            #[cfg(target_arch = "aarch64")]
14342            {
14343                // The pre-fix ARM loop, cell for cell.
14344                let gpr = cols / GROUP_SIZE;
14345                let v = Q4tpView::new(&bytes, rows, cols);
14346                let mut old = vec![0f32; b * rows];
14347                let mut sc = vec![0f32; gpr];
14348                let a8w8 = a8w8_enabled();
14349                let blocked = sdot_enabled() && blocked_enabled();
14350                let acts: Vec<SplitAct> = (0..b)
14351                    .map(|bi| split_act(&xs[bi * cols..(bi + 1) * cols]))
14352                    .collect();
14353                for r in 0..rows {
14354                    v.scales_into(r, gpr, &mut sc);
14355                    if !a8w8 {
14356                        for bi in 0..b {
14357                            let x = &xs[bi * cols..(bi + 1) * cols];
14358                            old[bi * rows + r] = q4tp_row_exact(v.nib, r, gpr, x, &sc);
14359                        }
14360                        continue;
14361                    }
14362                    let finish = |d: f32, act: &SplitAct| {
14363                        let mut acc = d * act.sx;
14364                        for &(j, xv) in &act.outliers {
14365                            let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
14366                            acc += w * s * xv;
14367                        }
14368                        acc
14369                    };
14370                    let mut bi = 0usize;
14371                    while blocked && bi + 4 <= b {
14372                        let xs4 = [
14373                            acts[bi].xq.as_slice(),
14374                            acts[bi + 1].xq.as_slice(),
14375                            acts[bi + 2].xq.as_slice(),
14376                            acts[bi + 3].xq.as_slice(),
14377                        ];
14378                        let d = unsafe { dot_q4tp_row_1x4_sdot(v.nib, r, gpr, xs4, &sc) };
14379                        for k in 0..4 {
14380                            old[(bi + k) * rows + r] = finish(d[k], &acts[bi + k]);
14381                        }
14382                        bi += 4;
14383                    }
14384                    for (bi, act) in acts.iter().enumerate().skip(bi) {
14385                        let d = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, &sc);
14386                        old[bi * rows + r] = finish(d, act);
14387                    }
14388                }
14389                assert_eq!(
14390                    bits(&fast),
14391                    bits(&old),
14392                    "{rows}x{cols} b={b}: the fast path changed outside row_exact"
14393                );
14394                // And it is still the fast tile that runs: its lane-parallel
14395                // fma sum rounds differently from the matvec somewhere.
14396                if blocked {
14397                    assert_ne!(
14398                        bits(&fast),
14399                        bits(&matvecs),
14400                        "{rows}x{cols} b={b}: the tuned tile no longer runs outside row_exact"
14401                    );
14402                }
14403            }
14404        }
14405        Q4TP_ALT.store(0, Relaxed);
14406    }
14407
14408    #[test]
14409    fn row_exact_scopes_survive_overlap_nesting_and_unwind() {
14410        use std::sync::{Barrier, atomic::{AtomicUsize, Ordering}};
14411        // A private counter makes this restoration test independent of
14412        // numerical tests concurrently using the production counter.
14413        let active = AtomicUsize::new(0);
14414        counted_row_exact_scope(&active, || {
14415            assert_eq!(active.load(Ordering::Acquire), 1);
14416            counted_row_exact_scope(&active, || {
14417                assert_eq!(active.load(Ordering::Acquire), 2);
14418            });
14419            assert_eq!(active.load(Ordering::Acquire), 1);
14420        });
14421        assert_eq!(active.load(Ordering::Acquire), 0);
14422
14423        let both_entered = Barrier::new(2);
14424        let release_last = Barrier::new(2);
14425        std::thread::scope(|s| {
14426            let first = s.spawn(|| counted_row_exact_scope(&active, || {
14427                both_entered.wait();
14428            }));
14429            let last = s.spawn(|| counted_row_exact_scope(&active, || {
14430                both_entered.wait();
14431                release_last.wait();
14432            }));
14433            first.join().unwrap();
14434            let after_first = active.load(Ordering::Acquire);
14435            release_last.wait();
14436            last.join().unwrap();
14437            assert_eq!(after_first, 1, "second request must remain exact");
14438        });
14439        assert_eq!(active.load(Ordering::Acquire), 0);
14440        let panic = std::panic::catch_unwind(|| {
14441            counted_row_exact_scope(&active, || panic!("scope unwind"));
14442        });
14443        assert!(panic.is_err());
14444        assert_eq!(active.load(Ordering::Acquire), 0);
14445    }
14446
14447    #[test]
14448    #[cfg(target_arch = "x86_64")]
14449    fn q4tp_float_avx2_is_bitwise_scalar() {
14450        if !avx2_enabled() {
14451            return;
14452        }
14453        for cols in [32, 64, 96, 2048, 4096] {
14454            let rows = 9;
14455            let bytes = synth_q4tp(rows, cols);
14456            let v = Q4tpView::new(&bytes, rows, cols);
14457            let gpr = cols / GROUP_SIZE;
14458            let mut sc = vec![0.0; gpr];
14459            for seed in 1..=5 {
14460                let xs: Vec<f32> = (0..cols)
14461                    .map(|i| (((i * 104729 + seed * 8191) % 100003) as f32 - 50001.0) / 7919.0)
14462                    .collect();
14463                for r in 0..rows {
14464                    v.scales_into(r, gpr, &mut sc);
14465                    let scalar = q4tp_row_float_scalar(v.nib, r, gpr, &xs, &sc);
14466                    let vector = unsafe { q4tp_row_float_avx2(v.nib, r, gpr, &xs, &sc) };
14467                    assert_eq!(
14468                        scalar.to_bits(),
14469                        vector.to_bits(),
14470                        "cols={cols} row={r} seed={seed}"
14471                    );
14472                }
14473            }
14474        }
14475    }
14476
14477    #[test]
14478    fn multi_token_moe_rows_float_equal_single_token_decode() {
14479        float_activations_scope(multi_token_moe_rows_equal_single_token_decode);
14480    }
14481
14482    #[test]
14483    fn full_gpu_q8_scope_is_nested_and_thread_local() {
14484        assert!(!FULL_GPU_Q8.get());
14485        let before = gpu_split_frac();
14486        {
14487            let _guard = enter_full_gpu_q8_scope();
14488            assert_eq!(gpu_split_frac(), 1.0);
14489            {
14490                let _nested = enter_full_gpu_q8_scope();
14491            }
14492            assert_eq!(gpu_split_frac(), 1.0);
14493            std::thread::spawn(|| assert!(!FULL_GPU_Q8.get())).join().unwrap();
14494        }
14495        assert!(!FULL_GPU_Q8.get());
14496        assert_eq!(gpu_split_frac(), before);
14497    }
14498
14499    #[test]
14500    fn float_activation_scope_is_nested_thread_local_and_unwind_safe() {
14501        assert!(!FLOAT_ACTIVATIONS.get());
14502        let before = a8w8_enabled();
14503        float_activations_scope(|| {
14504            assert!(!a8w8_enabled());
14505            float_activations_scope(|| assert!(!a8w8_enabled()));
14506            assert!(FLOAT_ACTIVATIONS.get());
14507            std::thread::spawn(|| assert!(!FLOAT_ACTIVATIONS.get()))
14508                .join()
14509                .unwrap();
14510        });
14511        assert!(!FLOAT_ACTIVATIONS.get());
14512        assert_eq!(a8w8_enabled(), before);
14513        let _ = std::panic::catch_unwind(|| float_activations_scope(|| panic!("test unwind")));
14514        assert!(!FLOAT_ACTIVATIONS.get());
14515    }
14516
14517    /// Batched q4/vbit matmat must equal per-position matvec calls
14518    /// exactly (the fallback it replaced) — same kernels, same order.
14519    #[test]
14520    fn batched_matmat_equals_per_position_matvec() {
14521        let (rows, cols, b) = (8, 64, 5);
14522        // q4 blob.
14523        let groups = rows * cols / GROUP_SIZE;
14524        let mut q4 = Vec::new();
14525        for i in 0..groups * 16 {
14526            q4.push((((i * 7 + 3) % 256) & 0xFF) as u8);
14527        }
14528        for g in 0..groups {
14529            q4.extend_from_slice(
14530                &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
14531            );
14532        }
14533        // vbit blob (mixed widths incl. 8).
14534        let ng = cols / GROUP_SIZE;
14535        let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4, 5, 3];
14536        let mut vb = bits.clone();
14537        for g in 0..rows * ng {
14538            vb.extend_from_slice(
14539                &cortiq_core::quant::f32_to_f16(0.02 + 0.001 * g as f32).to_le_bytes(),
14540            );
14541        }
14542        for r in 0..rows {
14543            let bw = bits[r] as usize;
14544            let (mut acc, mut nb) = (0u64, 0usize);
14545            let mut rowbytes = Vec::new();
14546            for i in 0..cols {
14547                let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
14548                acc = (acc << bw) | v;
14549                nb += bw;
14550                while nb >= 8 {
14551                    nb -= 8;
14552                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
14553                }
14554            }
14555            if nb > 0 {
14556                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
14557            }
14558            vb.extend_from_slice(&rowbytes);
14559        }
14560        let offsets = vbit_row_offsets(&vb, rows, cols);
14561
14562        let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
14563
14564        // q4: batch vs singles.
14565        let mut got = vec![0f32; b * rows];
14566        q4matmat(&q4, &xs, b, rows, cols, &mut got, None);
14567        for bi in 0..b {
14568            let mut expect = vec![0f32; rows];
14569            q4matvec(
14570                &q4,
14571                &xs[bi * cols..(bi + 1) * cols],
14572                rows,
14573                cols,
14574                &mut expect,
14575                None,
14576            );
14577            assert_eq!(
14578                &got[bi * rows..(bi + 1) * rows],
14579                &expect[..],
14580                "q4 batch pos {bi}"
14581            );
14582        }
14583
14584        // vbit: batch vs singles.
14585        let mut got = vec![0f32; b * rows];
14586        vbitmatmat(&vb, &offsets, &xs, b, rows, cols, &mut got, None);
14587        for bi in 0..b {
14588            let mut expect = vec![0f32; rows];
14589            vbitmatvec(
14590                &vb,
14591                &offsets,
14592                &xs[bi * cols..(bi + 1) * cols],
14593                rows,
14594                cols,
14595                &mut expect,
14596                None,
14597            );
14598            assert_eq!(
14599                &got[bi * rows..(bi + 1) * rows],
14600                &expect[..],
14601                "vbit batch pos {bi}"
14602            );
14603        }
14604    }
14605
14606    /// q4_tiled kernels must produce BIT-identical outputs to the q4
14607    /// split kernels on the same values (same ints, same order — only
14608    /// the byte placement differs).
14609    #[test]
14610    fn q4_tiled_matches_q4_block_bitexact() {
14611        let (rows, cols, b) = (8usize, 128usize, 3usize);
14612        let groups = rows * cols / GROUP_SIZE;
14613        let mut split = Vec::with_capacity(groups * 18);
14614        for i in 0..groups * 16 {
14615            split.push((((i * 7 + 3) % 256) & 0xFF) as u8);
14616        }
14617        for g in 0..groups {
14618            split.extend_from_slice(
14619                &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
14620            );
14621        }
14622        // Re-tile: [scale][nibbles] per group.
14623        let (packed, scales) = split.split_at(groups * 16);
14624        let mut tiled = Vec::with_capacity(groups * Q4_TILE);
14625        for g in 0..groups {
14626            tiled.extend_from_slice(&scales[g * 2..g * 2 + 2]);
14627            tiled.extend_from_slice(&packed[g * 16..(g + 1) * 16]);
14628        }
14629
14630        let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
14631        x1[9] = 250.0; // exercise the outlier path
14632        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
14633
14634        let (mut a, mut t) = (vec![0f32; rows], vec![0f32; rows]);
14635        q4matvec(&split, &x1, rows, cols, &mut a, None);
14636        q4t_matvec(&tiled, &x1, rows, cols, &mut t, None);
14637        assert_eq!(a, t, "q4t matvec must match q4 bit-for-bit");
14638
14639        let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
14640        let (mut t1, mut t2) = (vec![0f32; rows], vec![0f32; rows]);
14641        q4matvec2(&split, &x1, &x2, rows, cols, &mut a1, &mut a2, None);
14642        q4t_matvec2(&tiled, &x1, &x2, rows, cols, &mut t1, &mut t2, None);
14643        assert_eq!(a1, t1);
14644        assert_eq!(a2, t2);
14645
14646        let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
14647        let (mut am, mut tm) = (vec![0f32; b * rows], vec![0f32; b * rows]);
14648        q4matmat(&split, &xs, b, rows, cols, &mut am, None);
14649        q4t_matmat(&tiled, &xs, b, rows, cols, &mut tm, None);
14650        assert_eq!(am, tm, "q4t matmat must match q4 bit-for-bit");
14651    }
14652
14653    /// q4 SDOT outlier correction: a single huge activation channel
14654    /// (>8·rms → outlier, zeroed in xq) must still contribute its EXACT
14655    /// term. On-grid bulk (±1/0 → xq dequantizes exactly) isolates the
14656    /// correction from A8W8 noise. cols must exceed 64: at n=64 the
14657    /// 8·rms threshold equals sqrt(v²+rest) ≥ v, so a single outlier
14658    /// can never qualify (8² = n).
14659    #[test]
14660    fn q4matvec_sdot_outlier_exact() {
14661        let (rows, cols) = (4, 128);
14662        let groups = rows * cols / GROUP_SIZE;
14663        let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
14664        for i in 0..groups * 16 {
14665            bytes.push(((i * 11 + 5) % 256) as u8);
14666        }
14667        for g in 0..groups {
14668            let s = 0.02 + 0.002 * g as f32;
14669            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
14670        }
14671        let mut x: Vec<f32> = (0..cols)
14672            .map(|i| match i % 3 {
14673                0 => 1.0,
14674                1 => -1.0,
14675                _ => 0.0,
14676            })
14677            .collect();
14678        x[17] = 300.0; // ≫ 8·rms → outlier channel
14679
14680        let mut reference = vec![0.0f32; rows * cols];
14681        cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
14682        let mut expect = vec![0.0f32; rows];
14683        for r in 0..rows {
14684            expect[r] = reference[r * cols..(r + 1) * cols]
14685                .iter()
14686                .zip(&x)
14687                .map(|(w, xv)| w * xv)
14688                .sum();
14689        }
14690        let mut got = vec![0.0f32; rows];
14691        q4matvec(&bytes, &x, rows, cols, &mut got, None);
14692        let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
14693        for r in 0..rows {
14694            assert!(
14695                (got[r] - expect[r]).abs() < 2e-3 * scale,
14696                "row {r}: {} vs {} (outlier term must be exact)",
14697                got[r],
14698                expect[r]
14699            );
14700        }
14701    }
14702
14703    /// The fused q1t matvec must equal the reference (dequant_q1t → dot),
14704    /// including the ternary zero level and the binary-searched outlier
14705    /// overlay. Guards the mmap kernel that makes a 12B q1t runnable.
14706    #[test]
14707    fn q1t_matvec_matches_reference() {
14708        use cortiq_core::quant::{dequant_q1t, f32_to_f16};
14709        let (rows, cols) = (3usize, 64usize); // gpr = 2
14710        let gpr = cols / GROUP_SIZE;
14711        let scales = [0.5f32, 0.3, 0.7, 0.2, 0.6, 0.15];
14712        // Overlay (must be sorted by flat index): a few spikes across rows.
14713        let outliers: [(u32, f32); 3] = [(5, 9.0), (70, -4.5), (150, 3.25)];
14714        let is_out = |flat: usize| outliers.iter().any(|&(i, _)| i as usize == flat);
14715        let mut bytes = Vec::new();
14716        for r in 0..rows {
14717            for g in 0..gpr {
14718                bytes.extend_from_slice(&f32_to_f16(scales[r * gpr + g]).to_le_bytes());
14719                let mut c = [0u8; 7];
14720                for k in 0..GROUP_SIZE {
14721                    // Encoder invariant: code 0 at outlier positions.
14722                    let code = if is_out(r * cols + g * GROUP_SIZE + k) {
14723                        0
14724                    } else {
14725                        ((k + r * 3 + g) % 3) as u8 // 0,1,2
14726                    };
14727                    cortiq_core::quant::q1t_pack(&mut c, k, code);
14728                }
14729                bytes.extend_from_slice(&c);
14730            }
14731        }
14732        // Per-row overlay: [u32 row_ptr[rows+1]] then [(u16 col, f16 val)] by
14733        // row (outliers are sorted by flat index → already grouped by row).
14734        let mut row_ptr = vec![0u32; rows + 1];
14735        for &(idx, _) in &outliers {
14736            row_ptr[idx as usize / cols + 1] += 1;
14737        }
14738        for r in 0..rows {
14739            row_ptr[r + 1] += row_ptr[r];
14740        }
14741        for &p in &row_ptr {
14742            bytes.extend_from_slice(&p.to_le_bytes());
14743        }
14744        for &(idx, v) in &outliers {
14745            bytes.extend_from_slice(&((idx as usize % cols) as u16).to_le_bytes());
14746            bytes.extend_from_slice(&f32_to_f16(v).to_le_bytes());
14747        }
14748
14749        let mut refw = vec![0f32; rows * cols];
14750        dequant_q1t(&bytes, rows, cols, &mut refw);
14751        // On-grid activations (±1, amax 1) so the int8 SDOT path reconstructs
14752        // x exactly and matches the f32 reference (same trick as the q1 test).
14753        let x: Vec<f32> = (0..cols)
14754            .map(|j| if j % 3 == 0 { 1.0 } else { -1.0 })
14755            .collect();
14756        let mut expect = vec![0f32; rows];
14757        for r in 0..rows {
14758            let mut a = 0.0f32;
14759            for j in 0..cols {
14760                a += refw[r * cols + j] * x[j];
14761            }
14762            expect[r] = a;
14763        }
14764        let tol = |e: f32| 1e-3 * e.abs().max(1e-3);
14765        let mut got = vec![0f32; rows];
14766        q1t_matvec(&bytes, &x, rows, cols, &mut got, None);
14767        for r in 0..rows {
14768            assert!(
14769                (got[r] - expect[r]).abs() < tol(expect[r]),
14770                "row {r}: {} vs {}",
14771                got[r],
14772                expect[r]
14773            );
14774        }
14775        // matmat (b=2, f32 decode path) must agree too.
14776        let x2: Vec<f32> = x.iter().chain(x.iter()).copied().collect();
14777        let mut gm = vec![0f32; 2 * rows];
14778        q1t_matmat(&bytes, &x2, 2, rows, cols, &mut gm, None);
14779        for r in 0..rows {
14780            assert!((gm[r] - expect[r]).abs() < tol(expect[r]));
14781            assert!((gm[rows + r] - expect[r]).abs() < tol(expect[r]));
14782        }
14783        // Fused pair (q1t_matvec2) must equal two single matvecs
14784        // bit-for-bit: same unpack, same group order, same f32
14785        // accumulation per stream. Distinct x2 exercises both lanes.
14786        let xb: Vec<f32> = (0..cols)
14787            .map(|j| if j % 5 == 0 { -1.0 } else { 1.0 })
14788            .collect();
14789        let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
14790        q1t_matvec(&bytes, &x, rows, cols, &mut s1, None);
14791        q1t_matvec(&bytes, &xb, rows, cols, &mut s2, None);
14792        let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
14793        q1t_matvec2(&bytes, &x, &xb, rows, cols, &mut p1, &mut p2, None);
14794        assert_eq!(p1, s1, "q1t pair lane 1 ≠ single matvec");
14795        assert_eq!(p2, s2, "q1t pair lane 2 ≠ single matvec");
14796    }
14797
14798    /// Pair == 2×matvec with an ODD group count (the kernel's tail
14799    /// group) and no overlay section.
14800    #[test]
14801    fn q1t_matvec2_odd_gpr_matches_singles() {
14802        use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_pack};
14803        let (rows, cols) = (5usize, 96usize); // gpr = 3 → paired + tail
14804        let gpr = cols / GROUP_SIZE;
14805        let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE);
14806        for r in 0..rows {
14807            for g in 0..gpr {
14808                bytes.extend_from_slice(&f32_to_f16(0.1 + 0.05 * (r + g) as f32).to_le_bytes());
14809                let mut c = [0u8; 7];
14810                for k in 0..GROUP_SIZE {
14811                    q1t_pack(&mut c, k, ((k * 7 + r * 5 + g * 3) % 3) as u8);
14812                }
14813                bytes.extend_from_slice(&c);
14814            }
14815        }
14816        let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.31).sin()).collect();
14817        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).cos()).collect();
14818        let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
14819        q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
14820        q1t_matvec(&bytes, &x2, rows, cols, &mut s2, None);
14821        let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
14822        q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
14823        assert_eq!(p1, s1, "odd-gpr pair lane 1 ≠ single");
14824        assert_eq!(p2, s2, "odd-gpr pair lane 2 ≠ single");
14825    }
14826
14827    // Speed A/B: fused pair (one unpack, two streams) vs two single
14828    // matvecs. Single-threaded, FFN-sized, min-of paired in-process.
14829    //   cargo test -p cortiq-engine --release q1t_matvec2_speed -- --ignored --nocapture
14830    #[test]
14831    #[ignore]
14832    fn q1t_matvec2_speed() {
14833        use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_pack};
14834        use std::time::Instant;
14835        let (rows, cols) = (8192usize, 4096usize);
14836        let gpr = cols / GROUP_SIZE;
14837        let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE);
14838        for r in 0..rows {
14839            for g in 0..gpr {
14840                let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
14841                bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
14842                let mut c = [0u8; 7];
14843                for k in 0..GROUP_SIZE {
14844                    q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
14845                }
14846                bytes.extend_from_slice(&c);
14847            }
14848        }
14849        let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.31).sin()).collect();
14850        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).cos()).collect();
14851        let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
14852        let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
14853        // Warm both paths once.
14854        q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
14855        q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
14856        let (mut t_pair, mut t_two) = (f64::MAX, f64::MAX);
14857        for _ in 0..8 {
14858            let t0 = Instant::now();
14859            q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
14860            t_pair = t_pair.min(t0.elapsed().as_secs_f64() * 1000.0);
14861            let t1 = Instant::now();
14862            q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
14863            q1t_matvec(&bytes, &x2, rows, cols, &mut s2, None);
14864            t_two = t_two.min(t1.elapsed().as_secs_f64() * 1000.0);
14865        }
14866        assert_eq!(p1, s1);
14867        assert_eq!(p2, s2);
14868        println!("q1t pair {rows}x{cols}: fused {t_pair:.2} ms | two singles {t_two:.2} ms");
14869    }
14870
14871    // Speed A/B: the base-3-division decode (what the packing commit left in
14872    // place) vs the fused sign-LUT matvec. Both single-threaded, same bytes.
14873    //   cargo test -p cortiq-engine q1t_matvec_speed -- --ignored --nocapture
14874    #[test]
14875    #[ignore]
14876    fn q1t_matvec_speed() {
14877        use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_code, q1t_pack};
14878        use std::time::Instant;
14879        let (rows, cols) = (8192usize, 4096usize); // FFN-sized
14880        let gpr = cols / GROUP_SIZE;
14881        let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE + 16);
14882        for r in 0..rows {
14883            for g in 0..gpr {
14884                let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
14885                bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
14886                let mut c = [0u8; 7];
14887                for k in 0..GROUP_SIZE {
14888                    q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
14889                }
14890                bytes.extend_from_slice(&c);
14891            }
14892        }
14893        let (n, stride) = (rows * cols, 40usize); // ~2.5% outliers, per-row overlay
14894        let mut row_ptr = vec![0u32; rows + 1];
14895        let mut idx = 0usize;
14896        while idx < n {
14897            row_ptr[idx / cols + 1] += 1;
14898            idx += stride;
14899        }
14900        for r in 0..rows {
14901            row_ptr[r + 1] += row_ptr[r];
14902        }
14903        for &p in &row_ptr {
14904            bytes.extend_from_slice(&p.to_le_bytes());
14905        }
14906        let mut idx = 0usize;
14907        while idx < n {
14908            bytes.extend_from_slice(&((idx % cols) as u16).to_le_bytes());
14909            bytes.extend_from_slice(&f32_to_f16((idx % 13) as f32 * 0.1 - 0.6).to_le_bytes());
14910            idx += stride;
14911        }
14912        // On-grid ±1 so the fast path's int8 SDOT is exact vs the f32 "slow"
14913        // reference (the A/B is a timing check; values must still agree).
14914        let x: Vec<f32> = (0..cols)
14915            .map(|j| if j % 3 == 0 { 1.0 } else { -1.0 })
14916            .collect();
14917        let (rp_off, ent_off, has_ov) = q1t_overlay(&bytes, rows * gpr * Q1T_TILE, rows);
14918
14919        // "before": base-3 division decode into a buffer, then dot.
14920        let slow = |out: &mut [f32]| {
14921            let mut buf = vec![0f32; cols];
14922            for r in 0..rows {
14923                for g in 0..gpr {
14924                    let off = (r * gpr + g) * Q1T_TILE;
14925                    let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
14926                    let codes = &bytes[off + 2..off + Q1T_TILE];
14927                    for k in 0..GROUP_SIZE {
14928                        buf[g * GROUP_SIZE + k] = match q1t_code(codes, k) {
14929                            1 => s,
14930                            2 => -s,
14931                            _ => 0.0,
14932                        };
14933                    }
14934                }
14935                out[r] = q1t_row_outlier_correction(&bytes, r, rp_off, ent_off, has_ov, &x)
14936                    + (0..cols).map(|j| buf[j] * x[j]).sum::<f32>();
14937            }
14938        };
14939        let iters = 5;
14940        let mut a = vec![0f32; rows];
14941        slow(&mut a); // warm
14942        let t = Instant::now();
14943        for _ in 0..iters {
14944            slow(&mut a);
14945        }
14946        let slow_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
14947
14948        let mut b = vec![0f32; rows];
14949        q1t_matvec(&bytes, &x, rows, cols, &mut b, None); // warm
14950        let t = Instant::now();
14951        for _ in 0..iters {
14952            q1t_matvec(&bytes, &x, rows, cols, &mut b, None);
14953        }
14954        let fast_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
14955
14956        for r in 0..rows {
14957            assert!((a[r] - b[r]).abs() < 1e-2, "mismatch row {r}");
14958        }
14959        println!(
14960            "q1t matvec {rows}x{cols} (1 thread): div-decode {slow_ms:.2} ms  fused-LUT {fast_ms:.2} ms  => {:.2}x",
14961            slow_ms / fast_ms
14962        );
14963    }
14964}
14965
14966#[cfg(test)]
14967mod gemm_bench {
14968    /// `cargo test -p cortiq-engine --release q4tp_matmat_throughput -- --ignored --nocapture`
14969    /// Times the batched q4tp GEMM at the shapes the image DiT runs
14970    /// (b=296 tokens, 2304 -> 9216), on synthetic bytes: no model, no
14971    /// mmap, no thermal drift over minutes — a kernel change shows up
14972    /// here in seconds where a full render hides it in noise.
14973    ///
14974    /// On macOS add `CMF_ACCEL=0`: this shape is over the 500k-cell mark
14975    /// where the matmat hands off to Accelerate's dequant sgemm, and
14976    /// without the opt-out both rows below measure the AMX, not the
14977    /// kernel under test.
14978    #[test]
14979    #[ignore]
14980    fn q4tp_matmat_throughput() {
14981        let _alt = super::Q4TP_ALT_TEST_LOCK
14982            .lock()
14983            .unwrap_or_else(|e| e.into_inner());
14984        // 296 is a prompt-encode batch; the image DiT runs 2085 at
14985        // 512x512, where the activation panel stops fitting L2 and the
14986        // loop's shape starts to matter more than its instructions.
14987        let b: usize = std::env::var("CMF_BENCH_B")
14988            .ok()
14989            .and_then(|v| v.parse().ok())
14990            .unwrap_or(296);
14991        let (rows, cols) = (9216usize, 2304usize);
14992        let (_, _, _) = (rows, cols, b);
14993        let total =
14994            cortiq_core::quant::expected_nbytes(cortiq_core::TensorDtype::Q4TiledP, &[rows, cols])
14995                .unwrap();
14996        // Random nibbles are fine, but the row params are f16 (lo, step)
14997        // of a geometric ladder: garbage there gives exp2 of a huge
14998        // exponent, the scales come back inf, and the whole bench times
14999        // NaN arithmetic instead of the kernel.
15000        let (params_off, codes_off, _) = cortiq_core::quant::q4tp_sections(rows, cols);
15001        let mut bytes: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
15002        let lo = cortiq_core::quant::f32_to_f16(-4.0);
15003        let step = cortiq_core::quant::f32_to_f16(0.1);
15004        for r in 0..rows {
15005            let o = params_off + r * 4;
15006            bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
15007            bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
15008        }
15009        let _ = codes_off;
15010        let xs: Vec<f32> = (0..b * cols)
15011            .map(|i| ((i % 97) as f32 - 48.0) / 48.0)
15012            .collect();
15013        let mut out = vec![0f32; b * rows];
15014        let pool = crate::pool::Pool::from_env();
15015        // A shared 48-core stand drifts ±25% run to run, which is wider
15016        // than any kernel change worth making. So: alternate the two
15017        // kernels inside one process and keep the BEST time for
15018        // each. Interleaving makes both see the same interference, and a
15019        // minimum is the one statistic another tenant cannot inflate.
15020        super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
15021        let reps: usize = std::env::var("CMF_BENCH_REPS")
15022            .ok()
15023            .and_then(|v| v.parse().ok())
15024            .unwrap_or(10);
15025        let mut best = [f64::MAX; 2];
15026        let mut sums = [0f32; 2];
15027        for _ in 0..reps {
15028            for (k, w) in [(0usize, 1u8), (1usize, 2u8)] {
15029                super::Q4TP_ALT.store(w, std::sync::atomic::Ordering::Relaxed);
15030                let t = std::time::Instant::now();
15031                super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
15032                best[k] = best[k].min(t.elapsed().as_secs_f64());
15033                sums[k] = out.iter().take(64).sum::<f32>();
15034            }
15035        }
15036        let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
15037        for (k, name) in ["previous", "tuned   "].iter().enumerate() {
15038            println!(
15039                "q4tp matmat {rows}x{cols} b={b} {name}: {:.1} ms  {:.1} GFLOP/s  (checksum {:.3})",
15040                best[k] * 1e3,
15041                flops / best[k] / 1e9,
15042                sums[k]
15043            );
15044        }
15045        assert!(
15046            (sums[0] - sums[1]).abs() < 1e-2,
15047            "the tuned kernel changed the result: {} vs {}",
15048            sums[0],
15049            sums[1]
15050        );
15051    }
15052
15053    /// The blocked kernel must agree with the per-column path exactly —
15054    /// same weights, same activation split, only a different instruction
15055    /// mix. Shapes are chosen to hit the awkward cases: a column count
15056    /// that leaves an odd group (the 512-bit kernel does two at a time),
15057    /// and a batch that does not divide by four.
15058    #[test]
15059    fn q4tp_matmat_blocked_matches_scalar() {
15060        use std::sync::atomic::Ordering::Relaxed;
15061        let _alt = super::Q4TP_ALT_TEST_LOCK
15062            .lock()
15063            .unwrap_or_else(|e| e.into_inner());
15064        // The last shape carries the image DiT's column count — 2304, so
15065        // 72 groups of accumulation, which is where a reordered sum can
15066        // actually drift — and runs through the thread pool, since the
15067        // blocked path splits rows across workers. Its row count stays
15068        // under 500k cells on purpose: above that, macOS diverts the whole
15069        // matmat to the Accelerate/AMX dequant sgemm and neither kernel
15070        // here would run.
15071        for &(rows, cols, b) in &[
15072            (64usize, 128usize, 7usize),
15073            (33, 96, 4),
15074            (16, 256, 9),
15075            (192, 2304, 37),
15076        ] {
15077            let total = cortiq_core::quant::expected_nbytes(
15078                cortiq_core::TensorDtype::Q4TiledP,
15079                &[rows, cols],
15080            )
15081            .unwrap();
15082            let (params_off, _, _) = cortiq_core::quant::q4tp_sections(rows, cols);
15083            let mut bytes: Vec<u8> = (0..total).map(|i| (i * 61 % 251) as u8).collect();
15084            let lo = cortiq_core::quant::f32_to_f16(-4.0);
15085            let step = cortiq_core::quant::f32_to_f16(0.1);
15086            for r in 0..rows {
15087                let o = params_off + r * 4;
15088                bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
15089                bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
15090            }
15091            let xs: Vec<f32> = (0..b * cols)
15092                .map(|i| ((i % 89) as f32 - 44.0) / 44.0)
15093                .collect();
15094            let mut got = vec![0f32; b * rows];
15095            let mut want = vec![0f32; b * rows];
15096            let gpr = cols / 32;
15097            let view = super::Q4tpView::new(&bytes, rows, cols);
15098            let pool = crate::pool::Pool::from_env();
15099            super::Q4TP_ALT.store(2, Relaxed);
15100            super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut got, pool.as_deref());
15101            super::Q4TP_ALT.store(1, Relaxed);
15102            super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut want, pool.as_deref());
15103            super::Q4TP_ALT.store(0, Relaxed);
15104            // Measured against the output's scale, not cell by cell: a
15105            // dot product of 2304 terms lands near zero wherever the row
15106            // and the activation nearly cancel, and there a per-cell
15107            // ratio reports 1e-3 for an absolute error of 5e-6 — f32's
15108            // own rounding, reordered. What must stay small is the error
15109            // relative to what the layer actually outputs.
15110            let scale = want.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
15111            let (mut worst, mut at) = (0f32, 0usize);
15112            for (i, (g, w)) in got.iter().zip(&want).enumerate() {
15113                if (g - w).abs() > worst {
15114                    worst = (g - w).abs();
15115                    at = i;
15116                }
15117            }
15118            assert!(
15119                worst <= 1e-4 * scale,
15120                "{rows}x{cols} b={b}: blocked and scalar disagree by {worst:.3e} \
15121                 (scale {scale:.3e}) at cell {at}: {} vs {}",
15122                got[at],
15123                want[at]
15124            );
15125
15126            // "Same speed, no quality loss" is a claim about which answer
15127            // is RIGHT, not about which two agree. Both paths sum the same
15128            // 2304 products in different orders, so f64 decides: the
15129            // blocked kernel keeps sixteen partial sums and folds them at
15130            // the end, which is a shallower addition tree than the
15131            // per-column path's running scalar, and it must not be worse.
15132            let (mut e_blocked, mut e_scalar) = (0f64, 0f64);
15133            for bi in 0..b {
15134                let act = super::split_act(&xs[bi * cols..(bi + 1) * cols]);
15135                for r in 0..rows {
15136                    let mut sc = vec![0f32; gpr];
15137                    view.scales_into(r, gpr, &mut sc);
15138                    let mut exact = 0f64;
15139                    for j in 0..cols {
15140                        let (w, sq) = super::q4tp_outlier(view.nib, r, gpr, j, &sc);
15141                        exact += w as f64 * sq as f64 * act.xq[j] as f64;
15142                    }
15143                    exact *= act.sx as f64;
15144                    for &(j, xv) in &act.outliers {
15145                        let (w, sq) = super::q4tp_outlier(view.nib, r, gpr, j, &sc);
15146                        exact += w as f64 * sq as f64 * xv as f64;
15147                    }
15148                    let i = bi * rows + r;
15149                    e_blocked = e_blocked.max((got[i] as f64 - exact).abs());
15150                    e_scalar = e_scalar.max((want[i] as f64 - exact).abs());
15151                }
15152            }
15153            println!(
15154                "{rows}x{cols} b={b}: worst error vs f64 — blocked {e_blocked:.3e}, \
15155                 per-column {e_scalar:.3e}"
15156            );
15157            // An absolute bar, not a race between the two: at these
15158            // magnitudes both sit in f32's last bits, and on a small shape
15159            // whichever one happens to round the unluckiest cell "wins" by
15160            // a factor the next seed reverses.
15161            assert!(
15162                e_blocked <= 1e-5 * scale as f64 && e_scalar <= 1e-5 * scale as f64,
15163                "{rows}x{cols} b={b}: error against f64 too large — blocked \
15164                 {e_blocked:.3e}, per-column {e_scalar:.3e}, scale {scale:.3e}"
15165            );
15166        }
15167    }
15168
15169    /// What the row-exact contract costs a speculative-verify panel on the
15170    /// host: the same batch through the fast arms, through the row-exact
15171    /// arms, and as one matvec per token (the other way to be exact).
15172    /// Arms alternate inside one process and keep their best time.
15173    /// `CMF_GPU=0 cargo test -p cortiq-engine --release q4tp_matmat_row_exact_cost -- --ignored --nocapture`
15174    #[test]
15175    #[ignore]
15176    fn q4tp_matmat_row_exact_cost() {
15177        let _alt = super::Q4TP_ALT_TEST_LOCK
15178            .lock()
15179            .unwrap_or_else(|e| e.into_inner());
15180        let pool = crate::pool::Pool::from_env();
15181        let reps: usize = std::env::var("CMF_BENCH_REPS")
15182            .ok()
15183            .and_then(|v| v.parse().ok())
15184            .unwrap_or(30);
15185        for &(rows, cols) in &[(2048usize, 4096usize), (4096, 2048), (4096, 4096)] {
15186            let total = cortiq_core::quant::expected_nbytes(
15187                cortiq_core::TensorDtype::Q4TiledP,
15188                &[rows, cols],
15189            )
15190            .unwrap();
15191            let (params_off, _, _) = cortiq_core::quant::q4tp_sections(rows, cols);
15192            let mut bytes: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
15193            let lo = cortiq_core::quant::f32_to_f16(-4.0);
15194            let step = cortiq_core::quant::f32_to_f16(0.1);
15195            for r in 0..rows {
15196                let o = params_off + r * 4;
15197                bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
15198                bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
15199            }
15200            for &b in &[2usize, 4, 5, 8] {
15201                let xs: Vec<f32> = (0..b * cols)
15202                    .map(|i| ((i % 97) as f32 - 48.0) / 48.0)
15203                    .collect();
15204                let mut out = vec![0f32; b * rows];
15205                let mut best = [f64::MAX; 3];
15206                for _ in 0..reps {
15207                    for (k, best_k) in best.iter_mut().enumerate() {
15208                        let t = std::time::Instant::now();
15209                        match k {
15210                            0 | 1 => super::q4tp_matmat_with(
15211                                &bytes,
15212                                &xs,
15213                                b,
15214                                rows,
15215                                cols,
15216                                &mut out,
15217                                pool.as_deref(),
15218                                k == 1,
15219                            ),
15220                            _ => {
15221                                for (bi, o) in out.chunks_mut(rows).enumerate() {
15222                                    super::q4tp_matvec(
15223                                        &bytes,
15224                                        &xs[bi * cols..(bi + 1) * cols],
15225                                        rows,
15226                                        cols,
15227                                        o,
15228                                        pool.as_deref(),
15229                                    );
15230                                }
15231                            }
15232                        }
15233                        *best_k = best_k.min(t.elapsed().as_secs_f64());
15234                    }
15235                }
15236                println!(
15237                    "q4tp {rows}x{cols} b={b}: fast {:.3} ms, row-exact {:.3} ms, \
15238                     {b} matvecs {:.3} ms",
15239                    best[0] * 1e3,
15240                    best[1] * 1e3,
15241                    best[2] * 1e3
15242                );
15243            }
15244        }
15245    }
15246
15247    /// The q4t twin of the throughput bench, same shape and rules, so the
15248    /// two quantisations' batch kernels can be read against each other.
15249    /// `cargo test -p cortiq-engine --release q4t_matmat_throughput -- --ignored --nocapture`
15250    #[test]
15251    #[ignore]
15252    fn q4t_matmat_throughput() {
15253        let (rows, cols, b) = (9216usize, 2304usize, 296usize);
15254        let total =
15255            cortiq_core::quant::expected_nbytes(cortiq_core::TensorDtype::Q4Tiled, &[rows, cols])
15256                .unwrap();
15257        // q4t carries a per-group f16 scale in the tile's first two bytes;
15258        // random bytes there decode to inf and the bench would time NaNs.
15259        let mut bytes: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
15260        let sc = cortiq_core::quant::f32_to_f16(0.02);
15261        for t in bytes.chunks_mut(super::Q4_TILE) {
15262            t[..2].copy_from_slice(&sc.to_le_bytes());
15263        }
15264        let xs: Vec<f32> = (0..b * cols)
15265            .map(|i| ((i % 97) as f32 - 48.0) / 48.0)
15266            .collect();
15267        let mut out = vec![0f32; b * rows];
15268        let pool = crate::pool::Pool::from_env();
15269        super::q4t_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
15270        let reps: usize = std::env::var("CMF_BENCH_REPS")
15271            .ok()
15272            .and_then(|v| v.parse().ok())
15273            .unwrap_or(10);
15274        let mut best = f64::MAX;
15275        for _ in 0..reps {
15276            let t = std::time::Instant::now();
15277            super::q4t_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
15278            best = best.min(t.elapsed().as_secs_f64());
15279        }
15280        let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
15281        println!(
15282            "q4t matmat {rows}x{cols} b={b}: {:.1} ms  {:.1} GFLOP/s  (checksum {:.3})",
15283            best * 1e3,
15284            flops / best / 1e9,
15285            out.iter().take(64).sum::<f32>()
15286        );
15287    }
15288}