Skip to main content

cortiq_engine/
qtensor.rs

1//! QTensor — weight tensor with pluggable storage.
2//!
3//! Two backings, one interface:
4//! - `F32`   — owned dense floats (small models, tests). Every operation
5//!   is bit-identical to the historical `&[f32]` code paths.
6//! - `Mapped` — quantized bytes zero-copy from the CMF mmap (`q8_row` /
7//!   `q8_2f`). The matvec is fused: int8 rows × f32 activations, the
8//!   q8_2f column field folds into a pre-scale of the input
9//!   (`x'[i] = col[i]·x[i]`), so the inner loop is the same i8 dot as
10//!   q8_row. This is what lets a 15B file run in a few GB of RSS.
11//!
12//! Extension point: new dtypes = new match arm here, nothing else moves.
13
14use crate::pool::{Pool, matvec_rows, matvec_rows2};
15use cortiq_core::quant::{
16    GROUP_SIZE, Q1_TILE, Q2TP_CHUNK, Q4_TILE, Q4TP_NIB, f16_to_f32, q2tp_ladder, q2tp_sections,
17    q4tp_code, q4tp_ladder, q4tp_sections,
18};
19use cortiq_core::{CmfModel, TensorDtype};
20use std::cell::UnsafeCell;
21use std::sync::Arc;
22
23pub enum QTensor {
24    F32 {
25        data: Vec<f32>,
26        rows: usize,
27        cols: usize,
28    },
29    Mapped {
30        model: Arc<CmfModel>,
31        /// Index into the model's tensor directory.
32        idx: usize,
33        dtype: TensorDtype,
34        rows: usize,
35        cols: usize,
36        /// Per-row scales, dequantized to f32 up front (tiny).
37        row_scale: Vec<f32>,
38        /// q8_2f column field (θ), dequantized up front; empty for q8_row.
39        col_field: Vec<f32>,
40        /// Vbit only: byte offset of each row's packed data within the
41        /// tensor blob (`[rows + 1]`, computed once at load — the per-
42        /// matvec prefix scan over row bit-widths was O(rows) each call).
43        vbit_offsets: Vec<usize>,
44        /// q8-family decode repack (load-time, optional): rows in groups
45        /// of 4, interleaved in 16-byte units — one 64-byte line per
46        /// iteration feeds all 4 sdot lanes, ONE sequential weight
47        /// stream per worker instead of four (this is where llama.cpp's
48        /// repacked Q8 kernels get their bandwidth). Empty = off
49        /// (CMF_REPACK=0, non-SDOT arch, or an ineligible shape). Trades
50        /// an anonymous copy of the quants for mmap pages that go cold.
51        repack: Vec<u8>,
52    },
53    /// Binary low-rank (`blr`): ±1 paths, signs from the mmap, scales as
54    /// f32 (docs/BLR_NVG.md §3). Every device-graph accessor answers `None`
55    /// for it (whole-token graphs decline to per-op/host paths); the CPU
56    /// kernels live in `crate::blr`, the per-op wgpu kernels in
57    /// `gpu_wgpu::blr`.
58    Blr(Box<crate::blr::BlrTensor>),
59}
60
61/// Load-time q8 repack gate (see `Mapped::repack`). OPT-IN
62/// (`CMF_REPACK=1`): the single-stream hypothesis LOST on Apple Silicon
63/// (M4, interleaved A/B: decode 101 vs 94 tok/s — four adjacent row
64/// streams per worker feed the prefetcher MORE memory-level parallelism
65/// than one); kept as an experiment flag for x86, where the tradeoff
66/// may land differently.
67fn repack_enabled() -> bool {
68    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
69    *ON.get_or_init(|| {
70        std::env::var("CMF_REPACK")
71            .map(|v| v == "1")
72            .unwrap_or(cfg!(target_os = "android"))
73    })
74}
75
76/// Interleave q8 rows for the decode kernel: group g holds rows
77/// 4g..4g+4 as [r0[c], r1[c], r2[c], r3[c]] per 16-byte chunk c. Only
78/// full groups are packed — tail rows keep reading the mmap layout.
79fn q8_repack(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
80    #[cfg(target_arch = "aarch64")]
81    let arch_ok = sdot_enabled();
82    #[cfg(not(target_arch = "aarch64"))]
83    let arch_ok = false;
84    if !arch_ok || !repack_enabled() || rows < 256 || cols % 16 != 0 {
85        return Vec::new();
86    }
87    q8_repack_layout(bytes, rows, cols)
88}
89
90/// The pure layout transform behind `q8_repack` (tested directly —
91/// the gate depends on arch and env).
92fn q8_repack_layout(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
93    let groups = rows / 4;
94    let mut rep = vec![0u8; groups * 4 * cols];
95    for g in 0..groups {
96        let dst = &mut rep[g * 4 * cols..(g + 1) * 4 * cols];
97        for c in 0..cols / 16 {
98            for lane in 0..4 {
99                let src = (g * 4 + lane) * cols + c * 16;
100                dst[c * 64 + lane * 16..c * 64 + lane * 16 + 16]
101                    .copy_from_slice(&bytes[src..src + 16]);
102            }
103        }
104    }
105    rep
106}
107
108/// Prefix-sum of vbit row payload offsets (absolute within the tensor
109/// bytes). `offsets[r]..offsets[r+1]` is row r's packed data.
110fn vbit_row_offsets(bytes: &[u8], rows: usize, cols: usize) -> Vec<usize> {
111    let ng = cols / GROUP_SIZE;
112    let bits = &bytes[..rows];
113    let mut offsets = Vec::with_capacity(rows + 1);
114    let mut off = rows + rows * ng * 2;
115    for r in 0..rows {
116        offsets.push(off);
117        off += (cols * bits[r] as usize).div_ceil(8);
118    }
119    offsets.push(off);
120    offsets
121}
122
123/// `CMF_X86_BLOCKED` / `CMF_GPU_LMHEAD` / `CMF_GPU_SPLIT`, read once. They
124/// used to be read from the environment on every large matvec and on every
125/// matmat in six places — microseconds each, but also a knob that could
126/// change under a running process, which is not a thing a kernel choice
127/// should be able to do mid-sequence.
128fn blocked_enabled() -> bool {
129    use std::sync::atomic::Ordering::Relaxed;
130    match BLOCKED_OVERRIDE.load(Relaxed) {
131        1 => false,
132        2 => true,
133        _ => {
134            static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
135            *ON.get_or_init(|| {
136                std::env::var("CMF_X86_BLOCKED")
137                    .map(|v| v != "0")
138                    .unwrap_or(true)
139            })
140        }
141    }
142}
143
144static BLOCKED_OVERRIDE: std::sync::atomic::AtomicU8 = std::sync::atomic::AtomicU8::new(0);
145
146/// Force the blocked GEMM on or off, ignoring the environment; `None`
147/// restores it. For tests that need to run BOTH paths and compare them:
148/// `blocked_enabled` caches its answer for the life of the process, which
149/// is right when the environment is the only input, but leaves a test that
150/// flips `CMF_X86_BLOCKED` between two calls comparing a path against
151/// itself — or against whatever a test running in parallel latched first.
152pub fn set_blocked_override(on: Option<bool>) {
153    let v = match on {
154        None => 0,
155        Some(false) => 1,
156        Some(true) => 2,
157    };
158    BLOCKED_OVERRIDE.store(v, std::sync::atomic::Ordering::Relaxed);
159}
160
161fn gpu_lmhead_enabled() -> bool {
162    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
163    *ON.get_or_init(|| {
164        std::env::var("CMF_GPU_LMHEAD")
165            .map(|v| v != "0")
166            .unwrap_or(true)
167    })
168}
169
170fn gpu_split_frac() -> f32 {
171    // MiMo's banked verification projects every row on the device. Plain
172    // decode must not quantize half of its O/head activation on the CPU.
173    if FULL_GPU_Q8.get() {
174        return 1.0;
175    }
176    static F: std::sync::OnceLock<f32> = std::sync::OnceLock::new();
177    *F.get_or_init(|| {
178        std::env::var("CMF_GPU_SPLIT")
179            .ok()
180            .and_then(|v| v.parse::<f32>().ok())
181            .unwrap_or(0.5)
182            .clamp(0.0, 1.0)
183    })
184}
185
186impl QTensor {
187    pub fn from_f32(data: Vec<f32>, rows: usize, cols: usize) -> Self {
188        debug_assert_eq!(data.len(), rows * cols);
189        Self::F32 { data, rows, cols }
190    }
191
192    /// Wrap a directory tensor without dequantizing the payload.
193    /// Falls back to dequantized f32 for dtypes without a fused kernel.
194    pub fn from_model(model: &Arc<CmfModel>, name: &str) -> Result<Self, String> {
195        // Indexed lookup: the linear directory scan made pipeline build
196        // O(N²) on MoE/skills files with thousands of tensors.
197        let idx = model
198            .tensor_index(name)
199            .ok_or_else(|| format!("tensor '{name}' not found in CMF directory"))?;
200        let entry = &model.tensors[idx];
201        if entry.shape.len() != 2 {
202            return Err(format!("QTensor::from_model needs 2-D, got '{name}'"));
203        }
204        let (rows, cols) = (entry.shape[0], entry.shape[1]);
205        let bytes = model.entry_bytes(entry);
206
207        match entry.dtype {
208            TensorDtype::Q8Row | TensorDtype::Q8_2f => {
209                let n = rows * cols;
210                let scales_off = n;
211                let row_scale: Vec<f32> = (0..rows)
212                    .map(|o| {
213                        f16_to_f32(u16::from_le_bytes([
214                            bytes[scales_off + o * 2],
215                            bytes[scales_off + o * 2 + 1],
216                        ]))
217                    })
218                    .collect();
219                let col_field: Vec<f32> = if entry.dtype == TensorDtype::Q8_2f {
220                    let col_off = n + rows * 2;
221                    (0..cols)
222                        .map(|i| {
223                            f16_to_f32(u16::from_le_bytes([
224                                bytes[col_off + i * 2],
225                                bytes[col_off + i * 2 + 1],
226                            ]))
227                        })
228                        .collect()
229                } else {
230                    Vec::new()
231                };
232                Ok(Self::Mapped {
233                    model: model.clone(),
234                    idx,
235                    dtype: entry.dtype,
236                    rows,
237                    cols,
238                    row_scale,
239                    col_field,
240                    vbit_offsets: Vec::new(),
241                    repack: q8_repack(bytes, rows, cols),
242                })
243            }
244            // vbit: fused kernel unpacks variable-bit rows from mmap.
245            TensorDtype::Vbit if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
246                model: model.clone(),
247                idx,
248                dtype: entry.dtype,
249                rows,
250                cols,
251                row_scale: Vec::new(),
252                col_field: Vec::new(),
253                vbit_offsets: vbit_row_offsets(bytes, rows, cols),
254                repack: Vec::new(),
255            }),
256            // vbit_ro (§4.2): the offset table comes straight from the
257            // file — no load-time prefix scan; kernels are shared with
258            // legacy vbit (they consume absolute offsets either way).
259            TensorDtype::VbitRo if cols % GROUP_SIZE == 0 => {
260                let (_, off_off, packed_off) = cortiq_core::quant::vbit_ro_sections(rows, cols);
261                let offsets: Vec<usize> = (0..=rows)
262                    .map(|r| packed_off + cortiq_core::quant::vbit_ro_offset(bytes, off_off, r))
263                    .collect();
264                Ok(Self::Mapped {
265                    model: model.clone(),
266                    idx,
267                    dtype: entry.dtype,
268                    rows,
269                    cols,
270                    row_scale: Vec::new(),
271                    col_field: Vec::new(),
272                    vbit_offsets: offsets,
273                    repack: Vec::new(),
274                })
275            }
276            // q4_block: fused kernel reads nibbles straight from mmap —
277            // a 14B q4 file no longer explodes into ×8 f32 RAM.
278            // q4_tiled (§4.3): interleaved [scale][nibbles] tiles — one
279            // sequential memory stream (measured ×1.66 ARM / ×1.13 AVX2
280            // at kernel level over the split layout).
281            TensorDtype::Q4Tiled if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
282                model: model.clone(),
283                idx,
284                dtype: entry.dtype,
285                rows,
286                cols,
287                row_scale: Vec::new(),
288                col_field: Vec::new(),
289                vbit_offsets: Vec::new(),
290                repack: Vec::new(),
291            }),
292            // q4tp (§4.10): nibbles from mmap, scale from the row ladder —
293            // 7.3% less file than q4t at the same 4-bit grid.
294            TensorDtype::Q4TiledP if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
295                model: model.clone(),
296                idx,
297                dtype: entry.dtype,
298                rows,
299                cols,
300                row_scale: Vec::new(),
301                col_field: Vec::new(),
302                vbit_offsets: Vec::new(),
303                repack: Vec::new(),
304            }),
305            // q2tp: 2-bit chunks from mmap, scale from the same row ladder.
306            TensorDtype::Q2TiledP if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
307                model: model.clone(),
308                idx,
309                dtype: entry.dtype,
310                rows,
311                cols,
312                row_scale: Vec::new(),
313                col_field: Vec::new(),
314                vbit_offsets: Vec::new(),
315                repack: Vec::new(),
316            }),
317            TensorDtype::Q4Block if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
318                model: model.clone(),
319                idx,
320                dtype: entry.dtype,
321                rows,
322                cols,
323                row_scale: Vec::new(),
324                col_field: Vec::new(),
325                vbit_offsets: Vec::new(),
326                repack: Vec::new(),
327            }),
328            // q1: binary sign-bit tiles from mmap (1-bit-trained models).
329            // q1g: the same signs as planes with one scale per 128-group.
330            TensorDtype::Q1 if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
331                model: model.clone(),
332                idx,
333                dtype: entry.dtype,
334                rows,
335                cols,
336                row_scale: Vec::new(),
337                col_field: Vec::new(),
338                vbit_offsets: Vec::new(),
339                repack: Vec::new(),
340            }),
341            TensorDtype::Q1G if cols % cortiq_core::quant::Q1G_GROUP == 0 => Ok(Self::Mapped {
342                model: model.clone(),
343                idx,
344                dtype: entry.dtype,
345                rows,
346                cols,
347                row_scale: Vec::new(),
348                col_field: Vec::new(),
349                vbit_offsets: Vec::new(),
350                repack: Vec::new(),
351            }),
352            // q1t (ternary + outlier overlay): fused per-row dequant kernel
353            // reads straight from mmap — a 12B q1t stays ~its file size in
354            // RAM instead of dequantizing to ~48 GB of f32.
355            TensorDtype::Q1T if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
356                model: model.clone(),
357                idx,
358                dtype: entry.dtype,
359                rows,
360                cols,
361                row_scale: Vec::new(),
362                col_field: Vec::new(),
363                vbit_offsets: Vec::new(),
364                repack: Vec::new(),
365            }),
366            // blr: factored kernels over the packed signs. CMF_BLR_DENSE=1
367            // dequantizes to f32 instead (A/B parity runs).
368            TensorDtype::BinaryLowRank if std::env::var("CMF_BLR_DENSE").as_deref() != Ok("1") => {
369                Ok(Self::Blr(Box::new(crate::blr::BlrTensor::from_model(
370                    model, idx,
371                )?)))
372            }
373            // No fused kernel yet → dequantize once (correct, more RAM).
374            _ => {
375                let mut data = vec![0.0f32; rows * cols];
376                cortiq_core::quant::dequant_tensor(entry, bytes, &mut data)?;
377                Ok(Self::from_f32(data, rows, cols))
378            }
379        }
380    }
381
382    /// Binary low-rank tensor? Fused multi-tensor kernels refuse it and fall
383    /// back to per-tensor calls.
384    pub fn is_blr(&self) -> bool {
385        matches!(self, Self::Blr(_))
386    }
387
388    /// Binary-weight-mapped tensor, q1 or q1g? (GPU gates: the binary CPU
389    /// kernels are compute-bound, so offload pays at much smaller shapes
390    /// than q8.)
391    pub(crate) fn is_q1(&self) -> bool {
392        matches!(
393            self,
394            Self::Mapped {
395                dtype: TensorDtype::Q1 | TensorDtype::Q1G,
396                ..
397            }
398        )
399    }
400
401    /// The binary-weight view (q1 tiles or q1g planes) the shared CPU row
402    /// drivers take; None for every other storage.
403    fn bin_w(&self) -> Option<BinW<'_>> {
404        match self {
405            Self::Mapped {
406                dtype: TensorDtype::Q1,
407                ..
408            } => Some(BinW::Q1(self.quant_bytes())),
409            Self::Mapped {
410                dtype: TensorDtype::Q1G,
411                rows,
412                cols,
413                ..
414            } => Some(BinW::q1g(self.quant_bytes(), *rows, *cols)),
415            _ => None,
416        }
417    }
418
419    /// Owned-f32 view (data, rows, cols) — the GDN a/b gate projections
420    /// arrive dequantized (force-f16 in the converter → F32 in RAM).
421    pub(crate) fn f32_parts(&self) -> Option<(&[f32], usize, usize)> {
422        match self {
423            Self::F32 { data, rows, cols } => Some((data, *rows, *cols)),
424            _ => None,
425        }
426    }
427
428    /// (directory idx, rows, cols) of a q1-mapped tensor — the
429    /// whole-block GPU path resolves offsets itself.
430    /// (idx, rows, cols) of a mapped tensor the whole-token GPU graph can drive
431    /// — Q1, Q1T or Q4-block (it resolves the offset and picks the kernel by
432    /// dtype). Q4-block lets a precise down_proj/lm_head stay on-device.
433    /// Named `q1_parts` for historical reasons.
434    pub(crate) fn q1_parts(&self) -> Option<(usize, usize, usize)> {
435        if self.has_prism_contract() {
436            return None;
437        }
438        match self {
439            #[cfg(target_os = "macos")]
440            Self::Mapped {
441                dtype: TensorDtype::Q1T,
442                ..
443            } if !crate::gpu::metal_q1t_enabled() => None,
444            Self::Mapped {
445                idx,
446                dtype:
447                    TensorDtype::Q1
448                    | TensorDtype::Q1G
449                    | TensorDtype::Q1T
450                    | TensorDtype::Q4Block
451                    | TensorDtype::Q4Tiled
452                    // Q2TiledP deliberately absent: the Metal graph has no
453                    // q2tp kernel, and advertising it here made the block
454                    // plan truncate mid-run at the first q2tp layer.
455                    | TensorDtype::Q4TiledP
456                    | TensorDtype::Q8Row
457                    | TensorDtype::Q8_2f,
458                rows,
459                cols,
460                ..
461            } => Some((*idx, *rows, *cols)),
462            _ => None,
463        }
464    }
465
466    /// `(directory idx, rows, cols)` for the native Metal token graph.  The
467    /// historical q1 graph gate intentionally refuses every Prism tensor so
468    /// an untransformed q2 payload cannot slip into the resident path.  The
469    /// Metal2 graph is descriptor-aware and admits only the production
470    /// q2tp-affine forward targets; ordinary q1/q4 callers retain the old
471    /// `q1_parts` behaviour.
472    #[cfg(target_os = "macos")]
473    pub(crate) fn metal_graph_parts(&self) -> Option<(usize, usize, usize)> {
474        if let Some((model, idx, kind, _)) = self.graph_weight_descriptor() {
475            let name = &model.tensors[idx].name;
476            let forward = kind == 9 && crate::prism::is_forward_weight(model, name);
477            let affine = kind == 9 && crate::prism::is_affine_target(model, name);
478            if forward && affine {
479                let e = model.tensors.get(idx)?;
480                return Some((idx, *e.shape.first()?, *e.shape.get(1)?));
481            }
482        }
483        self.q1_parts()
484    }
485
486    /// (directory idx, rows, cols) of a q4_tiled mapped tensor. The
487    /// chunk-prefill graph takes it in the same 4-tuple slot as
488    /// `q8_row_parts` with an EMPTY row_scale — q4t carries its scales
489    /// inside the 18-byte tiles, and the empty slice is what tells the
490    /// encoder to reach for the q4t kernels.
491    pub(crate) fn q4t_parts(&self) -> Option<(usize, usize, usize)> {
492        if self.has_prism_contract() {
493            return None;
494        }
495        match self {
496            Self::Mapped {
497                idx,
498                dtype: TensorDtype::Q4Tiled,
499                rows,
500                cols,
501                ..
502            } => Some((*idx, *rows, *cols)),
503            _ => None,
504        }
505    }
506
507    /// (directory idx, rows, cols) of a q4tp mapped tensor. Same empty-scale
508    /// slot as `q4t_parts` in the chunk graph — the encoder tells the two
509    /// apart by the tensor's dtype, not by the slot.
510    pub(crate) fn q4tp_parts(&self) -> Option<(usize, usize, usize)> {
511        if self.has_prism_contract() {
512            return None;
513        }
514        match self {
515            Self::Mapped {
516                idx,
517                dtype: TensorDtype::Q4TiledP,
518                rows,
519                cols,
520                ..
521            } => Some((*idx, *rows, *cols)),
522            _ => None,
523        }
524    }
525
526    /// (directory idx, rows, cols, row_scale) of a plain q8_row mapped
527    /// tensor — the chunk-prefill GPU graph resolves offsets itself.
528    /// q8_2f is excluded on purpose: its column field would need a
529    /// prescale stage on the device.
530    pub(crate) fn q8_row_parts(&self) -> Option<(usize, usize, usize, &[f32])> {
531        if self.has_prism_contract() {
532            return None;
533        }
534        match self {
535            Self::Mapped {
536                idx,
537                dtype: TensorDtype::Q8Row,
538                rows,
539                cols,
540                row_scale,
541                col_field,
542                ..
543            } if col_field.is_empty() => Some((*idx, *rows, *cols, row_scale)),
544            _ => None,
545        }
546    }
547
548    /// (directory idx, rows, cols, row_scale, col_field) of a q8_2f mapped
549    /// tensor — the chunk-prefill graph stages x·col on the device first.
550    pub(crate) fn q8_2f_parts(&self) -> Option<(usize, usize, usize, &[f32], &[f32])> {
551        if self.has_prism_contract() {
552            return None;
553        }
554        match self {
555            Self::Mapped {
556                idx,
557                dtype: TensorDtype::Q8_2f,
558                rows,
559                cols,
560                row_scale,
561                col_field,
562                ..
563            } if col_field.len() == *cols => Some((*idx, *rows, *cols, row_scale, col_field)),
564            _ => None,
565        }
566    }
567
568    /// The layout this tensor is stored in, when it is mapped from a model.
569    /// The frames branch on it — a q2tp gate against a q4tp down is a real
570    /// combination in the 2-bit profile and needs a different kernel.
571    pub fn model_dtype(&self) -> Option<cortiq_core::TensorDtype> {
572        match self {
573            Self::Mapped { dtype, .. } => Some(*dtype),
574            _ => None,
575        }
576    }
577
578    /// The tensor's index in the model directory, when it is mapped from one.
579    /// The GPU frames bind by index rather than by name — a name lookup per
580    /// layer per token is not free, and the index is what the device cache is
581    /// keyed on anyway.
582    pub fn model_idx(&self) -> Option<usize> {
583        match self {
584            Self::Mapped { idx, .. } => Some(*idx),
585            _ => None,
586        }
587    }
588
589    /// The model this tensor is mapped from, when it is mapped at all. The
590    /// GPU frames need the container to reach the bytes; a QTensor already
591    /// holds it, and threading a second handle down every call site to say
592    /// the same thing invites the two to disagree.
593    pub fn model_arc(&self) -> Option<std::sync::Arc<cortiq_core::CmfModel>> {
594        match self {
595            Self::Mapped { model, .. } => Some(model.clone()),
596            _ => None,
597        }
598    }
599
600    /// Whether this mapped tensor belongs to the Prism/Bonsai transform
601    /// contract.  Device graphs do not carry the descriptor, so callers use
602    /// this conservative predicate to stay on the descriptor-aware CPU path
603    /// instead of silently executing an unrotated matrix.
604    pub(crate) fn has_prism_contract(&self) -> bool {
605        matches!(self, Self::Mapped { model, .. } if crate::prism::has_contract(model))
606    }
607
608    pub fn rows(&self) -> usize {
609        match self {
610            Self::F32 { rows, .. } | Self::Mapped { rows, .. } => *rows,
611            Self::Blr(t) => t.rows,
612        }
613    }
614
615    /// Mapped q4t handle (model + directory index) — the fused GPU FFN
616    /// needs the raw file coordinates of its three projections.
617    pub(crate) fn mapped_q4t(&self) -> Option<(&Arc<CmfModel>, usize)> {
618        if self.has_prism_contract() {
619            return None;
620        }
621        match self {
622            Self::Mapped {
623                model,
624                idx,
625                dtype: TensorDtype::Q4Tiled,
626                ..
627            } => Some((model, *idx)),
628            _ => None,
629        }
630    }
631
632    /// Same slot as `mapped_q4t` for a q4tp tensor — the fused DiT FFN picks
633    /// its kernels by which of the two answers.
634    pub fn mapped_q4tp(&self) -> Option<(&Arc<CmfModel>, usize)> {
635        if self.has_prism_contract() {
636            return None;
637        }
638        match self {
639            Self::Mapped {
640                model,
641                idx,
642                dtype: TensorDtype::Q4TiledP,
643                ..
644            } => Some((model, *idx)),
645            _ => None,
646        }
647    }
648
649    /// (model, tensor idx) for a mapped weight in ANY codec the fused device
650    /// paths can run — four-bit tiled or either int8 layout.
651    ///
652    /// The fused DiT chains asked for `mapped_q4tp` by name, so an eight-bit
653    /// container never reached them and rendered through per-op GEMMs even
654    /// after those kernels learned its codec. The gate is what the codec has
655    /// a device GEMM for, not which codec it is.
656    pub fn mapped_device_gemm(&self) -> Option<(&Arc<CmfModel>, usize)> {
657        if self.has_prism_contract() {
658            return None;
659        }
660        match self {
661            Self::Mapped {
662                model,
663                idx,
664                dtype: TensorDtype::Q4TiledP | TensorDtype::Q8Row | TensorDtype::Q8_2f,
665                ..
666            } => Some((model, *idx)),
667            _ => None,
668        }
669    }
670
671    /// (model, tensor idx) for a q8_2f mapped weight — the eight-bit slot of
672    /// the MoE graph ladder (int8 body, f16 row scales, f16 column field).
673    pub fn mapped_q8_2f(&self) -> Option<(&Arc<CmfModel>, usize)> {
674        if self.has_prism_contract() {
675            return None;
676        }
677        match self {
678            Self::Mapped {
679                model,
680                idx,
681                dtype: TensorDtype::Q8_2f,
682                ..
683            } => Some((model, *idx)),
684            _ => None,
685        }
686    }
687
688    /// (model, tensor idx) for a q2tp mapped weight — the 2-bit twin of
689    /// `mapped_q4tp`, used by the mixed MoE profile.
690    pub fn mapped_q2tp(&self) -> Option<(&Arc<CmfModel>, usize)> {
691        if self.has_prism_contract() {
692            return None;
693        }
694        match self {
695            Self::Mapped {
696                model,
697                idx,
698                dtype: TensorDtype::Q2TiledP,
699                ..
700            } => Some((model, *idx)),
701            _ => None,
702        }
703    }
704
705    pub fn cols(&self) -> usize {
706        match self {
707            Self::F32 { cols, .. } | Self::Mapped { cols, .. } => *cols,
708            Self::Blr(t) => t.cols,
709        }
710    }
711
712    /// (model, tensor idx) for a q1 mapped weight — the wgpu token graph
713    /// keys its resident VRAM cache by idx. None for any other dtype/kind.
714    pub fn mapped_q1(&self) -> Option<(&std::sync::Arc<CmfModel>, usize)> {
715        if self.has_prism_contract() {
716            return None;
717        }
718        match self {
719            Self::Mapped {
720                model,
721                idx,
722                dtype: TensorDtype::Q1,
723                ..
724            } => Some((model, *idx)),
725            _ => None,
726        }
727    }
728
729    /// (model, idx, kind, row_scale) for a graph-capable mapped weight.
730    /// kind: 0=q8_row (per-row scales), 1=q1, 2=q4_block, 3=q1t
731    /// (tile-embedded, no rs), 5=q4_tiled, 6=q4tp, 7=q8_2f (both scale
732    /// planes live inside the tensor), 9=q2tp, 10=q1g (sign + scale
733    /// planes). None only for `vbit`.
734    ///
735    /// The old comment here claimed q4_block was unhandled while the arm
736    /// right below mapped it, and it named q8_2f as unhandled after that
737    /// stopped being true — a stale comment on this function is how a
738    /// model silently loses the graph, so it is worth keeping honest.
739    pub fn graph_weight(&self) -> Option<(&std::sync::Arc<CmfModel>, usize, u8, &[f32])> {
740        if self.has_prism_contract() {
741            return None;
742        }
743        self.graph_weight_descriptor()
744    }
745
746    /// Descriptor-aware graph handle used only by the Prism token graph.
747    /// Ordinary graph callers continue to use [`graph_weight`] and therefore
748    /// remain fail-closed until they provide the same explicit transform
749    /// contract.
750    pub(crate) fn graph_weight_descriptor(
751        &self,
752    ) -> Option<(&std::sync::Arc<CmfModel>, usize, u8, &[f32])> {
753        match self {
754            Self::Mapped {
755                model,
756                idx,
757                dtype: TensorDtype::Q8Row,
758                row_scale,
759                ..
760            } => Some((model, *idx, 0, row_scale.as_slice())),
761            Self::Mapped {
762                model,
763                idx,
764                dtype: TensorDtype::Q1,
765                ..
766            } => Some((model, *idx, 1, &[])),
767            // Kind 10, never 1: q1g's sign and scale planes are a different
768            // walk from q1's 6-byte tiles (the kind-2/5 mixup taught that).
769            Self::Mapped {
770                model,
771                idx,
772                dtype: TensorDtype::Q1G,
773                ..
774            } => Some((model, *idx, 10, &[])),
775            // Q4Tiled is kind 5, NOT 2: both carried 2 historically, and
776            // the wgpu token graph fed 18B interleaved tiles to the
777            // split-layout q4b kernel — garbage output on q4t models
778            // (caught by an end-to-end answer check on real Vulkan).
779            Self::Mapped {
780                model,
781                idx,
782                dtype: TensorDtype::Q4Tiled,
783                ..
784            } => Some((model, *idx, 5, &[])),
785            // Kind 6, not 5: q4tp's nibble stride and scale planes differ,
786            // and feeding them to the q4t kernel is exactly the mistake that
787            // produced garbage when Q4Tiled shared kind 2 with Q4Block.
788            Self::Mapped {
789                model,
790                idx,
791                dtype: TensorDtype::Q4TiledP,
792                ..
793            } => Some((model, *idx, 6, &[])),
794            Self::Mapped {
795                model,
796                idx,
797                dtype: TensorDtype::Q4Block,
798                ..
799            } => Some((model, *idx, 2, &[])),
800            // q8_2f carries BOTH scale planes after the int8 body (rows
801            // f16, then cols f16), so the graph takes the whole tensor
802            // and the kernel reads them where they lie — no host-side
803            // prescale, which is what the per-op path does instead.
804            Self::Mapped {
805                model,
806                idx,
807                dtype: TensorDtype::Q8_2f,
808                ..
809            } => Some((model, *idx, 7, &[])),
810            Self::Mapped {
811                model,
812                idx,
813                dtype: TensorDtype::Q1T,
814                ..
815            } => Some((model, *idx, 3, &[])),
816            // Kind 9: the 2-bit plane on the q4tp ladder (dense FFN gate/up
817            // of the q2tp profile). Its own kernel — 8 bytes a group where
818            // q4tp has 16, and rung 0 is the exact zero.
819            Self::Mapped {
820                model,
821                idx,
822                dtype: TensorDtype::Q2TiledP,
823                ..
824            } => Some((model, *idx, 9, &[])),
825            _ => None,
826        }
827    }
828
829    /// Dense f32 view — only for owned tensors. Masked/sparse execution
830    /// paths require it; quantized weights don't support masks yet.
831    pub fn as_f32(&self) -> Option<&[f32]> {
832        match self {
833            Self::F32 { data, .. } => Some(data),
834            Self::Mapped { .. } | Self::Blr(_) => None,
835        }
836    }
837
838    fn quant_bytes(&self) -> &[u8] {
839        match self {
840            Self::Mapped { model, idx, .. } => model.entry_bytes(&model.tensors[*idx]),
841            Self::Blr(t) => t.bytes(),
842            Self::F32 { .. } => unreachable!("quant_bytes on F32"),
843        }
844    }
845
846    /// Dequantize one row into `dst` (embedding lookup).
847    pub fn row_f32(&self, r: usize, dst: &mut [f32]) {
848        let cols = self.cols();
849        debug_assert_eq!(dst.len(), cols);
850        match self {
851            Self::F32 { data, .. } => dst.copy_from_slice(&data[r * cols..(r + 1) * cols]),
852            Self::Blr(t) => t.row_f32(r, dst),
853            Self::Mapped {
854                model,
855                idx,
856                dtype,
857                row_scale,
858                col_field,
859                vbit_offsets,
860                ..
861            } => {
862                if *dtype == TensorDtype::Q4Tiled {
863                    let bytes = self.quant_bytes();
864                    let gpr = cols / GROUP_SIZE;
865                    for gi in 0..gpr {
866                        let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
867                        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
868                        for (k, &b) in tile[2..].iter().enumerate() {
869                            dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
870                            dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
871                        }
872                    }
873                    if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
874                        crate::prism::inverse_embedding(model, dst);
875                    }
876                    return;
877                }
878                if *dtype == TensorDtype::Q4TiledP {
879                    let bytes = self.quant_bytes();
880                    let gpr = cols / GROUP_SIZE;
881                    let v = Q4tpView::new(bytes, self.rows(), cols);
882                    let mut sc = vec![0f32; gpr];
883                    v.scales_into(r, gpr, &mut sc);
884                    for gi in 0..gpr {
885                        let tile = &v.nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
886                        let s = sc[gi];
887                        for (k, &b) in tile.iter().enumerate() {
888                            dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
889                            dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
890                        }
891                    }
892                    if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
893                        crate::prism::inverse_embedding(model, dst);
894                    }
895                    return;
896                }
897                if *dtype == TensorDtype::Q2TiledP {
898                    let bytes = self.quant_bytes();
899                    let gpr = cols / GROUP_SIZE;
900                    let v = Q4tpView::new_q2(bytes, self.rows(), cols);
901                    let mut sc = vec![0f32; gpr];
902                    v.scales_into(r, gpr, &mut sc);
903                    for gi in 0..gpr {
904                        let ch =
905                            &v.nib[(r * gpr + gi) * Q2TP_CHUNK..(r * gpr + gi + 1) * Q2TP_CHUNK];
906                        let s = sc[gi];
907                        for (k, &b) in ch.iter().enumerate() {
908                            for j in 0..4 {
909                                let center = if crate::prism::is_affine_target(
910                                    model,
911                                    &model.tensors[*idx].name,
912                                ) {
913                                    1.0
914                                } else {
915                                    1.5
916                                };
917                                dst[gi * GROUP_SIZE + k * 4 + j] =
918                                    (((b >> (2 * j)) & 3) as f32 - center) * s;
919                            }
920                        }
921                    }
922                    if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
923                        crate::prism::inverse_embedding(model, dst);
924                    }
925                    return;
926                }
927                if *dtype == TensorDtype::Q4Block {
928                    let (packed, scales) = q4_split(self.quant_bytes(), self.rows(), cols);
929                    let gpr = cols / GROUP_SIZE;
930                    for gi in 0..gpr {
931                        let g = r * gpr + gi;
932                        let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
933                        for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
934                            dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
935                            dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
936                        }
937                    }
938                    if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
939                        crate::prism::inverse_embedding(model, dst);
940                    }
941                    return;
942                }
943                if *dtype == TensorDtype::Q1 {
944                    let bytes = self.quant_bytes();
945                    let gpr = cols / GROUP_SIZE;
946                    for gi in 0..gpr {
947                        let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
948                        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
949                        for (j, &b) in tile[2..].iter().enumerate() {
950                            for k in 0..8 {
951                                dst[gi * GROUP_SIZE + j * 8 + k] =
952                                    (((b >> k) & 1) as f32 * 2.0 - 1.0) * s;
953                            }
954                        }
955                    }
956                    if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
957                        crate::prism::inverse_embedding(model, dst);
958                    }
959                    return;
960                }
961                if *dtype == TensorDtype::Q1G {
962                    // The tied Bonsai embedding is q1g too: one row = cols/8
963                    // sign bytes + cols/128 scales, same ±s values as q1.
964                    let bytes = self.quant_bytes();
965                    let sc_off = cortiq_core::quant::q1g_scale_offset(self.rows(), cols);
966                    let gpr = cols / GROUP_SIZE;
967                    for gi in 0..gpr {
968                        let t = r * gpr + gi;
969                        let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, t));
970                        for (j, &b) in bytes[t * 4..t * 4 + 4].iter().enumerate() {
971                            for k in 0..8 {
972                                dst[gi * GROUP_SIZE + j * 8 + k] =
973                                    (((b >> k) & 1) as f32 * 2.0 - 1.0) * s;
974                            }
975                        }
976                    }
977                    if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
978                        crate::prism::inverse_embedding(model, dst);
979                    }
980                    return;
981                }
982                if *dtype == TensorDtype::Q1T {
983                    let bytes = self.quant_bytes();
984                    let gpr = cols / GROUP_SIZE;
985                    let base_len = self.rows() * gpr * cortiq_core::quant::Q1T_TILE;
986                    for gi in 0..gpr {
987                        let off = (r * gpr + gi) * cortiq_core::quant::Q1T_TILE;
988                        let s = cortiq_core::quant::f16_to_f32(u16::from_le_bytes([
989                            bytes[off],
990                            bytes[off + 1],
991                        ]));
992                        let codes = &bytes[off + 2..off + cortiq_core::quant::Q1T_TILE];
993                        for k in 0..GROUP_SIZE {
994                            dst[gi * GROUP_SIZE + k] = match cortiq_core::quant::q1t_code(codes, k)
995                            {
996                                1 => s,
997                                2 => -s,
998                                _ => 0.0,
999                            };
1000                        }
1001                    }
1002                    // Overlay
1003                    let rows = self.rows();
1004                    let entries = base_len + (rows + 1) * 4;
1005                    if entries <= bytes.len() {
1006                        let ptrs = &bytes[base_len..base_len + (rows + 1) * 4];
1007                        let r0 = u32::from_le_bytes([
1008                            ptrs[r * 4],
1009                            ptrs[r * 4 + 1],
1010                            ptrs[r * 4 + 2],
1011                            ptrs[r * 4 + 3],
1012                        ]) as usize;
1013                        let r1 = u32::from_le_bytes([
1014                            ptrs[(r + 1) * 4],
1015                            ptrs[(r + 1) * 4 + 1],
1016                            ptrs[(r + 1) * 4 + 2],
1017                            ptrs[(r + 1) * 4 + 3],
1018                        ]) as usize;
1019                        let off = entries + r0 * 4;
1020                        for i in 0..r1 - r0 {
1021                            let item = &bytes[off + i * 4..off + i * 4 + 4];
1022                            let c = u16::from_le_bytes([item[0], item[1]]) as usize;
1023                            let v = cortiq_core::quant::f16_to_f32(u16::from_le_bytes([
1024                                item[2], item[3],
1025                            ]));
1026                            if c < cols {
1027                                dst[c] = v;
1028                            }
1029                        }
1030                    }
1031                    if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
1032                        crate::prism::inverse_embedding(model, dst);
1033                    }
1034                    return;
1035                }
1036                if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1037                    let bytes = self.quant_bytes();
1038                    let rows = self.rows();
1039                    let ng = cols / GROUP_SIZE;
1040                    let bits = &bytes[..rows];
1041                    let sc_off = rows;
1042                    // Precomputed at load — embedding lookup used to scan
1043                    // the bit-widths of every preceding row (O(token_id)).
1044                    let off = vbit_offsets[r];
1045                    let b = bits[r] as usize;
1046                    let l = ((1usize << (b - 1)) - 1) as f32;
1047                    let data = &bytes[off..];
1048                    let (mut acc, mut nbits, mut byte_idx) = (0u64, 0usize, 0usize);
1049                    for (i, d) in dst.iter_mut().enumerate() {
1050                        while nbits < b {
1051                            acc = (acc << 8) | data[byte_idx] as u64;
1052                            byte_idx += 1;
1053                            nbits += 8;
1054                        }
1055                        let u = ((acc >> (nbits - b)) & ((1u64 << b) - 1)) as f32;
1056                        nbits -= b;
1057                        let so = (r * ng + i / GROUP_SIZE) * 2;
1058                        let sv = f16_to_f32(u16::from_le_bytes([
1059                            bytes[sc_off + so],
1060                            bytes[sc_off + so + 1],
1061                        ]));
1062                        *d = (u - l) * sv;
1063                    }
1064                    if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
1065                        crate::prism::inverse_embedding(model, dst);
1066                    }
1067                    return;
1068                }
1069                let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
1070                let s = row_scale[r];
1071                match dtype {
1072                    TensorDtype::Q8Row => {
1073                        for (d, &b) in dst.iter_mut().zip(q) {
1074                            *d = (b as i8) as f32 * s;
1075                        }
1076                    }
1077                    TensorDtype::Q8_2f => {
1078                        for (i, (d, &b)) in dst.iter_mut().zip(q).enumerate() {
1079                            *d = (b as i8) as f32 * s * col_field[i];
1080                        }
1081                    }
1082                    _ => unreachable!(),
1083                }
1084                if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
1085                    crate::prism::inverse_embedding(model, dst);
1086                }
1087            }
1088        }
1089    }
1090
1091    /// Can this tensor's columns be read cheaply (for sparse down_proj)?
1092    /// True for F32/Q8Row/Q8_2f (per-row scale, direct strided access);
1093    /// false for group-packed q4/vbit (column access would unpack whole
1094    /// groups — sparse execution falls back to f32 for those).
1095    pub fn sparse_col_ok(&self) -> bool {
1096        match self {
1097            Self::F32 { .. } => true,
1098            Self::Blr(_) => false,
1099            Self::Mapped { dtype, .. } => {
1100                matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
1101            }
1102        }
1103    }
1104
1105    /// down_proj [hidden, inter]: accumulate `w · col(c)` into `out`
1106    /// [hidden] — reads ONLY column `c` (one neuron) from the mmap,
1107    /// no full-matrix dequant. `out[k] += w · down[k, c]`.
1108    pub fn add_col_scaled(&self, c: usize, w: f32, out: &mut [f32]) {
1109        let inter = self.cols();
1110        let hidden = self.rows();
1111        debug_assert_eq!(out.len(), hidden);
1112        match self {
1113            Self::F32 { data, .. } => {
1114                for (k, o) in out.iter_mut().enumerate() {
1115                    *o += w * data[k * inter + c];
1116                }
1117            }
1118            Self::Blr(t) => t.add_col_scaled(c, w, out),
1119            Self::Mapped {
1120                dtype,
1121                row_scale,
1122                col_field,
1123                ..
1124            } => {
1125                let q = self.quant_bytes();
1126                let colf = if *dtype == TensorDtype::Q8_2f {
1127                    col_field[c]
1128                } else {
1129                    1.0
1130                };
1131                let wc = w * colf;
1132                for (k, o) in out.iter_mut().enumerate() {
1133                    let b = q[k * inter + c] as i8 as f32;
1134                    *o += wc * b * row_scale[k];
1135                }
1136            }
1137        }
1138    }
1139
1140    /// Touch the head of row `r` so the DRAM latency of the next
1141    /// neuron's weights overlaps the current one's arithmetic.
1142    ///
1143    /// Scattered rows are what per-token sparsity reads, and a 2 KB
1144    /// stride is past what the hardware prefetcher follows: without this
1145    /// every row starts with a cold miss that nothing hides. One touch
1146    /// per 512 bytes is enough — the rest of the row is a sequential run
1147    /// the prefetcher does pick up.
1148    #[inline]
1149    pub fn prefetch_row(&self, r: usize) {
1150        let Self::Mapped { dtype, .. } = self else {
1151            return;
1152        };
1153        if !matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f) {
1154            return;
1155        }
1156        let cols = self.cols();
1157        let q = self.quant_bytes();
1158        let (a, b) = (r * cols, (r + 1) * cols);
1159        if b > q.len() {
1160            return;
1161        }
1162        let mut j = a;
1163        while j < b {
1164            unsafe { std::ptr::read_volatile(q.as_ptr().add(j)) };
1165            j += 512;
1166        }
1167    }
1168
1169    /// `out += w · row(r)` — the transposed twin of `add_col_scaled`.
1170    ///
1171    /// A neuron's `down` weights are a COLUMN of `[hidden, inter]`, and a
1172    /// column is strided: reading one costs a cache line per element, so
1173    /// per-neuron dynamic sparsity saves arithmetic and no bytes. Stored
1174    /// transposed (`down_proj.t.weight`, `[inter, hidden]`) the same
1175    /// weights are a contiguous ROW, and this accumulate reads exactly
1176    /// the neurons the token asked for.
1177    pub fn add_row_scaled(&self, r: usize, w: f32, out: &mut [f32], scratch: &mut [f32]) {
1178        let cols = self.cols();
1179        debug_assert_eq!(out.len(), cols);
1180        match self {
1181            Self::F32 { data, .. } => {
1182                let row = &data[r * cols..(r + 1) * cols];
1183                for (o, v) in out.iter_mut().zip(row) {
1184                    *o += w * v;
1185                }
1186            }
1187            Self::Blr(t) => {
1188                let tmp = &mut scratch[..cols];
1189                t.row_f32(r, tmp);
1190                for (o, v) in out.iter_mut().zip(tmp.iter()) {
1191                    *o += w * v;
1192                }
1193            }
1194            Self::Mapped {
1195                dtype,
1196                row_scale,
1197                col_field,
1198                ..
1199            } => match dtype {
1200                TensorDtype::Q8Row => {
1201                    let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
1202                    let ws = w * row_scale[r];
1203                    let row: &[i8] =
1204                        unsafe { std::slice::from_raw_parts(q.as_ptr() as *const i8, q.len()) };
1205                    axpy_i8_f32(out, row, ws);
1206                }
1207                TensorDtype::Q8_2f => {
1208                    let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
1209                    let ws = w * row_scale[r];
1210                    for ((o, b), c) in out.iter_mut().zip(q).zip(col_field) {
1211                        *o += ws * c * (*b as i8 as f32);
1212                    }
1213                }
1214                _ => {
1215                    self.row_f32(r, scratch);
1216                    for (o, v) in out.iter_mut().zip(scratch.iter()) {
1217                        *o += w * v;
1218                    }
1219                }
1220            },
1221        }
1222    }
1223
1224    /// Dot of row `r` with `x` (gate/up active-neuron path). Reads only
1225    /// row `r` from the mmap — no full dequant. q4/vbit dequant the row
1226    /// into `scratch` first (rare for active-FFN weights).
1227    pub fn row_dot(&self, r: usize, x: &[f32], scratch: &mut [f32]) -> f32 {
1228        let cols = self.cols();
1229        match self {
1230            Self::F32 { data, .. } => {
1231                let row = &data[r * cols..(r + 1) * cols];
1232                row.iter().zip(x).map(|(w, v)| w * v).sum()
1233            }
1234            Self::Blr(t) => {
1235                let tmp = &mut scratch[..cols];
1236                t.row_f32(r, tmp);
1237                tmp.iter().zip(x).map(|(w, v)| w * v).sum()
1238            }
1239            Self::Mapped {
1240                model,
1241                idx,
1242                dtype,
1243                row_scale,
1244                col_field,
1245                ..
1246            } => {
1247                let prism_forward =
1248                    crate::prism::is_forward_weight(model, &model.tensors[*idx].name);
1249                if prism_forward {
1250                    let transformed = crate::prism::forward(model, &x[..cols]);
1251                    let gpr = cols / GROUP_SIZE;
1252                    match dtype {
1253                        TensorDtype::Q2TiledP => {
1254                            let v = Q4tpView::new_q2(self.quant_bytes(), self.rows(), cols);
1255                            let mut sc = vec![0f32; gpr];
1256                            v.scales_into(r, gpr, &mut sc);
1257                            if crate::prism::is_affine_target(model, &model.tensors[*idx].name) {
1258                                return q2tp_affine_row_exact(v.nib, r, gpr, &transformed, &sc);
1259                            }
1260                            return q2tp_row_exact(v.nib, r, gpr, &transformed, &sc);
1261                        }
1262                        _ => {
1263                            self.row_f32(r, scratch);
1264                            return scratch.iter().zip(&transformed).map(|(w, v)| w * v).sum();
1265                        }
1266                    }
1267                }
1268                match dtype {
1269                    TensorDtype::Q8Row => {
1270                        let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
1271                        dot_i8_f32(q, x) * row_scale[r]
1272                    }
1273                    TensorDtype::Q8_2f => {
1274                        let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
1275                        dot_i8_col_f32(q, x, col_field) * row_scale[r]
1276                    }
1277                    _ => {
1278                        self.row_f32(r, scratch);
1279                        scratch.iter().zip(x).map(|(w, v)| w * v).sum()
1280                    }
1281                }
1282            }
1283        }
1284    }
1285
1286    /// `out = W · x` (row-major). F32 delegates to the historical
1287    /// bit-exact path; Mapped runs the fused int8 kernel.
1288    pub fn matvec(&self, x: &[f32], out: &mut [f32], pool: Option<&Pool>) {
1289        match self {
1290            // NOTE: `out.len()` DRIVES this arm — it computes that many rows,
1291            // and `x.len()` is the stride. A short `out` is legitimate here,
1292            // which is why the check below lives in the Mapped arm only.
1293            Self::F32 { data, .. } => {
1294                if !crate::f32_backend::matvec(data, x, out) {
1295                    matvec_rows(pool, data, x, out);
1296                }
1297            }
1298            Self::Blr(t) => {
1299                if crate::gpu::enabled_here()
1300                    && crate::blr::gpu_wanted()
1301                    && crate::gpu::blr_matmat(t, &x[..t.cols], 1, &mut out[..t.rows])
1302                {
1303                    return;
1304                }
1305                t.matvec(x, out, pool)
1306            }
1307            Self::Mapped {
1308                model,
1309                idx,
1310                dtype,
1311                rows,
1312                cols,
1313                row_scale,
1314                col_field,
1315                vbit_offsets,
1316                repack,
1317            } => {
1318                let _ = (model, idx);
1319                // Every kernel below writes `rows` entries through a raw
1320                // pointer, so a short `out` is an out-of-bounds WRITE, not a
1321                // wrong answer: it scribbles on the allocator's metadata and
1322                // the process aborts much later, somewhere innocent
1323                // (`double free or corruption`, `corrupted double-linked
1324                // list`). The debug_assert two of the kernels carried is
1325                // compiled out of the release — exactly the build where it
1326                // matters. Fail here instead, while the caller is still on
1327                // the stack to be named.
1328                assert!(
1329                    out.len() >= *rows && x.len() >= *cols,
1330                    "matvec {rows}x{cols}: out {} (need {rows}), x {} (need {cols})",
1331                    out.len(),
1332                    x.len(),
1333                );
1334                let prism_forward =
1335                    crate::prism::is_forward_weight(model, &model.tensors[*idx].name);
1336                if *dtype == TensorDtype::Q2TiledP
1337                    && std::env::var("CMF_Q2TP_TRACE").as_deref() == Ok("1")
1338                {
1339                    use std::sync::atomic::{AtomicUsize, Ordering};
1340                    static N: AtomicUsize = AtomicUsize::new(0);
1341                    let n = N.fetch_add(1, Ordering::Relaxed);
1342                    if n < 128 {
1343                        eprintln!(
1344                            "q2tp-dispatch #{n} name={} prism={} rows={} cols={} gpu={} optin={} layer={}",
1345                            model.tensors[*idx].name,
1346                            prism_forward,
1347                            rows,
1348                            cols,
1349                            crate::gpu::enabled_here(),
1350                            crate::gpu::q2tp_gpu_opt_in(),
1351                            crate::gpu::cur_layer(),
1352                        );
1353                    }
1354                }
1355                // Prism stores every manifest-listed forward matrix in the
1356                // signed-Hadamard basis.  The q2tp WGSL path receives that
1357                // transformed vector and an explicit affine bit; codecs
1358                // without a descriptor-aware kernel remain on CPU below.
1359                if prism_forward {
1360                    let transformed = crate::prism::forward(model, &x[..*cols]);
1361                    match dtype {
1362                        TensorDtype::Q4Block => {
1363                            q4matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1364                        }
1365                        TensorDtype::Q4Tiled => {
1366                            q4t_matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1367                        }
1368                        TensorDtype::Q4TiledP => {
1369                            q4tp_matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1370                        }
1371                        TensorDtype::Q2TiledP => {
1372                            let affine =
1373                                crate::prism::is_affine_target(model, &model.tensors[*idx].name);
1374                            if *rows * *cols >= 8_388_608
1375                                && crate::gpu::enabled_here()
1376                                && crate::gpu::q2tp_gpu_opt_in()
1377                            {
1378                                let gpu_ok = if affine {
1379                                    crate::gpu::q2tp_affine_matvec(
1380                                        model,
1381                                        *idx,
1382                                        &transformed,
1383                                        *rows,
1384                                        *cols,
1385                                        out,
1386                                    )
1387                                } else {
1388                                    crate::gpu::q2tp_matvec(
1389                                        model,
1390                                        *idx,
1391                                        &transformed,
1392                                        *rows,
1393                                        *cols,
1394                                        out,
1395                                    )
1396                                };
1397                                if gpu_ok {
1398                                    return;
1399                                }
1400                            }
1401                            if affine {
1402                                q2tp_affine_matvec(
1403                                    self.quant_bytes(),
1404                                    &transformed,
1405                                    *rows,
1406                                    *cols,
1407                                    out,
1408                                    pool,
1409                                )
1410                            } else {
1411                                q2tp_matvec(
1412                                    self.quant_bytes(),
1413                                    &transformed,
1414                                    *rows,
1415                                    *cols,
1416                                    out,
1417                                    pool,
1418                                )
1419                            }
1420                        }
1421                        TensorDtype::Q1 | TensorDtype::Q1G => bin_matvec(
1422                            self.bin_w().expect("binary dtype"),
1423                            &transformed,
1424                            *rows,
1425                            *cols,
1426                            out,
1427                            pool,
1428                        ),
1429                        TensorDtype::Q1T => {
1430                            q1t_matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1431                        }
1432                        TensorDtype::Vbit | TensorDtype::VbitRo => vbitmatvec(
1433                            self.quant_bytes(),
1434                            vbit_offsets,
1435                            &transformed,
1436                            *rows,
1437                            *cols,
1438                            out,
1439                            pool,
1440                        ),
1441                        TensorDtype::Q8Row | TensorDtype::Q8_2f => qmatvec(
1442                            self.quant_bytes(),
1443                            repack,
1444                            row_scale,
1445                            &transformed,
1446                            col_field,
1447                            *dtype,
1448                            *rows,
1449                            *cols,
1450                            out,
1451                            pool,
1452                        ),
1453                        _ => unreachable!("unsupported mapped Prism dtype {dtype:?}"),
1454                    }
1455                    return;
1456                }
1457                if *dtype == TensorDtype::Q4Block {
1458                    // GPU route (wgpu q4b kernel) for large q4_block matvecs —
1459                    // gives NVIDIA/AMD/Intel q4 models a GPU path. Probe keeps
1460                    // the winner; Metal returns false → the CPU kernel below.
1461                    if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1462                        let t0 = std::time::Instant::now();
1463                        match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
1464                            crate::gpu::ProbeArm::Gpu => {
1465                                if crate::gpu::q4b_matvec(model, *idx, x, *rows, *cols, out) {
1466                                    crate::gpu::probe_record(
1467                                        crate::gpu::OpClass::Matvec,
1468                                        true,
1469                                        t0.elapsed(),
1470                                    );
1471                                    return;
1472                                }
1473                            }
1474                            crate::gpu::ProbeArm::CpuTimed => {
1475                                q4matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1476                                crate::gpu::probe_record(
1477                                    crate::gpu::OpClass::Matvec,
1478                                    false,
1479                                    t0.elapsed(),
1480                                );
1481                                return;
1482                            }
1483                            crate::gpu::ProbeArm::Cpu => {}
1484                        }
1485                    }
1486                    q4matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1487                    return;
1488                }
1489                if *dtype == TensorDtype::Q4Tiled {
1490                    // GPU route for large q4t matvecs — the lm_head class,
1491                    // same shape as the q4tp arm below. The probe keeps the
1492                    // winner; a backend without the kernel refuses and the
1493                    // CPU path stays.
1494                    if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1495                        let t0 = std::time::Instant::now();
1496                        let cls = crate::gpu::matvec_class(*rows, *cols);
1497                        match crate::gpu::probe_arm(cls) {
1498                            crate::gpu::ProbeArm::Gpu => {
1499                                if crate::gpu::q4t_matvec(model, *idx, x, *rows, *cols, out) {
1500                                    crate::gpu::probe_record(cls, true, t0.elapsed());
1501                                    return;
1502                                }
1503                            }
1504                            crate::gpu::ProbeArm::CpuTimed => {
1505                                q4t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1506                                crate::gpu::probe_record(cls, false, t0.elapsed());
1507                                return;
1508                            }
1509                            crate::gpu::ProbeArm::Cpu => {}
1510                        }
1511                    }
1512                    q4t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1513                    return;
1514                }
1515                if *dtype == TensorDtype::Q4TiledP {
1516                    // GPU route for large q4tp matvecs — the lm_head class.
1517                    // On a q4tp checkpoint the head is the biggest single
1518                    // host matvec left in the decode step, and the batched
1519                    // kernel at b=1 already exists on both backends. Probe
1520                    // keeps the winner, same as q4_block above.
1521                    if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1522                        let t0 = std::time::Instant::now();
1523                        let cls = crate::gpu::matvec_class(*rows, *cols);
1524                        match crate::gpu::probe_arm(cls) {
1525                            crate::gpu::ProbeArm::Gpu => {
1526                                if crate::gpu::q4tp_matvec(model, *idx, x, *rows, *cols, out) {
1527                                    crate::gpu::probe_record(cls, true, t0.elapsed());
1528                                    return;
1529                                }
1530                            }
1531                            crate::gpu::ProbeArm::CpuTimed => {
1532                                q4tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1533                                crate::gpu::probe_record(cls, false, t0.elapsed());
1534                                return;
1535                            }
1536                            crate::gpu::ProbeArm::Cpu => {}
1537                        }
1538                    }
1539                    q4tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1540                    return;
1541                }
1542                if *dtype == TensorDtype::Q2TiledP {
1543                    q2tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1544                    return;
1545                }
1546                if *dtype == TensorDtype::Q1 || *dtype == TensorDtype::Q1G {
1547                    // GPU route for large q1 matvecs (out_proj / lm_head
1548                    // class): the CPU q1 kernel is load-port-bound at
1549                    // ~4 GB/s/core, the GPU one is bandwidth-bound — the
1550                    // probe measures both arms and keeps the winner.
1551                    let w = self.bin_w().expect("binary dtype");
1552                    if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1553                        let t0 = std::time::Instant::now();
1554                        let arm = if crate::gpu::q1_force() {
1555                            crate::gpu::ProbeArm::Gpu
1556                        } else {
1557                            crate::gpu::probe_arm(crate::gpu::OpClass::Matvec)
1558                        };
1559                        match arm {
1560                            crate::gpu::ProbeArm::Gpu => {
1561                                let ok = if *dtype == TensorDtype::Q1G {
1562                                    crate::gpu::q1g_matvec(model, *idx, x, *rows, *cols, out)
1563                                } else {
1564                                    crate::gpu::q1_matvec(model, *idx, x, *rows, *cols, out)
1565                                };
1566                                if ok {
1567                                    crate::gpu::probe_record(
1568                                        crate::gpu::OpClass::Matvec,
1569                                        true,
1570                                        t0.elapsed(),
1571                                    );
1572                                    return;
1573                                }
1574                            }
1575                            crate::gpu::ProbeArm::CpuTimed => {
1576                                bin_matvec(w, x, *rows, *cols, out, pool);
1577                                crate::gpu::probe_record(
1578                                    crate::gpu::OpClass::Matvec,
1579                                    false,
1580                                    t0.elapsed(),
1581                                );
1582                                return;
1583                            }
1584                            crate::gpu::ProbeArm::Cpu => {}
1585                        }
1586                    }
1587                    bin_matvec(w, x, *rows, *cols, out, pool);
1588                    return;
1589                }
1590                if *dtype == TensorDtype::Q1T {
1591                    // GPU route for large q1t matvecs: the ternary BASE dot runs
1592                    // on the GPU (load-port-bound on CPU, like q1), then the
1593                    // sparse overlay is added on the CPU. Probe keeps the winner.
1594                    if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1595                        let t0 = std::time::Instant::now();
1596                        match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
1597                            crate::gpu::ProbeArm::Gpu => {
1598                                if crate::gpu::q1t_matvec(model, *idx, x, *rows, *cols, out) {
1599                                    q1t_add_overlay(self.quant_bytes(), x, *rows, *cols, out, pool);
1600                                    crate::gpu::probe_record(
1601                                        crate::gpu::OpClass::Matvec,
1602                                        true,
1603                                        t0.elapsed(),
1604                                    );
1605                                    return;
1606                                }
1607                            }
1608                            crate::gpu::ProbeArm::CpuTimed => {
1609                                q1t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1610                                crate::gpu::probe_record(
1611                                    crate::gpu::OpClass::Matvec,
1612                                    false,
1613                                    t0.elapsed(),
1614                                );
1615                                return;
1616                            }
1617                            crate::gpu::ProbeArm::Cpu => {}
1618                        }
1619                    }
1620                    q1t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1621                    return;
1622                }
1623                if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1624                    vbitmatvec(self.quant_bytes(), vbit_offsets, x, *rows, *cols, out, pool);
1625                    return;
1626                }
1627                let xs = prescale(x, col_field, *dtype);
1628                // D5: large q8 matrices (lm_head-class) — hybrid
1629                // CPU∥GPU: split the rows, both sides compute
1630                // SIMULTANEOUSLY (same math, shared prescale).
1631                // GPU share: CMF_GPU_SPLIT (0..1, default 0.5).
1632                if *rows >= crate::gpu::min_rows()
1633                    && matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
1634                    && gpu_lmhead_enabled()
1635                    && crate::gpu::enabled_here()
1636                {
1637                    // Runtime probe: alternate the hybrid against the
1638                    // pure-CPU matvec, keep whichever is faster HERE.
1639                    let t0 = std::time::Instant::now();
1640                    match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
1641                        crate::gpu::ProbeArm::Gpu => {}
1642                        crate::gpu::ProbeArm::CpuTimed => {
1643                            qmatvec(
1644                                self.quant_bytes(),
1645                                repack,
1646                                row_scale,
1647                                x,
1648                                col_field,
1649                                *dtype,
1650                                *rows,
1651                                *cols,
1652                                out,
1653                                pool,
1654                            );
1655                            crate::gpu::probe_record(
1656                                crate::gpu::OpClass::Matvec,
1657                                false,
1658                                t0.elapsed(),
1659                            );
1660                            return;
1661                        }
1662                        crate::gpu::ProbeArm::Cpu => {
1663                            qmatvec(
1664                                self.quant_bytes(),
1665                                repack,
1666                                row_scale,
1667                                x,
1668                                col_field,
1669                                *dtype,
1670                                *rows,
1671                                *cols,
1672                                out,
1673                                pool,
1674                            );
1675                            return;
1676                        }
1677                    }
1678                    let frac = gpu_split_frac();
1679                    let cpu_rows = ((*rows as f32) * (1.0 - frac)) as usize;
1680                    let (out_cpu, out_gpu) = out.split_at_mut(cpu_rows);
1681                    let bytes = self.quant_bytes();
1682                    let ok = std::thread::scope(|sc| {
1683                        let g = sc.spawn(|| {
1684                            crate::gpu::q8_matvec_range(
1685                                model,
1686                                *idx,
1687                                cpu_rows,
1688                                &row_scale[cpu_rows..],
1689                                &xs,
1690                                *rows - cpu_rows,
1691                                *cols,
1692                                out_gpu,
1693                            )
1694                        });
1695                        if cpu_rows > 0 {
1696                            // Repack prefix covers the full groups of the
1697                            // CPU half (the split starts at row 0).
1698                            let rep_cpu = if repack.is_empty() {
1699                                &[][..]
1700                            } else {
1701                                &repack[..(cpu_rows / 4) * 4 * *cols]
1702                            };
1703                            qmatvec(
1704                                &bytes[..cpu_rows * *cols],
1705                                rep_cpu,
1706                                &row_scale[..cpu_rows],
1707                                x,
1708                                col_field,
1709                                *dtype,
1710                                cpu_rows,
1711                                *cols,
1712                                out_cpu,
1713                                pool,
1714                            );
1715                        }
1716                        g.join().unwrap_or(false)
1717                    });
1718                    if ok {
1719                        crate::gpu::probe_record(crate::gpu::OpClass::Matvec, true, t0.elapsed());
1720                        return;
1721                    }
1722                    // GPU failed — CPU finishes its half (rows rebased —
1723                    // group offsets don't line up, mmap layout only).
1724                    qmatvec(
1725                        &bytes[cpu_rows * *cols..(*rows) * *cols],
1726                        &[],
1727                        &row_scale[cpu_rows..],
1728                        x,
1729                        col_field,
1730                        *dtype,
1731                        *rows - cpu_rows,
1732                        *cols,
1733                        out_gpu,
1734                        pool,
1735                    );
1736                    return;
1737                }
1738                qmatvec(
1739                    self.quant_bytes(),
1740                    repack,
1741                    row_scale,
1742                    x,
1743                    col_field,
1744                    *dtype,
1745                    *rows,
1746                    *cols,
1747                    out,
1748                    pool,
1749                );
1750            }
1751        }
1752    }
1753
1754    /// Fused two-input matvec (MTP verify pair): weights streamed once.
1755    pub fn matvec2(
1756        &self,
1757        x1: &[f32],
1758        x2: &[f32],
1759        o1: &mut [f32],
1760        o2: &mut [f32],
1761        pool: Option<&Pool>,
1762    ) {
1763        match self {
1764            Self::F32 { data, .. } => matvec_rows2(pool, data, x1, x2, o1, o2),
1765            Self::Blr(_) => {
1766                self.matvec(x1, o1, pool);
1767                self.matvec(x2, o2, pool);
1768            }
1769            Self::Mapped {
1770                model,
1771                idx,
1772                dtype,
1773                rows,
1774                cols,
1775                row_scale,
1776                col_field,
1777                vbit_offsets,
1778                ..
1779            } => {
1780                if crate::prism::is_forward_weight(model, &model.tensors[*idx].name) {
1781                    let tx1 = crate::prism::forward(model, &x1[..*cols]);
1782                    let tx2 = crate::prism::forward(model, &x2[..*cols]);
1783                    match dtype {
1784                        TensorDtype::Q4Block => {
1785                            q4matvec2(self.quant_bytes(), &tx1, &tx2, *rows, *cols, o1, o2, pool)
1786                        }
1787                        TensorDtype::Q4Tiled => {
1788                            q4t_matvec2(self.quant_bytes(), &tx1, &tx2, *rows, *cols, o1, o2, pool)
1789                        }
1790                        TensorDtype::Q4TiledP => {
1791                            q4tp_matvec2(self.quant_bytes(), &tx1, &tx2, *rows, *cols, o1, o2, pool)
1792                        }
1793                        TensorDtype::Q2TiledP => {
1794                            if crate::prism::is_affine_target(model, &model.tensors[*idx].name) {
1795                                q2tp_affine_matvec2(
1796                                    self.quant_bytes(),
1797                                    &tx1,
1798                                    &tx2,
1799                                    *rows,
1800                                    *cols,
1801                                    o1,
1802                                    o2,
1803                                    pool,
1804                                )
1805                            } else {
1806                                q2tp_matvec2(
1807                                    self.quant_bytes(),
1808                                    &tx1,
1809                                    &tx2,
1810                                    *rows,
1811                                    *cols,
1812                                    o1,
1813                                    o2,
1814                                    pool,
1815                                )
1816                            }
1817                        }
1818                        TensorDtype::Q1 | TensorDtype::Q1G => bin_matvec2(
1819                            self.bin_w().expect("binary dtype"),
1820                            &tx1,
1821                            &tx2,
1822                            *rows,
1823                            *cols,
1824                            o1,
1825                            o2,
1826                            pool,
1827                        ),
1828                        TensorDtype::Q1T => {
1829                            q1t_matvec2(self.quant_bytes(), &tx1, &tx2, *rows, *cols, o1, o2, pool)
1830                        }
1831                        TensorDtype::Vbit | TensorDtype::VbitRo => vbitmatvec2(
1832                            self.quant_bytes(),
1833                            vbit_offsets,
1834                            &tx1,
1835                            &tx2,
1836                            *rows,
1837                            *cols,
1838                            o1,
1839                            o2,
1840                            pool,
1841                        ),
1842                        TensorDtype::Q8Row | TensorDtype::Q8_2f => qmatvec2(
1843                            self.quant_bytes(),
1844                            row_scale,
1845                            &tx1,
1846                            &tx2,
1847                            col_field,
1848                            *dtype,
1849                            *rows,
1850                            *cols,
1851                            o1,
1852                            o2,
1853                            pool,
1854                        ),
1855                        _ => unreachable!("unsupported mapped Prism dtype {dtype:?}"),
1856                    }
1857                    return;
1858                }
1859                if *dtype == TensorDtype::Q4Block {
1860                    q4matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1861                    return;
1862                }
1863                if *dtype == TensorDtype::Q4Tiled {
1864                    q4t_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1865                    return;
1866                }
1867                if *dtype == TensorDtype::Q4TiledP {
1868                    q4tp_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1869                    return;
1870                }
1871                if *dtype == TensorDtype::Q2TiledP {
1872                    q2tp_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1873                    return;
1874                }
1875                if let Some(w) = self.bin_w() {
1876                    bin_matvec2(w, x1, x2, *rows, *cols, o1, o2, pool);
1877                    return;
1878                }
1879                if *dtype == TensorDtype::Q1T {
1880                    // Fused ternary pair: one row pass, the register
1881                    // unpack shared across both streams on ARM. (Q1T
1882                    // lacks a row_scale array — scales live inline in
1883                    // the tiles — so it must not fall through to the
1884                    // q8 qmatvec2 below.)
1885                    q1t_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1886                    return;
1887                }
1888                if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1889                    vbitmatvec2(
1890                        self.quant_bytes(),
1891                        vbit_offsets,
1892                        x1,
1893                        x2,
1894                        *rows,
1895                        *cols,
1896                        o1,
1897                        o2,
1898                        pool,
1899                    );
1900                    return;
1901                }
1902                qmatvec2(
1903                    self.quant_bytes(),
1904                    row_scale,
1905                    x1,
1906                    x2,
1907                    col_field,
1908                    *dtype,
1909                    *rows,
1910                    *cols,
1911                    o1,
1912                    o2,
1913                    pool,
1914                );
1915            }
1916        }
1917    }
1918}
1919
1920impl QTensor {
1921    /// Batched matvec (prefill-GEMM): xs — row-major [b, cols],
1922    /// out — row-major [b, rows]. Element-wise semantics are IDENTICAL
1923    /// to b matvec calls (same dot kernels in the same order); the win —
1924    /// the weight row streams from DRAM once per batch, not b times.
1925    /// `(model, index)` when this is a memory-mapped q4tp tensor — the
1926    /// identity a device-resident chain needs to hand `tp_matmat` the
1927    /// weight without going through this struct's own dispatch.
1928    pub fn q4tp_mapped(&self) -> Option<(&std::sync::Arc<CmfModel>, usize)> {
1929        if self.has_prism_contract() {
1930            return None;
1931        }
1932        match self {
1933            Self::Mapped {
1934                model, idx, dtype, ..
1935            } if *dtype == TensorDtype::Q4TiledP => Some((model, *idx)),
1936            _ => None,
1937        }
1938    }
1939
1940    pub fn matmat(&self, xs_all: &[f32], b: usize, out: &mut [f32], pool: Option<&Pool>) {
1941        let cols = self.cols();
1942        let rows = self.rows();
1943        debug_assert_eq!(xs_all.len(), b * cols);
1944        debug_assert_eq!(out.len(), b * rows);
1945        let _prof = crate::cpuprof::time(crate::cpuprof::Slot::Matmat);
1946        // GPTQ calibration: fold this layer's inputs into its Hessian. Only
1947        // Mapped tensors carry a directory name; the check is a relaxed
1948        // atomic load, free when not calibrating.
1949        if crate::gptq_capture::capturing() {
1950            if let Self::Mapped { model, idx, .. } = self {
1951                crate::gptq_capture::accumulate(&model.tensors[*idx].name, xs_all, b, cols);
1952            }
1953            if let Self::Blr(t) = self {
1954                crate::gptq_capture::accumulate(t.name(), xs_all, b, cols);
1955            }
1956        }
1957        match self {
1958            Self::Blr(t) => {
1959                if crate::gpu::enabled_here()
1960                    && crate::blr::gpu_wanted()
1961                    && crate::gpu::blr_matmat(t, &xs_all[..b * cols], b, &mut out[..b * rows])
1962                {
1963                    return;
1964                }
1965                t.matmat(xs_all, b, out, pool)
1966            }
1967            Self::F32 { data, .. } => {
1968                if crate::f32_backend::matmat(data, xs_all, b, rows, cols, out) {
1969                    return;
1970                }
1971                let out_addr = SendMut(out.as_mut_ptr());
1972                // Over the batch, four positions a row at a time: four
1973                // independent chains, each still summing j = 0..cols in the
1974                // scalar order (bit-identical to the one-at-a-time loop and
1975                // to `matvec_rows`). The batch is the wide axis of a prefill;
1976                // the rows may be few (Spark-X2.5's 16-row g_proj spent
1977                // 17 ms a layer here split over its rows).
1978                let run = |start: usize, end: usize| {
1979                    let mut bi = start;
1980                    while bi < end {
1981                        let n = (end - bi).min(4);
1982                        for o in 0..rows {
1983                            let row = &data[o * cols..(o + 1) * cols];
1984                            if n == 4 {
1985                                let x = |k: usize| &xs_all[(bi + k) * cols..(bi + k + 1) * cols];
1986                                let (x0, x1, x2, x3) = (x(0), x(1), x(2), x(3));
1987                                let (mut a0, mut a1, mut a2, mut a3) = (0f32, 0f32, 0f32, 0f32);
1988                                for j in 0..cols {
1989                                    let w = row[j];
1990                                    a0 += w * x0[j];
1991                                    a1 += w * x1[j];
1992                                    a2 += w * x2[j];
1993                                    a3 += w * x3[j];
1994                                }
1995                                for (k, acc) in [a0, a1, a2, a3].into_iter().enumerate() {
1996                                    unsafe { *out_addr.at((bi + k) * rows + o) = acc };
1997                                }
1998                            } else {
1999                                for k in 0..n {
2000                                    let x = &xs_all[(bi + k) * cols..(bi + k + 1) * cols];
2001                                    let mut acc = 0f32;
2002                                    for j in 0..cols {
2003                                        acc += row[j] * x[j];
2004                                    }
2005                                    unsafe { *out_addr.at((bi + k) * rows + o) = acc };
2006                                }
2007                            }
2008                        }
2009                        bi += n;
2010                    }
2011                };
2012                dispatch_rows(pool, b, &run);
2013            }
2014            Self::Mapped {
2015                model,
2016                idx,
2017                dtype,
2018                row_scale,
2019                col_field,
2020                vbit_offsets,
2021                ..
2022            } => {
2023                if crate::prism::is_forward_weight(model, &model.tensors[*idx].name) {
2024                    let mut transformed = Vec::with_capacity(xs_all.len());
2025                    for bi in 0..b {
2026                        transformed.extend_from_slice(&crate::prism::forward(
2027                            model,
2028                            &xs_all[bi * cols..(bi + 1) * cols],
2029                        ));
2030                    }
2031                    match dtype {
2032                        TensorDtype::Q4Block => {
2033                            q4matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
2034                        }
2035                        TensorDtype::Q4Tiled => {
2036                            q4t_matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
2037                        }
2038                        TensorDtype::Q4TiledP => {
2039                            q4tp_matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
2040                        }
2041                        TensorDtype::Q2TiledP => {
2042                            let affine =
2043                                crate::prism::is_affine_target(model, &model.tensors[*idx].name);
2044                            // Affine Prism Q2TP has a descriptor-aware GPU
2045                            // kernel for short/tail batches too.  Unlike the
2046                            // ordinary Q2TP path, don't force b<32 back to a
2047                            // scalar CPU matmat: prefill chunks and the final
2048                            // tail both need to stay on the tested GPU arm.
2049                            let gpu_batch_ok = if affine {
2050                                b >= 2
2051                            } else {
2052                                b >= 32 && b * rows * cols >= 128_000_000
2053                            };
2054                            if gpu_batch_ok
2055                                && cols % 32 == 0
2056                                && crate::gpu::enabled_here()
2057                                && crate::gpu::q2tp_gpu_opt_in()
2058                            {
2059                                let gpu_ok = if affine {
2060                                    crate::gpu::q2tp_affine_matmat(
2061                                        model,
2062                                        *idx,
2063                                        &transformed,
2064                                        b,
2065                                        rows,
2066                                        cols,
2067                                        out,
2068                                    )
2069                                } else {
2070                                    crate::gpu::q2tp_matmat(
2071                                        model,
2072                                        *idx,
2073                                        &transformed,
2074                                        b,
2075                                        rows,
2076                                        cols,
2077                                        out,
2078                                    )
2079                                };
2080                                if gpu_ok {
2081                                    return;
2082                                }
2083                            }
2084                            // A one-token Prism decode is the other short
2085                            // case.  Use the descriptor-aware matvec kernel
2086                            // before falling back to the exact CPU path.
2087                            if affine
2088                                && b == 1
2089                                && cols % 32 == 0
2090                                && crate::gpu::enabled_here()
2091                                && crate::gpu::q2tp_gpu_opt_in()
2092                                && crate::gpu::q2tp_affine_matvec(
2093                                    model,
2094                                    *idx,
2095                                    &transformed[..cols],
2096                                    rows,
2097                                    cols,
2098                                    &mut out[..rows],
2099                                )
2100                            {
2101                                return;
2102                            }
2103                            if affine {
2104                                q2tp_affine_matmat(
2105                                    self.quant_bytes(),
2106                                    &transformed,
2107                                    b,
2108                                    rows,
2109                                    cols,
2110                                    out,
2111                                    pool,
2112                                )
2113                            } else {
2114                                q2tp_matmat(
2115                                    self.quant_bytes(),
2116                                    &transformed,
2117                                    b,
2118                                    rows,
2119                                    cols,
2120                                    out,
2121                                    pool,
2122                                )
2123                            }
2124                        }
2125                        TensorDtype::Q1 | TensorDtype::Q1G => bin_matmat(
2126                            self.bin_w().expect("binary dtype"),
2127                            &transformed,
2128                            b,
2129                            rows,
2130                            cols,
2131                            out,
2132                            pool,
2133                        ),
2134                        TensorDtype::Q1T => {
2135                            q1t_matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
2136                        }
2137                        TensorDtype::Vbit | TensorDtype::VbitRo => vbitmatmat(
2138                            self.quant_bytes(),
2139                            vbit_offsets,
2140                            &transformed,
2141                            b,
2142                            rows,
2143                            cols,
2144                            out,
2145                            pool,
2146                        ),
2147                        TensorDtype::Q8Row | TensorDtype::Q8_2f => {
2148                            let pre: Vec<std::borrow::Cow<'_, [f32]>> = (0..b)
2149                                .map(|bi| {
2150                                    prescale(
2151                                        &transformed[bi * cols..(bi + 1) * cols],
2152                                        col_field,
2153                                        *dtype,
2154                                    )
2155                                })
2156                                .collect();
2157                            qmatmat(self.quant_bytes(), row_scale, &pre, rows, cols, out, pool)
2158                        }
2159                        _ => unreachable!("unsupported mapped Prism dtype {dtype:?}"),
2160                    }
2161                    return;
2162                }
2163                if *dtype == TensorDtype::Q4Block {
2164                    q4matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2165                    return;
2166                }
2167                if *dtype == TensorDtype::Q4TiledP {
2168                    // GPU batched q4tp GEMM (dequant + f32nt mul_mm on the
2169                    // device); the probe keeps whichever beats the CPU arm.
2170                    // Narrow (prompt-encode) and wide (DiT) batches probe
2171                    // as separate classes — the regimes have opposite
2172                    // winners and one shared verdict locked the wrong arm.
2173                    // Kill switch (gpu::mm_kill): one grossly slow GPU op
2174                    // (a fair-condition op is ≤~100 ms even at 1024px)
2175                    // means the device is contended by another process
2176                    // (e.g. a simulator) — verdicts are per-process, so
2177                    // without the bail the whole render crawls behind
2178                    // someone else's queue.
2179                    // Row-exact batches stay on the host: the device GEMM
2180                    // is an f32 dequant-sgemm, not the host matvec's sum.
2181                    if b >= 32
2182                        && b * rows * cols >= 128_000_000
2183                        && cols % 32 == 0
2184                        && !row_exact()
2185                        && !crate::gpu::mm_killed()
2186                        && crate::gpu::enabled_here()
2187                    {
2188                        let class = if b >= 128 {
2189                            crate::gpu::OpClass::MatmatWide
2190                        } else {
2191                            crate::gpu::OpClass::Matmat
2192                        };
2193                        if let Self::Mapped { model, idx, .. } = self {
2194                            // In-process A/B (`CMF_MM_AB=1`). Three
2195                            // wall-clock A/Bs on a shared stand disagreed
2196                            // with each other by 25% on the same change,
2197                            // because the machine drifts between processes
2198                            // and interleaving whole renders does not fix
2199                            // that. Here both arms run back to back on the
2200                            // SAME data inside one call, so whatever the
2201                            // machine is doing, it does to both — and the
2202                            // disagreement between their outputs falls out
2203                            // for free. Doubles the work; a diagnostic,
2204                            // not a mode.
2205                            if crate::mm_ab::on() {
2206                                let mut g = vec![0f32; b * rows];
2207                                let t = std::time::Instant::now();
2208                                let took = crate::gpu::q4tp_matmat(
2209                                    model, *idx, xs_all, b, rows, cols, &mut g,
2210                                );
2211                                let dg = t.elapsed();
2212                                let t = std::time::Instant::now();
2213                                q4tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2214                                let dc = t.elapsed();
2215                                crate::mm_ab::record(b, rows, cols, took, dg, dc, &g, out);
2216                                return;
2217                            }
2218                            let t0 = std::time::Instant::now();
2219                            // A cold call takes the device arm: its sample
2220                            // is discarded either way, and the upload is
2221                            // what the next step needs.
2222                            let resident = crate::gpu::weight_is_resident(model, *idx);
2223                            match crate::gpu::probe_arm_cold_prefers_gpu(class, resident) {
2224                                crate::gpu::ProbeArm::Gpu => {
2225                                    if crate::gpu::q4tp_matmat(
2226                                        model, *idx, xs_all, b, rows, cols, out,
2227                                    ) {
2228                                        let el = t0.elapsed();
2229                                        // Work-proportional budget: ~8× the
2230                                        // fair-device estimate (+20 ms slack).
2231                                        // An absolute cap missed the worst
2232                                        // case — contended ops sit at
2233                                        // 100–240 ms each and still bury a
2234                                        // render whose fair op is 3–9 ms.
2235                                        // Cold ops (first PSO build, buffer
2236                                        // alloc) are exempt: a one-off
2237                                        // ~50 ms compile is not contention.
2238                                        let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
2239                                        let budget = std::time::Duration::from_secs_f64(
2240                                            flops / 1.5e12 * 8.0 + 0.020,
2241                                        );
2242                                        crate::gpu::mm_budget_check(
2243                                            "q4tp matmat",
2244                                            el,
2245                                            budget,
2246                                            crate::gpu::probe_was_cold() || !resident,
2247                                        );
2248                                        crate::gpu::probe_record(class, true, el);
2249                                        return;
2250                                    }
2251                                }
2252                                crate::gpu::ProbeArm::CpuTimed => {
2253                                    q4tp_matmat(
2254                                        self.quant_bytes(),
2255                                        xs_all,
2256                                        b,
2257                                        rows,
2258                                        cols,
2259                                        out,
2260                                        pool,
2261                                    );
2262                                    crate::gpu::probe_record(class, false, t0.elapsed());
2263                                    return;
2264                                }
2265                                crate::gpu::ProbeArm::Cpu => {}
2266                            }
2267                        }
2268                    }
2269                    q4tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2270                    return;
2271                }
2272                if *dtype == TensorDtype::Q2TiledP {
2273                    // Same device arm as q4tp, behind the same probe:
2274                    // the planes differ, the dispatch does not. Without
2275                    // this a q2tp file ran its widest projections on the
2276                    // host while the 4-bit one had the card, which is a
2277                    // codec paying for its size twice.
2278                    if b >= 32
2279                        && b * rows * cols >= 128_000_000
2280                        && cols % 32 == 0
2281                        && !crate::gpu::mm_killed()
2282                        && crate::gpu::enabled_here()
2283                    {
2284                        let class = if b >= 128 {
2285                            crate::gpu::OpClass::MatmatWide
2286                        } else {
2287                            crate::gpu::OpClass::Matmat
2288                        };
2289                        if let Self::Mapped { model, idx, .. } = self {
2290                            let t0 = std::time::Instant::now();
2291                            match crate::gpu::probe_arm(class) {
2292                                crate::gpu::ProbeArm::Gpu => {
2293                                    if crate::gpu::q2tp_matmat(
2294                                        model, *idx, xs_all, b, rows, cols, out,
2295                                    ) {
2296                                        crate::gpu::probe_record(class, true, t0.elapsed());
2297                                        return;
2298                                    }
2299                                }
2300                                crate::gpu::ProbeArm::CpuTimed => {
2301                                    q2tp_matmat(
2302                                        self.quant_bytes(),
2303                                        xs_all,
2304                                        b,
2305                                        rows,
2306                                        cols,
2307                                        out,
2308                                        pool,
2309                                    );
2310                                    crate::gpu::probe_record(class, false, t0.elapsed());
2311                                    return;
2312                                }
2313                                crate::gpu::ProbeArm::Cpu => {}
2314                            }
2315                        }
2316                    }
2317                    // Without a host arm a q2tp tensor falls through to
2318                    // the q8 fallback, which reads it at one BYTE per
2319                    // weight — a 2x overrun that killed pool workers
2320                    // mid-prefill while the dispatcher waited forever.
2321                    q2tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2322                    return;
2323                }
2324                if *dtype == TensorDtype::Q4Tiled {
2325                    // GPU batched q4t GEMM (dequant + f32nt mul_mm on the
2326                    // device); the probe keeps whichever beats the CPU arm.
2327                    // Narrow (prompt-encode) and wide (DiT) batches probe
2328                    // as separate classes — the regimes have opposite
2329                    // winners and one shared verdict locked the wrong arm.
2330                    // Kill switch (gpu::mm_kill): one grossly slow GPU op
2331                    // (a fair-condition op is ≤~100 ms even at 1024px)
2332                    // means the device is contended by another process
2333                    // (e.g. a simulator) — verdicts are per-process, so
2334                    // without the bail the whole render crawls behind
2335                    // someone else's queue.
2336                    if b >= 32
2337                        && b * rows * cols >= 128_000_000
2338                        && cols % 32 == 0
2339                        && !crate::gpu::mm_killed()
2340                        && crate::gpu::enabled_here()
2341                    {
2342                        let class = if b >= 128 {
2343                            crate::gpu::OpClass::MatmatWide
2344                        } else {
2345                            crate::gpu::OpClass::Matmat
2346                        };
2347                        if let Self::Mapped { model, idx, .. } = self {
2348                            let t0 = std::time::Instant::now();
2349                            match crate::gpu::probe_arm(class) {
2350                                crate::gpu::ProbeArm::Gpu => {
2351                                    if crate::gpu::q4t_matmat(
2352                                        model, *idx, xs_all, b, rows, cols, out,
2353                                    ) {
2354                                        let el = t0.elapsed();
2355                                        // Work-proportional budget: ~8× the
2356                                        // fair-device estimate (+20 ms slack).
2357                                        // An absolute cap missed the worst
2358                                        // case — contended ops sit at
2359                                        // 100–240 ms each and still bury a
2360                                        // render whose fair op is 3–9 ms.
2361                                        // Cold ops (first PSO build, buffer
2362                                        // alloc) are exempt: a one-off
2363                                        // ~50 ms compile is not contention.
2364                                        let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
2365                                        let budget = std::time::Duration::from_secs_f64(
2366                                            flops / 1.5e12 * 8.0 + 0.020,
2367                                        );
2368                                        crate::gpu::mm_budget_check(
2369                                            "q4t matmat",
2370                                            el,
2371                                            budget,
2372                                            crate::gpu::probe_was_cold(),
2373                                        );
2374                                        crate::gpu::probe_record(class, true, el);
2375                                        return;
2376                                    }
2377                                }
2378                                crate::gpu::ProbeArm::CpuTimed => {
2379                                    q4t_matmat(
2380                                        self.quant_bytes(),
2381                                        xs_all,
2382                                        b,
2383                                        rows,
2384                                        cols,
2385                                        out,
2386                                        pool,
2387                                    );
2388                                    crate::gpu::probe_record(class, false, t0.elapsed());
2389                                    return;
2390                                }
2391                                crate::gpu::ProbeArm::Cpu => {}
2392                            }
2393                        }
2394                    }
2395                    q4t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2396                    return;
2397                }
2398                if *dtype == TensorDtype::Q1 || *dtype == TensorDtype::Q1G {
2399                    // GPU batched q1 GEMM for wide prefill (q1_mul_mm on the
2400                    // device); the probe keeps whichever beats the CPU matmat.
2401                    // q1g has the twin `q1g_mul_mm` (cols % 128 by format).
2402                    let w = self.bin_w().expect("binary dtype");
2403                    if b >= 32
2404                        && b * rows * cols >= 128_000_000
2405                        && cols % 64 == 0
2406                        && crate::gpu::enabled_here()
2407                    {
2408                        if let Self::Mapped { model, idx, .. } = self {
2409                            let t0 = std::time::Instant::now();
2410                            match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
2411                                crate::gpu::ProbeArm::Gpu => {
2412                                    let ok = if *dtype == TensorDtype::Q1G {
2413                                        crate::gpu::q1g_matmat(
2414                                            model, *idx, xs_all, b, rows, cols, out,
2415                                        )
2416                                    } else {
2417                                        crate::gpu::q1_matmat(
2418                                            model, *idx, xs_all, b, rows, cols, out,
2419                                        )
2420                                    };
2421                                    if ok {
2422                                        crate::gpu::probe_record(
2423                                            crate::gpu::OpClass::Matmat,
2424                                            true,
2425                                            t0.elapsed(),
2426                                        );
2427                                        return;
2428                                    }
2429                                }
2430                                crate::gpu::ProbeArm::CpuTimed => {
2431                                    bin_matmat(w, xs_all, b, rows, cols, out, pool);
2432                                    crate::gpu::probe_record(
2433                                        crate::gpu::OpClass::Matmat,
2434                                        false,
2435                                        t0.elapsed(),
2436                                    );
2437                                    return;
2438                                }
2439                                crate::gpu::ProbeArm::Cpu => {}
2440                            }
2441                        }
2442                    }
2443                    bin_matmat(w, xs_all, b, rows, cols, out, pool);
2444                    return;
2445                }
2446                if *dtype == TensorDtype::Q1T {
2447                    // GPU batched GEMM for wide prefill (base + overlay on the
2448                    // device); probe keeps the winner vs the CPU matmat.
2449                    if b >= 32 && b * rows * cols >= 128_000_000 && crate::gpu::enabled_here() {
2450                        if let Self::Mapped { model, idx, .. } = self {
2451                            let t0 = std::time::Instant::now();
2452                            match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
2453                                crate::gpu::ProbeArm::Gpu => {
2454                                    if crate::gpu::q1t_matmat(
2455                                        model, *idx, xs_all, b, rows, cols, out,
2456                                    ) {
2457                                        crate::gpu::probe_record(
2458                                            crate::gpu::OpClass::Matmat,
2459                                            true,
2460                                            t0.elapsed(),
2461                                        );
2462                                        return;
2463                                    }
2464                                }
2465                                crate::gpu::ProbeArm::CpuTimed => {
2466                                    q1t_matmat(
2467                                        self.quant_bytes(),
2468                                        xs_all,
2469                                        b,
2470                                        rows,
2471                                        cols,
2472                                        out,
2473                                        pool,
2474                                    );
2475                                    crate::gpu::probe_record(
2476                                        crate::gpu::OpClass::Matmat,
2477                                        false,
2478                                        t0.elapsed(),
2479                                    );
2480                                    return;
2481                                }
2482                                crate::gpu::ProbeArm::Cpu => {}
2483                            }
2484                        }
2485                    }
2486                    q1t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2487                    return;
2488                }
2489                if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
2490                    vbitmatmat(
2491                        self.quant_bytes(),
2492                        vbit_offsets,
2493                        xs_all,
2494                        b,
2495                        rows,
2496                        cols,
2497                        out,
2498                        pool,
2499                    );
2500                    return;
2501                }
2502                let pre: Vec<std::borrow::Cow<'_, [f32]>> = (0..b)
2503                    .map(|bi| prescale(&xs_all[bi * cols..(bi + 1) * cols], col_field, *dtype))
2504                    .collect();
2505                // MiMo verification is a 2–4 row decode panel, not a wide
2506                // prompt GEMM. Keep q8 projections on the same device as
2507                // decode; the generic b>=8 gate otherwise silently moves
2508                // every projection back to CPU. The short wgpu matmat uses
2509                // the same 64-lane reduction as its single-token matvec.
2510                if row_exact()
2511                    && (1..=4).contains(&b)
2512                    && matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
2513                    && crate::gpu::enabled_here()
2514                    && crate::gpu::wgpu_active()
2515                {
2516                    let flat: Vec<f32> = pre.iter().flat_map(|v| v.iter().copied()).collect();
2517                    if crate::gpu::q8_matmat(model, *idx, row_scale, &flat, b, rows, cols, out) {
2518                        return;
2519                    }
2520                }
2521                // D5: large prefill-batch GEMMs — on the GPU (threshold by
2522                // work volume: submission carries b×rows×cols MACs).
2523                // Runtime probe: the naive GEMM shader + sync readback
2524                // lose to the CPU GEMM on slow driver stacks — alternate
2525                // both arms and keep the winner.
2526                if b >= 8 && b * rows * cols >= 128_000_000 && crate::gpu::enabled_here() {
2527                    if let Self::Mapped { model, idx, .. } = self {
2528                        let t0 = std::time::Instant::now();
2529                        match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
2530                            crate::gpu::ProbeArm::Gpu
2531                                if crate::gpu::probe_deciding(crate::gpu::OpClass::Matmat)
2532                                    && !crate::gpu::q8_resident_or_upload(model, *idx) =>
2533                            {
2534                                // Cold weights during probing: the upload
2535                                // has started, the count runs on the CPU —
2536                                // the GPU arm samples on the next touch.
2537                                let q = self.quant_bytes();
2538                                qmatmat(q, row_scale, &pre, rows, cols, out, pool);
2539                                return;
2540                            }
2541                            crate::gpu::ProbeArm::Gpu => {
2542                                // The two-field codec's own entry first, as
2543                                // `device_matmat` takes it: the column field
2544                                // folds into the weight plane and the GEMM
2545                                // runs on the matrix units. The plain int8
2546                                // entry below is the scalar f32 GEMM — on a
2547                                // Spark-X2.5 4B prefill it ran q|k|v|o at
2548                                // host speed.
2549                                if *dtype == TensorDtype::Q8_2f
2550                                    && std::env::var("CMF_Q8_2F_DEV").as_deref() != Ok("0")
2551                                    && crate::gpu::q8_matmat_2f(
2552                                        model, *idx, row_scale, col_field, xs_all, b, rows,
2553                                        cols, out,
2554                                    )
2555                                {
2556                                    crate::gpu::probe_record(
2557                                        crate::gpu::OpClass::Matmat,
2558                                        true,
2559                                        t0.elapsed(),
2560                                    );
2561                                    return;
2562                                }
2563                                let flat: Vec<f32> =
2564                                    pre.iter().flat_map(|v| v.iter().copied()).collect();
2565                                if crate::gpu::q8_matmat(
2566                                    model, *idx, row_scale, &flat, b, rows, cols, out,
2567                                ) {
2568                                    crate::gpu::probe_record(
2569                                        crate::gpu::OpClass::Matmat,
2570                                        true,
2571                                        t0.elapsed(),
2572                                    );
2573                                    return;
2574                                }
2575                            }
2576                            crate::gpu::ProbeArm::CpuTimed => {
2577                                let q = self.quant_bytes();
2578                                qmatmat(q, row_scale, &pre, rows, cols, out, pool);
2579                                crate::gpu::probe_record(
2580                                    crate::gpu::OpClass::Matmat,
2581                                    false,
2582                                    t0.elapsed(),
2583                                );
2584                                return;
2585                            }
2586                            crate::gpu::ProbeArm::Cpu => {}
2587                        }
2588                    }
2589                }
2590                let q = self.quant_bytes();
2591                qmatmat(q, row_scale, &pre, rows, cols, out, pool);
2592            }
2593        }
2594    }
2595}
2596
2597impl QTensor {
2598    /// The device GEMM this tensor would take, run once on the caller's
2599    /// data — the startup parity probe's arm, and the one place that knows
2600    /// which entry point each codec has.
2601    ///
2602    /// It exists because the probe used to look for a `q4tp` weight by
2603    /// name AND dtype, and a container packed any other way was declared
2604    /// "host path" for the whole render even though its codec had a device
2605    /// GEMM of its own. A gate that only recognizes one codec is a gate
2606    /// that silently downgrades every other one.
2607    pub fn device_matmat(&self, xs: &[f32], b: usize, out: &mut [f32]) -> bool {
2608        let (rows, cols) = (self.rows(), self.cols());
2609        let Self::Mapped {
2610            model,
2611            idx,
2612            dtype,
2613            row_scale,
2614            col_field,
2615            ..
2616        } = self
2617        else {
2618            return false;
2619        };
2620        if crate::prism::has_contract(model) {
2621            return false;
2622        }
2623        match *dtype {
2624            TensorDtype::Q4TiledP => crate::gpu::q4tp_matmat(model, *idx, xs, b, rows, cols, out),
2625            // The two-field codec folds its column field into the
2626            // activation, which leaves a plain per-row int8 GEMM — the
2627            // same kernel `q8_row` uses, on both backends.
2628            TensorDtype::Q8Row | TensorDtype::Q8_2f => {
2629                // The field belongs to the weight; only a backend that cannot
2630                // apply it there makes a scaled copy of the activation.
2631                if *dtype == TensorDtype::Q8_2f
2632                    && std::env::var("CMF_Q8_2F_DEV").as_deref() != Ok("0")
2633                    && crate::gpu::q8_matmat_2f(
2634                        model, *idx, row_scale, col_field, xs, b, rows, cols, out,
2635                    )
2636                {
2637                    return true;
2638                }
2639                let flat: Vec<f32> = (0..b)
2640                    .flat_map(|bi| {
2641                        prescale(&xs[bi * cols..(bi + 1) * cols], col_field, *dtype).into_owned()
2642                    })
2643                    .collect();
2644                crate::gpu::q8_matmat(model, *idx, row_scale, &flat, b, rows, cols, out)
2645            }
2646            _ => false,
2647        }
2648    }
2649
2650    /// Multi-matrix job (roadmap §3 P0): N tensors sharing one input
2651    /// run under a SINGLE pool dispatch — QKV or gate+up cost one
2652    /// barrier instead of N. Per-row math is the exact same kernel as
2653    /// `matvec` (bit-identical outputs); only the dispatch is fused.
2654    /// Falls back to N sequential matvecs when the set is not a uniform
2655    /// q8-family/F32 group or there is no pool.
2656    pub fn matvec_many<const N: usize>(
2657        ts: [&QTensor; N],
2658        x: &[f32],
2659        mut outs: [&mut [f32]; N],
2660        pool: Option<&Pool>,
2661    ) {
2662        let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
2663        if ts.iter().any(|t| t.has_prism_contract() || t.is_blr()) {
2664            // The fused range kernels have no transform descriptor.  Let
2665            // each tensor's ordinary matvec dispatch perform the explicit
2666            // signed FWHT (and retain CPU fallback for mixed q2tp/q4tp).
2667            // blr has no row-range kernel either: per-tensor calls.
2668            for (t, o) in ts.iter().zip(outs.iter_mut()) {
2669                t.matvec(x, o, pool);
2670            }
2671            return;
2672        }
2673        let uniform_q8 = ts.iter().all(|t| {
2674            matches!(
2675                t,
2676                Self::Mapped {
2677                    dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f,
2678                    ..
2679                }
2680            )
2681        });
2682        let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
2683        if uniform_f32 && crate::f32_backend::active() {
2684            for (t, o) in ts.iter().zip(outs.iter_mut()) {
2685                t.matvec(x, o, pool);
2686            }
2687            return;
2688        }
2689        let uniform_q4 = ts.iter().all(|t| {
2690            matches!(
2691                t,
2692                Self::Mapped {
2693                    dtype: TensorDtype::Q4Block,
2694                    ..
2695                }
2696            )
2697        });
2698        let uniform_vbit = ts.iter().all(|t| {
2699            matches!(
2700                t,
2701                Self::Mapped {
2702                    dtype: TensorDtype::Vbit | TensorDtype::VbitRo,
2703                    ..
2704                }
2705            )
2706        });
2707        // q1 and q1g may mix freely: the shared activation split and group
2708        // sums are layout-independent, only the row kernel differs.
2709        let uniform_q1 =
2710            ts.iter().all(|t| t.bin_w().is_some()) && ts.iter().all(|t| t.cols() == ts[0].cols());
2711        let uniform_q1t = ts.iter().all(|t| {
2712            matches!(
2713                t,
2714                Self::Mapped {
2715                    dtype: TensorDtype::Q1T,
2716                    ..
2717                }
2718            )
2719        });
2720        // q4tp is the skeleton dtype of the big MoE files, and without an arm
2721        // here every projection that shares an input paid its own pool
2722        // barrier: DeepSeek-V4's attention step alone hands this function
2723        // wq_a, wkv and both compressors' pairs off the same hidden state.
2724        let uniform_q4tp = ts.iter().all(|t| {
2725            matches!(
2726                t,
2727                Self::Mapped {
2728                    dtype: TensorDtype::Q4TiledP,
2729                    ..
2730                }
2731            )
2732        }) && ts
2733            .iter()
2734            .all(|t| t.cols() == ts[0].cols() && t.cols() % GROUP_SIZE == 0);
2735        let Some(pool) = pool else {
2736            for (t, o) in ts.iter().zip(outs.iter_mut()) {
2737                t.matvec(x, o, None);
2738            }
2739            return;
2740        };
2741        if total_rows < 256
2742            || !(uniform_q8
2743                || uniform_f32
2744                || uniform_q4
2745                || uniform_vbit
2746                || uniform_q1
2747                || uniform_q1t
2748                || uniform_q4tp)
2749        {
2750            for (t, o) in ts.iter().zip(outs.iter_mut()) {
2751                t.matvec(x, o, Some(pool));
2752            }
2753            return;
2754        }
2755
2756        if uniform_q4tp {
2757            // Every tensor's rows laid end to end in one virtual row space,
2758            // so the whole set is ONE dispatch. The per-row body is the
2759            // `q4tp_matvec` arm verbatim — same activation split, same
2760            // accumulation order — so the outputs are bit-identical to the
2761            // sequential calls this replaces.
2762            let cols = ts[0].cols();
2763            let gpr = cols / GROUP_SIZE;
2764            let views: [Q4tpView; N] =
2765                std::array::from_fn(|i| Q4tpView::new(ts[i].quant_bytes(), ts[i].rows(), cols));
2766            let rows_of: [usize; N] = std::array::from_fn(|i| ts[i].rows());
2767            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2768            // flat index -> (which tensor, which of its rows)
2769            let locate = |flat: usize| -> (usize, usize) {
2770                let mut acc = 0;
2771                for (i, &r) in rows_of.iter().enumerate() {
2772                    if flat < acc + r {
2773                        return (i, flat - acc);
2774                    }
2775                    acc += r;
2776                }
2777                (rows_of.len() - 1, 0)
2778            };
2779            let (views, outs_addr) = (&views, &outs_addr);
2780            if a8w8_enabled() {
2781                let act = split_act(x);
2782                let act = &act;
2783                let run = |start: usize, end: usize| {
2784                    with_krow(gpr, |sc| {
2785                        for flat in start..end {
2786                            let (t, r) = locate(flat);
2787                            let v = &views[t];
2788                            v.scales_into(r, gpr, sc);
2789                            let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, sc) * act.sx;
2790                            for &(j, xv) in &act.outliers {
2791                                let (w, s) = q4tp_outlier(v.nib, r, gpr, j, sc);
2792                                acc += w * s * xv;
2793                            }
2794                            // SAFETY: one worker owns each (tensor, row) pair.
2795                            unsafe { *outs_addr[t].at(r) = acc };
2796                        }
2797                    });
2798                };
2799                pool.run_rows(total_rows, &run);
2800            } else {
2801                let run = |start: usize, end: usize| {
2802                    with_krow(gpr, |sc| {
2803                        for flat in start..end {
2804                            let (t, r) = locate(flat);
2805                            let v = &views[t];
2806                            v.scales_into(r, gpr, sc);
2807                            // SAFETY: one worker owns each (tensor, row) pair.
2808                            unsafe { *outs_addr[t].at(r) = q4tp_row_exact(v.nib, r, gpr, x, sc) };
2809                        }
2810                    });
2811                };
2812                pool.run_rows(total_rows, &run);
2813            }
2814            return;
2815        }
2816
2817        if uniform_q1 {
2818            // One shared activation split + group sums (q1 has no col
2819            // field; the same input feeds every tensor).
2820            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2821            if a8w8_enabled() {
2822                let act = split_act(x);
2823                let gsum = q1_group_sums(&act.xq, ts[0].cols() / GROUP_SIZE);
2824                let (act, gsum) = (&act, &gsum);
2825                let closures: [_; N] = std::array::from_fn(|i| {
2826                    let (w, gpr, out) = (
2827                        ts[i].bin_w().expect("uniform_q1 checked"),
2828                        ts[i].cols() / GROUP_SIZE,
2829                        outs_addr[i],
2830                    );
2831                    move |s: usize, e: usize| bin_range_a8w8(w, gpr, act, gsum, out, s, e)
2832                });
2833                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2834                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2835                pool.run_many(&parts);
2836            } else {
2837                let closures: [_; N] = std::array::from_fn(|i| {
2838                    let (w, gpr, out) = (
2839                        ts[i].bin_w().expect("uniform_q1 checked"),
2840                        ts[i].cols() / GROUP_SIZE,
2841                        outs_addr[i],
2842                    );
2843                    move |s: usize, e: usize| bin_range_f32(w, gpr, x, out, s, e)
2844                });
2845                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2846                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2847                pool.run_many(&parts);
2848            }
2849            return;
2850        }
2851
2852        if uniform_q1t {
2853            // Q1T batched: one shared activation split + overlay decode,
2854            // all tensors' rows in ONE pool dispatch (saves N−1 dispatches
2855            // and N−1 redundant split_act calls per layer).
2856            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2857            const TILE: usize = cortiq_core::quant::Q1T_TILE;
2858            if a8w8_enabled() {
2859                let act = split_act(x);
2860                let act = &act;
2861                let x_ref = x;
2862                let closures: [_; N] = std::array::from_fn(|i| {
2863                    let bytes = ts[i].quant_bytes();
2864                    let (rows, cols) = (ts[i].rows(), ts[i].cols());
2865                    let gpr = cols / GROUP_SIZE;
2866                    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
2867                    let out = outs_addr[i];
2868                    move |s: usize, e: usize| {
2869                        q1t_range_a8w8(bytes, gpr, rp_off, ent_off, has_ov, act, x_ref, out, s, e)
2870                    }
2871                });
2872                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2873                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2874                pool.run_many(&parts);
2875            } else {
2876                let x_ref = x;
2877                let closures: [_; N] = std::array::from_fn(|i| {
2878                    let bytes = ts[i].quant_bytes();
2879                    let (rows, cols) = (ts[i].rows(), ts[i].cols());
2880                    let gpr = cols / GROUP_SIZE;
2881                    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
2882                    let out = outs_addr[i];
2883                    move |s: usize, e: usize| {
2884                        q1t_range_f32_batch(bytes, gpr, rp_off, ent_off, has_ov, x_ref, out, s, e)
2885                    }
2886                });
2887                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2888                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2889                pool.run_many(&parts);
2890            }
2891            return;
2892        }
2893
2894        if uniform_q4 || uniform_vbit {
2895            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2896            // q4/vbit share one activation split — no per-tensor col field.
2897            if a8w8_enabled() {
2898                let act = split_act(x);
2899                let act = &act;
2900                if uniform_q4 {
2901                    let closures: [_; N] = std::array::from_fn(|i| {
2902                        let (packed, scales) =
2903                            q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
2904                        let (gpr, cols, out) =
2905                            (ts[i].cols() / GROUP_SIZE, ts[i].cols(), outs_addr[i]);
2906                        move |s: usize, e: usize| {
2907                            q4_range_a8w8(packed, scales, gpr, cols, act, out, s, e)
2908                        }
2909                    });
2910                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2911                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2912                    pool.run_many(&parts);
2913                } else {
2914                    let closures: [_; N] = std::array::from_fn(|i| {
2915                        let Self::Mapped { vbit_offsets, .. } = ts[i] else {
2916                            unreachable!()
2917                        };
2918                        let (bytes, rows, cols, out) = (
2919                            ts[i].quant_bytes(),
2920                            ts[i].rows(),
2921                            ts[i].cols(),
2922                            outs_addr[i],
2923                        );
2924                        move |s: usize, e: usize| {
2925                            vbit_range_a8w8(bytes, vbit_offsets, x, act, rows, cols, out, s, e)
2926                        }
2927                    });
2928                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2929                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2930                    pool.run_many(&parts);
2931                }
2932                return;
2933            }
2934            if uniform_q4 {
2935                let closures: [_; N] = std::array::from_fn(|i| {
2936                    let (packed, scales) =
2937                        q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
2938                    let (gpr, out) = (ts[i].cols() / GROUP_SIZE, outs_addr[i]);
2939                    move |s: usize, e: usize| q4_range_f32(packed, scales, gpr, x, out, s, e)
2940                });
2941                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2942                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2943                pool.run_many(&parts);
2944            } else {
2945                let closures: [_; N] = std::array::from_fn(|i| {
2946                    let Self::Mapped { vbit_offsets, .. } = ts[i] else {
2947                        unreachable!()
2948                    };
2949                    let (bytes, rows, cols, out) = (
2950                        ts[i].quant_bytes(),
2951                        ts[i].rows(),
2952                        ts[i].cols(),
2953                        outs_addr[i],
2954                    );
2955                    move |s: usize, e: usize| {
2956                        vbit_range_f32(bytes, vbit_offsets, x, rows, cols, out, s, e)
2957                    }
2958                });
2959                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2960                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2961                pool.run_many(&parts);
2962            }
2963            return;
2964        }
2965
2966        if uniform_f32 {
2967            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2968            let closures: [_; N] = std::array::from_fn(|i| {
2969                let Self::F32 { data, cols, .. } = ts[i] else {
2970                    unreachable!()
2971                };
2972                let out = outs_addr[i];
2973                move |start: usize, end: usize| {
2974                    for o in start..end {
2975                        let row = &data[o * cols..(o + 1) * cols];
2976                        let mut sum = 0.0f32;
2977                        for j in 0..*cols {
2978                            sum += row[j] * x[j];
2979                        }
2980                        // SAFETY: disjoint (tensor, row) cells per worker.
2981                        unsafe { *out.at(o) = sum };
2982                    }
2983                }
2984            });
2985            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2986                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2987            pool.run_many(&parts);
2988            return;
2989        }
2990
2991        // Uniform q8-family: per-tensor prescale (q8_2f col fields
2992        // differ per tensor) + the shared range kernels.
2993        struct Ctx<'a> {
2994            bytes: &'a [u8],
2995            #[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
2996            rep: &'a [u8],
2997            row_scale: &'a [f32],
2998            cols: usize,
2999            xs: std::borrow::Cow<'a, [f32]>,
3000        }
3001        let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
3002            let Self::Mapped {
3003                dtype,
3004                cols,
3005                row_scale,
3006                col_field,
3007                repack,
3008                ..
3009            } = ts[i]
3010            else {
3011                unreachable!()
3012            };
3013            Ctx {
3014                bytes: ts[i].quant_bytes(),
3015                rep: repack,
3016                row_scale,
3017                cols: *cols,
3018                xs: prescale(x, col_field, *dtype),
3019            }
3020        });
3021        let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
3022        #[cfg(target_arch = "aarch64")]
3023        if sdot_enabled() {
3024            let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
3025            let closures: [_; N] = std::array::from_fn(|i| {
3026                let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
3027                move |start: usize, end: usize| {
3028                    q8_range_sdot(c.bytes, c.rep, c.row_scale, act, c.cols, out, start, end)
3029                }
3030            });
3031            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3032                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3033            pool.run_many(&parts);
3034            return;
3035        }
3036        #[cfg(target_arch = "x86_64")]
3037        if avx2_a8w8_enabled() {
3038            let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
3039            let closures: [_; N] = std::array::from_fn(|i| {
3040                let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
3041                move |start: usize, end: usize| {
3042                    q8_range_avx2(c.bytes, c.row_scale, act, c.cols, out, start, end)
3043                }
3044            });
3045            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3046                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3047            pool.run_many(&parts);
3048            return;
3049        }
3050        let closures: [_; N] = std::array::from_fn(|i| {
3051            let (c, out) = (&ctxs[i], outs_addr[i]);
3052            move |start: usize, end: usize| {
3053                q8_range_f32(c.bytes, c.row_scale, &c.xs, c.cols, out, start, end)
3054            }
3055        });
3056        let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3057            std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3058        pool.run_many(&parts);
3059    }
3060}
3061
3062impl QTensor {
3063    /// Pair-input multi-matrix job: N tensors × 2 shared inputs under a
3064    /// single pool dispatch — the MTP/pair decode path publishes one job
3065    /// for Q/K/V (and one for gate+up) instead of one per tensor.
3066    /// Per-row math is exactly `matvec2`'s kernels; bit-identical.
3067    #[allow(clippy::needless_range_loop)]
3068    pub fn matvec2_many<const N: usize>(
3069        ts: [&QTensor; N],
3070        x1: &[f32],
3071        x2: &[f32],
3072        mut o1s: [&mut [f32]; N],
3073        mut o2s: [&mut [f32]; N],
3074        pool: Option<&Pool>,
3075    ) {
3076        let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
3077        if ts.iter().any(|t| t.has_prism_contract() || t.is_blr()) {
3078            for i in 0..N {
3079                ts[i].matvec2(x1, x2, o1s[i], o2s[i], pool);
3080            }
3081            return;
3082        }
3083        let uniform_q8 = ts.iter().all(|t| {
3084            matches!(
3085                t,
3086                Self::Mapped {
3087                    dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f,
3088                    ..
3089                }
3090            )
3091        });
3092        let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
3093        let uniform_q4 = ts.iter().all(|t| {
3094            matches!(
3095                t,
3096                Self::Mapped {
3097                    dtype: TensorDtype::Q4Block,
3098                    ..
3099                }
3100            )
3101        });
3102        let uniform_vbit = ts.iter().all(|t| {
3103            matches!(
3104                t,
3105                Self::Mapped {
3106                    dtype: TensorDtype::Vbit | TensorDtype::VbitRo,
3107                    ..
3108                }
3109            )
3110        });
3111        let fusable = pool.is_some()
3112            && total_rows >= 256
3113            && (uniform_q8 || uniform_f32 || uniform_q4 || uniform_vbit);
3114        if !fusable {
3115            for i in 0..N {
3116                ts[i].matvec2(x1, x2, o1s[i], o2s[i], pool);
3117            }
3118            return;
3119        }
3120        let pool = pool.unwrap();
3121
3122        if uniform_q4 || uniform_vbit {
3123            let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
3124            let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
3125            // q4/vbit share activation splits — no per-tensor col field.
3126            if a8w8_enabled() {
3127                let a1 = split_act(x1);
3128                let a2 = split_act(x2);
3129                let (a1, a2) = (&a1, &a2);
3130                if uniform_q4 {
3131                    let closures: [_; N] = std::array::from_fn(|i| {
3132                        let (packed, scales) =
3133                            q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
3134                        let (gpr, cols, o1, o2) =
3135                            (ts[i].cols() / GROUP_SIZE, ts[i].cols(), p1[i], p2[i]);
3136                        move |s: usize, e: usize| {
3137                            q4_range2_a8w8(packed, scales, gpr, cols, a1, a2, o1, o2, s, e)
3138                        }
3139                    });
3140                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3141                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3142                    pool.run_many(&parts);
3143                } else {
3144                    let closures: [_; N] = std::array::from_fn(|i| {
3145                        let Self::Mapped { vbit_offsets, .. } = ts[i] else {
3146                            unreachable!()
3147                        };
3148                        let (bytes, rows, cols, o1, o2) = (
3149                            ts[i].quant_bytes(),
3150                            ts[i].rows(),
3151                            ts[i].cols(),
3152                            p1[i],
3153                            p2[i],
3154                        );
3155                        move |s: usize, e: usize| {
3156                            vbit_range2_a8w8(
3157                                bytes,
3158                                vbit_offsets,
3159                                x1,
3160                                x2,
3161                                a1,
3162                                a2,
3163                                rows,
3164                                cols,
3165                                o1,
3166                                o2,
3167                                s,
3168                                e,
3169                            )
3170                        }
3171                    });
3172                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3173                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3174                    pool.run_many(&parts);
3175                }
3176                return;
3177            }
3178            if uniform_q4 {
3179                let closures: [_; N] = std::array::from_fn(|i| {
3180                    let (packed, scales) =
3181                        q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
3182                    let (gpr, o1, o2) = (ts[i].cols() / GROUP_SIZE, p1[i], p2[i]);
3183                    move |s: usize, e: usize| {
3184                        q4_range2_f32(packed, scales, gpr, x1, x2, o1, o2, s, e)
3185                    }
3186                });
3187                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3188                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3189                pool.run_many(&parts);
3190            } else {
3191                let closures: [_; N] = std::array::from_fn(|i| {
3192                    let Self::Mapped { vbit_offsets, .. } = ts[i] else {
3193                        unreachable!()
3194                    };
3195                    let (bytes, rows, cols, o1, o2) = (
3196                        ts[i].quant_bytes(),
3197                        ts[i].rows(),
3198                        ts[i].cols(),
3199                        p1[i],
3200                        p2[i],
3201                    );
3202                    move |s: usize, e: usize| {
3203                        vbit_range2_f32(bytes, vbit_offsets, x1, x2, rows, cols, o1, o2, s, e)
3204                    }
3205                });
3206                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3207                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3208                pool.run_many(&parts);
3209            }
3210            return;
3211        }
3212
3213        if uniform_f32 {
3214            let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
3215            let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
3216            let closures: [_; N] = std::array::from_fn(|i| {
3217                let Self::F32 { data, cols, .. } = ts[i] else {
3218                    unreachable!()
3219                };
3220                let (o1, o2) = (p1[i], p2[i]);
3221                move |start: usize, end: usize| {
3222                    for o in start..end {
3223                        let row = &data[o * cols..(o + 1) * cols];
3224                        let (mut s1, mut s2) = (0.0f32, 0.0f32);
3225                        for j in 0..*cols {
3226                            s1 += row[j] * x1[j];
3227                            s2 += row[j] * x2[j];
3228                        }
3229                        // SAFETY: disjoint (tensor, row) cells per worker.
3230                        unsafe {
3231                            *o1.at(o) = s1;
3232                            *o2.at(o) = s2;
3233                        }
3234                    }
3235                }
3236            });
3237            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3238                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3239            pool.run_many(&parts);
3240            return;
3241        }
3242
3243        struct Ctx<'a> {
3244            bytes: &'a [u8],
3245            row_scale: &'a [f32],
3246            cols: usize,
3247            xs1: std::borrow::Cow<'a, [f32]>,
3248            xs2: std::borrow::Cow<'a, [f32]>,
3249        }
3250        let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
3251            let Self::Mapped {
3252                dtype,
3253                cols,
3254                row_scale,
3255                col_field,
3256                ..
3257            } = ts[i]
3258            else {
3259                unreachable!()
3260            };
3261            Ctx {
3262                bytes: ts[i].quant_bytes(),
3263                row_scale,
3264                cols: *cols,
3265                xs1: prescale(x1, col_field, *dtype),
3266                xs2: prescale(x2, col_field, *dtype),
3267            }
3268        });
3269        let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
3270        let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
3271        #[cfg(target_arch = "aarch64")]
3272        if sdot_enabled() {
3273            let acts: [(SplitAct, SplitAct); N] =
3274                std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
3275            let closures: [_; N] = std::array::from_fn(|i| {
3276                let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
3277                move |start: usize, end: usize| {
3278                    q8_range2_sdot(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
3279                }
3280            });
3281            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3282                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3283            pool.run_many(&parts);
3284            return;
3285        }
3286        #[cfg(target_arch = "x86_64")]
3287        if avx2_a8w8_enabled() {
3288            let acts: [(SplitAct, SplitAct); N] =
3289                std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
3290            let closures: [_; N] = std::array::from_fn(|i| {
3291                let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
3292                move |start: usize, end: usize| {
3293                    q8_range2_avx2(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
3294                }
3295            });
3296            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3297                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3298            pool.run_many(&parts);
3299            return;
3300        }
3301        let closures: [_; N] = std::array::from_fn(|i| {
3302            let (c, o1, o2) = (&ctxs[i], p1[i], p2[i]);
3303            move |start: usize, end: usize| {
3304                q8_range2_f32(
3305                    c.bytes,
3306                    c.row_scale,
3307                    &c.xs1,
3308                    &c.xs2,
3309                    c.cols,
3310                    o1,
3311                    o2,
3312                    start,
3313                    end,
3314                )
3315            }
3316        });
3317        let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3318            std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3319        pool.run_many(&parts);
3320    }
3321
3322    /// Fused gate+up matvec with SiLU·mul: for each row r, computes
3323    /// `silu(gate·x) * (up·x)` and writes to `out[r]`. ONE pool dispatch,
3324    /// no intermediate g/u buffers, no separate silu pass. Falls back
3325    /// (returns false) for unsupported dtype combos.
3326    pub fn matvec_silu_mul(
3327        gate: &QTensor,
3328        up: &QTensor,
3329        x: &[f32],
3330        out: &mut [f32],
3331        pool: Option<&Pool>,
3332    ) -> bool {
3333        Self::matvec_silu_mul_limited(gate, up, x, out, 0.0, pool)
3334    }
3335
3336    /// Fused gate+up+SiLU with the GLM asymmetrical clamp.  `limit == 0`
3337    /// preserves the historical unclamped helper; a positive limit clamps
3338    /// `up` to both sides and `gate` only from above, matching the GLM
3339    /// SwiGLU reference.  Keeping the limit in the row kernel avoids the two
3340    /// intermediate vectors and the extra combine pass on the Q2TP experts.
3341    pub fn matvec_silu_mul_limited(
3342        gate: &QTensor,
3343        up: &QTensor,
3344        x: &[f32],
3345        out: &mut [f32],
3346        limit: f32,
3347        pool: Option<&Pool>,
3348    ) -> bool {
3349        if gate.has_prism_contract() || up.has_prism_contract() || gate.is_blr() || up.is_blr() {
3350            // The fused gate/up kernels consume x directly.  Prism requires
3351            // a per-matrix signed FWHT, so the caller must use two ordinary
3352            // descriptor-aware matvecs instead of an unrotated fast path.
3353            return false;
3354        }
3355        let inter = gate.rows();
3356        debug_assert_eq!(up.rows(), inter);
3357        debug_assert_eq!(out.len(), inter);
3358        debug_assert_eq!(gate.cols(), up.cols());
3359        if !a8w8_enabled() {
3360            return false;
3361        }
3362        let act = split_act(x);
3363        let act = &act;
3364        let x_ref = x;
3365        let out_addr = SendMut(out.as_mut_ptr());
3366
3367        match (gate, up) {
3368            // Q4Block gate + Q4Block up (most common mobile q4 models)
3369            (
3370                Self::Mapped {
3371                    dtype: TensorDtype::Q4Block,
3372                    ..
3373                },
3374                Self::Mapped {
3375                    dtype: TensorDtype::Q4Block,
3376                    ..
3377                },
3378            ) => {
3379                let (gp, gs) = q4_split(gate.quant_bytes(), gate.rows(), gate.cols());
3380                let (up_p, up_s) = q4_split(up.quant_bytes(), up.rows(), up.cols());
3381                let gpr = gate.cols() / GROUP_SIZE;
3382                let cols = gate.cols();
3383                let run = move |start: usize, end: usize| {
3384                    for r in start..end {
3385                        let mut gv = dot_q4_row_i8(gp, gs, r * gpr, gpr, &act.xq) * act.sx;
3386                        let mut uv = dot_q4_row_i8(up_p, up_s, r * gpr, gpr, &act.xq) * act.sx;
3387                        for &(j, xv) in &act.outliers {
3388                            let flat = r * cols + j;
3389                            let gb = gp[flat / 2];
3390                            let gn = if flat & 1 == 0 { gb & 0x0F } else { gb >> 4 };
3391                            let gsc = f16_to_f32(u16::from_le_bytes([
3392                                gs[(flat / GROUP_SIZE) * 2],
3393                                gs[(flat / GROUP_SIZE) * 2 + 1],
3394                            ]));
3395                            gv += ((gn as i32 - 8) as f32) * gsc * xv;
3396                            let ub = up_p[flat / 2];
3397                            let un = if flat & 1 == 0 { ub & 0x0F } else { ub >> 4 };
3398                            let usc = f16_to_f32(u16::from_le_bytes([
3399                                up_s[(flat / GROUP_SIZE) * 2],
3400                                up_s[(flat / GROUP_SIZE) * 2 + 1],
3401                            ]));
3402                            uv += ((un as i32 - 8) as f32) * usc * xv;
3403                        }
3404                        // SAFETY: disjoint row ranges per worker.
3405                        unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3406                    }
3407                };
3408                dispatch_rows(pool, inter, &run);
3409                true
3410            }
3411            // Q4Tiled gate + Q4Tiled up — one row pass, both tile
3412            // streams sequential, silu·mul fused (same per-row math as
3413            // `q4t_matvec`).
3414            (
3415                Self::Mapped {
3416                    dtype: TensorDtype::Q4Tiled,
3417                    ..
3418                },
3419                Self::Mapped {
3420                    dtype: TensorDtype::Q4Tiled,
3421                    ..
3422                },
3423            ) => {
3424                let g_bytes = gate.quant_bytes();
3425                let u_bytes = up.quant_bytes();
3426                let gpr = gate.cols() / GROUP_SIZE;
3427                let run = move |start: usize, end: usize| {
3428                    for r in start..end {
3429                        let mut gv = dot_q4t_row_i8(g_bytes, r, gpr, &act.xq) * act.sx;
3430                        let mut uv = dot_q4t_row_i8(u_bytes, r, gpr, &act.xq) * act.sx;
3431                        for &(j, xv) in &act.outliers {
3432                            let (w, s) = q4t_outlier(g_bytes, r, gpr, j);
3433                            gv += w * s * xv;
3434                            let (w, s) = q4t_outlier(u_bytes, r, gpr, j);
3435                            uv += w * s * xv;
3436                        }
3437                        // SAFETY: disjoint row ranges per worker.
3438                        unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3439                    }
3440                };
3441                dispatch_rows(pool, inter, &run);
3442                true
3443            }
3444            // Q4TiledP gate + Q4TiledP up — the same fused row pass, with
3445            // each row's two ladders built once and spent on both streams.
3446            (
3447                Self::Mapped {
3448                    dtype: TensorDtype::Q4TiledP,
3449                    ..
3450                },
3451                Self::Mapped {
3452                    dtype: TensorDtype::Q4TiledP,
3453                    ..
3454                },
3455            ) => {
3456                let cols = gate.cols();
3457                let gpr = cols / GROUP_SIZE;
3458                let gv_view = Q4tpView::new(gate.quant_bytes(), inter, cols);
3459                let uv_view = Q4tpView::new(up.quant_bytes(), inter, cols);
3460                let run = |start: usize, end: usize| {
3461                    with_krows(gpr, |gsc, usc| {
3462                        for r in start..end {
3463                            gv_view.scales_into(r, gpr, gsc);
3464                            uv_view.scales_into(r, gpr, usc);
3465                            let mut gv =
3466                                dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsc) * act.sx;
3467                            let mut uv =
3468                                dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, usc) * act.sx;
3469                            for &(j, xv) in &act.outliers {
3470                                let (w, s) = q4tp_outlier(gv_view.nib, r, gpr, j, gsc);
3471                                gv += w * s * xv;
3472                                let (w, s) = q4tp_outlier(uv_view.nib, r, gpr, j, usc);
3473                                uv += w * s * xv;
3474                            }
3475                            // SAFETY: disjoint row ranges per worker.
3476                            unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3477                        }
3478                    });
3479                };
3480                dispatch_rows(pool, inter, &run);
3481                true
3482            }
3483            // Q1 gate + Q1 up — one row pass over both sign streams,
3484            // silu·mul fused (the per-row math of `q1_range_a8w8`); the
3485            // activation group sums are shared by both streams. Without
3486            // this arm a q1 dense FFN paid two dispatches + a combine
3487            // loop — the exact barrier this function exists to remove.
3488            // q1g (and a q1/q1g mix) rides the same arm through `BinW`.
3489            (
3490                Self::Mapped {
3491                    dtype: TensorDtype::Q1 | TensorDtype::Q1G,
3492                    ..
3493                },
3494                Self::Mapped {
3495                    dtype: TensorDtype::Q1 | TensorDtype::Q1G,
3496                    ..
3497                },
3498            ) => {
3499                let gw = gate.bin_w().expect("binary dtype matched");
3500                let uw = up.bin_w().expect("binary dtype matched");
3501                let gpr = gate.cols() / GROUP_SIZE;
3502                let gsum = q1_group_sums(&act.xq, gpr);
3503                let gsum = &gsum;
3504                let run = move |start: usize, end: usize| {
3505                    for r in start..end {
3506                        let mut gv = gw.dot_i8(r, gpr, &act.xq, gsum) * act.sx;
3507                        let mut uv = uw.dot_i8(r, gpr, &act.xq, gsum) * act.sx;
3508                        for &(j, xv) in &act.outliers {
3509                            let (w, s) = gw.outlier(r, gpr, j);
3510                            gv += w * s * xv;
3511                            let (w, s) = uw.outlier(r, gpr, j);
3512                            uv += w * s * xv;
3513                        }
3514                        // SAFETY: disjoint row ranges per worker.
3515                        unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3516                    }
3517                };
3518                dispatch_rows(pool, inter, &run);
3519                true
3520            }
3521            // Q2TiledP gate + Q2TiledP up — the 2-bit expert pair (MoE
3522            // FFNs of the W2 class): one row pass, both ladders built
3523            // once, integer code dots with shared group sums.
3524            (
3525                Self::Mapped {
3526                    dtype: TensorDtype::Q2TiledP,
3527                    ..
3528                },
3529                Self::Mapped {
3530                    dtype: TensorDtype::Q2TiledP,
3531                    ..
3532                },
3533            ) => {
3534                let cols = gate.cols();
3535                let gpr = cols / GROUP_SIZE;
3536                let gv_view = Q4tpView::new_q2(gate.quant_bytes(), inter, cols);
3537                let uv_view = Q4tpView::new_q2(up.quant_bytes(), inter, cols);
3538                let gsum = q1_group_sums(&act.xq, gpr);
3539                let gsum = &gsum;
3540                let run = move |start: usize, end: usize| {
3541                    with_krows(gpr, |gsc, usc| {
3542                        for r in start..end {
3543                            gv_view.scales_into(r, gpr, gsc);
3544                            uv_view.scales_into(r, gpr, usc);
3545                            let mut gv =
3546                                dot_q2tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsum, gsc) * act.sx;
3547                            let mut uv =
3548                                dot_q2tp_row_i8(uv_view.nib, r, gpr, &act.xq, gsum, usc) * act.sx;
3549                            for &(j, xv) in &act.outliers {
3550                                let (w, s) = q2tp_outlier(gv_view.nib, r, gpr, j, gsc);
3551                                gv += w * s * xv;
3552                                let (w, s) = q2tp_outlier(uv_view.nib, r, gpr, j, usc);
3553                                uv += w * s * xv;
3554                            }
3555                            // SAFETY: disjoint row ranges per worker.
3556                            unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3557                        }
3558                    });
3559                };
3560                dispatch_rows(pool, inter, &run);
3561                true
3562            }
3563            // Q8Row gate + Q8Row up — one row pass over both i8 streams.
3564            // Q8_2f stays out on purpose: its column field prescales the
3565            // activations PER TENSOR, which breaks this fn's shared
3566            // split_act contract — it keeps the two-dispatch path.
3567            (
3568                Self::Mapped {
3569                    dtype: TensorDtype::Q8Row,
3570                    row_scale: g_rs,
3571                    ..
3572                },
3573                Self::Mapped {
3574                    dtype: TensorDtype::Q8Row,
3575                    row_scale: u_rs,
3576                    ..
3577                },
3578            ) => {
3579                let g_bytes = gate.quant_bytes();
3580                let u_bytes = up.quant_bytes();
3581                let cols = gate.cols();
3582                let run = move |start: usize, end: usize| {
3583                    for r in start..end {
3584                        let gv = q8_row_dot(&g_bytes[r * cols..(r + 1) * cols], act) * g_rs[r];
3585                        let uv = q8_row_dot(&u_bytes[r * cols..(r + 1) * cols], act) * u_rs[r];
3586                        // SAFETY: disjoint row ranges per worker.
3587                        unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3588                    }
3589                };
3590                dispatch_rows(pool, inter, &run);
3591                true
3592            }
3593            // Q1T gate + Q1T up
3594            (
3595                Self::Mapped {
3596                    dtype: TensorDtype::Q1T,
3597                    ..
3598                },
3599                Self::Mapped {
3600                    dtype: TensorDtype::Q1T,
3601                    ..
3602                },
3603            ) => {
3604                const TILE: usize = cortiq_core::quant::Q1T_TILE;
3605                let g_bytes = gate.quant_bytes();
3606                let u_bytes = up.quant_bytes();
3607                let gpr = gate.cols() / GROUP_SIZE;
3608                let (g_rp, g_ent, g_ov) = q1t_overlay(g_bytes, inter * gpr * TILE, inter);
3609                let (u_rp, u_ent, u_ov) = q1t_overlay(u_bytes, inter * gpr * TILE, inter);
3610                let run = move |start: usize, end: usize| {
3611                    for r in start..end {
3612                        let mut gv = q1t_dot_row_i8(g_bytes, r, gpr, &act.xq) * act.sx;
3613                        let mut uv = q1t_dot_row_i8(u_bytes, r, gpr, &act.xq) * act.sx;
3614                        for &(j, xv) in &act.outliers {
3615                            gv += q1t_base_weight(g_bytes, r, gpr, j) * xv;
3616                            uv += q1t_base_weight(u_bytes, r, gpr, j) * xv;
3617                        }
3618                        gv += q1t_row_outlier_correction(g_bytes, r, g_rp, g_ent, g_ov, x_ref);
3619                        uv += q1t_row_outlier_correction(u_bytes, r, u_rp, u_ent, u_ov, x_ref);
3620                        // SAFETY: disjoint row ranges per worker.
3621                        unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3622                    }
3623                };
3624                dispatch_rows(pool, inter, &run);
3625                true
3626            }
3627            _ => false,
3628        }
3629    }
3630
3631    /// Every routed expert's fused gate/up/SiLU under ONE pool dispatch.
3632    ///
3633    /// The per-expert path pays a pool barrier per expert per stage: at 9
3634    /// experts over 40 layers that is ~720 barriers a token, and a decode
3635    /// profile of Qwen3.6-35B-A3B showed the pool parked in
3636    /// `psynch_cvwait` about twice as long as it spent computing. Laying
3637    /// every expert's rows end-to-end in one virtual row space collapses
3638    /// the stage to a single dispatch. The per-row body is the
3639    /// single-expert q4tp arm verbatim, so outputs are bit-identical.
3640    ///
3641    /// `false` = something is outside the fused q4tp kernel (dtype, shape,
3642    /// or a transformed tensor); the caller walks the ordinary per-expert
3643    /// path. Float activations use the same exact scalar rows, still fused
3644    /// under one pool dispatch.
3645    pub fn moe_gate_up_many(
3646        pairs: &[(&QTensor, &QTensor)],
3647        x: &[f32],
3648        outs: &mut [Vec<f32>],
3649        pool: Option<&Pool>,
3650    ) -> bool {
3651        Self::moe_gate_up_many_limited(pairs, x, outs, 0.0, pool)
3652    }
3653
3654    /// Batched gate/up/SiLU with the optional GLM clamp.  The public legacy
3655    /// helper above keeps its historical unclamped semantics; callers that
3656    /// implement a reference with a positive SwiGLU limit use this variant.
3657    pub fn moe_gate_up_many_limited(
3658        pairs: &[(&QTensor, &QTensor)],
3659        x: &[f32],
3660        outs: &mut [Vec<f32>],
3661        limit: f32,
3662        pool: Option<&Pool>,
3663    ) -> bool {
3664        if pairs.is_empty() || pairs.len() != outs.len() {
3665            return false;
3666        }
3667        if pairs.iter().any(|(g, u)| g.is_blr() || u.is_blr()) {
3668            return false;
3669        }
3670        if !a8w8_enabled() {
3671            if limit > 0.0 {
3672                // The exact-row fallback applies no SwiGLU clamp; the
3673                // caller's per-expert path carries it instead.
3674                return false;
3675            }
3676            let groups = vec![vec![0]; pairs.len()];
3677            return Self::moe_gate_up_rows(pairs, &groups, x, outs, pool);
3678        }
3679        let inter = pairs[0].0.rows();
3680        let cols = pairs[0].0.cols();
3681        if cols % GROUP_SIZE != 0 {
3682            return false;
3683        }
3684        let gpr = cols / GROUP_SIZE;
3685        // Uniform layout across every routed pair: q4tp, or the 2-bit
3686        // profile's q2tp gate/up (the W2 class). Mixed sets refuse.
3687        let q2 = matches!(
3688            pairs[0].0,
3689            Self::Mapped {
3690                dtype: TensorDtype::Q2TiledP,
3691                ..
3692            }
3693        );
3694        let want = if q2 {
3695            TensorDtype::Q2TiledP
3696        } else {
3697            TensorDtype::Q4TiledP
3698        };
3699        let mut views = Vec::with_capacity(pairs.len() * 2);
3700        for ((g, u), o) in pairs.iter().zip(outs.iter()) {
3701            let both = matches!(g, Self::Mapped { dtype, .. } if *dtype == want)
3702                && matches!(u, Self::Mapped { dtype, .. } if *dtype == want);
3703            if !both
3704                || g.rows() != inter
3705                || u.rows() != inter
3706                || g.cols() != cols
3707                || u.cols() != cols
3708                || o.len() != inter
3709            {
3710                return false;
3711            }
3712            let mk = if q2 { Q4tpView::new_q2 } else { Q4tpView::new };
3713            views.push(mk(g.quant_bytes(), inter, cols));
3714            views.push(mk(u.quant_bytes(), inter, cols));
3715        }
3716        let act = split_act(x);
3717        let gsum = if q2 {
3718            q1_group_sums(&act.xq, gpr)
3719        } else {
3720            Vec::new()
3721        };
3722        let (act, gsum) = (&act, &gsum);
3723        let ptrs: Vec<SendMut> = outs.iter_mut().map(|o| SendMut(o.as_mut_ptr())).collect();
3724        let (views, ptrs) = (&views, &ptrs);
3725        let run = |start: usize, end: usize| {
3726            with_krows(gpr, |gsc, usc| {
3727                for flat in start..end {
3728                    let (e, r) = (flat / inter, flat % inter);
3729                    let gv_view = &views[e * 2];
3730                    let uv_view = &views[e * 2 + 1];
3731                    gv_view.scales_into(r, gpr, gsc);
3732                    uv_view.scales_into(r, gpr, usc);
3733                    let (mut gv, mut uv) = if q2 {
3734                        (
3735                            dot_q2tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsum, gsc) * act.sx,
3736                            dot_q2tp_row_i8(uv_view.nib, r, gpr, &act.xq, gsum, usc) * act.sx,
3737                        )
3738                    } else {
3739                        (
3740                            dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsc) * act.sx,
3741                            dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, usc) * act.sx,
3742                        )
3743                    };
3744                    for &(j, xv) in &act.outliers {
3745                        let (og, ou) = if q2 {
3746                            (
3747                                q2tp_outlier(gv_view.nib, r, gpr, j, gsc),
3748                                q2tp_outlier(uv_view.nib, r, gpr, j, usc),
3749                            )
3750                        } else {
3751                            (
3752                                q4tp_outlier(gv_view.nib, r, gpr, j, gsc),
3753                                q4tp_outlier(uv_view.nib, r, gpr, j, usc),
3754                            )
3755                        };
3756                        gv += og.0 * og.1 * xv;
3757                        uv += ou.0 * ou.1 * xv;
3758                    }
3759                    // SAFETY: one worker owns each (expert, row) pair.
3760                    unsafe { *ptrs[e].at(r) = silu_mul_limited(gv, uv, limit) };
3761                }
3762            });
3763        };
3764        dispatch_rows(pool, pairs.len() * inter, &run);
3765        true
3766    }
3767
3768    /// Every routed expert's down projection, weighted and summed into
3769    /// `out`, under ONE pool dispatch.
3770    ///
3771    /// Partitioned by OUTPUT row rather than by expert: each row is owned
3772    /// by a single worker, so the experts are summed in the caller's order
3773    /// — the same sequence of f32 adds the serial `out[i] += w·eo[i]` loop
3774    /// performs, hence bit-identical. Partitioning by expert instead would
3775    /// race on the shared accumulator.
3776    pub fn moe_down_many(
3777        downs: &[&QTensor],
3778        gs: &[Vec<f32>],
3779        weights: &[f32],
3780        out: &mut [f32],
3781        pool: Option<&Pool>,
3782    ) -> bool {
3783        if downs.is_empty() || downs.len() != gs.len() || downs.len() != weights.len() {
3784            return false;
3785        }
3786        if downs.iter().any(|d| d.is_blr()) {
3787            return false;
3788        }
3789        if !a8w8_enabled() {
3790            let mut terms = vec![vec![0.0; out.len()]; downs.len()];
3791            if !Self::moe_down_rows(downs, &vec![1; downs.len()], gs, &mut terms, pool) {
3792                return false;
3793            }
3794            out.fill(0.0);
3795            for (row, &w) in terms.iter().zip(weights) {
3796                for (o, &v) in out.iter_mut().zip(row) {
3797                    *o += w * v;
3798                }
3799            }
3800            return true;
3801        }
3802        let rows = out.len();
3803        let cols = downs[0].cols();
3804        if cols % GROUP_SIZE != 0 {
3805            return false;
3806        }
3807        let gpr = cols / GROUP_SIZE;
3808        let mut views = Vec::with_capacity(downs.len());
3809        for (d, g) in downs.iter().zip(gs.iter()) {
3810            if !matches!(
3811                d,
3812                Self::Mapped {
3813                    dtype: TensorDtype::Q4TiledP,
3814                    ..
3815                }
3816            ) || d.rows() != rows
3817                || d.cols() != cols
3818                || g.len() != cols
3819            {
3820                return false;
3821            }
3822            views.push(Q4tpView::new(d.quant_bytes(), rows, cols));
3823        }
3824        // One int8 split per expert — the activation vectors differ.
3825        let acts: Vec<SplitAct> = gs.iter().map(|g| split_act(g)).collect();
3826        // Partitioned by OUTPUT row, with the experts folded inside: each
3827        // row is owned by one worker, so they are summed in the caller's
3828        // order — the same f32 sequence the serial `out[i] += w·eo[i]`
3829        // loop produces. Partitioning by expert instead would either race
3830        // on the accumulator or need a scratch plane and a second pass;
3831        // measured, that variant was a wash, so this keeps the simpler
3832        // shape.
3833        let out_addr = SendMut(out.as_mut_ptr());
3834        let (views, acts, weights) = (&views, &acts, &weights);
3835        let run = |start: usize, end: usize| {
3836            with_krow(gpr, |sc| {
3837                for r in start..end {
3838                    let mut acc = 0f32;
3839                    for (e, v) in views.iter().enumerate() {
3840                        v.scales_into(r, gpr, sc);
3841                        let a = &acts[e];
3842                        let mut d = dot_q4tp_row_i8(v.nib, r, gpr, &a.xq, sc) * a.sx;
3843                        for &(j, xv) in &a.outliers {
3844                            let (w, s) = q4tp_outlier(v.nib, r, gpr, j, sc);
3845                            d += w * s * xv;
3846                        }
3847                        acc += weights[e] * d;
3848                    }
3849                    // SAFETY: disjoint row ranges per worker.
3850                    unsafe { *out_addr.at(r) = acc };
3851                }
3852            });
3853        };
3854        dispatch_rows(pool, rows, &run);
3855        true
3856    }
3857
3858    /// `moe_gate_up_many` for SEVERAL tokens at once, decode-exact: expert
3859    /// `e` (`pairs[e]`, q4tp) serves the tokens `groups[e]` (row indices
3860    /// into `xs`, each `cols` wide). Every (expert, token) output is
3861    /// bit-identical to `moe_gate_up_many` run on that token alone — the
3862    /// same int8 activation split, VNNI dots, outlier terms and inline
3863    /// SiLU — while each weight row is read once for all the tokens routed
3864    /// to its expert (the speculative verify's expert sharing). `outs` is
3865    /// flat in (expert, token-of-group) order. False = not covered (not
3866    /// q4tp): the caller takes the per-token path. With float activations,
3867    /// the exact scalar row kernel replaces the int8 dot without changing
3868    /// the shared dispatch or route-order reduction.
3869    pub fn moe_gate_up_rows(
3870        pairs: &[(&QTensor, &QTensor)],
3871        groups: &[Vec<usize>],
3872        xs: &[f32],
3873        outs: &mut [Vec<f32>],
3874        pool: Option<&Pool>,
3875    ) -> bool {
3876        if pairs.is_empty() || pairs.len() != groups.len() {
3877            return false;
3878        }
3879        if pairs.iter().any(|(g, u)| g.is_blr() || u.is_blr()) {
3880            return false;
3881        }
3882        let inter = pairs[0].0.rows();
3883        let cols = pairs[0].0.cols();
3884        let n_pairs: usize = groups.iter().map(|g| g.len()).sum();
3885        if cols == 0 || cols % GROUP_SIZE != 0 || outs.len() != n_pairs || xs.len() % cols != 0 {
3886            return false;
3887        }
3888        let b = xs.len() / cols;
3889        let gpr = cols / GROUP_SIZE;
3890        let mut views = Vec::with_capacity(pairs.len() * 2);
3891        for (g, u) in pairs {
3892            let q4tp = |t: &QTensor| {
3893                matches!(
3894                    t,
3895                    Self::Mapped {
3896                        dtype: TensorDtype::Q4TiledP,
3897                        ..
3898                    }
3899                )
3900            };
3901            if g.has_prism_contract()
3902                || u.has_prism_contract()
3903                || !q4tp(g)
3904                || !q4tp(u)
3905                || g.rows() != inter
3906                || u.rows() != inter
3907                || g.cols() != cols
3908                || u.cols() != cols
3909            {
3910                return false;
3911            }
3912            views.push(Q4tpView::new(g.quant_bytes(), inter, cols));
3913            views.push(Q4tpView::new(u.quant_bytes(), inter, cols));
3914        }
3915        if outs.iter().any(|o| o.len() != inter) || groups.iter().flatten().any(|&t| t >= b) {
3916            return false;
3917        }
3918        let quantized = a8w8_enabled();
3919        let acts: Vec<SplitAct> = if quantized {
3920            (0..b)
3921                .map(|t| split_act(&xs[t * cols..(t + 1) * cols]))
3922                .collect()
3923        } else {
3924            Vec::new()
3925        };
3926        let mut offs = Vec::with_capacity(groups.len());
3927        let mut o = 0usize;
3928        for g in groups {
3929            offs.push(o);
3930            o += g.len();
3931        }
3932        let ptrs: Vec<SendMut> = outs.iter_mut().map(|o| SendMut(o.as_mut_ptr())).collect();
3933        let (views, ptrs, acts, offs) = (&views, &ptrs, &acts, &offs);
3934        let run = |start: usize, end: usize| {
3935            let (mut gsc, mut usc) = (vec![0f32; gpr], vec![0f32; gpr]);
3936            for flat in start..end {
3937                let (e, r) = (flat / inter, flat % inter);
3938                let (gv_view, uv_view) = (&views[e * 2], &views[e * 2 + 1]);
3939                gv_view.scales_into(r, gpr, &mut gsc);
3940                uv_view.scales_into(r, gpr, &mut usc);
3941                for (k, &t) in groups[e].iter().enumerate() {
3942                    if !quantized {
3943                        let x = &xs[t * cols..(t + 1) * cols];
3944                        let gv = q4tp_row_exact(gv_view.nib, r, gpr, x, &gsc);
3945                        let uv = q4tp_row_exact(uv_view.nib, r, gpr, x, &usc);
3946                        unsafe { *ptrs[offs[e] + k].at(r) = (gv / (1.0 + (-gv).exp())) * uv };
3947                        continue;
3948                    }
3949                    let act = &acts[t];
3950                    let mut gv = dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, &gsc) * act.sx;
3951                    let mut uv = dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, &usc) * act.sx;
3952                    for &(j, xv) in &act.outliers {
3953                        let og = q4tp_outlier(gv_view.nib, r, gpr, j, &gsc);
3954                        let ou = q4tp_outlier(uv_view.nib, r, gpr, j, &usc);
3955                        gv += og.0 * og.1 * xv;
3956                        uv += ou.0 * ou.1 * xv;
3957                    }
3958                    let silu_g = gv / (1.0 + (-gv).exp());
3959                    // SAFETY: one worker owns each (expert, row) cell of
3960                    // every output of the expert's group.
3961                    unsafe { *ptrs[offs[e] + k].at(r) = silu_g * uv };
3962                }
3963            }
3964        };
3965        dispatch_rows(pool, pairs.len() * inter, &run);
3966        true
3967    }
3968
3969    /// The per-(expert, token) down terms `moe_down_many` weights and sums,
3970    /// for SEVERAL tokens: `outs[p][o] = down_e[o] · gs[p]` (int8 split of
3971    /// `gs[p]`, VNNI dot, outlier terms — bit-identical to that kernel's
3972    /// `d`), each down row read once for its expert's whole group. The
3973    /// caller sums `w·d` per token in its route order, which reproduces
3974    /// `moe_down_many`'s f32 sequence exactly. Layout as `moe_gate_up_rows`.
3975    pub fn moe_down_rows(
3976        downs: &[&QTensor],
3977        group_lens: &[usize],
3978        gs: &[Vec<f32>],
3979        outs: &mut [Vec<f32>],
3980        pool: Option<&Pool>,
3981    ) -> bool {
3982        if downs.is_empty() || downs.len() != group_lens.len() {
3983            return false;
3984        }
3985        if downs.iter().any(|d| d.is_blr()) {
3986            return false;
3987        }
3988        let rows = downs[0].rows();
3989        let cols = downs[0].cols();
3990        let n_pairs: usize = group_lens.iter().sum();
3991        if cols == 0 || cols % GROUP_SIZE != 0 || gs.len() != n_pairs || outs.len() != n_pairs {
3992            return false;
3993        }
3994        let gpr = cols / GROUP_SIZE;
3995        let mut views = Vec::with_capacity(downs.len());
3996        for d in downs {
3997            if d.has_prism_contract()
3998                || !matches!(
3999                    d,
4000                    Self::Mapped {
4001                        dtype: TensorDtype::Q4TiledP,
4002                        ..
4003                    }
4004                )
4005                || d.rows() != rows
4006                || d.cols() != cols
4007            {
4008                return false;
4009            }
4010            views.push(Q4tpView::new(d.quant_bytes(), rows, cols));
4011        }
4012        if gs.iter().any(|g| g.len() != cols) || outs.iter().any(|o| o.len() != rows) {
4013            return false;
4014        }
4015        let quantized = a8w8_enabled();
4016        let acts: Vec<SplitAct> = if quantized {
4017            gs.iter().map(|g| split_act(g)).collect()
4018        } else {
4019            Vec::new()
4020        };
4021        let mut offs = Vec::with_capacity(group_lens.len());
4022        let mut o = 0usize;
4023        for &l in group_lens {
4024            offs.push(o);
4025            o += l;
4026        }
4027        let ptrs: Vec<SendMut> = outs.iter_mut().map(|o| SendMut(o.as_mut_ptr())).collect();
4028        let (views, ptrs, acts, offs) = (&views, &ptrs, &acts, &offs);
4029        let run = |start: usize, end: usize| {
4030            let mut sc = vec![0f32; gpr];
4031            for flat in start..end {
4032                let (e, r) = (flat / rows, flat % rows);
4033                let v = &views[e];
4034                v.scales_into(r, gpr, &mut sc);
4035                for k in 0..group_lens[e] {
4036                    if !quantized {
4037                        let d = q4tp_row_exact(v.nib, r, gpr, &gs[offs[e] + k], &sc);
4038                        unsafe { *ptrs[offs[e] + k].at(r) = d };
4039                        continue;
4040                    }
4041                    let a = &acts[offs[e] + k];
4042                    let mut d = dot_q4tp_row_i8(v.nib, r, gpr, &a.xq, &sc) * a.sx;
4043                    for &(j, xv) in &a.outliers {
4044                        let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
4045                        d += w * s * xv;
4046                    }
4047                    // SAFETY: one worker owns each (expert, row) cell.
4048                    unsafe { *ptrs[offs[e] + k].at(r) = d };
4049                }
4050            }
4051        };
4052        dispatch_rows(pool, downs.len() * rows, &run);
4053        true
4054    }
4055}
4056
4057/// Batched q8 kernel: same math as qmatvec, the row makes a single
4058/// pass from memory for the whole batch.
4059/// Accelerate CBLAS — the Apple AMX matrix units, the same engine
4060/// llama.cpp's `-ngl 0` prefill rides via ggml-blas.
4061#[cfg(target_os = "macos")]
4062mod accel_blas {
4063    #[link(name = "Accelerate", kind = "framework")]
4064    unsafe extern "C" {
4065        pub fn cblas_sgemm(
4066            order: i32,
4067            trans_a: i32,
4068            trans_b: i32,
4069            m: i32,
4070            n: i32,
4071            k: i32,
4072            alpha: f32,
4073            a: *const f32,
4074            lda: i32,
4075            b: *const f32,
4076            ldb: i32,
4077            beta: f32,
4078            c: *mut f32,
4079            ldc: i32,
4080        );
4081    }
4082}
4083
4084#[cfg(target_os = "macos")]
4085pub(crate) fn accel_gemm_enabled() -> bool {
4086    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4087    *ON.get_or_init(|| std::env::var("CMF_ACCEL").map(|v| v != "0").unwrap_or(true))
4088}
4089
4090/// Off macOS the "accel" GEMM is the portable NEON micro-kernel below —
4091/// same entry point, so the batched-attention path opens on mobile.
4092#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
4093pub(crate) fn accel_gemm_enabled() -> bool {
4094    true
4095}
4096
4097/// Portable NEON f32 GEMM (row-major, optional Bᵀ): a 4×8 fmla
4098/// micro-kernel with A broadcast against B panels — the mobile stand-in
4099/// for Accelerate in the batched causal attention (QKᵀ and P·V). Not a
4100/// BLAS: shapes here are the attention panels (m ≤ heads·chunk,
4101/// k = head_dim or context), and the goal is removing the per-position
4102/// quadratic wall, not peak GEMM.
4103#[cfg(target_arch = "aarch64")]
4104#[allow(clippy::too_many_arguments)]
4105pub(crate) fn neon_gemm_rm(
4106    m: usize,
4107    n: usize,
4108    k: usize,
4109    alpha: f32,
4110    a: &[f32],
4111    lda: usize,
4112    b_mat: &[f32],
4113    ldb: usize,
4114    b_rows_are_n: bool,
4115    c: &mut [f32],
4116    ldc: usize,
4117) {
4118    debug_assert!(a.len() >= (m - 1) * lda + k);
4119    debug_assert!(c.len() >= (m - 1) * ldc + n);
4120    // SAFETY: bounds asserted above; NEON is baseline on aarch64.
4121    unsafe {
4122        use core::arch::aarch64::*;
4123        let mut i = 0usize;
4124        while i < m {
4125            let mi = (m - i).min(4);
4126            let mut j = 0usize;
4127            while j < n {
4128                let nj = (n - j).min(8);
4129                if mi == 4 && nj == 8 {
4130                    let (mut c0a, mut c0b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
4131                    let (mut c1a, mut c1b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
4132                    let (mut c2a, mut c2b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
4133                    let (mut c3a, mut c3b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
4134                    for p in 0..k {
4135                        let (b0, b1) = if b_rows_are_n {
4136                            // B is [n, k]: column p of Bᵀ = element p of
4137                            // eight consecutive B rows — gathered.
4138                            let base = b_mat.as_ptr().add(j * ldb + p);
4139                            let g = |o: usize| *base.add(o * ldb);
4140                            ([g(0), g(1), g(2), g(3)], [g(4), g(5), g(6), g(7)])
4141                        } else {
4142                            let base = b_mat.as_ptr().add(p * ldb + j);
4143                            (
4144                                [*base, *base.add(1), *base.add(2), *base.add(3)],
4145                                [*base.add(4), *base.add(5), *base.add(6), *base.add(7)],
4146                            )
4147                        };
4148                        let bv0 = vld1q_f32(b0.as_ptr());
4149                        let bv1 = vld1q_f32(b1.as_ptr());
4150                        let a0 = vdupq_n_f32(*a.as_ptr().add(i * lda + p));
4151                        let a1 = vdupq_n_f32(*a.as_ptr().add((i + 1) * lda + p));
4152                        let a2 = vdupq_n_f32(*a.as_ptr().add((i + 2) * lda + p));
4153                        let a3 = vdupq_n_f32(*a.as_ptr().add((i + 3) * lda + p));
4154                        c0a = vfmaq_f32(c0a, a0, bv0);
4155                        c0b = vfmaq_f32(c0b, a0, bv1);
4156                        c1a = vfmaq_f32(c1a, a1, bv0);
4157                        c1b = vfmaq_f32(c1b, a1, bv1);
4158                        c2a = vfmaq_f32(c2a, a2, bv0);
4159                        c2b = vfmaq_f32(c2b, a2, bv1);
4160                        c3a = vfmaq_f32(c3a, a3, bv0);
4161                        c3b = vfmaq_f32(c3b, a3, bv1);
4162                    }
4163                    let al = vdupq_n_f32(alpha);
4164                    for (r, (ca, cb)) in [(c0a, c0b), (c1a, c1b), (c2a, c2b), (c3a, c3b)]
4165                        .iter()
4166                        .enumerate()
4167                    {
4168                        let dst = c.as_mut_ptr().add((i + r) * ldc + j);
4169                        vst1q_f32(dst, vmulq_f32(*ca, al));
4170                        vst1q_f32(dst.add(4), vmulq_f32(*cb, al));
4171                    }
4172                } else {
4173                    for r in 0..mi {
4174                        for q in 0..nj {
4175                            let mut acc = 0f32;
4176                            for p in 0..k {
4177                                let bv = if b_rows_are_n {
4178                                    b_mat[(j + q) * ldb + p]
4179                                } else {
4180                                    b_mat[p * ldb + j + q]
4181                                };
4182                                acc += a[(i + r) * lda + p] * bv;
4183                            }
4184                            c[(i + r) * ldc + j + q] = acc * alpha;
4185                        }
4186                    }
4187                }
4188                j += nj;
4189            }
4190            i += mi;
4191        }
4192    }
4193}
4194
4195/// Off-macOS aarch64: the batched attention rides the NEON micro-GEMM.
4196#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
4197#[allow(clippy::too_many_arguments)]
4198pub(crate) fn sgemm_rm(
4199    m: usize,
4200    n: usize,
4201    k: usize,
4202    alpha: f32,
4203    a: &[f32],
4204    lda: usize,
4205    b_mat: &[f32],
4206    ldb: usize,
4207    b_rows_are_n: bool,
4208    c: &mut [f32],
4209    ldc: usize,
4210) {
4211    neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
4212}
4213
4214/// Row-major f32 GEMM, exposed for offline tools (the AWNP pass builds a
4215/// per-layer projection and applies it to every expert; a naive triple loop
4216/// would turn a two-minute job into half an hour).
4217#[allow(clippy::too_many_arguments)]
4218pub fn sgemm_public(
4219    m: usize,
4220    n: usize,
4221    k: usize,
4222    alpha: f32,
4223    a: &[f32],
4224    lda: usize,
4225    b_mat: &[f32],
4226    ldb: usize,
4227    b_rows_are_n: bool,
4228    c: &mut [f32],
4229    ldc: usize,
4230) {
4231    #[cfg(any(target_os = "macos", target_arch = "aarch64"))]
4232    {
4233        sgemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
4234    }
4235    // x86 without Accelerate has no sgemm_rm (the specialized paths there are
4236    // quantized kernels): the packed AVX/FMA/AVX-512 kernels of
4237    // matrixmultiply, single-threaded — offline callers split rows across
4238    // their own threads.
4239    #[cfg(not(any(target_os = "macos", target_arch = "aarch64")))]
4240    {
4241        if m == 0 || n == 0 {
4242            return;
4243        }
4244        assert!(a.len() >= (m - 1) * lda + k && c.len() >= (m - 1) * ldc + n);
4245        let (rsb, csb) = if b_rows_are_n { (1, ldb as isize) } else { (ldb as isize, 1) };
4246        assert!(k == 0 || b_mat.len() > if b_rows_are_n { (n - 1) * ldb + k - 1 } else { (k - 1) * ldb + n - 1 });
4247        // SAFETY: the asserts above bound every element the strides reach.
4248        unsafe {
4249            matrixmultiply::sgemm(
4250                m,
4251                k,
4252                n,
4253                alpha,
4254                a.as_ptr(),
4255                lda as isize,
4256                1,
4257                b_mat.as_ptr(),
4258                rsb,
4259                csb,
4260                0.0,
4261                c.as_mut_ptr(),
4262                ldc as isize,
4263                1,
4264            );
4265        }
4266    }
4267}
4268
4269/// Row-major f32 GEMM on Accelerate: C[m,n] = alpha·A[m,k] × B(ᵀ).
4270/// `b_rows_are_n` = true multiplies by Bᵀ where B is stored [n, k].
4271#[cfg(target_os = "macos")]
4272#[allow(clippy::too_many_arguments)]
4273pub(crate) fn sgemm_rm(
4274    m: usize,
4275    n: usize,
4276    k: usize,
4277    alpha: f32,
4278    a: &[f32],
4279    lda: usize,
4280    b_mat: &[f32],
4281    ldb: usize,
4282    b_rows_are_n: bool,
4283    c: &mut [f32],
4284    ldc: usize,
4285) {
4286    debug_assert!(a.len() >= (m - 1) * lda + k);
4287    debug_assert!(c.len() >= (m - 1) * ldc + n);
4288    // Test hook: route the attention GEMMs through the portable NEON
4289    // micro-kernel ON APPLE SILICON — how the mobile batched attend is
4290    // measured without a phone in the loop. (Intel macOS has no NEON —
4291    // the hook is a no-op there, Accelerate continues below.)
4292    #[cfg(target_arch = "aarch64")]
4293    if std::env::var("CMF_FORCE_NEON_GEMM")
4294        .map(|v| v == "1")
4295        .unwrap_or(false)
4296    {
4297        return neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
4298    }
4299    unsafe {
4300        accel_blas::cblas_sgemm(
4301            101, // RowMajor
4302            111, // NoTrans A
4303            if b_rows_are_n { 112 } else { 111 },
4304            m as i32,
4305            n as i32,
4306            k as i32,
4307            alpha,
4308            a.as_ptr(),
4309            lda as i32,
4310            b_mat.as_ptr(),
4311            ldb as i32,
4312            0.0,
4313            c.as_mut_ptr(),
4314            ldc as i32,
4315        );
4316    }
4317}
4318
4319/// Prefill GEMM through Accelerate (macOS): dequantize q8 rows into
4320/// f32 tiles (scale folded in, pool-parallel) and multiply each tile
4321/// on the AMX with one row-major sgemm. Tiles live in cache, weights
4322/// stream once. Numerics are f32-GEMM (not the int8 dot): prefill
4323/// logits shift within f32 rounding — tolerance-class, like every
4324/// reduction-order change; decode (M=1) never takes this path.
4325#[cfg(target_os = "macos")]
4326fn qmatmat_accel(
4327    q: &[u8],
4328    row_scale: &[f32],
4329    pre: &[std::borrow::Cow<'_, [f32]>],
4330    rows: usize,
4331    cols: usize,
4332    out: &mut [f32],
4333    pool: Option<&Pool>,
4334) {
4335    // NOTE: double-buffering the dequant against the sgemm (a scoped
4336    // thread driving the pool on tile k+1 while the caller multiplies
4337    // tile k) was tried and LOST ~6%: Accelerate's sgemm is itself
4338    // multithreaded, and the dequant workers just steal its cores.
4339    const TR: usize = 2048;
4340    let b = pre.len();
4341    thread_local! {
4342        static XPANEL: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
4343        static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
4344    }
4345    XPANEL.with(|xp| {
4346        WTILE.with(|wt| {
4347            let mut xpanel = xp.borrow_mut();
4348            xpanel.clear();
4349            for x in pre {
4350                xpanel.extend_from_slice(x);
4351            }
4352            let mut wtile = wt.borrow_mut();
4353            wtile.resize(TR * cols, 0.0);
4354            let mut r0 = 0usize;
4355            while r0 < rows {
4356                let tr = TR.min(rows - r0);
4357                // Dequant the tile (scale folded) — pool-parallel.
4358                let wt_addr = SendMut(wtile.as_mut_ptr());
4359                let run = |start: usize, end: usize| {
4360                    for r in start..end {
4361                        let row = &q[(r0 + r) * cols..(r0 + r + 1) * cols];
4362                        let s = row_scale[r0 + r];
4363                        // SAFETY: workers cover disjoint r ranges.
4364                        let dst =
4365                            unsafe { std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols) };
4366                        for (d, &v) in dst.iter_mut().zip(row) {
4367                            *d = (v as i8) as f32 * s;
4368                        }
4369                    }
4370                };
4371                dispatch_rows(pool, tr, &run);
4372                // C[b, tr] (at column r0 of out[b, rows]) = X · Wtileᵀ
4373                unsafe {
4374                    accel_blas::cblas_sgemm(
4375                        101, // RowMajor
4376                        111, // NoTrans A
4377                        112, // Trans B
4378                        b as i32,
4379                        tr as i32,
4380                        cols as i32,
4381                        1.0,
4382                        xpanel.as_ptr(),
4383                        cols as i32,
4384                        wtile.as_ptr(),
4385                        cols as i32,
4386                        0.0,
4387                        out.as_mut_ptr().add(r0),
4388                        rows as i32,
4389                    );
4390                }
4391                r0 += tr;
4392            }
4393        })
4394    });
4395}
4396
4397fn qmatmat(
4398    q: &[u8],
4399    row_scale: &[f32],
4400    pre: &[std::borrow::Cow<'_, [f32]>],
4401    rows: usize,
4402    cols: usize,
4403    out: &mut [f32],
4404    pool: Option<&Pool>,
4405) {
4406    let b = pre.len();
4407    debug_assert_eq!(out.len(), b * rows);
4408    // Big prefill batches ride the AMX (roadmap PR3): the row×batch
4409    // SDOT loop below peaks near the CPU's dot throughput, an order
4410    // below the matrix units. Small tensors and tiny test models stay
4411    // on the exact integer path.
4412    #[cfg(target_os = "macos")]
4413    if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
4414        qmatmat_accel(q, row_scale, pre, rows, cols, out, pool);
4415        return;
4416    }
4417    #[cfg(target_arch = "aarch64")]
4418    if sdot_enabled() {
4419        let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
4420        let out_addr = SendMut(out.as_mut_ptr());
4421        // Blocked 2×4 (mobile prefill: no AMX to fall back on — this
4422        // path IS the ARM prefill GEMM off Apple silicon).
4423        let blocked_ok = blocked_enabled();
4424        let use_i8mm = i8mm_enabled();
4425        if blocked_ok {
4426            let run = |start: usize, end: usize| {
4427                let mut o = start;
4428                while o < end {
4429                    if o + 2 <= end {
4430                        let r0 = &q[o * cols..(o + 1) * cols];
4431                        let r1 = &q[(o + 1) * cols..(o + 2) * cols];
4432                        let mut bi = 0usize;
4433                        while bi + 4 <= acts.len() {
4434                            let xs = [
4435                                acts[bi].xq.as_slice(),
4436                                acts[bi + 1].xq.as_slice(),
4437                                acts[bi + 2].xq.as_slice(),
4438                                acts[bi + 3].xq.as_slice(),
4439                            ];
4440                            let d = if use_i8mm {
4441                                unsafe { dot_i8_smmla_2x4(r0, r1, xs) }
4442                            } else {
4443                                unsafe { dot_i8_sdot_2x4(r0, r1, xs) }
4444                            };
4445                            for (r, row) in [r0, r1].into_iter().enumerate() {
4446                                for k in 0..4 {
4447                                    let act = &acts[bi + k];
4448                                    let mut v = d[r][k] as f32 * act.sx;
4449                                    for &(j, xv) in &act.outliers {
4450                                        v += (row[j] as i8) as f32 * xv;
4451                                    }
4452                                    unsafe {
4453                                        *out_addr.at((bi + k) * rows + o + r) = v * row_scale[o + r]
4454                                    };
4455                                }
4456                            }
4457                            bi += 4;
4458                        }
4459                        while bi < acts.len() {
4460                            for (r, row) in [r0, r1].into_iter().enumerate() {
4461                                let v = row_dot_sdot(row, &acts[bi]) * row_scale[o + r];
4462                                unsafe { *out_addr.at(bi * rows + o + r) = v };
4463                            }
4464                            bi += 1;
4465                        }
4466                        o += 2;
4467                    } else {
4468                        let row = &q[o * cols..(o + 1) * cols];
4469                        for (bi, act) in acts.iter().enumerate() {
4470                            let v = row_dot_sdot(row, act) * row_scale[o];
4471                            unsafe { *out_addr.at(bi * rows + o) = v };
4472                        }
4473                        o += 1;
4474                    }
4475                }
4476            };
4477            dispatch_rows(pool, rows, &run);
4478            return;
4479        }
4480        let run = |start: usize, end: usize| {
4481            for o in start..end {
4482                let row = &q[o * cols..(o + 1) * cols];
4483                for (bi, act) in acts.iter().enumerate() {
4484                    let v = row_dot_sdot(row, act) * row_scale[o];
4485                    unsafe { *out_addr.at(bi * rows + o) = v };
4486                }
4487            }
4488        };
4489        dispatch_rows(pool, rows, &run);
4490        return;
4491    }
4492    // x86 A8W8 batch. Non-VNNI parts take the BLOCKED 2×4 kernel
4493    // (roadmap P0: two weight rows' abs() stay in registers across four
4494    // activation streams); VNNI machines keep the per-row bias-trick
4495    // dot, which is already throughput-bound there.
4496    #[cfg(target_arch = "x86_64")]
4497    if avx2_a8w8_enabled() {
4498        let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
4499        let out_addr = SendMut(out.as_mut_ptr());
4500        // CMF_X86_BLOCKED=0 forces the per-row path (paired in-process
4501        // A/B on noisy shared-vCPU hosts).
4502        let blocked_ok = blocked_enabled();
4503        if !avx512vnni_enabled() && blocked_ok && !row_exact() {
4504            let run = |start: usize, end: usize| {
4505                let mut o = start;
4506                while o < end {
4507                    if o + 2 <= end {
4508                        let r0 = &q[o * cols..(o + 1) * cols];
4509                        let r1 = &q[(o + 1) * cols..(o + 2) * cols];
4510                        let mut bi = 0usize;
4511                        while bi + 4 <= acts.len() {
4512                            let xs = [
4513                                acts[bi].xq.as_slice(),
4514                                acts[bi + 1].xq.as_slice(),
4515                                acts[bi + 2].xq.as_slice(),
4516                                acts[bi + 3].xq.as_slice(),
4517                            ];
4518                            let d = unsafe { dot_i8_i8_avx2_2x4(r0, r1, xs) };
4519                            for (r, row) in [r0, r1].into_iter().enumerate() {
4520                                for k in 0..4 {
4521                                    let act = &acts[bi + k];
4522                                    let mut v = d[r][k] as f32 * act.sx;
4523                                    for &(j, xv) in &act.outliers {
4524                                        v += (row[j] as i8) as f32 * xv;
4525                                    }
4526                                    unsafe {
4527                                        *out_addr.at((bi + k) * rows + o + r) = v * row_scale[o + r]
4528                                    };
4529                                }
4530                            }
4531                            bi += 4;
4532                        }
4533                        while bi < acts.len() {
4534                            for (r, row) in [r0, r1].into_iter().enumerate() {
4535                                let v = row_dot_avx2(row, &acts[bi]) * row_scale[o + r];
4536                                unsafe { *out_addr.at(bi * rows + o + r) = v };
4537                            }
4538                            bi += 1;
4539                        }
4540                        o += 2;
4541                    } else {
4542                        let row = &q[o * cols..(o + 1) * cols];
4543                        for (bi, act) in acts.iter().enumerate() {
4544                            let v = row_dot_avx2(row, act) * row_scale[o];
4545                            unsafe { *out_addr.at(bi * rows + o) = v };
4546                        }
4547                        o += 1;
4548                    }
4549                }
4550            };
4551            dispatch_rows(pool, rows, &run);
4552            return;
4553        }
4554        let run = |start: usize, end: usize| {
4555            for o in start..end {
4556                let row = &q[o * cols..(o + 1) * cols];
4557                for (bi, act) in acts.iter().enumerate() {
4558                    let v = row_dot_avx2(row, act) * row_scale[o];
4559                    unsafe { *out_addr.at(bi * rows + o) = v };
4560                }
4561            }
4562        };
4563        dispatch_rows(pool, rows, &run);
4564        return;
4565    }
4566    let out_addr = SendMut(out.as_mut_ptr());
4567    let run = |start: usize, end: usize| {
4568        for o in start..end {
4569            let row = &q[o * cols..(o + 1) * cols];
4570            for (bi, x) in pre.iter().enumerate() {
4571                let mut acc = 0f32;
4572                for j in 0..cols {
4573                    acc += (row[j] as i8) as f32 * x[j];
4574                }
4575                unsafe { *out_addr.at(bi * rows + o) = acc * row_scale[o] };
4576            }
4577        }
4578    };
4579    dispatch_rows(pool, rows, &run);
4580}
4581
4582/// Split rows across pool workers (shared qmatvec pattern). Self-balancing
4583/// — see `Pool::run_rows` for why a static 1/n split is wrong here.
4584fn dispatch_rows(pool: Option<&Pool>, rows: usize, run: &(dyn Fn(usize, usize) + Sync)) {
4585    match pool {
4586        Some(pool) if rows >= 256 => pool.run_rows(rows, run),
4587        _ => run(0, rows),
4588    }
4589}
4590
4591/// Split a q4_block blob into (packed nibbles, f16 group scales).
4592fn q4_split(bytes: &[u8], rows: usize, cols: usize) -> (&[u8], &[u8]) {
4593    let groups = rows * cols / GROUP_SIZE;
4594    bytes.split_at(groups * 16)
4595}
4596
4597/// SIMD unpack for the dominant vbit width B=4 (94% of rows on the
4598/// log2-shape calibration): 16 packed bytes -> 32 centered i8 values.
4599/// vbit packs MSB-first, so the HIGH nibble is the even element
4600/// (opposite of q4_block's lo-first interleave). Centering is u-7.
4601#[inline]
4602fn vbit_fill4(data: &[u8], buf: &mut [u8]) {
4603    #[cfg(target_arch = "aarch64")]
4604    unsafe {
4605        return vbit_fill4_neon(data, buf);
4606    }
4607    #[cfg(target_arch = "x86_64")]
4608    if avx2_enabled() {
4609        return unsafe { vbit_fill4_avx2(data, buf) };
4610    }
4611    #[allow(unreachable_code)]
4612    for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
4613        let u = unpack8::<4>(&data[blk * 4..]);
4614        for k in 0..8 {
4615            chunk[k] = (u[k] - 7) as i8 as u8;
4616        }
4617    }
4618}
4619
4620#[cfg(target_arch = "aarch64")]
4621#[target_feature(enable = "neon")]
4622unsafe fn vbit_fill4_neon(data: &[u8], buf: &mut [u8]) {
4623    // SAFETY: buf.len() is a multiple of GROUP_SIZE=32; data holds
4624    // buf.len()/2 packed bytes (validated at load).
4625    unsafe {
4626        use core::arch::aarch64::*;
4627        let n = buf.len();
4628        let mask = vdupq_n_u8(0x0F);
4629        let seven = vdupq_n_s8(7);
4630        let mut g = 0usize;
4631        while g * 32 + 32 <= n {
4632            let b = vld1q_u8(data.as_ptr().add(g * 16));
4633            let hi = vshrq_n_u8::<4>(b);
4634            let lo = vandq_u8(b, mask);
4635            let z0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(hi, lo)), seven);
4636            let z1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(hi, lo)), seven);
4637            vst1q_u8(buf.as_mut_ptr().add(g * 32), vreinterpretq_u8_s8(z0));
4638            vst1q_u8(buf.as_mut_ptr().add(g * 32 + 16), vreinterpretq_u8_s8(z1));
4639            g += 1;
4640        }
4641    }
4642}
4643
4644#[cfg(target_arch = "x86_64")]
4645#[target_feature(enable = "avx2")]
4646unsafe fn vbit_fill4_avx2(data: &[u8], buf: &mut [u8]) {
4647    // SAFETY: see vbit_fill4_neon.
4648    unsafe {
4649        use core::arch::x86_64::*;
4650        let n = buf.len();
4651        let mask = _mm_set1_epi8(0x0F);
4652        let seven = _mm256_set1_epi8(7);
4653        let mut g = 0usize;
4654        while g * 32 + 32 <= n {
4655            let b = _mm_loadu_si128(data.as_ptr().add(g * 16) as *const __m128i);
4656            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), mask);
4657            let lo = _mm_and_si128(b, mask);
4658            let z = _mm256_sub_epi8(
4659                _mm256_set_m128i(_mm_unpackhi_epi8(hi, lo), _mm_unpacklo_epi8(hi, lo)),
4660                seven,
4661            );
4662            _mm256_storeu_si256(buf.as_mut_ptr().add(g * 32) as *mut __m256i, z);
4663            g += 1;
4664        }
4665    }
4666}
4667
4668/// Unpack 8 MSB-first B-bit values from exactly B bytes (fixed shifts —
4669/// no serial bit-buffer, auto-vectorizable). Every 32-value group starts
4670/// byte-aligned (32·B/8 is integral for B∈3..8), so groups decompose
4671/// into 4 such blocks.
4672#[inline(always)]
4673fn unpack8<const B: usize>(data: &[u8]) -> [i32; 8] {
4674    let mut acc = 0u64;
4675    for i in 0..B {
4676        acc = (acc << 8) | data[i] as u64;
4677    }
4678    let mask = (1u64 << B) - 1;
4679    let mut out = [0i32; 8];
4680    for (k, o) in out.iter_mut().enumerate() {
4681        *o = ((acc >> ((7 - k) * B)) & mask) as i32;
4682    }
4683    out
4684}
4685
4686/// Fused vbit matvec straight from the mapped bytes (spec §3, P13
4687/// FIG.3): [u8 bits: rows][f16 scales: rows·cols/32][bit-packed rows,
4688/// MSB-first, byte-padded]. Row data offsets are precomputed at load
4689/// (`vbit_row_offsets`) — the per-call prefix scan was O(rows) pure
4690/// overhead on every matvec.
4691#[allow(clippy::too_many_arguments)]
4692fn vbitmatvec(
4693    bytes: &[u8],
4694    offsets: &[usize],
4695    x: &[f32],
4696    rows: usize,
4697    cols: usize,
4698    out: &mut [f32],
4699    pool: Option<&Pool>,
4700) {
4701    debug_assert_eq!(out.len(), rows);
4702    debug_assert_eq!(offsets.len(), rows + 1);
4703
4704    // SDOT path: unpack the row to centered i8 once, then per-group
4705    // int8 dot against the quantized activations — same A8W8 contract
4706    // as q8 (bounded noise; CMF_SDOT=0 keeps the exact scalar path).
4707    if a8w8_enabled() {
4708        let act = split_act(x);
4709        let out_addr = SendMut(out.as_mut_ptr());
4710        let run = move |start: usize, end: usize| {
4711            vbit_range_a8w8(bytes, offsets, x, &act, rows, cols, out_addr, start, end)
4712        };
4713        dispatch_rows(pool, rows, &run);
4714        return;
4715    }
4716
4717    let out_addr = SendMut(out.as_mut_ptr());
4718    let run = move |start: usize, end: usize| {
4719        vbit_range_f32(bytes, offsets, x, rows, cols, out_addr, start, end)
4720    };
4721    dispatch_rows(pool, rows, &run);
4722}
4723
4724/// One vbit row range via the A8W8 int8 path — kernel body of
4725/// `vbitmatvec`, extracted so multi-matrix jobs can drive it for
4726/// several tensors in one dispatch (b=8 rows go exact f32).
4727#[allow(clippy::too_many_arguments)]
4728fn vbit_range_a8w8(
4729    bytes: &[u8],
4730    offsets: &[usize],
4731    x: &[f32],
4732    act: &SplitAct,
4733    rows: usize,
4734    cols: usize,
4735    out: SendMut,
4736    start: usize,
4737    end: usize,
4738) {
4739    let ng = cols / GROUP_SIZE;
4740    let bits = &bytes[..rows];
4741    let sc_off = rows;
4742    let row_dot = |r: usize| -> f32 {
4743        let b = bits[r] as usize;
4744        let l = (1i32 << (b - 1)) - 1;
4745        let mask = (1u64 << b) - 1;
4746        let data = &bytes[offsets[r]..offsets[r + 1]];
4747        if b == 8 {
4748            // u−L reaches 128 → does not fit i8; exact f32 path.
4749            let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
4750            let mut dot = 0f32;
4751            for g in 0..ng {
4752                let so = (r * ng + g) * 2;
4753                let sgf = f16_to_f32(u16::from_le_bytes([
4754                    bytes[sc_off + so],
4755                    bytes[sc_off + so + 1],
4756                ]));
4757                let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
4758                let mut gd = 0f32;
4759                for &xv in xg.iter() {
4760                    if nbits < 8 {
4761                        acc = (acc << 8) | data[idx] as u64;
4762                        idx += 1;
4763                        nbits += 8;
4764                    }
4765                    let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
4766                    nbits -= 8;
4767                    gd += (u - l) as f32 * xv;
4768                }
4769                dot += gd * sgf;
4770            }
4771            return dot;
4772        }
4773        // Per-worker scratch: this closure runs for every row of the
4774        // tensor (lm_head ≈ 150k rows/token) — a heap allocation per
4775        // row was measurable pure overhead.
4776        thread_local! {
4777            static VBIT_SCRATCH: std::cell::RefCell<Vec<u8>> =
4778                const { std::cell::RefCell::new(Vec::new()) };
4779        }
4780        #[inline(always)]
4781        fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
4782            for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
4783                let u = unpack8::<B>(&data[blk * B..]);
4784                for k in 0..8 {
4785                    chunk[k] = (u[k] - l) as i8 as u8;
4786                }
4787            }
4788        }
4789        let _ = mask;
4790        VBIT_SCRATCH.with(|scratch| {
4791            let mut buf = scratch.borrow_mut();
4792            buf.resize(cols, 0);
4793            match b {
4794                3 => fill::<3>(data, l, &mut buf),
4795                4 => vbit_fill4(data, &mut buf),
4796                5 => fill::<5>(data, l, &mut buf),
4797                6 => fill::<6>(data, l, &mut buf),
4798                _ => unreachable!(),
4799            }
4800            let mut dot = 0f32;
4801            for g in 0..ng {
4802                let so = (r * ng + g) * 2;
4803                let s = f16_to_f32(u16::from_le_bytes([
4804                    bytes[sc_off + so],
4805                    bytes[sc_off + so + 1],
4806                ]));
4807                let d = dot_i8_i8(
4808                    &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
4809                    &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
4810                ) as f32
4811                    * act.sx;
4812                dot += d * s;
4813            }
4814            for &(j, xv) in &act.outliers {
4815                let so = (r * ng + j / GROUP_SIZE) * 2;
4816                let s = f16_to_f32(u16::from_le_bytes([
4817                    bytes[sc_off + so],
4818                    bytes[sc_off + so + 1],
4819                ]));
4820                // xq is zeroed at outlier slots — add the exact term.
4821                dot += (buf[j] as i8) as f32 * s * xv;
4822            }
4823            dot
4824        })
4825    };
4826    for r in start..end {
4827        // SAFETY: disjoint row ranges per worker.
4828        unsafe { *out.at(r) = row_dot(r) };
4829    }
4830}
4831
4832/// Exact scalar vbit row range (same extraction, non-SDOT path).
4833#[allow(clippy::too_many_arguments)]
4834fn vbit_range_f32(
4835    bytes: &[u8],
4836    offsets: &[usize],
4837    x: &[f32],
4838    rows: usize,
4839    cols: usize,
4840    out: SendMut,
4841    start: usize,
4842    end: usize,
4843) {
4844    let ng = cols / GROUP_SIZE;
4845    let bits = &bytes[..rows];
4846    let sc_off = rows;
4847    // Per-bit-width specialized inner loops: the compiler unrolls the
4848    // constant shifts (the generic bit-buffer loop was branch-bound —
4849    // 5.6 vs 13.2 tok/s q4 on the 0.8B).
4850    #[inline(always)]
4851    fn dot_row<const B: usize>(
4852        data: &[u8],
4853        bytes: &[u8],
4854        sc_off: usize,
4855        r: usize,
4856        ng: usize,
4857        x: &[f32],
4858    ) -> f32 {
4859        let l = ((1i32 << (B - 1)) - 1) as f32;
4860        let gbytes = GROUP_SIZE * B / 8;
4861        let mut dot = 0f32;
4862        for g in 0..ng {
4863            let so = (r * ng + g) * 2;
4864            let s = f16_to_f32(u16::from_le_bytes([
4865                bytes[sc_off + so],
4866                bytes[sc_off + so + 1],
4867            ]));
4868            let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
4869            let gd0 = &data[g * gbytes..(g + 1) * gbytes];
4870            let mut gd = 0f32;
4871            for blk in 0..GROUP_SIZE / 8 {
4872                let u = unpack8::<B>(&gd0[blk * B..]);
4873                let xb = &xg[blk * 8..blk * 8 + 8];
4874                for k in 0..8 {
4875                    gd += (u[k] as f32 - l) * xb[k];
4876                }
4877            }
4878            dot += gd * s;
4879        }
4880        dot
4881    }
4882    for r in start..end {
4883        let data = &bytes[offsets[r]..offsets[r + 1]];
4884        let v = match bits[r] {
4885            3 => dot_row::<3>(data, bytes, sc_off, r, ng, x),
4886            4 => dot_row::<4>(data, bytes, sc_off, r, ng, x),
4887            5 => dot_row::<5>(data, bytes, sc_off, r, ng, x),
4888            6 => dot_row::<6>(data, bytes, sc_off, r, ng, x),
4889            8 => dot_row::<8>(data, bytes, sc_off, r, ng, x),
4890            b => unreachable!("vbit bit-width {b} (validated at load)"),
4891        };
4892        // SAFETY: disjoint row ranges per worker.
4893        unsafe { *out.at(r) = v };
4894    }
4895}
4896
4897/// Fused two-input vbit matvec: each row is unpacked from the mmap ONCE
4898/// and dotted against BOTH activations (MTP verify / pair prefill used
4899/// to run two full matvecs — double weight traffic and double unpack).
4900/// Per-input math is identical to `vbitmatvec` → same accuracy contract.
4901#[allow(clippy::too_many_arguments)]
4902fn vbitmatvec2(
4903    bytes: &[u8],
4904    offsets: &[usize],
4905    x1: &[f32],
4906    x2: &[f32],
4907    rows: usize,
4908    cols: usize,
4909    o1: &mut [f32],
4910    o2: &mut [f32],
4911    pool: Option<&Pool>,
4912) {
4913    debug_assert_eq!(o1.len(), rows);
4914    debug_assert_eq!(o2.len(), rows);
4915
4916    if a8w8_enabled() {
4917        let a1 = split_act(x1);
4918        let a2 = split_act(x2);
4919        let p1 = SendMut(o1.as_mut_ptr());
4920        let p2 = SendMut(o2.as_mut_ptr());
4921        let run = move |start: usize, end: usize| {
4922            vbit_range2_a8w8(
4923                bytes, offsets, x1, x2, &a1, &a2, rows, cols, p1, p2, start, end,
4924            )
4925        };
4926        dispatch_rows(pool, rows, &run);
4927        return;
4928    }
4929
4930    let p1 = SendMut(o1.as_mut_ptr());
4931    let p2 = SendMut(o2.as_mut_ptr());
4932    let run = move |start: usize, end: usize| {
4933        vbit_range2_f32(bytes, offsets, x1, x2, rows, cols, p1, p2, start, end)
4934    };
4935    dispatch_rows(pool, rows, &run);
4936}
4937
4938/// Two-input vbit row range via the A8W8 int8 path — kernel body of
4939/// `vbitmatvec2`, extracted for pair multi-matrix jobs (b=8 rows go
4940/// exact f32 for both lanes, bits streamed once).
4941#[allow(clippy::too_many_arguments)]
4942fn vbit_range2_a8w8(
4943    bytes: &[u8],
4944    offsets: &[usize],
4945    x1: &[f32],
4946    x2: &[f32],
4947    a1: &SplitAct,
4948    a2: &SplitAct,
4949    rows: usize,
4950    cols: usize,
4951    p1: SendMut,
4952    p2: SendMut,
4953    start: usize,
4954    end: usize,
4955) {
4956    let ng = cols / GROUP_SIZE;
4957    let bits = &bytes[..rows];
4958    let sc_off = rows;
4959    let row_dots = |r: usize| -> (f32, f32) {
4960        let b = bits[r] as usize;
4961        let l = (1i32 << (b - 1)) - 1;
4962        let data = &bytes[offsets[r]..offsets[r + 1]];
4963        if b == 8 {
4964            // u−L reaches 128 → does not fit i8; exact f32 path,
4965            // bits still streamed once for both lanes.
4966            let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
4967            let (mut d1, mut d2) = (0f32, 0f32);
4968            for g in 0..ng {
4969                let so = (r * ng + g) * 2;
4970                let sgf = f16_to_f32(u16::from_le_bytes([
4971                    bytes[sc_off + so],
4972                    bytes[sc_off + so + 1],
4973                ]));
4974                let (mut g1, mut g2) = (0f32, 0f32);
4975                for k in 0..GROUP_SIZE {
4976                    if nbits < 8 {
4977                        acc = (acc << 8) | data[idx] as u64;
4978                        idx += 1;
4979                        nbits += 8;
4980                    }
4981                    let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
4982                    nbits -= 8;
4983                    let w = (u - l) as f32;
4984                    g1 += w * x1[g * GROUP_SIZE + k];
4985                    g2 += w * x2[g * GROUP_SIZE + k];
4986                }
4987                d1 += g1 * sgf;
4988                d2 += g2 * sgf;
4989            }
4990            return (d1, d2);
4991        }
4992        thread_local! {
4993            static VBIT_SCRATCH2: std::cell::RefCell<Vec<u8>> =
4994                const { std::cell::RefCell::new(Vec::new()) };
4995        }
4996        #[inline(always)]
4997        fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
4998            for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
4999                let u = unpack8::<B>(&data[blk * B..]);
5000                for k in 0..8 {
5001                    chunk[k] = (u[k] - l) as i8 as u8;
5002                }
5003            }
5004        }
5005        VBIT_SCRATCH2.with(|scratch| {
5006            let mut buf = scratch.borrow_mut();
5007            buf.resize(cols, 0);
5008            match b {
5009                3 => fill::<3>(data, l, &mut buf),
5010                4 => vbit_fill4(data, &mut buf),
5011                5 => fill::<5>(data, l, &mut buf),
5012                6 => fill::<6>(data, l, &mut buf),
5013                _ => unreachable!(),
5014            }
5015            let (mut d1, mut d2) = (0f32, 0f32);
5016            for g in 0..ng {
5017                let so = (r * ng + g) * 2;
5018                let s = f16_to_f32(u16::from_le_bytes([
5019                    bytes[sc_off + so],
5020                    bytes[sc_off + so + 1],
5021                ]));
5022                let wg = &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
5023                let v1 = dot_i8_i8(wg, &a1.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a1.sx;
5024                let v2 = dot_i8_i8(wg, &a2.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a2.sx;
5025                d1 += v1 * s;
5026                d2 += v2 * s;
5027            }
5028            for &(j, xv) in &a1.outliers {
5029                let so = (r * ng + j / GROUP_SIZE) * 2;
5030                let s = f16_to_f32(u16::from_le_bytes([
5031                    bytes[sc_off + so],
5032                    bytes[sc_off + so + 1],
5033                ]));
5034                d1 += (buf[j] as i8) as f32 * s * xv;
5035            }
5036            for &(j, xv) in &a2.outliers {
5037                let so = (r * ng + j / GROUP_SIZE) * 2;
5038                let s = f16_to_f32(u16::from_le_bytes([
5039                    bytes[sc_off + so],
5040                    bytes[sc_off + so + 1],
5041                ]));
5042                d2 += (buf[j] as i8) as f32 * s * xv;
5043            }
5044            (d1, d2)
5045        })
5046    };
5047    for r in start..end {
5048        let (v1, v2) = row_dots(r);
5049        // SAFETY: disjoint row ranges per worker.
5050        unsafe {
5051            *p1.at(r) = v1;
5052            *p2.at(r) = v2;
5053        }
5054    }
5055}
5056
5057/// Two-input exact scalar vbit row range (same extraction) —
5058/// per-bit-width specialized, two accumulators per row; per-lane
5059/// accumulation order matches `vbitmatvec` exactly.
5060#[allow(clippy::too_many_arguments)]
5061fn vbit_range2_f32(
5062    bytes: &[u8],
5063    offsets: &[usize],
5064    x1: &[f32],
5065    x2: &[f32],
5066    rows: usize,
5067    cols: usize,
5068    p1: SendMut,
5069    p2: SendMut,
5070    start: usize,
5071    end: usize,
5072) {
5073    let ng = cols / GROUP_SIZE;
5074    let bits = &bytes[..rows];
5075    let sc_off = rows;
5076    #[inline(always)]
5077    #[allow(clippy::too_many_arguments)]
5078    fn dot_row2<const B: usize>(
5079        data: &[u8],
5080        bytes: &[u8],
5081        sc_off: usize,
5082        r: usize,
5083        ng: usize,
5084        x1: &[f32],
5085        x2: &[f32],
5086    ) -> (f32, f32) {
5087        let l = ((1i32 << (B - 1)) - 1) as f32;
5088        let gbytes = GROUP_SIZE * B / 8;
5089        let (mut d1, mut d2) = (0f32, 0f32);
5090        for g in 0..ng {
5091            let so = (r * ng + g) * 2;
5092            let s = f16_to_f32(u16::from_le_bytes([
5093                bytes[sc_off + so],
5094                bytes[sc_off + so + 1],
5095            ]));
5096            let x1g = &x1[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
5097            let x2g = &x2[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
5098            let gd0 = &data[g * gbytes..(g + 1) * gbytes];
5099            let (mut g1, mut g2) = (0f32, 0f32);
5100            for blk in 0..GROUP_SIZE / 8 {
5101                let u = unpack8::<B>(&gd0[blk * B..]);
5102                for k in 0..8 {
5103                    let w = u[k] as f32 - l;
5104                    g1 += w * x1g[blk * 8 + k];
5105                    g2 += w * x2g[blk * 8 + k];
5106                }
5107            }
5108            d1 += g1 * s;
5109            d2 += g2 * s;
5110        }
5111        (d1, d2)
5112    }
5113    for r in start..end {
5114        let data = &bytes[offsets[r]..offsets[r + 1]];
5115        let (v1, v2) = match bits[r] {
5116            3 => dot_row2::<3>(data, bytes, sc_off, r, ng, x1, x2),
5117            4 => dot_row2::<4>(data, bytes, sc_off, r, ng, x1, x2),
5118            5 => dot_row2::<5>(data, bytes, sc_off, r, ng, x1, x2),
5119            6 => dot_row2::<6>(data, bytes, sc_off, r, ng, x1, x2),
5120            8 => dot_row2::<8>(data, bytes, sc_off, r, ng, x1, x2),
5121            b => unreachable!("vbit bit-width {b} (validated at load)"),
5122        };
5123        // SAFETY: disjoint row ranges per worker.
5124        unsafe {
5125            *p1.at(r) = v1;
5126            *p2.at(r) = v2;
5127        }
5128    }
5129}
5130
5131// ───────────────────── q4_tiled kernels (§4.3) ─────────────────────
5132
5133/// One q4_tiled row dot on the A8W8 int8 path: per 32-group the tile
5134/// is ONE sequential read — [f16 scale][16B nibbles] — versus the two
5135/// distant streams of the split layout. Values/order identical to the
5136/// split kernels.
5137#[inline]
5138#[allow(unreachable_code)]
5139fn dot_q4t_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
5140    #[cfg(target_arch = "aarch64")]
5141    unsafe {
5142        return dot_q4t_row_sdot(bytes, r, gpr, xq);
5143    }
5144    #[cfg(target_arch = "x86_64")]
5145    unsafe {
5146        if vnni_tiles_enabled() {
5147            return dot_q4t_row_vnni(bytes, r, gpr, xq);
5148        }
5149        return dot_q4t_row_avx2(bytes, r, gpr, xq);
5150    }
5151    let mut acc = 0f32;
5152    for gi in 0..gpr {
5153        let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
5154        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
5155        let mut d = 0i32;
5156        for (k, &b) in tile[2..].iter().enumerate() {
5157            d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
5158                + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
5159        }
5160        acc += d as f32 * s;
5161    }
5162    acc
5163}
5164
5165#[cfg(target_arch = "aarch64")]
5166#[target_feature(enable = "neon,dotprod")]
5167unsafe fn dot_q4t_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
5168    // SAFETY: callers uphold slice-length contracts (18B tile per group,
5169    // xq.len() == gpr·GROUP_SIZE).
5170    unsafe {
5171        use core::arch::aarch64::*;
5172        use core::arch::asm;
5173        let lomask = vdupq_n_u8(0x0F);
5174        let eight = vdupq_n_s8(8);
5175        let mut acc = 0f32;
5176        for gi in 0..gpr {
5177            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
5178            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5179            let b = vld1q_u8(t.add(2));
5180            let lo = vandq_u8(b, lomask);
5181            let hi = vshrq_n_u8::<4>(b);
5182            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
5183            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
5184            let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
5185            let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
5186            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5187            asm!(
5188                "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
5189                "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
5190                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5191                e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
5192                options(pure, nomem, nostack),
5193            );
5194            acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
5195        }
5196        acc
5197    }
5198}
5199
5200#[cfg(target_arch = "x86_64")]
5201#[target_feature(enable = "avx2")]
5202unsafe fn dot_q4t_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
5203    // SAFETY: see dot_q4t_row_sdot.
5204    unsafe {
5205        use core::arch::x86_64::*;
5206        let lomask = _mm_set1_epi8(0x0F);
5207        let eight = _mm256_set1_epi8(8);
5208        let ones = _mm256_set1_epi16(1);
5209        let mut acc = 0f32;
5210        for gi in 0..gpr {
5211            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
5212            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5213            let b = _mm_loadu_si128(t.add(2) as *const __m128i);
5214            let lo = _mm_and_si128(b, lomask);
5215            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
5216            let w = _mm256_sub_epi8(
5217                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5218                eight,
5219            );
5220            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5221            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
5222            let d = _mm256_madd_epi16(p16, ones);
5223            let hi128 = _mm256_extracti128_si256::<1>(d);
5224            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
5225            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
5226            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
5227            acc += _mm_cvtsi128_si32(s32) as f32 * s;
5228        }
5229        acc
5230    }
5231}
5232
5233/// VNNI twin of `dot_q4t_row_avx2`: same unpack, `vpdpbusd` replaces
5234/// the maddubs+madd pair (see `dpbusd_hsum` — sums are bit-identical).
5235/// 256-bit VL encoding, so the VEX `vpsignb` stays usable.
5236#[cfg(target_arch = "x86_64")]
5237#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
5238unsafe fn dot_q4t_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
5239    // SAFETY: see dot_q4t_row_sdot.
5240    unsafe {
5241        use core::arch::x86_64::*;
5242        let lomask = _mm_set1_epi8(0x0F);
5243        let eight = _mm256_set1_epi8(8);
5244        let mut acc = 0f32;
5245        for gi in 0..gpr {
5246            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
5247            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5248            let b = _mm_loadu_si128(t.add(2) as *const __m128i);
5249            let lo = _mm_and_si128(b, lomask);
5250            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
5251            let w = _mm256_sub_epi8(
5252                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5253                eight,
5254            );
5255            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5256            let d = dpbusd_hsum(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
5257            acc += d as f32 * s;
5258        }
5259        acc
5260    }
5261}
5262
5263/// One q4_tiled row against FOUR activation streams: the nibble unpack
5264/// and abs() happen once per group instead of once per (group,
5265/// activation) — the unpack is the dominant per-element cost of the
5266/// tiled format (roadmap P0 portable blocking, q4t leg).
5267#[cfg(target_arch = "x86_64")]
5268// `fma` is NOT implied by `avx2`: without it LLVM lowers _mm256_fmadd_ps
5269// to a libm call per lane — measured 2x slower than the reduction this
5270// kernel replaces. The runtime gate (`avx2_enabled`) already requires
5271// both features, so declaring it here is safe.
5272#[target_feature(enable = "avx2,fma")]
5273unsafe fn dot_q4t_row_1x4_avx2(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
5274    // SAFETY: callers uphold the 18B-tile and xq-length contracts.
5275    unsafe {
5276        use core::arch::x86_64::*;
5277        let lomask = _mm_set1_epi8(0x0F);
5278        let eight = _mm256_set1_epi8(8);
5279        let ones = _mm256_set1_epi16(1);
5280        // One f32 accumulator VECTOR per activation, reduced once at the
5281        // end. Folding each group's i32 lanes to a scalar inside the loop
5282        // costs an extracti128 + three shift/add + a movd — a cross-lane
5283        // dependency chain per (group, activation), 288 of them per row at
5284        // cols=2304. The per-group scale is what forces a float
5285        // accumulator; it does not force a horizontal sum.
5286        //
5287        // The four accumulators are NAMED, not an array: as `[__m256; 4]`
5288        // indexed by a loop variable LLVM keeps them in memory and every
5289        // group pays four 32-byte loads and stores. That alone made this
5290        // kernel 2x SLOWER than the per-group reduction it replaces
5291        // (measured on the EPYC box: 150 s vs 71 s for two 256² steps).
5292        let mut f0 = _mm256_setzero_ps();
5293        let mut f1 = _mm256_setzero_ps();
5294        let mut f2 = _mm256_setzero_ps();
5295        let mut f3 = _mm256_setzero_ps();
5296        for gi in 0..gpr {
5297            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
5298            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5299            let sv = _mm256_set1_ps(s);
5300            let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
5301            let lo = _mm_and_si128(bb, lomask);
5302            let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
5303            let w = _mm256_sub_epi8(
5304                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5305                eight,
5306            );
5307            let aw = _mm256_abs_epi8(w);
5308            let off = gi * GROUP_SIZE;
5309            let dot = |xq: &[i8]| {
5310                let x = _mm256_loadu_si256(xq.as_ptr().add(off) as *const __m256i);
5311                let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
5312                _mm256_cvtepi32_ps(_mm256_madd_epi16(p16, ones))
5313            };
5314            f0 = _mm256_fmadd_ps(dot(xs[0]), sv, f0);
5315            f1 = _mm256_fmadd_ps(dot(xs[1]), sv, f1);
5316            f2 = _mm256_fmadd_ps(dot(xs[2]), sv, f2);
5317            f3 = _mm256_fmadd_ps(dot(xs[3]), sv, f3);
5318        }
5319        [
5320            hsum256_ps(f0),
5321            hsum256_ps(f1),
5322            hsum256_ps(f2),
5323            hsum256_ps(f3),
5324        ]
5325    }
5326}
5327
5328/// Horizontal sum of eight f32 lanes — the one cross-lane reduction the
5329/// blocked kernels pay, once per row instead of once per group.
5330#[cfg(target_arch = "x86_64")]
5331#[target_feature(enable = "avx2")]
5332#[inline]
5333unsafe fn hsum256_ps(v: core::arch::x86_64::__m256) -> f32 {
5334    // SAFETY: pure register arithmetic on the caller's vector.
5335    unsafe {
5336        use core::arch::x86_64::*;
5337        let hi = _mm256_extractf128_ps::<1>(v);
5338        let s = _mm_add_ps(_mm256_castps256_ps128(v), hi);
5339        let s = _mm_add_ps(s, _mm_movehl_ps(s, s));
5340        let s = _mm_add_ss(s, _mm_shuffle_ps::<0x55>(s, s));
5341        _mm_cvtss_f32(s)
5342    }
5343}
5344
5345/// VNNI twin of `dot_q4t_row_1x4_avx2` (see `dpbusd_hsum`).
5346#[cfg(target_arch = "x86_64")]
5347#[target_feature(enable = "avx2,fma,avx512f,avx512bw,avx512vl,avx512vnni")]
5348unsafe fn dot_q4t_row_1x4_vnni(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
5349    // SAFETY: callers uphold the 18B-tile and xq-length contracts.
5350    unsafe {
5351        use core::arch::x86_64::*;
5352        let lomask = _mm_set1_epi8(0x0F);
5353        let eight = _mm256_set1_epi8(8);
5354        // Same shape as the AVX2 twin: accumulate in f32 vectors and pay
5355        // one cross-lane reduction per row, not per (group, activation).
5356        let mut f0 = _mm256_setzero_ps();
5357        let mut f1 = _mm256_setzero_ps();
5358        let mut f2 = _mm256_setzero_ps();
5359        let mut f3 = _mm256_setzero_ps();
5360        for gi in 0..gpr {
5361            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
5362            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5363            let sv = _mm256_set1_ps(s);
5364            let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
5365            let lo = _mm_and_si128(bb, lomask);
5366            let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
5367            let w = _mm256_sub_epi8(
5368                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5369                eight,
5370            );
5371            let aw = _mm256_abs_epi8(w);
5372            let off = gi * GROUP_SIZE;
5373            let dot = |xq: &[i8]| {
5374                let x = _mm256_loadu_si256(xq.as_ptr().add(off) as *const __m256i);
5375                _mm256_cvtepi32_ps(_mm256_dpbusd_epi32(
5376                    _mm256_setzero_si256(),
5377                    aw,
5378                    _mm256_sign_epi8(x, w),
5379                ))
5380            };
5381            f0 = _mm256_fmadd_ps(dot(xs[0]), sv, f0);
5382            f1 = _mm256_fmadd_ps(dot(xs[1]), sv, f1);
5383            f2 = _mm256_fmadd_ps(dot(xs[2]), sv, f2);
5384            f3 = _mm256_fmadd_ps(dot(xs[3]), sv, f3);
5385        }
5386        let acc = [
5387            hsum256_ps(f0),
5388            hsum256_ps(f1),
5389            hsum256_ps(f2),
5390            hsum256_ps(f3),
5391        ];
5392        acc
5393    }
5394}
5395
5396/// ARM twin of `dot_q4t_row_1x4_avx2`: one nibble unpack per group
5397/// serves FOUR activation streams. Per stream the group order and f32
5398/// accumulation match `dot_q4t_row_sdot` exactly — batch == matvec
5399/// bit-for-bit.
5400#[cfg(target_arch = "aarch64")]
5401#[target_feature(enable = "neon,dotprod")]
5402unsafe fn dot_q4t_row_1x4_sdot(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
5403    // SAFETY: callers uphold the 18B-tile and xq-length contracts.
5404    unsafe {
5405        use core::arch::aarch64::*;
5406        use core::arch::asm;
5407        let lomask = vdupq_n_u8(0x0F);
5408        let eight = vdupq_n_s8(8);
5409        let mut acc = [0f32; 4];
5410        for gi in 0..gpr {
5411            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
5412            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5413            let b = vld1q_u8(t.add(2));
5414            let lo = vandq_u8(b, lomask);
5415            let hi = vshrq_n_u8::<4>(b);
5416            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
5417            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
5418            for (k, xq) in xs.iter().enumerate() {
5419                let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
5420                let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
5421                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5422                asm!(
5423                    "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
5424                    "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
5425                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5426                    e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
5427                    options(pure, nomem, nostack),
5428                );
5429                acc[k] += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
5430            }
5431        }
5432        acc
5433    }
5434}
5435
5436/// Exact-term correction for A8W8 outliers on a tiled row.
5437#[inline]
5438fn q4t_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
5439    let gi = j / GROUP_SIZE;
5440    let k = j % GROUP_SIZE;
5441    let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
5442    let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
5443    let byte = tile[2 + k / 2];
5444    let nib = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
5445    ((nib as i32 - 8) as f32, s)
5446}
5447
5448/// Exact scalar q4_tiled row (CMF_SDOT=0 contract) — same pairwise
5449/// accumulation shape as `q4_range_f32`.
5450#[inline]
5451fn q4t_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
5452    let mut acc = 0f32;
5453    for gi in 0..gpr {
5454        let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
5455        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
5456        let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
5457        let mut ga = 0f32;
5458        for (k, &b) in tile[2..].iter().enumerate() {
5459            ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
5460                + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
5461        }
5462        acc += ga * s;
5463    }
5464    acc
5465}
5466
5467/// Split view of a `q4tp` payload. The three planes are resolved once per
5468/// matvec instead of per row — `q4tp_sections` is cheap, but doing it inside
5469/// the row loop would put a division on the hot path for nothing.
5470struct Q4tpView<'a> {
5471    nib: &'a [u8],
5472    params: &'a [u8],
5473    codes: &'a [u8],
5474    stride: usize,
5475    /// q2tp reads the ladder with rung 0 = exact zero.
5476    zero_rung: bool,
5477}
5478
5479impl<'a> Q4tpView<'a> {
5480    fn new(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
5481        let (params_off, codes_off, stride) = q4tp_sections(rows, cols);
5482        Self {
5483            nib: &bytes[..params_off],
5484            params: &bytes[params_off..codes_off],
5485            codes: &bytes[codes_off..],
5486            stride,
5487            zero_rung: false,
5488        }
5489    }
5490
5491    /// The q2tp view: identical params/codes planes, 8 B weight chunks.
5492    fn new_q2(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
5493        let (params_off, codes_off, stride) = q2tp_sections(rows, cols);
5494        Self {
5495            nib: &bytes[..params_off],
5496            params: &bytes[params_off..codes_off],
5497            codes: &bytes[codes_off..],
5498            stride,
5499            zero_rung: true,
5500        }
5501    }
5502
5503    /// Expand row `r`'s per-tile scales into `out` (length `gpr`).
5504    ///
5505    /// Doing this once per row — rather than decoding a 5-bit code inside the
5506    /// tile loop — is what makes the format free at runtime. Random access to
5507    /// a packed 5-bit field costs a division, two bounds checks and a branch;
5508    /// the tile's actual work is two `sdot`s, so per-tile decoding dominated
5509    /// the kernel and cost 5x (measured: 1.4 vs 6.9 tok/s on Nanbeige-3B).
5510    /// Walking the plane sequentially with a bit accumulator is ~3 ops.
5511    /// Eight 5-bit codes are exactly five bytes, so a whole group of
5512    /// eight decodes from one little-endian word at fixed shifts. The
5513    /// bit-accumulator this replaces carried a data-dependent `while
5514    /// have < 5` refill whose branch sat in the innermost loop of every
5515    /// q4tp row; a decode profile put this function above the dot
5516    /// products it feeds. Same bitstream, same codes — just no branch
5517    /// and eight independent extractions.
5518    #[inline]
5519    fn scales_into(&self, r: usize, gpr: usize, out: &mut [f32]) {
5520        let tab = if self.zero_rung {
5521            q2tp_ladder(self.params, r)
5522        } else {
5523            q4tp_ladder(self.params, r)
5524        };
5525        let codes = &self.codes[r * self.stride..(r + 1) * self.stride];
5526        let out = &mut out[..gpr];
5527        let mut chunks = out.chunks_exact_mut(8);
5528        let mut ci = 0usize;
5529        for c in &mut chunks {
5530            let w = u64::from(codes[ci])
5531                | u64::from(codes[ci + 1]) << 8
5532                | u64::from(codes[ci + 2]) << 16
5533                | u64::from(codes[ci + 3]) << 24
5534                | u64::from(codes[ci + 4]) << 32;
5535            for (k, o) in c.iter_mut().enumerate() {
5536                *o = tab[((w >> (5 * k)) & 31) as usize];
5537            }
5538            ci += 5;
5539        }
5540        // Fewer than eight codes left: the shared total accessor, which
5541        // tolerates a 5-bit field whose spill byte is past the stride.
5542        let tail = &codes[ci..];
5543        for (k, o) in chunks.into_remainder().iter_mut().enumerate() {
5544            *o = tab[q4tp_code(tail, k)];
5545        }
5546    }
5547}
5548
5549#[inline]
5550fn dot_q4tp_row_i8(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
5551    #[cfg(target_arch = "aarch64")]
5552    unsafe {
5553        return dot_q4tp_row_sdot(nib, r, gpr, xq, scales);
5554    }
5555    #[cfg(target_arch = "x86_64")]
5556    unsafe {
5557        if vnni_tiles_enabled() {
5558            return dot_q4tp_row_vnni(nib, r, gpr, xq, scales);
5559        }
5560        return dot_q4tp_row_avx2(nib, r, gpr, xq, scales);
5561    }
5562    #[allow(unreachable_code)]
5563    {
5564        let mut acc = 0f32;
5565        for gi in 0..gpr {
5566            let tile = &nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
5567            let s = scales[gi];
5568            let mut d = 0i32;
5569            for (k, &b) in tile.iter().enumerate() {
5570                d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
5571                    + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
5572            }
5573            acc += d as f32 * s;
5574        }
5575        acc
5576    }
5577}
5578
5579/// q4tp twin of `dot_q4t_row_sdot`: identical nibble math, but the tile
5580/// stride is 16 B (no inline scale) and the scale is a ladder lookup.
5581#[cfg(target_arch = "aarch64")]
5582#[target_feature(enable = "neon,dotprod")]
5583unsafe fn dot_q4tp_row_sdot(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
5584    // SAFETY: callers uphold slice-length contracts (16B tile per group,
5585    // xq.len() == gpr·GROUP_SIZE, codes covering gpr 5-bit fields).
5586    unsafe {
5587        use core::arch::aarch64::*;
5588        use core::arch::asm;
5589        let lomask = vdupq_n_u8(0x0F);
5590        let eight = vdupq_n_s8(8);
5591        let mut acc = 0f32;
5592        for gi in 0..gpr {
5593            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5594            let s = *scales.get_unchecked(gi);
5595            let b = vld1q_u8(t);
5596            let lo = vandq_u8(b, lomask);
5597            let hi = vshrq_n_u8::<4>(b);
5598            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
5599            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
5600            let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
5601            let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
5602            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5603            asm!(
5604                "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
5605                "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
5606                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5607                e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
5608                options(pure, nomem, nostack),
5609            );
5610            acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
5611        }
5612        acc
5613    }
5614}
5615
5616#[cfg(target_arch = "x86_64")]
5617#[target_feature(enable = "avx2")]
5618unsafe fn dot_q4tp_row_avx2(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
5619    // SAFETY: see dot_q4tp_row_sdot.
5620    unsafe {
5621        use core::arch::x86_64::*;
5622        let lomask = _mm_set1_epi8(0x0F);
5623        let eight = _mm256_set1_epi8(8);
5624        let ones = _mm256_set1_epi16(1);
5625        let mut acc = 0f32;
5626        for gi in 0..gpr {
5627            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5628            let s = *scales.get_unchecked(gi);
5629            let b = _mm_loadu_si128(t as *const __m128i);
5630            let lo = _mm_and_si128(b, lomask);
5631            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
5632            let w = _mm256_sub_epi8(
5633                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5634                eight,
5635            );
5636            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5637            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
5638            let d = _mm256_madd_epi16(p16, ones);
5639            let hi128 = _mm256_extracti128_si256::<1>(d);
5640            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
5641            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
5642            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
5643            acc += _mm_cvtsi128_si32(s32) as f32 * s;
5644        }
5645        acc
5646    }
5647}
5648
5649/// VNNI twin of `dot_q4tp_row_avx2` (see `dot_q4t_row_vnni` for why the
5650/// 256-bit VL encoding is the one to use here).
5651#[cfg(target_arch = "x86_64")]
5652#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
5653unsafe fn dot_q4tp_row_vnni(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
5654    // SAFETY: see dot_q4tp_row_sdot.
5655    unsafe {
5656        use core::arch::x86_64::*;
5657        let lomask = _mm_set1_epi8(0x0F);
5658        let eight = _mm256_set1_epi8(8);
5659        let mut acc = 0f32;
5660        for gi in 0..gpr {
5661            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5662            let s = *scales.get_unchecked(gi);
5663            let b = _mm_loadu_si128(t as *const __m128i);
5664            let lo = _mm_and_si128(b, lomask);
5665            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
5666            let w = _mm256_sub_epi8(
5667                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5668                eight,
5669            );
5670            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5671            acc += dpbusd_hsum(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w)) as f32 * s;
5672        }
5673        acc
5674    }
5675}
5676
5677/// Exact scalar q4tp row — the `CMF_SDOT=0` contract, same pairwise
5678/// accumulation shape as `q4t_row_exact`.
5679#[inline]
5680fn q4tp_row_exact(nib: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
5681    #[cfg(target_arch = "x86_64")]
5682    if avx2_enabled() {
5683        // Keep the scalar pair/group reduction order, not a vector sum.
5684        return unsafe { q4tp_row_float_avx2(nib, r, gpr, x, scales) };
5685    }
5686    q4tp_row_float_scalar(nib, r, gpr, x, scales)
5687}
5688
5689#[inline]
5690fn q4tp_row_float_scalar(nib: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
5691    let mut acc = 0f32;
5692    for gi in 0..gpr {
5693        let tile = &nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
5694        let s = scales[gi];
5695        let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
5696        let mut ga = 0f32;
5697        for (k, &b) in tile.iter().enumerate() {
5698            ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
5699                + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
5700        }
5701        acc += ga * s;
5702    }
5703    acc
5704}
5705
5706/// Vectorize unpack, conversion and multiplication, but preserve every
5707/// pair addition and the scalar accumulation order. No activation rounding
5708/// or FMA: bit-identical to the float scalar row, including its group scale.
5709#[cfg(target_arch = "x86_64")]
5710#[target_feature(enable = "avx2")]
5711unsafe fn q4tp_row_float_avx2(nib: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
5712    // SAFETY: caller checks AVX2 and provides the same complete 32-element
5713    // groups as the scalar row. Loads/stores are explicitly unaligned.
5714    unsafe {
5715        use core::arch::x86_64::*;
5716        let mask = _mm_set1_epi8(15);
5717        let eight = _mm_set1_epi8(8);
5718        let order = _mm256_setr_epi32(0, 1, 4, 5, 2, 3, 6, 7);
5719        let mut acc = 0.0f32;
5720        for gi in 0..gpr {
5721            let packed = _mm_loadu_si128(nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB).cast());
5722            let lo = _mm_and_si128(packed, mask);
5723            let hi = _mm_and_si128(_mm_srli_epi16::<4>(packed), mask);
5724            let w0 = _mm_sub_epi8(_mm_unpacklo_epi8(lo, hi), eight);
5725            let w1 = _mm_sub_epi8(_mm_unpackhi_epi8(lo, hi), eight);
5726            let xp = x.as_ptr().add(gi * GROUP_SIZE);
5727            let a = _mm256_mul_ps(
5728                _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(w0)),
5729                _mm256_loadu_ps(xp),
5730            );
5731            let b = _mm256_mul_ps(
5732                _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(w0))),
5733                _mm256_loadu_ps(xp.add(8)),
5734            );
5735            let c = _mm256_mul_ps(
5736                _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(w1)),
5737                _mm256_loadu_ps(xp.add(16)),
5738            );
5739            let d = _mm256_mul_ps(
5740                _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(w1))),
5741                _mm256_loadu_ps(xp.add(24)),
5742            );
5743            let mut pairs = [0.0f32; 16];
5744            _mm256_storeu_ps(
5745                pairs.as_mut_ptr(),
5746                _mm256_permutevar8x32_ps(_mm256_hadd_ps(a, b), order),
5747            );
5748            _mm256_storeu_ps(
5749                pairs.as_mut_ptr().add(8),
5750                _mm256_permutevar8x32_ps(_mm256_hadd_ps(c, d), order),
5751            );
5752            let mut ga = 0.0f32;
5753            for v in pairs {
5754                ga += v;
5755            }
5756            acc += ga * scales[gi];
5757        }
5758        acc
5759    }
5760}
5761
5762/// Single weight of a q4tp tensor — the a8w8 outlier path, which restores
5763/// activation outliers at full precision after the int8 pass.
5764#[inline]
5765fn q4tp_outlier(nib: &[u8], r: usize, gpr: usize, j: usize, scales: &[f32]) -> (f32, f32) {
5766    let (gi, k) = (j / GROUP_SIZE, j % GROUP_SIZE);
5767    let byte = nib[(r * gpr + gi) * Q4TP_NIB + k / 2];
5768    let n = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
5769    ((n as i32 - 8) as f32, scales[gi])
5770}
5771
5772/// Fused q4tp matvec (dispatch mirrors `q4t_matvec`).
5773fn q4tp_matvec(
5774    bytes: &[u8],
5775    x: &[f32],
5776    rows: usize,
5777    cols: usize,
5778    out: &mut [f32],
5779    pool: Option<&Pool>,
5780) {
5781    debug_assert_eq!(out.len(), rows);
5782    let gpr = cols / GROUP_SIZE;
5783    let v = Q4tpView::new(bytes, rows, cols);
5784    let out_addr = SendMut(out.as_mut_ptr());
5785    if a8w8_enabled() {
5786        let act = split_act(x);
5787        let run = |start: usize, end: usize| {
5788            // One scratch row of scales per worker — borrowed, not minted.
5789            with_krow(gpr, |sc| {
5790                for r in start..end {
5791                    v.scales_into(r, gpr, sc);
5792                    let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, sc) * act.sx;
5793                    for &(j, xv) in &act.outliers {
5794                        let (w, s) = q4tp_outlier(v.nib, r, gpr, j, sc);
5795                        acc += w * s * xv;
5796                    }
5797                    // SAFETY: disjoint row ranges per worker.
5798                    unsafe { *out_addr.at(r) = acc };
5799                }
5800            })
5801        };
5802        dispatch_rows(pool, rows, &run);
5803        return;
5804    }
5805    let run = |start: usize, end: usize| {
5806        with_krow(gpr, |sc| {
5807            for r in start..end {
5808                v.scales_into(r, gpr, sc);
5809                // SAFETY: disjoint row ranges per worker.
5810                unsafe { *out_addr.at(r) = q4tp_row_exact(v.nib, r, gpr, x, sc) };
5811            }
5812        })
5813    };
5814    dispatch_rows(pool, rows, &run);
5815}
5816
5817/// Fused two-input q4tp matvec — the SwiGLU gate/up pair. Weights and the
5818/// row ladder are read once and spent on both activation streams.
5819#[allow(clippy::too_many_arguments)]
5820fn q4tp_matvec2(
5821    bytes: &[u8],
5822    x1: &[f32],
5823    x2: &[f32],
5824    rows: usize,
5825    cols: usize,
5826    o1: &mut [f32],
5827    o2: &mut [f32],
5828    pool: Option<&Pool>,
5829) {
5830    let gpr = cols / GROUP_SIZE;
5831    let v = Q4tpView::new(bytes, rows, cols);
5832    let (p1, p2) = (SendMut(o1.as_mut_ptr()), SendMut(o2.as_mut_ptr()));
5833    let run = |start: usize, end: usize| {
5834        let mut sc = vec![0f32; gpr];
5835        for r in start..end {
5836            v.scales_into(r, gpr, &mut sc);
5837            // SAFETY: disjoint row ranges per worker.
5838            unsafe {
5839                *p1.at(r) = q4tp_row_exact(v.nib, r, gpr, x1, &sc);
5840                *p2.at(r) = q4tp_row_exact(v.nib, r, gpr, x2, &sc);
5841            }
5842        }
5843    };
5844    dispatch_rows(pool, rows, &run);
5845}
5846
5847/// One q2tp outlier weight at column `j` of row `r`: the 2-bit code and
5848/// its group scale, mirrored on `q4tp_outlier`.
5849#[inline]
5850fn q2tp_outlier(chunks: &[u8], r: usize, gpr: usize, j: usize, scales: &[f32]) -> (f32, f32) {
5851    let (gi, k) = (j / GROUP_SIZE, j % GROUP_SIZE);
5852    let byte = chunks[(r * gpr + gi) * Q2TP_CHUNK + k / 4];
5853    let c = (byte >> (2 * (k % 4))) & 3;
5854    (c as f32 - 1.5, scales[gi])
5855}
5856
5857#[cfg(target_arch = "x86_64")]
5858const Q2TP_DECODE_U32: [u32; 256] = {
5859    let mut tab = [0u32; 256];
5860    let mut b = 0usize;
5861    while b < 256 {
5862        tab[b] = ((b as u32) & 3)
5863            | ((((b as u32) >> 2) & 3) << 8)
5864            | ((((b as u32) >> 4) & 3) << 16)
5865            | ((((b as u32) >> 6) & 3) << 24);
5866        b += 1;
5867    }
5868    tab
5869};
5870
5871/// Eight packed q2tp bytes against 32 signed activation bytes. `maddubs`
5872/// exactly computes unsigned 2-bit code × signed i8; its pair sums cannot
5873/// saturate (2 × 3 × 127 < i16::MAX), and the second madd widens to i32.
5874#[cfg(target_arch = "x86_64")]
5875#[target_feature(enable = "avx2")]
5876unsafe fn q2tp_code_dot_avx2(ch: &[u8], x: &[i8]) -> i32 {
5877    use core::arch::x86_64::*;
5878    debug_assert!(ch.len() >= Q2TP_CHUNK && x.len() >= GROUP_SIZE);
5879    let codes = _mm256_setr_epi32(
5880        Q2TP_DECODE_U32[ch[0] as usize] as i32,
5881        Q2TP_DECODE_U32[ch[1] as usize] as i32,
5882        Q2TP_DECODE_U32[ch[2] as usize] as i32,
5883        Q2TP_DECODE_U32[ch[3] as usize] as i32,
5884        Q2TP_DECODE_U32[ch[4] as usize] as i32,
5885        Q2TP_DECODE_U32[ch[5] as usize] as i32,
5886        Q2TP_DECODE_U32[ch[6] as usize] as i32,
5887        Q2TP_DECODE_U32[ch[7] as usize] as i32,
5888    );
5889    let xv = unsafe { _mm256_loadu_si256(x.as_ptr().cast()) };
5890    let pair = _mm256_maddubs_epi16(codes, xv);
5891    let quad = _mm256_madd_epi16(pair, _mm256_set1_epi16(1));
5892    let sum128 = _mm_add_epi32(
5893        _mm256_castsi256_si128(quad),
5894        _mm256_extracti128_si256(quad, 1),
5895    );
5896    let sum64 = _mm_hadd_epi32(sum128, sum128);
5897    _mm_cvtsi128_si32(_mm_hadd_epi32(sum64, sum64))
5898}
5899
5900#[cfg(target_arch = "x86_64")]
5901#[inline]
5902fn q2tp_avx2_enabled() -> bool {
5903    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
5904    *ON.get_or_init(|| std::arch::is_x86_feature_detected!("avx2"))
5905}
5906
5907/// Integer dot of one q2tp row against pre-quantized activations:
5908/// Σ_g s_g · (Σ c·xq − 1.5·Σ xq). The half-integer grid (c − 1.5)
5909/// becomes exact integer math through the group sums — the same trick
5910/// every a8w8 kernel in this file rides. The codes decode into a
5911/// 32-byte scratch in natural order and the dot itself is the shared
5912/// SDOT primitive; elsewhere a scalar integer loop.
5913#[inline]
5914fn dot_q2tp_row_i8(
5915    chunks: &[u8],
5916    r: usize,
5917    gpr: usize,
5918    xq: &[i8],
5919    gsum: &[i32],
5920    scales: &[f32],
5921) -> f32 {
5922    let mut acc = 0f32;
5923    let base = r * gpr * Q2TP_CHUNK;
5924    #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
5925    let mut codes = [0i8; GROUP_SIZE];
5926    #[cfg(target_arch = "x86_64")]
5927    // Cache the process-wide AVX2 decision; probing it for every 32-weight
5928    // group adds a branch to the hottest Q2TP row loop while retaining the
5929    // established table-load AVX2 arithmetic (and without coupling this path
5930    // to the separate FMA-gated A8W8 switch).
5931    let avx2 = q2tp_avx2_enabled();
5932    for gi in 0..gpr {
5933        let ch = &chunks[base + gi * Q2TP_CHUNK..base + (gi + 1) * Q2TP_CHUNK];
5934        let xg = &xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
5935        #[cfg(target_arch = "aarch64")]
5936        // NEON: the byte's four 2-bit fields land in four lane vectors
5937        // (shift+mask), vld4 de-interleaves xq to match (xj[k] =
5938        // xq[4k+j]), widening MACs accumulate exactly in i32. A scalar
5939        // decode here cost as much as the dot it fed — the profile put
5940        // it at the top of the whole W2 decode.
5941        let dot = unsafe {
5942            use core::arch::aarch64::*;
5943            let b = vld1_u8(ch.as_ptr());
5944            let three = vdup_n_u8(3);
5945            let c0 = vreinterpret_s8_u8(vand_u8(b, three));
5946            let c1 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 2), three));
5947            let c2 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 4), three));
5948            let c3 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 6), three));
5949            let x4 = vld4_s8(xg.as_ptr());
5950            let mut acc4 = vdupq_n_s32(0);
5951            acc4 = vpadalq_s16(acc4, vmull_s8(c0, x4.0));
5952            acc4 = vpadalq_s16(acc4, vmull_s8(c1, x4.1));
5953            acc4 = vpadalq_s16(acc4, vmull_s8(c2, x4.2));
5954            acc4 = vpadalq_s16(acc4, vmull_s8(c3, x4.3));
5955            vaddvq_s32(acc4)
5956        };
5957        #[cfg(target_arch = "x86_64")]
5958        let dot: i32 = if avx2 {
5959            // SAFETY: the runtime feature check gates the target-feature body;
5960            // the group slices above are exactly 8 and 32 bytes long.
5961            unsafe { q2tp_code_dot_avx2(ch, xg) }
5962        } else {
5963            ch.iter()
5964                .enumerate()
5965                .map(|(k, &b)| {
5966                    ((b & 3) as i32) * xg[k * 4] as i32
5967                        + (((b >> 2) & 3) as i32) * xg[k * 4 + 1] as i32
5968                        + (((b >> 4) & 3) as i32) * xg[k * 4 + 2] as i32
5969                        + (((b >> 6) & 3) as i32) * xg[k * 4 + 3] as i32
5970                })
5971                .sum()
5972        };
5973        #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
5974        let dot: i32 = {
5975            for (k, &b) in ch.iter().enumerate() {
5976                codes[k * 4] = (b & 3) as i8;
5977                codes[k * 4 + 1] = ((b >> 2) & 3) as i8;
5978                codes[k * 4 + 2] = ((b >> 4) & 3) as i8;
5979                codes[k * 4 + 3] = ((b >> 6) & 3) as i8;
5980            }
5981            codes
5982                .iter()
5983                .zip(xg)
5984                .map(|(&c, &x)| c as i32 * x as i32)
5985                .sum()
5986        };
5987        acc += scales[gi] * (dot as f32 - 1.5 * gsum[gi] as f32);
5988    }
5989    acc
5990}
5991
5992/// Exact f32 dot of one q2tp row: 2-bit fields LSB-first, (c − 1.5)·s.
5993/// Scalar on purpose — the 2-bit class targets the GPU graph; the CPU
5994/// path exists for parity gates and small-machine fallback.
5995fn q2tp_row_exact(chunks: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
5996    q2tp_row_exact_center(chunks, r, gpr, x, scales, 1.5)
5997}
5998
5999/// Fused Prism affine row: the derived correction is applied inside the
6000/// decoded code, avoiding a second accumulated dot and avoiding cancellation
6001/// between `B=(c-1.5)s` and `+.5s` for long 5120/17408 rows.
6002#[inline]
6003fn q2tp_affine_row_exact(chunks: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
6004    q2tp_row_exact_center(chunks, r, gpr, x, scales, 1.0)
6005}
6006
6007#[inline]
6008fn q2tp_row_exact_center(
6009    chunks: &[u8],
6010    r: usize,
6011    gpr: usize,
6012    x: &[f32],
6013    scales: &[f32],
6014    center: f32,
6015) -> f32 {
6016    let mut acc = 0f32;
6017    for gi in 0..gpr {
6018        let ch = &chunks[(r * gpr + gi) * Q2TP_CHUNK..(r * gpr + gi + 1) * Q2TP_CHUNK];
6019        let s = scales[gi];
6020        let xb = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
6021        let mut g = 0f32;
6022        for (k, &b) in ch.iter().enumerate() {
6023            g += ((b & 3) as f32 - center) * xb[k * 4]
6024                + (((b >> 2) & 3) as f32 - center) * xb[k * 4 + 1]
6025                + (((b >> 4) & 3) as f32 - center) * xb[k * 4 + 2]
6026                + (((b >> 6) & 3) as f32 - center) * xb[k * 4 + 3];
6027        }
6028        acc += s * g;
6029    }
6030    acc
6031}
6032
6033fn q2tp_matvec(
6034    bytes: &[u8],
6035    x: &[f32],
6036    rows: usize,
6037    cols: usize,
6038    out: &mut [f32],
6039    pool: Option<&Pool>,
6040) {
6041    q2tp_matvec_mode(bytes, x, rows, cols, out, pool, false);
6042}
6043
6044fn q2tp_affine_matvec(
6045    bytes: &[u8],
6046    x: &[f32],
6047    rows: usize,
6048    cols: usize,
6049    out: &mut [f32],
6050    pool: Option<&Pool>,
6051) {
6052    q2tp_matvec_mode(bytes, x, rows, cols, out, pool, true);
6053}
6054
6055fn q2tp_matvec_mode(
6056    bytes: &[u8],
6057    x: &[f32],
6058    rows: usize,
6059    cols: usize,
6060    out: &mut [f32],
6061    pool: Option<&Pool>,
6062    affine: bool,
6063) {
6064    debug_assert_eq!(out.len(), rows);
6065    let gpr = cols / GROUP_SIZE;
6066    let v = Q4tpView::new_q2(bytes, rows, cols);
6067    let out_addr = SendMut(out.as_mut_ptr());
6068    // a8w8 fast path (CMF_SDOT=0 keeps the exact scalar walk): integer
6069    // code dots + group sums, exact outlier correction — the same
6070    // contract as every sibling kernel; measured 2-bit rows were the
6071    // only scalar holdout in the family.
6072    if !affine && a8w8_enabled() {
6073        let act = split_act(x);
6074        let gsum = q1_group_sums(&act.xq, gpr);
6075        let (act, gsum) = (&act, &gsum);
6076        let run = move |start: usize, end: usize| {
6077            with_krow(gpr, |sc| {
6078                for r in start..end {
6079                    v.scales_into(r, gpr, sc);
6080                    let mut acc = dot_q2tp_row_i8(v.nib, r, gpr, &act.xq, gsum, sc) * act.sx;
6081                    for &(j, xv) in &act.outliers {
6082                        let (w, s) = q2tp_outlier(v.nib, r, gpr, j, sc);
6083                        acc += w * s * xv;
6084                    }
6085                    // SAFETY: disjoint row ranges per worker.
6086                    unsafe { *out_addr.at(r) = acc };
6087                }
6088            })
6089        };
6090        dispatch_rows(pool, rows, &run);
6091        return;
6092    }
6093    let run = |start: usize, end: usize| {
6094        with_krow(gpr, |sc| {
6095            for r in start..end {
6096                v.scales_into(r, gpr, sc);
6097                // SAFETY: disjoint row ranges per worker.
6098                unsafe {
6099                    *out_addr.at(r) = if affine {
6100                        q2tp_affine_row_exact(v.nib, r, gpr, x, sc)
6101                    } else {
6102                        q2tp_row_exact(v.nib, r, gpr, x, sc)
6103                    }
6104                };
6105            }
6106        })
6107    };
6108    dispatch_rows(pool, rows, &run);
6109}
6110
6111/// Fused two-input q2tp matvec — the SwiGLU gate/up pair.
6112#[allow(clippy::too_many_arguments)]
6113fn q2tp_matvec2(
6114    bytes: &[u8],
6115    x1: &[f32],
6116    x2: &[f32],
6117    rows: usize,
6118    cols: usize,
6119    o1: &mut [f32],
6120    o2: &mut [f32],
6121    pool: Option<&Pool>,
6122) {
6123    q2tp_matvec2_mode(bytes, x1, x2, rows, cols, o1, o2, pool, false);
6124}
6125
6126#[allow(clippy::too_many_arguments)]
6127fn q2tp_affine_matvec2(
6128    bytes: &[u8],
6129    x1: &[f32],
6130    x2: &[f32],
6131    rows: usize,
6132    cols: usize,
6133    o1: &mut [f32],
6134    o2: &mut [f32],
6135    pool: Option<&Pool>,
6136) {
6137    q2tp_matvec2_mode(bytes, x1, x2, rows, cols, o1, o2, pool, true);
6138}
6139
6140#[allow(clippy::too_many_arguments)]
6141fn q2tp_matvec2_mode(
6142    bytes: &[u8],
6143    x1: &[f32],
6144    x2: &[f32],
6145    rows: usize,
6146    cols: usize,
6147    o1: &mut [f32],
6148    o2: &mut [f32],
6149    pool: Option<&Pool>,
6150    affine: bool,
6151) {
6152    let gpr = cols / GROUP_SIZE;
6153    let v = Q4tpView::new_q2(bytes, rows, cols);
6154    let (p1, p2) = (SendMut(o1.as_mut_ptr()), SendMut(o2.as_mut_ptr()));
6155    let run = |start: usize, end: usize| {
6156        let mut sc = vec![0f32; gpr];
6157        for r in start..end {
6158            v.scales_into(r, gpr, &mut sc);
6159            // SAFETY: disjoint row ranges per worker.
6160            unsafe {
6161                *p1.at(r) = if affine {
6162                    q2tp_affine_row_exact(v.nib, r, gpr, x1, &sc)
6163                } else {
6164                    q2tp_row_exact(v.nib, r, gpr, x1, &sc)
6165                };
6166                *p2.at(r) = if affine {
6167                    q2tp_affine_row_exact(v.nib, r, gpr, x2, &sc)
6168                } else {
6169                    q2tp_row_exact(v.nib, r, gpr, x2, &sc)
6170                };
6171            }
6172        }
6173    };
6174    dispatch_rows(pool, rows, &run);
6175}
6176
6177/// Batched q2tp matmat: scalar row kernel over every batch column. CPU
6178/// prefill only — decode rides the graph, so plain and correct beats
6179/// clever here.
6180/// Test doors into the host 2-bit kernels: the stand's heap corruption
6181/// pointed at down-shaped tensors, and the private fns need a way to be
6182/// held to a reference without a model file around them.
6183pub fn q2tp_matvec_for_test(bytes: &[u8], x: &[f32], rows: usize, cols: usize, out: &mut [f32]) {
6184    // The facade IS the reference: encoder oracles hold requant output
6185    // to the exact scalar walk. The production dispatch may take the i8
6186    // fast path, whose error scale is the ACTIVATIONS' — a different
6187    // claim than the encoder correctness these tests pin.
6188    let gpr = cols / GROUP_SIZE;
6189    let v = Q4tpView::new_q2(bytes, rows, cols);
6190    with_krow(gpr, |sc| {
6191        for r in 0..rows {
6192            v.scales_into(r, gpr, sc);
6193            out[r] = q2tp_row_exact(v.nib, r, gpr, x, sc);
6194        }
6195    });
6196}
6197
6198/// Test door for the descriptor-specific fused affine decode.  Production
6199/// callers select this through a validated Prism header, never by dtype alone.
6200pub fn q2tp_affine_matvec_for_test(
6201    bytes: &[u8],
6202    x: &[f32],
6203    rows: usize,
6204    cols: usize,
6205    out: &mut [f32],
6206) {
6207    q2tp_affine_matvec(bytes, x, rows, cols, out, None);
6208}
6209
6210pub fn q2tp_matmat_for_test(
6211    bytes: &[u8],
6212    xs_all: &[f32],
6213    b: usize,
6214    rows: usize,
6215    cols: usize,
6216    out: &mut [f32],
6217) {
6218    q2tp_matmat(bytes, xs_all, b, rows, cols, out, None);
6219}
6220
6221fn q2tp_matmat(
6222    bytes: &[u8],
6223    xs_all: &[f32],
6224    b: usize,
6225    rows: usize,
6226    cols: usize,
6227    out: &mut [f32],
6228    pool: Option<&Pool>,
6229) {
6230    q2tp_matmat_mode(bytes, xs_all, b, rows, cols, out, pool, false);
6231}
6232
6233fn q2tp_affine_matmat(
6234    bytes: &[u8],
6235    xs_all: &[f32],
6236    b: usize,
6237    rows: usize,
6238    cols: usize,
6239    out: &mut [f32],
6240    pool: Option<&Pool>,
6241) {
6242    q2tp_matmat_mode(bytes, xs_all, b, rows, cols, out, pool, true);
6243}
6244
6245fn q2tp_matmat_mode(
6246    bytes: &[u8],
6247    xs_all: &[f32],
6248    b: usize,
6249    rows: usize,
6250    cols: usize,
6251    out: &mut [f32],
6252    pool: Option<&Pool>,
6253    affine: bool,
6254) {
6255    debug_assert_eq!(out.len(), b * rows);
6256    let gpr = cols / GROUP_SIZE;
6257    let v = Q4tpView::new_q2(bytes, rows, cols);
6258    let out_addr = SendMut(out.as_mut_ptr());
6259    let run = |start: usize, end: usize| {
6260        let mut sc = vec![0f32; gpr];
6261        for r in start..end {
6262            v.scales_into(r, gpr, &mut sc);
6263            for bi in 0..b {
6264                let x = &xs_all[bi * cols..(bi + 1) * cols];
6265                // SAFETY: disjoint row ranges per worker.
6266                unsafe {
6267                    *out_addr.at(bi * rows + r) = if affine {
6268                        q2tp_affine_row_exact(v.nib, r, gpr, x, &sc)
6269                    } else {
6270                        q2tp_row_exact(v.nib, r, gpr, x, &sc)
6271                    }
6272                };
6273            }
6274        }
6275    };
6276    dispatch_rows(pool, rows, &run);
6277}
6278
6279/// The pre-vectorised shape, kept for A/B (`CMF_Q4TP_V1=1`): the
6280/// horizontal add lands once per group per column instead of once per
6281/// row. Same weights, same activations — only the reduction differs.
6282/// It is also the row-exact batch kernel: per group and column it forms
6283/// `int dot as f32 * scale` and adds it to a scalar running sum, which is
6284/// `dot_q4tp_row_sdot` step for step (Rust never contracts to an fma),
6285/// so each column is bit-identical to that column's matvec.
6286#[cfg(target_arch = "aarch64")]
6287#[target_feature(enable = "neon,dotprod")]
6288unsafe fn dot_q4tp_row_1x4_sdot_v1(
6289    nib: &[u8],
6290    r: usize,
6291    gpr: usize,
6292    xs: [&[i8]; 4],
6293    scales: &[f32],
6294) -> [f32; 4] {
6295    unsafe {
6296        use core::arch::aarch64::*;
6297        use core::arch::asm;
6298        let lomask = vdupq_n_u8(0x0F);
6299        let eight = vdupq_n_s8(8);
6300        let (mut f0, mut f1, mut f2, mut f3) = (0f32, 0f32, 0f32, 0f32);
6301        for gi in 0..gpr {
6302            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
6303            let s = *scales.get_unchecked(gi);
6304            let bb = vld1q_u8(t);
6305            let lo = vandq_u8(bb, lomask);
6306            let hi = vshrq_n_u8::<4>(bb);
6307            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
6308            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
6309            let mut d = [0f32; 4];
6310            for (k, dk) in d.iter_mut().enumerate() {
6311                let x0 = vld1q_s8(xs[k].as_ptr().add(gi * GROUP_SIZE));
6312                let x1 = vld1q_s8(xs[k].as_ptr().add(gi * GROUP_SIZE + 16));
6313                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6314                asm!(
6315                    "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
6316                    "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
6317                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6318                    e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
6319                    options(pure, nomem, nostack),
6320                );
6321                *dk = vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
6322            }
6323            f0 += d[0];
6324            f1 += d[1];
6325            f2 += d[2];
6326            f3 += d[3];
6327        }
6328        [f0, f1, f2, f3]
6329    }
6330}
6331
6332/// Which q4tp batch kernel to run: 1 = the previous one, 2 = the tuned
6333/// one, 0 = decide from the CPU. An atomic rather than a `OnceLock` so a
6334/// benchmark can alternate the two inside one process, where the machine's
6335/// mood — a shared box drifts ±25% between runs — is the same for both.
6336/// What the two mean is per-architecture: on x86 the blocked AVX-512 path
6337/// against the per-column one, on ARM the two reduction shapes.
6338#[allow(dead_code)]
6339static Q4TP_ALT: std::sync::atomic::AtomicU8 = std::sync::atomic::AtomicU8::new(0);
6340
6341/// Tests that store `Q4TP_ALT` hold this, so one test's kernel pick does
6342/// not leak into another's bit-exact comparison running in parallel.
6343#[cfg(test)]
6344static Q4TP_ALT_TEST_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
6345
6346/// Blocking pays on x86 only with 512-bit VNNI. With AVX2 alone, four
6347/// columns sharing an unpack still measured slower than the per-column
6348/// path (23.2 ms against 19.4 on a 48-thread EPYC), because that path
6349/// already dequantizes the row once — so the blocked kernel bought a
6350/// second unpack-free pass at the price of half the vector width.
6351#[cfg(target_arch = "x86_64")]
6352fn q4tp_blocked_x86() -> bool {
6353    match Q4TP_ALT.load(std::sync::atomic::Ordering::Relaxed) {
6354        1 => false,
6355        // A forced ON still asks the CPU. The switch exists so a bench can
6356        // pick a kernel, not so it can promise instructions the machine
6357        // does not have — CI caught that as a SIGILL on a runner without
6358        // AVX-512, where the parity test had turned the path on by hand.
6359        2 => avx512vnni_enabled(),
6360        // Deliberately not cached back into the switch: both gates below
6361        // hold their own `OnceLock`, and latching their answer here would
6362        // make a test's override outlive the test that set it.
6363        _ => blocked_enabled() && avx512vnni_enabled(),
6364    }
6365}
6366
6367/// `CMF_Q4TP_V1=1` picks the old reduction shape (A/B only).
6368#[cfg(target_arch = "aarch64")]
6369#[allow(dead_code)]
6370fn q4tp_v1() -> bool {
6371    match Q4TP_ALT.load(std::sync::atomic::Ordering::Relaxed) {
6372        1 => true,
6373        2 => false,
6374        _ => {
6375            static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6376            *ON.get_or_init(|| std::env::var("CMF_Q4TP_V1").is_ok_and(|v| v != "0"))
6377        }
6378    }
6379}
6380
6381/// Two weight rows against eight columns. The activation load is the
6382/// same for both rows, so it is paid once for twice the arithmetic, and
6383/// sixteen accumulator chains run where eight did — which is what a kernel
6384/// retiring 0.29 instructions a cycle is short of. Register pressure is
6385/// the limit: sixteen `zmm` accumulators, two weight tiles, one
6386/// activation, of thirty-two.
6387///
6388/// Four rows by four columns spends the same sixteen accumulators the
6389/// other way and measured worse — 1488 GFLOP/s against 1644 — so the
6390/// unpack, which four rows pay twice as often, costs more than the extra
6391/// sharing of one activation load buys.
6392#[cfg(target_arch = "x86_64")]
6393#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
6394unsafe fn dot_q4tp_2x8_avx512(
6395    nib: &[u8],
6396    r0: usize,
6397    gpr: usize,
6398    xs: [&[i8]; 8],
6399    sc0: &[f32],
6400    sc1: &[f32],
6401) -> [[f32; 8]; 2] {
6402    // SAFETY: as dot_q4tp_row_1x8_avx512, two adjacent rows at once; the
6403    // caller guarantees r0 + 1 < rows and the ISA.
6404    unsafe {
6405        use core::arch::x86_64::*;
6406        let lomask = _mm256_set1_epi8(0x0F);
6407        let eight = _mm256_set1_epi8(8);
6408        let zero = _mm512_setzero_si512();
6409        let mut v0 = [_mm512_setzero_ps(); 8];
6410        let mut v1 = [_mm512_setzero_ps(); 8];
6411        let pairs = gpr / 2;
6412        let unpack = |r: usize, gi: usize| -> (__m512i, __mmask64) {
6413            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
6414            let bb = _mm256_loadu_si256(t as *const __m256i);
6415            let lo = _mm256_and_si256(bb, lomask);
6416            let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
6417            let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
6418            let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
6419            let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
6420            let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
6421            (_mm512_abs_epi8(w), _mm512_movepi8_mask(w))
6422        };
6423        for gp in 0..pairs {
6424            let gi = gp * 2;
6425            let (wa0, neg0) = unpack(r0, gi);
6426            let (wa1, neg1) = unpack(r0 + 1, gi);
6427            let off = gi * GROUP_SIZE;
6428            let sv = |sc: &[f32]| {
6429                _mm512_insertf32x8::<1>(
6430                    _mm512_castps256_ps512(_mm256_set1_ps(*sc.get_unchecked(gi))),
6431                    _mm256_set1_ps(*sc.get_unchecked(gi + 1)),
6432                )
6433            };
6434            let s0 = sv(sc0);
6435            let s1 = sv(sc1);
6436            for k in 0..8 {
6437                let xv = _mm512_loadu_si512(xs[k].as_ptr().add(off) as *const __m512i);
6438                let d0 = _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(
6439                    zero,
6440                    wa0,
6441                    _mm512_mask_sub_epi8(xv, neg0, zero, xv),
6442                ));
6443                let d1 = _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(
6444                    zero,
6445                    wa1,
6446                    _mm512_mask_sub_epi8(xv, neg1, zero, xv),
6447                ));
6448                v0[k] = _mm512_fmadd_ps(d0, s0, v0[k]);
6449                v1[k] = _mm512_fmadd_ps(d1, s1, v1[k]);
6450            }
6451        }
6452        let mut acc = [[0f32; 8]; 2];
6453        for k in 0..8 {
6454            acc[0][k] = _mm512_reduce_add_ps(v0[k]);
6455            acc[1][k] = _mm512_reduce_add_ps(v1[k]);
6456        }
6457        if gpr % 2 == 1 {
6458            let off = (gpr - 1) * GROUP_SIZE;
6459            for j in off..off + GROUP_SIZE {
6460                let (w0, sa) = q4tp_outlier(nib, r0, gpr, j, sc0);
6461                let (w1, sb) = q4tp_outlier(nib, r0 + 1, gpr, j, sc1);
6462                for k in 0..8 {
6463                    let x = *xs[k].get_unchecked(j) as f32;
6464                    acc[0][k] += w0 * sa * x;
6465                    acc[1][k] += w1 * sb * x;
6466                }
6467            }
6468        }
6469        acc
6470    }
6471}
6472
6473/// The same, eight columns at a time. One unpack then feeds twice as many
6474/// activation streams, so a wide batch reads the weight tile half as
6475/// often; the price is eight accumulators live at once. Measured 9.0 ->
6476/// 8.3 ms at 9216x2304, b=296 on a 48-thread EPYC 9B45.
6477#[cfg(target_arch = "x86_64")]
6478#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
6479unsafe fn dot_q4tp_row_1x8_avx512(
6480    nib: &[u8],
6481    r: usize,
6482    gpr: usize,
6483    xs: [&[i8]; 8],
6484    scales: &[f32],
6485) -> [f32; 8] {
6486    // SAFETY: as dot_q4tp_row_1x4_avx2; caller guarantees the ISA.
6487    unsafe {
6488        use core::arch::x86_64::*;
6489        let lomask = _mm256_set1_epi8(0x0F);
6490        let eight = _mm256_set1_epi8(8);
6491        let zero = _mm512_setzero_si512();
6492        let (mut v0, mut v1, mut v2, mut v3) = (
6493            _mm512_setzero_ps(),
6494            _mm512_setzero_ps(),
6495            _mm512_setzero_ps(),
6496            _mm512_setzero_ps(),
6497        );
6498        let (mut v4, mut v5, mut v6, mut v7) = (
6499            _mm512_setzero_ps(),
6500            _mm512_setzero_ps(),
6501            _mm512_setzero_ps(),
6502            _mm512_setzero_ps(),
6503        );
6504        let pairs = gpr / 2;
6505        for gp in 0..pairs {
6506            let gi = gp * 2;
6507            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
6508            let bb = _mm256_loadu_si256(t as *const __m256i);
6509            let lo = _mm256_and_si256(bb, lomask);
6510            let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
6511            // `unpack` works per 128-bit lane, so the halves come out as
6512            // [A.lo, B.lo] and [A.hi, B.hi]; the shuffle reorders the four
6513            // 128-bit lanes into the weights' natural order, which is what
6514            // the straight activation load expects.
6515            let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
6516            let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
6517            let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
6518            let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
6519            let wabs = _mm512_abs_epi8(w);
6520            let neg = _mm512_movepi8_mask(w);
6521            let off = gi * GROUP_SIZE;
6522            let sv = _mm512_insertf32x8::<1>(
6523                _mm512_castps256_ps512(_mm256_set1_ps(*scales.get_unchecked(gi))),
6524                _mm256_set1_ps(*scales.get_unchecked(gi + 1)),
6525            );
6526            let dot = |x: &[i8]| -> __m512 {
6527                let xv = _mm512_loadu_si512(x.as_ptr().add(off) as *const __m512i);
6528                let sx = _mm512_mask_sub_epi8(xv, neg, zero, xv);
6529                _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(zero, wabs, sx))
6530            };
6531            v0 = _mm512_fmadd_ps(dot(xs[0]), sv, v0);
6532            v1 = _mm512_fmadd_ps(dot(xs[1]), sv, v1);
6533            v2 = _mm512_fmadd_ps(dot(xs[2]), sv, v2);
6534            v3 = _mm512_fmadd_ps(dot(xs[3]), sv, v3);
6535            v4 = _mm512_fmadd_ps(dot(xs[4]), sv, v4);
6536            v5 = _mm512_fmadd_ps(dot(xs[5]), sv, v5);
6537            v6 = _mm512_fmadd_ps(dot(xs[6]), sv, v6);
6538            v7 = _mm512_fmadd_ps(dot(xs[7]), sv, v7);
6539        }
6540        let mut acc = [
6541            _mm512_reduce_add_ps(v0),
6542            _mm512_reduce_add_ps(v1),
6543            _mm512_reduce_add_ps(v2),
6544            _mm512_reduce_add_ps(v3),
6545            _mm512_reduce_add_ps(v4),
6546            _mm512_reduce_add_ps(v5),
6547            _mm512_reduce_add_ps(v6),
6548            _mm512_reduce_add_ps(v7),
6549        ];
6550        // An odd group count leaves one group over; the narrow kernel
6551        // finishes it rather than the tail being a special case here.
6552        if gpr % 2 == 1 {
6553            let off = (gpr - 1) * GROUP_SIZE;
6554            for j in off..off + GROUP_SIZE {
6555                let (w, s) = q4tp_outlier(nib, r, gpr, j, scales);
6556                let ws = w * s;
6557                for k in 0..8 {
6558                    acc[k] += ws * *xs[k].get_unchecked(j) as f32;
6559                }
6560            }
6561        }
6562        acc
6563    }
6564}
6565
6566/// The same four columns, 512 bits wide. Two groups (64 weights) ride one
6567/// unpack and one `vpdpbusd`, where AVX2 needs two unpacks and four
6568/// `maddubs`/`madd` pairs — about 2.3x fewer instructions for the same
6569/// arithmetic. The two groups carry different scales, so the fma takes a
6570/// vector whose halves hold each group's scale rather than a broadcast.
6571///
6572/// There is no 512-bit `vpsignb`, so the activation's sign is applied by
6573/// negating under a mask taken from the weight's sign bits. That mask is
6574/// per-tile, so it is hoisted out of the column loop and the per-column
6575/// cost stays exactly one instruction, as with `sign_epi8`. Weights of
6576/// zero are not zeroed by the mask trick and do not need to be: their
6577/// magnitude is zero, so the product is.
6578#[cfg(target_arch = "x86_64")]
6579#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
6580unsafe fn dot_q4tp_row_1x4_avx512(
6581    nib: &[u8],
6582    r: usize,
6583    gpr: usize,
6584    xs: [&[i8]; 4],
6585    scales: &[f32],
6586) -> [f32; 4] {
6587    // SAFETY: as dot_q4tp_row_1x4_avx2; caller guarantees the ISA.
6588    unsafe {
6589        use core::arch::x86_64::*;
6590        let lomask = _mm256_set1_epi8(0x0F);
6591        let eight = _mm256_set1_epi8(8);
6592        let zero = _mm512_setzero_si512();
6593        let (mut v0, mut v1, mut v2, mut v3) = (
6594            _mm512_setzero_ps(),
6595            _mm512_setzero_ps(),
6596            _mm512_setzero_ps(),
6597            _mm512_setzero_ps(),
6598        );
6599        let pairs = gpr / 2;
6600        for gp in 0..pairs {
6601            let gi = gp * 2;
6602            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
6603            let bb = _mm256_loadu_si256(t as *const __m256i);
6604            let lo = _mm256_and_si256(bb, lomask);
6605            let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
6606            // `unpack` works per 128-bit lane, so the halves come out as
6607            // [A.lo, B.lo] and [A.hi, B.hi]; the shuffle reorders the four
6608            // 128-bit lanes into the weights' natural order, which is what
6609            // the straight activation load expects.
6610            let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
6611            let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
6612            let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
6613            let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
6614            let wabs = _mm512_abs_epi8(w);
6615            let neg = _mm512_movepi8_mask(w);
6616            let off = gi * GROUP_SIZE;
6617            let sv = _mm512_insertf32x8::<1>(
6618                _mm512_castps256_ps512(_mm256_set1_ps(*scales.get_unchecked(gi))),
6619                _mm256_set1_ps(*scales.get_unchecked(gi + 1)),
6620            );
6621            let dot = |x: &[i8]| -> __m512 {
6622                let xv = _mm512_loadu_si512(x.as_ptr().add(off) as *const __m512i);
6623                let sx = _mm512_mask_sub_epi8(xv, neg, zero, xv);
6624                _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(zero, wabs, sx))
6625            };
6626            v0 = _mm512_fmadd_ps(dot(xs[0]), sv, v0);
6627            v1 = _mm512_fmadd_ps(dot(xs[1]), sv, v1);
6628            v2 = _mm512_fmadd_ps(dot(xs[2]), sv, v2);
6629            v3 = _mm512_fmadd_ps(dot(xs[3]), sv, v3);
6630        }
6631        let mut acc = [
6632            _mm512_reduce_add_ps(v0),
6633            _mm512_reduce_add_ps(v1),
6634            _mm512_reduce_add_ps(v2),
6635            _mm512_reduce_add_ps(v3),
6636        ];
6637        // An odd group count leaves one group over; the narrow kernel
6638        // finishes it rather than the tail being a special case here.
6639        if gpr % 2 == 1 {
6640            let off = (gpr - 1) * GROUP_SIZE;
6641            for j in off..off + GROUP_SIZE {
6642                let (w, s) = q4tp_outlier(nib, r, gpr, j, scales);
6643                let ws = w * s;
6644                for k in 0..4 {
6645                    acc[k] += ws * *xs[k].get_unchecked(j) as f32;
6646                }
6647            }
6648        }
6649        acc
6650    }
6651}
6652
6653/// Four batch columns against one q4tp row: the tile is unpacked ONCE and
6654/// spent on four activation streams, which is where a prefill batch stops
6655/// being weight-bandwidth-bound. Twin of `dot_q4t_row_1x4_sdot`.
6656#[cfg(target_arch = "aarch64")]
6657#[target_feature(enable = "neon,dotprod")]
6658unsafe fn dot_q4tp_row_1x4_sdot(
6659    nib: &[u8],
6660    r: usize,
6661    gpr: usize,
6662    xs: [&[i8]; 4],
6663    scales: &[f32],
6664) -> [f32; 4] {
6665    // SAFETY: see dot_q4tp_row_sdot; every xs[k] is gpr·GROUP_SIZE long.
6666    unsafe {
6667        use core::arch::aarch64::*;
6668        use core::arch::asm;
6669        let lomask = vdupq_n_u8(0x0F);
6670        let eight = vdupq_n_s8(8);
6671        // Named accumulators, NOT an array indexed by a loop variable: the
6672        // latter does not stay in registers (the same defect cost 2x in the
6673        // AVX2 q4t kernel and again in WGSL).
6674        //
6675        // They are VECTORS, and the horizontal add happens once at the end
6676        // instead of once per group per column. `vaddvq` is a cross-lane
6677        // reduction — with 72 groups and four columns the old shape paid
6678        // 288 of them per row, each one a dependency stall the pipeline
6679        // cannot hide, to save four float adds. The group's scale now
6680        // rides an fma into the lane accumulators, so the arithmetic per
6681        // group is one convert and one fma. Summation order changes (the
6682        // lanes carry independent partial sums), which is the same
6683        // round-off class the SDOT path already lives in — the strict
6684        // kernel (`CMF_SDOT=0`, what `cortiq ppl` runs) is unchanged and
6685        // stays the reference.
6686        let (mut v0, mut v1, mut v2, mut v3) = (
6687            vdupq_n_f32(0.0),
6688            vdupq_n_f32(0.0),
6689            vdupq_n_f32(0.0),
6690            vdupq_n_f32(0.0),
6691        );
6692        for gi in 0..gpr {
6693            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
6694            let s = *scales.get_unchecked(gi);
6695            let bb = vld1q_u8(t);
6696            let lo = vandq_u8(bb, lomask);
6697            let hi = vshrq_n_u8::<4>(bb);
6698            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
6699            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
6700            let off = gi * GROUP_SIZE;
6701            let dot4 = |x: &[i8]| -> int32x4_t {
6702                let x0 = vld1q_s8(x.as_ptr().add(off));
6703                let x1 = vld1q_s8(x.as_ptr().add(off + 16));
6704                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6705                asm!(
6706                    "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
6707                    "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
6708                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6709                    e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
6710                    options(pure, nomem, nostack),
6711                );
6712                vaddq_s32(a0, a1)
6713            };
6714            v0 = vfmaq_n_f32(v0, vcvtq_f32_s32(dot4(xs[0])), s);
6715            v1 = vfmaq_n_f32(v1, vcvtq_f32_s32(dot4(xs[1])), s);
6716            v2 = vfmaq_n_f32(v2, vcvtq_f32_s32(dot4(xs[2])), s);
6717            v3 = vfmaq_n_f32(v3, vcvtq_f32_s32(dot4(xs[3])), s);
6718        }
6719        [
6720            vaddvq_f32(v0),
6721            vaddvq_f32(v1),
6722            vaddvq_f32(v2),
6723            vaddvq_f32(v3),
6724        ]
6725    }
6726}
6727
6728/// Fused q4tp matmat — the same three arms `q4t_matmat` has. Shipping only
6729/// the scalar one made Nanbeige-3B decode at 1.2 tok/s against q4t's 5.9:
6730/// the format was fine, the missing arms were the whole regression.
6731///
6732/// Under `row_exact()` every cell is summed in `q4tp_matvec`'s order on
6733/// every architecture; the mode is read once, so one call never mixes
6734/// the two contracts when a concurrent scope opens or closes mid-call.
6735fn q4tp_matmat(
6736    bytes: &[u8],
6737    xs_all: &[f32],
6738    b: usize,
6739    rows: usize,
6740    cols: usize,
6741    out: &mut [f32],
6742    pool: Option<&Pool>,
6743) {
6744    q4tp_matmat_with(bytes, xs_all, b, rows, cols, out, pool, row_exact())
6745}
6746
6747/// `q4tp_matmat` with the row-exact mode passed in rather than read from
6748/// the shared counter: `exact` = each cell equals its token's matvec bit
6749/// for bit, otherwise the fast blocked / AMX arms are free to reorder.
6750#[allow(clippy::too_many_arguments)]
6751fn q4tp_matmat_with(
6752    bytes: &[u8],
6753    xs_all: &[f32],
6754    b: usize,
6755    rows: usize,
6756    cols: usize,
6757    out: &mut [f32],
6758    pool: Option<&Pool>,
6759    exact: bool,
6760) {
6761    debug_assert_eq!(out.len(), b * rows);
6762    let gpr = cols / GROUP_SIZE;
6763    let v = Q4tpView::new(bytes, rows, cols);
6764
6765    // Wide batches ride the AMX through a dequant-tile sgemm, as in q4t.
6766    // An f32 GEMM over dequantized weights is not the matvec's int8 sum,
6767    // so a row-exact batch never takes it.
6768    #[cfg(target_os = "macos")]
6769    if !exact && b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
6770        dequant_matmat_accel(
6771            &|r, dst| {
6772                let mut sc = [0f32; 32];
6773                let mut scv;
6774                let s: &[f32] = if gpr <= 32 {
6775                    v.scales_into(r, gpr, &mut sc);
6776                    &sc[..gpr]
6777                } else {
6778                    scv = vec![0f32; gpr];
6779                    v.scales_into(r, gpr, &mut scv);
6780                    &scv
6781                };
6782                for gi in 0..gpr {
6783                    let tile = &v.nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
6784                    for (k, &bb) in tile.iter().enumerate() {
6785                        dst[gi * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s[gi];
6786                        dst[gi * GROUP_SIZE + k * 2 + 1] =
6787                            (((bb >> 4) & 0x0F) as f32 - 8.0) * s[gi];
6788                    }
6789                }
6790            },
6791            xs_all,
6792            b,
6793            rows,
6794            cols,
6795            out,
6796            pool,
6797        );
6798        return;
6799    }
6800
6801    let out_addr = SendMut(out.as_mut_ptr());
6802    if a8w8_enabled() {
6803        let acts: Vec<SplitAct> = (0..b)
6804            .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
6805            .collect();
6806        let acts = &acts;
6807        // ARM stays blocked under `exact` too: the 1x4 kernel then runs in
6808        // its v1 shape, whose per-group `int dot as f32 * scale` and scalar
6809        // running sum are exactly `dot_q4tp_row_sdot`'s, so the tile is
6810        // still unpacked once for four columns and each column equals its
6811        // matvec. The tuned shape (fma into lane partials, one horizontal
6812        // add per row) is 1-16 ulp off the matvec and is kept for `!exact`.
6813        #[cfg(target_arch = "aarch64")]
6814        let blocked_ok = sdot_enabled() && blocked_enabled();
6815        // x86 gets the same blocking: one tile unpack spent on four
6816        // columns. Without it every column re-decoded the row, which is
6817        // why a 48-core EPYC measured a sixth of an M4's per-core rate.
6818        // The gate is `avx2_enabled`, as in q4t — `sdot_enabled` answers
6819        // for ARM's dotprod and is hard-wired false everywhere else, so
6820        // asking it here left the whole blocked path unreachable on x86.
6821        #[cfg(target_arch = "x86_64")]
6822        let blocked_ok = q4tp_blocked_x86() && !exact;
6823        #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
6824        let blocked_ok = {
6825            let _ = exact;
6826            false
6827        };
6828        // Columns are swept in panels that fit L2. Without this a
6829        // row-pair walks every activation in the batch — 4.8 MB at
6830        // 512x512 — and does it again for the next pair, so the whole
6831        // batch streams out of the shared cache once per row. Measured
6832        // 800 GB/s of it, flat across batch sizes, which is the signature
6833        // of a loop bound by traffic rather than by arithmetic. A panel of
6834        // 256 columns is 590 KB beside 221 KB of this worker's weights:
6835        // both stay resident and the batch crosses L3 once instead of
6836        // once per row.
6837        let panel_cols: usize = std::env::var("CMF_Q4TP_PANEL")
6838            .ok()
6839            .and_then(|v| v.parse().ok())
6840            .filter(|v| *v > 0)
6841            .unwrap_or(256);
6842        let run = |start: usize, end: usize| {
6843            for abase in (0..acts.len()).step_by(panel_cols) {
6844                let alen = (acts.len() - abase).min(panel_cols);
6845                let mut sc = vec![0f32; gpr];
6846                #[cfg(target_arch = "x86_64")]
6847                let mut r_lo = start;
6848                #[cfg(target_arch = "x86_64")]
6849                if blocked_ok && alen >= 8 {
6850                    let mut sc1 = vec![0f32; gpr];
6851                    while r_lo + 2 <= end {
6852                        v.scales_into(r_lo, gpr, &mut sc);
6853                        v.scales_into(r_lo + 1, gpr, &mut sc1);
6854                        let mut bi = 0usize;
6855                        while bi + 8 <= alen {
6856                            let xs = [
6857                                acts[abase + bi].xq.as_slice(),
6858                                acts[abase + bi + 1].xq.as_slice(),
6859                                acts[abase + bi + 2].xq.as_slice(),
6860                                acts[abase + bi + 3].xq.as_slice(),
6861                                acts[abase + bi + 4].xq.as_slice(),
6862                                acts[abase + bi + 5].xq.as_slice(),
6863                                acts[abase + bi + 6].xq.as_slice(),
6864                                acts[abase + bi + 7].xq.as_slice(),
6865                            ];
6866                            let d = unsafe { dot_q4tp_2x8_avx512(v.nib, r_lo, gpr, xs, &sc, &sc1) };
6867                            for (row, dr, scr) in [(r_lo, &d[0], &sc), (r_lo + 1, &d[1], &sc1)] {
6868                                for k in 0..8 {
6869                                    let act = &acts[abase + bi + k];
6870                                    let mut acc = dr[k] * act.sx;
6871                                    for &(j, xv) in &act.outliers {
6872                                        let (w, s) = q4tp_outlier(v.nib, row, gpr, j, scr);
6873                                        acc += w * s * xv;
6874                                    }
6875                                    // SAFETY: disjoint (bi, r) cells per worker.
6876                                    unsafe { *out_addr.at((abase + bi + k) * rows + row) = acc };
6877                                }
6878                            }
6879                            bi += 8;
6880                        }
6881                        // Columns past the last group of eight, both rows —
6882                        // the same single-row kernel the tail below uses.
6883                        for row in [r_lo, r_lo + 1] {
6884                            let scr: &[f32] = if row == r_lo { &sc } else { &sc1 };
6885                            for b2 in bi..alen {
6886                                let act = &acts[abase + b2];
6887                                let xs4 = [
6888                                    act.xq.as_slice(),
6889                                    act.xq.as_slice(),
6890                                    act.xq.as_slice(),
6891                                    act.xq.as_slice(),
6892                                ];
6893                                let d =
6894                                    unsafe { dot_q4tp_row_1x4_avx512(v.nib, row, gpr, xs4, scr) };
6895                                let mut acc = d[0] * act.sx;
6896                                for &(j, xv) in &act.outliers {
6897                                    let (w, s) = q4tp_outlier(v.nib, row, gpr, j, scr);
6898                                    acc += w * s * xv;
6899                                }
6900                                // SAFETY: disjoint (bi, r) cells per worker.
6901                                unsafe { *out_addr.at((abase + b2) * rows + row) = acc };
6902                            }
6903                        }
6904                        r_lo += 2;
6905                    }
6906                }
6907                #[cfg(target_arch = "x86_64")]
6908                let row_start = r_lo;
6909                #[cfg(not(target_arch = "x86_64"))]
6910                let row_start = start;
6911                for r in row_start..end {
6912                    v.scales_into(r, gpr, &mut sc);
6913                    let mut bi = 0usize;
6914                    #[cfg(target_arch = "x86_64")]
6915                    if blocked_ok {
6916                        while bi + 8 <= alen {
6917                            let xs = [
6918                                acts[abase + bi].xq.as_slice(),
6919                                acts[abase + bi + 1].xq.as_slice(),
6920                                acts[abase + bi + 2].xq.as_slice(),
6921                                acts[abase + bi + 3].xq.as_slice(),
6922                                acts[abase + bi + 4].xq.as_slice(),
6923                                acts[abase + bi + 5].xq.as_slice(),
6924                                acts[abase + bi + 6].xq.as_slice(),
6925                                acts[abase + bi + 7].xq.as_slice(),
6926                            ];
6927                            let d = unsafe { dot_q4tp_row_1x8_avx512(v.nib, r, gpr, xs, &sc) };
6928                            for k in 0..8 {
6929                                let act = &acts[abase + bi + k];
6930                                let mut acc = d[k] * act.sx;
6931                                for &(j, xv) in &act.outliers {
6932                                    let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
6933                                    acc += w * s * xv;
6934                                }
6935                                // SAFETY: disjoint (bi, r) cells per worker.
6936                                unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
6937                            }
6938                            bi += 8;
6939                        }
6940                        while bi + 4 <= alen {
6941                            let xs = [
6942                                acts[abase + bi].xq.as_slice(),
6943                                acts[abase + bi + 1].xq.as_slice(),
6944                                acts[abase + bi + 2].xq.as_slice(),
6945                                acts[abase + bi + 3].xq.as_slice(),
6946                            ];
6947                            let d = unsafe { dot_q4tp_row_1x4_avx512(v.nib, r, gpr, xs, &sc) };
6948                            for k in 0..4 {
6949                                let act = &acts[abase + bi + k];
6950                                let mut acc = d[k] * act.sx;
6951                                for &(j, xv) in &act.outliers {
6952                                    let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
6953                                    acc += w * s * xv;
6954                                }
6955                                // SAFETY: disjoint (bi, r) cells per worker.
6956                                unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
6957                            }
6958                            bi += 4;
6959                        }
6960                    }
6961                    #[cfg(target_arch = "aarch64")]
6962                    if blocked_ok {
6963                        while bi + 4 <= alen {
6964                            let xs = [
6965                                acts[abase + bi].xq.as_slice(),
6966                                acts[abase + bi + 1].xq.as_slice(),
6967                                acts[abase + bi + 2].xq.as_slice(),
6968                                acts[abase + bi + 3].xq.as_slice(),
6969                            ];
6970                            let d = unsafe {
6971                                if exact || q4tp_v1() {
6972                                    dot_q4tp_row_1x4_sdot_v1(v.nib, r, gpr, xs, &sc)
6973                                } else {
6974                                    dot_q4tp_row_1x4_sdot(v.nib, r, gpr, xs, &sc)
6975                                }
6976                            };
6977                            for k in 0..4 {
6978                                let act = &acts[abase + bi + k];
6979                                let mut acc = d[k] * act.sx;
6980                                for &(j, xv) in &act.outliers {
6981                                    let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
6982                                    acc += w * s * xv;
6983                                }
6984                                // SAFETY: disjoint (bi, r) cells per worker.
6985                                unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
6986                            }
6987                            bi += 4;
6988                        }
6989                    }
6990                    let _ = blocked_ok;
6991                    while bi < alen {
6992                        let act = &acts[abase + bi];
6993                        let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, &sc) * act.sx;
6994                        for &(j, xv) in &act.outliers {
6995                            let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
6996                            acc += w * s * xv;
6997                        }
6998                        // SAFETY: disjoint (bi, r) cells per worker range.
6999                        unsafe { *out_addr.at((abase + bi) * rows + r) = acc };
7000                        bi += 1;
7001                    }
7002                }
7003            }
7004        };
7005        dispatch_rows(pool, rows, &run);
7006        return;
7007    }
7008
7009    let run = |start: usize, end: usize| {
7010        let mut sc = vec![0f32; gpr];
7011        for r in start..end {
7012            v.scales_into(r, gpr, &mut sc);
7013            for bi in 0..b {
7014                let x = &xs_all[bi * cols..(bi + 1) * cols];
7015                // SAFETY: disjoint (bi, r) cells per worker range.
7016                unsafe { *out_addr.at(bi * rows + r) = q4tp_row_exact(v.nib, r, gpr, x, &sc) };
7017            }
7018        }
7019    };
7020    dispatch_rows(pool, rows, &run);
7021}
7022
7023/// Fused q4_tiled matvec (dispatch mirrors `q4matvec`).
7024fn q4t_matvec(
7025    bytes: &[u8],
7026    x: &[f32],
7027    rows: usize,
7028    cols: usize,
7029    out: &mut [f32],
7030    pool: Option<&Pool>,
7031) {
7032    debug_assert_eq!(out.len(), rows);
7033    let gpr = cols / GROUP_SIZE;
7034    let out_addr = SendMut(out.as_mut_ptr());
7035    if a8w8_enabled() {
7036        let act = split_act(x);
7037        let run = move |start: usize, end: usize| {
7038            for r in start..end {
7039                let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
7040                for &(j, xv) in &act.outliers {
7041                    let (w, s) = q4t_outlier(bytes, r, gpr, j);
7042                    acc += w * s * xv;
7043                }
7044                // SAFETY: disjoint row ranges per worker.
7045                unsafe { *out_addr.at(r) = acc };
7046            }
7047        };
7048        dispatch_rows(pool, rows, &run);
7049        return;
7050    }
7051    let run = move |start: usize, end: usize| {
7052        for r in start..end {
7053            // SAFETY: disjoint row ranges per worker.
7054            unsafe { *out_addr.at(r) = q4t_row_exact(bytes, r, gpr, x) };
7055        }
7056    };
7057    dispatch_rows(pool, rows, &run);
7058}
7059
7060/// Fused two-input q4_tiled matvec (weights read once per pair).
7061#[allow(clippy::too_many_arguments)]
7062fn q4t_matvec2(
7063    bytes: &[u8],
7064    x1: &[f32],
7065    x2: &[f32],
7066    rows: usize,
7067    cols: usize,
7068    o1: &mut [f32],
7069    o2: &mut [f32],
7070    pool: Option<&Pool>,
7071) {
7072    let gpr = cols / GROUP_SIZE;
7073    let p1 = SendMut(o1.as_mut_ptr());
7074    let p2 = SendMut(o2.as_mut_ptr());
7075    if a8w8_enabled() {
7076        let a1 = split_act(x1);
7077        let a2 = split_act(x2);
7078        let run = move |start: usize, end: usize| {
7079            for r in start..end {
7080                let mut v1 = dot_q4t_row_i8(bytes, r, gpr, &a1.xq) * a1.sx;
7081                let mut v2 = dot_q4t_row_i8(bytes, r, gpr, &a2.xq) * a2.sx;
7082                for &(j, xv) in &a1.outliers {
7083                    let (w, s) = q4t_outlier(bytes, r, gpr, j);
7084                    v1 += w * s * xv;
7085                }
7086                for &(j, xv) in &a2.outliers {
7087                    let (w, s) = q4t_outlier(bytes, r, gpr, j);
7088                    v2 += w * s * xv;
7089                }
7090                // SAFETY: disjoint row ranges per worker.
7091                unsafe {
7092                    *p1.at(r) = v1;
7093                    *p2.at(r) = v2;
7094                }
7095            }
7096        };
7097        dispatch_rows(pool, rows, &run);
7098        return;
7099    }
7100    let run = move |start: usize, end: usize| {
7101        for r in start..end {
7102            // SAFETY: disjoint row ranges per worker.
7103            unsafe {
7104                *p1.at(r) = q4t_row_exact(bytes, r, gpr, x1);
7105                *p2.at(r) = q4t_row_exact(bytes, r, gpr, x2);
7106            }
7107        }
7108    };
7109    dispatch_rows(pool, rows, &run);
7110}
7111
7112/// Batched q4_tiled matmat: each row's tiles stream once per microbatch.
7113#[allow(clippy::too_many_arguments)]
7114/// Prefill GEMM through Accelerate for group-quantized codecs: a
7115/// caller-supplied row dequantizer fills f32 tiles (pool-parallel) and
7116/// each tile rides the AMX with one sgemm — the generic sibling of
7117/// `qmatmat_accel` (q8). Numerics are f32-GEMM (tolerance class);
7118/// decode (b=1) never takes this path.
7119#[cfg(target_os = "macos")]
7120fn dequant_matmat_accel(
7121    dequant_row: &(dyn Fn(usize, &mut [f32]) + Sync),
7122    xs_all: &[f32],
7123    b: usize,
7124    rows: usize,
7125    cols: usize,
7126    out: &mut [f32],
7127    pool: Option<&Pool>,
7128) {
7129    const TR: usize = 2048;
7130    thread_local! {
7131        static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
7132    }
7133    WTILE.with(|wt| {
7134        let mut wtile = wt.borrow_mut();
7135        wtile.resize(TR * cols, 0.0);
7136        let mut r0 = 0usize;
7137        while r0 < rows {
7138            let tr = TR.min(rows - r0);
7139            let wt_addr = SendMut(wtile.as_mut_ptr());
7140            let run = |start: usize, end: usize| {
7141                for r in start..end {
7142                    // SAFETY: workers cover disjoint r ranges.
7143                    let dst = unsafe { std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols) };
7144                    dequant_row(r0 + r, dst);
7145                }
7146            };
7147            dispatch_rows(pool, tr, &run);
7148            unsafe {
7149                accel_blas::cblas_sgemm(
7150                    101, // RowMajor
7151                    111, // NoTrans A
7152                    112, // Trans B
7153                    b as i32,
7154                    tr as i32,
7155                    cols as i32,
7156                    1.0,
7157                    xs_all.as_ptr(),
7158                    cols as i32,
7159                    wtile.as_ptr(),
7160                    cols as i32,
7161                    0.0,
7162                    out.as_mut_ptr().add(r0),
7163                    rows as i32,
7164                );
7165            }
7166            r0 += tr;
7167        }
7168    });
7169}
7170
7171fn q4t_matmat(
7172    bytes: &[u8],
7173    xs_all: &[f32],
7174    b: usize,
7175    rows: usize,
7176    cols: usize,
7177    out: &mut [f32],
7178    pool: Option<&Pool>,
7179) {
7180    debug_assert_eq!(out.len(), b * rows);
7181    let gpr = cols / GROUP_SIZE;
7182    // Wide batches ride the AMX like q8's qmatmat: on Apple silicon
7183    // the dequant-tile sgemm is an order above the SDOT row loop for
7184    // prefill shapes (imagegen DiT forwards are exactly this).
7185    #[cfg(target_os = "macos")]
7186    if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
7187        dequant_matmat_accel(
7188            &|r, dst| {
7189                for gi in 0..gpr {
7190                    let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
7191                    let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
7192                    for (k, &bb) in tile[2..].iter().enumerate() {
7193                        dst[gi * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s;
7194                        dst[gi * GROUP_SIZE + k * 2 + 1] = (((bb >> 4) & 0x0F) as f32 - 8.0) * s;
7195                    }
7196                }
7197            },
7198            xs_all,
7199            b,
7200            rows,
7201            cols,
7202            out,
7203            pool,
7204        );
7205        return;
7206    }
7207    let out_addr = SendMut(out.as_mut_ptr());
7208    if a8w8_enabled() {
7209        let acts: Vec<SplitAct> = (0..b)
7210            .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
7211            .collect();
7212        let acts = &acts;
7213        #[cfg(target_arch = "x86_64")]
7214        let blocked_ok = avx2_enabled() && blocked_enabled();
7215        #[cfg(target_arch = "aarch64")]
7216        let blocked_ok = sdot_enabled() && blocked_enabled();
7217        #[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))]
7218        let blocked_ok = false;
7219        let run = move |start: usize, end: usize| {
7220            for r in start..end {
7221                let mut bi = 0usize;
7222                #[cfg(target_arch = "aarch64")]
7223                if blocked_ok {
7224                    while bi + 4 <= acts.len() {
7225                        let xs = [
7226                            acts[bi].xq.as_slice(),
7227                            acts[bi + 1].xq.as_slice(),
7228                            acts[bi + 2].xq.as_slice(),
7229                            acts[bi + 3].xq.as_slice(),
7230                        ];
7231                        let d = unsafe { dot_q4t_row_1x4_sdot(bytes, r, gpr, xs) };
7232                        for k in 0..4 {
7233                            let act = &acts[bi + k];
7234                            let mut acc = d[k] * act.sx;
7235                            for &(j, xv) in &act.outliers {
7236                                let (w, sc) = q4t_outlier(bytes, r, gpr, j);
7237                                acc += w * sc * xv;
7238                            }
7239                            // SAFETY: disjoint (bi, r) cells per worker.
7240                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
7241                        }
7242                        bi += 4;
7243                    }
7244                }
7245                #[cfg(target_arch = "x86_64")]
7246                if blocked_ok {
7247                    while bi + 4 <= acts.len() {
7248                        let xs = [
7249                            acts[bi].xq.as_slice(),
7250                            acts[bi + 1].xq.as_slice(),
7251                            acts[bi + 2].xq.as_slice(),
7252                            acts[bi + 3].xq.as_slice(),
7253                        ];
7254                        let d = unsafe {
7255                            if vnni_tiles_enabled() {
7256                                dot_q4t_row_1x4_vnni(bytes, r, gpr, xs)
7257                            } else {
7258                                dot_q4t_row_1x4_avx2(bytes, r, gpr, xs)
7259                            }
7260                        };
7261                        for k in 0..4 {
7262                            let act = &acts[bi + k];
7263                            let mut acc = d[k] * act.sx;
7264                            for &(j, xv) in &act.outliers {
7265                                let (w, sc) = q4t_outlier(bytes, r, gpr, j);
7266                                acc += w * sc * xv;
7267                            }
7268                            // SAFETY: disjoint (bi, r) cells per worker.
7269                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
7270                        }
7271                        bi += 4;
7272                    }
7273                }
7274                let _ = blocked_ok;
7275                while bi < acts.len() {
7276                    let act = &acts[bi];
7277                    let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
7278                    for &(j, xv) in &act.outliers {
7279                        let (w, s) = q4t_outlier(bytes, r, gpr, j);
7280                        acc += w * s * xv;
7281                    }
7282                    // SAFETY: disjoint (bi, r) cells per worker range.
7283                    unsafe { *out_addr.at(bi * rows + r) = acc };
7284                    bi += 1;
7285                }
7286            }
7287        };
7288        dispatch_rows(pool, rows, &run);
7289        return;
7290    }
7291    let run = move |start: usize, end: usize| {
7292        for r in start..end {
7293            for bi in 0..b {
7294                let x = &xs_all[bi * cols..(bi + 1) * cols];
7295                // SAFETY: disjoint (bi, r) cells per worker range.
7296                unsafe { *out_addr.at(bi * rows + r) = q4t_row_exact(bytes, r, gpr, x) };
7297            }
7298        }
7299    };
7300    dispatch_rows(pool, rows, &run);
7301}
7302
7303// ── q1 (dtype 12): binary weights, [f16 scale][4B sign bits] per
7304// 32-group tile. The kernel family mirrors q4_tiled: one sequential
7305// stream of 6-byte tiles, per-tile integer dot × scale, exact outlier
7306// correction (A8W8 contract), exact scalar path under CMF_SDOT=0. ──
7307
7308/// Per-32-group sums of the quantized activation — the ±1 identity's
7309/// shared half: `dot = −2·sdot(mask, x) − gsum[g]`, computed ONCE per
7310/// matvec and reused by every row.
7311fn q1_group_sums(xq: &[i8], gpr: usize) -> Vec<i32> {
7312    (0..gpr)
7313        .map(|gi| {
7314            xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE]
7315                .iter()
7316                .map(|&v| v as i32)
7317                .sum()
7318        })
7319        .collect()
7320}
7321
7322/// One q1 row via the A8W8 int8 path — mask-SDOT on ARM (no ±1
7323/// expansion at all), scalar bit loop elsewhere (AVX2 queued with the
7324/// x86 pass).
7325#[inline]
7326#[allow(unreachable_code)]
7327/// AVX2 q1 row via the same ±1 identity as the ARM sdot kernel: the
7328/// sign bits expand to a {0, −1} byte mask through shuffle+cmpeq, the
7329/// masked activation sums through maddubs(1, x&mask), and
7330/// `dot = −(2·masked_sum + Σx_group)` — bit-identical integer math.
7331#[cfg(target_arch = "x86_64")]
7332#[target_feature(enable = "avx2")]
7333unsafe fn dot_q1_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
7334    // SAFETY: callers uphold the 6B-tile and xq/gsum length contracts.
7335    unsafe {
7336        use core::arch::x86_64::*;
7337        // Byte j of the mask must replicate bits-byte j/8.
7338        let expand = _mm256_setr_epi8(
7339            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,
7340            3, 3, 3,
7341        );
7342        let bitsel = _mm256_setr_epi8(
7343            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
7344            -128, 1, 2, 4, 8, 16, 32, 64, -128,
7345        );
7346        let ones8 = _mm256_set1_epi8(1);
7347        let ones16 = _mm256_set1_epi16(1);
7348        let mut acc = 0f32;
7349        for gi in 0..gpr {
7350            let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
7351            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
7352            let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
7353            let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
7354            let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
7355            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7356            let sel = _mm256_and_si256(x, mask);
7357            // Σ of selected i8 lanes: maddubs(1u8, sel_i8) pairs → madd.
7358            let p16 = _mm256_maddubs_epi16(ones8, sel);
7359            let d32 = _mm256_madd_epi16(p16, ones16);
7360            let hi128 = _mm256_extracti128_si256::<1>(d32);
7361            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
7362            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7363            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7364            let msum = _mm_cvtsi128_si32(s32);
7365            // The and-select keeps x UN-negated (unlike ARM's −1-mask
7366            // sdot): d = Σ_set − Σ_unset = 2·Σ_set − Σ_all.
7367            let d = 2 * msum - gsum[gi];
7368            acc += d as f32 * s;
7369        }
7370        acc
7371    }
7372}
7373
7374/// VNNI twin of `dot_q1_row_avx2`: the masked-select sum goes through
7375/// one `vpdpbusd(1u8, sel)` (see `dpbusd_hsum` — bit-identical).
7376#[cfg(target_arch = "x86_64")]
7377#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
7378unsafe fn dot_q1_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
7379    // SAFETY: callers uphold the 6B-tile and xq/gsum length contracts.
7380    unsafe {
7381        use core::arch::x86_64::*;
7382        let expand = _mm256_setr_epi8(
7383            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,
7384            3, 3, 3,
7385        );
7386        let bitsel = _mm256_setr_epi8(
7387            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
7388            -128, 1, 2, 4, 8, 16, 32, 64, -128,
7389        );
7390        let ones8 = _mm256_set1_epi8(1);
7391        let mut acc = 0f32;
7392        for gi in 0..gpr {
7393            let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
7394            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
7395            let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
7396            let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
7397            let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
7398            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7399            let msum = dpbusd_hsum(ones8, _mm256_and_si256(x, mask));
7400            let d = 2 * msum - gsum[gi];
7401            acc += d as f32 * s;
7402        }
7403        acc
7404    }
7405}
7406
7407/// VNNI twin of `dot_q1_row_1x4_avx2` (see `dpbusd_hsum`).
7408#[cfg(target_arch = "x86_64")]
7409#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
7410unsafe fn dot_q1_row_1x4_vnni(
7411    bytes: &[u8],
7412    r: usize,
7413    gpr: usize,
7414    xs: [&[i8]; 4],
7415    gsums: [&[i32]; 4],
7416) -> [f32; 4] {
7417    // SAFETY: callers uphold the 6B-tile and xq/gsum length contracts.
7418    unsafe {
7419        use core::arch::x86_64::*;
7420        let expand = _mm256_setr_epi8(
7421            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,
7422            3, 3, 3,
7423        );
7424        let bitsel = _mm256_setr_epi8(
7425            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
7426            -128, 1, 2, 4, 8, 16, 32, 64, -128,
7427        );
7428        let ones8 = _mm256_set1_epi8(1);
7429        let mut acc = [0f32; 4];
7430        for gi in 0..gpr {
7431            let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
7432            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
7433            let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
7434            let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
7435            let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
7436            for (k, xq) in xs.iter().enumerate() {
7437                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7438                let msum = dpbusd_hsum(ones8, _mm256_and_si256(x, mask));
7439                let d = 2 * msum - gsums[k][gi];
7440                acc[k] += d as f32 * s;
7441            }
7442        }
7443        acc
7444    }
7445}
7446
7447/// The blocked 1×4 flavor: the expanded bit mask serves four activation
7448/// streams per group (mask build once, four select+reduce chains).
7449#[cfg(target_arch = "x86_64")]
7450#[target_feature(enable = "avx2")]
7451unsafe fn dot_q1_row_1x4_avx2(
7452    bytes: &[u8],
7453    r: usize,
7454    gpr: usize,
7455    xs: [&[i8]; 4],
7456    gsums: [&[i32]; 4],
7457) -> [f32; 4] {
7458    // SAFETY: callers uphold the 6B-tile and xq/gsum length contracts.
7459    unsafe {
7460        use core::arch::x86_64::*;
7461        let expand = _mm256_setr_epi8(
7462            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,
7463            3, 3, 3,
7464        );
7465        let bitsel = _mm256_setr_epi8(
7466            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
7467            -128, 1, 2, 4, 8, 16, 32, 64, -128,
7468        );
7469        let ones8 = _mm256_set1_epi8(1);
7470        let ones16 = _mm256_set1_epi16(1);
7471        let mut acc = [0f32; 4];
7472        for gi in 0..gpr {
7473            let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
7474            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
7475            let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
7476            let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
7477            let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
7478            for (k, xq) in xs.iter().enumerate() {
7479                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7480                let sel = _mm256_and_si256(x, mask);
7481                let p16 = _mm256_maddubs_epi16(ones8, sel);
7482                let d32 = _mm256_madd_epi16(p16, ones16);
7483                let hi128 = _mm256_extracti128_si256::<1>(d32);
7484                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
7485                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7486                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7487                let msum = _mm_cvtsi128_si32(s32);
7488                let d = 2 * msum - gsums[k][gi];
7489                acc[k] += d as f32 * s;
7490            }
7491        }
7492        acc
7493    }
7494}
7495
7496#[allow(unreachable_code)]
7497fn dot_q1_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
7498    #[cfg(target_arch = "aarch64")]
7499    unsafe {
7500        return dot_q1_row_sdot(bytes, r, gpr, xq, gsum);
7501    }
7502    #[cfg(target_arch = "x86_64")]
7503    if avx2_enabled() {
7504        unsafe {
7505            if vnni_tiles_enabled() {
7506                return dot_q1_row_vnni(bytes, r, gpr, xq, gsum);
7507            }
7508            return dot_q1_row_avx2(bytes, r, gpr, xq, gsum);
7509        }
7510    }
7511    let _ = gsum;
7512    let mut acc = 0f32;
7513    for gi in 0..gpr {
7514        let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
7515        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
7516        let mut d = 0i32;
7517        for (j, &b) in tile[2..].iter().enumerate() {
7518            for k in 0..8 {
7519                let w = ((b >> k) & 1) as i32 * 2 - 1;
7520                d += w * xq[gi * GROUP_SIZE + j * 8 + k] as i32;
7521            }
7522        }
7523        acc += d as f32 * s;
7524    }
7525    acc
7526}
7527
7528/// SDOT q1 row via the ±1 identity: the vtst mask (0xFF where the bit
7529/// is set, i.e. −1 as i8) feeds `sdot` DIRECTLY — no expansion to ±1
7530/// lanes at all — and `dot = −(2·sdot(mask, x) + Σx_group)`, with the
7531/// per-group activation sums shared across every row of the matvec.
7532/// Four tiles (128 weights) per iteration: integer dots reduce through
7533/// a vpaddq tree into ONE i32x4 that meets its four scales in a single
7534/// fused f32 multiply-add. Integer math throughout — bit-identical to
7535/// the scalar ±1 reference.
7536#[cfg(target_arch = "aarch64")]
7537#[target_feature(enable = "neon,dotprod")]
7538unsafe fn dot_q1_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
7539    // SAFETY: callers uphold slice-length contracts (6B tile per group,
7540    // xq.len() == gpr·GROUP_SIZE, gsum.len() == gpr).
7541    unsafe {
7542        use core::arch::aarch64::*;
7543        use core::arch::asm;
7544        const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
7545        let m = vld1q_u8(MASKS.as_ptr());
7546        // One tile's −Σ_set(x) as an UNREDUCED i32x4 (two mask-sdots).
7547        macro_rules! tile_dot {
7548            ($t:expr, $x:expr) => {{
7549                let v0 = vcombine_u8(vdup_n_u8(*$t.add(2)), vdup_n_u8(*$t.add(3)));
7550                let v1 = vcombine_u8(vdup_n_u8(*$t.add(4)), vdup_n_u8(*$t.add(5)));
7551                let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
7552                let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
7553                let x0 = vld1q_s8($x);
7554                let x1 = vld1q_s8($x.add(16));
7555                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
7556                asm!(
7557                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
7558                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
7559                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
7560                    w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
7561                    options(pure, nomem, nostack),
7562                );
7563                vaddq_s32(a0, a1)
7564            }};
7565        }
7566        // TBL unpack over PAIR loads: one vld1q covers two 6B tiles
7567        // ([s s b b b b][s s b b b b] + 4B slack), TBL replicates each
7568        // bit-byte across 8 lanes for vtst, and the four scales gather
7569        // through tbl2 into one fcvtl — the 16 ld1r broadcast loads and
7570        // 4 branchy software f16 conversions per 128 weights (the
7571        // measured load-port wall of this kernel) become 2 vector
7572        // loads + 9 table lookups. Integer math order is unchanged —
7573        // bit-identical results (FCVTL is exact on every f16).
7574        const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
7575        const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
7576        const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
7577        const IW11: [u8; 16] = [
7578            10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
7579        ];
7580        const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
7581        let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
7582        let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
7583        let isc = vld1_u8(ISC.as_ptr());
7584        // One tile's −Σ_set(x) from a TBL-unpacked pair load.
7585        macro_rules! tile_dot_tbl {
7586            ($ld:expr, $i0:expr, $i1:expr, $x:expr) => {{
7587                let w0 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i0), m));
7588                let w1 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i1), m));
7589                let x0 = vld1q_s8($x);
7590                let x1 = vld1q_s8($x.add(16));
7591                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
7592                asm!(
7593                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
7594                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
7595                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
7596                    w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
7597                    options(pure, nomem, nostack),
7598                );
7599                vaddq_s32(a0, a1)
7600            }};
7601        }
7602        let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
7603        let row_base = r * gpr * Q1_TILE;
7604        let abs_end = bytes.len();
7605        let xp = xq.as_ptr();
7606        let gp = gsum.as_ptr();
7607        let mut accv = vdupq_n_f32(0.0);
7608        let mut gi = 0;
7609        // The second pair load reads 4B past tile gi+3 — stay inside
7610        // the payload slice (only the file's final tiles fall back).
7611        while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
7612            let t0 = base.add(gi * Q1_TILE);
7613            let ld_a = vld1q_u8(t0);
7614            let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
7615            let d0 = tile_dot_tbl!(ld_a, iw00, iw01, xp.add(gi * GROUP_SIZE));
7616            let d1 = tile_dot_tbl!(ld_a, iw10, iw11, xp.add((gi + 1) * GROUP_SIZE));
7617            let d2 = tile_dot_tbl!(ld_b, iw00, iw01, xp.add((gi + 2) * GROUP_SIZE));
7618            let d3 = tile_dot_tbl!(ld_b, iw10, iw11, xp.add((gi + 3) * GROUP_SIZE));
7619            // [−Σ0, −Σ1, −Σ2, −Σ3] → dots = −(2·Σset_neg + gsum)
7620            let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
7621            let g = vld1q_s32(gp.add(gi));
7622            let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
7623            let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
7624            let scf: float32x4_t;
7625            asm!(
7626                "fcvtl {o:v}.4s, {i:v}.4h",
7627                o = out(vreg) scf, i = in(vreg) sc16,
7628                options(pure, nomem, nostack),
7629            );
7630            accv = vfmaq_f32(accv, vcvtq_f32_s32(dots), scf);
7631            gi += 4;
7632        }
7633        let mut acc = vaddvq_f32(accv);
7634        while gi < gpr {
7635            let t = base.add(gi * Q1_TILE);
7636            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
7637            let d = vaddvq_s32(tile_dot!(t, xp.add(gi * GROUP_SIZE)));
7638            acc += (-(2 * d + *gp.add(gi))) as f32 * s;
7639            gi += 1;
7640        }
7641        acc
7642    }
7643}
7644
7645/// Blocked q1 1×4: one TBL unpack of the tile pair serves FOUR
7646/// activation streams (prefill amortization — the same idea as the
7647/// AVX2 twin; per stream the group order, fma order and tail match the
7648/// single-row kernel exactly, so batch == matvec bit-for-bit).
7649#[cfg(target_arch = "aarch64")]
7650#[target_feature(enable = "neon,dotprod")]
7651unsafe fn dot_q1_row_1x4_sdot(
7652    bytes: &[u8],
7653    r: usize,
7654    gpr: usize,
7655    xs: [&[i8]; 4],
7656    gs: [&[i32]; 4],
7657) -> [f32; 4] {
7658    // SAFETY: same slice-length contracts as `dot_q1_row_sdot`, ×4.
7659    unsafe {
7660        use core::arch::aarch64::*;
7661        use core::arch::asm;
7662        const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
7663        const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
7664        const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
7665        const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
7666        const IW11: [u8; 16] = [
7667            10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
7668        ];
7669        const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
7670        let m = vld1q_u8(MASKS.as_ptr());
7671        let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
7672        let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
7673        let isc = vld1_u8(ISC.as_ptr());
7674        macro_rules! sdot2 {
7675            ($w0:expr, $w1:expr, $x:expr) => {{
7676                let x0 = vld1q_s8($x);
7677                let x1 = vld1q_s8($x.add(16));
7678                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
7679                asm!(
7680                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
7681                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
7682                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
7683                    w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
7684                    options(pure, nomem, nostack),
7685                );
7686                vaddq_s32(a0, a1)
7687            }};
7688        }
7689        let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
7690        let row_base = r * gpr * Q1_TILE;
7691        let abs_end = bytes.len();
7692        let mut accv = [vdupq_n_f32(0.0); 4];
7693        let mut gi = 0;
7694        while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
7695            let t0 = base.add(gi * Q1_TILE);
7696            let ld_a = vld1q_u8(t0);
7697            let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
7698            // Unpack ONCE — eight ±mask vectors serve all four streams.
7699            let w00 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw00), m));
7700            let w01 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw01), m));
7701            let w10 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw10), m));
7702            let w11 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw11), m));
7703            let w20 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw00), m));
7704            let w21 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw01), m));
7705            let w30 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw10), m));
7706            let w31 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw11), m));
7707            let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
7708            let scf: float32x4_t;
7709            asm!(
7710                "fcvtl {o:v}.4s, {i:v}.4h",
7711                o = out(vreg) scf, i = in(vreg) sc16,
7712                options(pure, nomem, nostack),
7713            );
7714            for k in 0..4 {
7715                let xp = xs[k].as_ptr();
7716                let d0 = sdot2!(w00, w01, xp.add(gi * GROUP_SIZE));
7717                let d1 = sdot2!(w10, w11, xp.add((gi + 1) * GROUP_SIZE));
7718                let d2 = sdot2!(w20, w21, xp.add((gi + 2) * GROUP_SIZE));
7719                let d3 = sdot2!(w30, w31, xp.add((gi + 3) * GROUP_SIZE));
7720                let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
7721                let g = vld1q_s32(gs[k].as_ptr().add(gi));
7722                let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
7723                accv[k] = vfmaq_f32(accv[k], vcvtq_f32_s32(dots), scf);
7724            }
7725            gi += 4;
7726        }
7727        let mut acc = [
7728            vaddvq_f32(accv[0]),
7729            vaddvq_f32(accv[1]),
7730            vaddvq_f32(accv[2]),
7731            vaddvq_f32(accv[3]),
7732        ];
7733        while gi < gpr {
7734            let t = base.add(gi * Q1_TILE);
7735            let sc = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
7736            let v0 = vcombine_u8(vdup_n_u8(*t.add(2)), vdup_n_u8(*t.add(3)));
7737            let v1 = vcombine_u8(vdup_n_u8(*t.add(4)), vdup_n_u8(*t.add(5)));
7738            let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
7739            let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
7740            for k in 0..4 {
7741                let d = vaddvq_s32(sdot2!(w0, w1, xs[k].as_ptr().add(gi * GROUP_SIZE)));
7742                acc[k] += (-(2 * d + *gs[k].as_ptr().add(gi))) as f32 * sc;
7743            }
7744            gi += 1;
7745        }
7746        acc
7747    }
7748}
7749
7750/// (weight ±1, scale) of one q1 element — the exact outlier term.
7751#[inline]
7752fn q1_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
7753    let gi = j / GROUP_SIZE;
7754    let k = j % GROUP_SIZE;
7755    let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
7756    let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
7757    let bit = (tile[2 + k / 8] >> (k % 8)) & 1;
7758    ((bit as i32 * 2 - 1) as f32, s)
7759}
7760
7761/// Exact scalar q1 row (CMF_SDOT=0 contract).
7762#[inline]
7763fn q1_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
7764    let mut acc = 0f32;
7765    for gi in 0..gpr {
7766        let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
7767        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
7768        let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
7769        let mut ga = 0f32;
7770        for (j, &b) in tile[2..].iter().enumerate() {
7771            for k in 0..8 {
7772                ga += (((b >> k) & 1) as f32 * 2.0 - 1.0) * xg[j * 8 + k];
7773            }
7774        }
7775        acc += ga * s;
7776    }
7777    acc
7778}
7779
7780/// One binary-weight row range via A8W8 (the body of `bin_matvec`'s hot
7781/// loop, extracted so multi-matrix jobs drive the same kernel).
7782#[allow(clippy::too_many_arguments)]
7783fn bin_range_a8w8(
7784    w: BinW,
7785    gpr: usize,
7786    act: &SplitAct,
7787    gsum: &[i32],
7788    out: SendMut,
7789    start: usize,
7790    end: usize,
7791) {
7792    for r in start..end {
7793        let mut acc = w.dot_i8(r, gpr, &act.xq, gsum) * act.sx;
7794        for &(j, xv) in &act.outliers {
7795            let (wv, s) = w.outlier(r, gpr, j);
7796            acc += wv * s * xv;
7797        }
7798        // SAFETY: disjoint row ranges per worker.
7799        unsafe { *out.at(r) = acc };
7800    }
7801}
7802
7803/// Exact-scalar binary-weight row range (CMF_SDOT=0 contract).
7804fn bin_range_f32(w: BinW, gpr: usize, x: &[f32], out: SendMut, start: usize, end: usize) {
7805    for r in start..end {
7806        // SAFETY: disjoint row ranges per worker.
7807        unsafe { *out.at(r) = w.row_exact(r, gpr, x) };
7808    }
7809}
7810
7811/// A binary (±s) weight in either of its layouts: `q1` 6-byte tiles, or
7812/// `q1g` planes with one scale per 128-group. The row drivers — matvec,
7813/// matvec2, matmat, the multi-projection job and the fused gate/up pair —
7814/// are written once against this; only the per-row kernels differ, and
7815/// the match costs one predictable branch per row. `gpr` is always the
7816/// row's count of 32-tiles (cols / 32): the activation group sums and the
7817/// tile walk are shared by both layouts.
7818#[derive(Clone, Copy)]
7819pub(crate) enum BinW<'a> {
7820    Q1(&'a [u8]),
7821    Q1g {
7822        bytes: &'a [u8],
7823        /// Byte offset of the f16 scale plane (= rows · cols/8).
7824        sc_off: usize,
7825        /// Row count: the q1 kernels treat the payload's final row
7826        /// specially (see `dot_q1g_row_sdot`), and q1g mirrors that.
7827        rows: usize,
7828    },
7829}
7830
7831impl<'a> BinW<'a> {
7832    fn q1g(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
7833        Self::Q1g {
7834            bytes,
7835            sc_off: cortiq_core::quant::q1g_scale_offset(rows, cols),
7836            rows,
7837        }
7838    }
7839
7840    #[inline]
7841    fn dot_i8(self, r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
7842        match self {
7843            Self::Q1(b) => dot_q1_row_i8(b, r, gpr, xq, gsum),
7844            Self::Q1g {
7845                bytes,
7846                sc_off,
7847                rows,
7848            } => dot_q1g_row_i8(bytes, sc_off, rows, r, gpr, xq, gsum),
7849        }
7850    }
7851
7852    #[inline]
7853    fn row_exact(self, r: usize, gpr: usize, x: &[f32]) -> f32 {
7854        match self {
7855            Self::Q1(b) => q1_row_exact(b, r, gpr, x),
7856            Self::Q1g { bytes, sc_off, .. } => q1g_row_exact(bytes, sc_off, r, gpr, x),
7857        }
7858    }
7859
7860    #[inline]
7861    fn outlier(self, r: usize, gpr: usize, j: usize) -> (f32, f32) {
7862        match self {
7863            Self::Q1(b) => q1_outlier(b, r, gpr, j),
7864            Self::Q1g { bytes, sc_off, .. } => q1g_outlier(bytes, sc_off, r, gpr, j),
7865        }
7866    }
7867
7868    /// Blocked 1×4 row (one unpack, four activation streams).
7869    ///
7870    /// # Safety
7871    /// The SDOT/dotprod target feature must be present (`sdot_enabled`).
7872    #[cfg(target_arch = "aarch64")]
7873    #[inline]
7874    unsafe fn dot_1x4_sdot(
7875        self,
7876        r: usize,
7877        gpr: usize,
7878        xs: [&[i8]; 4],
7879        gs: [&[i32]; 4],
7880    ) -> [f32; 4] {
7881        // SAFETY: forwarded caller contract.
7882        unsafe {
7883            match self {
7884                Self::Q1(b) => dot_q1_row_1x4_sdot(b, r, gpr, xs, gs),
7885                Self::Q1g {
7886                    bytes,
7887                    sc_off,
7888                    rows,
7889                } => dot_q1g_row_1x4_sdot(bytes, sc_off, rows, r, gpr, xs, gs),
7890            }
7891        }
7892    }
7893
7894    /// Blocked 1×4 row on x86 (VNNI when the tile switch is on).
7895    ///
7896    /// # Safety
7897    /// AVX2 must be present (`avx2_enabled`).
7898    #[cfg(target_arch = "x86_64")]
7899    #[inline]
7900    unsafe fn dot_1x4_x86(self, r: usize, gpr: usize, xs: [&[i8]; 4], gs: [&[i32]; 4]) -> [f32; 4] {
7901        // SAFETY: forwarded caller contract.
7902        unsafe {
7903            match (self, vnni_tiles_enabled()) {
7904                (Self::Q1(b), true) => dot_q1_row_1x4_vnni(b, r, gpr, xs, gs),
7905                (Self::Q1(b), false) => dot_q1_row_1x4_avx2(b, r, gpr, xs, gs),
7906                (Self::Q1g { bytes, sc_off, .. }, true) => {
7907                    dot_q1g_row_1x4_vnni(bytes, sc_off, r, gpr, xs, gs)
7908                }
7909                (Self::Q1g { bytes, sc_off, .. }, false) => {
7910                    dot_q1g_row_1x4_avx2(bytes, sc_off, r, gpr, xs, gs)
7911                }
7912            }
7913        }
7914    }
7915}
7916
7917// ── q1g (dtype 20): binary weights, ONE f16 scale per 128-group, in sign
7918// and scale PLANES (`[bits: rows·cols/8][f16: rows·cols/128]`). Same
7919// kernel family as q1 and — deliberately — the same float accumulation
7920// order (one `acc += d·s` per 32-tile, the group's scale reused four
7921// times), so a q1g repack of a q1 tensor whose tile scales agree gives
7922// the same f32 outputs bit for bit on every path. The win is the stream:
7923// 18 bytes per 128 weights instead of 24, the bits 16-byte aligned. ──
7924
7925/// Sign bits of q1g tile `t` (flat 32-tile index `r·gpr + gi`).
7926#[inline]
7927fn q1g_bits(bytes: &[u8], t: usize) -> u32 {
7928    u32::from_le_bytes([
7929        bytes[t * 4],
7930        bytes[t * 4 + 1],
7931        bytes[t * 4 + 2],
7932        bytes[t * 4 + 3],
7933    ])
7934}
7935
7936/// Raw f16 bits of the scale of the 128-group holding q1g tile `t`.
7937#[inline]
7938fn q1g_scale_bits(bytes: &[u8], sc_off: usize, t: usize) -> u16 {
7939    let o = sc_off + (t / 4) * 2;
7940    u16::from_le_bytes([bytes[o], bytes[o + 1]])
7941}
7942
7943/// One q1g row via the A8W8 int8 path (dispatch twin of `dot_q1_row_i8`).
7944#[allow(unreachable_code)]
7945#[allow(clippy::too_many_arguments)]
7946fn dot_q1g_row_i8(
7947    bytes: &[u8],
7948    sc_off: usize,
7949    rows: usize,
7950    r: usize,
7951    gpr: usize,
7952    xq: &[i8],
7953    gsum: &[i32],
7954) -> f32 {
7955    #[cfg(target_arch = "aarch64")]
7956    unsafe {
7957        return dot_q1g_row_sdot(bytes, sc_off, rows, r, gpr, xq, gsum);
7958    }
7959    #[cfg(target_arch = "x86_64")]
7960    if avx2_enabled() {
7961        unsafe {
7962            if vnni_tiles_enabled() {
7963                return dot_q1g_row_vnni(bytes, sc_off, r, gpr, xq, gsum);
7964            }
7965            return dot_q1g_row_avx2(bytes, sc_off, r, gpr, xq, gsum);
7966        }
7967    }
7968    let _ = (gsum, rows);
7969    let mut acc = 0f32;
7970    for gi in 0..gpr {
7971        let t = r * gpr + gi;
7972        let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, t));
7973        let bits = q1g_bits(bytes, t);
7974        let mut d = 0i32;
7975        for k in 0..GROUP_SIZE {
7976            let w = ((bits >> k) & 1) as i32 * 2 - 1;
7977            d += w * xq[gi * GROUP_SIZE + k] as i32;
7978        }
7979        acc += d as f32 * s;
7980    }
7981    acc
7982}
7983
7984/// Four tiles' unreduced `i32×8` partial sums → `[Σ0, Σ1, Σ2, Σ3]` in one
7985/// hadd tree: 3 hadds + 1 add for four tiles where the per-tile reduce of
7986/// the q1 kernels costs 7 ops each. Integer math, so any order gives the
7987/// q1 kernel's per-tile sums exactly.
7988#[cfg(target_arch = "x86_64")]
7989#[target_feature(enable = "avx2")]
7990#[inline]
7991unsafe fn q1g_hsum4(
7992    d0: core::arch::x86_64::__m256i,
7993    d1: core::arch::x86_64::__m256i,
7994    d2: core::arch::x86_64::__m256i,
7995    d3: core::arch::x86_64::__m256i,
7996) -> core::arch::x86_64::__m128i {
7997    // SAFETY: pure register math.
7998    unsafe {
7999        use core::arch::x86_64::*;
8000        let h = _mm256_hadd_epi32(_mm256_hadd_epi32(d0, d1), _mm256_hadd_epi32(d2, d3));
8001        _mm_add_epi32(_mm256_castsi256_si128(h), _mm256_extracti128_si256::<1>(h))
8002    }
8003}
8004
8005/// One 128-group's epilogue, in q1's float order: `d = 2·Σset − Σx` per
8006/// tile (exact i32 → exact f32), then `acc += d·s` tile by tile — four
8007/// separately rounded multiply/add pairs, the q1 kernels' sequence.
8008#[cfg(target_arch = "x86_64")]
8009#[target_feature(enable = "avx2")]
8010#[inline]
8011unsafe fn q1g_group_acc(
8012    acc: f32,
8013    msum4: core::arch::x86_64::__m128i,
8014    gsum4: &[i32],
8015    s: f32,
8016) -> f32 {
8017    // SAFETY: gsum4 holds the group's four tile sums (callers slice it).
8018    unsafe {
8019        use core::arch::x86_64::*;
8020        let g = _mm_loadu_si128(gsum4.as_ptr() as *const __m128i);
8021        let d = _mm_cvtepi32_ps(_mm_sub_epi32(_mm_slli_epi32::<1>(msum4), g));
8022        let mut f = [0f32; 4];
8023        _mm_storeu_ps(f.as_mut_ptr(), d);
8024        let mut acc = acc;
8025        for v in f {
8026            acc += v * s;
8027        }
8028        acc
8029    }
8030}
8031
8032/// AVX2 q1g row: `dot_q1_row_avx2`'s mask/select/maddubs per tile, but the
8033/// four tiles of a 128-group reduce together (`q1g_hsum4`) and the scale
8034/// converts once per group — bit-identical results (integer sums, then
8035/// q1's per-tile float sequence), fewer horizontal ops and f16 decodes.
8036#[cfg(target_arch = "x86_64")]
8037#[target_feature(enable = "avx2")]
8038unsafe fn dot_q1g_row_avx2(
8039    bytes: &[u8],
8040    sc_off: usize,
8041    r: usize,
8042    gpr: usize,
8043    xq: &[i8],
8044    gsum: &[i32],
8045) -> f32 {
8046    // SAFETY: callers uphold the plane-length and xq/gsum contracts
8047    // (gpr % 4 == 0, xq.len() == gpr·32, gsum.len() == gpr).
8048    unsafe {
8049        use core::arch::x86_64::*;
8050        let expand = _mm256_setr_epi8(
8051            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,
8052            3, 3, 3,
8053        );
8054        let bitsel = _mm256_setr_epi8(
8055            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
8056            -128, 1, 2, 4, 8, 16, 32, 64, -128,
8057        );
8058        let ones8 = _mm256_set1_epi8(1);
8059        let ones16 = _mm256_set1_epi16(1);
8060        let bp = bytes.as_ptr().add(r * gpr * 4);
8061        let mut acc = 0f32;
8062        let mut gi = 0;
8063        while gi < gpr {
8064            let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
8065            let mut d = [_mm256_setzero_si256(); 4];
8066            for (k, dk) in d.iter_mut().enumerate() {
8067                let t = gi + k;
8068                let bits = (bp.add(t * 4) as *const u32).read_unaligned();
8069                let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
8070                let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
8071                let x = _mm256_loadu_si256(xq.as_ptr().add(t * GROUP_SIZE) as *const __m256i);
8072                let sel = _mm256_and_si256(x, mask);
8073                *dk = _mm256_madd_epi16(_mm256_maddubs_epi16(ones8, sel), ones16);
8074            }
8075            let m4 = q1g_hsum4(d[0], d[1], d[2], d[3]);
8076            acc = q1g_group_acc(acc, m4, &gsum[gi..gi + 4], s);
8077            gi += 4;
8078        }
8079        acc
8080    }
8081}
8082
8083/// VNNI twin of `dot_q1g_row_avx2`: `vpdpbusd(1u8, sel)` per tile (the
8084/// same i32 lanes as maddubs+madd — bounded products), same group reduce.
8085#[cfg(target_arch = "x86_64")]
8086#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
8087unsafe fn dot_q1g_row_vnni(
8088    bytes: &[u8],
8089    sc_off: usize,
8090    r: usize,
8091    gpr: usize,
8092    xq: &[i8],
8093    gsum: &[i32],
8094) -> f32 {
8095    // SAFETY: callers uphold the plane-length and xq/gsum contracts.
8096    unsafe {
8097        use core::arch::x86_64::*;
8098        let expand = _mm256_setr_epi8(
8099            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,
8100            3, 3, 3,
8101        );
8102        let bitsel = _mm256_setr_epi8(
8103            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
8104            -128, 1, 2, 4, 8, 16, 32, 64, -128,
8105        );
8106        let ones8 = _mm256_set1_epi8(1);
8107        let bp = bytes.as_ptr().add(r * gpr * 4);
8108        let mut acc = 0f32;
8109        let mut gi = 0;
8110        while gi < gpr {
8111            let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
8112            let mut d = [_mm256_setzero_si256(); 4];
8113            for (k, dk) in d.iter_mut().enumerate() {
8114                let t = gi + k;
8115                let bits = (bp.add(t * 4) as *const u32).read_unaligned();
8116                let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
8117                let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
8118                let x = _mm256_loadu_si256(xq.as_ptr().add(t * GROUP_SIZE) as *const __m256i);
8119                *dk = _mm256_dpbusd_epi32(_mm256_setzero_si256(), ones8, _mm256_and_si256(x, mask));
8120            }
8121            let m4 = q1g_hsum4(d[0], d[1], d[2], d[3]);
8122            acc = q1g_group_acc(acc, m4, &gsum[gi..gi + 4], s);
8123            gi += 4;
8124        }
8125        acc
8126    }
8127}
8128
8129/// Blocked 1×4 AVX2 q1g row: each tile's mask is built once and serves
8130/// four activation streams; per stream, the group reduce and float order
8131/// of `dot_q1g_row_avx2` (so batch == matvec bit for bit).
8132#[cfg(target_arch = "x86_64")]
8133#[target_feature(enable = "avx2")]
8134unsafe fn dot_q1g_row_1x4_avx2(
8135    bytes: &[u8],
8136    sc_off: usize,
8137    r: usize,
8138    gpr: usize,
8139    xs: [&[i8]; 4],
8140    gsums: [&[i32]; 4],
8141) -> [f32; 4] {
8142    // SAFETY: callers uphold the plane-length and xq/gsum contracts.
8143    unsafe {
8144        use core::arch::x86_64::*;
8145        let expand = _mm256_setr_epi8(
8146            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,
8147            3, 3, 3,
8148        );
8149        let bitsel = _mm256_setr_epi8(
8150            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
8151            -128, 1, 2, 4, 8, 16, 32, 64, -128,
8152        );
8153        let ones8 = _mm256_set1_epi8(1);
8154        let ones16 = _mm256_set1_epi16(1);
8155        let bp = bytes.as_ptr().add(r * gpr * 4);
8156        let mut acc = [0f32; 4];
8157        let mut gi = 0;
8158        while gi < gpr {
8159            let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
8160            // d[stream][tile]
8161            let mut d = [[_mm256_setzero_si256(); 4]; 4];
8162            for t in 0..4 {
8163                let bits = (bp.add((gi + t) * 4) as *const u32).read_unaligned();
8164                let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
8165                let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
8166                for (k, xq) in xs.iter().enumerate() {
8167                    let x = _mm256_loadu_si256(
8168                        xq.as_ptr().add((gi + t) * GROUP_SIZE) as *const __m256i
8169                    );
8170                    let sel = _mm256_and_si256(x, mask);
8171                    d[k][t] = _mm256_madd_epi16(_mm256_maddubs_epi16(ones8, sel), ones16);
8172                }
8173            }
8174            for k in 0..4 {
8175                let m4 = q1g_hsum4(d[k][0], d[k][1], d[k][2], d[k][3]);
8176                acc[k] = q1g_group_acc(acc[k], m4, &gsums[k][gi..gi + 4], s);
8177            }
8178            gi += 4;
8179        }
8180        acc
8181    }
8182}
8183
8184/// VNNI twin of `dot_q1g_row_1x4_avx2`.
8185#[cfg(target_arch = "x86_64")]
8186#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
8187unsafe fn dot_q1g_row_1x4_vnni(
8188    bytes: &[u8],
8189    sc_off: usize,
8190    r: usize,
8191    gpr: usize,
8192    xs: [&[i8]; 4],
8193    gsums: [&[i32]; 4],
8194) -> [f32; 4] {
8195    // SAFETY: callers uphold the plane-length and xq/gsum contracts.
8196    unsafe {
8197        use core::arch::x86_64::*;
8198        let expand = _mm256_setr_epi8(
8199            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,
8200            3, 3, 3,
8201        );
8202        let bitsel = _mm256_setr_epi8(
8203            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
8204            -128, 1, 2, 4, 8, 16, 32, 64, -128,
8205        );
8206        let ones8 = _mm256_set1_epi8(1);
8207        let bp = bytes.as_ptr().add(r * gpr * 4);
8208        let mut acc = [0f32; 4];
8209        let mut gi = 0;
8210        while gi < gpr {
8211            let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
8212            let mut d = [[_mm256_setzero_si256(); 4]; 4];
8213            for t in 0..4 {
8214                let bits = (bp.add((gi + t) * 4) as *const u32).read_unaligned();
8215                let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
8216                let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
8217                for (k, xq) in xs.iter().enumerate() {
8218                    let x = _mm256_loadu_si256(
8219                        xq.as_ptr().add((gi + t) * GROUP_SIZE) as *const __m256i
8220                    );
8221                    d[k][t] = _mm256_dpbusd_epi32(
8222                        _mm256_setzero_si256(),
8223                        ones8,
8224                        _mm256_and_si256(x, mask),
8225                    );
8226                }
8227            }
8228            for k in 0..4 {
8229                let m4 = q1g_hsum4(d[k][0], d[k][1], d[k][2], d[k][3]);
8230                acc[k] = q1g_group_acc(acc[k], m4, &gsums[k][gi..gi + 4], s);
8231            }
8232            gi += 4;
8233        }
8234        acc
8235    }
8236}
8237
8238/// SDOT q1g row: the `dot_q1_row_sdot` math on the plane layout. One
8239/// aligned 16-byte load carries a whole 128-group's signs (q1 needs two
8240/// pair loads and a TBL gather of four scales for the same weights), TBL
8241/// replicates each sign byte across 8 lanes for vtst, and the group's ONE
8242/// scale converts once (FCVTL, exact) and broadcasts. Lane order, the
8243/// vpaddq tree, the fma into the four lane accumulators and the tail are
8244/// the q1 kernel's — including its quirk: q1's pair load reads 4 bytes
8245/// past the 4th tile, so the payload's final four tiles (last row) take
8246/// the scalar tail there. q1g has no overread, but takes the same tail on
8247/// the last row to stay bit-identical to q1.
8248#[cfg(target_arch = "aarch64")]
8249#[target_feature(enable = "neon,dotprod")]
8250#[allow(clippy::too_many_arguments)]
8251unsafe fn dot_q1g_row_sdot(
8252    bytes: &[u8],
8253    sc_off: usize,
8254    rows: usize,
8255    r: usize,
8256    gpr: usize,
8257    xq: &[i8],
8258    gsum: &[i32],
8259) -> f32 {
8260    // SAFETY: callers uphold the plane-length and xq/gsum contracts
8261    // (gpr % 4 == 0, xq.len() == gpr·32, gsum.len() == gpr).
8262    unsafe {
8263        use core::arch::aarch64::*;
8264        use core::arch::asm;
8265        const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
8266        // Tile t of the group = sign bytes 4t..4t+4 of the 16-byte load.
8267        const IW: [[u8; 16]; 8] = [
8268            [0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1],
8269            [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3],
8270            [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5],
8271            [6, 6, 6, 6, 6, 6, 6, 6, 7, 7, 7, 7, 7, 7, 7, 7],
8272            [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9],
8273            [
8274                10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
8275            ],
8276            [
8277                12, 12, 12, 12, 12, 12, 12, 12, 13, 13, 13, 13, 13, 13, 13, 13,
8278            ],
8279            [
8280                14, 14, 14, 14, 14, 14, 14, 14, 15, 15, 15, 15, 15, 15, 15, 15,
8281            ],
8282        ];
8283        let m = vld1q_u8(MASKS.as_ptr());
8284        let iw: [uint8x16_t; 8] = core::array::from_fn(|i| vld1q_u8(IW[i].as_ptr()));
8285        macro_rules! tile_dot {
8286            ($w0:expr, $w1:expr, $x:expr) => {{
8287                let x0 = vld1q_s8($x);
8288                let x1 = vld1q_s8($x.add(16));
8289                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
8290                asm!(
8291                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
8292                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
8293                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
8294                    w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
8295                    options(pure, nomem, nostack),
8296                );
8297                vaddq_s32(a0, a1)
8298            }};
8299        }
8300        macro_rules! tbl_mask {
8301            ($ld:expr, $i:expr) => {
8302                vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, iw[$i]), m))
8303            };
8304        }
8305        let bp = bytes.as_ptr().add(r * gpr * 4);
8306        let xp = xq.as_ptr();
8307        let gp = gsum.as_ptr();
8308        let vec_end = if r + 1 == rows { gpr - 4 } else { gpr };
8309        let mut accv = vdupq_n_f32(0.0);
8310        let mut gi = 0;
8311        while gi + 4 <= vec_end {
8312            let ld = vld1q_u8(bp.add(gi * 4));
8313            let d0 = tile_dot!(tbl_mask!(ld, 0), tbl_mask!(ld, 1), xp.add(gi * GROUP_SIZE));
8314            let d1 = tile_dot!(
8315                tbl_mask!(ld, 2),
8316                tbl_mask!(ld, 3),
8317                xp.add((gi + 1) * GROUP_SIZE)
8318            );
8319            let d2 = tile_dot!(
8320                tbl_mask!(ld, 4),
8321                tbl_mask!(ld, 5),
8322                xp.add((gi + 2) * GROUP_SIZE)
8323            );
8324            let d3 = tile_dot!(
8325                tbl_mask!(ld, 6),
8326                tbl_mask!(ld, 7),
8327                xp.add((gi + 3) * GROUP_SIZE)
8328            );
8329            let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
8330            let g = vld1q_s32(gp.add(gi));
8331            let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
8332            let h = vdup_n_u16(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
8333            let scf: float32x4_t;
8334            asm!(
8335                "fcvtl {o:v}.4s, {i:v}.4h",
8336                o = out(vreg) scf, i = in(vreg) h,
8337                options(pure, nomem, nostack),
8338            );
8339            accv = vfmaq_f32(accv, vcvtq_f32_s32(dots), scf);
8340            gi += 4;
8341        }
8342        let mut acc = vaddvq_f32(accv);
8343        while gi < gpr {
8344            let t = bp.add(gi * 4);
8345            let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
8346            let v0 = vcombine_u8(vdup_n_u8(*t), vdup_n_u8(*t.add(1)));
8347            let v1 = vcombine_u8(vdup_n_u8(*t.add(2)), vdup_n_u8(*t.add(3)));
8348            let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
8349            let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
8350            let d = vaddvq_s32(tile_dot!(w0, w1, xp.add(gi * GROUP_SIZE)));
8351            acc += (-(2 * d + *gp.add(gi))) as f32 * s;
8352            gi += 1;
8353        }
8354        acc
8355    }
8356}
8357
8358/// Blocked q1g 1×4: one unpack of the group serves four activation
8359/// streams; per stream bit-identical to `dot_q1g_row_sdot` (and so to the
8360/// q1 1×4 kernel, same last-row tail).
8361#[cfg(target_arch = "aarch64")]
8362#[target_feature(enable = "neon,dotprod")]
8363#[allow(clippy::too_many_arguments)]
8364unsafe fn dot_q1g_row_1x4_sdot(
8365    bytes: &[u8],
8366    sc_off: usize,
8367    rows: usize,
8368    r: usize,
8369    gpr: usize,
8370    xs: [&[i8]; 4],
8371    gs: [&[i32]; 4],
8372) -> [f32; 4] {
8373    // SAFETY: same contracts as `dot_q1g_row_sdot`, ×4.
8374    unsafe {
8375        use core::arch::aarch64::*;
8376        use core::arch::asm;
8377        const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
8378        const IW: [[u8; 16]; 8] = [
8379            [0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1],
8380            [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3],
8381            [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5],
8382            [6, 6, 6, 6, 6, 6, 6, 6, 7, 7, 7, 7, 7, 7, 7, 7],
8383            [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9],
8384            [
8385                10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
8386            ],
8387            [
8388                12, 12, 12, 12, 12, 12, 12, 12, 13, 13, 13, 13, 13, 13, 13, 13,
8389            ],
8390            [
8391                14, 14, 14, 14, 14, 14, 14, 14, 15, 15, 15, 15, 15, 15, 15, 15,
8392            ],
8393        ];
8394        let m = vld1q_u8(MASKS.as_ptr());
8395        let iw: [uint8x16_t; 8] = core::array::from_fn(|i| vld1q_u8(IW[i].as_ptr()));
8396        macro_rules! sdot2 {
8397            ($w0:expr, $w1:expr, $x:expr) => {{
8398                let x0 = vld1q_s8($x);
8399                let x1 = vld1q_s8($x.add(16));
8400                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
8401                asm!(
8402                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
8403                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
8404                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
8405                    w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
8406                    options(pure, nomem, nostack),
8407                );
8408                vaddq_s32(a0, a1)
8409            }};
8410        }
8411        let bp = bytes.as_ptr().add(r * gpr * 4);
8412        let vec_end = if r + 1 == rows { gpr - 4 } else { gpr };
8413        let mut accv = [vdupq_n_f32(0.0); 4];
8414        let mut gi = 0;
8415        while gi + 4 <= vec_end {
8416            let ld = vld1q_u8(bp.add(gi * 4));
8417            // Unpack ONCE — eight ±mask vectors serve all four streams.
8418            let w: [int8x16_t; 8] =
8419                core::array::from_fn(|i| vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld, iw[i]), m)));
8420            let h = vdup_n_u16(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
8421            let scf: float32x4_t;
8422            asm!(
8423                "fcvtl {o:v}.4s, {i:v}.4h",
8424                o = out(vreg) scf, i = in(vreg) h,
8425                options(pure, nomem, nostack),
8426            );
8427            for k in 0..4 {
8428                let xp = xs[k].as_ptr();
8429                let d0 = sdot2!(w[0], w[1], xp.add(gi * GROUP_SIZE));
8430                let d1 = sdot2!(w[2], w[3], xp.add((gi + 1) * GROUP_SIZE));
8431                let d2 = sdot2!(w[4], w[5], xp.add((gi + 2) * GROUP_SIZE));
8432                let d3 = sdot2!(w[6], w[7], xp.add((gi + 3) * GROUP_SIZE));
8433                let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
8434                let g = vld1q_s32(gs[k].as_ptr().add(gi));
8435                let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
8436                accv[k] = vfmaq_f32(accv[k], vcvtq_f32_s32(dots), scf);
8437            }
8438            gi += 4;
8439        }
8440        let mut acc = [
8441            vaddvq_f32(accv[0]),
8442            vaddvq_f32(accv[1]),
8443            vaddvq_f32(accv[2]),
8444            vaddvq_f32(accv[3]),
8445        ];
8446        while gi < gpr {
8447            let t = bp.add(gi * 4);
8448            let sc = f16_to_f32(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
8449            let v0 = vcombine_u8(vdup_n_u8(*t), vdup_n_u8(*t.add(1)));
8450            let v1 = vcombine_u8(vdup_n_u8(*t.add(2)), vdup_n_u8(*t.add(3)));
8451            let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
8452            let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
8453            for k in 0..4 {
8454                let d = vaddvq_s32(sdot2!(w0, w1, xs[k].as_ptr().add(gi * GROUP_SIZE)));
8455                acc[k] += (-(2 * d + *gs[k].as_ptr().add(gi))) as f32 * sc;
8456            }
8457            gi += 1;
8458        }
8459        acc
8460    }
8461}
8462
8463/// (weight ±1, scale) of one q1g element — the exact outlier term.
8464#[inline]
8465fn q1g_outlier(bytes: &[u8], sc_off: usize, r: usize, gpr: usize, j: usize) -> (f32, f32) {
8466    let t = r * gpr + j / GROUP_SIZE;
8467    let k = j % GROUP_SIZE;
8468    let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, t));
8469    let bit = (bytes[t * 4 + k / 8] >> (k % 8)) & 1;
8470    ((bit as i32 * 2 - 1) as f32, s)
8471}
8472
8473/// Exact scalar q1g row (CMF_SDOT=0 contract) — `q1_row_exact`'s order.
8474#[inline]
8475fn q1g_row_exact(bytes: &[u8], sc_off: usize, r: usize, gpr: usize, x: &[f32]) -> f32 {
8476    let mut acc = 0f32;
8477    for gi in 0..gpr {
8478        let t = r * gpr + gi;
8479        let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, t));
8480        let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
8481        let mut ga = 0f32;
8482        for (j, &b) in bytes[t * 4..t * 4 + 4].iter().enumerate() {
8483            for k in 0..8 {
8484                ga += (((b >> k) & 1) as f32 * 2.0 - 1.0) * xg[j * 8 + k];
8485            }
8486        }
8487        acc += ga * s;
8488    }
8489    acc
8490}
8491
8492/// q1t per-row overlay locator. After the base (`base_len`) come
8493/// `[u32 row_ptr[rows+1]]` then `[(u16 col, f16 val)]` grouped by row (row
8494/// `r`'s entries are `[row_ptr[r], row_ptr[r+1])`). Returns
8495/// `(row_ptr offset, entries offset, present)`.
8496fn q1t_overlay(bytes: &[u8], base_len: usize, rows: usize) -> (usize, usize, bool) {
8497    let entries = base_len + (rows + 1) * 4;
8498    (base_len, entries, entries <= bytes.len())
8499}
8500
8501/// Read `row_ptr[r]` from the overlay's prefix-sum table.
8502#[inline]
8503fn q1t_rowptr(bytes: &[u8], rp_off: usize, r: usize) -> usize {
8504    let o = rp_off + r * 4;
8505    u32::from_le_bytes([bytes[o], bytes[o + 1], bytes[o + 2], bytes[o + 3]]) as usize
8506}
8507
8508/// Byte → the 5 ternary signs it packs `{−1,0,+1}` as f32, precomputed so
8509/// decoding a q1t code is a table load, not the base-3 divide/modulo per
8510/// weight (division is ~20–40× the cost of a load). Built at compile time.
8511const SIGN5: [[f32; 5]; 256] = {
8512    let mut lut = [[0.0f32; 5]; 256];
8513    let pow3 = [1u16, 3, 9, 27, 81];
8514    let mut byte = 0usize;
8515    while byte < 256 {
8516        let mut i = 0usize;
8517        while i < 5 {
8518            let code = (byte as u16 / pow3[i]) % 3;
8519            lut[byte][i] = if code == 1 {
8520                1.0
8521            } else if code == 2 {
8522                -1.0
8523            } else {
8524                0.0
8525            };
8526            i += 1;
8527        }
8528        byte += 1;
8529    }
8530    lut
8531};
8532
8533/// Same table, as i8 signs — the operand for the int8 SDOT base kernel.
8534const SIGN5_I8: [[i8; 5]; 256] = {
8535    let mut lut = [[0i8; 5]; 256];
8536    let pow3 = [1u16, 3, 9, 27, 81];
8537    let mut byte = 0usize;
8538    while byte < 256 {
8539        let mut i = 0usize;
8540        while i < 5 {
8541            let code = (byte as u16 / pow3[i]) % 3;
8542            lut[byte][i] = if code == 1 {
8543                1
8544            } else if code == 2 {
8545                -1
8546            } else {
8547                0
8548            };
8549            i += 1;
8550        }
8551        byte += 1;
8552    }
8553    lut
8554};
8555
8556/// The same 5 i8 signs packed into a u64 (`[s0 s1 s2 s3 s4 0 0 0]`, LE) so the
8557/// group unpack is 7 unaligned u64 stores at offsets 0,5,10,…,30 instead of
8558/// six 5-byte copies + LUT indexing — each store's trailing zeros are fixed by
8559/// the next store, and the last one runs 6 B past the 32nd weight (the unpack
8560/// buffer is padded to 40). This is the decode/prefill hot inner op.
8561const SIGN5_U64: [u64; 256] = {
8562    let mut lut = [0u64; 256];
8563    let pow3 = [1u16, 3, 9, 27, 81];
8564    let mut byte = 0usize;
8565    while byte < 256 {
8566        let mut v = 0u64;
8567        let mut i = 0usize;
8568        while i < 5 {
8569            let code = (byte as u16 / pow3[i]) % 3;
8570            let s: u8 = if code == 1 {
8571                1
8572            } else if code == 2 {
8573                0xFF
8574            } else {
8575                0
8576            };
8577            v |= (s as u64) << (i * 8);
8578            i += 1;
8579        }
8580        lut[byte] = v;
8581        byte += 1;
8582    }
8583    lut
8584};
8585
8586/// Ternary base weight at `(row r, col j)` = `sign(code)·s_group`. Used to add
8587/// back activation-outlier columns, whose `x` was zeroed for the int8 bulk dot
8588/// (`split_act`). At a weight-outlier position the code is 0, so this is 0 and
8589/// the overlay correction owns that column — no double counting.
8590#[inline]
8591fn q1t_base_weight(bytes: &[u8], r: usize, gpr: usize, j: usize) -> f32 {
8592    const TILE: usize = cortiq_core::quant::Q1T_TILE;
8593    let off = (r * gpr + j / GROUP_SIZE) * TILE;
8594    let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
8595    let within = j % GROUP_SIZE;
8596    SIGN5[bytes[off + 2 + within / 5] as usize][within % 5] * s
8597}
8598
8599/// One 32-group int8 dot via two SDOTs. Bit-exact vs the scalar i8 sum
8600/// (integer accumulation is order-independent).
8601#[cfg(target_arch = "aarch64")]
8602#[target_feature(enable = "neon,dotprod")]
8603#[inline]
8604unsafe fn sdot32_i8(w: *const i8, x: *const i8) -> i32 {
8605    // SAFETY: caller guarantees 32 readable i8 at each pointer.
8606    unsafe {
8607        use core::arch::aarch64::*;
8608        use core::arch::asm;
8609        let w0 = vld1q_s8(w);
8610        let w1 = vld1q_s8(w.add(16));
8611        let x0 = vld1q_s8(x);
8612        let x1 = vld1q_s8(x.add(16));
8613        let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
8614        asm!(
8615            "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
8616            "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
8617            a0 = inout(vreg) a0, a1 = inout(vreg) a1,
8618            w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
8619            options(pure, nomem, nostack),
8620        );
8621        vaddvq_s32(vaddq_s32(a0, a1))
8622    }
8623}
8624
8625/// One 32-group int8 dot via AVX2: signed·signed as `maddubs(|w|, sign(x,w))`
8626/// then `madd` and a horizontal reduce (the same idiom as `dot_q4t_row_avx2`).
8627#[cfg(target_arch = "x86_64")]
8628#[target_feature(enable = "avx2")]
8629#[inline]
8630unsafe fn i8dot32_avx2(w: *const i8, x: *const i8) -> i32 {
8631    // SAFETY: caller guarantees 32 readable i8 at each pointer.
8632    unsafe {
8633        use core::arch::x86_64::*;
8634        let wv = _mm256_loadu_si256(w as *const __m256i);
8635        let xv = _mm256_loadu_si256(x as *const __m256i);
8636        let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
8637        let d = _mm256_madd_epi16(p16, _mm256_set1_epi16(1));
8638        let hi128 = _mm256_extracti128_si256::<1>(d);
8639        let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
8640        let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8641        let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8642        _mm_cvtsi128_si32(s32)
8643    }
8644}
8645
8646/// Unpack one q1t group's base-3 codes into 32 i8 signs via 7 unaligned u64
8647/// stores (see `SIGN5_U64`). `dst` MUST have ≥ 40 bytes: the 7th store writes
8648/// `dst[30..38]`. Stores go in order so each one's trailing zeros are
8649/// overwritten by the next; the final 6 padding bytes are unused by the dot.
8650#[inline]
8651fn q1t_unpack_group_i8(codes: *const u8, dst: &mut [i8]) {
8652    debug_assert!(dst.len() >= 40);
8653    // SAFETY: codes points at 7 readable bytes; dst has ≥ 40 bytes so every
8654    // 8-byte store at offset bi*5 (bi ≤ 6 → ≤ 30) stays in bounds.
8655    unsafe {
8656        let p = dst.as_mut_ptr();
8657        for bi in 0..7 {
8658            core::ptr::write_unaligned(
8659                p.add(bi * 5) as *mut u64,
8660                SIGN5_U64[*codes.add(bi) as usize],
8661            );
8662        }
8663    }
8664}
8665
8666/// One 32-group int8 dot, arch-dispatched (the matmat inner loop, where the
8667/// row's signs are unpacked once and dotted against every batch input).
8668/// Callers are gated by `a8w8_enabled()`, so the target-feature arms are
8669/// reachable; the scalar arm is a non-SIMD-arch fallback.
8670#[inline]
8671fn q1t_i8dot32(w: *const i8, x: *const i8) -> i32 {
8672    #[cfg(target_arch = "aarch64")]
8673    unsafe {
8674        return sdot32_i8(w, x);
8675    }
8676    #[cfg(target_arch = "x86_64")]
8677    unsafe {
8678        return i8dot32_avx2(w, x);
8679    }
8680    #[allow(unreachable_code)]
8681    unsafe {
8682        let mut s = 0i32;
8683        for k in 0..GROUP_SIZE {
8684            s += *w.add(k) as i32 * *x.add(k) as i32;
8685        }
8686        s
8687    }
8688}
8689
8690#[inline]
8691unsafe fn q1t_unpack_reg_u64s(codes: *const u8) -> (u64, u64, u64, u64) {
8692    let (s0, s1, s2, s3, s4, s5, s6) = unsafe {
8693        (
8694            SIGN5_U64[*codes as usize],
8695            SIGN5_U64[*codes.add(1) as usize],
8696            SIGN5_U64[*codes.add(2) as usize],
8697            SIGN5_U64[*codes.add(3) as usize],
8698            SIGN5_U64[*codes.add(4) as usize],
8699            SIGN5_U64[*codes.add(5) as usize],
8700            SIGN5_U64[*codes.add(6) as usize],
8701        )
8702    };
8703
8704    let u0 = s0 | (s1 << 40);
8705    let u1 = (s1 >> 24) | (s2 << 16) | (s3 << 56);
8706    let u2 = (s3 >> 8) | (s4 << 32);
8707    let u3 = (s4 >> 32) | (s5 << 8) | (s6 << 48);
8708
8709    (u0, u1, u2, u3)
8710}
8711
8712/// One q1t row's int8 base dot: `Σ_group s·dot(signs, xq)` (before the shared
8713/// `sx`). Direct register unpacking (zero stack stores/loads, no STLF stalls).
8714/// ARM SDOT.
8715#[cfg(target_arch = "aarch64")]
8716#[target_feature(enable = "neon,dotprod")]
8717unsafe fn q1t_dot_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
8718    use core::arch::aarch64::*;
8719    use core::arch::asm;
8720    unsafe {
8721        const TILE: usize = cortiq_core::quant::Q1T_TILE;
8722        let mut acc = 0f32;
8723        let bytes_ptr = bytes.as_ptr();
8724        let xq_ptr = xq.as_ptr();
8725        let row_off = r * gpr * TILE;
8726
8727        let gpr2 = gpr & !1;
8728        let mut gi = 0;
8729        while gi < gpr2 {
8730            let off0 = row_off + gi * TILE;
8731            let off1 = off0 + TILE;
8732            let s0 = f16_to_f32(u16::from_le_bytes([
8733                *bytes_ptr.add(off0),
8734                *bytes_ptr.add(off0 + 1),
8735            ]));
8736            let s1 = f16_to_f32(u16::from_le_bytes([
8737                *bytes_ptr.add(off1),
8738                *bytes_ptr.add(off1 + 1),
8739            ]));
8740
8741            let (u0_0, u1_0, u2_0, u3_0) = q1t_unpack_reg_u64s(bytes_ptr.add(off0 + 2));
8742            let (u0_1, u1_1, u2_1, u3_1) = q1t_unpack_reg_u64s(bytes_ptr.add(off1 + 2));
8743
8744            let w0_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_0), vcreate_u64(u1_0)));
8745            let w1_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_0), vcreate_u64(u3_0)));
8746            let w0_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_1), vcreate_u64(u1_1)));
8747            let w1_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_1), vcreate_u64(u3_1)));
8748
8749            let x0_0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE));
8750            let x1_0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE + 16));
8751            let x0_1 = vld1q_s8(xq_ptr.add((gi + 1) * GROUP_SIZE));
8752            let x1_1 = vld1q_s8(xq_ptr.add((gi + 1) * GROUP_SIZE + 16));
8753
8754            let (mut a0_0, mut a1_0) = (vdupq_n_s32(0), vdupq_n_s32(0));
8755            let (mut a0_1, mut a1_1) = (vdupq_n_s32(0), vdupq_n_s32(0));
8756            asm!(
8757                "sdot {a0_0:v}.4s, {w0_0:v}.16b, {x0_0:v}.16b",
8758                "sdot {a1_0:v}.4s, {w1_0:v}.16b, {x1_0:v}.16b",
8759                "sdot {a0_1:v}.4s, {w0_1:v}.16b, {x0_1:v}.16b",
8760                "sdot {a1_1:v}.4s, {w1_1:v}.16b, {x1_1:v}.16b",
8761                a0_0 = inout(vreg) a0_0, a1_0 = inout(vreg) a1_0,
8762                a0_1 = inout(vreg) a0_1, a1_1 = inout(vreg) a1_1,
8763                w0_0 = in(vreg) w0_0, x0_0 = in(vreg) x0_0, w1_0 = in(vreg) w1_0, x1_0 = in(vreg) x1_0,
8764                w0_1 = in(vreg) w0_1, x0_1 = in(vreg) x0_1, w1_1 = in(vreg) w1_1, x1_1 = in(vreg) x1_1,
8765                options(pure, nomem, nostack),
8766            );
8767            let d0 = vaddvq_s32(vaddq_s32(a0_0, a1_0));
8768            let d1 = vaddvq_s32(vaddq_s32(a0_1, a1_1));
8769            acc += d0 as f32 * s0 + d1 as f32 * s1;
8770            gi += 2;
8771        }
8772
8773        if gi < gpr {
8774            let off = row_off + gi * TILE;
8775            let s = f16_to_f32(u16::from_le_bytes([
8776                *bytes_ptr.add(off),
8777                *bytes_ptr.add(off + 1),
8778            ]));
8779            let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
8780            let w0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0), vcreate_u64(u1)));
8781            let w1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2), vcreate_u64(u3)));
8782            let x0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE));
8783            let x1 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE + 16));
8784            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
8785            asm!(
8786                "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
8787                "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
8788                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
8789                w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
8790                options(pure, nomem, nostack),
8791            );
8792            let d = vaddvq_s32(vaddq_s32(a0, a1));
8793            acc += d as f32 * s;
8794        }
8795        acc
8796    }
8797}
8798
8799/// x86 AVX2 mirror of `q1t_dot_row_sdot` (maddubs int8 dot per group).
8800#[cfg(target_arch = "x86_64")]
8801#[target_feature(enable = "avx2")]
8802unsafe fn q1t_dot_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
8803    use core::arch::x86_64::*;
8804    unsafe {
8805        const TILE: usize = cortiq_core::quant::Q1T_TILE;
8806        let mut acc = 0f32;
8807        let bytes_ptr = bytes.as_ptr();
8808        let xq_ptr = xq.as_ptr();
8809        let row_off = r * gpr * TILE;
8810
8811        let ones = _mm256_set1_epi16(1);
8812        for gi in 0..gpr {
8813            let off = row_off + gi * TILE;
8814            let s = f16_to_f32(u16::from_le_bytes([
8815                *bytes_ptr.add(off),
8816                *bytes_ptr.add(off + 1),
8817            ]));
8818            let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
8819            let wv = _mm256_set_epi64x(u3 as i64, u2 as i64, u1 as i64, u0 as i64);
8820            let xv = _mm256_loadu_si256(xq_ptr.add(gi * GROUP_SIZE) as *const __m256i);
8821            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
8822            let d256 = _mm256_madd_epi16(p16, ones);
8823            let d128 = _mm_add_epi32(
8824                _mm256_castsi256_si128(d256),
8825                _mm256_extracti128_si256(d256, 1),
8826            );
8827            let d64 = _mm_add_epi32(d128, _mm_shuffle_epi32(d128, 0xee));
8828            let d32 = _mm_cvtsi128_si32(_mm_add_epi32(d64, _mm_shuffle_epi32(d64, 0x55)));
8829            acc += d32 as f32 * s;
8830        }
8831        acc
8832    }
8833}
8834
8835/// VNNI twin of `q1t_dot_row_avx2` (see `dpbusd_hsum`).
8836#[cfg(target_arch = "x86_64")]
8837#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
8838unsafe fn q1t_dot_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
8839    use core::arch::x86_64::*;
8840    // SAFETY: same tile/xq contracts as `q1t_dot_row_avx2`.
8841    unsafe {
8842        const TILE: usize = cortiq_core::quant::Q1T_TILE;
8843        let mut acc = 0f32;
8844        let bytes_ptr = bytes.as_ptr();
8845        let xq_ptr = xq.as_ptr();
8846        let row_off = r * gpr * TILE;
8847        for gi in 0..gpr {
8848            let off = row_off + gi * TILE;
8849            let s = f16_to_f32(u16::from_le_bytes([
8850                *bytes_ptr.add(off),
8851                *bytes_ptr.add(off + 1),
8852            ]));
8853            let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
8854            let wv = _mm256_set_epi64x(u3 as i64, u2 as i64, u1 as i64, u0 as i64);
8855            let xv = _mm256_loadu_si256(xq_ptr.add(gi * GROUP_SIZE) as *const __m256i);
8856            let d = dpbusd_hsum(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
8857            acc += d as f32 * s;
8858        }
8859        acc
8860    }
8861}
8862
8863/// Per-row int8 base dot, dispatched once per row (matvec decode hot path).
8864/// Callers are gated by `a8w8_enabled()`, so the target-feature kernels are
8865/// reachable.
8866#[inline]
8867fn q1t_dot_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
8868    #[cfg(target_arch = "aarch64")]
8869    unsafe {
8870        return q1t_dot_row_sdot(bytes, r, gpr, xq);
8871    }
8872    #[cfg(target_arch = "x86_64")]
8873    unsafe {
8874        if vnni_tiles_enabled() {
8875            return q1t_dot_row_vnni(bytes, r, gpr, xq);
8876        }
8877        return q1t_dot_row_avx2(bytes, r, gpr, xq);
8878    }
8879    #[allow(unreachable_code)]
8880    {
8881        const TILE: usize = cortiq_core::quant::Q1T_TILE;
8882        let mut acc = 0f32;
8883        let mut sg = [0i8; GROUP_SIZE + 8]; // +8 slack for the u64-store unpack
8884        for gi in 0..gpr {
8885            let off = (r * gpr + gi) * TILE;
8886            let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
8887            q1t_unpack_group_i8(bytes.as_ptr().wrapping_add(off + 2), &mut sg);
8888            let mut d = 0i32;
8889            for k in 0..GROUP_SIZE {
8890                d += sg[k] as i32 * xq[gi * GROUP_SIZE + k] as i32;
8891            }
8892            acc += d as f32 * s;
8893        }
8894        acc
8895    }
8896}
8897
8898/// Σ over a row's outliers of `value·x[col]` — the correction that adds the
8899/// overlay's exact weights on top of the base dot. INVARIANT: the encoder
8900/// writes ternary code 0 at every outlier position (`quantize_q1t`), so the
8901/// base contributes nothing there and this is a plain `value·x`, not
8902/// `(value − base)·x` — no scattered per-outlier scale read. Row `r`'s entries
8903/// are the contiguous slice `[row_ptr[r], row_ptr[r+1])`, so no binary search.
8904fn q1t_row_outlier_correction(
8905    bytes: &[u8],
8906    r: usize,
8907    rp_off: usize,
8908    entries_off: usize,
8909    has_ov: bool,
8910    x: &[f32],
8911) -> f32 {
8912    if !has_ov {
8913        return 0.0;
8914    }
8915    let (c0, c1) = (
8916        q1t_rowptr(bytes, rp_off, r),
8917        q1t_rowptr(bytes, rp_off, r + 1),
8918    );
8919    let mut corr = 0f32;
8920    for p in c0..c1 {
8921        let e = entries_off + p * 4;
8922        let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
8923        let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
8924        corr += val * x[col];
8925    }
8926    corr
8927}
8928
8929/// Dequantize one q1t row into `buf[..cols]` via the sign LUT (no division),
8930/// then apply the row's outliers (its `[row_ptr[r], row_ptr[r+1])` slice).
8931/// Used by the batched (prefill) path where the decode amortizes over the batch.
8932fn q1t_dequant_row(
8933    bytes: &[u8],
8934    r: usize,
8935    gpr: usize,
8936    rp_off: usize,
8937    entries_off: usize,
8938    has_ov: bool,
8939    buf: &mut [f32],
8940) {
8941    const TILE: usize = cortiq_core::quant::Q1T_TILE;
8942    for g in 0..gpr {
8943        let off = (r * gpr + g) * TILE;
8944        let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
8945        let codes = &bytes[off + 2..off + TILE];
8946        let bc = g * GROUP_SIZE;
8947        // 6 full bytes (30 codes) + a 7th byte holding the last 2.
8948        for bi in 0..6 {
8949            let lut = &SIGN5[codes[bi] as usize];
8950            let d = &mut buf[bc + bi * 5..bc + bi * 5 + 5];
8951            for i in 0..5 {
8952                d[i] = lut[i] * s;
8953            }
8954        }
8955        let lut = &SIGN5[codes[6] as usize];
8956        buf[bc + 30] = lut[0] * s;
8957        buf[bc + 31] = lut[1] * s;
8958    }
8959    if !has_ov {
8960        return;
8961    }
8962    let (c0, c1) = (
8963        q1t_rowptr(bytes, rp_off, r),
8964        q1t_rowptr(bytes, rp_off, r + 1),
8965    );
8966    for p in c0..c1 {
8967        let e = entries_off + p * 4;
8968        let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
8969        buf[col] = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
8970    }
8971}
8972
8973/// Add the sparse outlier overlay onto a base dot already in `out` (the GPU
8974/// computes the ternary base; the overlay stays on the CPU — its entries are
8975/// few and its per-row gather doesn't vectorize on the GPU). Row-parallel.
8976fn q1t_add_overlay(
8977    bytes: &[u8],
8978    x: &[f32],
8979    rows: usize,
8980    cols: usize,
8981    out: &mut [f32],
8982    pool: Option<&Pool>,
8983) {
8984    const TILE: usize = cortiq_core::quant::Q1T_TILE;
8985    let gpr = cols / GROUP_SIZE;
8986    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
8987    if !has_ov {
8988        return;
8989    }
8990    let out_addr = SendMut(out.as_mut_ptr());
8991    let run = move |start: usize, end: usize| {
8992        for r in start..end {
8993            let corr = q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
8994            // SAFETY: disjoint rows; add onto the base the GPU already wrote.
8995            unsafe { *out_addr.at(r) += corr };
8996        }
8997    };
8998    dispatch_rows(pool, rows, &run);
8999}
9000
9001/// Q1T row range via the A8W8 int8 path — shared activation split,
9002/// per-row: base SDOT dot + outlier correction + overlay.
9003#[allow(clippy::too_many_arguments)]
9004fn q1t_range_a8w8(
9005    bytes: &[u8],
9006    gpr: usize,
9007    rp_off: usize,
9008    ent_off: usize,
9009    has_ov: bool,
9010    act: &SplitAct,
9011    x: &[f32],
9012    out: SendMut,
9013    start: usize,
9014    end: usize,
9015) {
9016    for r in start..end {
9017        let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
9018        for &(j, xv) in &act.outliers {
9019            acc += q1t_base_weight(bytes, r, gpr, j) * xv;
9020        }
9021        acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
9022        // SAFETY: disjoint row ranges per worker.
9023        unsafe { *out.at(r) = acc };
9024    }
9025}
9026
9027/// Q1T row range via the f32 path (no SDOT) — for matvec_many batched
9028/// dispatch when a8w8 is unavailable.
9029#[allow(clippy::too_many_arguments)]
9030fn q1t_range_f32_batch(
9031    bytes: &[u8],
9032    gpr: usize,
9033    rp_off: usize,
9034    ent_off: usize,
9035    has_ov: bool,
9036    x: &[f32],
9037    out: SendMut,
9038    start: usize,
9039    end: usize,
9040) {
9041    const TILE: usize = cortiq_core::quant::Q1T_TILE;
9042    let mut sg = [0f32; GROUP_SIZE];
9043    for r in start..end {
9044        let mut acc = 0f32;
9045        for g in 0..gpr {
9046            let off = (r * gpr + g) * TILE;
9047            let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
9048            let codes = &bytes[off + 2..off + TILE];
9049            let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
9050            for bi in 0..6 {
9051                sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
9052            }
9053            let lut = &SIGN5[codes[6] as usize];
9054            sg[30] = lut[0];
9055            sg[31] = lut[1];
9056            let mut gsum = 0f32;
9057            for k in 0..GROUP_SIZE {
9058                gsum += sg[k] * xg[k];
9059            }
9060            acc += s * gsum;
9061        }
9062        acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
9063        // SAFETY: disjoint row ranges per worker.
9064        unsafe { *out.at(r) = acc };
9065    }
9066}
9067
9068/// Ternary (q1t) matvec — decode+dot straight from mmap, one group at a time:
9069/// no per-ROW buffer, no division (the sign LUT), and a tiny per-group sign
9070/// buffer so the 32-wide dot vectorizes. This is the decode hot path.
9071fn q1t_matvec(
9072    bytes: &[u8],
9073    x: &[f32],
9074    rows: usize,
9075    cols: usize,
9076    out: &mut [f32],
9077    pool: Option<&Pool>,
9078) {
9079    debug_assert_eq!(out.len(), rows);
9080    const TILE: usize = cortiq_core::quant::Q1T_TILE;
9081    let gpr = cols / GROUP_SIZE;
9082    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
9083    let out_addr = SendMut(out.as_mut_ptr());
9084    // int8 SDOT base dot (ARM dotprod): ~4× the f32 arithmetic. x → i8 once
9085    // (`split_act`), activation outliers added back exactly in f32, weight
9086    // overlay on top. ARM SDOT / x86 AVX2; CMF_SDOT=0 keeps the exact f32 path.
9087    if a8w8_enabled() {
9088        let act = split_act(x);
9089        let act = &act;
9090        let run = move |start: usize, end: usize| {
9091            for r in start..end {
9092                let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
9093                for &(j, xv) in &act.outliers {
9094                    acc += q1t_base_weight(bytes, r, gpr, j) * xv;
9095                }
9096                acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
9097                // SAFETY: disjoint row ranges per worker.
9098                unsafe { *out_addr.at(r) = acc };
9099            }
9100        };
9101        dispatch_rows(pool, rows, &run);
9102        return;
9103    }
9104    let run = move |start: usize, end: usize| {
9105        // Per-group signs, unpacked contiguously so the dot below is a clean
9106        // 32-wide reduction the autovectorizer turns into f32x4 FMAs — the
9107        // 5-values-per-byte base-3 layout won't SIMD in place.
9108        let mut sg = [0f32; GROUP_SIZE];
9109        for r in start..end {
9110            let mut acc = 0f32;
9111            for g in 0..gpr {
9112                let off = (r * gpr + g) * TILE;
9113                let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
9114                let codes = &bytes[off + 2..off + TILE];
9115                let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
9116                for bi in 0..6 {
9117                    sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
9118                }
9119                let lut = &SIGN5[codes[6] as usize];
9120                sg[30] = lut[0];
9121                sg[31] = lut[1];
9122                let mut gsum = 0f32;
9123                for k in 0..GROUP_SIZE {
9124                    gsum += sg[k] * xg[k];
9125                }
9126                acc += s * gsum;
9127            }
9128            acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
9129            unsafe { *out_addr.at(r) = acc };
9130        }
9131    };
9132    dispatch_rows(pool, rows, &run);
9133}
9134
9135/// Fused-pair twin of `q1t_dot_row_sdot`: ONE register unpack of the
9136/// ternary codes serves BOTH activation streams (the unpack chain is
9137/// the dominant per-row cost — MTP verify pairs paid it twice). Per
9138/// stream the group order and f32 accumulation match the single-row
9139/// kernel exactly, so pair == 2×matvec bit-for-bit.
9140#[cfg(target_arch = "aarch64")]
9141#[target_feature(enable = "neon,dotprod")]
9142unsafe fn q1t_dot_row_sdot2(bytes: &[u8], r: usize, gpr: usize, xa: &[i8], xb: &[i8]) -> [f32; 2] {
9143    use core::arch::aarch64::*;
9144    use core::arch::asm;
9145    // SAFETY: same slice-length contracts as `q1t_dot_row_sdot`, ×2.
9146    unsafe {
9147        const TILE: usize = cortiq_core::quant::Q1T_TILE;
9148        let bytes_ptr = bytes.as_ptr();
9149        let row_off = r * gpr * TILE;
9150        let xp = [xa.as_ptr(), xb.as_ptr()];
9151        let mut acc = [0f32; 2];
9152        macro_rules! sdot2 {
9153            ($w0:expr, $w1:expr, $x:expr) => {{
9154                let x0 = vld1q_s8($x);
9155                let x1 = vld1q_s8($x.add(16));
9156                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
9157                asm!(
9158                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
9159                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
9160                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
9161                    w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
9162                    options(pure, nomem, nostack),
9163                );
9164                vaddvq_s32(vaddq_s32(a0, a1))
9165            }};
9166        }
9167        let gpr2 = gpr & !1;
9168        let mut gi = 0;
9169        while gi < gpr2 {
9170            let off0 = row_off + gi * TILE;
9171            let off1 = off0 + TILE;
9172            let s0 = f16_to_f32(u16::from_le_bytes([
9173                *bytes_ptr.add(off0),
9174                *bytes_ptr.add(off0 + 1),
9175            ]));
9176            let s1 = f16_to_f32(u16::from_le_bytes([
9177                *bytes_ptr.add(off1),
9178                *bytes_ptr.add(off1 + 1),
9179            ]));
9180            let (u0_0, u1_0, u2_0, u3_0) = q1t_unpack_reg_u64s(bytes_ptr.add(off0 + 2));
9181            let (u0_1, u1_1, u2_1, u3_1) = q1t_unpack_reg_u64s(bytes_ptr.add(off1 + 2));
9182            let w0_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_0), vcreate_u64(u1_0)));
9183            let w1_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_0), vcreate_u64(u3_0)));
9184            let w0_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_1), vcreate_u64(u1_1)));
9185            let w1_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_1), vcreate_u64(u3_1)));
9186            for k in 0..2 {
9187                let d0 = sdot2!(w0_0, w1_0, xp[k].add(gi * GROUP_SIZE));
9188                let d1 = sdot2!(w0_1, w1_1, xp[k].add((gi + 1) * GROUP_SIZE));
9189                acc[k] += d0 as f32 * s0 + d1 as f32 * s1;
9190            }
9191            gi += 2;
9192        }
9193        if gi < gpr {
9194            let off = row_off + gi * TILE;
9195            let s = f16_to_f32(u16::from_le_bytes([
9196                *bytes_ptr.add(off),
9197                *bytes_ptr.add(off + 1),
9198            ]));
9199            let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
9200            let w0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0), vcreate_u64(u1)));
9201            let w1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2), vcreate_u64(u3)));
9202            for k in 0..2 {
9203                let d = sdot2!(w0, w1, xp[k].add(gi * GROUP_SIZE));
9204                acc[k] += d as f32 * s;
9205            }
9206        }
9207        acc
9208    }
9209}
9210
9211/// Fused Q1T pair matvec: ONE pass over the rows serves both
9212/// activation streams — on ARM the ternary register unpack happens
9213/// once per tile pair (`q1t_dot_row_sdot2`); elsewhere the second dot
9214/// rides the row's L1-warm tile bytes. Per stream the math matches
9215/// `q1t_matvec` exactly.
9216fn q1t_matvec2(
9217    bytes: &[u8],
9218    x1: &[f32],
9219    x2: &[f32],
9220    rows: usize,
9221    cols: usize,
9222    o1: &mut [f32],
9223    o2: &mut [f32],
9224    pool: Option<&Pool>,
9225) {
9226    debug_assert_eq!(o1.len(), rows);
9227    debug_assert_eq!(o2.len(), rows);
9228    const TILE: usize = cortiq_core::quant::Q1T_TILE;
9229    let gpr = cols / GROUP_SIZE;
9230    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
9231    let out1 = SendMut(o1.as_mut_ptr());
9232    let out2 = SendMut(o2.as_mut_ptr());
9233    if a8w8_enabled() {
9234        let a1 = split_act(x1);
9235        let a2 = split_act(x2);
9236        let (a1, a2) = (&a1, &a2);
9237        let run = move |start: usize, end: usize| {
9238            for r in start..end {
9239                #[cfg(target_arch = "aarch64")]
9240                // a8w8 on aarch64 ⇔ sdot_enabled(), so the kernel's
9241                // target features are present.
9242                let ds = unsafe { q1t_dot_row_sdot2(bytes, r, gpr, &a1.xq, &a2.xq) };
9243                #[cfg(not(target_arch = "aarch64"))]
9244                let ds = [
9245                    q1t_dot_row_i8(bytes, r, gpr, &a1.xq),
9246                    q1t_dot_row_i8(bytes, r, gpr, &a2.xq),
9247                ];
9248                let mut acc1 = ds[0] * a1.sx;
9249                for &(j, xv) in &a1.outliers {
9250                    acc1 += q1t_base_weight(bytes, r, gpr, j) * xv;
9251                }
9252                acc1 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x1);
9253                let mut acc2 = ds[1] * a2.sx;
9254                for &(j, xv) in &a2.outliers {
9255                    acc2 += q1t_base_weight(bytes, r, gpr, j) * xv;
9256                }
9257                acc2 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x2);
9258                // SAFETY: disjoint row ranges per worker.
9259                unsafe {
9260                    *out1.at(r) = acc1;
9261                    *out2.at(r) = acc2;
9262                }
9263            }
9264        };
9265        dispatch_rows(pool, rows, &run);
9266        return;
9267    }
9268    let run = move |start: usize, end: usize| {
9269        // Exact path (CMF_SDOT=0): unpack the sign LUT once per group,
9270        // dot both streams — same op order per stream as `q1t_matvec`.
9271        let mut sg = [0f32; GROUP_SIZE];
9272        for r in start..end {
9273            let mut acc1 = 0f32;
9274            let mut acc2 = 0f32;
9275            for g in 0..gpr {
9276                let off = (r * gpr + g) * TILE;
9277                let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
9278                let codes = &bytes[off + 2..off + TILE];
9279                for bi in 0..6 {
9280                    sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
9281                }
9282                let lut = &SIGN5[codes[6] as usize];
9283                sg[30] = lut[0];
9284                sg[31] = lut[1];
9285                let xg1 = &x1[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
9286                let xg2 = &x2[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
9287                let mut gsum1 = 0f32;
9288                for k in 0..GROUP_SIZE {
9289                    gsum1 += sg[k] * xg1[k];
9290                }
9291                acc1 += s * gsum1;
9292                let mut gsum2 = 0f32;
9293                for k in 0..GROUP_SIZE {
9294                    gsum2 += sg[k] * xg2[k];
9295                }
9296                acc2 += s * gsum2;
9297            }
9298            acc1 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x1);
9299            acc2 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x2);
9300            // SAFETY: disjoint row ranges per worker.
9301            unsafe {
9302                *out1.at(r) = acc1;
9303                *out2.at(r) = acc2;
9304            }
9305        }
9306    };
9307    dispatch_rows(pool, rows, &run);
9308}
9309
9310/// Ternary (q1t) matmat (prefill) — dequant each row once, dot the whole
9311/// batch against it (amortizes the per-row decode).
9312fn q1t_matmat(
9313    bytes: &[u8],
9314    xs: &[f32],
9315    b: usize,
9316    rows: usize,
9317    cols: usize,
9318    out: &mut [f32],
9319    pool: Option<&Pool>,
9320) {
9321    debug_assert_eq!(out.len(), b * rows);
9322    const TILE: usize = cortiq_core::quant::Q1T_TILE;
9323    let gpr = cols / GROUP_SIZE;
9324    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
9325    let out_addr = SendMut(out.as_mut_ptr());
9326    // int8 prefill (ARM SDOT / x86 AVX2): quantize the B inputs once, unpack
9327    // each weight row's signs to i8 ONCE, then int8-dot against every input —
9328    // the row sign-decode amortizes over the whole batch. CMF_SDOT=0 → f32.
9329    if a8w8_enabled() {
9330        let acts: Vec<SplitAct> = (0..b)
9331            .map(|bi| split_act(&xs[bi * cols..(bi + 1) * cols]))
9332            .collect();
9333        let acts = &acts;
9334        let run = move |start: usize, end: usize| {
9335            let mut sg = vec![0i8; cols + 8]; // row signs, i8 (+8 unpack slack)
9336            let mut sc = vec![0f32; gpr]; // per-group scales
9337            let mut accs = vec![0f32; b]; // per-batch accumulators, reused per row
9338            for r in start..end {
9339                for g in 0..gpr {
9340                    let off = (r * gpr + g) * TILE;
9341                    sc[g] = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
9342                    q1t_unpack_group_i8(
9343                        bytes.as_ptr().wrapping_add(off + 2),
9344                        &mut sg[g * GROUP_SIZE..],
9345                    );
9346                }
9347                for bi in 0..b {
9348                    let act = &acts[bi];
9349                    let mut isum = 0f32;
9350                    for g in 0..gpr {
9351                        let d = q1t_i8dot32(
9352                            sg.as_ptr().wrapping_add(g * GROUP_SIZE),
9353                            act.xq.as_ptr().wrapping_add(g * GROUP_SIZE),
9354                        );
9355                        isum += d as f32 * sc[g];
9356                    }
9357                    let mut acc = isum * act.sx;
9358                    for &(j, xv) in &act.outliers {
9359                        acc += q1t_base_weight(bytes, r, gpr, j) * xv;
9360                    }
9361                    accs[bi] = acc;
9362                }
9363                // Overlay ONCE per row for the whole batch: read each (col, val)
9364                // from mmap a single time (was b× — the re-read dominated prefill)
9365                // and fan it out over the batch via the cached inputs.
9366                if has_ov {
9367                    let (c0, c1) = (
9368                        q1t_rowptr(bytes, rp_off, r),
9369                        q1t_rowptr(bytes, rp_off, r + 1),
9370                    );
9371                    for p in c0..c1 {
9372                        let e = ent_off + p * 4;
9373                        let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
9374                        let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
9375                        for bi in 0..b {
9376                            accs[bi] += val * xs[bi * cols + col];
9377                        }
9378                    }
9379                }
9380                for bi in 0..b {
9381                    unsafe { *out_addr.at(bi * rows + r) = accs[bi] };
9382                }
9383            }
9384        };
9385        dispatch_rows(pool, rows, &run);
9386        return;
9387    }
9388    let run = move |start: usize, end: usize| {
9389        let mut buf = vec![0f32; cols];
9390        for r in start..end {
9391            q1t_dequant_row(bytes, r, gpr, rp_off, ent_off, has_ov, &mut buf);
9392            for bi in 0..b {
9393                let xr = &xs[bi * cols..(bi + 1) * cols];
9394                let mut acc = 0f32;
9395                for j in 0..cols {
9396                    acc += buf[j] * xr[j];
9397                }
9398                unsafe { *out_addr.at(bi * rows + r) = acc };
9399            }
9400        }
9401    };
9402    dispatch_rows(pool, rows, &run);
9403}
9404
9405/// q1-tile entry points of the binary drivers (the tests' handles; the
9406/// runtime goes through `QTensor::bin_w`).
9407#[cfg(test)]
9408fn q1_matvec(
9409    bytes: &[u8],
9410    x: &[f32],
9411    rows: usize,
9412    cols: usize,
9413    out: &mut [f32],
9414    pool: Option<&Pool>,
9415) {
9416    bin_matvec(BinW::Q1(bytes), x, rows, cols, out, pool)
9417}
9418
9419/// Binary-weight matvec (q1 or q1g): A8W8 rows with the shared group
9420/// sums, or the exact scalar walk under CMF_SDOT=0.
9421fn bin_matvec(w: BinW, x: &[f32], rows: usize, cols: usize, out: &mut [f32], pool: Option<&Pool>) {
9422    debug_assert_eq!(out.len(), rows);
9423    let gpr = cols / GROUP_SIZE;
9424    let out_addr = SendMut(out.as_mut_ptr());
9425    if a8w8_enabled() {
9426        let act = split_act(x);
9427        let gsum = q1_group_sums(&act.xq, gpr);
9428        let (act, gsum) = (&act, &gsum);
9429        let run =
9430            move |start: usize, end: usize| bin_range_a8w8(w, gpr, act, gsum, out_addr, start, end);
9431        dispatch_rows(pool, rows, &run);
9432        return;
9433    }
9434    let run = move |start: usize, end: usize| bin_range_f32(w, gpr, x, out_addr, start, end);
9435    dispatch_rows(pool, rows, &run);
9436}
9437
9438/// Fused two-input q1 matvec (weights read once per pair).
9439#[cfg(test)]
9440#[allow(clippy::too_many_arguments)]
9441fn q1_matvec2(
9442    bytes: &[u8],
9443    x1: &[f32],
9444    x2: &[f32],
9445    rows: usize,
9446    cols: usize,
9447    o1: &mut [f32],
9448    o2: &mut [f32],
9449    pool: Option<&Pool>,
9450) {
9451    bin_matvec2(BinW::Q1(bytes), x1, x2, rows, cols, o1, o2, pool)
9452}
9453
9454/// Fused two-input binary-weight matvec (q1 or q1g).
9455#[allow(clippy::too_many_arguments)]
9456fn bin_matvec2(
9457    w: BinW,
9458    x1: &[f32],
9459    x2: &[f32],
9460    rows: usize,
9461    cols: usize,
9462    o1: &mut [f32],
9463    o2: &mut [f32],
9464    pool: Option<&Pool>,
9465) {
9466    let gpr = cols / GROUP_SIZE;
9467    let p1 = SendMut(o1.as_mut_ptr());
9468    let p2 = SendMut(o2.as_mut_ptr());
9469    if a8w8_enabled() {
9470        let a1 = split_act(x1);
9471        let a2 = split_act(x2);
9472        let g1 = q1_group_sums(&a1.xq, gpr);
9473        let g2 = q1_group_sums(&a2.xq, gpr);
9474        let (a1, a2, g1, g2) = (&a1, &a2, &g1, &g2);
9475        let run = move |start: usize, end: usize| {
9476            for r in start..end {
9477                let mut v1 = w.dot_i8(r, gpr, &a1.xq, g1) * a1.sx;
9478                let mut v2 = w.dot_i8(r, gpr, &a2.xq, g2) * a2.sx;
9479                for &(j, xv) in &a1.outliers {
9480                    let (wv, s) = w.outlier(r, gpr, j);
9481                    v1 += wv * s * xv;
9482                }
9483                for &(j, xv) in &a2.outliers {
9484                    let (wv, s) = w.outlier(r, gpr, j);
9485                    v2 += wv * s * xv;
9486                }
9487                // SAFETY: disjoint row ranges per worker.
9488                unsafe {
9489                    *p1.at(r) = v1;
9490                    *p2.at(r) = v2;
9491                }
9492            }
9493        };
9494        dispatch_rows(pool, rows, &run);
9495        return;
9496    }
9497    let run = move |start: usize, end: usize| {
9498        for r in start..end {
9499            // SAFETY: disjoint row ranges per worker.
9500            unsafe {
9501                *p1.at(r) = w.row_exact(r, gpr, x1);
9502                *p2.at(r) = w.row_exact(r, gpr, x2);
9503            }
9504        }
9505    };
9506    dispatch_rows(pool, rows, &run);
9507}
9508
9509/// Batched q1 matmat: each row's tiles stream once per microbatch.
9510#[cfg(test)]
9511#[allow(clippy::too_many_arguments)]
9512fn q1_matmat(
9513    bytes: &[u8],
9514    xs_all: &[f32],
9515    b: usize,
9516    rows: usize,
9517    cols: usize,
9518    out: &mut [f32],
9519    pool: Option<&Pool>,
9520) {
9521    bin_matmat(BinW::Q1(bytes), xs_all, b, rows, cols, out, pool)
9522}
9523
9524/// Batched binary-weight matmat (q1 or q1g).
9525#[allow(clippy::too_many_arguments)]
9526fn bin_matmat(
9527    w: BinW,
9528    xs_all: &[f32],
9529    b: usize,
9530    rows: usize,
9531    cols: usize,
9532    out: &mut [f32],
9533    pool: Option<&Pool>,
9534) {
9535    debug_assert_eq!(out.len(), b * rows);
9536    let gpr = cols / GROUP_SIZE;
9537    let out_addr = SendMut(out.as_mut_ptr());
9538    if a8w8_enabled() {
9539        let acts: Vec<(SplitAct, Vec<i32>)> = (0..b)
9540            .map(|bi| {
9541                let act = split_act(&xs_all[bi * cols..(bi + 1) * cols]);
9542                let gsum = q1_group_sums(&act.xq, gpr);
9543                (act, gsum)
9544            })
9545            .collect();
9546        let acts = &acts;
9547        #[cfg(target_arch = "x86_64")]
9548        let blocked_ok = avx2_enabled() && blocked_enabled();
9549        #[cfg(target_arch = "aarch64")]
9550        let blocked_ok = sdot_enabled() && blocked_enabled();
9551        let run = move |start: usize, end: usize| {
9552            for r in start..end {
9553                let mut bi = 0usize;
9554                // Blocked 1×4: the unpacked bit mask serves four
9555                // activation streams per group.
9556                #[cfg(target_arch = "aarch64")]
9557                if blocked_ok {
9558                    while bi + 4 <= acts.len() {
9559                        let xs = [
9560                            acts[bi].0.xq.as_slice(),
9561                            acts[bi + 1].0.xq.as_slice(),
9562                            acts[bi + 2].0.xq.as_slice(),
9563                            acts[bi + 3].0.xq.as_slice(),
9564                        ];
9565                        let gs = [
9566                            acts[bi].1.as_slice(),
9567                            acts[bi + 1].1.as_slice(),
9568                            acts[bi + 2].1.as_slice(),
9569                            acts[bi + 3].1.as_slice(),
9570                        ];
9571                        // SAFETY: blocked_ok checked SDOT.
9572                        let d = unsafe { w.dot_1x4_sdot(r, gpr, xs, gs) };
9573                        for k in 0..4 {
9574                            let (act, _) = &acts[bi + k];
9575                            let mut acc = d[k] * act.sx;
9576                            for &(j, xv) in &act.outliers {
9577                                let (wv, sc) = w.outlier(r, gpr, j);
9578                                acc += wv * sc * xv;
9579                            }
9580                            // SAFETY: disjoint (bi, r) cells per worker.
9581                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
9582                        }
9583                        bi += 4;
9584                    }
9585                }
9586                #[cfg(target_arch = "x86_64")]
9587                if blocked_ok {
9588                    while bi + 4 <= acts.len() {
9589                        let xs = [
9590                            acts[bi].0.xq.as_slice(),
9591                            acts[bi + 1].0.xq.as_slice(),
9592                            acts[bi + 2].0.xq.as_slice(),
9593                            acts[bi + 3].0.xq.as_slice(),
9594                        ];
9595                        let gs = [
9596                            acts[bi].1.as_slice(),
9597                            acts[bi + 1].1.as_slice(),
9598                            acts[bi + 2].1.as_slice(),
9599                            acts[bi + 3].1.as_slice(),
9600                        ];
9601                        // SAFETY: blocked_ok checked AVX2.
9602                        let d = unsafe { w.dot_1x4_x86(r, gpr, xs, gs) };
9603                        for k in 0..4 {
9604                            let (act, _) = &acts[bi + k];
9605                            let mut acc = d[k] * act.sx;
9606                            for &(j, xv) in &act.outliers {
9607                                let (wv, sc) = w.outlier(r, gpr, j);
9608                                acc += wv * sc * xv;
9609                            }
9610                            // SAFETY: disjoint (bi, r) cells per worker.
9611                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
9612                        }
9613                        bi += 4;
9614                    }
9615                }
9616                while bi < acts.len() {
9617                    let (act, gsum) = &acts[bi];
9618                    let mut acc = w.dot_i8(r, gpr, &act.xq, gsum) * act.sx;
9619                    for &(j, xv) in &act.outliers {
9620                        let (wv, s) = w.outlier(r, gpr, j);
9621                        acc += wv * s * xv;
9622                    }
9623                    // SAFETY: disjoint (bi, r) cells per worker range.
9624                    unsafe { *out_addr.at(bi * rows + r) = acc };
9625                    bi += 1;
9626                }
9627            }
9628        };
9629        dispatch_rows(pool, rows, &run);
9630        return;
9631    }
9632    let run = move |start: usize, end: usize| {
9633        for r in start..end {
9634            for bi in 0..b {
9635                let x = &xs_all[bi * cols..(bi + 1) * cols];
9636                // SAFETY: disjoint (bi, r) cells per worker range.
9637                unsafe { *out_addr.at(bi * rows + r) = w.row_exact(r, gpr, x) };
9638            }
9639        }
9640    };
9641    dispatch_rows(pool, rows, &run);
9642}
9643
9644/// Fused q4_block matvec straight from the mapped bytes. SDOT path when
9645/// dotprod is available (port of vmfcore `dot_q4_block_sdot`, measured
9646/// +23% on q4 decode): nibbles → centered i8, int8×int8 `sdot` per
9647/// 32-group, exact outlier correction — the same A8W8 contract as q8.
9648/// `CMF_SDOT=0` keeps the exact scalar path.
9649fn q4matvec(
9650    bytes: &[u8],
9651    x: &[f32],
9652    rows: usize,
9653    cols: usize,
9654    out: &mut [f32],
9655    pool: Option<&Pool>,
9656) {
9657    debug_assert_eq!(out.len(), rows);
9658    let (packed, scales) = q4_split(bytes, rows, cols);
9659    let gpr = cols / GROUP_SIZE;
9660    let out_addr = SendMut(out.as_mut_ptr());
9661
9662    if a8w8_enabled() {
9663        let act = split_act(x);
9664        let run = move |start: usize, end: usize| {
9665            q4_range_a8w8(packed, scales, gpr, cols, &act, out_addr, start, end)
9666        };
9667        dispatch_rows(pool, rows, &run);
9668        return;
9669    }
9670
9671    let run =
9672        move |start: usize, end: usize| q4_range_f32(packed, scales, gpr, x, out_addr, start, end);
9673    dispatch_rows(pool, rows, &run);
9674}
9675
9676/// One q4 row via the A8W8 int8 path — SDOT on ARM, AVX2 maddubs on
9677/// x86 (scalar fallback is unreachable: callers gate on a8w8_enabled).
9678#[inline]
9679#[allow(unreachable_code)]
9680/// One UNPACKED q4 row (centered i8 in `buf`) against four activation
9681/// streams: the 32-byte weight chunk and its abs() load once per group,
9682/// the per-group f16 scale decodes once — four maddubs+reduce chains
9683/// instead of four full (load, abs, dot) rounds.
9684#[cfg(target_arch = "x86_64")]
9685#[target_feature(enable = "avx2")]
9686unsafe fn dot_q4b_row_1x4_avx2(
9687    buf: &[u8],
9688    scales: &[u8],
9689    g0: usize,
9690    gpr: usize,
9691    xs: [&[i8]; 4],
9692) -> [f32; 4] {
9693    // SAFETY: callers uphold buffer contracts (buf.len() == gpr·32).
9694    unsafe {
9695        use core::arch::x86_64::*;
9696        let ones = _mm256_set1_epi16(1);
9697        let mut acc = [0f32; 4];
9698        for gi in 0..gpr {
9699            let s = f16_to_f32(u16::from_le_bytes([
9700                scales[(g0 + gi) * 2],
9701                scales[(g0 + gi) * 2 + 1],
9702            ]));
9703            let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9704            let aw = _mm256_abs_epi8(w);
9705            for (k, xq) in xs.iter().enumerate() {
9706                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9707                let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
9708                let d = _mm256_madd_epi16(p16, ones);
9709                let hi128 = _mm256_extracti128_si256::<1>(d);
9710                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
9711                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
9712                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
9713                acc[k] += _mm_cvtsi128_si32(s32) as f32 * s;
9714            }
9715        }
9716        acc
9717    }
9718}
9719
9720/// VNNI twin of `dot_q4b_row_1x4_avx2` (see `dpbusd_hsum`).
9721#[cfg(target_arch = "x86_64")]
9722#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
9723unsafe fn dot_q4b_row_1x4_vnni(
9724    buf: &[u8],
9725    scales: &[u8],
9726    g0: usize,
9727    gpr: usize,
9728    xs: [&[i8]; 4],
9729) -> [f32; 4] {
9730    // SAFETY: callers uphold buffer contracts (buf.len() == gpr·32).
9731    unsafe {
9732        use core::arch::x86_64::*;
9733        let mut acc = [0f32; 4];
9734        for gi in 0..gpr {
9735            let s = f16_to_f32(u16::from_le_bytes([
9736                scales[(g0 + gi) * 2],
9737                scales[(g0 + gi) * 2 + 1],
9738            ]));
9739            let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9740            let aw = _mm256_abs_epi8(w);
9741            for (k, xq) in xs.iter().enumerate() {
9742                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9743                let d = dpbusd_hsum(aw, _mm256_sign_epi8(x, w));
9744                acc[k] += d as f32 * s;
9745            }
9746        }
9747        acc
9748    }
9749}
9750
9751/// The vbit flavor of the blocked 1×4: the per-activation A8W8 scale
9752/// folds in PER GROUP as `(d·sx)·s` — bit-matching the single-matvec
9753/// accumulation order (the q4_block flavor applies sx once at the end,
9754/// matching ITS single path; the two conventions are historical and
9755/// each blocked leg must mirror its own).
9756#[cfg(target_arch = "x86_64")]
9757#[target_feature(enable = "avx2")]
9758unsafe fn dot_q4b_row_1x4_sx_avx2(
9759    buf: &[u8],
9760    scales: &[u8],
9761    g0: usize,
9762    gpr: usize,
9763    xs: [&[i8]; 4],
9764    sxs: [f32; 4],
9765) -> [f32; 4] {
9766    // SAFETY: callers uphold buffer contracts (buf.len() == gpr·32).
9767    unsafe {
9768        use core::arch::x86_64::*;
9769        let ones = _mm256_set1_epi16(1);
9770        let mut acc = [0f32; 4];
9771        for gi in 0..gpr {
9772            let s = f16_to_f32(u16::from_le_bytes([
9773                scales[(g0 + gi) * 2],
9774                scales[(g0 + gi) * 2 + 1],
9775            ]));
9776            let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9777            let aw = _mm256_abs_epi8(w);
9778            for (k, xq) in xs.iter().enumerate() {
9779                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9780                let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
9781                let d = _mm256_madd_epi16(p16, ones);
9782                let hi128 = _mm256_extracti128_si256::<1>(d);
9783                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
9784                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
9785                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
9786                acc[k] += (_mm_cvtsi128_si32(s32) as f32 * sxs[k]) * s;
9787            }
9788        }
9789        acc
9790    }
9791}
9792
9793/// VNNI twin of `dot_q4b_row_1x4_sx_avx2` (see `dpbusd_hsum`; the
9794/// per-group `(d·sx)·s` fold mirrors the vbit single path).
9795#[cfg(target_arch = "x86_64")]
9796#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
9797unsafe fn dot_q4b_row_1x4_sx_vnni(
9798    buf: &[u8],
9799    scales: &[u8],
9800    g0: usize,
9801    gpr: usize,
9802    xs: [&[i8]; 4],
9803    sxs: [f32; 4],
9804) -> [f32; 4] {
9805    // SAFETY: callers uphold buffer contracts (buf.len() == gpr·32).
9806    unsafe {
9807        use core::arch::x86_64::*;
9808        let mut acc = [0f32; 4];
9809        for gi in 0..gpr {
9810            let s = f16_to_f32(u16::from_le_bytes([
9811                scales[(g0 + gi) * 2],
9812                scales[(g0 + gi) * 2 + 1],
9813            ]));
9814            let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9815            let aw = _mm256_abs_epi8(w);
9816            for (k, xq) in xs.iter().enumerate() {
9817                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9818                let d = dpbusd_hsum(aw, _mm256_sign_epi8(x, w));
9819                acc[k] += (d as f32 * sxs[k]) * s;
9820            }
9821        }
9822        acc
9823    }
9824}
9825
9826#[allow(unreachable_code)]
9827fn dot_q4_row_i8(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
9828    #[cfg(target_arch = "aarch64")]
9829    unsafe {
9830        return dot_q4_row_sdot(packed, scales, g0, gpr, xq);
9831    }
9832    #[cfg(target_arch = "x86_64")]
9833    unsafe {
9834        return dot_q4_row_avx2(packed, scales, g0, gpr, xq);
9835    }
9836    let mut acc = 0f32;
9837    for gi in 0..gpr {
9838        let g = g0 + gi;
9839        let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
9840        let mut d = 0i32;
9841        for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
9842            d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
9843                + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
9844        }
9845        acc += d as f32 * s;
9846    }
9847    acc
9848}
9849
9850/// Two-activation q4 row via the A8W8 int8 path (see `dot_q4_row_i8`).
9851#[inline]
9852#[allow(unreachable_code)]
9853fn dot_q4_row_i8_2(
9854    packed: &[u8],
9855    scales: &[u8],
9856    g0: usize,
9857    gpr: usize,
9858    xq1: &[i8],
9859    xq2: &[i8],
9860) -> (f32, f32) {
9861    #[cfg(target_arch = "aarch64")]
9862    unsafe {
9863        return dot_q4_row_sdot2(packed, scales, g0, gpr, xq1, xq2);
9864    }
9865    #[cfg(target_arch = "x86_64")]
9866    unsafe {
9867        return dot_q4_row_avx2_2(packed, scales, g0, gpr, xq1, xq2);
9868    }
9869    (
9870        dot_q4_row_i8(packed, scales, g0, gpr, xq1),
9871        dot_q4_row_i8(packed, scales, g0, gpr, xq2),
9872    )
9873}
9874
9875/// One q4 row range via SDOT (kernel body of `q4matvec`, extracted so
9876/// multi-matrix jobs can drive it for several tensors in one dispatch).
9877#[allow(clippy::too_many_arguments)]
9878fn q4_range_a8w8(
9879    packed: &[u8],
9880    scales: &[u8],
9881    gpr: usize,
9882    cols: usize,
9883    act: &SplitAct,
9884    out: SendMut,
9885    start: usize,
9886    end: usize,
9887) {
9888    for r in start..end {
9889        let mut acc = dot_q4_row_i8(packed, scales, r * gpr, gpr, &act.xq) * act.sx;
9890        // xq is zeroed at outlier slots — add the exact terms.
9891        for &(j, xv) in &act.outliers {
9892            let flat = r * cols + j;
9893            let byte = packed[flat / 2];
9894            let nib = if flat & 1 == 0 {
9895                byte & 0x0F
9896            } else {
9897                byte >> 4
9898            };
9899            let s = f16_to_f32(u16::from_le_bytes([
9900                scales[(flat / GROUP_SIZE) * 2],
9901                scales[(flat / GROUP_SIZE) * 2 + 1],
9902            ]));
9903            acc += ((nib as i32 - 8) as f32) * s * xv;
9904        }
9905        // SAFETY: disjoint row ranges per worker.
9906        unsafe { *out.at(r) = acc };
9907    }
9908}
9909
9910/// Two-input q4 row range via the A8W8 int8 path — kernel body of
9911/// `q4matvec2`, extracted for pair multi-matrix jobs.
9912#[allow(clippy::too_many_arguments)]
9913fn q4_range2_a8w8(
9914    packed: &[u8],
9915    scales: &[u8],
9916    gpr: usize,
9917    cols: usize,
9918    a1: &SplitAct,
9919    a2: &SplitAct,
9920    p1: SendMut,
9921    p2: SendMut,
9922    start: usize,
9923    end: usize,
9924) {
9925    for r in start..end {
9926        let (s1, s2) = dot_q4_row_i8_2(packed, scales, r * gpr, gpr, &a1.xq, &a2.xq);
9927        let mut acc1 = s1 * a1.sx;
9928        let mut acc2 = s2 * a2.sx;
9929        // xq is zeroed at outlier slots — add the exact terms.
9930        let fix = |outliers: &[(usize, f32)], acc: &mut f32| {
9931            for &(j, xv) in outliers {
9932                let flat = r * cols + j;
9933                let byte = packed[flat / 2];
9934                let nib = if flat & 1 == 0 {
9935                    byte & 0x0F
9936                } else {
9937                    byte >> 4
9938                };
9939                let s = f16_to_f32(u16::from_le_bytes([
9940                    scales[(flat / GROUP_SIZE) * 2],
9941                    scales[(flat / GROUP_SIZE) * 2 + 1],
9942                ]));
9943                *acc += ((nib as i32 - 8) as f32) * s * xv;
9944            }
9945        };
9946        fix(&a1.outliers, &mut acc1);
9947        fix(&a2.outliers, &mut acc2);
9948        // SAFETY: disjoint row ranges per worker.
9949        unsafe {
9950            *p1.at(r) = acc1;
9951            *p2.at(r) = acc2;
9952        }
9953    }
9954}
9955
9956/// Exact scalar q4 row range (same extraction, non-SDOT path).
9957fn q4_range_f32(
9958    packed: &[u8],
9959    scales: &[u8],
9960    gpr: usize,
9961    x: &[f32],
9962    out: SendMut,
9963    start: usize,
9964    end: usize,
9965) {
9966    for r in start..end {
9967        let mut acc = 0f32;
9968        for gi in 0..gpr {
9969            let g = r * gpr + gi;
9970            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
9971            let pk = &packed[g * 16..(g + 1) * 16];
9972            let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
9973            let mut ga = 0f32;
9974            for (k, &b) in pk.iter().enumerate() {
9975                ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
9976                    + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
9977            }
9978            acc += ga * s;
9979        }
9980        // SAFETY: disjoint row ranges per worker.
9981        unsafe { *out.at(r) = acc };
9982    }
9983}
9984
9985/// Fused two-input q4 matvec: nibbles are unpacked ONCE per group and
9986/// dotted against both activations (was: two full matvecs — double
9987/// weight traffic). Per-lane math matches `q4matvec` exactly.
9988#[allow(clippy::too_many_arguments)]
9989fn q4matvec2(
9990    bytes: &[u8],
9991    x1: &[f32],
9992    x2: &[f32],
9993    rows: usize,
9994    cols: usize,
9995    o1: &mut [f32],
9996    o2: &mut [f32],
9997    pool: Option<&Pool>,
9998) {
9999    debug_assert_eq!(o1.len(), rows);
10000    debug_assert_eq!(o2.len(), rows);
10001    let (packed, scales) = q4_split(bytes, rows, cols);
10002    let gpr = cols / GROUP_SIZE;
10003
10004    if a8w8_enabled() {
10005        let a1 = split_act(x1);
10006        let a2 = split_act(x2);
10007        let p1 = SendMut(o1.as_mut_ptr());
10008        let p2 = SendMut(o2.as_mut_ptr());
10009        let run = move |start: usize, end: usize| {
10010            q4_range2_a8w8(packed, scales, gpr, cols, &a1, &a2, p1, p2, start, end)
10011        };
10012        dispatch_rows(pool, rows, &run);
10013        return;
10014    }
10015
10016    let p1 = SendMut(o1.as_mut_ptr());
10017    let p2 = SendMut(o2.as_mut_ptr());
10018    let run = move |start: usize, end: usize| {
10019        q4_range2_f32(packed, scales, gpr, x1, x2, p1, p2, start, end)
10020    };
10021    dispatch_rows(pool, rows, &run);
10022}
10023
10024/// Two-input exact scalar q4 row range (same extraction).
10025#[allow(clippy::too_many_arguments)]
10026fn q4_range2_f32(
10027    packed: &[u8],
10028    scales: &[u8],
10029    gpr: usize,
10030    x1: &[f32],
10031    x2: &[f32],
10032    p1: SendMut,
10033    p2: SendMut,
10034    start: usize,
10035    end: usize,
10036) {
10037    for r in start..end {
10038        let (mut acc1, mut acc2) = (0f32, 0f32);
10039        for gi in 0..gpr {
10040            let g = r * gpr + gi;
10041            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
10042            let pk = &packed[g * 16..(g + 1) * 16];
10043            let x1g = &x1[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
10044            let x2g = &x2[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
10045            let (mut g1, mut g2) = (0f32, 0f32);
10046            for (k, &b) in pk.iter().enumerate() {
10047                let wl = (b & 0x0F) as f32 - 8.0;
10048                let wh = ((b >> 4) & 0x0F) as f32 - 8.0;
10049                g1 += wl * x1g[k * 2] + wh * x1g[k * 2 + 1];
10050                g2 += wl * x2g[k * 2] + wh * x2g[k * 2 + 1];
10051            }
10052            acc1 += g1 * s;
10053            acc2 += g2 * s;
10054        }
10055        // SAFETY: disjoint row ranges per worker.
10056        unsafe {
10057            *p1.at(r) = acc1;
10058            *p2.at(r) = acc2;
10059        }
10060    }
10061}
10062
10063thread_local! {
10064    /// Per-worker decoded-row scratch for the batched q4/vbit kernels
10065    /// (centered i8 for SDOT, f32 for the exact/scalar paths).
10066    static ROW_I8: std::cell::RefCell<Vec<u8>> = const { std::cell::RefCell::new(Vec::new()) };
10067    static ROW_F32: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
10068}
10069
10070/// Batched q4 matmat: each weight row is unpacked from the mmap ONCE
10071/// and dotted against ALL b activations (prefill used to fall back to b
10072/// full matvecs — b× weight traffic and b× nibble decode). Per-position
10073/// math matches `q4matvec` exactly: same group order, same accumulation.
10074/// `out` is row-major [b, rows] like `qmatmat`.
10075#[allow(clippy::too_many_arguments)]
10076fn q4matmat(
10077    bytes: &[u8],
10078    xs_all: &[f32],
10079    b: usize,
10080    rows: usize,
10081    cols: usize,
10082    out: &mut [f32],
10083    pool: Option<&Pool>,
10084) {
10085    debug_assert_eq!(xs_all.len(), b * cols);
10086    debug_assert_eq!(out.len(), b * rows);
10087    let (packed, scales) = q4_split(bytes, rows, cols);
10088    let gpr = cols / GROUP_SIZE;
10089    let gscale = |g: usize| f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
10090
10091    if a8w8_enabled() {
10092        let acts: Vec<SplitAct> = (0..b)
10093            .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
10094            .collect();
10095        let acts = &acts;
10096        let out_addr = SendMut(out.as_mut_ptr());
10097        let run = move |start: usize, end: usize| {
10098            ROW_I8.with(|rb| {
10099                let mut buf = rb.borrow_mut();
10100                buf.resize(cols, 0);
10101                for r in start..end {
10102                    // Unpack the row's nibbles to centered i8 once
10103                    // (element 2k = low nibble, 2k+1 = high — flat order,
10104                    // same as dot_q4_row_sdot's zip).
10105                    for gi in 0..gpr {
10106                        let g = r * gpr + gi;
10107                        for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
10108                            buf[gi * GROUP_SIZE + k * 2] = ((bt & 0x0F) as i32 - 8) as i8 as u8;
10109                            buf[gi * GROUP_SIZE + k * 2 + 1] =
10110                                (((bt >> 4) & 0x0F) as i32 - 8) as i8 as u8;
10111                        }
10112                    }
10113                    let mut bi = 0usize;
10114                    #[cfg(target_arch = "x86_64")]
10115                    if avx2_enabled() && blocked_enabled() {
10116                        while bi + 4 <= acts.len() {
10117                            let xs = [
10118                                acts[bi].xq.as_slice(),
10119                                acts[bi + 1].xq.as_slice(),
10120                                acts[bi + 2].xq.as_slice(),
10121                                acts[bi + 3].xq.as_slice(),
10122                            ];
10123                            let d = unsafe {
10124                                if vnni_tiles_enabled() {
10125                                    dot_q4b_row_1x4_vnni(&buf, scales, r * gpr, gpr, xs)
10126                                } else {
10127                                    dot_q4b_row_1x4_avx2(&buf, scales, r * gpr, gpr, xs)
10128                                }
10129                            };
10130                            for k in 0..4 {
10131                                let act = &acts[bi + k];
10132                                let mut acc = d[k] * act.sx;
10133                                for &(j, xv) in &act.outliers {
10134                                    acc += (buf[j] as i8) as f32
10135                                        * gscale((r * cols + j) / GROUP_SIZE)
10136                                        * xv;
10137                                }
10138                                // SAFETY: disjoint (bi, r) cells per worker.
10139                                unsafe { *out_addr.at((bi + k) * rows + r) = acc };
10140                            }
10141                            bi += 4;
10142                        }
10143                    }
10144                    while bi < acts.len() {
10145                        let act = &acts[bi];
10146                        let mut acc = 0f32;
10147                        for gi in 0..gpr {
10148                            let d = dot_i8_i8(
10149                                &buf[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
10150                                &act.xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
10151                            );
10152                            acc += d as f32 * gscale(r * gpr + gi);
10153                        }
10154                        acc *= act.sx;
10155                        // xq is zeroed at outlier slots — exact terms.
10156                        for &(j, xv) in &act.outliers {
10157                            acc += (buf[j] as i8) as f32 * gscale((r * cols + j) / GROUP_SIZE) * xv;
10158                        }
10159                        // SAFETY: disjoint (bi, r) cells per worker row range.
10160                        unsafe { *out_addr.at(bi * rows + r) = acc };
10161                        bi += 1;
10162                    }
10163                }
10164            })
10165        };
10166        dispatch_rows(pool, rows, &run);
10167        return;
10168    }
10169
10170    let out_addr = SendMut(out.as_mut_ptr());
10171    let run = move |start: usize, end: usize| {
10172        ROW_F32.with(|rb| {
10173            let mut buf = rb.borrow_mut();
10174            buf.resize(cols, 0.0);
10175            for r in start..end {
10176                // Decode raw (nib − 8) values once; scales stay per-group
10177                // so the accumulation order matches q4matvec bit-for-bit.
10178                for gi in 0..gpr {
10179                    let g = r * gpr + gi;
10180                    for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
10181                        buf[gi * GROUP_SIZE + k * 2] = (bt & 0x0F) as f32 - 8.0;
10182                        buf[gi * GROUP_SIZE + k * 2 + 1] = ((bt >> 4) & 0x0F) as f32 - 8.0;
10183                    }
10184                }
10185                for bi in 0..b {
10186                    let x = &xs_all[bi * cols..(bi + 1) * cols];
10187                    let mut acc = 0f32;
10188                    for gi in 0..gpr {
10189                        let mut ga = 0f32;
10190                        // Pairwise (lo + hi) addition, matching
10191                        // q4matvec's `ga += lo·x + hi·x` shape exactly —
10192                        // a flat one-per-element loop rounds differently
10193                        // and broke bit-parity on the scalar (x86) path.
10194                        for k in 0..GROUP_SIZE / 2 {
10195                            let e = gi * GROUP_SIZE + k * 2;
10196                            ga += buf[e] * x[e] + buf[e + 1] * x[e + 1];
10197                        }
10198                        acc += ga * gscale(r * gpr + gi);
10199                    }
10200                    // SAFETY: disjoint (bi, r) cells per worker row range.
10201                    unsafe { *out_addr.at(bi * rows + r) = acc };
10202                }
10203            }
10204        })
10205    };
10206    dispatch_rows(pool, rows, &run);
10207}
10208
10209/// Batched vbit matmat: each variable-bit row is decoded from the mmap
10210/// ONCE for the whole microbatch. Same per-position math as
10211/// `vbitmatvec` (SDOT A8W8 with exact outliers / exact f32 for b=8 rows
10212/// and the scalar path).
10213#[allow(clippy::too_many_arguments)]
10214fn vbitmatmat(
10215    bytes: &[u8],
10216    offsets: &[usize],
10217    xs_all: &[f32],
10218    b: usize,
10219    rows: usize,
10220    cols: usize,
10221    out: &mut [f32],
10222    pool: Option<&Pool>,
10223) {
10224    debug_assert_eq!(xs_all.len(), b * cols);
10225    debug_assert_eq!(out.len(), b * rows);
10226    debug_assert_eq!(offsets.len(), rows + 1);
10227    let ng = cols / GROUP_SIZE;
10228    let bits = &bytes[..rows];
10229    let sc_off = rows;
10230    let gscale = |r: usize, g: usize| {
10231        let so = (r * ng + g) * 2;
10232        f16_to_f32(u16::from_le_bytes([
10233            bytes[sc_off + so],
10234            bytes[sc_off + so + 1],
10235        ]))
10236    };
10237
10238    // Decode row r's raw (u − L) values into `dst` (f32, unscaled).
10239    let decode_f32 = |r: usize, dst: &mut [f32]| {
10240        let bw = bits[r] as usize;
10241        let l = ((1i32 << (bw - 1)) - 1) as f32;
10242        let data = &bytes[offsets[r]..offsets[r + 1]];
10243        let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
10244        for d in dst.iter_mut() {
10245            while nbits < bw {
10246                acc = (acc << 8) | data[idx] as u64;
10247                idx += 1;
10248                nbits += 8;
10249            }
10250            let u = ((acc >> (nbits - bw)) & ((1u64 << bw) - 1)) as f32;
10251            nbits -= bw;
10252            *d = u - l;
10253        }
10254    };
10255
10256    if a8w8_enabled() {
10257        let acts: Vec<SplitAct> = (0..b)
10258            .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
10259            .collect();
10260        let acts = &acts;
10261        let out_addr = SendMut(out.as_mut_ptr());
10262        let run = move |start: usize, end: usize| {
10263            for r in start..end {
10264                let bw = bits[r] as usize;
10265                if bw == 8 {
10266                    // u−L reaches 128 → no i8 path; decode once, exact
10267                    // f32 dots for every position (same as vbitmatvec).
10268                    ROW_F32.with(|rb| {
10269                        let mut buf = rb.borrow_mut();
10270                        buf.resize(cols, 0.0);
10271                        decode_f32(r, &mut buf);
10272                        for bi in 0..b {
10273                            let x = &xs_all[bi * cols..(bi + 1) * cols];
10274                            let mut dot = 0f32;
10275                            for g in 0..ng {
10276                                let mut gd = 0f32;
10277                                for k in 0..GROUP_SIZE {
10278                                    gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
10279                                }
10280                                dot += gd * gscale(r, g);
10281                            }
10282                            // SAFETY: disjoint (bi, r) cells per worker range.
10283                            unsafe { *out_addr.at(bi * rows + r) = dot };
10284                        }
10285                    });
10286                    continue;
10287                }
10288                let l = (1i32 << (bw - 1)) - 1;
10289                let data = &bytes[offsets[r]..offsets[r + 1]];
10290                ROW_I8.with(|rb| {
10291                    let mut buf = rb.borrow_mut();
10292                    buf.resize(cols, 0);
10293                    #[inline(always)]
10294                    fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
10295                        for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
10296                            let u = unpack8::<B>(&data[blk * B..]);
10297                            for k in 0..8 {
10298                                chunk[k] = (u[k] - l) as i8 as u8;
10299                            }
10300                        }
10301                    }
10302                    match bw {
10303                        3 => fill::<3>(data, l, &mut buf),
10304                        4 => vbit_fill4(data, &mut buf),
10305                        5 => fill::<5>(data, l, &mut buf),
10306                        6 => fill::<6>(data, l, &mut buf),
10307                        _ => unreachable!("vbit bit-width {bw} (validated at load)"),
10308                    }
10309                    let mut bi = 0usize;
10310                    // The vbit scale table shares q4_block's layout
10311                    // (contiguous f16 per (row·ng + g)), so the same
10312                    // blocked 1×4 kernel serves the decoded row.
10313                    #[cfg(target_arch = "x86_64")]
10314                    if avx2_enabled() && blocked_enabled() {
10315                        while bi + 4 <= acts.len() {
10316                            let xs = [
10317                                acts[bi].xq.as_slice(),
10318                                acts[bi + 1].xq.as_slice(),
10319                                acts[bi + 2].xq.as_slice(),
10320                                acts[bi + 3].xq.as_slice(),
10321                            ];
10322                            let sxs = [
10323                                acts[bi].sx,
10324                                acts[bi + 1].sx,
10325                                acts[bi + 2].sx,
10326                                acts[bi + 3].sx,
10327                            ];
10328                            let d = unsafe {
10329                                if vnni_tiles_enabled() {
10330                                    dot_q4b_row_1x4_sx_vnni(
10331                                        &buf,
10332                                        &bytes[sc_off..],
10333                                        r * ng,
10334                                        ng,
10335                                        xs,
10336                                        sxs,
10337                                    )
10338                                } else {
10339                                    dot_q4b_row_1x4_sx_avx2(
10340                                        &buf,
10341                                        &bytes[sc_off..],
10342                                        r * ng,
10343                                        ng,
10344                                        xs,
10345                                        sxs,
10346                                    )
10347                                }
10348                            };
10349                            for k in 0..4 {
10350                                let act = &acts[bi + k];
10351                                let mut dot = d[k];
10352                                for &(j, xv) in &act.outliers {
10353                                    dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
10354                                }
10355                                // SAFETY: disjoint (bi, r) cells per worker.
10356                                unsafe { *out_addr.at((bi + k) * rows + r) = dot };
10357                            }
10358                            bi += 4;
10359                        }
10360                    }
10361                    while bi < acts.len() {
10362                        let act = &acts[bi];
10363                        let mut dot = 0f32;
10364                        for g in 0..ng {
10365                            let d = dot_i8_i8(
10366                                &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
10367                                &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
10368                            ) as f32
10369                                * act.sx;
10370                            dot += d * gscale(r, g);
10371                        }
10372                        for &(j, xv) in &act.outliers {
10373                            dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
10374                        }
10375                        // SAFETY: disjoint (bi, r) cells per worker range.
10376                        unsafe { *out_addr.at(bi * rows + r) = dot };
10377                        bi += 1;
10378                    }
10379                });
10380            }
10381        };
10382        dispatch_rows(pool, rows, &run);
10383        return;
10384    }
10385
10386    let out_addr = SendMut(out.as_mut_ptr());
10387    let run = move |start: usize, end: usize| {
10388        ROW_F32.with(|rb| {
10389            let mut buf = rb.borrow_mut();
10390            buf.resize(cols, 0.0);
10391            for r in start..end {
10392                decode_f32(r, &mut buf);
10393                for bi in 0..b {
10394                    let x = &xs_all[bi * cols..(bi + 1) * cols];
10395                    let mut dot = 0f32;
10396                    for g in 0..ng {
10397                        let mut gd = 0f32;
10398                        for k in 0..GROUP_SIZE {
10399                            gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
10400                        }
10401                        dot += gd * gscale(r, g);
10402                    }
10403                    // SAFETY: disjoint (bi, r) cells per worker range.
10404                    unsafe { *out_addr.at(bi * rows + r) = dot };
10405                }
10406            }
10407        })
10408    };
10409    dispatch_rows(pool, rows, &run);
10410}
10411
10412/// Build a GPU batch job for a q8-family mapped tensor (primary
10413/// shard): prescaled input + directory coordinates. None → not
10414/// GPU-eligible, caller stays on the CPU.
10415pub(crate) fn gpu_batch_job<'a>(
10416    t: &'a QTensor,
10417    x: &[f32],
10418) -> Option<(std::sync::Arc<CmfModel>, crate::gpu::BatchJob<'a>)> {
10419    match t {
10420        QTensor::Mapped {
10421            model,
10422            idx,
10423            dtype: dt @ (TensorDtype::Q8Row | TensorDtype::Q8_2f),
10424            rows,
10425            cols,
10426            row_scale,
10427            col_field,
10428            ..
10429        } => Some((
10430            model.clone(),
10431            crate::gpu::BatchJob {
10432                idx: *idx,
10433                rows: *rows,
10434                cols: *cols,
10435                row_scale,
10436                xs: prescale(x, col_field, *dt).into_owned(),
10437                layout: crate::gpu::BatchLayout::Q8,
10438            },
10439        )),
10440        // q1: raw f32 activations, tile-embedded scales.
10441        QTensor::Mapped {
10442            model,
10443            idx,
10444            dtype: TensorDtype::Q1,
10445            rows,
10446            cols,
10447            ..
10448        } => Some((
10449            model.clone(),
10450            crate::gpu::BatchJob {
10451                idx: *idx,
10452                rows: *rows,
10453                cols: *cols,
10454                row_scale: &[],
10455                xs: x.to_vec(),
10456                layout: crate::gpu::BatchLayout::Q1,
10457            },
10458        )),
10459        // q1g: same raw-f32 contract, sign + scale planes.
10460        QTensor::Mapped {
10461            model,
10462            idx,
10463            dtype: TensorDtype::Q1G,
10464            rows,
10465            cols,
10466            ..
10467        } => Some((
10468            model.clone(),
10469            crate::gpu::BatchJob {
10470                idx: *idx,
10471                rows: *rows,
10472                cols: *cols,
10473                row_scale: &[],
10474                xs: x.to_vec(),
10475                layout: crate::gpu::BatchLayout::Q1g,
10476            },
10477        )),
10478        // q4_tiled / q4tp: raw f32 activations; the scales live in the
10479        // payload (inline tiles / row ladder), so row_scale stays empty.
10480        // The GDN projection batch already runs these layouts on Metal —
10481        // this arm lets the attention QKV batch reach the same kernels.
10482        QTensor::Mapped {
10483            model,
10484            idx,
10485            dtype: dt @ (TensorDtype::Q4Tiled | TensorDtype::Q4TiledP),
10486            rows,
10487            cols,
10488            ..
10489        } => Some((
10490            model.clone(),
10491            crate::gpu::BatchJob {
10492                idx: *idx,
10493                rows: *rows,
10494                cols: *cols,
10495                row_scale: &[],
10496                xs: x.to_vec(),
10497                layout: if *dt == TensorDtype::Q4Tiled {
10498                    crate::gpu::BatchLayout::Q4t
10499                } else {
10500                    crate::gpu::BatchLayout::Q4tp
10501                },
10502            },
10503        )),
10504        _ => None,
10505    }
10506}
10507
10508thread_local! {
10509    static PRESCALE_BUF1: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
10510    static PRESCALE_BUF2: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
10511}
10512
10513pub(crate) fn prescale<'a>(
10514    x: &'a [f32],
10515    col_field: &[f32],
10516    dtype: TensorDtype,
10517) -> std::borrow::Cow<'a, [f32]> {
10518    if dtype == TensorDtype::Q8_2f {
10519        x.iter().zip(col_field).map(|(a, c)| a * c).collect()
10520    } else {
10521        std::borrow::Cow::Borrowed(x)
10522    }
10523}
10524
10525/// θ col-field fold for q8_2f activations. Borrowed pass-through for
10526/// every other dtype, using thread-local buffers to eliminate per-matvec allocations.
10527pub(crate) fn prescale_with<R, F: FnOnce(&[f32]) -> R>(
10528    x: &[f32],
10529    col_field: &[f32],
10530    dtype: TensorDtype,
10531    buf_id: u8,
10532    f: F,
10533) -> R {
10534    if dtype == TensorDtype::Q8_2f {
10535        if buf_id == 1 {
10536            PRESCALE_BUF1.with(|b| {
10537                let mut buf = b.borrow_mut();
10538                buf.clear();
10539                buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
10540                f(&buf)
10541            })
10542        } else {
10543            PRESCALE_BUF2.with(|b| {
10544                let mut buf = b.borrow_mut();
10545                buf.clear();
10546                buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
10547                f(&buf)
10548            })
10549        }
10550    } else {
10551        f(x)
10552    }
10553}
10554
10555// ───────────────────── x86-64 AVX2 kernels (roadmap этап 2) ─────────────────────
10556
10557/// AVX2+FMA available? Default ON when the CPU supports both;
10558/// `CMF_AVX2=0` disables (falls back to the autovectorized loops).
10559#[cfg(target_arch = "x86_64")]
10560pub(crate) fn avx2_enabled() -> bool {
10561    use std::sync::OnceLock;
10562    static ON: OnceLock<bool> = OnceLock::new();
10563    *ON.get_or_init(|| {
10564        std::env::var("CMF_AVX2").map(|v| v != "0").unwrap_or(true)
10565            && std::arch::is_x86_feature_detected!("avx2")
10566            && std::arch::is_x86_feature_detected!("fma")
10567    })
10568}
10569
10570/// AVX2 A8W8 allowed? The quantized-activation contract is switched by
10571/// the SAME env as the ARM SDOT path: `CMF_SDOT=0` keeps exact kernels
10572/// (the golden-parity exact gate relies on it) — AVX2 f32 kernels stay
10573/// active either way, they are exact (regrouped sums only).
10574#[cfg(target_arch = "x86_64")]
10575fn avx2_a8w8_enabled() -> bool {
10576    if FLOAT_ACTIVATIONS.get() {
10577        return false;
10578    }
10579    use std::sync::OnceLock;
10580    static ON: OnceLock<bool> = OnceLock::new();
10581    *ON.get_or_init(|| {
10582        avx2_enabled() && std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true)
10583    })
10584}
10585
10586thread_local! {
10587    static FULL_GPU_Q8: std::cell::Cell<bool> = const { std::cell::Cell::new(false) };
10588}
10589
10590/// Match the graph's full-device q8 projection precision on MiMo's host
10591/// tail. Does not enable the GPU or bypass a CPU-only/device-refusal gate.
10592pub(crate) fn enter_full_gpu_q8_scope() -> impl Drop {
10593    struct Restore(bool, std::marker::PhantomData<std::rc::Rc<()>>);
10594    impl Drop for Restore {
10595        fn drop(&mut self) {
10596            FULL_GPU_Q8.set(self.0);
10597        }
10598    }
10599    Restore(FULL_GPU_Q8.replace(true), std::marker::PhantomData)
10600}
10601
10602// Dynamic MiMo experts must not change activation precision when a cache
10603// fill moves them from CPU to GPU. Thread-local: only the cold-expert
10604// dispatch selects float kernels; concurrent pipelines keep their policy.
10605thread_local! {
10606    static FLOAT_ACTIVATIONS: std::cell::Cell<bool> = const { std::cell::Cell::new(false) };
10607}
10608
10609pub(crate) fn float_activations_scope<R>(f: impl FnOnce() -> R) -> R {
10610    struct Restore(bool);
10611    impl Drop for Restore {
10612        fn drop(&mut self) {
10613            FLOAT_ACTIVATIONS.set(self.0);
10614        }
10615    }
10616    let _restore = Restore(FLOAT_ACTIVATIONS.replace(true));
10617    f()
10618}
10619
10620/// Row-exact batching: while set, the x86 batched kernels (`qmatmat`,
10621/// `q4tp_matmat`) compute every (weight row, token) cell with the
10622/// single-token kernel instead of the blocked 2×4 / 1×4 / 1×8 tiles, so a
10623/// token's result does not depend on the batch it rides in and equals its
10624/// matvec. On ARM `q4tp_matmat` keeps its 1×4 tile but in the matvec's
10625/// reduction order, and no q4tp batch takes the AMX or device GEMM. The
10626/// MiMo speculative verify holds it (`row_exact_scope`) — its accepted
10627/// rows must be the rows plain decode would have produced.
10628// Shared with pool workers, so overlapping requests must keep the mode
10629// enabled until the LAST scope leaves. Saving/restoring a global bool is
10630// incorrect when two threads enter and leave in a non-LIFO order.
10631static ROW_EXACT: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
10632
10633pub(crate) fn row_exact() -> bool {
10634    ROW_EXACT.load(std::sync::atomic::Ordering::Acquire) != 0
10635}
10636
10637fn counted_row_exact_scope<R>(active: &std::sync::atomic::AtomicUsize, f: impl FnOnce() -> R) -> R {
10638    struct Restore<'a>(&'a std::sync::atomic::AtomicUsize);
10639    impl Drop for Restore<'_> {
10640        fn drop(&mut self) {
10641            self.0.fetch_sub(1, std::sync::atomic::Ordering::AcqRel);
10642        }
10643    }
10644    active.fetch_add(1, std::sync::atomic::Ordering::AcqRel);
10645    let _restore = Restore(active);
10646    f()
10647}
10648
10649/// Run `f` with row-exact batching on (also released on unwind).
10650pub(crate) fn row_exact_scope<R>(f: impl FnOnce() -> R) -> R {
10651    counted_row_exact_scope(&ROW_EXACT, f)
10652}
10653
10654/// A8W8 quantized-activation path available on THIS machine? One
10655/// switch across architectures: ARM dotprod (CMF_SDOT) or x86 AVX2
10656/// (CMF_AVX2 + the same CMF_SDOT exact-contract override).
10657#[inline]
10658pub(crate) fn a8w8_enabled() -> bool {
10659    #[cfg(target_arch = "aarch64")]
10660    {
10661        sdot_enabled()
10662    }
10663    #[cfg(target_arch = "x86_64")]
10664    {
10665        avx2_a8w8_enabled()
10666    }
10667    #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
10668    {
10669        false
10670    }
10671}
10672
10673/// int8·int8 dot dispatch: SDOT on ARM; AVX-512 VNNI (vpdpbusd) or AVX2
10674/// maddubs on x86. Callers are gated by `a8w8_enabled()`.
10675#[inline]
10676#[allow(unreachable_code)]
10677fn dot_i8_i8(w: &[u8], xq: &[i8]) -> i32 {
10678    #[cfg(target_arch = "aarch64")]
10679    unsafe {
10680        return dot_i8_sdot(w, xq);
10681    }
10682    #[cfg(target_arch = "x86_64")]
10683    unsafe {
10684        if avx512vnni_enabled() {
10685            return dot_i8_i8_vnni(w, xq);
10686        }
10687        return dot_i8_i8_avx2(w, xq);
10688    }
10689    w.iter()
10690        .zip(xq)
10691        .map(|(&a, &b)| (a as i8) as i32 * b as i32)
10692        .sum()
10693}
10694
10695/// AVX-512 VNNI available? (F+BW+VL+VNNI; `CMF_AVX512=0` falls back to
10696/// AVX2.) VL matters: short 32-byte groups (q4/vbit) ride the 256-bit
10697/// `vpdpbusd` encoding.
10698#[cfg(target_arch = "x86_64")]
10699fn avx512vnni_enabled() -> bool {
10700    use std::sync::OnceLock;
10701    static ON: OnceLock<bool> = OnceLock::new();
10702    *ON.get_or_init(|| {
10703        std::env::var("CMF_AVX512")
10704            .map(|v| v != "0")
10705            .unwrap_or(true)
10706            && std::arch::is_x86_feature_detected!("avx512f")
10707            && std::arch::is_x86_feature_detected!("avx512bw")
10708            && std::arch::is_x86_feature_detected!("avx512vl")
10709            && std::arch::is_x86_feature_detected!("avx512vnni")
10710    })
10711}
10712
10713/// Grouped-codec VNNI arms (the q4t/q4b/q1/q1t tile kernels): default
10714/// ON where AVX-512 VNNI exists (`CMF_VNNI_TILES=0` opt-out). Measured
10715/// on Ryzen 7950X (Zen4, 3 alternating process pairs, blocked GEMM
10716/// 4864×896 b=256): q4t 63→68 GF/s (+8%), q1 53→56 (+6%), q4b 72→75
10717/// (+4%) — consistent, no leg regressed. The tile kernels keep a
10718/// horizontal reduce per 32-weight group, so the `vpdpbusd` saving is
10719/// smaller than the long-dot q8 win (+13%), but it is real and free.
10720#[cfg(target_arch = "x86_64")]
10721fn vnni_tiles_enabled() -> bool {
10722    use std::sync::OnceLock;
10723    static ON: OnceLock<bool> = OnceLock::new();
10724    *ON.get_or_init(|| {
10725        std::env::var("CMF_VNNI_TILES")
10726            .map(|v| v != "0")
10727            .unwrap_or(true)
10728            && avx512vnni_enabled()
10729    })
10730}
10731
10732/// One 256-bit u8×i8 dot → i32 via `vpdpbusd` into a fresh accumulator
10733/// plus the same horizontal reduce the AVX2 kernels use. Products are
10734/// bounded (|w| ≤ 8 or ≤ 1), so maddubs never saturated — the i32 sum
10735/// is bit-identical to the maddubs+madd pair it replaces.
10736#[cfg(target_arch = "x86_64")]
10737#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
10738#[inline]
10739unsafe fn dpbusd_hsum(aw: core::arch::x86_64::__m256i, xs: core::arch::x86_64::__m256i) -> i32 {
10740    // SAFETY: pure register math.
10741    unsafe {
10742        use core::arch::x86_64::*;
10743        let d = _mm256_dpbusd_epi32(_mm256_setzero_si256(), aw, xs);
10744        let hi128 = _mm256_extracti128_si256::<1>(d);
10745        let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
10746        let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
10747        let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
10748        _mm_cvtsi128_si32(s32)
10749    }
10750}
10751
10752/// int8·int8 via AVX-512 VNNI: `vpdpbusd` fuses the maddubs+madd+add
10753/// triple into one u8×i8 dot-accumulate. AVX-512 has no vpsignb, so the
10754/// |w|·sign(x,w) trick becomes |w| × (x negated where w<0) via a mask
10755/// subtract — w==0 lanes contribute 0 through |w|=0 either way.
10756#[cfg(target_arch = "x86_64")]
10757#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
10758unsafe fn dot_i8_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
10759    // SAFETY: callers uphold slice-length contracts (see call sites).
10760    unsafe {
10761        use core::arch::x86_64::*;
10762        let n = w.len();
10763        let mut j = 0usize;
10764        let mut total: i32;
10765        // 4 independent accumulators: vpdpbusd is its own loop-carried
10766        // dependency (~5-cycle latency) — a single-acc loop runs
10767        // latency-bound and LOSES to the AVX2 maddubs kernel, measured
10768        // on Granite Rapids.
10769        {
10770            #[inline(always)]
10771            unsafe fn step(
10772                w: *const u8,
10773                x: *const i8,
10774                acc: core::arch::x86_64::__m512i,
10775            ) -> core::arch::x86_64::__m512i {
10776                unsafe {
10777                    use core::arch::x86_64::*;
10778                    let wv = _mm512_loadu_si512(w as *const _);
10779                    let xv = _mm512_loadu_si512(x as *const _);
10780                    let aw = _mm512_abs_epi8(wv);
10781                    let neg = _mm512_movepi8_mask(wv);
10782                    let sx = _mm512_mask_sub_epi8(xv, neg, _mm512_setzero_si512(), xv);
10783                    _mm512_dpbusd_epi32(acc, aw, sx)
10784                }
10785            }
10786            let (mut a0, mut a1, mut a2, mut a3) = (
10787                _mm512_setzero_si512(),
10788                _mm512_setzero_si512(),
10789                _mm512_setzero_si512(),
10790                _mm512_setzero_si512(),
10791            );
10792            while j + 256 <= n {
10793                a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
10794                a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), a1);
10795                a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), a2);
10796                a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), a3);
10797                j += 256;
10798            }
10799            while j + 64 <= n {
10800                a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
10801                j += 64;
10802            }
10803            let s01 = _mm512_add_epi32(a0, a1);
10804            let s23 = _mm512_add_epi32(a2, a3);
10805            total = _mm512_reduce_add_epi32(_mm512_add_epi32(s01, s23));
10806        }
10807        // 32-wide (q4/vbit groups are exactly 32 bytes).
10808        if j + 32 <= n {
10809            let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
10810            let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
10811            let d = _mm256_dpbusd_epi32(
10812                _mm256_setzero_si256(),
10813                _mm256_abs_epi8(wv),
10814                _mm256_sign_epi8(xv, wv),
10815            );
10816            let hi128 = _mm256_extracti128_si256::<1>(d);
10817            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
10818            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
10819            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
10820            total += _mm_cvtsi128_si32(s32);
10821            j += 32;
10822        }
10823        while j < n {
10824            total += (w[j] as i8) as i32 * xq[j] as i32;
10825            j += 1;
10826        }
10827        total
10828    }
10829}
10830
10831/// i8 row · f32 x via AVX2/FMA (x86 mirror of `dot_i8_f32_neon`).
10832#[cfg(target_arch = "x86_64")]
10833#[target_feature(enable = "avx2,fma")]
10834unsafe fn dot_i8_f32_avx2(w: &[u8], x: &[f32]) -> f32 {
10835    // SAFETY: callers uphold slice-length contracts (see call sites).
10836    unsafe {
10837        use core::arch::x86_64::*;
10838        let n = x.len();
10839        let wp = w.as_ptr();
10840        let xp = x.as_ptr();
10841        let (mut a0, mut a1) = (_mm256_setzero_ps(), _mm256_setzero_ps());
10842        let mut j = 0usize;
10843        while j + 16 <= n {
10844            let wb = _mm_loadu_si128(wp.add(j) as *const __m128i);
10845            let lo = _mm256_cvtepi8_epi32(wb);
10846            let hi = _mm256_cvtepi8_epi32(_mm_srli_si128::<8>(wb));
10847            a0 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(lo), _mm256_loadu_ps(xp.add(j)), a0);
10848            a1 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(hi), _mm256_loadu_ps(xp.add(j + 8)), a1);
10849            j += 16;
10850        }
10851        let acc = _mm256_add_ps(a0, a1);
10852        let hi128 = _mm256_extractf128_ps::<1>(acc);
10853        let s128 = _mm_add_ps(_mm256_castps256_ps128(acc), hi128);
10854        let s64 = _mm_add_ps(s128, _mm_movehl_ps(s128, s128));
10855        let s32 = _mm_add_ss(s64, _mm_shuffle_ps::<1>(s64, s64));
10856        let mut sum = _mm_cvtss_f32(s32);
10857        while j < n {
10858            sum += (*wp.add(j) as i8) as f32 * *xp.add(j);
10859            j += 1;
10860        }
10861        sum
10862    }
10863}
10864
10865/// int8(weight)·int8(activation) → i32 via AVX2 maddubs — the x86
10866/// analogue of the SDOT A8W8 path. `maddubs` takes u8×i8, so the
10867/// standard sign trick applies: |w| × sign(x, w) ≡ w × x per lane.
10868/// Pair saturation is safe: |w|≤128, |x|≤127 → 2·128·127 < 32767.
10869#[cfg(target_arch = "x86_64")]
10870#[target_feature(enable = "avx2")]
10871unsafe fn dot_i8_i8_avx2(w: &[u8], xq: &[i8]) -> i32 {
10872    // SAFETY: callers uphold slice-length contracts (see call sites).
10873    unsafe {
10874        use core::arch::x86_64::*;
10875        let n = w.len();
10876        let ones = _mm256_set1_epi16(1);
10877        let mut acc = _mm256_setzero_si256();
10878        let mut j = 0usize;
10879        while j + 32 <= n {
10880            let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
10881            let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
10882            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
10883            acc = _mm256_add_epi32(acc, _mm256_madd_epi16(p16, ones));
10884            j += 32;
10885        }
10886        let hi128 = _mm256_extracti128_si256::<1>(acc);
10887        let s128 = _mm_add_epi32(_mm256_castsi256_si128(acc), hi128);
10888        let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
10889        let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
10890        let mut s = _mm_cvtsi128_si32(s32);
10891        while j < n {
10892            s += (w[j] as i8) as i32 * xq[j] as i32;
10893            j += 1;
10894        }
10895        s
10896    }
10897}
10898
10899/// smmla 2×4: one instruction covers a 2-row × 2-activation × 8-deep
10900/// tile (32 MACs vs sdot's 16) — the weight pair loads once per 8-k
10901/// slice as a combined 2×8 register and meets two activation pairs.
10902#[cfg(target_arch = "aarch64")]
10903#[target_feature(enable = "neon,i8mm")]
10904unsafe fn dot_i8_smmla_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
10905    // SAFETY: callers uphold slice-length contracts.
10906    unsafe {
10907        use core::arch::aarch64::*;
10908        use core::arch::asm;
10909        let n = w0.len();
10910        let w0p = w0.as_ptr() as *const i8;
10911        let w1p = w1.as_ptr() as *const i8;
10912        // acc01 holds [c(r0,x0) c(r0,x1) c(r1,x0) c(r1,x1)]; acc23 the
10913        // same for x2/x3.
10914        let mut acc01 = vdupq_n_s32(0);
10915        let mut acc23 = vdupq_n_s32(0);
10916        let mut i = 0usize;
10917        while i + 8 <= n {
10918            let wa = vcombine_s8(vld1_s8(w0p.add(i)), vld1_s8(w1p.add(i)));
10919            let xb01 = vcombine_s8(
10920                vld1_s8(xs[0].as_ptr().add(i)),
10921                vld1_s8(xs[1].as_ptr().add(i)),
10922            );
10923            let xb23 = vcombine_s8(
10924                vld1_s8(xs[2].as_ptr().add(i)),
10925                vld1_s8(xs[3].as_ptr().add(i)),
10926            );
10927            asm!(
10928                "smmla {a01:v}.4s, {w:v}.16b, {x01:v}.16b",
10929                "smmla {a23:v}.4s, {w:v}.16b, {x23:v}.16b",
10930                a01 = inout(vreg) acc01, a23 = inout(vreg) acc23,
10931                w = in(vreg) wa, x01 = in(vreg) xb01, x23 = in(vreg) xb23,
10932                options(pure, nomem, nostack),
10933            );
10934            i += 8;
10935        }
10936        let mut out = [[0i32; 4]; 2];
10937        let a01: [i32; 4] = core::mem::transmute(acc01);
10938        let a23: [i32; 4] = core::mem::transmute(acc23);
10939        out[0][0] = a01[0];
10940        out[0][1] = a01[1];
10941        out[1][0] = a01[2];
10942        out[1][1] = a01[3];
10943        out[0][2] = a23[0];
10944        out[0][3] = a23[1];
10945        out[1][2] = a23[2];
10946        out[1][3] = a23[3];
10947        if i < n {
10948            for (k, x) in xs.iter().enumerate() {
10949                for j in i..n {
10950                    out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
10951                    out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
10952                }
10953            }
10954        }
10955        out
10956    }
10957}
10958
10959/// ARM twin of the x86 blocked prefill GEMM: two weight rows stay in
10960/// registers across four activation streams, eight sdot accumulators.
10961/// (The per-row form re-read each W row once per activation.)
10962#[cfg(target_arch = "aarch64")]
10963#[target_feature(enable = "neon,dotprod")]
10964unsafe fn dot_i8_sdot_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
10965    // SAFETY: callers uphold slice-length contracts.
10966    unsafe {
10967        use core::arch::aarch64::*;
10968        use core::arch::asm;
10969        let n = w0.len();
10970        let w0p = w0.as_ptr() as *const i8;
10971        let w1p = w1.as_ptr() as *const i8;
10972        let mut acc = [[vdupq_n_s32(0); 4]; 2];
10973        let mut i = 0usize;
10974        while i + 16 <= n {
10975            let wv0 = vld1q_s8(w0p.add(i));
10976            let wv1 = vld1q_s8(w1p.add(i));
10977            for (k, x) in xs.iter().enumerate() {
10978                let xv = vld1q_s8(x.as_ptr().add(i));
10979                let (mut a0, mut a1) = (acc[0][k], acc[1][k]);
10980                asm!(
10981                    "sdot {a0:v}.4s, {w0:v}.16b, {x:v}.16b",
10982                    "sdot {a1:v}.4s, {w1:v}.16b, {x:v}.16b",
10983                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
10984                    w0 = in(vreg) wv0, w1 = in(vreg) wv1, x = in(vreg) xv,
10985                    options(pure, nomem, nostack),
10986                );
10987                acc[0][k] = a0;
10988                acc[1][k] = a1;
10989            }
10990            i += 16;
10991        }
10992        let mut out = [[0i32; 4]; 2];
10993        for r in 0..2 {
10994            for k in 0..4 {
10995                out[r][k] = vaddvq_s32(acc[r][k]);
10996            }
10997        }
10998        if i < n {
10999            for (k, x) in xs.iter().enumerate() {
11000                for j in i..n {
11001                    out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
11002                    out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
11003                }
11004            }
11005        }
11006        out
11007    }
11008}
11009
11010/// Blocked 2 weight rows × 4 activations for the prefill GEMM
11011/// (roadmap P0: packed panels + multi-row accumulators). The two rows'
11012/// abs() live in registers across all four activation streams; the
11013/// sign-fixup is recomputed per pair (the price of the maddubs trick).
11014/// Returns raw i8·i8 dots; the caller applies scales and outliers.
11015#[cfg(target_arch = "x86_64")]
11016#[target_feature(enable = "avx2")]
11017unsafe fn dot_i8_i8_avx2_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
11018    // SAFETY: callers uphold slice-length contracts.
11019    unsafe {
11020        use core::arch::x86_64::*;
11021        let n = w0.len();
11022        let ones = _mm256_set1_epi16(1);
11023        let mut acc = [[_mm256_setzero_si256(); 4]; 2];
11024        let mut j = 0usize;
11025        while j + 32 <= n {
11026            let wv0 = _mm256_loadu_si256(w0.as_ptr().add(j) as *const __m256i);
11027            let wv1 = _mm256_loadu_si256(w1.as_ptr().add(j) as *const __m256i);
11028            let aw0 = _mm256_abs_epi8(wv0);
11029            let aw1 = _mm256_abs_epi8(wv1);
11030            for (k, x) in xs.iter().enumerate() {
11031                let xv = _mm256_loadu_si256(x.as_ptr().add(j) as *const __m256i);
11032                let p0 = _mm256_maddubs_epi16(aw0, _mm256_sign_epi8(xv, wv0));
11033                acc[0][k] = _mm256_add_epi32(acc[0][k], _mm256_madd_epi16(p0, ones));
11034                let p1 = _mm256_maddubs_epi16(aw1, _mm256_sign_epi8(xv, wv1));
11035                acc[1][k] = _mm256_add_epi32(acc[1][k], _mm256_madd_epi16(p1, ones));
11036            }
11037            j += 32;
11038        }
11039        let mut out = [[0i32; 4]; 2];
11040        for r in 0..2 {
11041            for k in 0..4 {
11042                let a = acc[r][k];
11043                let hi128 = _mm256_extracti128_si256::<1>(a);
11044                let s128 = _mm_add_epi32(_mm256_castsi256_si128(a), hi128);
11045                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
11046                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
11047                out[r][k] = _mm_cvtsi128_si32(s32);
11048            }
11049        }
11050        if j < n {
11051            for (k, x) in xs.iter().enumerate() {
11052                for i in j..n {
11053                    out[0][k] += (w0[i] as i8) as i32 * x[i] as i32;
11054                    out[1][k] += (w1[i] as i8) as i32 * x[i] as i32;
11055                }
11056            }
11057        }
11058        out
11059    }
11060}
11061
11062/// AVX2/VNNI q8 row dot with exact outlier correction (x86 mirror of
11063/// `row_dot_sdot` — same A8W8 contract). With AVX-512 VNNI the row goes
11064/// through the bias trick: Σ(w+128)·x via pure `vpdpbusd` (no per-lane
11065/// sign fixups), corrected by −128·Σx with Σx precomputed per split.
11066#[cfg(target_arch = "x86_64")]
11067#[inline]
11068fn row_dot_avx2(row: &[u8], act: &SplitAct) -> f32 {
11069    let dot = if avx512vnni_enabled() && row.len() >= 64 {
11070        (unsafe { dot_u8p128_i8_vnni(row, &act.xq) }) - 128 * act.xsum
11071    } else {
11072        unsafe { dot_i8_i8_avx2(row, &act.xq) }
11073    };
11074    let mut acc = dot as f32 * act.sx;
11075    for &(j, xv) in &act.outliers {
11076        acc += (row[j] as i8) as f32 * xv;
11077    }
11078    acc
11079}
11080
11081/// Σ (w[i]+128)·x[i] via pure `vpdpbusd` — the caller subtracts
11082/// 128·Σx. Four independent accumulators (dpbusd is ~5-cycle latency;
11083/// a single-acc loop runs latency-bound, measured on Granite Rapids).
11084#[cfg(target_arch = "x86_64")]
11085#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
11086unsafe fn dot_u8p128_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
11087    // SAFETY: callers uphold slice-length contracts (see call sites).
11088    unsafe {
11089        use core::arch::x86_64::*;
11090        let n = w.len();
11091        let flip = _mm512_set1_epi8(-128); // XOR 0x80: i8 w → u8 (w+128)
11092        #[inline(always)]
11093        unsafe fn step(
11094            w: *const u8,
11095            x: *const i8,
11096            flip: core::arch::x86_64::__m512i,
11097            acc: core::arch::x86_64::__m512i,
11098        ) -> core::arch::x86_64::__m512i {
11099            unsafe {
11100                use core::arch::x86_64::*;
11101                let wv = _mm512_xor_si512(_mm512_loadu_si512(w as *const _), flip);
11102                _mm512_dpbusd_epi32(acc, wv, _mm512_loadu_si512(x as *const _))
11103            }
11104        }
11105        let (mut a0, mut a1, mut a2, mut a3) = (
11106            _mm512_setzero_si512(),
11107            _mm512_setzero_si512(),
11108            _mm512_setzero_si512(),
11109            _mm512_setzero_si512(),
11110        );
11111        let mut j = 0usize;
11112        while j + 256 <= n {
11113            a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
11114            a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), flip, a1);
11115            a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), flip, a2);
11116            a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), flip, a3);
11117            j += 256;
11118        }
11119        while j + 64 <= n {
11120            a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
11121            j += 64;
11122        }
11123        let mut total = _mm512_reduce_add_epi32(_mm512_add_epi32(
11124            _mm512_add_epi32(a0, a1),
11125            _mm512_add_epi32(a2, a3),
11126        ));
11127        // Scalar tail: (w as i8) + 128 ≡ (w as u8) ^ 0x80.
11128        while j < n {
11129            total += ((w[j] ^ 0x80) as i32) * xq[j] as i32;
11130            j += 1;
11131        }
11132        total
11133    }
11134}
11135
11136/// One q4 row via AVX2: nibbles → centered i8 (unpacklo/hi restores the
11137/// writer's flat order, same as the NEON vzip pair), maddubs against
11138/// the pre-quantized activation group, × the group's f16 scale. Pair
11139/// saturation safe: |w|≤8, |x|≤127 → 2·8·127 ≪ 32767. Mirror of
11140/// `dot_q4_row_sdot`.
11141#[cfg(target_arch = "x86_64")]
11142#[target_feature(enable = "avx2")]
11143unsafe fn dot_q4_row_avx2(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
11144    // SAFETY: callers uphold slice-length contracts (16 packed bytes and
11145    // 2 scale bytes per group; xq.len() == gpr·GROUP_SIZE).
11146    unsafe {
11147        use core::arch::x86_64::*;
11148        let lomask = _mm_set1_epi8(0x0F);
11149        let eight = _mm256_set1_epi8(8);
11150        let ones = _mm256_set1_epi16(1);
11151        let mut acc = 0f32;
11152        for gi in 0..gpr {
11153            let g = g0 + gi;
11154            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
11155            let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
11156            let lo = _mm_and_si128(b, lomask);
11157            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
11158            let w = _mm256_sub_epi8(
11159                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
11160                eight,
11161            );
11162            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
11163            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
11164            let d = _mm256_madd_epi16(p16, ones);
11165            let hi128 = _mm256_extracti128_si256::<1>(d);
11166            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
11167            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
11168            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
11169            acc += _mm_cvtsi128_si32(s32) as f32 * s;
11170        }
11171        acc
11172    }
11173}
11174
11175/// Two-activation q4 row via AVX2: nibbles unpacked ONCE per group,
11176/// both activations dotted against the same centered i8 register.
11177#[cfg(target_arch = "x86_64")]
11178#[target_feature(enable = "avx2")]
11179unsafe fn dot_q4_row_avx2_2(
11180    packed: &[u8],
11181    scales: &[u8],
11182    g0: usize,
11183    gpr: usize,
11184    xq1: &[i8],
11185    xq2: &[i8],
11186) -> (f32, f32) {
11187    // SAFETY: callers uphold slice-length contracts (see dot_q4_row_avx2).
11188    unsafe {
11189        use core::arch::x86_64::*;
11190        let lomask = _mm_set1_epi8(0x0F);
11191        let eight = _mm256_set1_epi8(8);
11192        let ones = _mm256_set1_epi16(1);
11193        let (mut acc1, mut acc2) = (0f32, 0f32);
11194        #[inline(always)]
11195        unsafe fn hsum(d: core::arch::x86_64::__m256i) -> i32 {
11196            unsafe {
11197                use core::arch::x86_64::*;
11198                let hi128 = _mm256_extracti128_si256::<1>(d);
11199                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
11200                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
11201                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
11202                _mm_cvtsi128_si32(s32)
11203            }
11204        }
11205        for gi in 0..gpr {
11206            let g = g0 + gi;
11207            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
11208            let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
11209            let lo = _mm_and_si128(b, lomask);
11210            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
11211            let w = _mm256_sub_epi8(
11212                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
11213                eight,
11214            );
11215            let aw = _mm256_abs_epi8(w);
11216            let x1 = _mm256_loadu_si256(xq1.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
11217            let x2 = _mm256_loadu_si256(xq2.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
11218            let d1 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x1, w)), ones);
11219            let d2 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x2, w)), ones);
11220            acc1 += hsum(d1) as f32 * s;
11221            acc2 += hsum(d2) as f32 * s;
11222        }
11223        (acc1, acc2)
11224    }
11225}
11226
11227/// One q8 row range via AVX2 (x86 mirror of `q8_range_sdot`).
11228#[cfg(target_arch = "x86_64")]
11229fn q8_range_avx2(
11230    q: &[u8],
11231    row_scale: &[f32],
11232    act: &SplitAct,
11233    cols: usize,
11234    out_addr: SendMut,
11235    start: usize,
11236    end: usize,
11237) {
11238    for o in start..end {
11239        let v = row_dot_avx2(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
11240        // SAFETY: disjoint row ranges per worker.
11241        unsafe { *out_addr.at(o) = v };
11242    }
11243}
11244
11245/// Two-input q8 row range via AVX2 (x86 mirror of `q8_range2_sdot`).
11246#[cfg(target_arch = "x86_64")]
11247#[allow(clippy::too_many_arguments)]
11248fn q8_range2_avx2(
11249    q: &[u8],
11250    row_scale: &[f32],
11251    a1: &SplitAct,
11252    a2: &SplitAct,
11253    cols: usize,
11254    p1: SendMut,
11255    p2: SendMut,
11256    start: usize,
11257    end: usize,
11258) {
11259    for o in start..end {
11260        let row = &q[o * cols..(o + 1) * cols];
11261        // SAFETY: disjoint row ranges per worker.
11262        unsafe {
11263            *p1.at(o) = row_dot_avx2(row, a1) * row_scale[o];
11264            *p2.at(o) = row_dot_avx2(row, a2) * row_scale[o];
11265        }
11266    }
11267}
11268
11269// ───────────────────── A8W8 SDOT path (port of vmfcore, ×1.78 decode) ─────────────────────
11270
11271/// ARMv8.6 i8mm (smmla): 32 int8 MACs per instruction vs sdot's 16 —
11272/// yet MEASURED 2.4× SLOWER than the blocked sdot on Apple silicon
11273/// (108 vs 264 GF/s): the on-the-fly vcombine packing and the two-
11274/// accumulator dependency chain swamp the MAC advantage, and Apple's
11275/// four SIMD pipes already keep sdot fed. OPT-IN (CMF_I8MM=1) for
11276/// field trials on Cortex-A710/X-class parts with two pipes, where the
11277/// balance may differ; a pre-interleaved weight layout (repack infra)
11278/// is the known path if it ever earns its keep.
11279#[cfg(target_arch = "aarch64")]
11280fn i8mm_enabled() -> bool {
11281    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
11282    *ON.get_or_init(|| {
11283        std::env::var("CMF_I8MM").map(|v| v == "1").unwrap_or(false)
11284            && std::arch::is_aarch64_feature_detected!("i8mm")
11285    })
11286}
11287
11288/// SDOT enabled? Default ON when the CPU has ARMv8.2 dotprod;
11289/// `CMF_SDOT=0` disables (falls back to i8×f32 NEON).
11290/// (On non-ARM release builds only the test tolerance switch calls it.)
11291#[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
11292fn sdot_enabled() -> bool {
11293    if FLOAT_ACTIVATIONS.get() {
11294        return false;
11295    }
11296    use std::sync::OnceLock;
11297    static ON: OnceLock<bool> = OnceLock::new();
11298    *ON.get_or_init(|| {
11299        let want = std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true);
11300        if !want {
11301            return false;
11302        }
11303
11304        #[cfg(target_arch = "aarch64")]
11305        {
11306            if std::arch::is_aarch64_feature_detected!("dotprod") {
11307                return true;
11308            }
11309            #[cfg(target_os = "android")]
11310            {
11311                if let Ok(cpuinfo) = std::fs::read_to_string("/proc/cpuinfo") {
11312                    if cpuinfo.lines().any(|l| {
11313                        (l.starts_with("Features") || l.starts_with("features"))
11314                            && l.contains("asimddp")
11315                    }) {
11316                        return true;
11317                    }
11318                }
11319            }
11320            false
11321        }
11322        #[cfg(not(target_arch = "aarch64"))]
11323        {
11324            false
11325        }
11326    })
11327}
11328
11329/// Two-field activation split (≡ vmfcore `q8_split_prep`): outlier
11330/// channels (>8·rms) are computed exactly in f32; the bulk (outliers
11331/// zeroed → clean absmax) goes through int8 SDOT. Computed ONCE per
11332/// matvec, shared by all rows/workers.
11333struct SplitAct {
11334    xq: Vec<i8>,
11335    sx: f32,
11336    outliers: Vec<(usize, f32)>,
11337    /// Σ xq — the VNNI bias-trick correction (`(w+128)·x` sums need
11338    /// `−128·Σx`); one i32 per split, computed once per matvec.
11339    #[cfg_attr(not(target_arch = "x86_64"), allow(dead_code))]
11340    xsum: i32,
11341}
11342
11343thread_local! {
11344    /// Recycled xq buffers: split_act runs for every matvec (~200/token)
11345    /// and its hidden-size allocation was steady-state heap churn.
11346    static XQ_FREE: std::cell::RefCell<Vec<Vec<i8>>> =
11347        const { std::cell::RefCell::new(Vec::new()) };
11348}
11349
11350impl Drop for SplitAct {
11351    fn drop(&mut self) {
11352        let buf = std::mem::take(&mut self.xq);
11353        if buf.capacity() > 0 {
11354            XQ_FREE.with(|f| {
11355                let mut f = f.borrow_mut();
11356                if f.len() < 16 {
11357                    f.push(buf);
11358                }
11359            });
11360        }
11361    }
11362}
11363
11364thread_local! {
11365    /// One scratch row per WORKER, kept for the life of the thread.
11366    ///
11367    /// The kernels take a row of group scales per dispatch, and a fresh
11368    /// `vec![0f32; gpr]` inside the closure is one allocation per worker per
11369    /// dispatch — on the release checkpoint about six thousand a token, a
11370    /// quarter of everything the benchmark counts.
11371    static KROW: UnsafeCell<Vec<f32>> = const { UnsafeCell::new(Vec::new()) };
11372    /// Two scratch rows for a gate/up pair.  Q2TP and Q4TP fused SwiGLU
11373    /// decode each projection's ladder independently, but both ladders have
11374    /// the same lifetime; retaining them per worker removes two heap trips
11375    /// from every dispatched expert pair.
11376    static KROWS: UnsafeCell<[Vec<f32>; 2]> =
11377        const { UnsafeCell::new([Vec::new(), Vec::new()]) };
11378}
11379
11380/// Borrow `n` floats of the calling worker's scratch. Nothing inside a
11381/// kernel body borrows it again, which is what keeps the RefCell honest.
11382#[inline]
11383fn with_krow<R>(n: usize, f: impl FnOnce(&mut [f32]) -> R) -> R {
11384    KROW.with(|s| {
11385        // SAFETY: `KROW` is thread-local and this function does not recurse;
11386        // each worker has exclusive access to its own scratch row.
11387        let b = unsafe { &mut *s.get() };
11388        if b.len() < n {
11389            b.resize(n, 0.0);
11390        }
11391        f(&mut b[..n])
11392    })
11393}
11394
11395/// `t.round().clamp(-127.0, 127.0) as i8`, bit for bit, without the libm
11396/// call. On baseline x86-64 (no SSE4.1 `roundps`) `f32::round` is a
11397/// function call per element, and split_act runs it over every hidden
11398/// state before every matvec: measured 27 us a call on a 2048-wide
11399/// activation on an EPYC 7763 — 5.4 ms of a 55 ms decode token, all of
11400/// it on the caller's thread while thirty workers wait. Clamping first is
11401/// equivalent (round is monotonic and ±127 are integers), and after the
11402/// clamp `t - trunc(t)` is exact, so the half-away-from-zero decision is
11403/// the one `round` makes. NaN clamps to NaN and converts to 0, as before.
11404/// The loop vectorizes (cvttps2dq + compare/select).
11405#[inline(always)]
11406fn q8_round(t: f32) -> i8 {
11407    let t = t.clamp(-127.0, 127.0);
11408    let i = t as i32;
11409    let f = t - i as f32;
11410    let r = if f >= 0.5 {
11411        i + 1
11412    } else if f <= -0.5 {
11413        i - 1
11414    } else {
11415        i
11416    };
11417    r as i8
11418}
11419
11420#[inline]
11421fn with_krows<R>(n: usize, f: impl FnOnce(&mut [f32], &mut [f32]) -> R) -> R {
11422    KROWS.with(|s| {
11423        // SAFETY: KROWS is thread-local and no kernel body recursively calls
11424        // this helper; each participant owns its two vectors exclusively.
11425        let b = unsafe { &mut *s.get() };
11426        for row in b.iter_mut() {
11427            if row.len() < n {
11428                row.resize(n, 0.0);
11429            }
11430        }
11431        let (a, b) = b.split_at_mut(1);
11432        f(&mut a[0][..n], &mut b[0][..n])
11433    })
11434}
11435
11436#[inline]
11437fn silu_mul_limited(mut gate: f32, mut up: f32, limit: f32) -> f32 {
11438    if limit > 0.0 {
11439        up = up.clamp(-limit, limit);
11440        gate = gate.min(limit);
11441    }
11442    gate / (1.0 + (-gate).exp()) * up
11443}
11444
11445fn split_act(x: &[f32]) -> SplitAct {
11446    let _prof = crate::cpuprof::time(crate::cpuprof::Slot::SplitAct);
11447    let n = x.len();
11448    let rms = (x.iter().map(|&v| (v * v) as f64).sum::<f64>() / n.max(1) as f64).sqrt() as f32;
11449    let thr = 8.0 * rms;
11450    // One pass: collect outliers and the bulk absmax (outliers excluded —
11451    // identical to the old zero-then-fold over a copied buffer, minus the
11452    // full-vector copy).
11453    let mut outliers: Vec<(usize, f32)> = Vec::new();
11454    let mut amax = 0f32;
11455    for (j, &v) in x.iter().enumerate() {
11456        let a = v.abs();
11457        if a > thr {
11458            outliers.push((j, v));
11459        } else if a > amax {
11460            amax = a;
11461        }
11462    }
11463    let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
11464    let inv = 1.0 / sx;
11465    let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
11466    xq.clear();
11467    xq.reserve(n);
11468    if outliers.is_empty() {
11469        xq.extend(
11470            x.iter()
11471                .map(|&v| q8_round(v * inv)),
11472        );
11473    } else {
11474        // Outlier slots quantize to 0 (their exact term is added later).
11475        xq.extend(x.iter().map(|&v| {
11476            if v.abs() > thr {
11477                0
11478            } else {
11479                q8_round(v * inv)
11480            }
11481        }));
11482    }
11483    let xsum = xq.iter().map(|&v| v as i32).sum();
11484    SplitAct {
11485        xq,
11486        sx,
11487        outliers,
11488        xsum,
11489    }
11490}
11491
11492fn split_act_q8_2f(x: &[f32], col: &[f32]) -> SplitAct {
11493    let _prof = crate::cpuprof::time(crate::cpuprof::Slot::SplitAct);
11494    let n = x.len();
11495    let rms = (x
11496        .iter()
11497        .zip(col)
11498        .map(|(&a, &c)| {
11499            let v = a * c;
11500            (v * v) as f64
11501        })
11502        .sum::<f64>()
11503        / n.max(1) as f64)
11504        .sqrt() as f32;
11505    let thr = 8.0 * rms;
11506
11507    let mut outliers = Vec::new();
11508    let mut amax = 0f32;
11509    for (j, (&a, &c)) in x.iter().zip(col).enumerate() {
11510        let v = a * c;
11511        let s = v.abs();
11512        if s > thr {
11513            outliers.push((j, v));
11514        } else if s > amax {
11515            amax = s;
11516        }
11517    }
11518
11519    let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
11520    let inv = 1.0 / sx;
11521    let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
11522    xq.clear();
11523    xq.reserve(n);
11524    if outliers.is_empty() {
11525        xq.extend(
11526            x.iter()
11527                .zip(col)
11528                .map(|(&a, &c)| q8_round((a * c) * inv)),
11529        );
11530    } else {
11531        xq.extend(x.iter().zip(col).map(|(&a, &c)| {
11532            let v = a * c;
11533            if v.abs() > thr {
11534                0
11535            } else {
11536                q8_round(v * inv)
11537            }
11538        }));
11539    }
11540    let xsum = xq.iter().map(|&v| v as i32).sum();
11541    SplitAct {
11542        xq,
11543        sx,
11544        outliers,
11545        xsum,
11546    }
11547}
11548
11549/// int8(weight)·int8(activation) → i32 via `sdot` (inline asm — the
11550/// vdotq intrinsic is unstable; port of vmfcore `dot_i8_sdot`).
11551#[cfg(target_arch = "aarch64")]
11552#[target_feature(enable = "neon,dotprod")]
11553unsafe fn dot_i8_sdot(w: &[u8], xq: &[i8]) -> i32 {
11554    // SAFETY: callers uphold slice-length contracts (see call sites).
11555    unsafe {
11556        use core::arch::aarch64::*;
11557        use core::arch::asm;
11558        let wp = w.as_ptr() as *const i8;
11559        let n = w.len();
11560        let (mut a0, mut a1, mut a2, mut a3) = (
11561            vdupq_n_s32(0),
11562            vdupq_n_s32(0),
11563            vdupq_n_s32(0),
11564            vdupq_n_s32(0),
11565        );
11566        let mut i = 0;
11567        while i + 64 <= n {
11568            let (w0, x0) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
11569            let (w1, x1) = (vld1q_s8(wp.add(i + 16)), vld1q_s8(xq.as_ptr().add(i + 16)));
11570            let (w2, x2) = (vld1q_s8(wp.add(i + 32)), vld1q_s8(xq.as_ptr().add(i + 32)));
11571            let (w3, x3) = (vld1q_s8(wp.add(i + 48)), vld1q_s8(xq.as_ptr().add(i + 48)));
11572            asm!(
11573                "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
11574                "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
11575                "sdot {a2:v}.4s, {w2:v}.16b, {x2:v}.16b",
11576                "sdot {a3:v}.4s, {w3:v}.16b, {x3:v}.16b",
11577                a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
11578                w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
11579                w2 = in(vreg) w2, x2 = in(vreg) x2, w3 = in(vreg) w3, x3 = in(vreg) x3,
11580                options(pure, nomem, nostack),
11581            );
11582            i += 64;
11583        }
11584        while i + 16 <= n {
11585            let (wv, xv) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
11586            asm!("sdot {a:v}.4s, {w:v}.16b, {x:v}.16b",
11587                 a = inout(vreg) a0, w = in(vreg) wv, x = in(vreg) xv, options(pure, nomem, nostack));
11588            i += 16;
11589        }
11590        let mut s = vaddvq_s32(vaddq_s32(vaddq_s32(a0, a1), vaddq_s32(a2, a3)));
11591        while i < n {
11592            s += (*wp.add(i)) as i32 * xq[i] as i32;
11593            i += 1;
11594        }
11595        s
11596    }
11597}
11598
11599/// Row-blocked SDOT: 4 output rows per pass — the activation chunk is
11600/// loaded once and reused, 4 independent accumulators hide sdot latency
11601/// (port of vmfcore `dot_i8_sdot_4rows`).
11602#[cfg(target_arch = "aarch64")]
11603#[target_feature(enable = "neon,dotprod")]
11604unsafe fn dot_i8_sdot_4rows(w0: &[u8], w1: &[u8], w2: &[u8], w3: &[u8], xq: &[i8]) -> [i32; 4] {
11605    // SAFETY: callers uphold slice-length contracts (see call sites).
11606    unsafe {
11607        use core::arch::aarch64::*;
11608        use core::arch::asm;
11609        let n = xq.len();
11610        let px = xq.as_ptr();
11611        let (p0, p1, p2, p3) = (
11612            w0.as_ptr() as *const i8,
11613            w1.as_ptr() as *const i8,
11614            w2.as_ptr() as *const i8,
11615            w3.as_ptr() as *const i8,
11616        );
11617        let (mut a0, mut a1, mut a2, mut a3) = (
11618            vdupq_n_s32(0),
11619            vdupq_n_s32(0),
11620            vdupq_n_s32(0),
11621            vdupq_n_s32(0),
11622        );
11623        let mut i = 0;
11624        while i + 16 <= n {
11625            let x = vld1q_s8(px.add(i));
11626            let v0 = vld1q_s8(p0.add(i));
11627            let v1 = vld1q_s8(p1.add(i));
11628            let v2 = vld1q_s8(p2.add(i));
11629            let v3 = vld1q_s8(p3.add(i));
11630            asm!(
11631                "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
11632                "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
11633                "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
11634                "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
11635                a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
11636                v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
11637                options(pure, nomem, nostack),
11638            );
11639            i += 16;
11640        }
11641        let mut r = [
11642            vaddvq_s32(a0),
11643            vaddvq_s32(a1),
11644            vaddvq_s32(a2),
11645            vaddvq_s32(a3),
11646        ];
11647        while i < n {
11648            let xi = *px.add(i) as i32;
11649            r[0] += (*p0.add(i)) as i32 * xi;
11650            r[1] += (*p1.add(i)) as i32 * xi;
11651            r[2] += (*p2.add(i)) as i32 * xi;
11652            r[3] += (*p3.add(i)) as i32 * xi;
11653            i += 1;
11654        }
11655        r
11656    }
11657}
11658
11659/// 4 interleaved rows in one pass: the repacked group is [r0[c], r1[c],
11660/// r2[c], r3[c]] per 16-byte chunk, so each iteration reads ONE 64-byte
11661/// line plus the shared activation chunk — a single sequential weight
11662/// stream per worker. Per-row accumulation is the same one-accumulator
11663/// scheme as `dot_i8_sdot_4rows`; integer sums are exact, so outputs
11664/// are bit-identical to the mmap-layout kernel.
11665#[cfg(target_arch = "aarch64")]
11666#[target_feature(enable = "neon,dotprod")]
11667unsafe fn dot_i8_sdot_4rows_il(g: &[u8], xq: &[i8]) -> [i32; 4] {
11668    // SAFETY: callers uphold slice-length contracts (g.len() == 4·n,
11669    // n % 16 == 0 — guaranteed by the repack gate).
11670    unsafe {
11671        use core::arch::aarch64::*;
11672        use core::arch::asm;
11673        let n = xq.len();
11674        let px = xq.as_ptr();
11675        let pg = g.as_ptr() as *const i8;
11676        let (mut a0, mut a1, mut a2, mut a3) = (
11677            vdupq_n_s32(0),
11678            vdupq_n_s32(0),
11679            vdupq_n_s32(0),
11680            vdupq_n_s32(0),
11681        );
11682        let mut i = 0;
11683        while i + 16 <= n {
11684            let x = vld1q_s8(px.add(i));
11685            let base = pg.add(4 * i);
11686            let v0 = vld1q_s8(base);
11687            let v1 = vld1q_s8(base.add(16));
11688            let v2 = vld1q_s8(base.add(32));
11689            let v3 = vld1q_s8(base.add(48));
11690            asm!(
11691                "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
11692                "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
11693                "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
11694                "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
11695                a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
11696                v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
11697                options(pure, nomem, nostack),
11698            );
11699            i += 16;
11700        }
11701        [
11702            vaddvq_s32(a0),
11703            vaddvq_s32(a1),
11704            vaddvq_s32(a2),
11705            vaddvq_s32(a3),
11706        ]
11707    }
11708}
11709
11710/// One q8 row range via SDOT (4-row blocks + tail) — the body of
11711/// `qmatvec`'s hot loop, extracted so multi-matrix jobs can drive the
11712/// SAME kernel for several tensors under one pool dispatch. `rep` — the
11713/// load-time interleaved repack (empty = mmap layout only); rows outside
11714/// full 4-row groups always come from the mmap layout.
11715#[cfg(target_arch = "aarch64")]
11716fn q8_range_sdot(
11717    q: &[u8],
11718    rep: &[u8],
11719    row_scale: &[f32],
11720    act: &SplitAct,
11721    cols: usize,
11722    out_addr: SendMut,
11723    start: usize,
11724    end: usize,
11725) {
11726    let mut o = start;
11727    // Leading rows to the group boundary (repack path only): the pool
11728    // splits row ranges arbitrarily, groups are absolute.
11729    if !rep.is_empty() {
11730        while o < end && o % 4 != 0 {
11731            let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
11732            unsafe { *out_addr.at(o) = v };
11733            o += 1;
11734        }
11735    }
11736    while o + 4 <= end {
11737        let r = if rep.is_empty() {
11738            unsafe {
11739                dot_i8_sdot_4rows(
11740                    &q[o * cols..(o + 1) * cols],
11741                    &q[(o + 1) * cols..(o + 2) * cols],
11742                    &q[(o + 2) * cols..(o + 3) * cols],
11743                    &q[(o + 3) * cols..(o + 4) * cols],
11744                    &act.xq,
11745                )
11746            }
11747        } else {
11748            unsafe { dot_i8_sdot_4rows_il(&rep[o * cols..(o + 4) * cols], &act.xq) }
11749        };
11750        for k in 0..4 {
11751            let mut acc = r[k] as f32 * act.sx;
11752            for &(j, xv) in &act.outliers {
11753                acc += (q[(o + k) * cols + j] as i8) as f32 * xv;
11754            }
11755            // SAFETY: disjoint row ranges per worker.
11756            unsafe { *out_addr.at(o + k) = acc * row_scale[o + k] };
11757        }
11758        o += 4;
11759    }
11760    while o < end {
11761        let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
11762        unsafe { *out_addr.at(o) = v };
11763        o += 1;
11764    }
11765}
11766
11767/// Two-input q8 row range via SDOT — `qmatvec2`'s hot loop, extracted
11768/// for the fused pair multi-matrix job (`matvec2_many`).
11769#[cfg(target_arch = "aarch64")]
11770#[allow(clippy::too_many_arguments)]
11771fn q8_range2_sdot(
11772    q: &[u8],
11773    row_scale: &[f32],
11774    a1: &SplitAct,
11775    a2: &SplitAct,
11776    cols: usize,
11777    p1: SendMut,
11778    p2: SendMut,
11779    start: usize,
11780    end: usize,
11781) {
11782    for o in start..end {
11783        let row = &q[o * cols..(o + 1) * cols];
11784        // SAFETY: disjoint row ranges per worker.
11785        unsafe {
11786            *p1.at(o) = row_dot_sdot(row, a1) * row_scale[o];
11787            *p2.at(o) = row_dot_sdot(row, a2) * row_scale[o];
11788        }
11789    }
11790}
11791
11792/// Two-input q8 row range, f32 kernel (non-SDOT) — same extraction.
11793#[allow(clippy::too_many_arguments)]
11794fn q8_range2_f32(
11795    q: &[u8],
11796    row_scale: &[f32],
11797    x1: &[f32],
11798    x2: &[f32],
11799    cols: usize,
11800    p1: SendMut,
11801    p2: SendMut,
11802    start: usize,
11803    end: usize,
11804) {
11805    for o in start..end {
11806        let row = &q[o * cols..(o + 1) * cols];
11807        // SAFETY: disjoint row ranges per worker.
11808        unsafe {
11809            *p1.at(o) = dot_i8_f32(row, x1) * row_scale[o];
11810            *p2.at(o) = dot_i8_f32(row, x2) * row_scale[o];
11811        }
11812    }
11813}
11814
11815/// Scalar/NEON-f32 q8 row range (non-SDOT platforms) — same extraction.
11816fn q8_range_f32(
11817    q: &[u8],
11818    row_scale: &[f32],
11819    xs: &[f32],
11820    cols: usize,
11821    out_addr: SendMut,
11822    start: usize,
11823    end: usize,
11824) {
11825    for o in start..end {
11826        let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
11827        // SAFETY: disjoint row ranges per worker.
11828        unsafe { *out_addr.at(o) = v };
11829    }
11830}
11831
11832/// One q8 row against a split activation, portable: the per-arch fast
11833/// dots where they exist, the exact scalar loop elsewhere. The scalar
11834/// arm is also the test oracle for both fast arms.
11835#[inline]
11836fn q8_row_dot(row: &[u8], act: &SplitAct) -> f32 {
11837    #[cfg(target_arch = "aarch64")]
11838    return row_dot_sdot(row, act);
11839    #[cfg(target_arch = "x86_64")]
11840    return row_dot_avx2(row, act);
11841    #[allow(unreachable_code)]
11842    q8_row_dot_scalar(row, act)
11843}
11844
11845#[allow(dead_code)]
11846fn q8_row_dot_scalar(row: &[u8], act: &SplitAct) -> f32 {
11847    let mut acc = 0i32;
11848    for (k, &b) in row.iter().enumerate() {
11849        acc += (b as i8) as i32 * act.xq[k] as i32;
11850    }
11851    let mut acc = acc as f32 * act.sx;
11852    for &(j, xv) in &act.outliers {
11853        acc += (row[j] as i8) as f32 * xv;
11854    }
11855    acc
11856}
11857
11858/// SDOT row dot with exact outlier correction:
11859/// `dot = sdot(w, xq)·sx + Σ_outl w[j]·x[j]` (then × row_scale by caller).
11860#[cfg(target_arch = "aarch64")]
11861#[inline]
11862fn row_dot_sdot(row: &[u8], act: &SplitAct) -> f32 {
11863    let mut acc = unsafe { dot_i8_sdot(row, &act.xq) } as f32 * act.sx;
11864    for &(j, xv) in &act.outliers {
11865        acc += (row[j] as i8) as f32 * xv;
11866    }
11867    acc
11868}
11869
11870/// One q4 row via SDOT: each 32-group's nibbles unpack to centered i8
11871/// (nib−8 ∈ [−8,7]), int8×int8 `sdot` against the pre-quantized
11872/// activation group, × the group's f16 scale. Returns Σ_g dot_g·s_g;
11873/// the caller multiplies by the activation scale and adds the exact
11874/// outlier terms (port of vmfcore `dot_q4_block_sdot`, +23% measured).
11875/// Nibble order matches the writer: element 2k = low nibble, 2k+1 = high
11876/// → zip(lo,hi) restores flat order.
11877#[cfg(target_arch = "aarch64")]
11878#[target_feature(enable = "neon,dotprod")]
11879unsafe fn dot_q4_row_sdot(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
11880    // SAFETY: callers uphold slice-length contracts (16 packed bytes and
11881    // 2 scale bytes per group; xq.len() == gpr·GROUP_SIZE).
11882    unsafe {
11883        use core::arch::aarch64::*;
11884        use core::arch::asm;
11885        let lomask = vdupq_n_u8(0x0F);
11886        let eight = vdupq_n_s8(8);
11887        let mut acc = 0f32;
11888        for gi in 0..gpr {
11889            let g = g0 + gi;
11890            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
11891            let b = vld1q_u8(packed.as_ptr().add(g * 16));
11892            let lo = vandq_u8(b, lomask);
11893            let hi = vshrq_n_u8::<4>(b);
11894            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
11895            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
11896            let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
11897            let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
11898            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
11899            asm!(
11900                "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
11901                "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
11902                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
11903                e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
11904                options(pure, nomem, nostack),
11905            );
11906            acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
11907        }
11908        acc
11909    }
11910}
11911
11912/// Two-activation q4 row via SDOT: the nibble unpack (the expensive
11913/// part) happens ONCE per group; both pre-quantized activations are
11914/// dotted against the same centered i8 registers. Per-lane math matches
11915/// `dot_q4_row_sdot` exactly.
11916#[cfg(target_arch = "aarch64")]
11917#[target_feature(enable = "neon,dotprod")]
11918unsafe fn dot_q4_row_sdot2(
11919    packed: &[u8],
11920    scales: &[u8],
11921    g0: usize,
11922    gpr: usize,
11923    xq1: &[i8],
11924    xq2: &[i8],
11925) -> (f32, f32) {
11926    // SAFETY: callers uphold slice-length contracts (16 packed bytes and
11927    // 2 scale bytes per group; xq*.len() == gpr·GROUP_SIZE).
11928    unsafe {
11929        use core::arch::aarch64::*;
11930        use core::arch::asm;
11931        let lomask = vdupq_n_u8(0x0F);
11932        let eight = vdupq_n_s8(8);
11933        let (mut acc1, mut acc2) = (0f32, 0f32);
11934        for gi in 0..gpr {
11935            let g = g0 + gi;
11936            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
11937            let b = vld1q_u8(packed.as_ptr().add(g * 16));
11938            let lo = vandq_u8(b, lomask);
11939            let hi = vshrq_n_u8::<4>(b);
11940            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
11941            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
11942            let x10 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE));
11943            let x11 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE + 16));
11944            let x20 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE));
11945            let x21 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE + 16));
11946            let (mut a0, mut a1, mut b0, mut b1) = (
11947                vdupq_n_s32(0),
11948                vdupq_n_s32(0),
11949                vdupq_n_s32(0),
11950                vdupq_n_s32(0),
11951            );
11952            asm!(
11953                "sdot {a0:v}.4s, {e0:v}.16b, {x10:v}.16b",
11954                "sdot {a1:v}.4s, {e1:v}.16b, {x11:v}.16b",
11955                "sdot {b0:v}.4s, {e0:v}.16b, {x20:v}.16b",
11956                "sdot {b1:v}.4s, {e1:v}.16b, {x21:v}.16b",
11957                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
11958                b0 = inout(vreg) b0, b1 = inout(vreg) b1,
11959                e0 = in(vreg) e0, e1 = in(vreg) e1,
11960                x10 = in(vreg) x10, x11 = in(vreg) x11,
11961                x20 = in(vreg) x20, x21 = in(vreg) x21,
11962                options(pure, nomem, nostack),
11963            );
11964            acc1 += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
11965            acc2 += vaddvq_s32(vaddq_s32(b0, b1)) as f32 * s;
11966        }
11967        (acc1, acc2)
11968    }
11969}
11970
11971// ───────────────────── fused int8 kernels ─────────────────────
11972
11973/// `acc += w · row` where the row is centered i8 — NEON widen+fma on
11974/// aarch64, scalar elsewhere. The KV-cache q8 value path rides on this.
11975#[inline]
11976pub(crate) fn axpy_i8_f32(acc: &mut [f32], row: &[i8], w: f32) {
11977    #[cfg(target_arch = "aarch64")]
11978    unsafe {
11979        return axpy_i8_f32_neon(acc, row, w);
11980    }
11981    #[cfg(target_arch = "x86_64")]
11982    if avx2_enabled() {
11983        return unsafe { axpy_i8_f32_avx2(acc, row, w) };
11984    }
11985    #[allow(unreachable_code)]
11986    {
11987        for (a, &b) in acc.iter_mut().zip(row) {
11988            *a += w * b as f32;
11989        }
11990    }
11991}
11992
11993/// i8→f32 axpy via AVX2/FMA (x86 mirror of `axpy_i8_f32_neon`).
11994#[cfg(target_arch = "x86_64")]
11995#[target_feature(enable = "avx2,fma")]
11996unsafe fn axpy_i8_f32_avx2(acc: &mut [f32], row: &[i8], w: f32) {
11997    // SAFETY: callers uphold slice-length contracts (see call sites).
11998    unsafe {
11999        use core::arch::x86_64::*;
12000        let n = acc.len().min(row.len());
12001        let ap = acc.as_mut_ptr();
12002        let rp = row.as_ptr();
12003        let wv = _mm256_set1_ps(w);
12004        let mut j = 0usize;
12005        while j + 16 <= n {
12006            let rb = _mm_loadu_si128(rp.add(j) as *const __m128i);
12007            let lo = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(rb));
12008            let hi = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(rb)));
12009            let v0 = _mm256_fmadd_ps(wv, lo, _mm256_loadu_ps(ap.add(j)));
12010            let v1 = _mm256_fmadd_ps(wv, hi, _mm256_loadu_ps(ap.add(j + 8)));
12011            _mm256_storeu_ps(ap.add(j), v0);
12012            _mm256_storeu_ps(ap.add(j + 8), v1);
12013            j += 16;
12014        }
12015        while j < n {
12016            *ap.add(j) += w * (*rp.add(j)) as f32;
12017            j += 1;
12018        }
12019    }
12020}
12021
12022#[cfg(target_arch = "aarch64")]
12023#[target_feature(enable = "neon")]
12024unsafe fn axpy_i8_f32_neon(acc: &mut [f32], row: &[i8], w: f32) {
12025    // SAFETY: callers uphold slice-length contracts (see call sites).
12026    unsafe {
12027        use core::arch::aarch64::*;
12028        let n = acc.len().min(row.len());
12029        let ap = acc.as_mut_ptr();
12030        let rp = row.as_ptr();
12031        let wv = vdupq_n_f32(w);
12032        let mut j = 0usize;
12033        while j + 16 <= n {
12034            let rb = vld1q_s8(rp.add(j));
12035            let lo = vmovl_s8(vget_low_s8(rb));
12036            let hi = vmovl_s8(vget_high_s8(rb));
12037            for (off, half) in [(0, lo), (8, hi)] {
12038                let f0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(half)));
12039                let f1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(half)));
12040                let o = j + off;
12041                vst1q_f32(ap.add(o), vfmaq_f32(vld1q_f32(ap.add(o)), wv, f0));
12042                vst1q_f32(ap.add(o + 4), vfmaq_f32(vld1q_f32(ap.add(o + 4)), wv, f1));
12043            }
12044            j += 16;
12045        }
12046        while j < n {
12047            *ap.add(j) += w * (*rp.add(j)) as f32;
12048            j += 1;
12049        }
12050    }
12051}
12052
12053/// i8 row · f32 x. NEON on aarch64 (ported from vmfcore `dot_i8_f32_neon`,
12054/// ≈9× scalar), scalar elsewhere.
12055#[inline]
12056pub(crate) fn dot_i8_f32(w: &[u8], x: &[f32]) -> f32 {
12057    #[cfg(target_arch = "aarch64")]
12058    unsafe {
12059        return dot_i8_f32_neon(w, x);
12060    }
12061    #[cfg(target_arch = "x86_64")]
12062    if avx2_enabled() {
12063        return unsafe { dot_i8_f32_avx2(w, x) };
12064    }
12065    #[allow(unreachable_code)]
12066    {
12067        let mut sum = 0.0f32;
12068        for (j, &b) in w.iter().enumerate() {
12069            sum += (b as i8) as f32 * x[j];
12070        }
12071        sum
12072    }
12073}
12074
12075/// i8 row · (x ⊙ col_field) — the q8_2f row dot with the θ col-field
12076/// folded into the product (no prescaled copy of x). NEON on aarch64,
12077/// scalar elsewhere. Used by the active-neuron path `row_dot`.
12078#[inline]
12079fn dot_i8_col_f32(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
12080    #[cfg(target_arch = "aarch64")]
12081    unsafe {
12082        return dot_i8_col_f32_neon(w, x, col);
12083    }
12084    #[allow(unreachable_code)]
12085    {
12086        let mut sum = 0.0f32;
12087        for (j, &b) in w.iter().enumerate() {
12088            sum += (b as i8) as f32 * x[j] * col[j];
12089        }
12090        sum
12091    }
12092}
12093
12094#[cfg(target_arch = "aarch64")]
12095#[target_feature(enable = "neon")]
12096unsafe fn dot_i8_col_f32_neon(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
12097    // SAFETY: callers uphold slice-length contracts (see call sites).
12098    unsafe {
12099        use core::arch::aarch64::*;
12100        let n = x.len();
12101        let wp = w.as_ptr() as *const i8;
12102        let xp = x.as_ptr();
12103        let cp = col.as_ptr();
12104        let (mut a0, mut a1, mut a2, mut a3) = (
12105            vdupq_n_f32(0.0),
12106            vdupq_n_f32(0.0),
12107            vdupq_n_f32(0.0),
12108            vdupq_n_f32(0.0),
12109        );
12110        let mut j = 0usize;
12111        while j + 16 <= n {
12112            let wb = vld1q_s8(wp.add(j));
12113            let lo = vmovl_s8(vget_low_s8(wb));
12114            let hi = vmovl_s8(vget_high_s8(wb));
12115            let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
12116            let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
12117            let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
12118            let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
12119            a0 = vfmaq_f32(
12120                a0,
12121                w0,
12122                vmulq_f32(vld1q_f32(xp.add(j)), vld1q_f32(cp.add(j))),
12123            );
12124            a1 = vfmaq_f32(
12125                a1,
12126                w1,
12127                vmulq_f32(vld1q_f32(xp.add(j + 4)), vld1q_f32(cp.add(j + 4))),
12128            );
12129            a2 = vfmaq_f32(
12130                a2,
12131                w2,
12132                vmulq_f32(vld1q_f32(xp.add(j + 8)), vld1q_f32(cp.add(j + 8))),
12133            );
12134            a3 = vfmaq_f32(
12135                a3,
12136                w3,
12137                vmulq_f32(vld1q_f32(xp.add(j + 12)), vld1q_f32(cp.add(j + 12))),
12138            );
12139            j += 16;
12140        }
12141        let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
12142        while j < n {
12143            sum += (*wp.add(j)) as f32 * *xp.add(j) * *cp.add(j);
12144            j += 1;
12145        }
12146        sum
12147    }
12148}
12149
12150#[cfg(target_arch = "aarch64")]
12151#[target_feature(enable = "neon")]
12152unsafe fn dot_i8_f32_neon(w: &[u8], x: &[f32]) -> f32 {
12153    // SAFETY: callers uphold slice-length contracts (see call sites).
12154    unsafe {
12155        use core::arch::aarch64::*;
12156        let n = x.len();
12157        let wp = w.as_ptr() as *const i8;
12158        let xp = x.as_ptr();
12159        let (mut a0, mut a1, mut a2, mut a3) = (
12160            vdupq_n_f32(0.0),
12161            vdupq_n_f32(0.0),
12162            vdupq_n_f32(0.0),
12163            vdupq_n_f32(0.0),
12164        );
12165        let mut j = 0usize;
12166        while j + 16 <= n {
12167            let wb = vld1q_s8(wp.add(j));
12168            let lo = vmovl_s8(vget_low_s8(wb));
12169            let hi = vmovl_s8(vget_high_s8(wb));
12170            let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
12171            let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
12172            let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
12173            let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
12174            a0 = vfmaq_f32(a0, w0, vld1q_f32(xp.add(j)));
12175            a1 = vfmaq_f32(a1, w1, vld1q_f32(xp.add(j + 4)));
12176            a2 = vfmaq_f32(a2, w2, vld1q_f32(xp.add(j + 8)));
12177            a3 = vfmaq_f32(a3, w3, vld1q_f32(xp.add(j + 12)));
12178            j += 16;
12179        }
12180        let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
12181        while j < n {
12182            sum += (*wp.add(j)) as f32 * *xp.add(j);
12183            j += 1;
12184        }
12185        sum
12186    }
12187}
12188
12189#[allow(clippy::too_many_arguments)]
12190fn qmatvec(
12191    q: &[u8],
12192    rep: &[u8],
12193    row_scale: &[f32],
12194    x: &[f32],
12195    col_field: &[f32],
12196    dtype: TensorDtype,
12197    rows: usize,
12198    cols: usize,
12199    out: &mut [f32],
12200    pool: Option<&Pool>,
12201) {
12202    debug_assert_eq!(out.len(), rows);
12203    #[cfg(not(target_arch = "aarch64"))]
12204    let _ = rep;
12205
12206    #[cfg(target_arch = "aarch64")]
12207    if sdot_enabled() {
12208        let act = if dtype == TensorDtype::Q8_2f {
12209            split_act_q8_2f(x, col_field)
12210        } else {
12211            split_act(x)
12212        };
12213        let out_addr = SendMut(out.as_mut_ptr());
12214        let run_range = |start: usize, end: usize| {
12215            q8_range_sdot(q, rep, row_scale, &act, cols, out_addr, start, end)
12216        };
12217        match pool {
12218            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
12219            _ => run_range(0, rows),
12220        }
12221        return;
12222    }
12223    // x86 A8W8 via AVX2 maddubs — same quantized-activation contract as
12224    // the SDOT path (CMF_AVX2=0 keeps the exact i8×f32 loop).
12225    #[cfg(target_arch = "x86_64")]
12226    if avx2_a8w8_enabled() {
12227        let act = if dtype == TensorDtype::Q8_2f {
12228            split_act_q8_2f(x, col_field)
12229        } else {
12230            split_act(x)
12231        };
12232        let out_addr = SendMut(out.as_mut_ptr());
12233        let run_range = |start: usize, end: usize| {
12234            q8_range_avx2(q, row_scale, &act, cols, out_addr, start, end)
12235        };
12236        match pool {
12237            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
12238            _ => run_range(0, rows),
12239        }
12240        return;
12241    }
12242
12243    prescale_with(x, col_field, dtype, 1, |xs| {
12244        let out_addr = SendMut(out.as_mut_ptr());
12245        let run_range = move |start: usize, end: usize| {
12246            for o in start..end {
12247                let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
12248                // SAFETY: disjoint row ranges per worker.
12249                unsafe { *out_addr.at(o) = v };
12250            }
12251        };
12252        match pool {
12253            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
12254            _ => run_range(0, rows),
12255        }
12256    });
12257}
12258
12259#[allow(clippy::too_many_arguments)]
12260fn qmatvec2(
12261    q: &[u8],
12262    row_scale: &[f32],
12263    x1: &[f32],
12264    x2: &[f32],
12265    col_field: &[f32],
12266    dtype: TensorDtype,
12267    rows: usize,
12268    cols: usize,
12269    o1: &mut [f32],
12270    o2: &mut [f32],
12271    pool: Option<&Pool>,
12272) {
12273    #[cfg(target_arch = "aarch64")]
12274    if sdot_enabled() {
12275        let a1s = if dtype == TensorDtype::Q8_2f {
12276            split_act_q8_2f(x1, col_field)
12277        } else {
12278            split_act(x1)
12279        };
12280        let a2s = if dtype == TensorDtype::Q8_2f {
12281            split_act_q8_2f(x2, col_field)
12282        } else {
12283            split_act(x2)
12284        };
12285        let p1 = SendMut(o1.as_mut_ptr());
12286        let p2 = SendMut(o2.as_mut_ptr());
12287        let run_range = |start: usize, end: usize| {
12288            q8_range2_sdot(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
12289        };
12290        match pool {
12291            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
12292            _ => run_range(0, rows),
12293        }
12294        return;
12295    }
12296    #[cfg(target_arch = "x86_64")]
12297    if avx2_a8w8_enabled() {
12298        let a1s = if dtype == TensorDtype::Q8_2f {
12299            split_act_q8_2f(x1, col_field)
12300        } else {
12301            split_act(x1)
12302        };
12303        let a2s = if dtype == TensorDtype::Q8_2f {
12304            split_act_q8_2f(x2, col_field)
12305        } else {
12306            split_act(x2)
12307        };
12308        let p1 = SendMut(o1.as_mut_ptr());
12309        let p2 = SendMut(o2.as_mut_ptr());
12310        let run_range = |start: usize, end: usize| {
12311            q8_range2_avx2(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
12312        };
12313        match pool {
12314            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
12315            _ => run_range(0, rows),
12316        }
12317        return;
12318    }
12319
12320    prescale_with(x1, col_field, dtype, 1, |x1s| {
12321        prescale_with(x2, col_field, dtype, 2, |x2s| {
12322            let p1 = SendMut(o1.as_mut_ptr());
12323            let p2 = SendMut(o2.as_mut_ptr());
12324            let run_range = move |start: usize, end: usize| {
12325                for o in start..end {
12326                    let row = &q[o * cols..(o + 1) * cols];
12327                    let s1 = dot_i8_f32(row, x1s) * row_scale[o];
12328                    let s2 = dot_i8_f32(row, x2s) * row_scale[o];
12329                    // SAFETY: disjoint row ranges per worker.
12330                    unsafe {
12331                        *p1.at(o) = s1;
12332                        *p2.at(o) = s2;
12333                    }
12334                }
12335            };
12336            match pool {
12337                Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
12338                _ => run_range(0, rows),
12339            }
12340        });
12341    });
12342}
12343
12344#[derive(Clone, Copy)]
12345struct SendMut(*mut f32);
12346unsafe impl Send for SendMut {}
12347unsafe impl Sync for SendMut {}
12348
12349impl SendMut {
12350    #[inline]
12351    fn at(self, i: usize) -> *mut f32 {
12352        unsafe { self.0.add(i) }
12353    }
12354}
12355
12356#[cfg(test)]
12357mod tests {
12358    /// `q8_round` must be `round().clamp(±127) as i8` bit for bit: every
12359    /// half-integer, their neighbours one ulp either side, the clamp
12360    /// boundary, huge values, infinities and NaN, plus a dense sweep.
12361    #[test]
12362    fn q8_round_is_round_clamp() {
12363        let reference = |t: f32| t.round().clamp(-127.0, 127.0) as i8;
12364        let mut probe = vec![
12365            0.0f32,
12366            -0.0,
12367            f32::NAN,
12368            f32::INFINITY,
12369            f32::NEG_INFINITY,
12370            f32::MAX,
12371            f32::MIN,
12372            1e30,
12373            -1e30,
12374            f32::MIN_POSITIVE,
12375            -f32::MIN_POSITIVE,
12376        ];
12377        for k in -300i32..=300 {
12378            let h = k as f32 * 0.5;
12379            let up = f32::from_bits(h.to_bits() + 1);
12380            let down = f32::from_bits(h.to_bits().wrapping_sub(1));
12381            for t in [h, up, down] {
12382                probe.push(t);
12383                probe.push(-t);
12384            }
12385        }
12386        let mut t = -140.0f32;
12387        while t < 140.0 {
12388            probe.push(t);
12389            t += 0.000_731;
12390        }
12391        for t in probe {
12392            assert_eq!(q8_round(t), reference(t), "t = {t:e} ({:#x})", t.to_bits());
12393        }
12394    }
12395
12396    use super::*;
12397
12398    /// The f32 `matmat` walks the batch four positions a row at a time; every
12399    /// output must still be the scalar-order sum `matvec` gives, bit for bit
12400    /// (a 7-row batch covers the four-wide block and the remainder).
12401    #[test]
12402    fn f32_matmat_equals_matvec_bitwise() {
12403        let (rows, cols, b) = (5usize, 37usize, 7usize);
12404        let w: Vec<f32> = (0..rows * cols).map(|i| ((i * 7919 % 113) as f32 - 56.0) / 37.0).collect();
12405        let xs: Vec<f32> = (0..b * cols).map(|i| ((i * 104_729 % 97) as f32 - 48.0) / 29.0).collect();
12406        let t = QTensor::from_f32(w, rows, cols);
12407        let mut out = vec![0f32; b * rows];
12408        t.matmat(&xs, b, &mut out, None);
12409        for bi in 0..b {
12410            let mut one = vec![0f32; rows];
12411            t.matvec(&xs[bi * cols..(bi + 1) * cols], &mut one, None);
12412            for o in 0..rows {
12413                assert_eq!(out[bi * rows + o].to_bits(), one[o].to_bits(), "bi {bi} row {o}");
12414            }
12415        }
12416    }
12417
12418    #[test]
12419    fn q2tp_i8_dot_matches_exact_on_grid() {
12420        // On-grid activations (±1 → sx=1/127, xq=±127 dequantizes
12421        // exactly, no outliers) must make the integer path agree with
12422        // the exact scalar walk to f32 rounding.
12423        let (rows, cols) = (5, 64);
12424        let gpr = cols / GROUP_SIZE;
12425        // Synthetic codes plane + a flat ladder: scales_into is not under
12426        // test here, so drive dot_q2tp_row_i8 / q2tp_row_exact directly
12427        // with hand-made scales.
12428        let chunks: Vec<u8> = (0..rows * gpr * Q2TP_CHUNK)
12429            .map(|i| (i as u32).wrapping_mul(2654435761) as u8)
12430            .collect();
12431        let scales: Vec<f32> = (0..gpr).map(|g| 0.5 + g as f32 * 0.25).collect();
12432        let x: Vec<f32> = (0..cols)
12433            .map(|i| if i % 3 == 0 { -1.0 } else { 1.0 })
12434            .collect();
12435        let act = split_act(&x);
12436        assert!(
12437            act.outliers.is_empty(),
12438            "on-grid input must have no outliers"
12439        );
12440        let gsum = q1_group_sums(&act.xq, gpr);
12441        for r in 0..rows {
12442            let exact = q2tp_row_exact(&chunks, r, gpr, &x, &scales);
12443            let fast = dot_q2tp_row_i8(&chunks, r, gpr, &act.xq, &gsum, &scales) * act.sx;
12444            assert!(
12445                (exact - fast).abs() <= exact.abs() * 1e-5 + 1e-5,
12446                "row {r}: exact {exact} vs i8 {fast}"
12447            );
12448        }
12449    }
12450
12451    #[cfg(target_arch = "x86_64")]
12452    #[test]
12453    fn q2tp_avx2_dot_matches_scalar_for_random_patterns() {
12454        // Compare the release AVX2 integer dot against the scalar oracle over
12455        // arbitrary packed bytes/activation signs.  This guards the exact
12456        // table-load path used after rejecting a faster-looking decoder whose
12457        // full-checkpoint greedy output drifted.
12458        if !std::arch::is_x86_feature_detected!("avx2") {
12459            return;
12460        }
12461        let mut seed = 0x9e3779b9u32;
12462        let mut next = || {
12463            seed = seed.wrapping_mul(1664525).wrapping_add(1013904223);
12464            seed
12465        };
12466        for _ in 0..20_000 {
12467            let mut ch = [0u8; Q2TP_CHUNK];
12468            let mut x = [0i8; GROUP_SIZE];
12469            for b in &mut ch {
12470                *b = next() as u8;
12471            }
12472            for v in &mut x {
12473                *v = (next() >> 24) as i8;
12474            }
12475            let mut reference = 0i32;
12476            for (k, &b) in ch.iter().enumerate() {
12477                reference += (b & 3) as i32 * x[k * 4] as i32;
12478                reference += ((b >> 2) & 3) as i32 * x[k * 4 + 1] as i32;
12479                reference += ((b >> 4) & 3) as i32 * x[k * 4 + 2] as i32;
12480                reference += ((b >> 6) & 3) as i32 * x[k * 4 + 3] as i32;
12481            }
12482            // SAFETY: guarded by the runtime AVX2 feature check and fixed
12483            // 8-byte/32-byte slice lengths above.
12484            let got = unsafe { q2tp_code_dot_avx2(&ch, &x) };
12485            assert_eq!(got, reference, "packed q2 lane mismatch");
12486        }
12487    }
12488
12489    #[test]
12490    fn q2tp_affine_fuses_half_scale_correction_without_changing_raw_decode() {
12491        let (rows, cols) = (1usize, GROUP_SIZE);
12492        let mut bytes = vec![0u8; Q2TP_CHUNK + 4 + 1];
12493        // Repeating symbols 0,1,2,0 at unit scale.  q2tp's raw B is
12494        // (c-1.5), while the affine Prism operator is (c-1.0).
12495        bytes[..Q2TP_CHUNK].fill(0x24); // codes 0,1,2,0 in LSB-first order
12496        bytes[Q2TP_CHUNK..Q2TP_CHUNK + 2].copy_from_slice(&0u16.to_le_bytes());
12497        bytes[Q2TP_CHUNK + 2..Q2TP_CHUNK + 4].copy_from_slice(&0u16.to_le_bytes());
12498        bytes[Q2TP_CHUNK + 4] = 1; // dtype16 rung 1 = 1.0
12499        let x = vec![1.0f32; cols];
12500        let mut raw = vec![0.0f32; rows];
12501        let mut affine = vec![0.0f32; rows];
12502        q2tp_matvec_for_test(&bytes, &x, rows, cols, &mut raw);
12503        q2tp_affine_matvec_for_test(&bytes, &x, rows, cols, &mut affine);
12504        assert_eq!(raw, vec![-24.0]);
12505        assert_eq!(affine, vec![-8.0]);
12506        assert!((affine[0] - (raw[0] + 0.5 * cols as f32)).abs() < 1e-6);
12507    }
12508
12509    #[test]
12510    fn q8_row_dot_fast_matches_scalar() {
12511        // The per-arch fast dot must agree with the exact scalar oracle
12512        // (same contract the fused q8 FFN arm rides on).
12513        let cols = 96;
12514        let row: Vec<u8> = (0..cols)
12515            .map(|i| ((i * 37 % 251) - 125) as i8 as u8)
12516            .collect();
12517        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.13).sin()).collect();
12518        let act = split_act(&x);
12519        let fast = q8_row_dot(&row, &act);
12520        let scalar = q8_row_dot_scalar(&row, &act);
12521        assert!(
12522            (fast - scalar).abs() <= scalar.abs() * 1e-5 + 1e-5,
12523            "fast {fast} vs scalar {scalar}"
12524        );
12525    }
12526
12527    #[test]
12528    fn f32_matvec_matches_matvec_rows_bitexact() {
12529        let (rows, cols) = (300, 40);
12530        let w: Vec<f32> = (0..rows * cols).map(|i| (i as f32 * 0.017).sin()).collect();
12531        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.05).cos()).collect();
12532        let qt = QTensor::from_f32(w.clone(), rows, cols);
12533
12534        let mut a = vec![0.0f32; rows];
12535        matvec_rows(None, &w, &x, &mut a);
12536        let mut b = vec![0.0f32; rows];
12537        qt.matvec(&x, &mut b, None);
12538        assert_eq!(a, b);
12539    }
12540
12541    #[test]
12542    fn sdot_kernel_exact_on_grid() {
12543        // Activations already on the i8 grid (±1 with amax=1 → sx=1/127,
12544        // xq=±127 dequantizes EXACTLY) → the SDOT path must match the
12545        // exact f32 dot to float rounding. This isolates kernel
12546        // correctness from quantization noise.
12547        eprintln!("sdot_enabled = {}", sdot_enabled());
12548        let (rows, cols) = (9, 80); // odd rows → exercises 4-row + tail
12549        let w: Vec<u8> = (0..rows * cols)
12550            .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
12551            .collect();
12552        let scales: Vec<f32> = (0..rows).map(|o| 0.005 + o as f32 * 0.001).collect();
12553        let x: Vec<f32> = (0..cols)
12554            .map(|i| match i % 3 {
12555                0 => 1.0,
12556                1 => -1.0,
12557                _ => 0.0,
12558            })
12559            .collect();
12560        let mut a = vec![0.0f32; rows];
12561        qmatvec(
12562            &w,
12563            &[],
12564            &scales,
12565            &x,
12566            &[],
12567            TensorDtype::Q8Row,
12568            rows,
12569            cols,
12570            &mut a,
12571            None,
12572        );
12573        for o in 0..rows {
12574            let mut acc = 0.0f32;
12575            for j in 0..cols {
12576                acc += (w[o * cols + j] as i8) as f32 * x[j];
12577            }
12578            let expect = acc * scales[o];
12579            assert!(
12580                (a[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
12581                "row {o}: {} vs {expect}",
12582                a[o]
12583            );
12584        }
12585    }
12586
12587    #[test]
12588    fn q1_tbl_fast_path_matches_reference() {
12589        // gpr = 8 exercises the TBL pair-load fast loop, and the LAST
12590        // row's final 4-tile window trips the 4B-overread guard (the
12591        // payload ends exactly at the last tile) — both paths must
12592        // agree with the dequant reference.
12593        let (rows, cols) = (5, 256);
12594        let gpr = cols / GROUP_SIZE;
12595        let mut bytes = Vec::new();
12596        for t in 0..rows * gpr {
12597            let s = 0.007 + (t % 11) as f32 * 0.004;
12598            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
12599            for j in 0..4 {
12600                bytes.push(((t * 53 + j * 89 + 7) % 249) as u8);
12601            }
12602        }
12603        let x: Vec<f32> = (0..cols)
12604            .map(|i| if (i * 5) % 7 < 3 { 1.0 } else { -1.0 })
12605            .collect();
12606        let mut w = vec![0.0f32; rows * cols];
12607        cortiq_core::quant::dequant_q1(&bytes, &mut w);
12608        let mut got = vec![0.0f32; rows];
12609        q1_matvec(&bytes, &x, rows, cols, &mut got, None);
12610        for o in 0..rows {
12611            let expect: f32 = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
12612            assert!(
12613                (got[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
12614                "row {o}: {} vs {expect}",
12615                got[o]
12616            );
12617        }
12618        // Blocked 1×4 batch (b=5: one quad + remainder) must equal the
12619        // single-matvec path bit-for-bit.
12620        let b = 5usize;
12621        let mut xs_all = Vec::new();
12622        for bi in 0..b {
12623            xs_all.extend(x.iter().map(|v| if bi % 2 == 0 { *v } else { -*v }));
12624        }
12625        let mut mm = vec![0.0f32; b * rows];
12626        q1_matmat(&bytes, &xs_all, b, rows, cols, &mut mm, None);
12627        for bi in 0..b {
12628            let mut single = vec![0.0f32; rows];
12629            q1_matvec(
12630                &bytes,
12631                &xs_all[bi * cols..(bi + 1) * cols],
12632                rows,
12633                cols,
12634                &mut single,
12635                None,
12636            );
12637            assert_eq!(&mm[bi * rows..(bi + 1) * rows], &single[..], "stream {bi}");
12638        }
12639    }
12640
12641    #[test]
12642    fn q1_kernels_match_exact_reference() {
12643        // Synthetic q1 payload: 6-byte tiles [f16 scale][4B bits].
12644        let (rows, cols) = (7, 96);
12645        let gpr = cols / GROUP_SIZE;
12646        let mut bytes = Vec::new();
12647        for t in 0..rows * gpr {
12648            let s = 0.01 + (t % 13) as f32 * 0.003;
12649            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
12650            for j in 0..4 {
12651                bytes.push(((t * 31 + j * 97) % 251) as u8);
12652            }
12653        }
12654        // On-grid activations (±1, amax 1) → the SDOT path is exact.
12655        let x: Vec<f32> = (0..cols)
12656            .map(|i| if i % 3 == 0 { 1.0 } else { -1.0 })
12657            .collect();
12658        // Reference through the core dequant.
12659        let mut w = vec![0.0f32; rows * cols];
12660        cortiq_core::quant::dequant_q1(&bytes, &mut w);
12661        let mut expect = vec![0.0f32; rows];
12662        for o in 0..rows {
12663            expect[o] = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
12664        }
12665        let mut got = vec![0.0f32; rows];
12666        q1_matvec(&bytes, &x, rows, cols, &mut got, None);
12667        for o in 0..rows {
12668            assert!(
12669                (got[o] - expect[o]).abs() < 1e-3 * expect[o].abs().max(1e-3),
12670                "row {o}: {} vs {}",
12671                got[o],
12672                expect[o]
12673            );
12674        }
12675        // Pair and batch paths agree with the single path.
12676        let x2: Vec<f32> = x.iter().map(|v| -v).collect();
12677        let (mut a1, mut a2) = (vec![0.0f32; rows], vec![0.0f32; rows]);
12678        q1_matvec2(&bytes, &x, &x2, rows, cols, &mut a1, &mut a2, None);
12679        assert_eq!(a1, got);
12680        let mut xs = x.clone();
12681        xs.extend_from_slice(&x2);
12682        let mut mm = vec![0.0f32; 2 * rows];
12683        q1_matmat(&bytes, &xs, 2, rows, cols, &mut mm, None);
12684        assert_eq!(&mm[..rows], got.as_slice());
12685        assert_eq!(&mm[rows..], a2.as_slice());
12686    }
12687
12688    /// A Bonsai-shaped q1 payload (one f16 scale per 128-group, stored in
12689    /// all four tiles; pseudo-random signs) and its q1g repack.
12690    fn q1_q1g_pair(rows: usize, cols: usize, salt: usize) -> (Vec<u8>, Vec<u8>) {
12691        let tiles = rows * cols / GROUP_SIZE;
12692        let mut q1 = Vec::with_capacity(tiles * Q1_TILE);
12693        for t in 0..tiles {
12694            let s = 0.004 + ((t / 4 * 37 + salt) % 29) as f32 * 0.0011;
12695            q1.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
12696            let bits = ((t * 2_654_435_761 + salt * 40_503) as u32).rotate_left((t % 31) as u32);
12697            q1.extend_from_slice(&bits.to_le_bytes());
12698        }
12699        let rp = cortiq_core::quant::repack_q1_to_q1g(&q1, rows, cols, 0.0);
12700        assert_eq!(rp.differing, 0, "uniform scales repack exactly");
12701        (q1, rp.data)
12702    }
12703
12704    /// Activations off the int8 grid with two outlier channels, so the
12705    /// A8W8 path quantizes AND takes the exact outlier correction.
12706    fn q1g_test_x(cols: usize, salt: usize) -> Vec<f32> {
12707        let mut x: Vec<f32> = (0..cols)
12708            .map(|i| ((i * 13 + salt * 7) as f32 * 0.071).sin() * 0.9)
12709            .collect();
12710        x[3 % cols] = 41.0;
12711        x[(cols / 2 + 5) % cols] = -37.5;
12712        x
12713    }
12714
12715    /// Every CPU row kernel of q1g returns the q1 kernel's f32 BIT FOR BIT
12716    /// on a uniform-scale repack: A8W8 rows (SIMD of this machine, and the
12717    /// x86 AVX2/VNNI twins directly), the exact scalar rows (CMF_SDOT=0
12718    /// contract), the outlier term, the drivers (matvec, matvec2, matmat
12719    /// with the blocked 1×4 quad + remainder) — including gpr = 4 and the
12720    /// last row, where q1's pair load falls back to its scalar tail.
12721    #[test]
12722    fn q1g_cpu_kernels_equal_q1_bit_for_bit() {
12723        for &(rows, cols) in &[(3usize, 128usize), (5, 256), (9, 2048), (2, 6144)] {
12724            let gpr = cols / GROUP_SIZE;
12725            let (q1, q1g) = q1_q1g_pair(rows, cols, rows + cols);
12726            let (wa, wb) = (BinW::Q1(&q1), BinW::q1g(&q1g, rows, cols));
12727            let x = q1g_test_x(cols, rows);
12728            let act = split_act(&x);
12729            assert!(!act.outliers.is_empty(), "test wants the outlier path");
12730            let gsum = q1_group_sums(&act.xq, gpr);
12731            for r in 0..rows {
12732                let (a, b) = (wa.row_exact(r, gpr, &x), wb.row_exact(r, gpr, &x));
12733                assert_eq!(a.to_bits(), b.to_bits(), "{rows}x{cols} exact row {r}");
12734                let (a, b) = (
12735                    wa.dot_i8(r, gpr, &act.xq, &gsum),
12736                    wb.dot_i8(r, gpr, &act.xq, &gsum),
12737                );
12738                assert_eq!(a.to_bits(), b.to_bits(), "{rows}x{cols} i8 row {r}");
12739                for j in [0, 31, 32, 127, cols - 1] {
12740                    let (a, b) = (wa.outlier(r, gpr, j), wb.outlier(r, gpr, j));
12741                    assert_eq!(
12742                        (a.0.to_bits(), a.1.to_bits()),
12743                        (b.0.to_bits(), b.1.to_bits())
12744                    );
12745                }
12746                #[cfg(target_arch = "x86_64")]
12747                if avx2_enabled() {
12748                    let sc = cortiq_core::quant::q1g_scale_offset(rows, cols);
12749                    // SAFETY: AVX2 checked; payloads sized for (rows, cols).
12750                    unsafe {
12751                        let a = dot_q1_row_avx2(&q1, r, gpr, &act.xq, &gsum);
12752                        let b = dot_q1g_row_avx2(&q1g, sc, r, gpr, &act.xq, &gsum);
12753                        assert_eq!(a.to_bits(), b.to_bits(), "avx2 row {r}");
12754                    }
12755                    if std::is_x86_feature_detected!("avx512vnni")
12756                        && std::is_x86_feature_detected!("avx512bw")
12757                        && std::is_x86_feature_detected!("avx512vl")
12758                    {
12759                        // SAFETY: VNNI family detected just above.
12760                        unsafe {
12761                            let a = dot_q1_row_vnni(&q1, r, gpr, &act.xq, &gsum);
12762                            let b = dot_q1g_row_vnni(&q1g, sc, r, gpr, &act.xq, &gsum);
12763                            assert_eq!(a.to_bits(), b.to_bits(), "vnni row {r}");
12764                        }
12765                    }
12766                }
12767            }
12768            // Drivers: single, pair, batch (b = 5: one blocked quad + one).
12769            let (mut ga, mut gb) = (vec![0f32; rows], vec![0f32; rows]);
12770            bin_matvec(wa, &x, rows, cols, &mut ga, None);
12771            bin_matvec(wb, &x, rows, cols, &mut gb, None);
12772            assert_eq!(ga, gb, "{rows}x{cols} matvec");
12773            let x2 = q1g_test_x(cols, rows + 1);
12774            let (mut a1, mut a2, mut b1, mut b2) = (
12775                vec![0f32; rows],
12776                vec![0f32; rows],
12777                vec![0f32; rows],
12778                vec![0f32; rows],
12779            );
12780            bin_matvec2(wa, &x, &x2, rows, cols, &mut a1, &mut a2, None);
12781            bin_matvec2(wb, &x, &x2, rows, cols, &mut b1, &mut b2, None);
12782            assert_eq!((&a1, &a2), (&b1, &b2), "{rows}x{cols} matvec2");
12783            let b = 5usize;
12784            let xs: Vec<f32> = (0..b).flat_map(|k| q1g_test_x(cols, k * 3 + 2)).collect();
12785            let (mut ma, mut mb) = (vec![0f32; b * rows], vec![0f32; b * rows]);
12786            bin_matmat(wa, &xs, b, rows, cols, &mut ma, None);
12787            bin_matmat(wb, &xs, b, rows, cols, &mut mb, None);
12788            assert_eq!(ma, mb, "{rows}x{cols} matmat");
12789            // The blocked 1×4 kernels directly (the driver only takes them
12790            // under CMF_X86_BLOCKED / SDOT, a process-wide switch).
12791            let acts: Vec<(SplitAct, Vec<i32>)> = (0..4)
12792                .map(|k| {
12793                    let a = split_act(&xs[k * cols..(k + 1) * cols]);
12794                    let g = q1_group_sums(&a.xq, gpr);
12795                    (a, g)
12796                })
12797                .collect();
12798            let xq4: [&[i8]; 4] = std::array::from_fn(|k| acts[k].0.xq.as_slice());
12799            let gs4: [&[i32]; 4] = std::array::from_fn(|k| acts[k].1.as_slice());
12800            for r in 0..rows {
12801                #[cfg(target_arch = "x86_64")]
12802                if avx2_enabled() {
12803                    // SAFETY: AVX2 checked; payloads sized for (rows, cols).
12804                    let (a, b) = unsafe {
12805                        (
12806                            wa.dot_1x4_x86(r, gpr, xq4, gs4),
12807                            wb.dot_1x4_x86(r, gpr, xq4, gs4),
12808                        )
12809                    };
12810                    assert_eq!(a.map(f32::to_bits), b.map(f32::to_bits), "1x4 row {r}");
12811                }
12812                #[cfg(target_arch = "aarch64")]
12813                if sdot_enabled() {
12814                    // SAFETY: SDOT checked; payloads sized for (rows, cols).
12815                    let (a, b) = unsafe {
12816                        (
12817                            wa.dot_1x4_sdot(r, gpr, xq4, gs4),
12818                            wb.dot_1x4_sdot(r, gpr, xq4, gs4),
12819                        )
12820                    };
12821                    assert_eq!(a.map(f32::to_bits), b.map(f32::to_bits), "1x4 row {r}");
12822                }
12823                let _ = (r, &xq4, &gs4);
12824            }
12825        }
12826    }
12827
12828    /// CPU kernel A/B, q1 vs q1g matvec (A8W8 path of this machine, the
12829    /// exact path too), Bonsai-1.7B decode shapes, arms alternated every
12830    /// round, medians printed. `CMF_THREADS` sets the pool (default 8).
12831    /// `cargo test -p cortiq-engine --release --lib q1g_cpu_matvec_speed --
12832    /// --ignored --nocapture`
12833    #[test]
12834    #[ignore]
12835    fn q1g_cpu_matvec_speed() {
12836        use crate::pool::Pool;
12837        let threads = std::env::var("CMF_THREADS")
12838            .ok()
12839            .and_then(|v| v.parse().ok())
12840            .unwrap_or(8usize);
12841        let pool = Pool::new(threads);
12842        for &(name, rows, cols, reps) in &[
12843            ("gate/up_proj", 6144usize, 2048usize, 41usize),
12844            ("down_proj", 2048, 6144, 41),
12845            ("lm_head", 151_669, 2048, 9),
12846        ] {
12847            let (q1, q1g) = q1_q1g_pair(rows, cols, 3);
12848            let (wa, wb) = (BinW::Q1(&q1), BinW::q1g(&q1g, rows, cols));
12849            let x = q1g_test_x(cols, 1);
12850            let mut out = vec![0f32; rows];
12851            let mut time = |w: BinW| {
12852                let t0 = std::time::Instant::now();
12853                bin_matvec(w, &x, rows, cols, &mut out, Some(&pool));
12854                t0.elapsed().as_secs_f64()
12855            };
12856            let _ = (time(wa), time(wb));
12857            let (mut ta, mut tb) = (Vec::new(), Vec::new());
12858            for r in 0..reps {
12859                if r % 2 == 0 {
12860                    ta.push(time(wa));
12861                    tb.push(time(wb));
12862                } else {
12863                    tb.push(time(wb));
12864                    ta.push(time(wa));
12865                }
12866            }
12867            ta.sort_by(|a, b| a.partial_cmp(b).unwrap());
12868            tb.sort_by(|a, b| a.partial_cmp(b).unwrap());
12869            let (ma, mb) = (ta[reps / 2], tb[reps / 2]);
12870            eprintln!(
12871                "cpu x{threads} a8w8={} {name:>12} {rows}x{cols}: q1 {:8.1} us | q1g {:8.1} us | q1g/q1 time {:.3}",
12872                a8w8_enabled(),
12873                ma * 1e6,
12874                mb * 1e6,
12875                mb / ma
12876            );
12877        }
12878    }
12879
12880    /// The mapped QTensor paths a real Bonsai q1g file takes — matvec, the
12881    /// multi-projection job (all-q1g and q1/q1g mixed), the fused gate/up
12882    /// with SiLU, matmat, matvec2 and the embedding row — equal the q1
12883    /// tensor they were repacked from, bit for bit.
12884    #[test]
12885    fn q1g_mapped_tensor_paths_equal_q1() {
12886        use crate::pool::Pool;
12887        use cortiq_core::{CmfHeader, CmfModel, TensorSpec};
12888        let (rows, cols) = (300usize, 256usize);
12889        let (gq1, gq1g) = q1_q1g_pair(rows, cols, 1);
12890        let (uq1, uq1g) = q1_q1g_pair(rows, cols, 2);
12891        let header: CmfHeader = serde_json::from_value(serde_json::json!({
12892            "format": "cmf", "version": cortiq_core::CMF_VERSION,
12893            "arch": {"arch_name": "tiny-q1g", "hidden_size": cols,
12894                "intermediate_size": rows, "num_layers": 1, "num_attention_heads": 2,
12895                "num_kv_heads": 1, "head_dim": 128, "vocab_size": rows,
12896                "layer_types": ["FullAttention"], "rms_norm_eps": 1e-6,
12897                "max_position_embeddings": 8},
12898            "quant_type": "VBIT"
12899        }))
12900        .unwrap();
12901        let spec = |name: &str, dtype, data: &Vec<u8>| TensorSpec {
12902            name: name.into(),
12903            dtype,
12904            shape: vec![rows, cols],
12905            data: data.clone(),
12906        };
12907        let specs = [
12908            spec("g1", TensorDtype::Q1, &gq1),
12909            spec("u1", TensorDtype::Q1, &uq1),
12910            spec("gg", TensorDtype::Q1G, &gq1g),
12911            spec("ug", TensorDtype::Q1G, &uq1g),
12912        ];
12913        let dir = std::env::temp_dir().join(format!("cmf-q1g-mapped-{}", std::process::id()));
12914        std::fs::create_dir_all(&dir).unwrap();
12915        let path = dir.join("m.cmf");
12916        CmfModel::write(&path, &header, &specs, None, None).unwrap();
12917        let model = Arc::new(CmfModel::open(&path).unwrap());
12918        assert!(model.verify().is_empty(), "{:?}", model.verify());
12919        let t = |n: &str| QTensor::from_model(&model, n).unwrap();
12920        let (g1, u1, gg, ug) = (t("g1"), t("u1"), t("gg"), t("ug"));
12921        assert_eq!(gg.model_dtype(), Some(TensorDtype::Q1G));
12922        assert!(gg.is_q1() && gg.graph_weight().map(|w| w.2) == Some(10));
12923        let pool = Pool::new(3);
12924        let x = q1g_test_x(cols, 4);
12925        let run = |a: &QTensor| {
12926            let mut o = vec![0f32; rows];
12927            a.matvec(&x, &mut o, Some(&pool));
12928            o
12929        };
12930        assert_eq!(run(&g1), run(&gg), "matvec");
12931        let many = |a: &QTensor, b: &QTensor| {
12932            let (mut oa, mut ob) = (vec![0f32; rows], vec![0f32; rows]);
12933            QTensor::matvec_many([a, b], &x, [&mut oa, &mut ob], Some(&pool));
12934            (oa, ob)
12935        };
12936        let want = many(&g1, &u1);
12937        assert_eq!(want, many(&gg, &ug), "matvec_many q1g");
12938        assert_eq!(want, many(&g1, &ug), "matvec_many mixed");
12939        assert_eq!(want, (run(&g1), run(&u1)), "matvec_many == singles");
12940        let fused = |g: &QTensor, u: &QTensor| {
12941            let mut o = vec![0f32; rows];
12942            let ok = QTensor::matvec_silu_mul(g, u, &x, &mut o, Some(&pool));
12943            (ok, o)
12944        };
12945        let (ok_a, fa) = fused(&g1, &u1);
12946        let (ok_b, fb) = fused(&gg, &ug);
12947        assert_eq!(ok_a, ok_b, "q1g takes the fused arm whenever q1 does");
12948        assert_eq!(fa, fb, "fused gate/up");
12949        let b = 6usize;
12950        let xs: Vec<f32> = (0..b).flat_map(|k| q1g_test_x(cols, k + 9)).collect();
12951        let mm = |a: &QTensor| {
12952            let mut o = vec![0f32; b * rows];
12953            a.matmat(&xs, b, &mut o, Some(&pool));
12954            o
12955        };
12956        assert_eq!(mm(&g1), mm(&gg), "matmat");
12957        let x2 = q1g_test_x(cols, 77);
12958        let pair = |a: &QTensor| {
12959            let (mut o1, mut o2) = (vec![0f32; rows], vec![0f32; rows]);
12960            a.matvec2(&x, &x2, &mut o1, &mut o2, Some(&pool));
12961            (o1, o2)
12962        };
12963        assert_eq!(pair(&g1), pair(&gg), "matvec2");
12964        for r in [0, 1, rows / 2, rows - 1] {
12965            let (mut ra, mut rb) = (vec![0f32; cols], vec![0f32; cols]);
12966            g1.row_f32(r, &mut ra);
12967            gg.row_f32(r, &mut rb);
12968            assert_eq!(ra, rb, "embedding row {r}");
12969        }
12970        drop((g1, u1, gg, ug, model));
12971        let _ = std::fs::remove_dir_all(&dir);
12972    }
12973
12974    /// A `blr` tensor read from a real CMF file: it maps to the factored
12975    /// arm (not a dense copy), declines every device-graph accessor, and
12976    /// every QTensor entry point a layer may take — matvec, matvec2,
12977    /// matmat (table loop and GEMM sizes), the multi-tensor job mixed with
12978    /// a q8_row sibling, the fused gate/up refusal, rows and columns — equals
12979    /// the dense reference built by the core dequant.
12980    #[test]
12981    fn blr_mapped_tensor_paths_equal_dense() {
12982        use crate::pool::Pool;
12983        use cortiq_core::{CmfHeader, CmfModel, TensorSpec};
12984        let (rows, cols, r) = (96usize, 128usize, 16usize);
12985        let mut rng = crate::blr::tests::Rng(41);
12986        let gate = crate::blr::tests::make(&mut rng, rows, cols, r, 2);
12987        let up = crate::blr::tests::make(&mut rng, rows, cols, r, 1);
12988        // q8_row sibling: [int8 rows·cols][f16 row scale · rows].
12989        let mut q8: Vec<u8> = (0..rows * cols)
12990            .map(|i| (((i * 89) % 253) as i32 - 126) as i8 as u8)
12991            .collect();
12992        for o in 0..rows {
12993            q8.extend_from_slice(&cortiq_core::quant::f32_to_f16(0.002 + o as f32 * 1e-4).to_le_bytes());
12994        }
12995        let header: CmfHeader = serde_json::from_value(serde_json::json!({
12996            "format": "cmf", "version": cortiq_core::CMF_VERSION,
12997            "arch": {"arch_name": "tiny-blr", "hidden_size": cols,
12998                "intermediate_size": rows, "num_layers": 1, "num_attention_heads": 2,
12999                "num_kv_heads": 1, "head_dim": 64, "vocab_size": rows,
13000                "layer_types": ["FullAttention"], "rms_norm_eps": 1e-6,
13001                "max_position_embeddings": 8},
13002            "quant_type": "Q8_2F"
13003        }))
13004        .unwrap();
13005        let specs = [
13006            TensorSpec {
13007                name: "g".into(),
13008                dtype: TensorDtype::BinaryLowRank,
13009                shape: vec![rows, cols],
13010                data: gate.clone(),
13011            },
13012            TensorSpec {
13013                name: "u".into(),
13014                dtype: TensorDtype::BinaryLowRank,
13015                shape: vec![rows, cols],
13016                data: up.clone(),
13017            },
13018            TensorSpec {
13019                name: "q".into(),
13020                dtype: TensorDtype::Q8Row,
13021                shape: vec![rows, cols],
13022                data: q8,
13023            },
13024        ];
13025        let dir = std::env::temp_dir().join(format!("cmf-blr-mapped-{}", std::process::id()));
13026        std::fs::create_dir_all(&dir).unwrap();
13027        let path = dir.join("m.cmf");
13028        CmfModel::write(&path, &header, &specs, None, None).unwrap();
13029        let model = Arc::new(CmfModel::open(&path).unwrap());
13030        assert!(model.verify().is_empty(), "{:?}", model.verify());
13031        let t = |n: &str| QTensor::from_model(&model, n).unwrap();
13032        let (g, u, q) = (t("g"), t("u"), t("q"));
13033        assert!(g.is_blr() && u.is_blr() && !q.is_blr());
13034        assert!(g.graph_weight().is_none() && g.graph_weight_descriptor().is_none());
13035        assert!(g.as_f32().is_none() && g.model_idx().is_none());
13036        assert!(g.mapped_device_gemm().is_none() && g.q1_parts().is_none());
13037        assert!(!g.sparse_col_ok());
13038        assert_eq!((g.rows(), g.cols()), (rows, cols));
13039        let deq = |bytes: &[u8]| {
13040            let mut w = vec![0f32; rows * cols];
13041            cortiq_core::blr::dequant_blr(bytes, rows, cols, &mut w).unwrap();
13042            w
13043        };
13044        let (wg, wu) = (deq(&gate), deq(&up));
13045        let mv = |w: &[f32], x: &[f32]| -> Vec<f32> {
13046            (0..rows)
13047                .map(|i| (0..cols).map(|j| w[i * cols + j] * x[j]).sum())
13048                .collect()
13049        };
13050        let close = |a: &[f32], b: &[f32], what: &str| {
13051            let s = b.iter().map(|v| v * v).sum::<f32>().sqrt().max(1e-30);
13052            let e = a.iter().zip(b).map(|(x, y)| (x - y) * (x - y)).sum::<f32>().sqrt();
13053            assert!(e / s < 1e-5, "{what}: rel {}", e / s);
13054        };
13055        let pool = Pool::new(3);
13056        let x: Vec<f32> = (0..cols).map(|_| rng.f()).collect();
13057        let x2: Vec<f32> = (0..cols).map(|_| rng.f()).collect();
13058        let mut o = vec![0f32; rows];
13059        g.matvec(&x, &mut o, Some(&pool));
13060        close(&o, &mv(&wg, &x), "matvec");
13061        let (mut o1, mut o2) = (vec![0f32; rows], vec![0f32; rows]);
13062        g.matvec2(&x, &x2, &mut o1, &mut o2, Some(&pool));
13063        close(&o1, &mv(&wg, &x), "matvec2 a");
13064        close(&o2, &mv(&wg, &x2), "matvec2 b");
13065        for b in [2usize, 5] {
13066            let xs: Vec<f32> = (0..b * cols).map(|_| rng.f()).collect();
13067            let mut om = vec![0f32; b * rows];
13068            g.matmat(&xs, b, &mut om, Some(&pool));
13069            for t in 0..b {
13070                close(
13071                    &om[t * rows..(t + 1) * rows],
13072                    &mv(&wg, &xs[t * cols..(t + 1) * cols]),
13073                    "matmat",
13074                );
13075            }
13076        }
13077        // Multi-tensor job mixed with a q8_row sibling: per-tensor fallback,
13078        // the q8_row output identical to its own matvec.
13079        let (mut a, mut bq, mut c) = (vec![0f32; rows], vec![0f32; rows], vec![0f32; rows]);
13080        QTensor::matvec_many([&g, &q, &u], &x, [&mut a, &mut bq, &mut c], Some(&pool));
13081        close(&a, &mv(&wg, &x), "matvec_many blr gate");
13082        close(&c, &mv(&wu, &x), "matvec_many blr up");
13083        let mut qo = vec![0f32; rows];
13084        q.matvec(&x, &mut qo, Some(&pool));
13085        assert_eq!(bq, qo, "matvec_many q8_row sibling");
13086        let mut f = vec![0f32; rows];
13087        assert!(!QTensor::matvec_silu_mul(&g, &u, &x, &mut f, Some(&pool)));
13088        let mut row = vec![0f32; cols];
13089        g.row_f32(rows / 3, &mut row);
13090        close(&row, &wg[(rows / 3) * cols..(rows / 3 + 1) * cols], "row_f32");
13091        let mut scratch = vec![0f32; cols];
13092        let d = g.row_dot(7, &x, &mut scratch);
13093        let want: f32 = (0..cols).map(|j| wg[7 * cols + j] * x[j]).sum();
13094        assert!((d - want).abs() < 1e-4 * (1.0 + want.abs()), "row_dot");
13095        let mut col = vec![0f32; rows];
13096        g.add_col_scaled(cols / 2, 1.5, &mut col);
13097        let want: Vec<f32> = (0..rows).map(|i| 1.5 * wg[i * cols + cols / 2]).collect();
13098        close(&col, &want, "add_col_scaled");
13099        drop((g, u, q, model));
13100        let _ = std::fs::remove_dir_all(&dir);
13101    }
13102
13103    #[test]
13104    fn repack_is_bit_identical() {
13105        // The interleaved-repack kernel must produce EXACTLY the same
13106        // bits as the mmap-layout kernel: integer accumulation is order-
13107        // exact, the f32 epilogue is identical. Odd rows exercise the
13108        // tail; direct range calls exercise unaligned pool splits.
13109        let (rows, cols) = (267, 96); // 66 groups + 3 tail rows, cols % 16 == 0
13110        let w: Vec<u8> = (0..rows * cols)
13111            .map(|i| (((i * 89) % 253) as i32 - 126) as i8 as u8)
13112            .collect();
13113        let scales: Vec<f32> = (0..rows).map(|o| 0.003 + o as f32 * 0.0007).collect();
13114        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.37).sin() * 2.0).collect();
13115        let rep = q8_repack_layout(&w, rows, cols);
13116        // Group interleave round-trips.
13117        for g in 0..rows / 4 {
13118            for c in 0..cols / 16 {
13119                for lane in 0..4 {
13120                    assert_eq!(
13121                        &rep[g * 4 * cols + c * 64 + lane * 16
13122                            ..g * 4 * cols + c * 64 + lane * 16 + 16],
13123                        &w[(g * 4 + lane) * cols + c * 16..(g * 4 + lane) * cols + c * 16 + 16],
13124                    );
13125                }
13126            }
13127        }
13128        let mut a = vec![0.0f32; rows];
13129        qmatvec(
13130            &w,
13131            &[],
13132            &scales,
13133            &x,
13134            &[],
13135            TensorDtype::Q8Row,
13136            rows,
13137            cols,
13138            &mut a,
13139            None,
13140        );
13141        let mut b = vec![0.0f32; rows];
13142        qmatvec(
13143            &w,
13144            &rep,
13145            &scales,
13146            &x,
13147            &[],
13148            TensorDtype::Q8Row,
13149            rows,
13150            cols,
13151            &mut b,
13152            None,
13153        );
13154        assert_eq!(a, b, "full-range repack output diverged");
13155
13156        #[cfg(target_arch = "aarch64")]
13157        if sdot_enabled() {
13158            // Unaligned range split (pool workers get arbitrary bounds).
13159            let act = split_act(&x);
13160            let mut c1 = vec![0.0f32; rows];
13161            let mut c2 = vec![0.0f32; rows];
13162            q8_range_sdot(
13163                &w,
13164                &[],
13165                &scales,
13166                &act,
13167                cols,
13168                SendMut(c1.as_mut_ptr()),
13169                3,
13170                rows - 2,
13171            );
13172            q8_range_sdot(
13173                &w,
13174                &rep,
13175                &scales,
13176                &act,
13177                cols,
13178                SendMut(c2.as_mut_ptr()),
13179                3,
13180                rows - 2,
13181            );
13182            assert_eq!(c1, c2, "unaligned-range repack output diverged");
13183        }
13184    }
13185
13186    #[test]
13187    fn sdot_a8w8_noise_is_bounded() {
13188        // Off-grid activations: A8 quantization noise must stay small in
13189        // relative L2 over the whole output (realistic accuracy contract;
13190        // vmfcore measured argmax-identical decode on real models).
13191        let (rows, cols) = (16, 512);
13192        let w: Vec<u8> = (0..rows * cols)
13193            .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
13194            .collect();
13195        let scales = vec![0.01f32; rows];
13196        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.21).sin()).collect();
13197        let mut a = vec![0.0f32; rows];
13198        qmatvec(
13199            &w,
13200            &[],
13201            &scales,
13202            &x,
13203            &[],
13204            TensorDtype::Q8Row,
13205            rows,
13206            cols,
13207            &mut a,
13208            None,
13209        );
13210        let (mut num, mut den) = (0f64, 0f64);
13211        for o in 0..rows {
13212            let mut acc = 0.0f32;
13213            for j in 0..cols {
13214                acc += (w[o * cols + j] as i8) as f32 * x[j];
13215            }
13216            let expect = acc * scales[o];
13217            num += ((a[o] - expect) as f64).powi(2);
13218            den += (expect as f64).powi(2);
13219        }
13220        let rel = (num / den.max(1e-12)).sqrt();
13221        assert!(rel < 0.05, "A8W8 relative L2 error too high: {rel}");
13222    }
13223
13224    #[test]
13225    fn i8_dot_neon_matches_scalar() {
13226        let n = 100;
13227        let w: Vec<u8> = (0..n).map(|i| ((i * 37 + 11) % 251) as u8).collect();
13228        let x: Vec<f32> = (0..n).map(|i| (i as f32 * 0.13).sin()).collect();
13229        let mut scalar = 0.0f32;
13230        for j in 0..n {
13231            scalar += (w[j] as i8) as f32 * x[j];
13232        }
13233        let fast = dot_i8_f32(&w, &x);
13234        assert!((scalar - fast).abs() < 1e-3 * scalar.abs().max(1.0));
13235    }
13236
13237    /// Fused vbit matvec must match full dequant_vbit + dense matvec.
13238    #[test]
13239    fn vbitmatvec_matches_full_dequant() {
13240        let (rows, cols) = (6, 64);
13241        let ng = cols / GROUP_SIZE;
13242        // Hand-craft: bits per row, f16 scales, packed rows.
13243        let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
13244        let mut bytes = bits.clone();
13245        for g in 0..rows * ng {
13246            let s = 0.02 + 0.001 * g as f32;
13247            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
13248        }
13249        for r in 0..rows {
13250            let b = bits[r] as usize;
13251            let (mut acc, mut nb) = (0u64, 0usize);
13252            let mut rowbytes = Vec::new();
13253            for i in 0..cols {
13254                let v = ((i * 7 + r * 13) % (1 << b)) as u64;
13255                acc = (acc << b) | v;
13256                nb += b;
13257                while nb >= 8 {
13258                    nb -= 8;
13259                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
13260                }
13261            }
13262            if nb > 0 {
13263                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
13264            }
13265            bytes.extend_from_slice(&rowbytes);
13266        }
13267        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
13268
13269        let mut reference = vec![0f32; rows * cols];
13270        cortiq_core::quant::dequant_vbit(&bytes, rows, cols, &mut reference).unwrap();
13271        let mut expect = vec![0f32; rows];
13272        for r in 0..rows {
13273            expect[r] = reference[r * cols..(r + 1) * cols]
13274                .iter()
13275                .zip(&x)
13276                .map(|(w, xv)| w * xv)
13277                .sum();
13278        }
13279        let mut got = vec![0f32; rows];
13280        let offsets = vbit_row_offsets(&bytes, rows, cols);
13281        vbitmatvec(&bytes, &offsets, &x, rows, cols, &mut got, None);
13282        // SDOT path quantizes activations to i8 (A8W8): bounded noise,
13283        // same contract as q8 (exact path is pinned by CMF_SDOT=0 in
13284        // the golden-parity gate).
13285        let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
13286        let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
13287        for r in 0..rows {
13288            assert!(
13289                (got[r] - expect[r]).abs() < tol * scale,
13290                "row {r}: {} vs {}",
13291                got[r],
13292                expect[r]
13293            );
13294        }
13295    }
13296
13297    /// Fused q4 matvec must match the reference full-dequant + dense
13298    /// matvec bit-for-bit in structure (same f32 math, group order).
13299    /// vbit matmat: the blocked 1×4 leg must match the per-row path
13300    /// (paired env toggle; larger shape so both code paths engage).
13301    #[test]
13302    #[cfg(target_arch = "x86_64")]
13303    fn vbit_matmat_blocked_matches_per_row() {
13304        let (rows, cols, b) = (64usize, 128usize, 9usize);
13305        let ng = cols / GROUP_SIZE;
13306        let bits: Vec<u8> = (0..rows).map(|r| [3u8, 4, 5, 6][r % 4]).collect();
13307        let mut bytes = bits.clone();
13308        for g in 0..rows * ng {
13309            let sc = 0.02 + 0.0005 * g as f32;
13310            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
13311        }
13312        for r in 0..rows {
13313            let bw = bits[r] as usize;
13314            let (mut acc, mut nb) = (0u64, 0usize);
13315            let mut rowbytes = Vec::new();
13316            for i in 0..cols {
13317                let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
13318                acc = (acc << bw) | v;
13319                nb += bw;
13320                while nb >= 8 {
13321                    nb -= 8;
13322                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
13323                }
13324            }
13325            if nb > 0 {
13326                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
13327            }
13328            bytes.extend_from_slice(&rowbytes);
13329        }
13330        let x: Vec<f32> = (0..b * cols)
13331            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
13332            .collect();
13333        let offsets = vbit_row_offsets(&bytes, rows, cols);
13334        let mut y_a = vec![0f32; b * rows];
13335        let mut y_b = vec![0f32; b * rows];
13336        unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
13337        vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_a, None);
13338        unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
13339        vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_b, None);
13340        unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
13341        let max_d = y_a
13342            .iter()
13343            .zip(&y_b)
13344            .map(|(p, q)| (p - q).abs())
13345            .fold(0.0f32, f32::max);
13346        assert!(max_d < 1e-4, "vbit blocked ≠ per-row: max|Δ| = {max_d}");
13347    }
13348
13349    /// q4t blocked 1×4 (SDOT on ARM, AVX2 on x86) must equal the
13350    /// per-row path exactly: same nibble unpack, same group order,
13351    /// same f32 accumulation — batch == matvec bit-for-bit. b=9 covers
13352    /// two full 1×4 blocks plus a remainder through the single-row
13353    /// kernel. (Both paths produce identical output, so the shared
13354    /// CMF_X86_BLOCKED env var racing with other tests cannot flip
13355    /// the verdict — worst case both sides take the same path.)
13356    #[test]
13357    fn q4t_matmat_blocked_matches_per_row() {
13358        let (rows, cols, b) = (16usize, 64usize, 9usize);
13359        let gpr = cols / GROUP_SIZE;
13360        let mut bytes = vec![0u8; rows * gpr * Q4_TILE];
13361        for r in 0..rows {
13362            for g in 0..gpr {
13363                let t = (r * gpr + g) * Q4_TILE;
13364                let sc = 0.02 + 0.001 * (r * gpr + g) as f32;
13365                bytes[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
13366                for k in 0..16 {
13367                    bytes[t + 2 + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
13368                }
13369            }
13370        }
13371        let x: Vec<f32> = (0..b * cols)
13372            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
13373            .collect();
13374        let mut y_blk = vec![0f32; b * rows];
13375        let mut y_row = vec![0f32; b * rows];
13376        unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
13377        q4t_matmat(&bytes, &x, b, rows, cols, &mut y_blk, None);
13378        unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
13379        q4t_matmat(&bytes, &x, b, rows, cols, &mut y_row, None);
13380        unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
13381        assert_eq!(y_blk, y_row, "q4t blocked 1x4 ≠ per-row");
13382    }
13383
13384    /// The wide-batch Accelerate arm of q4t_matmat vs a brute-force
13385    /// f32 dequant matmul: both are f32 GEMMs, so only reduction
13386    /// order differs — tight tolerance.
13387    /// A synthetic q4tp payload: random nibbles plus a per-row ladder whose
13388    /// span varies row to row, so the codes actually exercise the full 0..31
13389    /// range rather than clustering on one rung.
13390    fn synth_q4tp(rows: usize, cols: usize) -> Vec<u8> {
13391        use cortiq_core::quant::{f32_to_f16, q4tp_code_stride, q4tp_put_code};
13392        let gpr = cols / GROUP_SIZE;
13393        let stride = q4tp_code_stride(gpr);
13394        let (params_off, codes_off, _) = q4tp_sections(rows, cols);
13395        let mut b = vec![0u8; codes_off + rows * stride];
13396        for r in 0..rows {
13397            for g in 0..gpr {
13398                let t = (r * gpr + g) * Q4TP_NIB;
13399                for k in 0..16 {
13400                    b[t + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
13401                }
13402            }
13403            let lo = -6.0 - 0.03 * (r % 17) as f32;
13404            let step = 0.01 + 0.004 * (r % 11) as f32;
13405            let p = params_off + r * 4;
13406            b[p..p + 2].copy_from_slice(&f32_to_f16(lo).to_le_bytes());
13407            b[p + 2..p + 4].copy_from_slice(&f32_to_f16(step).to_le_bytes());
13408            let crow = &mut b[codes_off + r * stride..codes_off + (r + 1) * stride];
13409            for g in 0..gpr {
13410                q4tp_put_code(crow, g, (r * 5 + g * 3) % 32);
13411            }
13412        }
13413        b
13414    }
13415
13416    /// The same weights re-expressed as q4_tiled, so the proven kernel can
13417    /// be the reference: each tile stores the ladder scale its code selects.
13418    /// Only the f16 rounding of that scale separates the two payloads.
13419    fn q4tp_as_q4t(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
13420        let gpr = cols / GROUP_SIZE;
13421        let v = Q4tpView::new(bytes, rows, cols);
13422        let mut out = vec![0u8; rows * gpr * Q4_TILE];
13423        let mut sc = vec![0f32; gpr];
13424        for r in 0..rows {
13425            v.scales_into(r, gpr, &mut sc);
13426            for g in 0..gpr {
13427                let t = (r * gpr + g) * Q4_TILE;
13428                let s = sc[g];
13429                out[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
13430                let src = (r * gpr + g) * Q4TP_NIB;
13431                out[t + 2..t + Q4_TILE].copy_from_slice(&v.nib[src..src + Q4TP_NIB]);
13432            }
13433        }
13434        out
13435    }
13436
13437    /// The exact (`CMF_SDOT=0`) path must reproduce `dequant_q4tp` to f32
13438    /// rounding — that scalar routine is the format's definition, and the
13439    /// kernels re-derive the scale from the ladder independently. Call the
13440    /// row kernel directly: `matmat` picks the int8 arm when a8w8 is on,
13441    /// so routing through it would test the other path by accident.
13442    #[test]
13443    fn q4tp_exact_path_matches_dequant_reference() {
13444        let (rows, cols) = (256usize, 512usize);
13445        let gpr = cols / GROUP_SIZE;
13446        let bytes = synth_q4tp(rows, cols);
13447        let mut w = vec![0f32; rows * cols];
13448        cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
13449
13450        let x: Vec<f32> = (0..cols)
13451            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
13452            .collect();
13453        let v = Q4tpView::new(&bytes, rows, cols);
13454        let mut sc = vec![0f32; gpr];
13455        for r in 0..rows {
13456            v.scales_into(r, gpr, &mut sc);
13457            let got = q4tp_row_exact(v.nib, r, gpr, &x, &sc);
13458            let want: f32 = (0..cols).map(|c| w[r * cols + c] * x[c]).sum();
13459            // These dot products cancel down to ~1e-3 from terms of ~5e-2, so
13460            // the meaningful yardstick is the summed magnitude, not the result:
13461            // against the result any reordering of a 512-term f32 sum "fails".
13462            let mag: f32 = (0..cols).map(|c| (w[r * cols + c] * x[c]).abs()).sum();
13463            assert!(
13464                (got - want).abs() <= 1e-5 * mag,
13465                "row {r}: kernel {got} vs dequant {want}"
13466            );
13467        }
13468    }
13469
13470    /// The int8 (a8w8) path can't be checked against an f32 reference — the
13471    /// activation quantization dominates. Check it against the q4t kernel it
13472    /// was ported from instead, on payloads holding the same weights: that
13473    /// isolates exactly what the port could break (16 B stride, ladder
13474    /// lookup, nibble unpack) from what it deliberately shares.
13475    #[test]
13476    fn q4tp_matvec_matches_the_q4t_kernel_it_was_ported_from() {
13477        let (rows, cols) = (256usize, 512usize);
13478        let bytes = synth_q4tp(rows, cols);
13479        let twin = q4tp_as_q4t(&bytes, rows, cols);
13480        let x: Vec<f32> = (0..cols)
13481            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
13482            .collect();
13483
13484        let mut got = vec![0f32; rows];
13485        q4tp_matvec(&bytes, &x, rows, cols, &mut got, None);
13486        let mut want = vec![0f32; rows];
13487        q4t_matvec(&twin, &x, rows, cols, &mut want, None);
13488
13489        // Scale is f16 in the twin and f32 here, so allow that rounding on
13490        // top of the summed magnitude (same cancellation argument as above).
13491        let mut w = vec![0f32; rows * cols];
13492        cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
13493        for r in 0..rows {
13494            let mag: f32 = (0..cols).map(|c| (w[r * cols + c] * x[c]).abs()).sum();
13495            assert!(
13496                (got[r] - want[r]).abs() <= 1e-3 * mag,
13497                "row {r}: q4tp {} vs q4t {}",
13498                got[r],
13499                want[r]
13500            );
13501        }
13502    }
13503
13504    /// `matmat` carries three arms (Accelerate, blocked int8 1x4, scalar).
13505    /// Batch 5 crosses the blocked kernel's stride, so this exercises the
13506    /// 1x4 path AND its scalar tail in one run — the blocked kernel is new
13507    /// code and its four accumulators are exactly what tends to go wrong.
13508    #[test]
13509    fn q4tp_matmat_matches_the_q4t_kernel_it_was_ported_from() {
13510        let (rows, cols, b) = (256usize, 512usize, 5usize);
13511        let bytes = synth_q4tp(rows, cols);
13512        let twin = q4tp_as_q4t(&bytes, rows, cols);
13513        let xs: Vec<f32> = (0..b * cols)
13514            .map(|i| ((i * 29 + 11) % 89) as f32 / 89.0 - 0.5)
13515            .collect();
13516
13517        let mut got = vec![0f32; b * rows];
13518        q4tp_matmat(&bytes, &xs, b, rows, cols, &mut got, None);
13519        let mut want = vec![0f32; b * rows];
13520        q4t_matmat(&twin, &xs, b, rows, cols, &mut want, None);
13521
13522        let mut w = vec![0f32; rows * cols];
13523        cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
13524        for t in 0..b {
13525            for r in 0..rows {
13526                let mag: f32 = (0..cols)
13527                    .map(|c| (w[r * cols + c] * xs[t * cols + c]).abs())
13528                    .sum();
13529                let (g, wa) = (got[t * rows + r], want[t * rows + r]);
13530                assert!(
13531                    (g - wa).abs() <= 1e-3 * mag,
13532                    "batch {t} row {r}: q4tp {g} vs q4t {wa}"
13533                );
13534            }
13535        }
13536    }
13537
13538    #[test]
13539    fn q4tp_matvec2_matches_the_single_stream_kernel() {
13540        let (rows, cols) = (128usize, 256usize);
13541        let gpr = cols / GROUP_SIZE;
13542        let bytes = synth_q4tp(rows, cols);
13543        let xs: Vec<f32> = (0..2 * cols)
13544            .map(|i| ((i * 29 + 11) % 89) as f32 / 89.0 - 0.5)
13545            .collect();
13546
13547        let (mut o1, mut o2) = (vec![0f32; rows], vec![0f32; rows]);
13548        q4tp_matvec2(
13549            &bytes,
13550            &xs[..cols],
13551            &xs[cols..],
13552            rows,
13553            cols,
13554            &mut o1,
13555            &mut o2,
13556            None,
13557        );
13558
13559        // matvec2 takes the exact path for both streams, so the single-row
13560        // kernel is an exact reference — no tolerance for path differences.
13561        let v = Q4tpView::new(&bytes, rows, cols);
13562        let mut sc = vec![0f32; gpr];
13563        for r in 0..rows {
13564            v.scales_into(r, gpr, &mut sc);
13565            assert_eq!(o1[r], q4tp_row_exact(v.nib, r, gpr, &xs[..cols], &sc));
13566            assert_eq!(o2[r], q4tp_row_exact(v.nib, r, gpr, &xs[cols..], &sc));
13567        }
13568    }
13569
13570    /// q4tp must not COST speed — it exists to save bytes, and a format that
13571    /// trades 7% of a file for a slower model is a bad trade. This guard is
13572    /// here because correctness tests happily passed while `q4tp_matmat` was
13573    /// missing its int8 and Accelerate arms and the model ran 5x slower.
13574    /// Measured on M-series: 0.97-1.04x, i.e. parity (16 B tiles are better
13575    /// aligned than q4t's 18 B, which pays for the scale indirection).
13576    #[test]
13577    fn q4tp_matvec_keeps_pace_with_q4t() {
13578        let (rows, cols) = (4096usize, 3072usize);
13579        let bytes = synth_q4tp(rows, cols);
13580        let twin = q4tp_as_q4t(&bytes, rows, cols);
13581        let x: Vec<f32> = (0..cols).map(|i| (i % 97) as f32 / 97.0 - 0.5).collect();
13582        let mut o = vec![0f32; rows];
13583        let n = 12;
13584        let mut best = (f64::MAX, f64::MAX);
13585        // Interleaved A/B, minimum statistic: this machine throttles, and a
13586        // mean over a thermal ramp reliably indicts whichever ran second.
13587        for _ in 0..3 {
13588            let t0 = std::time::Instant::now();
13589            for _ in 0..n {
13590                q4t_matvec(&twin, &x, rows, cols, &mut o, None);
13591            }
13592            best.0 = best.0.min(t0.elapsed().as_secs_f64());
13593            let t0 = std::time::Instant::now();
13594            for _ in 0..n {
13595                q4tp_matvec(&bytes, &x, rows, cols, &mut o, None);
13596            }
13597            best.1 = best.1.min(t0.elapsed().as_secs_f64());
13598        }
13599        let ratio = best.1 / best.0;
13600        println!(
13601            "q4t {:.3} ms | q4tp {:.3} ms | {ratio:.2}x",
13602            best.0 * 1e3 / n as f64,
13603            best.1 * 1e3 / n as f64
13604        );
13605        assert!(ratio < 2.0, "q4tp matvec {ratio:.2}x slower than q4t");
13606    }
13607
13608    #[cfg(target_os = "macos")]
13609    #[test]
13610    fn q4t_matmat_accel_matches_dequant_reference() {
13611        if !accel_gemm_enabled() {
13612            return; // CMF_ACCEL=0
13613        }
13614        let (rows, cols, b) = (512usize, 1024usize, 8usize); // ≥500K → accel arm
13615        let gpr = cols / GROUP_SIZE;
13616        let mut bytes = vec![0u8; rows * gpr * Q4_TILE];
13617        for r in 0..rows {
13618            for g in 0..gpr {
13619                let t = (r * gpr + g) * Q4_TILE;
13620                let sc = 0.02 + 0.0005 * ((r * gpr + g) % 64) as f32;
13621                bytes[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
13622                for k in 0..16 {
13623                    bytes[t + 2 + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
13624                }
13625            }
13626        }
13627        let x: Vec<f32> = (0..b * cols)
13628            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
13629            .collect();
13630        let mut got = vec![0f32; b * rows];
13631        q4t_matmat(&bytes, &x, b, rows, cols, &mut got, None);
13632        // Brute-force reference off the same tiles.
13633        let mut w = vec![0f32; rows * cols];
13634        for r in 0..rows {
13635            for g in 0..gpr {
13636                let t = (r * gpr + g) * Q4_TILE;
13637                let s = f16_to_f32(u16::from_le_bytes([bytes[t], bytes[t + 1]]));
13638                for (k, &bb) in bytes[t + 2..t + Q4_TILE].iter().enumerate() {
13639                    w[r * cols + g * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s;
13640                    w[r * cols + g * GROUP_SIZE + k * 2 + 1] =
13641                        (((bb >> 4) & 0x0F) as f32 - 8.0) * s;
13642                }
13643            }
13644        }
13645        for bi in 0..b {
13646            for r in 0..rows {
13647                let want: f32 = (0..cols).map(|j| x[bi * cols + j] * w[r * cols + j]).sum();
13648                let d = (got[bi * rows + r] - want).abs();
13649                assert!(
13650                    d <= want.abs().max(1.0) * 1e-4,
13651                    "accel q4t GEMM diverged at ({bi},{r}): {} vs {want}",
13652                    got[bi * rows + r]
13653                );
13654            }
13655        }
13656    }
13657
13658    #[test]
13659    fn q4matvec_matches_full_dequant() {
13660        let (rows, cols) = (8, 64);
13661        let groups = rows * cols / GROUP_SIZE;
13662        // Hand-craft a q4_block blob: nibbles then f16 scales.
13663        let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
13664        for i in 0..groups * 16 {
13665            bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
13666        }
13667        for g in 0..groups {
13668            let s = 0.01 + 0.003 * g as f32;
13669            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
13670        }
13671        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
13672
13673        let mut reference = vec![0.0f32; rows * cols];
13674        cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
13675        let mut expect = vec![0.0f32; rows];
13676        for r in 0..rows {
13677            expect[r] = reference[r * cols..(r + 1) * cols]
13678                .iter()
13679                .zip(&x)
13680                .map(|(w, xv)| w * xv)
13681                .sum();
13682        }
13683
13684        let mut got = vec![0.0f32; rows];
13685        q4matvec(&bytes, &x, rows, cols, &mut got, None);
13686        // SDOT path quantizes activations to i8 (A8W8): bounded noise,
13687        // same contract as q8/vbit (exact path is pinned by CMF_SDOT=0
13688        // in the golden-parity gate).
13689        let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
13690        let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
13691        for r in 0..rows {
13692            assert!(
13693                (got[r] - expect[r]).abs() < tol * scale,
13694                "row {r}: {} vs {}",
13695                got[r],
13696                expect[r]
13697            );
13698        }
13699    }
13700
13701    /// Fused two-input vbit matvec must equal two single matvecs exactly
13702    /// (same per-lane accumulation order on both scalar and SDOT paths).
13703    #[test]
13704    fn vbitmatvec2_equals_two_singles() {
13705        let (rows, cols) = (6, 64);
13706        let ng = cols / GROUP_SIZE;
13707        let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
13708        let mut bytes = bits.clone();
13709        for g in 0..rows * ng {
13710            let s = 0.02 + 0.001 * g as f32;
13711            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
13712        }
13713        for r in 0..rows {
13714            let b = bits[r] as usize;
13715            let (mut acc, mut nb) = (0u64, 0usize);
13716            let mut rowbytes = Vec::new();
13717            for i in 0..cols {
13718                let v = ((i * 7 + r * 13) % (1 << b)) as u64;
13719                acc = (acc << b) | v;
13720                nb += b;
13721                while nb >= 8 {
13722                    nb -= 8;
13723                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
13724                }
13725            }
13726            if nb > 0 {
13727                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
13728            }
13729            bytes.extend_from_slice(&rowbytes);
13730        }
13731        let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
13732        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).cos()).collect();
13733        let offsets = vbit_row_offsets(&bytes, rows, cols);
13734
13735        let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
13736        vbitmatvec(&bytes, &offsets, &x1, rows, cols, &mut a1, None);
13737        vbitmatvec(&bytes, &offsets, &x2, rows, cols, &mut a2, None);
13738        let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
13739        vbitmatvec2(
13740            &bytes, &offsets, &x1, &x2, rows, cols, &mut b1, &mut b2, None,
13741        );
13742        assert_eq!(a1, b1, "fused vbit lane 1 must be bit-identical");
13743        assert_eq!(a2, b2, "fused vbit lane 2 must be bit-identical");
13744    }
13745
13746    /// Fused two-input q4 matvec must equal two single matvecs exactly.
13747    #[test]
13748    fn q4matvec2_equals_two_singles() {
13749        let (rows, cols) = (8, 128);
13750        let groups = rows * cols / GROUP_SIZE;
13751        let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
13752        for i in 0..groups * 16 {
13753            bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
13754        }
13755        for g in 0..groups {
13756            let s = 0.01 + 0.003 * g as f32;
13757            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
13758        }
13759        // Include an outlier channel so the SDOT correction path is
13760        // exercised in the pair kernel too.
13761        let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
13762        x1[9] = 250.0;
13763        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
13764
13765        let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
13766        q4matvec(&bytes, &x1, rows, cols, &mut a1, None);
13767        q4matvec(&bytes, &x2, rows, cols, &mut a2, None);
13768        let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
13769        q4matvec2(&bytes, &x1, &x2, rows, cols, &mut b1, &mut b2, None);
13770        assert_eq!(a1, b1, "fused q4 lane 1 must be bit-identical");
13771        assert_eq!(a2, b2, "fused q4 lane 2 must be bit-identical");
13772    }
13773
13774    /// Multi-matrix job must equal separate matvecs exactly — same
13775    /// kernels, only the dispatch is fused.
13776    #[test]
13777    fn matvec_many_equals_separate_matvecs() {
13778        use crate::pool::Pool;
13779        let (r1, r2, cols) = (300, 200, 64);
13780        let mk = |salt: usize, rows: usize| {
13781            QTensor::from_f32(
13782                (0..rows * cols)
13783                    .map(|i| ((i * 7 + salt) % 97) as f32 / 97.0 - 0.5)
13784                    .collect(),
13785                rows,
13786                cols,
13787            )
13788        };
13789        let (a, b) = (mk(1, r1), mk(5, r2));
13790        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).sin()).collect();
13791        let pool = Pool::new(3);
13792
13793        let (mut ea, mut eb) = (vec![0f32; r1], vec![0f32; r2]);
13794        a.matvec(&x, &mut ea, Some(&pool));
13795        b.matvec(&x, &mut eb, Some(&pool));
13796        let (mut ga, mut gb) = (vec![0f32; r1], vec![0f32; r2]);
13797        QTensor::matvec_many([&a, &b], &x, [&mut ga, &mut gb], Some(&pool));
13798        assert_eq!(ea, ga, "fused multi-matrix lane 1 must be bit-identical");
13799        assert_eq!(eb, gb, "fused multi-matrix lane 2 must be bit-identical");
13800    }
13801
13802    /// The public Q4TP operator must take the real mapped matvec_many arm,
13803    /// rather than the F32 fallback above.  Build a tiny valid CMF so both
13804    /// handles retain their mmap payloads, then compare the fused dispatch
13805    /// with two ordinary mapped matvec calls bit-for-bit.
13806    #[test]
13807    fn q4tp_matvec_many_equals_separate_matvecs() {
13808        use crate::pool::Pool;
13809        use cortiq_core::{CMF_VERSION, CmfHeader, CmfModel, QuantType, TensorSpec};
13810
13811        let (r1, r2, cols) = (300usize, 200usize, 64usize);
13812        let arch: cortiq_core::ModelArch = serde_json::from_value(serde_json::json!({
13813            "arch_name": "tiny-q4tp",
13814            "hidden_size": cols,
13815            "intermediate_size": cols * 2,
13816            "num_layers": 1,
13817            "num_attention_heads": 2,
13818            "num_kv_heads": 1,
13819            "head_dim": 32,
13820            "vocab_size": r1,
13821            "layer_types": ["FullAttention"],
13822            "rms_norm_eps": 1e-6,
13823            "max_position_embeddings": 8,
13824            "linear_conv_kernel_dim": 0,
13825            "linear_num_key_heads": 0,
13826            "linear_num_value_heads": 0
13827        }))
13828        .unwrap();
13829        let header = CmfHeader {
13830            format: "cmf".into(),
13831            version: CMF_VERSION,
13832            arch,
13833            quant_type: QuantType::Q4Block,
13834            provenance: None,
13835            tokenizer_config: None,
13836            section_hashes: None,
13837            skills: Vec::new(),
13838            shard: None,
13839            calibration: None,
13840            routing: None,
13841            genome: None,
13842            lineage: Vec::new(),
13843            router: None,
13844            segments: Vec::new(),
13845        };
13846        let specs = [
13847            TensorSpec {
13848                name: "q".into(),
13849                dtype: TensorDtype::Q4TiledP,
13850                shape: vec![r1, cols],
13851                data: synth_q4tp(r1, cols),
13852            },
13853            TensorSpec {
13854                name: "kv".into(),
13855                dtype: TensorDtype::Q4TiledP,
13856                shape: vec![r2, cols],
13857                data: synth_q4tp(r2, cols),
13858            },
13859        ];
13860        let dir = std::env::temp_dir().join(format!("cmf-q4tp-many-{}", std::process::id()));
13861        std::fs::create_dir_all(&dir).unwrap();
13862        let path = dir.join("m.cmf");
13863        CmfModel::write(&path, &header, &specs, None, None).unwrap();
13864        let model = Arc::new(CmfModel::open(&path).unwrap());
13865        let (a, b) = (
13866            QTensor::from_model(&model, "q").unwrap(),
13867            QTensor::from_model(&model, "kv").unwrap(),
13868        );
13869        assert_eq!(a.model_dtype(), Some(TensorDtype::Q4TiledP));
13870        assert_eq!(b.model_dtype(), Some(TensorDtype::Q4TiledP));
13871        let x: Vec<f32> = (0..cols)
13872            .map(|i| ((i * 17 + 3) % 97) as f32 / 97.0 - 0.5)
13873            .collect();
13874        let pool = Pool::new(3);
13875        let (mut ea, mut eb) = (vec![0.0f32; r1], vec![0.0f32; r2]);
13876        a.matvec(&x, &mut ea, Some(&pool));
13877        b.matvec(&x, &mut eb, Some(&pool));
13878        let (mut ga, mut gb) = (vec![0.0f32; r1], vec![0.0f32; r2]);
13879        QTensor::matvec_many([&a, &b], &x, [&mut ga, &mut gb], Some(&pool));
13880        assert_eq!(ea, ga, "Q4TP fused lane 1 must be bit-identical");
13881        assert_eq!(eb, gb, "Q4TP fused lane 2 must be bit-identical");
13882        let _ = std::fs::remove_dir_all(&dir);
13883    }
13884
13885    /// The MiMo speculative verify's kernels: several tokens' MoE through
13886    /// `moe_gate_up_rows` / `moe_down_rows` (+ the caller's route-order sum)
13887    /// is bit-identical to each token alone through `moe_gate_up_many` /
13888    /// `moe_down_many` (decode), and a row-exact `q4tp_matmat` of five
13889    /// tokens (wide enough for the blocked tiles) equals five matvecs.
13890    #[test]
13891    fn multi_token_moe_rows_equal_single_token_decode() {
13892        use crate::pool::Pool;
13893        use cortiq_core::{CMF_VERSION, CmfHeader, CmfModel, QuantType, TensorSpec};
13894
13895        let (h, inter, ne) = (64usize, 128usize, 3usize);
13896        let arch: cortiq_core::ModelArch = serde_json::from_value(serde_json::json!({
13897            "arch_name": "tiny-q4tp-moe",
13898            "hidden_size": h,
13899            "intermediate_size": inter,
13900            "num_layers": 1,
13901            "num_attention_heads": 2,
13902            "num_kv_heads": 1,
13903            "head_dim": 32,
13904            "vocab_size": 8,
13905            "layer_types": ["FullAttention"],
13906            "rms_norm_eps": 1e-6,
13907            "max_position_embeddings": 8,
13908            "linear_conv_kernel_dim": 0,
13909            "linear_num_key_heads": 0,
13910            "linear_num_value_heads": 0
13911        }))
13912        .unwrap();
13913        let header = CmfHeader {
13914            format: "cmf".into(),
13915            version: CMF_VERSION,
13916            arch,
13917            quant_type: QuantType::Q4Block,
13918            provenance: None,
13919            tokenizer_config: None,
13920            section_hashes: None,
13921            skills: Vec::new(),
13922            shard: None,
13923            calibration: None,
13924            routing: None,
13925            genome: None,
13926            lineage: Vec::new(),
13927            router: None,
13928            segments: Vec::new(),
13929        };
13930        let mut specs = Vec::new();
13931        for e in 0..ne {
13932            for (k, (n, r, c)) in [("g", inter, h), ("u", inter, h), ("d", h, inter)]
13933                .into_iter()
13934                .enumerate()
13935            {
13936                // Distinct experts: perturb only the nibble plane (any byte
13937                // is a valid pair of codes; the ladder stays intact).
13938                let mut data = synth_q4tp(r, c);
13939                for (i, byte) in data[..r * (c / GROUP_SIZE) * Q4TP_NIB]
13940                    .iter_mut()
13941                    .enumerate()
13942                {
13943                    *byte ^= ((i * (e * 3 + k + 1)) % 251) as u8;
13944                }
13945                specs.push(TensorSpec {
13946                    name: format!("{n}{e}"),
13947                    dtype: TensorDtype::Q4TiledP,
13948                    shape: vec![r, c],
13949                    data,
13950                });
13951            }
13952        }
13953        let dir = std::env::temp_dir().join(format!(
13954            "cmf-moe-rows-{}-{}",
13955            std::process::id(),
13956            FLOAT_ACTIVATIONS.get()
13957        ));
13958        std::fs::create_dir_all(&dir).unwrap();
13959        let path = dir.join("m.cmf");
13960        CmfModel::write(&path, &header, &specs, None, None).unwrap();
13961        let model = Arc::new(CmfModel::open(&path).unwrap());
13962        let t = |n: String| QTensor::from_model(&model, &n).unwrap();
13963        let g: Vec<QTensor> = (0..ne).map(|e| t(format!("g{e}"))).collect();
13964        let u: Vec<QTensor> = (0..ne).map(|e| t(format!("u{e}"))).collect();
13965        let d: Vec<QTensor> = (0..ne).map(|e| t(format!("d{e}"))).collect();
13966        let b = 4usize;
13967        let mut xs: Vec<f32> = (0..b * h)
13968            .map(|i| ((i * 31 + 7) % 89) as f32 / 89.0 - 0.5)
13969            .collect();
13970        xs[5] = 9.0; // an activation outlier on token 0
13971        // Token -> (experts in route order, weights).
13972        let routes: Vec<(Vec<usize>, Vec<f32>)> = vec![
13973            (vec![2, 0], vec![0.6, 0.4]),
13974            (vec![0, 1, 2], vec![0.2, 0.5, 0.3]),
13975            (vec![1], vec![1.0]),
13976            (vec![2, 1, 0], vec![0.25, 0.25, 0.5]),
13977        ];
13978        let pool = Pool::new(3);
13979        // Decode reference, token by token.
13980        let mut want = vec![0f32; b * h];
13981        for (tk, (idx, w)) in routes.iter().enumerate() {
13982            let x = &xs[tk * h..(tk + 1) * h];
13983            let pairs: Vec<(&QTensor, &QTensor)> = idx.iter().map(|&e| (&g[e], &u[e])).collect();
13984            let mut gs: Vec<Vec<f32>> = idx.iter().map(|_| vec![0f32; inter]).collect();
13985            assert!(QTensor::moe_gate_up_many(&pairs, x, &mut gs, Some(&pool)));
13986            if FLOAT_ACTIVATIONS.get() {
13987                for (slot, &e) in idx.iter().enumerate() {
13988                    let (mut gate, mut up) = (vec![0.0; inter], vec![0.0; inter]);
13989                    g[e].matvec(x, &mut gate, Some(&pool));
13990                    u[e].matvec(x, &mut up, Some(&pool));
13991                    for (v, u) in gate.iter_mut().zip(up) {
13992                        *v = (*v / (1.0 + (-*v).exp())) * u;
13993                    }
13994                    assert_eq!(gs[slot], gate, "float gate/up must equal ordinary matvecs");
13995                }
13996            }
13997            let downs: Vec<&QTensor> = idx.iter().map(|&e| &d[e]).collect();
13998            assert!(QTensor::moe_down_many(
13999                &downs,
14000                &gs,
14001                w,
14002                &mut want[tk * h..(tk + 1) * h],
14003                Some(&pool)
14004            ));
14005        }
14006        if FLOAT_ACTIVATIONS.get() {
14007            for (tk, (idx, w)) in routes.iter().enumerate() {
14008                let mut scalar = vec![0.0; h];
14009                for (&e, &weight) in idx.iter().zip(w) {
14010                    let (mut gate, mut up, mut down) =
14011                        (vec![0.0; inter], vec![0.0; inter], vec![0.0; h]);
14012                    g[e].matvec(&xs[tk * h..(tk + 1) * h], &mut gate, Some(&pool));
14013                    u[e].matvec(&xs[tk * h..(tk + 1) * h], &mut up, Some(&pool));
14014                    for (v, u) in gate.iter_mut().zip(up) {
14015                        *v = (*v / (1.0 + (-*v).exp())) * u;
14016                    }
14017                    d[e].matvec(&gate, &mut down, Some(&pool));
14018                    for (v, d) in scalar.iter_mut().zip(down) {
14019                        *v += weight * d;
14020                    }
14021                }
14022                assert_eq!(
14023                    &want[tk * h..(tk + 1) * h],
14024                    scalar,
14025                    "float many equals scalar experts"
14026                );
14027            }
14028        }
14029        // All four tokens at once, grouped by expert.
14030        let mut experts: Vec<usize> = Vec::new();
14031        let mut groups: Vec<Vec<usize>> = Vec::new();
14032        for (tk, (idx, _)) in routes.iter().enumerate() {
14033            for &e in idx {
14034                match experts.iter().position(|&x| x == e) {
14035                    Some(k) => groups[k].push(tk),
14036                    None => {
14037                        experts.push(e);
14038                        groups.push(vec![tk]);
14039                    }
14040                }
14041            }
14042        }
14043        let n_pairs: usize = groups.iter().map(|g| g.len()).sum();
14044        let pairs: Vec<(&QTensor, &QTensor)> = experts.iter().map(|&e| (&g[e], &u[e])).collect();
14045        let mut gs: Vec<Vec<f32>> = (0..n_pairs).map(|_| vec![0f32; inter]).collect();
14046        assert!(QTensor::moe_gate_up_rows(
14047            &pairs,
14048            &groups,
14049            &xs,
14050            &mut gs,
14051            Some(&pool)
14052        ));
14053        let downs: Vec<&QTensor> = experts.iter().map(|&e| &d[e]).collect();
14054        let lens: Vec<usize> = groups.iter().map(|g| g.len()).collect();
14055        let mut ds: Vec<Vec<f32>> = (0..n_pairs).map(|_| vec![0f32; h]).collect();
14056        assert!(QTensor::moe_down_rows(
14057            &downs,
14058            &lens,
14059            &gs,
14060            &mut ds,
14061            Some(&pool)
14062        ));
14063        let slot = |tk: usize, e: usize| {
14064            let k = experts.iter().position(|&x| x == e).unwrap();
14065            groups[..k].iter().map(|g| g.len()).sum::<usize>()
14066                + groups[k].iter().position(|&x| x == tk).unwrap()
14067        };
14068        let mut got = vec![0f32; b * h];
14069        for (tk, (idx, w)) in routes.iter().enumerate() {
14070            for i in 0..h {
14071                let mut acc = 0f32;
14072                for (&e, &we) in idx.iter().zip(w) {
14073                    acc += we * ds[slot(tk, e)][i];
14074                }
14075                got[tk * h + i] = acc;
14076            }
14077        }
14078        assert!(want.iter().any(|v| *v != 0.0));
14079        assert_eq!(
14080            want.iter().map(|v| v.to_bits()).collect::<Vec<_>>(),
14081            got.iter().map(|v| v.to_bits()).collect::<Vec<_>>(),
14082            "multi-token MoE must equal decode bit for bit"
14083        );
14084
14085        // Row-exact q4tp_matmat: five tokens (a blocked 1x4 tile + tail
14086        // otherwise) equal five matvecs.
14087        let b5 = 5usize;
14088        let x5: Vec<f32> = (0..b5 * h)
14089            .map(|i| ((i * 13 + 5) % 71) as f32 / 71.0 - 0.5)
14090            .collect();
14091        let mut mm = vec![0f32; b5 * inter];
14092        row_exact_scope(|| g[1].matmat(&x5, b5, &mut mm, Some(&pool)));
14093        for tk in 0..b5 {
14094            let mut mv = vec![0f32; inter];
14095            g[1].matvec(&x5[tk * h..(tk + 1) * h], &mut mv, Some(&pool));
14096            assert_eq!(
14097                mv.iter().map(|v| v.to_bits()).collect::<Vec<_>>(),
14098                mm[tk * inter..(tk + 1) * inter]
14099                    .iter()
14100                    .map(|v| v.to_bits())
14101                    .collect::<Vec<_>>(),
14102                "row-exact matmat token {tk}"
14103            );
14104        }
14105        // Other concurrent tests/requests may still hold the shared mode.
14106        // Nested, overlapping and unwind restoration is checked separately.
14107        let _ = std::fs::remove_dir_all(&dir);
14108    }
14109
14110    /// The row-exact fix must not touch the fast path. Outside the scope
14111    /// `q4tp_matmat` has to produce exactly what it did before, and on ARM
14112    /// "before" is spelled out below: the tuned 1x4 SDOT tile for every
14113    /// four columns and the single-row kernel for the tail. Inside the
14114    /// scope every column equals its token's matvec. `q4tp_matmat_with`
14115    /// takes the mode as an argument, so a concurrent test holding the
14116    /// shared scope cannot flip it under this one.
14117    #[test]
14118    fn q4tp_matmat_fast_path_unchanged_outside_row_exact() {
14119        use crate::pool::Pool;
14120        use std::sync::atomic::Ordering::Relaxed;
14121        let _alt = Q4TP_ALT_TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner());
14122        // The tuned ARM shape, which is also what an unset switch picks
14123        // unless CMF_Q4TP_V1 is exported.
14124        Q4TP_ALT.store(2, Relaxed);
14125        let pool = Pool::new(3);
14126        // Under 500k cells, so macOS keeps the matmat off the AMX; the
14127        // second shape runs across pool workers (rows >= 256) with 32
14128        // groups of accumulation and a tail after two 1x4 tiles.
14129        for &(rows, cols, b) in &[(64usize, 256usize, 7usize), (320, 1024, 9)] {
14130            let bytes = synth_q4tp(rows, cols);
14131            let mut xs: Vec<f32> = (0..b * cols)
14132                .map(|i| ((i * 29 + 11) % 83) as f32 / 83.0 - 0.5)
14133                .collect();
14134            xs[3] = 7.5; // an activation outlier on token 0
14135            let run = |exact: bool| {
14136                let mut out = vec![0f32; b * rows];
14137                q4tp_matmat_with(&bytes, &xs, b, rows, cols, &mut out, Some(&pool), exact);
14138                out
14139            };
14140            let (fast, exact) = (run(false), run(true));
14141            #[cfg(not(target_arch = "aarch64"))]
14142            let _ = fast;
14143            let bits = |v: &[f32]| v.iter().map(|x| x.to_bits()).collect::<Vec<_>>();
14144            let mut matvecs = vec![0f32; b * rows];
14145            for (bi, o) in matvecs.chunks_mut(rows).enumerate() {
14146                q4tp_matvec(
14147                    &bytes,
14148                    &xs[bi * cols..(bi + 1) * cols],
14149                    rows,
14150                    cols,
14151                    o,
14152                    Some(&pool),
14153                );
14154            }
14155            assert!(matvecs.iter().any(|v| *v != 0.0));
14156            assert_eq!(
14157                bits(&exact),
14158                bits(&matvecs),
14159                "{rows}x{cols} b={b}: row-exact matmat must equal per-token matvecs"
14160            );
14161            #[cfg(target_arch = "aarch64")]
14162            {
14163                // The pre-fix ARM loop, cell for cell.
14164                let gpr = cols / GROUP_SIZE;
14165                let v = Q4tpView::new(&bytes, rows, cols);
14166                let mut old = vec![0f32; b * rows];
14167                let mut sc = vec![0f32; gpr];
14168                let a8w8 = a8w8_enabled();
14169                let blocked = sdot_enabled() && blocked_enabled();
14170                let acts: Vec<SplitAct> = (0..b)
14171                    .map(|bi| split_act(&xs[bi * cols..(bi + 1) * cols]))
14172                    .collect();
14173                for r in 0..rows {
14174                    v.scales_into(r, gpr, &mut sc);
14175                    if !a8w8 {
14176                        for bi in 0..b {
14177                            let x = &xs[bi * cols..(bi + 1) * cols];
14178                            old[bi * rows + r] = q4tp_row_exact(v.nib, r, gpr, x, &sc);
14179                        }
14180                        continue;
14181                    }
14182                    let finish = |d: f32, act: &SplitAct| {
14183                        let mut acc = d * act.sx;
14184                        for &(j, xv) in &act.outliers {
14185                            let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
14186                            acc += w * s * xv;
14187                        }
14188                        acc
14189                    };
14190                    let mut bi = 0usize;
14191                    while blocked && bi + 4 <= b {
14192                        let xs4 = [
14193                            acts[bi].xq.as_slice(),
14194                            acts[bi + 1].xq.as_slice(),
14195                            acts[bi + 2].xq.as_slice(),
14196                            acts[bi + 3].xq.as_slice(),
14197                        ];
14198                        let d = unsafe { dot_q4tp_row_1x4_sdot(v.nib, r, gpr, xs4, &sc) };
14199                        for k in 0..4 {
14200                            old[(bi + k) * rows + r] = finish(d[k], &acts[bi + k]);
14201                        }
14202                        bi += 4;
14203                    }
14204                    for (bi, act) in acts.iter().enumerate().skip(bi) {
14205                        let d = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, &sc);
14206                        old[bi * rows + r] = finish(d, act);
14207                    }
14208                }
14209                assert_eq!(
14210                    bits(&fast),
14211                    bits(&old),
14212                    "{rows}x{cols} b={b}: the fast path changed outside row_exact"
14213                );
14214                // And it is still the fast tile that runs: its lane-parallel
14215                // fma sum rounds differently from the matvec somewhere.
14216                if blocked {
14217                    assert_ne!(
14218                        bits(&fast),
14219                        bits(&matvecs),
14220                        "{rows}x{cols} b={b}: the tuned tile no longer runs outside row_exact"
14221                    );
14222                }
14223            }
14224        }
14225        Q4TP_ALT.store(0, Relaxed);
14226    }
14227
14228    #[test]
14229    fn row_exact_scopes_survive_overlap_nesting_and_unwind() {
14230        use std::sync::{Barrier, atomic::{AtomicUsize, Ordering}};
14231        // A private counter makes this restoration test independent of
14232        // numerical tests concurrently using the production counter.
14233        let active = AtomicUsize::new(0);
14234        counted_row_exact_scope(&active, || {
14235            assert_eq!(active.load(Ordering::Acquire), 1);
14236            counted_row_exact_scope(&active, || {
14237                assert_eq!(active.load(Ordering::Acquire), 2);
14238            });
14239            assert_eq!(active.load(Ordering::Acquire), 1);
14240        });
14241        assert_eq!(active.load(Ordering::Acquire), 0);
14242
14243        let both_entered = Barrier::new(2);
14244        let release_last = Barrier::new(2);
14245        std::thread::scope(|s| {
14246            let first = s.spawn(|| counted_row_exact_scope(&active, || {
14247                both_entered.wait();
14248            }));
14249            let last = s.spawn(|| counted_row_exact_scope(&active, || {
14250                both_entered.wait();
14251                release_last.wait();
14252            }));
14253            first.join().unwrap();
14254            let after_first = active.load(Ordering::Acquire);
14255            release_last.wait();
14256            last.join().unwrap();
14257            assert_eq!(after_first, 1, "second request must remain exact");
14258        });
14259        assert_eq!(active.load(Ordering::Acquire), 0);
14260        let panic = std::panic::catch_unwind(|| {
14261            counted_row_exact_scope(&active, || panic!("scope unwind"));
14262        });
14263        assert!(panic.is_err());
14264        assert_eq!(active.load(Ordering::Acquire), 0);
14265    }
14266
14267    #[test]
14268    #[cfg(target_arch = "x86_64")]
14269    fn q4tp_float_avx2_is_bitwise_scalar() {
14270        if !avx2_enabled() {
14271            return;
14272        }
14273        for cols in [32, 64, 96, 2048, 4096] {
14274            let rows = 9;
14275            let bytes = synth_q4tp(rows, cols);
14276            let v = Q4tpView::new(&bytes, rows, cols);
14277            let gpr = cols / GROUP_SIZE;
14278            let mut sc = vec![0.0; gpr];
14279            for seed in 1..=5 {
14280                let xs: Vec<f32> = (0..cols)
14281                    .map(|i| (((i * 104729 + seed * 8191) % 100003) as f32 - 50001.0) / 7919.0)
14282                    .collect();
14283                for r in 0..rows {
14284                    v.scales_into(r, gpr, &mut sc);
14285                    let scalar = q4tp_row_float_scalar(v.nib, r, gpr, &xs, &sc);
14286                    let vector = unsafe { q4tp_row_float_avx2(v.nib, r, gpr, &xs, &sc) };
14287                    assert_eq!(
14288                        scalar.to_bits(),
14289                        vector.to_bits(),
14290                        "cols={cols} row={r} seed={seed}"
14291                    );
14292                }
14293            }
14294        }
14295    }
14296
14297    #[test]
14298    fn multi_token_moe_rows_float_equal_single_token_decode() {
14299        float_activations_scope(multi_token_moe_rows_equal_single_token_decode);
14300    }
14301
14302    #[test]
14303    fn full_gpu_q8_scope_is_nested_and_thread_local() {
14304        assert!(!FULL_GPU_Q8.get());
14305        let before = gpu_split_frac();
14306        {
14307            let _guard = enter_full_gpu_q8_scope();
14308            assert_eq!(gpu_split_frac(), 1.0);
14309            {
14310                let _nested = enter_full_gpu_q8_scope();
14311            }
14312            assert_eq!(gpu_split_frac(), 1.0);
14313            std::thread::spawn(|| assert!(!FULL_GPU_Q8.get())).join().unwrap();
14314        }
14315        assert!(!FULL_GPU_Q8.get());
14316        assert_eq!(gpu_split_frac(), before);
14317    }
14318
14319    #[test]
14320    fn float_activation_scope_is_nested_thread_local_and_unwind_safe() {
14321        assert!(!FLOAT_ACTIVATIONS.get());
14322        let before = a8w8_enabled();
14323        float_activations_scope(|| {
14324            assert!(!a8w8_enabled());
14325            float_activations_scope(|| assert!(!a8w8_enabled()));
14326            assert!(FLOAT_ACTIVATIONS.get());
14327            std::thread::spawn(|| assert!(!FLOAT_ACTIVATIONS.get()))
14328                .join()
14329                .unwrap();
14330        });
14331        assert!(!FLOAT_ACTIVATIONS.get());
14332        assert_eq!(a8w8_enabled(), before);
14333        let _ = std::panic::catch_unwind(|| float_activations_scope(|| panic!("test unwind")));
14334        assert!(!FLOAT_ACTIVATIONS.get());
14335    }
14336
14337    /// Batched q4/vbit matmat must equal per-position matvec calls
14338    /// exactly (the fallback it replaced) — same kernels, same order.
14339    #[test]
14340    fn batched_matmat_equals_per_position_matvec() {
14341        let (rows, cols, b) = (8, 64, 5);
14342        // q4 blob.
14343        let groups = rows * cols / GROUP_SIZE;
14344        let mut q4 = Vec::new();
14345        for i in 0..groups * 16 {
14346            q4.push((((i * 7 + 3) % 256) & 0xFF) as u8);
14347        }
14348        for g in 0..groups {
14349            q4.extend_from_slice(
14350                &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
14351            );
14352        }
14353        // vbit blob (mixed widths incl. 8).
14354        let ng = cols / GROUP_SIZE;
14355        let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4, 5, 3];
14356        let mut vb = bits.clone();
14357        for g in 0..rows * ng {
14358            vb.extend_from_slice(
14359                &cortiq_core::quant::f32_to_f16(0.02 + 0.001 * g as f32).to_le_bytes(),
14360            );
14361        }
14362        for r in 0..rows {
14363            let bw = bits[r] as usize;
14364            let (mut acc, mut nb) = (0u64, 0usize);
14365            let mut rowbytes = Vec::new();
14366            for i in 0..cols {
14367                let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
14368                acc = (acc << bw) | v;
14369                nb += bw;
14370                while nb >= 8 {
14371                    nb -= 8;
14372                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
14373                }
14374            }
14375            if nb > 0 {
14376                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
14377            }
14378            vb.extend_from_slice(&rowbytes);
14379        }
14380        let offsets = vbit_row_offsets(&vb, rows, cols);
14381
14382        let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
14383
14384        // q4: batch vs singles.
14385        let mut got = vec![0f32; b * rows];
14386        q4matmat(&q4, &xs, b, rows, cols, &mut got, None);
14387        for bi in 0..b {
14388            let mut expect = vec![0f32; rows];
14389            q4matvec(
14390                &q4,
14391                &xs[bi * cols..(bi + 1) * cols],
14392                rows,
14393                cols,
14394                &mut expect,
14395                None,
14396            );
14397            assert_eq!(
14398                &got[bi * rows..(bi + 1) * rows],
14399                &expect[..],
14400                "q4 batch pos {bi}"
14401            );
14402        }
14403
14404        // vbit: batch vs singles.
14405        let mut got = vec![0f32; b * rows];
14406        vbitmatmat(&vb, &offsets, &xs, b, rows, cols, &mut got, None);
14407        for bi in 0..b {
14408            let mut expect = vec![0f32; rows];
14409            vbitmatvec(
14410                &vb,
14411                &offsets,
14412                &xs[bi * cols..(bi + 1) * cols],
14413                rows,
14414                cols,
14415                &mut expect,
14416                None,
14417            );
14418            assert_eq!(
14419                &got[bi * rows..(bi + 1) * rows],
14420                &expect[..],
14421                "vbit batch pos {bi}"
14422            );
14423        }
14424    }
14425
14426    /// q4_tiled kernels must produce BIT-identical outputs to the q4
14427    /// split kernels on the same values (same ints, same order — only
14428    /// the byte placement differs).
14429    #[test]
14430    fn q4_tiled_matches_q4_block_bitexact() {
14431        let (rows, cols, b) = (8usize, 128usize, 3usize);
14432        let groups = rows * cols / GROUP_SIZE;
14433        let mut split = Vec::with_capacity(groups * 18);
14434        for i in 0..groups * 16 {
14435            split.push((((i * 7 + 3) % 256) & 0xFF) as u8);
14436        }
14437        for g in 0..groups {
14438            split.extend_from_slice(
14439                &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
14440            );
14441        }
14442        // Re-tile: [scale][nibbles] per group.
14443        let (packed, scales) = split.split_at(groups * 16);
14444        let mut tiled = Vec::with_capacity(groups * Q4_TILE);
14445        for g in 0..groups {
14446            tiled.extend_from_slice(&scales[g * 2..g * 2 + 2]);
14447            tiled.extend_from_slice(&packed[g * 16..(g + 1) * 16]);
14448        }
14449
14450        let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
14451        x1[9] = 250.0; // exercise the outlier path
14452        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
14453
14454        let (mut a, mut t) = (vec![0f32; rows], vec![0f32; rows]);
14455        q4matvec(&split, &x1, rows, cols, &mut a, None);
14456        q4t_matvec(&tiled, &x1, rows, cols, &mut t, None);
14457        assert_eq!(a, t, "q4t matvec must match q4 bit-for-bit");
14458
14459        let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
14460        let (mut t1, mut t2) = (vec![0f32; rows], vec![0f32; rows]);
14461        q4matvec2(&split, &x1, &x2, rows, cols, &mut a1, &mut a2, None);
14462        q4t_matvec2(&tiled, &x1, &x2, rows, cols, &mut t1, &mut t2, None);
14463        assert_eq!(a1, t1);
14464        assert_eq!(a2, t2);
14465
14466        let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
14467        let (mut am, mut tm) = (vec![0f32; b * rows], vec![0f32; b * rows]);
14468        q4matmat(&split, &xs, b, rows, cols, &mut am, None);
14469        q4t_matmat(&tiled, &xs, b, rows, cols, &mut tm, None);
14470        assert_eq!(am, tm, "q4t matmat must match q4 bit-for-bit");
14471    }
14472
14473    /// q4 SDOT outlier correction: a single huge activation channel
14474    /// (>8·rms → outlier, zeroed in xq) must still contribute its EXACT
14475    /// term. On-grid bulk (±1/0 → xq dequantizes exactly) isolates the
14476    /// correction from A8W8 noise. cols must exceed 64: at n=64 the
14477    /// 8·rms threshold equals sqrt(v²+rest) ≥ v, so a single outlier
14478    /// can never qualify (8² = n).
14479    #[test]
14480    fn q4matvec_sdot_outlier_exact() {
14481        let (rows, cols) = (4, 128);
14482        let groups = rows * cols / GROUP_SIZE;
14483        let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
14484        for i in 0..groups * 16 {
14485            bytes.push(((i * 11 + 5) % 256) as u8);
14486        }
14487        for g in 0..groups {
14488            let s = 0.02 + 0.002 * g as f32;
14489            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
14490        }
14491        let mut x: Vec<f32> = (0..cols)
14492            .map(|i| match i % 3 {
14493                0 => 1.0,
14494                1 => -1.0,
14495                _ => 0.0,
14496            })
14497            .collect();
14498        x[17] = 300.0; // ≫ 8·rms → outlier channel
14499
14500        let mut reference = vec![0.0f32; rows * cols];
14501        cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
14502        let mut expect = vec![0.0f32; rows];
14503        for r in 0..rows {
14504            expect[r] = reference[r * cols..(r + 1) * cols]
14505                .iter()
14506                .zip(&x)
14507                .map(|(w, xv)| w * xv)
14508                .sum();
14509        }
14510        let mut got = vec![0.0f32; rows];
14511        q4matvec(&bytes, &x, rows, cols, &mut got, None);
14512        let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
14513        for r in 0..rows {
14514            assert!(
14515                (got[r] - expect[r]).abs() < 2e-3 * scale,
14516                "row {r}: {} vs {} (outlier term must be exact)",
14517                got[r],
14518                expect[r]
14519            );
14520        }
14521    }
14522
14523    /// The fused q1t matvec must equal the reference (dequant_q1t → dot),
14524    /// including the ternary zero level and the binary-searched outlier
14525    /// overlay. Guards the mmap kernel that makes a 12B q1t runnable.
14526    #[test]
14527    fn q1t_matvec_matches_reference() {
14528        use cortiq_core::quant::{dequant_q1t, f32_to_f16};
14529        let (rows, cols) = (3usize, 64usize); // gpr = 2
14530        let gpr = cols / GROUP_SIZE;
14531        let scales = [0.5f32, 0.3, 0.7, 0.2, 0.6, 0.15];
14532        // Overlay (must be sorted by flat index): a few spikes across rows.
14533        let outliers: [(u32, f32); 3] = [(5, 9.0), (70, -4.5), (150, 3.25)];
14534        let is_out = |flat: usize| outliers.iter().any(|&(i, _)| i as usize == flat);
14535        let mut bytes = Vec::new();
14536        for r in 0..rows {
14537            for g in 0..gpr {
14538                bytes.extend_from_slice(&f32_to_f16(scales[r * gpr + g]).to_le_bytes());
14539                let mut c = [0u8; 7];
14540                for k in 0..GROUP_SIZE {
14541                    // Encoder invariant: code 0 at outlier positions.
14542                    let code = if is_out(r * cols + g * GROUP_SIZE + k) {
14543                        0
14544                    } else {
14545                        ((k + r * 3 + g) % 3) as u8 // 0,1,2
14546                    };
14547                    cortiq_core::quant::q1t_pack(&mut c, k, code);
14548                }
14549                bytes.extend_from_slice(&c);
14550            }
14551        }
14552        // Per-row overlay: [u32 row_ptr[rows+1]] then [(u16 col, f16 val)] by
14553        // row (outliers are sorted by flat index → already grouped by row).
14554        let mut row_ptr = vec![0u32; rows + 1];
14555        for &(idx, _) in &outliers {
14556            row_ptr[idx as usize / cols + 1] += 1;
14557        }
14558        for r in 0..rows {
14559            row_ptr[r + 1] += row_ptr[r];
14560        }
14561        for &p in &row_ptr {
14562            bytes.extend_from_slice(&p.to_le_bytes());
14563        }
14564        for &(idx, v) in &outliers {
14565            bytes.extend_from_slice(&((idx as usize % cols) as u16).to_le_bytes());
14566            bytes.extend_from_slice(&f32_to_f16(v).to_le_bytes());
14567        }
14568
14569        let mut refw = vec![0f32; rows * cols];
14570        dequant_q1t(&bytes, rows, cols, &mut refw);
14571        // On-grid activations (±1, amax 1) so the int8 SDOT path reconstructs
14572        // x exactly and matches the f32 reference (same trick as the q1 test).
14573        let x: Vec<f32> = (0..cols)
14574            .map(|j| if j % 3 == 0 { 1.0 } else { -1.0 })
14575            .collect();
14576        let mut expect = vec![0f32; rows];
14577        for r in 0..rows {
14578            let mut a = 0.0f32;
14579            for j in 0..cols {
14580                a += refw[r * cols + j] * x[j];
14581            }
14582            expect[r] = a;
14583        }
14584        let tol = |e: f32| 1e-3 * e.abs().max(1e-3);
14585        let mut got = vec![0f32; rows];
14586        q1t_matvec(&bytes, &x, rows, cols, &mut got, None);
14587        for r in 0..rows {
14588            assert!(
14589                (got[r] - expect[r]).abs() < tol(expect[r]),
14590                "row {r}: {} vs {}",
14591                got[r],
14592                expect[r]
14593            );
14594        }
14595        // matmat (b=2, f32 decode path) must agree too.
14596        let x2: Vec<f32> = x.iter().chain(x.iter()).copied().collect();
14597        let mut gm = vec![0f32; 2 * rows];
14598        q1t_matmat(&bytes, &x2, 2, rows, cols, &mut gm, None);
14599        for r in 0..rows {
14600            assert!((gm[r] - expect[r]).abs() < tol(expect[r]));
14601            assert!((gm[rows + r] - expect[r]).abs() < tol(expect[r]));
14602        }
14603        // Fused pair (q1t_matvec2) must equal two single matvecs
14604        // bit-for-bit: same unpack, same group order, same f32
14605        // accumulation per stream. Distinct x2 exercises both lanes.
14606        let xb: Vec<f32> = (0..cols)
14607            .map(|j| if j % 5 == 0 { -1.0 } else { 1.0 })
14608            .collect();
14609        let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
14610        q1t_matvec(&bytes, &x, rows, cols, &mut s1, None);
14611        q1t_matvec(&bytes, &xb, rows, cols, &mut s2, None);
14612        let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
14613        q1t_matvec2(&bytes, &x, &xb, rows, cols, &mut p1, &mut p2, None);
14614        assert_eq!(p1, s1, "q1t pair lane 1 ≠ single matvec");
14615        assert_eq!(p2, s2, "q1t pair lane 2 ≠ single matvec");
14616    }
14617
14618    /// Pair == 2×matvec with an ODD group count (the kernel's tail
14619    /// group) and no overlay section.
14620    #[test]
14621    fn q1t_matvec2_odd_gpr_matches_singles() {
14622        use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_pack};
14623        let (rows, cols) = (5usize, 96usize); // gpr = 3 → paired + tail
14624        let gpr = cols / GROUP_SIZE;
14625        let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE);
14626        for r in 0..rows {
14627            for g in 0..gpr {
14628                bytes.extend_from_slice(&f32_to_f16(0.1 + 0.05 * (r + g) as f32).to_le_bytes());
14629                let mut c = [0u8; 7];
14630                for k in 0..GROUP_SIZE {
14631                    q1t_pack(&mut c, k, ((k * 7 + r * 5 + g * 3) % 3) as u8);
14632                }
14633                bytes.extend_from_slice(&c);
14634            }
14635        }
14636        let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.31).sin()).collect();
14637        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).cos()).collect();
14638        let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
14639        q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
14640        q1t_matvec(&bytes, &x2, rows, cols, &mut s2, None);
14641        let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
14642        q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
14643        assert_eq!(p1, s1, "odd-gpr pair lane 1 ≠ single");
14644        assert_eq!(p2, s2, "odd-gpr pair lane 2 ≠ single");
14645    }
14646
14647    // Speed A/B: fused pair (one unpack, two streams) vs two single
14648    // matvecs. Single-threaded, FFN-sized, min-of paired in-process.
14649    //   cargo test -p cortiq-engine --release q1t_matvec2_speed -- --ignored --nocapture
14650    #[test]
14651    #[ignore]
14652    fn q1t_matvec2_speed() {
14653        use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_pack};
14654        use std::time::Instant;
14655        let (rows, cols) = (8192usize, 4096usize);
14656        let gpr = cols / GROUP_SIZE;
14657        let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE);
14658        for r in 0..rows {
14659            for g in 0..gpr {
14660                let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
14661                bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
14662                let mut c = [0u8; 7];
14663                for k in 0..GROUP_SIZE {
14664                    q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
14665                }
14666                bytes.extend_from_slice(&c);
14667            }
14668        }
14669        let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.31).sin()).collect();
14670        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).cos()).collect();
14671        let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
14672        let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
14673        // Warm both paths once.
14674        q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
14675        q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
14676        let (mut t_pair, mut t_two) = (f64::MAX, f64::MAX);
14677        for _ in 0..8 {
14678            let t0 = Instant::now();
14679            q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
14680            t_pair = t_pair.min(t0.elapsed().as_secs_f64() * 1000.0);
14681            let t1 = Instant::now();
14682            q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
14683            q1t_matvec(&bytes, &x2, rows, cols, &mut s2, None);
14684            t_two = t_two.min(t1.elapsed().as_secs_f64() * 1000.0);
14685        }
14686        assert_eq!(p1, s1);
14687        assert_eq!(p2, s2);
14688        println!("q1t pair {rows}x{cols}: fused {t_pair:.2} ms | two singles {t_two:.2} ms");
14689    }
14690
14691    // Speed A/B: the base-3-division decode (what the packing commit left in
14692    // place) vs the fused sign-LUT matvec. Both single-threaded, same bytes.
14693    //   cargo test -p cortiq-engine q1t_matvec_speed -- --ignored --nocapture
14694    #[test]
14695    #[ignore]
14696    fn q1t_matvec_speed() {
14697        use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_code, q1t_pack};
14698        use std::time::Instant;
14699        let (rows, cols) = (8192usize, 4096usize); // FFN-sized
14700        let gpr = cols / GROUP_SIZE;
14701        let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE + 16);
14702        for r in 0..rows {
14703            for g in 0..gpr {
14704                let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
14705                bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
14706                let mut c = [0u8; 7];
14707                for k in 0..GROUP_SIZE {
14708                    q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
14709                }
14710                bytes.extend_from_slice(&c);
14711            }
14712        }
14713        let (n, stride) = (rows * cols, 40usize); // ~2.5% outliers, per-row overlay
14714        let mut row_ptr = vec![0u32; rows + 1];
14715        let mut idx = 0usize;
14716        while idx < n {
14717            row_ptr[idx / cols + 1] += 1;
14718            idx += stride;
14719        }
14720        for r in 0..rows {
14721            row_ptr[r + 1] += row_ptr[r];
14722        }
14723        for &p in &row_ptr {
14724            bytes.extend_from_slice(&p.to_le_bytes());
14725        }
14726        let mut idx = 0usize;
14727        while idx < n {
14728            bytes.extend_from_slice(&((idx % cols) as u16).to_le_bytes());
14729            bytes.extend_from_slice(&f32_to_f16((idx % 13) as f32 * 0.1 - 0.6).to_le_bytes());
14730            idx += stride;
14731        }
14732        // On-grid ±1 so the fast path's int8 SDOT is exact vs the f32 "slow"
14733        // reference (the A/B is a timing check; values must still agree).
14734        let x: Vec<f32> = (0..cols)
14735            .map(|j| if j % 3 == 0 { 1.0 } else { -1.0 })
14736            .collect();
14737        let (rp_off, ent_off, has_ov) = q1t_overlay(&bytes, rows * gpr * Q1T_TILE, rows);
14738
14739        // "before": base-3 division decode into a buffer, then dot.
14740        let slow = |out: &mut [f32]| {
14741            let mut buf = vec![0f32; cols];
14742            for r in 0..rows {
14743                for g in 0..gpr {
14744                    let off = (r * gpr + g) * Q1T_TILE;
14745                    let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
14746                    let codes = &bytes[off + 2..off + Q1T_TILE];
14747                    for k in 0..GROUP_SIZE {
14748                        buf[g * GROUP_SIZE + k] = match q1t_code(codes, k) {
14749                            1 => s,
14750                            2 => -s,
14751                            _ => 0.0,
14752                        };
14753                    }
14754                }
14755                out[r] = q1t_row_outlier_correction(&bytes, r, rp_off, ent_off, has_ov, &x)
14756                    + (0..cols).map(|j| buf[j] * x[j]).sum::<f32>();
14757            }
14758        };
14759        let iters = 5;
14760        let mut a = vec![0f32; rows];
14761        slow(&mut a); // warm
14762        let t = Instant::now();
14763        for _ in 0..iters {
14764            slow(&mut a);
14765        }
14766        let slow_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
14767
14768        let mut b = vec![0f32; rows];
14769        q1t_matvec(&bytes, &x, rows, cols, &mut b, None); // warm
14770        let t = Instant::now();
14771        for _ in 0..iters {
14772            q1t_matvec(&bytes, &x, rows, cols, &mut b, None);
14773        }
14774        let fast_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
14775
14776        for r in 0..rows {
14777            assert!((a[r] - b[r]).abs() < 1e-2, "mismatch row {r}");
14778        }
14779        println!(
14780            "q1t matvec {rows}x{cols} (1 thread): div-decode {slow_ms:.2} ms  fused-LUT {fast_ms:.2} ms  => {:.2}x",
14781            slow_ms / fast_ms
14782        );
14783    }
14784}
14785
14786#[cfg(test)]
14787mod gemm_bench {
14788    /// `cargo test -p cortiq-engine --release q4tp_matmat_throughput -- --ignored --nocapture`
14789    /// Times the batched q4tp GEMM at the shapes the image DiT runs
14790    /// (b=296 tokens, 2304 -> 9216), on synthetic bytes: no model, no
14791    /// mmap, no thermal drift over minutes — a kernel change shows up
14792    /// here in seconds where a full render hides it in noise.
14793    ///
14794    /// On macOS add `CMF_ACCEL=0`: this shape is over the 500k-cell mark
14795    /// where the matmat hands off to Accelerate's dequant sgemm, and
14796    /// without the opt-out both rows below measure the AMX, not the
14797    /// kernel under test.
14798    #[test]
14799    #[ignore]
14800    fn q4tp_matmat_throughput() {
14801        let _alt = super::Q4TP_ALT_TEST_LOCK
14802            .lock()
14803            .unwrap_or_else(|e| e.into_inner());
14804        // 296 is a prompt-encode batch; the image DiT runs 2085 at
14805        // 512x512, where the activation panel stops fitting L2 and the
14806        // loop's shape starts to matter more than its instructions.
14807        let b: usize = std::env::var("CMF_BENCH_B")
14808            .ok()
14809            .and_then(|v| v.parse().ok())
14810            .unwrap_or(296);
14811        let (rows, cols) = (9216usize, 2304usize);
14812        let (_, _, _) = (rows, cols, b);
14813        let total =
14814            cortiq_core::quant::expected_nbytes(cortiq_core::TensorDtype::Q4TiledP, &[rows, cols])
14815                .unwrap();
14816        // Random nibbles are fine, but the row params are f16 (lo, step)
14817        // of a geometric ladder: garbage there gives exp2 of a huge
14818        // exponent, the scales come back inf, and the whole bench times
14819        // NaN arithmetic instead of the kernel.
14820        let (params_off, codes_off, _) = cortiq_core::quant::q4tp_sections(rows, cols);
14821        let mut bytes: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
14822        let lo = cortiq_core::quant::f32_to_f16(-4.0);
14823        let step = cortiq_core::quant::f32_to_f16(0.1);
14824        for r in 0..rows {
14825            let o = params_off + r * 4;
14826            bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
14827            bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
14828        }
14829        let _ = codes_off;
14830        let xs: Vec<f32> = (0..b * cols)
14831            .map(|i| ((i % 97) as f32 - 48.0) / 48.0)
14832            .collect();
14833        let mut out = vec![0f32; b * rows];
14834        let pool = crate::pool::Pool::from_env();
14835        // A shared 48-core stand drifts ±25% run to run, which is wider
14836        // than any kernel change worth making. So: alternate the two
14837        // kernels inside one process and keep the BEST time for
14838        // each. Interleaving makes both see the same interference, and a
14839        // minimum is the one statistic another tenant cannot inflate.
14840        super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
14841        let reps: usize = std::env::var("CMF_BENCH_REPS")
14842            .ok()
14843            .and_then(|v| v.parse().ok())
14844            .unwrap_or(10);
14845        let mut best = [f64::MAX; 2];
14846        let mut sums = [0f32; 2];
14847        for _ in 0..reps {
14848            for (k, w) in [(0usize, 1u8), (1usize, 2u8)] {
14849                super::Q4TP_ALT.store(w, std::sync::atomic::Ordering::Relaxed);
14850                let t = std::time::Instant::now();
14851                super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
14852                best[k] = best[k].min(t.elapsed().as_secs_f64());
14853                sums[k] = out.iter().take(64).sum::<f32>();
14854            }
14855        }
14856        let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
14857        for (k, name) in ["previous", "tuned   "].iter().enumerate() {
14858            println!(
14859                "q4tp matmat {rows}x{cols} b={b} {name}: {:.1} ms  {:.1} GFLOP/s  (checksum {:.3})",
14860                best[k] * 1e3,
14861                flops / best[k] / 1e9,
14862                sums[k]
14863            );
14864        }
14865        assert!(
14866            (sums[0] - sums[1]).abs() < 1e-2,
14867            "the tuned kernel changed the result: {} vs {}",
14868            sums[0],
14869            sums[1]
14870        );
14871    }
14872
14873    /// The blocked kernel must agree with the per-column path exactly —
14874    /// same weights, same activation split, only a different instruction
14875    /// mix. Shapes are chosen to hit the awkward cases: a column count
14876    /// that leaves an odd group (the 512-bit kernel does two at a time),
14877    /// and a batch that does not divide by four.
14878    #[test]
14879    fn q4tp_matmat_blocked_matches_scalar() {
14880        use std::sync::atomic::Ordering::Relaxed;
14881        let _alt = super::Q4TP_ALT_TEST_LOCK
14882            .lock()
14883            .unwrap_or_else(|e| e.into_inner());
14884        // The last shape carries the image DiT's column count — 2304, so
14885        // 72 groups of accumulation, which is where a reordered sum can
14886        // actually drift — and runs through the thread pool, since the
14887        // blocked path splits rows across workers. Its row count stays
14888        // under 500k cells on purpose: above that, macOS diverts the whole
14889        // matmat to the Accelerate/AMX dequant sgemm and neither kernel
14890        // here would run.
14891        for &(rows, cols, b) in &[
14892            (64usize, 128usize, 7usize),
14893            (33, 96, 4),
14894            (16, 256, 9),
14895            (192, 2304, 37),
14896        ] {
14897            let total = cortiq_core::quant::expected_nbytes(
14898                cortiq_core::TensorDtype::Q4TiledP,
14899                &[rows, cols],
14900            )
14901            .unwrap();
14902            let (params_off, _, _) = cortiq_core::quant::q4tp_sections(rows, cols);
14903            let mut bytes: Vec<u8> = (0..total).map(|i| (i * 61 % 251) as u8).collect();
14904            let lo = cortiq_core::quant::f32_to_f16(-4.0);
14905            let step = cortiq_core::quant::f32_to_f16(0.1);
14906            for r in 0..rows {
14907                let o = params_off + r * 4;
14908                bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
14909                bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
14910            }
14911            let xs: Vec<f32> = (0..b * cols)
14912                .map(|i| ((i % 89) as f32 - 44.0) / 44.0)
14913                .collect();
14914            let mut got = vec![0f32; b * rows];
14915            let mut want = vec![0f32; b * rows];
14916            let gpr = cols / 32;
14917            let view = super::Q4tpView::new(&bytes, rows, cols);
14918            let pool = crate::pool::Pool::from_env();
14919            super::Q4TP_ALT.store(2, Relaxed);
14920            super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut got, pool.as_deref());
14921            super::Q4TP_ALT.store(1, Relaxed);
14922            super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut want, pool.as_deref());
14923            super::Q4TP_ALT.store(0, Relaxed);
14924            // Measured against the output's scale, not cell by cell: a
14925            // dot product of 2304 terms lands near zero wherever the row
14926            // and the activation nearly cancel, and there a per-cell
14927            // ratio reports 1e-3 for an absolute error of 5e-6 — f32's
14928            // own rounding, reordered. What must stay small is the error
14929            // relative to what the layer actually outputs.
14930            let scale = want.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
14931            let (mut worst, mut at) = (0f32, 0usize);
14932            for (i, (g, w)) in got.iter().zip(&want).enumerate() {
14933                if (g - w).abs() > worst {
14934                    worst = (g - w).abs();
14935                    at = i;
14936                }
14937            }
14938            assert!(
14939                worst <= 1e-4 * scale,
14940                "{rows}x{cols} b={b}: blocked and scalar disagree by {worst:.3e} \
14941                 (scale {scale:.3e}) at cell {at}: {} vs {}",
14942                got[at],
14943                want[at]
14944            );
14945
14946            // "Same speed, no quality loss" is a claim about which answer
14947            // is RIGHT, not about which two agree. Both paths sum the same
14948            // 2304 products in different orders, so f64 decides: the
14949            // blocked kernel keeps sixteen partial sums and folds them at
14950            // the end, which is a shallower addition tree than the
14951            // per-column path's running scalar, and it must not be worse.
14952            let (mut e_blocked, mut e_scalar) = (0f64, 0f64);
14953            for bi in 0..b {
14954                let act = super::split_act(&xs[bi * cols..(bi + 1) * cols]);
14955                for r in 0..rows {
14956                    let mut sc = vec![0f32; gpr];
14957                    view.scales_into(r, gpr, &mut sc);
14958                    let mut exact = 0f64;
14959                    for j in 0..cols {
14960                        let (w, sq) = super::q4tp_outlier(view.nib, r, gpr, j, &sc);
14961                        exact += w as f64 * sq as f64 * act.xq[j] as f64;
14962                    }
14963                    exact *= act.sx as f64;
14964                    for &(j, xv) in &act.outliers {
14965                        let (w, sq) = super::q4tp_outlier(view.nib, r, gpr, j, &sc);
14966                        exact += w as f64 * sq as f64 * xv as f64;
14967                    }
14968                    let i = bi * rows + r;
14969                    e_blocked = e_blocked.max((got[i] as f64 - exact).abs());
14970                    e_scalar = e_scalar.max((want[i] as f64 - exact).abs());
14971                }
14972            }
14973            println!(
14974                "{rows}x{cols} b={b}: worst error vs f64 — blocked {e_blocked:.3e}, \
14975                 per-column {e_scalar:.3e}"
14976            );
14977            // An absolute bar, not a race between the two: at these
14978            // magnitudes both sit in f32's last bits, and on a small shape
14979            // whichever one happens to round the unluckiest cell "wins" by
14980            // a factor the next seed reverses.
14981            assert!(
14982                e_blocked <= 1e-5 * scale as f64 && e_scalar <= 1e-5 * scale as f64,
14983                "{rows}x{cols} b={b}: error against f64 too large — blocked \
14984                 {e_blocked:.3e}, per-column {e_scalar:.3e}, scale {scale:.3e}"
14985            );
14986        }
14987    }
14988
14989    /// What the row-exact contract costs a speculative-verify panel on the
14990    /// host: the same batch through the fast arms, through the row-exact
14991    /// arms, and as one matvec per token (the other way to be exact).
14992    /// Arms alternate inside one process and keep their best time.
14993    /// `CMF_GPU=0 cargo test -p cortiq-engine --release q4tp_matmat_row_exact_cost -- --ignored --nocapture`
14994    #[test]
14995    #[ignore]
14996    fn q4tp_matmat_row_exact_cost() {
14997        let _alt = super::Q4TP_ALT_TEST_LOCK
14998            .lock()
14999            .unwrap_or_else(|e| e.into_inner());
15000        let pool = crate::pool::Pool::from_env();
15001        let reps: usize = std::env::var("CMF_BENCH_REPS")
15002            .ok()
15003            .and_then(|v| v.parse().ok())
15004            .unwrap_or(30);
15005        for &(rows, cols) in &[(2048usize, 4096usize), (4096, 2048), (4096, 4096)] {
15006            let total = cortiq_core::quant::expected_nbytes(
15007                cortiq_core::TensorDtype::Q4TiledP,
15008                &[rows, cols],
15009            )
15010            .unwrap();
15011            let (params_off, _, _) = cortiq_core::quant::q4tp_sections(rows, cols);
15012            let mut bytes: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
15013            let lo = cortiq_core::quant::f32_to_f16(-4.0);
15014            let step = cortiq_core::quant::f32_to_f16(0.1);
15015            for r in 0..rows {
15016                let o = params_off + r * 4;
15017                bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
15018                bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
15019            }
15020            for &b in &[2usize, 4, 5, 8] {
15021                let xs: Vec<f32> = (0..b * cols)
15022                    .map(|i| ((i % 97) as f32 - 48.0) / 48.0)
15023                    .collect();
15024                let mut out = vec![0f32; b * rows];
15025                let mut best = [f64::MAX; 3];
15026                for _ in 0..reps {
15027                    for (k, best_k) in best.iter_mut().enumerate() {
15028                        let t = std::time::Instant::now();
15029                        match k {
15030                            0 | 1 => super::q4tp_matmat_with(
15031                                &bytes,
15032                                &xs,
15033                                b,
15034                                rows,
15035                                cols,
15036                                &mut out,
15037                                pool.as_deref(),
15038                                k == 1,
15039                            ),
15040                            _ => {
15041                                for (bi, o) in out.chunks_mut(rows).enumerate() {
15042                                    super::q4tp_matvec(
15043                                        &bytes,
15044                                        &xs[bi * cols..(bi + 1) * cols],
15045                                        rows,
15046                                        cols,
15047                                        o,
15048                                        pool.as_deref(),
15049                                    );
15050                                }
15051                            }
15052                        }
15053                        *best_k = best_k.min(t.elapsed().as_secs_f64());
15054                    }
15055                }
15056                println!(
15057                    "q4tp {rows}x{cols} b={b}: fast {:.3} ms, row-exact {:.3} ms, \
15058                     {b} matvecs {:.3} ms",
15059                    best[0] * 1e3,
15060                    best[1] * 1e3,
15061                    best[2] * 1e3
15062                );
15063            }
15064        }
15065    }
15066
15067    /// The q4t twin of the throughput bench, same shape and rules, so the
15068    /// two quantisations' batch kernels can be read against each other.
15069    /// `cargo test -p cortiq-engine --release q4t_matmat_throughput -- --ignored --nocapture`
15070    #[test]
15071    #[ignore]
15072    fn q4t_matmat_throughput() {
15073        let (rows, cols, b) = (9216usize, 2304usize, 296usize);
15074        let total =
15075            cortiq_core::quant::expected_nbytes(cortiq_core::TensorDtype::Q4Tiled, &[rows, cols])
15076                .unwrap();
15077        // q4t carries a per-group f16 scale in the tile's first two bytes;
15078        // random bytes there decode to inf and the bench would time NaNs.
15079        let mut bytes: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
15080        let sc = cortiq_core::quant::f32_to_f16(0.02);
15081        for t in bytes.chunks_mut(super::Q4_TILE) {
15082            t[..2].copy_from_slice(&sc.to_le_bytes());
15083        }
15084        let xs: Vec<f32> = (0..b * cols)
15085            .map(|i| ((i % 97) as f32 - 48.0) / 48.0)
15086            .collect();
15087        let mut out = vec![0f32; b * rows];
15088        let pool = crate::pool::Pool::from_env();
15089        super::q4t_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
15090        let reps: usize = std::env::var("CMF_BENCH_REPS")
15091            .ok()
15092            .and_then(|v| v.parse().ok())
15093            .unwrap_or(10);
15094        let mut best = f64::MAX;
15095        for _ in 0..reps {
15096            let t = std::time::Instant::now();
15097            super::q4t_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
15098            best = best.min(t.elapsed().as_secs_f64());
15099        }
15100        let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
15101        println!(
15102            "q4t matmat {rows}x{cols} b={b}: {:.1} ms  {:.1} GFLOP/s  (checksum {:.3})",
15103            best * 1e3,
15104            flops / best / 1e9,
15105            out.iter().take(64).sum::<f32>()
15106        );
15107    }
15108}