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}
54
55/// Load-time q8 repack gate (see `Mapped::repack`). OPT-IN
56/// (`CMF_REPACK=1`): the single-stream hypothesis LOST on Apple Silicon
57/// (M4, interleaved A/B: decode 101 vs 94 tok/s — four adjacent row
58/// streams per worker feed the prefetcher MORE memory-level parallelism
59/// than one); kept as an experiment flag for x86, where the tradeoff
60/// may land differently.
61fn repack_enabled() -> bool {
62    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
63    *ON.get_or_init(|| {
64        std::env::var("CMF_REPACK")
65            .map(|v| v == "1")
66            .unwrap_or(cfg!(target_os = "android"))
67    })
68}
69
70/// Interleave q8 rows for the decode kernel: group g holds rows
71/// 4g..4g+4 as [r0[c], r1[c], r2[c], r3[c]] per 16-byte chunk c. Only
72/// full groups are packed — tail rows keep reading the mmap layout.
73fn q8_repack(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
74    #[cfg(target_arch = "aarch64")]
75    let arch_ok = sdot_enabled();
76    #[cfg(not(target_arch = "aarch64"))]
77    let arch_ok = false;
78    if !arch_ok || !repack_enabled() || rows < 256 || cols % 16 != 0 {
79        return Vec::new();
80    }
81    q8_repack_layout(bytes, rows, cols)
82}
83
84/// The pure layout transform behind `q8_repack` (tested directly —
85/// the gate depends on arch and env).
86fn q8_repack_layout(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
87    let groups = rows / 4;
88    let mut rep = vec![0u8; groups * 4 * cols];
89    for g in 0..groups {
90        let dst = &mut rep[g * 4 * cols..(g + 1) * 4 * cols];
91        for c in 0..cols / 16 {
92            for lane in 0..4 {
93                let src = (g * 4 + lane) * cols + c * 16;
94                dst[c * 64 + lane * 16..c * 64 + lane * 16 + 16]
95                    .copy_from_slice(&bytes[src..src + 16]);
96            }
97        }
98    }
99    rep
100}
101
102/// Prefix-sum of vbit row payload offsets (absolute within the tensor
103/// bytes). `offsets[r]..offsets[r+1]` is row r's packed data.
104fn vbit_row_offsets(bytes: &[u8], rows: usize, cols: usize) -> Vec<usize> {
105    let ng = cols / GROUP_SIZE;
106    let bits = &bytes[..rows];
107    let mut offsets = Vec::with_capacity(rows + 1);
108    let mut off = rows + rows * ng * 2;
109    for r in 0..rows {
110        offsets.push(off);
111        off += (cols * bits[r] as usize).div_ceil(8);
112    }
113    offsets.push(off);
114    offsets
115}
116
117/// `CMF_X86_BLOCKED` / `CMF_GPU_LMHEAD` / `CMF_GPU_SPLIT`, read once. They
118/// used to be read from the environment on every large matvec and on every
119/// matmat in six places — microseconds each, but also a knob that could
120/// change under a running process, which is not a thing a kernel choice
121/// should be able to do mid-sequence.
122fn blocked_enabled() -> bool {
123    use std::sync::atomic::Ordering::Relaxed;
124    match BLOCKED_OVERRIDE.load(Relaxed) {
125        1 => false,
126        2 => true,
127        _ => {
128            static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
129            *ON.get_or_init(|| {
130                std::env::var("CMF_X86_BLOCKED")
131                    .map(|v| v != "0")
132                    .unwrap_or(true)
133            })
134        }
135    }
136}
137
138static BLOCKED_OVERRIDE: std::sync::atomic::AtomicU8 = std::sync::atomic::AtomicU8::new(0);
139
140/// Force the blocked GEMM on or off, ignoring the environment; `None`
141/// restores it. For tests that need to run BOTH paths and compare them:
142/// `blocked_enabled` caches its answer for the life of the process, which
143/// is right when the environment is the only input, but leaves a test that
144/// flips `CMF_X86_BLOCKED` between two calls comparing a path against
145/// itself — or against whatever a test running in parallel latched first.
146pub fn set_blocked_override(on: Option<bool>) {
147    let v = match on {
148        None => 0,
149        Some(false) => 1,
150        Some(true) => 2,
151    };
152    BLOCKED_OVERRIDE.store(v, std::sync::atomic::Ordering::Relaxed);
153}
154
155fn gpu_lmhead_enabled() -> bool {
156    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
157    *ON.get_or_init(|| {
158        std::env::var("CMF_GPU_LMHEAD")
159            .map(|v| v != "0")
160            .unwrap_or(true)
161    })
162}
163
164fn gpu_split_frac() -> f32 {
165    // MiMo's banked verification projects every row on the device. Plain
166    // decode must not quantize half of its O/head activation on the CPU.
167    if FULL_GPU_Q8.get() {
168        return 1.0;
169    }
170    static F: std::sync::OnceLock<f32> = std::sync::OnceLock::new();
171    *F.get_or_init(|| {
172        std::env::var("CMF_GPU_SPLIT")
173            .ok()
174            .and_then(|v| v.parse::<f32>().ok())
175            .unwrap_or(0.5)
176            .clamp(0.0, 1.0)
177    })
178}
179
180impl QTensor {
181    pub fn from_f32(data: Vec<f32>, rows: usize, cols: usize) -> Self {
182        debug_assert_eq!(data.len(), rows * cols);
183        Self::F32 { data, rows, cols }
184    }
185
186    /// Wrap a directory tensor without dequantizing the payload.
187    /// Falls back to dequantized f32 for dtypes without a fused kernel.
188    pub fn from_model(model: &Arc<CmfModel>, name: &str) -> Result<Self, String> {
189        // Indexed lookup: the linear directory scan made pipeline build
190        // O(N²) on MoE/skills files with thousands of tensors.
191        let idx = model
192            .tensor_index(name)
193            .ok_or_else(|| format!("tensor '{name}' not found in CMF directory"))?;
194        let entry = &model.tensors[idx];
195        if entry.shape.len() != 2 {
196            return Err(format!("QTensor::from_model needs 2-D, got '{name}'"));
197        }
198        let (rows, cols) = (entry.shape[0], entry.shape[1]);
199        let bytes = model.entry_bytes(entry);
200
201        match entry.dtype {
202            TensorDtype::Q8Row | TensorDtype::Q8_2f => {
203                let n = rows * cols;
204                let scales_off = n;
205                let row_scale: Vec<f32> = (0..rows)
206                    .map(|o| {
207                        f16_to_f32(u16::from_le_bytes([
208                            bytes[scales_off + o * 2],
209                            bytes[scales_off + o * 2 + 1],
210                        ]))
211                    })
212                    .collect();
213                let col_field: Vec<f32> = if entry.dtype == TensorDtype::Q8_2f {
214                    let col_off = n + rows * 2;
215                    (0..cols)
216                        .map(|i| {
217                            f16_to_f32(u16::from_le_bytes([
218                                bytes[col_off + i * 2],
219                                bytes[col_off + i * 2 + 1],
220                            ]))
221                        })
222                        .collect()
223                } else {
224                    Vec::new()
225                };
226                Ok(Self::Mapped {
227                    model: model.clone(),
228                    idx,
229                    dtype: entry.dtype,
230                    rows,
231                    cols,
232                    row_scale,
233                    col_field,
234                    vbit_offsets: Vec::new(),
235                    repack: q8_repack(bytes, rows, cols),
236                })
237            }
238            // vbit: fused kernel unpacks variable-bit rows from mmap.
239            TensorDtype::Vbit if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
240                model: model.clone(),
241                idx,
242                dtype: entry.dtype,
243                rows,
244                cols,
245                row_scale: Vec::new(),
246                col_field: Vec::new(),
247                vbit_offsets: vbit_row_offsets(bytes, rows, cols),
248                repack: Vec::new(),
249            }),
250            // vbit_ro (§4.2): the offset table comes straight from the
251            // file — no load-time prefix scan; kernels are shared with
252            // legacy vbit (they consume absolute offsets either way).
253            TensorDtype::VbitRo if cols % GROUP_SIZE == 0 => {
254                let (_, off_off, packed_off) = cortiq_core::quant::vbit_ro_sections(rows, cols);
255                let offsets: Vec<usize> = (0..=rows)
256                    .map(|r| packed_off + cortiq_core::quant::vbit_ro_offset(bytes, off_off, r))
257                    .collect();
258                Ok(Self::Mapped {
259                    model: model.clone(),
260                    idx,
261                    dtype: entry.dtype,
262                    rows,
263                    cols,
264                    row_scale: Vec::new(),
265                    col_field: Vec::new(),
266                    vbit_offsets: offsets,
267                    repack: Vec::new(),
268                })
269            }
270            // q4_block: fused kernel reads nibbles straight from mmap —
271            // a 14B q4 file no longer explodes into ×8 f32 RAM.
272            // q4_tiled (§4.3): interleaved [scale][nibbles] tiles — one
273            // sequential memory stream (measured ×1.66 ARM / ×1.13 AVX2
274            // at kernel level over the split layout).
275            TensorDtype::Q4Tiled if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
276                model: model.clone(),
277                idx,
278                dtype: entry.dtype,
279                rows,
280                cols,
281                row_scale: Vec::new(),
282                col_field: Vec::new(),
283                vbit_offsets: Vec::new(),
284                repack: Vec::new(),
285            }),
286            // q4tp (§4.10): nibbles from mmap, scale from the row ladder —
287            // 7.3% less file than q4t at the same 4-bit grid.
288            TensorDtype::Q4TiledP if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
289                model: model.clone(),
290                idx,
291                dtype: entry.dtype,
292                rows,
293                cols,
294                row_scale: Vec::new(),
295                col_field: Vec::new(),
296                vbit_offsets: Vec::new(),
297                repack: Vec::new(),
298            }),
299            // q2tp: 2-bit chunks from mmap, scale from the same row ladder.
300            TensorDtype::Q2TiledP if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
301                model: model.clone(),
302                idx,
303                dtype: entry.dtype,
304                rows,
305                cols,
306                row_scale: Vec::new(),
307                col_field: Vec::new(),
308                vbit_offsets: Vec::new(),
309                repack: Vec::new(),
310            }),
311            TensorDtype::Q4Block if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
312                model: model.clone(),
313                idx,
314                dtype: entry.dtype,
315                rows,
316                cols,
317                row_scale: Vec::new(),
318                col_field: Vec::new(),
319                vbit_offsets: Vec::new(),
320                repack: Vec::new(),
321            }),
322            // q1: binary sign-bit tiles from mmap (1-bit-trained models).
323            TensorDtype::Q1 if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
324                model: model.clone(),
325                idx,
326                dtype: entry.dtype,
327                rows,
328                cols,
329                row_scale: Vec::new(),
330                col_field: Vec::new(),
331                vbit_offsets: Vec::new(),
332                repack: Vec::new(),
333            }),
334            // q1t (ternary + outlier overlay): fused per-row dequant kernel
335            // reads straight from mmap — a 12B q1t stays ~its file size in
336            // RAM instead of dequantizing to ~48 GB of f32.
337            TensorDtype::Q1T if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
338                model: model.clone(),
339                idx,
340                dtype: entry.dtype,
341                rows,
342                cols,
343                row_scale: Vec::new(),
344                col_field: Vec::new(),
345                vbit_offsets: Vec::new(),
346                repack: Vec::new(),
347            }),
348            // No fused kernel yet → dequantize once (correct, more RAM).
349            _ => {
350                let mut data = vec![0.0f32; rows * cols];
351                cortiq_core::quant::dequant_tensor(entry, bytes, &mut data)?;
352                Ok(Self::from_f32(data, rows, cols))
353            }
354        }
355    }
356
357    /// q1-mapped tensor? (GPU gates: the q1 CPU kernel is
358    /// compute-bound, so offload pays at much smaller shapes than q8.)
359    pub(crate) fn is_q1(&self) -> bool {
360        matches!(
361            self,
362            Self::Mapped {
363                dtype: TensorDtype::Q1,
364                ..
365            }
366        )
367    }
368
369    /// Owned-f32 view (data, rows, cols) — the GDN a/b gate projections
370    /// arrive dequantized (force-f16 in the converter → F32 in RAM).
371    pub(crate) fn f32_parts(&self) -> Option<(&[f32], usize, usize)> {
372        match self {
373            Self::F32 { data, rows, cols } => Some((data, *rows, *cols)),
374            _ => None,
375        }
376    }
377
378    /// (directory idx, rows, cols) of a q1-mapped tensor — the
379    /// whole-block GPU path resolves offsets itself.
380    /// (idx, rows, cols) of a mapped tensor the whole-token GPU graph can drive
381    /// — Q1, Q1T or Q4-block (it resolves the offset and picks the kernel by
382    /// dtype). Q4-block lets a precise down_proj/lm_head stay on-device.
383    /// Named `q1_parts` for historical reasons.
384    pub(crate) fn q1_parts(&self) -> Option<(usize, usize, usize)> {
385        if self.has_prism_contract() {
386            return None;
387        }
388        match self {
389            #[cfg(target_os = "macos")]
390            Self::Mapped {
391                dtype: TensorDtype::Q1T,
392                ..
393            } if !crate::gpu::metal_q1t_enabled() => None,
394            Self::Mapped {
395                idx,
396                dtype:
397                    TensorDtype::Q1
398                    | TensorDtype::Q1T
399                    | TensorDtype::Q4Block
400                    | TensorDtype::Q4Tiled
401                    // Q2TiledP deliberately absent: the Metal graph has no
402                    // q2tp kernel, and advertising it here made the block
403                    // plan truncate mid-run at the first q2tp layer.
404                    | TensorDtype::Q4TiledP
405                    | TensorDtype::Q8Row
406                    | TensorDtype::Q8_2f,
407                rows,
408                cols,
409                ..
410            } => Some((*idx, *rows, *cols)),
411            _ => None,
412        }
413    }
414
415    /// `(directory idx, rows, cols)` for the native Metal token graph.  The
416    /// historical q1 graph gate intentionally refuses every Prism tensor so
417    /// an untransformed q2 payload cannot slip into the resident path.  The
418    /// Metal2 graph is descriptor-aware and admits only the production
419    /// q2tp-affine forward targets; ordinary q1/q4 callers retain the old
420    /// `q1_parts` behaviour.
421    #[cfg(target_os = "macos")]
422    pub(crate) fn metal_graph_parts(&self) -> Option<(usize, usize, usize)> {
423        if let Some((model, idx, kind, _)) = self.graph_weight_descriptor() {
424            let name = &model.tensors[idx].name;
425            let forward = kind == 9 && crate::prism::is_forward_weight(model, name);
426            let affine = kind == 9 && crate::prism::is_affine_target(model, name);
427            if forward && affine {
428                let e = model.tensors.get(idx)?;
429                return Some((idx, *e.shape.first()?, *e.shape.get(1)?));
430            }
431        }
432        self.q1_parts()
433    }
434
435    /// (directory idx, rows, cols) of a q4_tiled mapped tensor. The
436    /// chunk-prefill graph takes it in the same 4-tuple slot as
437    /// `q8_row_parts` with an EMPTY row_scale — q4t carries its scales
438    /// inside the 18-byte tiles, and the empty slice is what tells the
439    /// encoder to reach for the q4t kernels.
440    pub(crate) fn q4t_parts(&self) -> Option<(usize, usize, usize)> {
441        if self.has_prism_contract() {
442            return None;
443        }
444        match self {
445            Self::Mapped {
446                idx,
447                dtype: TensorDtype::Q4Tiled,
448                rows,
449                cols,
450                ..
451            } => Some((*idx, *rows, *cols)),
452            _ => None,
453        }
454    }
455
456    /// (directory idx, rows, cols) of a q4tp mapped tensor. Same empty-scale
457    /// slot as `q4t_parts` in the chunk graph — the encoder tells the two
458    /// apart by the tensor's dtype, not by the slot.
459    pub(crate) fn q4tp_parts(&self) -> Option<(usize, usize, usize)> {
460        if self.has_prism_contract() {
461            return None;
462        }
463        match self {
464            Self::Mapped {
465                idx,
466                dtype: TensorDtype::Q4TiledP,
467                rows,
468                cols,
469                ..
470            } => Some((*idx, *rows, *cols)),
471            _ => None,
472        }
473    }
474
475    /// (directory idx, rows, cols, row_scale) of a plain q8_row mapped
476    /// tensor — the chunk-prefill GPU graph resolves offsets itself.
477    /// q8_2f is excluded on purpose: its column field would need a
478    /// prescale stage on the device.
479    pub(crate) fn q8_row_parts(&self) -> Option<(usize, usize, usize, &[f32])> {
480        if self.has_prism_contract() {
481            return None;
482        }
483        match self {
484            Self::Mapped {
485                idx,
486                dtype: TensorDtype::Q8Row,
487                rows,
488                cols,
489                row_scale,
490                col_field,
491                ..
492            } if col_field.is_empty() => Some((*idx, *rows, *cols, row_scale)),
493            _ => None,
494        }
495    }
496
497    /// (directory idx, rows, cols, row_scale, col_field) of a q8_2f mapped
498    /// tensor — the chunk-prefill graph stages x·col on the device first.
499    pub(crate) fn q8_2f_parts(&self) -> Option<(usize, usize, usize, &[f32], &[f32])> {
500        if self.has_prism_contract() {
501            return None;
502        }
503        match self {
504            Self::Mapped {
505                idx,
506                dtype: TensorDtype::Q8_2f,
507                rows,
508                cols,
509                row_scale,
510                col_field,
511                ..
512            } if col_field.len() == *cols => Some((*idx, *rows, *cols, row_scale, col_field)),
513            _ => None,
514        }
515    }
516
517    /// The layout this tensor is stored in, when it is mapped from a model.
518    /// The frames branch on it — a q2tp gate against a q4tp down is a real
519    /// combination in the 2-bit profile and needs a different kernel.
520    pub fn model_dtype(&self) -> Option<cortiq_core::TensorDtype> {
521        match self {
522            Self::Mapped { dtype, .. } => Some(*dtype),
523            _ => None,
524        }
525    }
526
527    /// The tensor's index in the model directory, when it is mapped from one.
528    /// The GPU frames bind by index rather than by name — a name lookup per
529    /// layer per token is not free, and the index is what the device cache is
530    /// keyed on anyway.
531    pub fn model_idx(&self) -> Option<usize> {
532        match self {
533            Self::Mapped { idx, .. } => Some(*idx),
534            _ => None,
535        }
536    }
537
538    /// The model this tensor is mapped from, when it is mapped at all. The
539    /// GPU frames need the container to reach the bytes; a QTensor already
540    /// holds it, and threading a second handle down every call site to say
541    /// the same thing invites the two to disagree.
542    pub fn model_arc(&self) -> Option<std::sync::Arc<cortiq_core::CmfModel>> {
543        match self {
544            Self::Mapped { model, .. } => Some(model.clone()),
545            _ => None,
546        }
547    }
548
549    /// Whether this mapped tensor belongs to the Prism/Bonsai transform
550    /// contract.  Device graphs do not carry the descriptor, so callers use
551    /// this conservative predicate to stay on the descriptor-aware CPU path
552    /// instead of silently executing an unrotated matrix.
553    pub(crate) fn has_prism_contract(&self) -> bool {
554        matches!(self, Self::Mapped { model, .. } if crate::prism::has_contract(model))
555    }
556
557    pub fn rows(&self) -> usize {
558        match self {
559            Self::F32 { rows, .. } | Self::Mapped { rows, .. } => *rows,
560        }
561    }
562
563    /// Mapped q4t handle (model + directory index) — the fused GPU FFN
564    /// needs the raw file coordinates of its three projections.
565    pub(crate) fn mapped_q4t(&self) -> Option<(&Arc<CmfModel>, usize)> {
566        if self.has_prism_contract() {
567            return None;
568        }
569        match self {
570            Self::Mapped {
571                model,
572                idx,
573                dtype: TensorDtype::Q4Tiled,
574                ..
575            } => Some((model, *idx)),
576            _ => None,
577        }
578    }
579
580    /// Same slot as `mapped_q4t` for a q4tp tensor — the fused DiT FFN picks
581    /// its kernels by which of the two answers.
582    pub fn mapped_q4tp(&self) -> Option<(&Arc<CmfModel>, usize)> {
583        if self.has_prism_contract() {
584            return None;
585        }
586        match self {
587            Self::Mapped {
588                model,
589                idx,
590                dtype: TensorDtype::Q4TiledP,
591                ..
592            } => Some((model, *idx)),
593            _ => None,
594        }
595    }
596
597    /// (model, tensor idx) for a mapped weight in ANY codec the fused device
598    /// paths can run — four-bit tiled or either int8 layout.
599    ///
600    /// The fused DiT chains asked for `mapped_q4tp` by name, so an eight-bit
601    /// container never reached them and rendered through per-op GEMMs even
602    /// after those kernels learned its codec. The gate is what the codec has
603    /// a device GEMM for, not which codec it is.
604    pub fn mapped_device_gemm(&self) -> Option<(&Arc<CmfModel>, usize)> {
605        if self.has_prism_contract() {
606            return None;
607        }
608        match self {
609            Self::Mapped {
610                model,
611                idx,
612                dtype: TensorDtype::Q4TiledP | TensorDtype::Q8Row | TensorDtype::Q8_2f,
613                ..
614            } => Some((model, *idx)),
615            _ => None,
616        }
617    }
618
619    /// (model, tensor idx) for a q8_2f mapped weight — the eight-bit slot of
620    /// the MoE graph ladder (int8 body, f16 row scales, f16 column field).
621    pub fn mapped_q8_2f(&self) -> Option<(&Arc<CmfModel>, usize)> {
622        if self.has_prism_contract() {
623            return None;
624        }
625        match self {
626            Self::Mapped {
627                model,
628                idx,
629                dtype: TensorDtype::Q8_2f,
630                ..
631            } => Some((model, *idx)),
632            _ => None,
633        }
634    }
635
636    /// (model, tensor idx) for a q2tp mapped weight — the 2-bit twin of
637    /// `mapped_q4tp`, used by the mixed MoE profile.
638    pub fn mapped_q2tp(&self) -> Option<(&Arc<CmfModel>, usize)> {
639        if self.has_prism_contract() {
640            return None;
641        }
642        match self {
643            Self::Mapped {
644                model,
645                idx,
646                dtype: TensorDtype::Q2TiledP,
647                ..
648            } => Some((model, *idx)),
649            _ => None,
650        }
651    }
652
653    pub fn cols(&self) -> usize {
654        match self {
655            Self::F32 { cols, .. } | Self::Mapped { cols, .. } => *cols,
656        }
657    }
658
659    /// (model, tensor idx) for a q1 mapped weight — the wgpu token graph
660    /// keys its resident VRAM cache by idx. None for any other dtype/kind.
661    pub fn mapped_q1(&self) -> Option<(&std::sync::Arc<CmfModel>, usize)> {
662        if self.has_prism_contract() {
663            return None;
664        }
665        match self {
666            Self::Mapped {
667                model,
668                idx,
669                dtype: TensorDtype::Q1,
670                ..
671            } => Some((model, *idx)),
672            _ => None,
673        }
674    }
675
676    /// (model, idx, kind, row_scale) for a graph-capable mapped weight.
677    /// kind: 0=q8_row (per-row scales), 1=q1, 2=q4_block, 3=q1t
678    /// (tile-embedded, no rs), 5=q4_tiled, 6=q4tp, 7=q8_2f (both scale
679    /// planes live inside the tensor). None only for `vbit`.
680    ///
681    /// The old comment here claimed q4_block was unhandled while the arm
682    /// right below mapped it, and it named q8_2f as unhandled after that
683    /// stopped being true — a stale comment on this function is how a
684    /// model silently loses the graph, so it is worth keeping honest.
685    pub fn graph_weight(&self) -> Option<(&std::sync::Arc<CmfModel>, usize, u8, &[f32])> {
686        if self.has_prism_contract() {
687            return None;
688        }
689        self.graph_weight_descriptor()
690    }
691
692    /// Descriptor-aware graph handle used only by the Prism token graph.
693    /// Ordinary graph callers continue to use [`graph_weight`] and therefore
694    /// remain fail-closed until they provide the same explicit transform
695    /// contract.
696    pub(crate) fn graph_weight_descriptor(
697        &self,
698    ) -> Option<(&std::sync::Arc<CmfModel>, usize, u8, &[f32])> {
699        match self {
700            Self::Mapped {
701                model,
702                idx,
703                dtype: TensorDtype::Q8Row,
704                row_scale,
705                ..
706            } => Some((model, *idx, 0, row_scale.as_slice())),
707            Self::Mapped {
708                model,
709                idx,
710                dtype: TensorDtype::Q1,
711                ..
712            } => Some((model, *idx, 1, &[])),
713            // Q4Tiled is kind 5, NOT 2: both carried 2 historically, and
714            // the wgpu token graph fed 18B interleaved tiles to the
715            // split-layout q4b kernel — garbage output on q4t models
716            // (caught by an end-to-end answer check on real Vulkan).
717            Self::Mapped {
718                model,
719                idx,
720                dtype: TensorDtype::Q4Tiled,
721                ..
722            } => Some((model, *idx, 5, &[])),
723            // Kind 6, not 5: q4tp's nibble stride and scale planes differ,
724            // and feeding them to the q4t kernel is exactly the mistake that
725            // produced garbage when Q4Tiled shared kind 2 with Q4Block.
726            Self::Mapped {
727                model,
728                idx,
729                dtype: TensorDtype::Q4TiledP,
730                ..
731            } => Some((model, *idx, 6, &[])),
732            Self::Mapped {
733                model,
734                idx,
735                dtype: TensorDtype::Q4Block,
736                ..
737            } => Some((model, *idx, 2, &[])),
738            // q8_2f carries BOTH scale planes after the int8 body (rows
739            // f16, then cols f16), so the graph takes the whole tensor
740            // and the kernel reads them where they lie — no host-side
741            // prescale, which is what the per-op path does instead.
742            Self::Mapped {
743                model,
744                idx,
745                dtype: TensorDtype::Q8_2f,
746                ..
747            } => Some((model, *idx, 7, &[])),
748            Self::Mapped {
749                model,
750                idx,
751                dtype: TensorDtype::Q1T,
752                ..
753            } => Some((model, *idx, 3, &[])),
754            // Kind 9: the 2-bit plane on the q4tp ladder (dense FFN gate/up
755            // of the q2tp profile). Its own kernel — 8 bytes a group where
756            // q4tp has 16, and rung 0 is the exact zero.
757            Self::Mapped {
758                model,
759                idx,
760                dtype: TensorDtype::Q2TiledP,
761                ..
762            } => Some((model, *idx, 9, &[])),
763            _ => None,
764        }
765    }
766
767    /// Dense f32 view — only for owned tensors. Masked/sparse execution
768    /// paths require it; quantized weights don't support masks yet.
769    pub fn as_f32(&self) -> Option<&[f32]> {
770        match self {
771            Self::F32 { data, .. } => Some(data),
772            Self::Mapped { .. } => None,
773        }
774    }
775
776    fn quant_bytes(&self) -> &[u8] {
777        match self {
778            Self::Mapped { model, idx, .. } => model.entry_bytes(&model.tensors[*idx]),
779            Self::F32 { .. } => unreachable!("quant_bytes on F32"),
780        }
781    }
782
783    /// Dequantize one row into `dst` (embedding lookup).
784    pub fn row_f32(&self, r: usize, dst: &mut [f32]) {
785        let cols = self.cols();
786        debug_assert_eq!(dst.len(), cols);
787        match self {
788            Self::F32 { data, .. } => dst.copy_from_slice(&data[r * cols..(r + 1) * cols]),
789            Self::Mapped {
790                model,
791                idx,
792                dtype,
793                row_scale,
794                col_field,
795                vbit_offsets,
796                ..
797            } => {
798                if *dtype == TensorDtype::Q4Tiled {
799                    let bytes = self.quant_bytes();
800                    let gpr = cols / GROUP_SIZE;
801                    for gi in 0..gpr {
802                        let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
803                        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
804                        for (k, &b) in tile[2..].iter().enumerate() {
805                            dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
806                            dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
807                        }
808                    }
809                    if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
810                        crate::prism::inverse_embedding(model, dst);
811                    }
812                    return;
813                }
814                if *dtype == TensorDtype::Q4TiledP {
815                    let bytes = self.quant_bytes();
816                    let gpr = cols / GROUP_SIZE;
817                    let v = Q4tpView::new(bytes, self.rows(), cols);
818                    let mut sc = vec![0f32; gpr];
819                    v.scales_into(r, gpr, &mut sc);
820                    for gi in 0..gpr {
821                        let tile = &v.nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
822                        let s = sc[gi];
823                        for (k, &b) in tile.iter().enumerate() {
824                            dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
825                            dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
826                        }
827                    }
828                    if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
829                        crate::prism::inverse_embedding(model, dst);
830                    }
831                    return;
832                }
833                if *dtype == TensorDtype::Q2TiledP {
834                    let bytes = self.quant_bytes();
835                    let gpr = cols / GROUP_SIZE;
836                    let v = Q4tpView::new_q2(bytes, self.rows(), cols);
837                    let mut sc = vec![0f32; gpr];
838                    v.scales_into(r, gpr, &mut sc);
839                    for gi in 0..gpr {
840                        let ch =
841                            &v.nib[(r * gpr + gi) * Q2TP_CHUNK..(r * gpr + gi + 1) * Q2TP_CHUNK];
842                        let s = sc[gi];
843                        for (k, &b) in ch.iter().enumerate() {
844                            for j in 0..4 {
845                                let center = if crate::prism::is_affine_target(
846                                    model,
847                                    &model.tensors[*idx].name,
848                                ) {
849                                    1.0
850                                } else {
851                                    1.5
852                                };
853                                dst[gi * GROUP_SIZE + k * 4 + j] =
854                                    (((b >> (2 * j)) & 3) as f32 - center) * s;
855                            }
856                        }
857                    }
858                    if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
859                        crate::prism::inverse_embedding(model, dst);
860                    }
861                    return;
862                }
863                if *dtype == TensorDtype::Q4Block {
864                    let (packed, scales) = q4_split(self.quant_bytes(), self.rows(), cols);
865                    let gpr = cols / GROUP_SIZE;
866                    for gi in 0..gpr {
867                        let g = r * gpr + gi;
868                        let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
869                        for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
870                            dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
871                            dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
872                        }
873                    }
874                    if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
875                        crate::prism::inverse_embedding(model, dst);
876                    }
877                    return;
878                }
879                if *dtype == TensorDtype::Q1 {
880                    let bytes = self.quant_bytes();
881                    let gpr = cols / GROUP_SIZE;
882                    for gi in 0..gpr {
883                        let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
884                        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
885                        for (j, &b) in tile[2..].iter().enumerate() {
886                            for k in 0..8 {
887                                dst[gi * GROUP_SIZE + j * 8 + k] =
888                                    (((b >> k) & 1) as f32 * 2.0 - 1.0) * s;
889                            }
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::Q1T {
898                    let bytes = self.quant_bytes();
899                    let gpr = cols / GROUP_SIZE;
900                    let base_len = self.rows() * gpr * cortiq_core::quant::Q1T_TILE;
901                    for gi in 0..gpr {
902                        let off = (r * gpr + gi) * cortiq_core::quant::Q1T_TILE;
903                        let s = cortiq_core::quant::f16_to_f32(u16::from_le_bytes([
904                            bytes[off],
905                            bytes[off + 1],
906                        ]));
907                        let codes = &bytes[off + 2..off + cortiq_core::quant::Q1T_TILE];
908                        for k in 0..GROUP_SIZE {
909                            dst[gi * GROUP_SIZE + k] = match cortiq_core::quant::q1t_code(codes, k)
910                            {
911                                1 => s,
912                                2 => -s,
913                                _ => 0.0,
914                            };
915                        }
916                    }
917                    // Overlay
918                    let rows = self.rows();
919                    let entries = base_len + (rows + 1) * 4;
920                    if entries <= bytes.len() {
921                        let ptrs = &bytes[base_len..base_len + (rows + 1) * 4];
922                        let r0 = u32::from_le_bytes([
923                            ptrs[r * 4],
924                            ptrs[r * 4 + 1],
925                            ptrs[r * 4 + 2],
926                            ptrs[r * 4 + 3],
927                        ]) as usize;
928                        let r1 = u32::from_le_bytes([
929                            ptrs[(r + 1) * 4],
930                            ptrs[(r + 1) * 4 + 1],
931                            ptrs[(r + 1) * 4 + 2],
932                            ptrs[(r + 1) * 4 + 3],
933                        ]) as usize;
934                        let off = entries + r0 * 4;
935                        for i in 0..r1 - r0 {
936                            let item = &bytes[off + i * 4..off + i * 4 + 4];
937                            let c = u16::from_le_bytes([item[0], item[1]]) as usize;
938                            let v = cortiq_core::quant::f16_to_f32(u16::from_le_bytes([
939                                item[2], item[3],
940                            ]));
941                            if c < cols {
942                                dst[c] = v;
943                            }
944                        }
945                    }
946                    if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
947                        crate::prism::inverse_embedding(model, dst);
948                    }
949                    return;
950                }
951                if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
952                    let bytes = self.quant_bytes();
953                    let rows = self.rows();
954                    let ng = cols / GROUP_SIZE;
955                    let bits = &bytes[..rows];
956                    let sc_off = rows;
957                    // Precomputed at load — embedding lookup used to scan
958                    // the bit-widths of every preceding row (O(token_id)).
959                    let off = vbit_offsets[r];
960                    let b = bits[r] as usize;
961                    let l = ((1usize << (b - 1)) - 1) as f32;
962                    let data = &bytes[off..];
963                    let (mut acc, mut nbits, mut byte_idx) = (0u64, 0usize, 0usize);
964                    for (i, d) in dst.iter_mut().enumerate() {
965                        while nbits < b {
966                            acc = (acc << 8) | data[byte_idx] as u64;
967                            byte_idx += 1;
968                            nbits += 8;
969                        }
970                        let u = ((acc >> (nbits - b)) & ((1u64 << b) - 1)) as f32;
971                        nbits -= b;
972                        let so = (r * ng + i / GROUP_SIZE) * 2;
973                        let sv = f16_to_f32(u16::from_le_bytes([
974                            bytes[sc_off + so],
975                            bytes[sc_off + so + 1],
976                        ]));
977                        *d = (u - l) * sv;
978                    }
979                    if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
980                        crate::prism::inverse_embedding(model, dst);
981                    }
982                    return;
983                }
984                let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
985                let s = row_scale[r];
986                match dtype {
987                    TensorDtype::Q8Row => {
988                        for (d, &b) in dst.iter_mut().zip(q) {
989                            *d = (b as i8) as f32 * s;
990                        }
991                    }
992                    TensorDtype::Q8_2f => {
993                        for (i, (d, &b)) in dst.iter_mut().zip(q).enumerate() {
994                            *d = (b as i8) as f32 * s * col_field[i];
995                        }
996                    }
997                    _ => unreachable!(),
998                }
999                if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
1000                    crate::prism::inverse_embedding(model, dst);
1001                }
1002            }
1003        }
1004    }
1005
1006    /// Can this tensor's columns be read cheaply (for sparse down_proj)?
1007    /// True for F32/Q8Row/Q8_2f (per-row scale, direct strided access);
1008    /// false for group-packed q4/vbit (column access would unpack whole
1009    /// groups — sparse execution falls back to f32 for those).
1010    pub fn sparse_col_ok(&self) -> bool {
1011        match self {
1012            Self::F32 { .. } => true,
1013            Self::Mapped { dtype, .. } => {
1014                matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
1015            }
1016        }
1017    }
1018
1019    /// down_proj [hidden, inter]: accumulate `w · col(c)` into `out`
1020    /// [hidden] — reads ONLY column `c` (one neuron) from the mmap,
1021    /// no full-matrix dequant. `out[k] += w · down[k, c]`.
1022    pub fn add_col_scaled(&self, c: usize, w: f32, out: &mut [f32]) {
1023        let inter = self.cols();
1024        let hidden = self.rows();
1025        debug_assert_eq!(out.len(), hidden);
1026        match self {
1027            Self::F32 { data, .. } => {
1028                for (k, o) in out.iter_mut().enumerate() {
1029                    *o += w * data[k * inter + c];
1030                }
1031            }
1032            Self::Mapped {
1033                dtype,
1034                row_scale,
1035                col_field,
1036                ..
1037            } => {
1038                let q = self.quant_bytes();
1039                let colf = if *dtype == TensorDtype::Q8_2f {
1040                    col_field[c]
1041                } else {
1042                    1.0
1043                };
1044                let wc = w * colf;
1045                for (k, o) in out.iter_mut().enumerate() {
1046                    let b = q[k * inter + c] as i8 as f32;
1047                    *o += wc * b * row_scale[k];
1048                }
1049            }
1050        }
1051    }
1052
1053    /// Touch the head of row `r` so the DRAM latency of the next
1054    /// neuron's weights overlaps the current one's arithmetic.
1055    ///
1056    /// Scattered rows are what per-token sparsity reads, and a 2 KB
1057    /// stride is past what the hardware prefetcher follows: without this
1058    /// every row starts with a cold miss that nothing hides. One touch
1059    /// per 512 bytes is enough — the rest of the row is a sequential run
1060    /// the prefetcher does pick up.
1061    #[inline]
1062    pub fn prefetch_row(&self, r: usize) {
1063        let Self::Mapped { dtype, .. } = self else {
1064            return;
1065        };
1066        if !matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f) {
1067            return;
1068        }
1069        let cols = self.cols();
1070        let q = self.quant_bytes();
1071        let (a, b) = (r * cols, (r + 1) * cols);
1072        if b > q.len() {
1073            return;
1074        }
1075        let mut j = a;
1076        while j < b {
1077            unsafe { std::ptr::read_volatile(q.as_ptr().add(j)) };
1078            j += 512;
1079        }
1080    }
1081
1082    /// `out += w · row(r)` — the transposed twin of `add_col_scaled`.
1083    ///
1084    /// A neuron's `down` weights are a COLUMN of `[hidden, inter]`, and a
1085    /// column is strided: reading one costs a cache line per element, so
1086    /// per-neuron dynamic sparsity saves arithmetic and no bytes. Stored
1087    /// transposed (`down_proj.t.weight`, `[inter, hidden]`) the same
1088    /// weights are a contiguous ROW, and this accumulate reads exactly
1089    /// the neurons the token asked for.
1090    pub fn add_row_scaled(&self, r: usize, w: f32, out: &mut [f32], scratch: &mut [f32]) {
1091        let cols = self.cols();
1092        debug_assert_eq!(out.len(), cols);
1093        match self {
1094            Self::F32 { data, .. } => {
1095                let row = &data[r * cols..(r + 1) * cols];
1096                for (o, v) in out.iter_mut().zip(row) {
1097                    *o += w * v;
1098                }
1099            }
1100            Self::Mapped {
1101                dtype,
1102                row_scale,
1103                col_field,
1104                ..
1105            } => match dtype {
1106                TensorDtype::Q8Row => {
1107                    let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
1108                    let ws = w * row_scale[r];
1109                    let row: &[i8] =
1110                        unsafe { std::slice::from_raw_parts(q.as_ptr() as *const i8, q.len()) };
1111                    axpy_i8_f32(out, row, ws);
1112                }
1113                TensorDtype::Q8_2f => {
1114                    let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
1115                    let ws = w * row_scale[r];
1116                    for ((o, b), c) in out.iter_mut().zip(q).zip(col_field) {
1117                        *o += ws * c * (*b as i8 as f32);
1118                    }
1119                }
1120                _ => {
1121                    self.row_f32(r, scratch);
1122                    for (o, v) in out.iter_mut().zip(scratch.iter()) {
1123                        *o += w * v;
1124                    }
1125                }
1126            },
1127        }
1128    }
1129
1130    /// Dot of row `r` with `x` (gate/up active-neuron path). Reads only
1131    /// row `r` from the mmap — no full dequant. q4/vbit dequant the row
1132    /// into `scratch` first (rare for active-FFN weights).
1133    pub fn row_dot(&self, r: usize, x: &[f32], scratch: &mut [f32]) -> f32 {
1134        let cols = self.cols();
1135        match self {
1136            Self::F32 { data, .. } => {
1137                let row = &data[r * cols..(r + 1) * cols];
1138                row.iter().zip(x).map(|(w, v)| w * v).sum()
1139            }
1140            Self::Mapped {
1141                model,
1142                idx,
1143                dtype,
1144                row_scale,
1145                col_field,
1146                ..
1147            } => {
1148                let prism_forward =
1149                    crate::prism::is_forward_weight(model, &model.tensors[*idx].name);
1150                if prism_forward {
1151                    let transformed = crate::prism::forward(model, &x[..cols]);
1152                    let gpr = cols / GROUP_SIZE;
1153                    match dtype {
1154                        TensorDtype::Q2TiledP => {
1155                            let v = Q4tpView::new_q2(self.quant_bytes(), self.rows(), cols);
1156                            let mut sc = vec![0f32; gpr];
1157                            v.scales_into(r, gpr, &mut sc);
1158                            if crate::prism::is_affine_target(model, &model.tensors[*idx].name) {
1159                                return q2tp_affine_row_exact(v.nib, r, gpr, &transformed, &sc);
1160                            }
1161                            return q2tp_row_exact(v.nib, r, gpr, &transformed, &sc);
1162                        }
1163                        _ => {
1164                            self.row_f32(r, scratch);
1165                            return scratch.iter().zip(&transformed).map(|(w, v)| w * v).sum();
1166                        }
1167                    }
1168                }
1169                match dtype {
1170                    TensorDtype::Q8Row => {
1171                        let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
1172                        dot_i8_f32(q, x) * row_scale[r]
1173                    }
1174                    TensorDtype::Q8_2f => {
1175                        let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
1176                        dot_i8_col_f32(q, x, col_field) * row_scale[r]
1177                    }
1178                    _ => {
1179                        self.row_f32(r, scratch);
1180                        scratch.iter().zip(x).map(|(w, v)| w * v).sum()
1181                    }
1182                }
1183            }
1184        }
1185    }
1186
1187    /// `out = W · x` (row-major). F32 delegates to the historical
1188    /// bit-exact path; Mapped runs the fused int8 kernel.
1189    pub fn matvec(&self, x: &[f32], out: &mut [f32], pool: Option<&Pool>) {
1190        match self {
1191            // NOTE: `out.len()` DRIVES this arm — it computes that many rows,
1192            // and `x.len()` is the stride. A short `out` is legitimate here,
1193            // which is why the check below lives in the Mapped arm only.
1194            Self::F32 { data, .. } => {
1195                if !crate::f32_backend::matvec(data, x, out) {
1196                    matvec_rows(pool, data, x, out);
1197                }
1198            }
1199            Self::Mapped {
1200                model,
1201                idx,
1202                dtype,
1203                rows,
1204                cols,
1205                row_scale,
1206                col_field,
1207                vbit_offsets,
1208                repack,
1209            } => {
1210                let _ = (model, idx);
1211                // Every kernel below writes `rows` entries through a raw
1212                // pointer, so a short `out` is an out-of-bounds WRITE, not a
1213                // wrong answer: it scribbles on the allocator's metadata and
1214                // the process aborts much later, somewhere innocent
1215                // (`double free or corruption`, `corrupted double-linked
1216                // list`). The debug_assert two of the kernels carried is
1217                // compiled out of the release — exactly the build where it
1218                // matters. Fail here instead, while the caller is still on
1219                // the stack to be named.
1220                assert!(
1221                    out.len() >= *rows && x.len() >= *cols,
1222                    "matvec {rows}x{cols}: out {} (need {rows}), x {} (need {cols})",
1223                    out.len(),
1224                    x.len(),
1225                );
1226                let prism_forward =
1227                    crate::prism::is_forward_weight(model, &model.tensors[*idx].name);
1228                if *dtype == TensorDtype::Q2TiledP
1229                    && std::env::var("CMF_Q2TP_TRACE").as_deref() == Ok("1")
1230                {
1231                    use std::sync::atomic::{AtomicUsize, Ordering};
1232                    static N: AtomicUsize = AtomicUsize::new(0);
1233                    let n = N.fetch_add(1, Ordering::Relaxed);
1234                    if n < 128 {
1235                        eprintln!(
1236                            "q2tp-dispatch #{n} name={} prism={} rows={} cols={} gpu={} optin={} layer={}",
1237                            model.tensors[*idx].name,
1238                            prism_forward,
1239                            rows,
1240                            cols,
1241                            crate::gpu::enabled_here(),
1242                            crate::gpu::q2tp_gpu_opt_in(),
1243                            crate::gpu::cur_layer(),
1244                        );
1245                    }
1246                }
1247                // Prism stores every manifest-listed forward matrix in the
1248                // signed-Hadamard basis.  The q2tp WGSL path receives that
1249                // transformed vector and an explicit affine bit; codecs
1250                // without a descriptor-aware kernel remain on CPU below.
1251                if prism_forward {
1252                    let transformed = crate::prism::forward(model, &x[..*cols]);
1253                    match dtype {
1254                        TensorDtype::Q4Block => {
1255                            q4matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1256                        }
1257                        TensorDtype::Q4Tiled => {
1258                            q4t_matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1259                        }
1260                        TensorDtype::Q4TiledP => {
1261                            q4tp_matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1262                        }
1263                        TensorDtype::Q2TiledP => {
1264                            let affine =
1265                                crate::prism::is_affine_target(model, &model.tensors[*idx].name);
1266                            if *rows * *cols >= 8_388_608
1267                                && crate::gpu::enabled_here()
1268                                && crate::gpu::q2tp_gpu_opt_in()
1269                            {
1270                                let gpu_ok = if affine {
1271                                    crate::gpu::q2tp_affine_matvec(
1272                                        model,
1273                                        *idx,
1274                                        &transformed,
1275                                        *rows,
1276                                        *cols,
1277                                        out,
1278                                    )
1279                                } else {
1280                                    crate::gpu::q2tp_matvec(
1281                                        model,
1282                                        *idx,
1283                                        &transformed,
1284                                        *rows,
1285                                        *cols,
1286                                        out,
1287                                    )
1288                                };
1289                                if gpu_ok {
1290                                    return;
1291                                }
1292                            }
1293                            if affine {
1294                                q2tp_affine_matvec(
1295                                    self.quant_bytes(),
1296                                    &transformed,
1297                                    *rows,
1298                                    *cols,
1299                                    out,
1300                                    pool,
1301                                )
1302                            } else {
1303                                q2tp_matvec(
1304                                    self.quant_bytes(),
1305                                    &transformed,
1306                                    *rows,
1307                                    *cols,
1308                                    out,
1309                                    pool,
1310                                )
1311                            }
1312                        }
1313                        TensorDtype::Q1 => {
1314                            q1_matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1315                        }
1316                        TensorDtype::Q1T => {
1317                            q1t_matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1318                        }
1319                        TensorDtype::Vbit | TensorDtype::VbitRo => vbitmatvec(
1320                            self.quant_bytes(),
1321                            vbit_offsets,
1322                            &transformed,
1323                            *rows,
1324                            *cols,
1325                            out,
1326                            pool,
1327                        ),
1328                        TensorDtype::Q8Row | TensorDtype::Q8_2f => qmatvec(
1329                            self.quant_bytes(),
1330                            repack,
1331                            row_scale,
1332                            &transformed,
1333                            col_field,
1334                            *dtype,
1335                            *rows,
1336                            *cols,
1337                            out,
1338                            pool,
1339                        ),
1340                        _ => unreachable!("unsupported mapped Prism dtype {dtype:?}"),
1341                    }
1342                    return;
1343                }
1344                if *dtype == TensorDtype::Q4Block {
1345                    // GPU route (wgpu q4b kernel) for large q4_block matvecs —
1346                    // gives NVIDIA/AMD/Intel q4 models a GPU path. Probe keeps
1347                    // the winner; Metal returns false → the CPU kernel below.
1348                    if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1349                        let t0 = std::time::Instant::now();
1350                        match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
1351                            crate::gpu::ProbeArm::Gpu => {
1352                                if crate::gpu::q4b_matvec(model, *idx, x, *rows, *cols, out) {
1353                                    crate::gpu::probe_record(
1354                                        crate::gpu::OpClass::Matvec,
1355                                        true,
1356                                        t0.elapsed(),
1357                                    );
1358                                    return;
1359                                }
1360                            }
1361                            crate::gpu::ProbeArm::CpuTimed => {
1362                                q4matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1363                                crate::gpu::probe_record(
1364                                    crate::gpu::OpClass::Matvec,
1365                                    false,
1366                                    t0.elapsed(),
1367                                );
1368                                return;
1369                            }
1370                            crate::gpu::ProbeArm::Cpu => {}
1371                        }
1372                    }
1373                    q4matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1374                    return;
1375                }
1376                if *dtype == TensorDtype::Q4Tiled {
1377                    // GPU route for large q4t matvecs — the lm_head class,
1378                    // same shape as the q4tp arm below. The probe keeps the
1379                    // winner; a backend without the kernel refuses and the
1380                    // CPU path stays.
1381                    if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1382                        let t0 = std::time::Instant::now();
1383                        let cls = crate::gpu::matvec_class(*rows, *cols);
1384                        match crate::gpu::probe_arm(cls) {
1385                            crate::gpu::ProbeArm::Gpu => {
1386                                if crate::gpu::q4t_matvec(model, *idx, x, *rows, *cols, out) {
1387                                    crate::gpu::probe_record(cls, true, t0.elapsed());
1388                                    return;
1389                                }
1390                            }
1391                            crate::gpu::ProbeArm::CpuTimed => {
1392                                q4t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1393                                crate::gpu::probe_record(cls, false, t0.elapsed());
1394                                return;
1395                            }
1396                            crate::gpu::ProbeArm::Cpu => {}
1397                        }
1398                    }
1399                    q4t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1400                    return;
1401                }
1402                if *dtype == TensorDtype::Q4TiledP {
1403                    // GPU route for large q4tp matvecs — the lm_head class.
1404                    // On a q4tp checkpoint the head is the biggest single
1405                    // host matvec left in the decode step, and the batched
1406                    // kernel at b=1 already exists on both backends. Probe
1407                    // keeps the winner, same as q4_block above.
1408                    if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1409                        let t0 = std::time::Instant::now();
1410                        let cls = crate::gpu::matvec_class(*rows, *cols);
1411                        match crate::gpu::probe_arm(cls) {
1412                            crate::gpu::ProbeArm::Gpu => {
1413                                if crate::gpu::q4tp_matvec(model, *idx, x, *rows, *cols, out) {
1414                                    crate::gpu::probe_record(cls, true, t0.elapsed());
1415                                    return;
1416                                }
1417                            }
1418                            crate::gpu::ProbeArm::CpuTimed => {
1419                                q4tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1420                                crate::gpu::probe_record(cls, false, t0.elapsed());
1421                                return;
1422                            }
1423                            crate::gpu::ProbeArm::Cpu => {}
1424                        }
1425                    }
1426                    q4tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1427                    return;
1428                }
1429                if *dtype == TensorDtype::Q2TiledP {
1430                    q2tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1431                    return;
1432                }
1433                if *dtype == TensorDtype::Q1 {
1434                    // GPU route for large q1 matvecs (out_proj / lm_head
1435                    // class): the CPU q1 kernel is load-port-bound at
1436                    // ~4 GB/s/core, the GPU one is bandwidth-bound — the
1437                    // probe measures both arms and keeps the winner.
1438                    if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1439                        let t0 = std::time::Instant::now();
1440                        let arm = if crate::gpu::q1_force() {
1441                            crate::gpu::ProbeArm::Gpu
1442                        } else {
1443                            crate::gpu::probe_arm(crate::gpu::OpClass::Matvec)
1444                        };
1445                        match arm {
1446                            crate::gpu::ProbeArm::Gpu => {
1447                                if crate::gpu::q1_matvec(model, *idx, x, *rows, *cols, out) {
1448                                    crate::gpu::probe_record(
1449                                        crate::gpu::OpClass::Matvec,
1450                                        true,
1451                                        t0.elapsed(),
1452                                    );
1453                                    return;
1454                                }
1455                            }
1456                            crate::gpu::ProbeArm::CpuTimed => {
1457                                q1_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1458                                crate::gpu::probe_record(
1459                                    crate::gpu::OpClass::Matvec,
1460                                    false,
1461                                    t0.elapsed(),
1462                                );
1463                                return;
1464                            }
1465                            crate::gpu::ProbeArm::Cpu => {}
1466                        }
1467                    }
1468                    q1_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1469                    return;
1470                }
1471                if *dtype == TensorDtype::Q1T {
1472                    // GPU route for large q1t matvecs: the ternary BASE dot runs
1473                    // on the GPU (load-port-bound on CPU, like q1), then the
1474                    // sparse overlay is added on the CPU. Probe keeps the winner.
1475                    if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1476                        let t0 = std::time::Instant::now();
1477                        match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
1478                            crate::gpu::ProbeArm::Gpu => {
1479                                if crate::gpu::q1t_matvec(model, *idx, x, *rows, *cols, out) {
1480                                    q1t_add_overlay(self.quant_bytes(), x, *rows, *cols, out, pool);
1481                                    crate::gpu::probe_record(
1482                                        crate::gpu::OpClass::Matvec,
1483                                        true,
1484                                        t0.elapsed(),
1485                                    );
1486                                    return;
1487                                }
1488                            }
1489                            crate::gpu::ProbeArm::CpuTimed => {
1490                                q1t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1491                                crate::gpu::probe_record(
1492                                    crate::gpu::OpClass::Matvec,
1493                                    false,
1494                                    t0.elapsed(),
1495                                );
1496                                return;
1497                            }
1498                            crate::gpu::ProbeArm::Cpu => {}
1499                        }
1500                    }
1501                    q1t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1502                    return;
1503                }
1504                if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1505                    vbitmatvec(self.quant_bytes(), vbit_offsets, x, *rows, *cols, out, pool);
1506                    return;
1507                }
1508                let xs = prescale(x, col_field, *dtype);
1509                // D5: large q8 matrices (lm_head-class) — hybrid
1510                // CPU∥GPU: split the rows, both sides compute
1511                // SIMULTANEOUSLY (same math, shared prescale).
1512                // GPU share: CMF_GPU_SPLIT (0..1, default 0.5).
1513                if *rows >= crate::gpu::min_rows()
1514                    && matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
1515                    && gpu_lmhead_enabled()
1516                    && crate::gpu::enabled_here()
1517                {
1518                    // Runtime probe: alternate the hybrid against the
1519                    // pure-CPU matvec, keep whichever is faster HERE.
1520                    let t0 = std::time::Instant::now();
1521                    match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
1522                        crate::gpu::ProbeArm::Gpu => {}
1523                        crate::gpu::ProbeArm::CpuTimed => {
1524                            qmatvec(
1525                                self.quant_bytes(),
1526                                repack,
1527                                row_scale,
1528                                x,
1529                                col_field,
1530                                *dtype,
1531                                *rows,
1532                                *cols,
1533                                out,
1534                                pool,
1535                            );
1536                            crate::gpu::probe_record(
1537                                crate::gpu::OpClass::Matvec,
1538                                false,
1539                                t0.elapsed(),
1540                            );
1541                            return;
1542                        }
1543                        crate::gpu::ProbeArm::Cpu => {
1544                            qmatvec(
1545                                self.quant_bytes(),
1546                                repack,
1547                                row_scale,
1548                                x,
1549                                col_field,
1550                                *dtype,
1551                                *rows,
1552                                *cols,
1553                                out,
1554                                pool,
1555                            );
1556                            return;
1557                        }
1558                    }
1559                    let frac = gpu_split_frac();
1560                    let cpu_rows = ((*rows as f32) * (1.0 - frac)) as usize;
1561                    let (out_cpu, out_gpu) = out.split_at_mut(cpu_rows);
1562                    let bytes = self.quant_bytes();
1563                    let ok = std::thread::scope(|sc| {
1564                        let g = sc.spawn(|| {
1565                            crate::gpu::q8_matvec_range(
1566                                model,
1567                                *idx,
1568                                cpu_rows,
1569                                &row_scale[cpu_rows..],
1570                                &xs,
1571                                *rows - cpu_rows,
1572                                *cols,
1573                                out_gpu,
1574                            )
1575                        });
1576                        if cpu_rows > 0 {
1577                            // Repack prefix covers the full groups of the
1578                            // CPU half (the split starts at row 0).
1579                            let rep_cpu = if repack.is_empty() {
1580                                &[][..]
1581                            } else {
1582                                &repack[..(cpu_rows / 4) * 4 * *cols]
1583                            };
1584                            qmatvec(
1585                                &bytes[..cpu_rows * *cols],
1586                                rep_cpu,
1587                                &row_scale[..cpu_rows],
1588                                x,
1589                                col_field,
1590                                *dtype,
1591                                cpu_rows,
1592                                *cols,
1593                                out_cpu,
1594                                pool,
1595                            );
1596                        }
1597                        g.join().unwrap_or(false)
1598                    });
1599                    if ok {
1600                        crate::gpu::probe_record(crate::gpu::OpClass::Matvec, true, t0.elapsed());
1601                        return;
1602                    }
1603                    // GPU failed — CPU finishes its half (rows rebased —
1604                    // group offsets don't line up, mmap layout only).
1605                    qmatvec(
1606                        &bytes[cpu_rows * *cols..(*rows) * *cols],
1607                        &[],
1608                        &row_scale[cpu_rows..],
1609                        x,
1610                        col_field,
1611                        *dtype,
1612                        *rows - cpu_rows,
1613                        *cols,
1614                        out_gpu,
1615                        pool,
1616                    );
1617                    return;
1618                }
1619                qmatvec(
1620                    self.quant_bytes(),
1621                    repack,
1622                    row_scale,
1623                    x,
1624                    col_field,
1625                    *dtype,
1626                    *rows,
1627                    *cols,
1628                    out,
1629                    pool,
1630                );
1631            }
1632        }
1633    }
1634
1635    /// Fused two-input matvec (MTP verify pair): weights streamed once.
1636    pub fn matvec2(
1637        &self,
1638        x1: &[f32],
1639        x2: &[f32],
1640        o1: &mut [f32],
1641        o2: &mut [f32],
1642        pool: Option<&Pool>,
1643    ) {
1644        match self {
1645            Self::F32 { data, .. } => matvec_rows2(pool, data, x1, x2, o1, o2),
1646            Self::Mapped {
1647                model,
1648                idx,
1649                dtype,
1650                rows,
1651                cols,
1652                row_scale,
1653                col_field,
1654                vbit_offsets,
1655                ..
1656            } => {
1657                if crate::prism::is_forward_weight(model, &model.tensors[*idx].name) {
1658                    let tx1 = crate::prism::forward(model, &x1[..*cols]);
1659                    let tx2 = crate::prism::forward(model, &x2[..*cols]);
1660                    match dtype {
1661                        TensorDtype::Q4Block => {
1662                            q4matvec2(self.quant_bytes(), &tx1, &tx2, *rows, *cols, o1, o2, pool)
1663                        }
1664                        TensorDtype::Q4Tiled => {
1665                            q4t_matvec2(self.quant_bytes(), &tx1, &tx2, *rows, *cols, o1, o2, pool)
1666                        }
1667                        TensorDtype::Q4TiledP => {
1668                            q4tp_matvec2(self.quant_bytes(), &tx1, &tx2, *rows, *cols, o1, o2, pool)
1669                        }
1670                        TensorDtype::Q2TiledP => {
1671                            if crate::prism::is_affine_target(model, &model.tensors[*idx].name) {
1672                                q2tp_affine_matvec2(
1673                                    self.quant_bytes(),
1674                                    &tx1,
1675                                    &tx2,
1676                                    *rows,
1677                                    *cols,
1678                                    o1,
1679                                    o2,
1680                                    pool,
1681                                )
1682                            } else {
1683                                q2tp_matvec2(
1684                                    self.quant_bytes(),
1685                                    &tx1,
1686                                    &tx2,
1687                                    *rows,
1688                                    *cols,
1689                                    o1,
1690                                    o2,
1691                                    pool,
1692                                )
1693                            }
1694                        }
1695                        TensorDtype::Q1 => {
1696                            q1_matvec2(self.quant_bytes(), &tx1, &tx2, *rows, *cols, o1, o2, pool)
1697                        }
1698                        TensorDtype::Q1T => {
1699                            q1t_matvec2(self.quant_bytes(), &tx1, &tx2, *rows, *cols, o1, o2, pool)
1700                        }
1701                        TensorDtype::Vbit | TensorDtype::VbitRo => vbitmatvec2(
1702                            self.quant_bytes(),
1703                            vbit_offsets,
1704                            &tx1,
1705                            &tx2,
1706                            *rows,
1707                            *cols,
1708                            o1,
1709                            o2,
1710                            pool,
1711                        ),
1712                        TensorDtype::Q8Row | TensorDtype::Q8_2f => qmatvec2(
1713                            self.quant_bytes(),
1714                            row_scale,
1715                            &tx1,
1716                            &tx2,
1717                            col_field,
1718                            *dtype,
1719                            *rows,
1720                            *cols,
1721                            o1,
1722                            o2,
1723                            pool,
1724                        ),
1725                        _ => unreachable!("unsupported mapped Prism dtype {dtype:?}"),
1726                    }
1727                    return;
1728                }
1729                if *dtype == TensorDtype::Q4Block {
1730                    q4matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1731                    return;
1732                }
1733                if *dtype == TensorDtype::Q4Tiled {
1734                    q4t_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1735                    return;
1736                }
1737                if *dtype == TensorDtype::Q4TiledP {
1738                    q4tp_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1739                    return;
1740                }
1741                if *dtype == TensorDtype::Q2TiledP {
1742                    q2tp_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1743                    return;
1744                }
1745                if *dtype == TensorDtype::Q1 {
1746                    q1_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1747                    return;
1748                }
1749                if *dtype == TensorDtype::Q1T {
1750                    // Fused ternary pair: one row pass, the register
1751                    // unpack shared across both streams on ARM. (Q1T
1752                    // lacks a row_scale array — scales live inline in
1753                    // the tiles — so it must not fall through to the
1754                    // q8 qmatvec2 below.)
1755                    q1t_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1756                    return;
1757                }
1758                if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1759                    vbitmatvec2(
1760                        self.quant_bytes(),
1761                        vbit_offsets,
1762                        x1,
1763                        x2,
1764                        *rows,
1765                        *cols,
1766                        o1,
1767                        o2,
1768                        pool,
1769                    );
1770                    return;
1771                }
1772                qmatvec2(
1773                    self.quant_bytes(),
1774                    row_scale,
1775                    x1,
1776                    x2,
1777                    col_field,
1778                    *dtype,
1779                    *rows,
1780                    *cols,
1781                    o1,
1782                    o2,
1783                    pool,
1784                );
1785            }
1786        }
1787    }
1788}
1789
1790impl QTensor {
1791    /// Batched matvec (prefill-GEMM): xs — row-major [b, cols],
1792    /// out — row-major [b, rows]. Element-wise semantics are IDENTICAL
1793    /// to b matvec calls (same dot kernels in the same order); the win —
1794    /// the weight row streams from DRAM once per batch, not b times.
1795    /// `(model, index)` when this is a memory-mapped q4tp tensor — the
1796    /// identity a device-resident chain needs to hand `tp_matmat` the
1797    /// weight without going through this struct's own dispatch.
1798    pub fn q4tp_mapped(&self) -> Option<(&std::sync::Arc<CmfModel>, usize)> {
1799        if self.has_prism_contract() {
1800            return None;
1801        }
1802        match self {
1803            Self::Mapped {
1804                model, idx, dtype, ..
1805            } if *dtype == TensorDtype::Q4TiledP => Some((model, *idx)),
1806            _ => None,
1807        }
1808    }
1809
1810    pub fn matmat(&self, xs_all: &[f32], b: usize, out: &mut [f32], pool: Option<&Pool>) {
1811        let cols = self.cols();
1812        let rows = self.rows();
1813        debug_assert_eq!(xs_all.len(), b * cols);
1814        debug_assert_eq!(out.len(), b * rows);
1815        let _prof = crate::cpuprof::time(crate::cpuprof::Slot::Matmat);
1816        // GPTQ calibration: fold this layer's inputs into its Hessian. Only
1817        // Mapped tensors carry a directory name; the check is a relaxed
1818        // atomic load, free when not calibrating.
1819        if crate::gptq_capture::capturing() {
1820            if let Self::Mapped { model, idx, .. } = self {
1821                crate::gptq_capture::accumulate(&model.tensors[*idx].name, xs_all, b, cols);
1822            }
1823        }
1824        match self {
1825            Self::F32 { data, .. } => {
1826                if crate::f32_backend::matmat(data, xs_all, b, rows, cols, out) {
1827                    return;
1828                }
1829                let out_addr = SendMut(out.as_mut_ptr());
1830                // Over the batch, four positions a row at a time: four
1831                // independent chains, each still summing j = 0..cols in the
1832                // scalar order (bit-identical to the one-at-a-time loop and
1833                // to `matvec_rows`). The batch is the wide axis of a prefill;
1834                // the rows may be few (Spark-X2.5's 16-row g_proj spent
1835                // 17 ms a layer here split over its rows).
1836                let run = |start: usize, end: usize| {
1837                    let mut bi = start;
1838                    while bi < end {
1839                        let n = (end - bi).min(4);
1840                        for o in 0..rows {
1841                            let row = &data[o * cols..(o + 1) * cols];
1842                            if n == 4 {
1843                                let x = |k: usize| &xs_all[(bi + k) * cols..(bi + k + 1) * cols];
1844                                let (x0, x1, x2, x3) = (x(0), x(1), x(2), x(3));
1845                                let (mut a0, mut a1, mut a2, mut a3) = (0f32, 0f32, 0f32, 0f32);
1846                                for j in 0..cols {
1847                                    let w = row[j];
1848                                    a0 += w * x0[j];
1849                                    a1 += w * x1[j];
1850                                    a2 += w * x2[j];
1851                                    a3 += w * x3[j];
1852                                }
1853                                for (k, acc) in [a0, a1, a2, a3].into_iter().enumerate() {
1854                                    unsafe { *out_addr.at((bi + k) * rows + o) = acc };
1855                                }
1856                            } else {
1857                                for k in 0..n {
1858                                    let x = &xs_all[(bi + k) * cols..(bi + k + 1) * cols];
1859                                    let mut acc = 0f32;
1860                                    for j in 0..cols {
1861                                        acc += row[j] * x[j];
1862                                    }
1863                                    unsafe { *out_addr.at((bi + k) * rows + o) = acc };
1864                                }
1865                            }
1866                        }
1867                        bi += n;
1868                    }
1869                };
1870                dispatch_rows(pool, b, &run);
1871            }
1872            Self::Mapped {
1873                model,
1874                idx,
1875                dtype,
1876                row_scale,
1877                col_field,
1878                vbit_offsets,
1879                ..
1880            } => {
1881                if crate::prism::is_forward_weight(model, &model.tensors[*idx].name) {
1882                    let mut transformed = Vec::with_capacity(xs_all.len());
1883                    for bi in 0..b {
1884                        transformed.extend_from_slice(&crate::prism::forward(
1885                            model,
1886                            &xs_all[bi * cols..(bi + 1) * cols],
1887                        ));
1888                    }
1889                    match dtype {
1890                        TensorDtype::Q4Block => {
1891                            q4matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
1892                        }
1893                        TensorDtype::Q4Tiled => {
1894                            q4t_matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
1895                        }
1896                        TensorDtype::Q4TiledP => {
1897                            q4tp_matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
1898                        }
1899                        TensorDtype::Q2TiledP => {
1900                            let affine =
1901                                crate::prism::is_affine_target(model, &model.tensors[*idx].name);
1902                            // Affine Prism Q2TP has a descriptor-aware GPU
1903                            // kernel for short/tail batches too.  Unlike the
1904                            // ordinary Q2TP path, don't force b<32 back to a
1905                            // scalar CPU matmat: prefill chunks and the final
1906                            // tail both need to stay on the tested GPU arm.
1907                            let gpu_batch_ok = if affine {
1908                                b >= 2
1909                            } else {
1910                                b >= 32 && b * rows * cols >= 128_000_000
1911                            };
1912                            if gpu_batch_ok
1913                                && cols % 32 == 0
1914                                && crate::gpu::enabled_here()
1915                                && crate::gpu::q2tp_gpu_opt_in()
1916                            {
1917                                let gpu_ok = if affine {
1918                                    crate::gpu::q2tp_affine_matmat(
1919                                        model,
1920                                        *idx,
1921                                        &transformed,
1922                                        b,
1923                                        rows,
1924                                        cols,
1925                                        out,
1926                                    )
1927                                } else {
1928                                    crate::gpu::q2tp_matmat(
1929                                        model,
1930                                        *idx,
1931                                        &transformed,
1932                                        b,
1933                                        rows,
1934                                        cols,
1935                                        out,
1936                                    )
1937                                };
1938                                if gpu_ok {
1939                                    return;
1940                                }
1941                            }
1942                            // A one-token Prism decode is the other short
1943                            // case.  Use the descriptor-aware matvec kernel
1944                            // before falling back to the exact CPU path.
1945                            if affine
1946                                && b == 1
1947                                && cols % 32 == 0
1948                                && crate::gpu::enabled_here()
1949                                && crate::gpu::q2tp_gpu_opt_in()
1950                                && crate::gpu::q2tp_affine_matvec(
1951                                    model,
1952                                    *idx,
1953                                    &transformed[..cols],
1954                                    rows,
1955                                    cols,
1956                                    &mut out[..rows],
1957                                )
1958                            {
1959                                return;
1960                            }
1961                            if affine {
1962                                q2tp_affine_matmat(
1963                                    self.quant_bytes(),
1964                                    &transformed,
1965                                    b,
1966                                    rows,
1967                                    cols,
1968                                    out,
1969                                    pool,
1970                                )
1971                            } else {
1972                                q2tp_matmat(
1973                                    self.quant_bytes(),
1974                                    &transformed,
1975                                    b,
1976                                    rows,
1977                                    cols,
1978                                    out,
1979                                    pool,
1980                                )
1981                            }
1982                        }
1983                        TensorDtype::Q1 => {
1984                            q1_matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
1985                        }
1986                        TensorDtype::Q1T => {
1987                            q1t_matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
1988                        }
1989                        TensorDtype::Vbit | TensorDtype::VbitRo => vbitmatmat(
1990                            self.quant_bytes(),
1991                            vbit_offsets,
1992                            &transformed,
1993                            b,
1994                            rows,
1995                            cols,
1996                            out,
1997                            pool,
1998                        ),
1999                        TensorDtype::Q8Row | TensorDtype::Q8_2f => {
2000                            let pre: Vec<std::borrow::Cow<'_, [f32]>> = (0..b)
2001                                .map(|bi| {
2002                                    prescale(
2003                                        &transformed[bi * cols..(bi + 1) * cols],
2004                                        col_field,
2005                                        *dtype,
2006                                    )
2007                                })
2008                                .collect();
2009                            qmatmat(self.quant_bytes(), row_scale, &pre, rows, cols, out, pool)
2010                        }
2011                        _ => unreachable!("unsupported mapped Prism dtype {dtype:?}"),
2012                    }
2013                    return;
2014                }
2015                if *dtype == TensorDtype::Q4Block {
2016                    q4matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2017                    return;
2018                }
2019                if *dtype == TensorDtype::Q4TiledP {
2020                    // GPU batched q4tp GEMM (dequant + f32nt mul_mm on the
2021                    // device); the probe keeps whichever beats the CPU arm.
2022                    // Narrow (prompt-encode) and wide (DiT) batches probe
2023                    // as separate classes — the regimes have opposite
2024                    // winners and one shared verdict locked the wrong arm.
2025                    // Kill switch (gpu::mm_kill): one grossly slow GPU op
2026                    // (a fair-condition op is ≤~100 ms even at 1024px)
2027                    // means the device is contended by another process
2028                    // (e.g. a simulator) — verdicts are per-process, so
2029                    // without the bail the whole render crawls behind
2030                    // someone else's queue.
2031                    // Row-exact batches stay on the host: the device GEMM
2032                    // is an f32 dequant-sgemm, not the host matvec's sum.
2033                    if b >= 32
2034                        && b * rows * cols >= 128_000_000
2035                        && cols % 32 == 0
2036                        && !row_exact()
2037                        && !crate::gpu::mm_killed()
2038                        && crate::gpu::enabled_here()
2039                    {
2040                        let class = if b >= 128 {
2041                            crate::gpu::OpClass::MatmatWide
2042                        } else {
2043                            crate::gpu::OpClass::Matmat
2044                        };
2045                        if let Self::Mapped { model, idx, .. } = self {
2046                            // In-process A/B (`CMF_MM_AB=1`). Three
2047                            // wall-clock A/Bs on a shared stand disagreed
2048                            // with each other by 25% on the same change,
2049                            // because the machine drifts between processes
2050                            // and interleaving whole renders does not fix
2051                            // that. Here both arms run back to back on the
2052                            // SAME data inside one call, so whatever the
2053                            // machine is doing, it does to both — and the
2054                            // disagreement between their outputs falls out
2055                            // for free. Doubles the work; a diagnostic,
2056                            // not a mode.
2057                            if crate::mm_ab::on() {
2058                                let mut g = vec![0f32; b * rows];
2059                                let t = std::time::Instant::now();
2060                                let took = crate::gpu::q4tp_matmat(
2061                                    model, *idx, xs_all, b, rows, cols, &mut g,
2062                                );
2063                                let dg = t.elapsed();
2064                                let t = std::time::Instant::now();
2065                                q4tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2066                                let dc = t.elapsed();
2067                                crate::mm_ab::record(b, rows, cols, took, dg, dc, &g, out);
2068                                return;
2069                            }
2070                            let t0 = std::time::Instant::now();
2071                            // A cold call takes the device arm: its sample
2072                            // is discarded either way, and the upload is
2073                            // what the next step needs.
2074                            let resident = crate::gpu::weight_is_resident(model, *idx);
2075                            match crate::gpu::probe_arm_cold_prefers_gpu(class, resident) {
2076                                crate::gpu::ProbeArm::Gpu => {
2077                                    if crate::gpu::q4tp_matmat(
2078                                        model, *idx, xs_all, b, rows, cols, out,
2079                                    ) {
2080                                        let el = t0.elapsed();
2081                                        // Work-proportional budget: ~8× the
2082                                        // fair-device estimate (+20 ms slack).
2083                                        // An absolute cap missed the worst
2084                                        // case — contended ops sit at
2085                                        // 100–240 ms each and still bury a
2086                                        // render whose fair op is 3–9 ms.
2087                                        // Cold ops (first PSO build, buffer
2088                                        // alloc) are exempt: a one-off
2089                                        // ~50 ms compile is not contention.
2090                                        let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
2091                                        let budget = std::time::Duration::from_secs_f64(
2092                                            flops / 1.5e12 * 8.0 + 0.020,
2093                                        );
2094                                        crate::gpu::mm_budget_check(
2095                                            "q4tp matmat",
2096                                            el,
2097                                            budget,
2098                                            crate::gpu::probe_was_cold() || !resident,
2099                                        );
2100                                        crate::gpu::probe_record(class, true, el);
2101                                        return;
2102                                    }
2103                                }
2104                                crate::gpu::ProbeArm::CpuTimed => {
2105                                    q4tp_matmat(
2106                                        self.quant_bytes(),
2107                                        xs_all,
2108                                        b,
2109                                        rows,
2110                                        cols,
2111                                        out,
2112                                        pool,
2113                                    );
2114                                    crate::gpu::probe_record(class, false, t0.elapsed());
2115                                    return;
2116                                }
2117                                crate::gpu::ProbeArm::Cpu => {}
2118                            }
2119                        }
2120                    }
2121                    q4tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2122                    return;
2123                }
2124                if *dtype == TensorDtype::Q2TiledP {
2125                    // Same device arm as q4tp, behind the same probe:
2126                    // the planes differ, the dispatch does not. Without
2127                    // this a q2tp file ran its widest projections on the
2128                    // host while the 4-bit one had the card, which is a
2129                    // codec paying for its size twice.
2130                    if b >= 32
2131                        && b * rows * cols >= 128_000_000
2132                        && cols % 32 == 0
2133                        && !crate::gpu::mm_killed()
2134                        && crate::gpu::enabled_here()
2135                    {
2136                        let class = if b >= 128 {
2137                            crate::gpu::OpClass::MatmatWide
2138                        } else {
2139                            crate::gpu::OpClass::Matmat
2140                        };
2141                        if let Self::Mapped { model, idx, .. } = self {
2142                            let t0 = std::time::Instant::now();
2143                            match crate::gpu::probe_arm(class) {
2144                                crate::gpu::ProbeArm::Gpu => {
2145                                    if crate::gpu::q2tp_matmat(
2146                                        model, *idx, xs_all, b, rows, cols, out,
2147                                    ) {
2148                                        crate::gpu::probe_record(class, true, t0.elapsed());
2149                                        return;
2150                                    }
2151                                }
2152                                crate::gpu::ProbeArm::CpuTimed => {
2153                                    q2tp_matmat(
2154                                        self.quant_bytes(),
2155                                        xs_all,
2156                                        b,
2157                                        rows,
2158                                        cols,
2159                                        out,
2160                                        pool,
2161                                    );
2162                                    crate::gpu::probe_record(class, false, t0.elapsed());
2163                                    return;
2164                                }
2165                                crate::gpu::ProbeArm::Cpu => {}
2166                            }
2167                        }
2168                    }
2169                    // Without a host arm a q2tp tensor falls through to
2170                    // the q8 fallback, which reads it at one BYTE per
2171                    // weight — a 2x overrun that killed pool workers
2172                    // mid-prefill while the dispatcher waited forever.
2173                    q2tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2174                    return;
2175                }
2176                if *dtype == TensorDtype::Q4Tiled {
2177                    // GPU batched q4t GEMM (dequant + f32nt mul_mm on the
2178                    // device); the probe keeps whichever beats the CPU arm.
2179                    // Narrow (prompt-encode) and wide (DiT) batches probe
2180                    // as separate classes — the regimes have opposite
2181                    // winners and one shared verdict locked the wrong arm.
2182                    // Kill switch (gpu::mm_kill): one grossly slow GPU op
2183                    // (a fair-condition op is ≤~100 ms even at 1024px)
2184                    // means the device is contended by another process
2185                    // (e.g. a simulator) — verdicts are per-process, so
2186                    // without the bail the whole render crawls behind
2187                    // someone else's queue.
2188                    if b >= 32
2189                        && b * rows * cols >= 128_000_000
2190                        && cols % 32 == 0
2191                        && !crate::gpu::mm_killed()
2192                        && crate::gpu::enabled_here()
2193                    {
2194                        let class = if b >= 128 {
2195                            crate::gpu::OpClass::MatmatWide
2196                        } else {
2197                            crate::gpu::OpClass::Matmat
2198                        };
2199                        if let Self::Mapped { model, idx, .. } = self {
2200                            let t0 = std::time::Instant::now();
2201                            match crate::gpu::probe_arm(class) {
2202                                crate::gpu::ProbeArm::Gpu => {
2203                                    if crate::gpu::q4t_matmat(
2204                                        model, *idx, xs_all, b, rows, cols, out,
2205                                    ) {
2206                                        let el = t0.elapsed();
2207                                        // Work-proportional budget: ~8× the
2208                                        // fair-device estimate (+20 ms slack).
2209                                        // An absolute cap missed the worst
2210                                        // case — contended ops sit at
2211                                        // 100–240 ms each and still bury a
2212                                        // render whose fair op is 3–9 ms.
2213                                        // Cold ops (first PSO build, buffer
2214                                        // alloc) are exempt: a one-off
2215                                        // ~50 ms compile is not contention.
2216                                        let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
2217                                        let budget = std::time::Duration::from_secs_f64(
2218                                            flops / 1.5e12 * 8.0 + 0.020,
2219                                        );
2220                                        crate::gpu::mm_budget_check(
2221                                            "q4t matmat",
2222                                            el,
2223                                            budget,
2224                                            crate::gpu::probe_was_cold(),
2225                                        );
2226                                        crate::gpu::probe_record(class, true, el);
2227                                        return;
2228                                    }
2229                                }
2230                                crate::gpu::ProbeArm::CpuTimed => {
2231                                    q4t_matmat(
2232                                        self.quant_bytes(),
2233                                        xs_all,
2234                                        b,
2235                                        rows,
2236                                        cols,
2237                                        out,
2238                                        pool,
2239                                    );
2240                                    crate::gpu::probe_record(class, false, t0.elapsed());
2241                                    return;
2242                                }
2243                                crate::gpu::ProbeArm::Cpu => {}
2244                            }
2245                        }
2246                    }
2247                    q4t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2248                    return;
2249                }
2250                if *dtype == TensorDtype::Q1 {
2251                    // GPU batched q1 GEMM for wide prefill (q1_mul_mm on the
2252                    // device); the probe keeps whichever beats the CPU matmat.
2253                    if b >= 32
2254                        && b * rows * cols >= 128_000_000
2255                        && cols % 64 == 0
2256                        && crate::gpu::enabled_here()
2257                    {
2258                        if let Self::Mapped { model, idx, .. } = self {
2259                            let t0 = std::time::Instant::now();
2260                            match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
2261                                crate::gpu::ProbeArm::Gpu => {
2262                                    if crate::gpu::q1_matmat(
2263                                        model, *idx, xs_all, b, rows, cols, out,
2264                                    ) {
2265                                        crate::gpu::probe_record(
2266                                            crate::gpu::OpClass::Matmat,
2267                                            true,
2268                                            t0.elapsed(),
2269                                        );
2270                                        return;
2271                                    }
2272                                }
2273                                crate::gpu::ProbeArm::CpuTimed => {
2274                                    q1_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2275                                    crate::gpu::probe_record(
2276                                        crate::gpu::OpClass::Matmat,
2277                                        false,
2278                                        t0.elapsed(),
2279                                    );
2280                                    return;
2281                                }
2282                                crate::gpu::ProbeArm::Cpu => {}
2283                            }
2284                        }
2285                    }
2286                    q1_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2287                    return;
2288                }
2289                if *dtype == TensorDtype::Q1T {
2290                    // GPU batched GEMM for wide prefill (base + overlay on the
2291                    // device); probe keeps the winner vs the CPU matmat.
2292                    if b >= 32 && b * rows * cols >= 128_000_000 && crate::gpu::enabled_here() {
2293                        if let Self::Mapped { model, idx, .. } = self {
2294                            let t0 = std::time::Instant::now();
2295                            match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
2296                                crate::gpu::ProbeArm::Gpu => {
2297                                    if crate::gpu::q1t_matmat(
2298                                        model, *idx, xs_all, b, rows, cols, out,
2299                                    ) {
2300                                        crate::gpu::probe_record(
2301                                            crate::gpu::OpClass::Matmat,
2302                                            true,
2303                                            t0.elapsed(),
2304                                        );
2305                                        return;
2306                                    }
2307                                }
2308                                crate::gpu::ProbeArm::CpuTimed => {
2309                                    q1t_matmat(
2310                                        self.quant_bytes(),
2311                                        xs_all,
2312                                        b,
2313                                        rows,
2314                                        cols,
2315                                        out,
2316                                        pool,
2317                                    );
2318                                    crate::gpu::probe_record(
2319                                        crate::gpu::OpClass::Matmat,
2320                                        false,
2321                                        t0.elapsed(),
2322                                    );
2323                                    return;
2324                                }
2325                                crate::gpu::ProbeArm::Cpu => {}
2326                            }
2327                        }
2328                    }
2329                    q1t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2330                    return;
2331                }
2332                if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
2333                    vbitmatmat(
2334                        self.quant_bytes(),
2335                        vbit_offsets,
2336                        xs_all,
2337                        b,
2338                        rows,
2339                        cols,
2340                        out,
2341                        pool,
2342                    );
2343                    return;
2344                }
2345                let pre: Vec<std::borrow::Cow<'_, [f32]>> = (0..b)
2346                    .map(|bi| prescale(&xs_all[bi * cols..(bi + 1) * cols], col_field, *dtype))
2347                    .collect();
2348                // MiMo verification is a 2–4 row decode panel, not a wide
2349                // prompt GEMM. Keep q8 projections on the same device as
2350                // decode; the generic b>=8 gate otherwise silently moves
2351                // every projection back to CPU. The short wgpu matmat uses
2352                // the same 64-lane reduction as its single-token matvec.
2353                if row_exact()
2354                    && (1..=4).contains(&b)
2355                    && matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
2356                    && crate::gpu::enabled_here()
2357                    && crate::gpu::wgpu_active()
2358                {
2359                    let flat: Vec<f32> = pre.iter().flat_map(|v| v.iter().copied()).collect();
2360                    if crate::gpu::q8_matmat(model, *idx, row_scale, &flat, b, rows, cols, out) {
2361                        return;
2362                    }
2363                }
2364                // D5: large prefill-batch GEMMs — on the GPU (threshold by
2365                // work volume: submission carries b×rows×cols MACs).
2366                // Runtime probe: the naive GEMM shader + sync readback
2367                // lose to the CPU GEMM on slow driver stacks — alternate
2368                // both arms and keep the winner.
2369                if b >= 8 && b * rows * cols >= 128_000_000 && crate::gpu::enabled_here() {
2370                    if let Self::Mapped { model, idx, .. } = self {
2371                        let t0 = std::time::Instant::now();
2372                        match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
2373                            crate::gpu::ProbeArm::Gpu
2374                                if crate::gpu::probe_deciding(crate::gpu::OpClass::Matmat)
2375                                    && !crate::gpu::q8_resident_or_upload(model, *idx) =>
2376                            {
2377                                // Cold weights during probing: the upload
2378                                // has started, the count runs on the CPU —
2379                                // the GPU arm samples on the next touch.
2380                                let q = self.quant_bytes();
2381                                qmatmat(q, row_scale, &pre, rows, cols, out, pool);
2382                                return;
2383                            }
2384                            crate::gpu::ProbeArm::Gpu => {
2385                                // The two-field codec's own entry first, as
2386                                // `device_matmat` takes it: the column field
2387                                // folds into the weight plane and the GEMM
2388                                // runs on the matrix units. The plain int8
2389                                // entry below is the scalar f32 GEMM — on a
2390                                // Spark-X2.5 4B prefill it ran q|k|v|o at
2391                                // host speed.
2392                                if *dtype == TensorDtype::Q8_2f
2393                                    && std::env::var("CMF_Q8_2F_DEV").as_deref() != Ok("0")
2394                                    && crate::gpu::q8_matmat_2f(
2395                                        model, *idx, row_scale, col_field, xs_all, b, rows,
2396                                        cols, out,
2397                                    )
2398                                {
2399                                    crate::gpu::probe_record(
2400                                        crate::gpu::OpClass::Matmat,
2401                                        true,
2402                                        t0.elapsed(),
2403                                    );
2404                                    return;
2405                                }
2406                                let flat: Vec<f32> =
2407                                    pre.iter().flat_map(|v| v.iter().copied()).collect();
2408                                if crate::gpu::q8_matmat(
2409                                    model, *idx, row_scale, &flat, b, rows, cols, out,
2410                                ) {
2411                                    crate::gpu::probe_record(
2412                                        crate::gpu::OpClass::Matmat,
2413                                        true,
2414                                        t0.elapsed(),
2415                                    );
2416                                    return;
2417                                }
2418                            }
2419                            crate::gpu::ProbeArm::CpuTimed => {
2420                                let q = self.quant_bytes();
2421                                qmatmat(q, row_scale, &pre, rows, cols, out, pool);
2422                                crate::gpu::probe_record(
2423                                    crate::gpu::OpClass::Matmat,
2424                                    false,
2425                                    t0.elapsed(),
2426                                );
2427                                return;
2428                            }
2429                            crate::gpu::ProbeArm::Cpu => {}
2430                        }
2431                    }
2432                }
2433                let q = self.quant_bytes();
2434                qmatmat(q, row_scale, &pre, rows, cols, out, pool);
2435            }
2436        }
2437    }
2438}
2439
2440impl QTensor {
2441    /// The device GEMM this tensor would take, run once on the caller's
2442    /// data — the startup parity probe's arm, and the one place that knows
2443    /// which entry point each codec has.
2444    ///
2445    /// It exists because the probe used to look for a `q4tp` weight by
2446    /// name AND dtype, and a container packed any other way was declared
2447    /// "host path" for the whole render even though its codec had a device
2448    /// GEMM of its own. A gate that only recognizes one codec is a gate
2449    /// that silently downgrades every other one.
2450    pub fn device_matmat(&self, xs: &[f32], b: usize, out: &mut [f32]) -> bool {
2451        let (rows, cols) = (self.rows(), self.cols());
2452        let Self::Mapped {
2453            model,
2454            idx,
2455            dtype,
2456            row_scale,
2457            col_field,
2458            ..
2459        } = self
2460        else {
2461            return false;
2462        };
2463        if crate::prism::has_contract(model) {
2464            return false;
2465        }
2466        match *dtype {
2467            TensorDtype::Q4TiledP => crate::gpu::q4tp_matmat(model, *idx, xs, b, rows, cols, out),
2468            // The two-field codec folds its column field into the
2469            // activation, which leaves a plain per-row int8 GEMM — the
2470            // same kernel `q8_row` uses, on both backends.
2471            TensorDtype::Q8Row | TensorDtype::Q8_2f => {
2472                // The field belongs to the weight; only a backend that cannot
2473                // apply it there makes a scaled copy of the activation.
2474                if *dtype == TensorDtype::Q8_2f
2475                    && std::env::var("CMF_Q8_2F_DEV").as_deref() != Ok("0")
2476                    && crate::gpu::q8_matmat_2f(
2477                        model, *idx, row_scale, col_field, xs, b, rows, cols, out,
2478                    )
2479                {
2480                    return true;
2481                }
2482                let flat: Vec<f32> = (0..b)
2483                    .flat_map(|bi| {
2484                        prescale(&xs[bi * cols..(bi + 1) * cols], col_field, *dtype).into_owned()
2485                    })
2486                    .collect();
2487                crate::gpu::q8_matmat(model, *idx, row_scale, &flat, b, rows, cols, out)
2488            }
2489            _ => false,
2490        }
2491    }
2492
2493    /// Multi-matrix job (roadmap §3 P0): N tensors sharing one input
2494    /// run under a SINGLE pool dispatch — QKV or gate+up cost one
2495    /// barrier instead of N. Per-row math is the exact same kernel as
2496    /// `matvec` (bit-identical outputs); only the dispatch is fused.
2497    /// Falls back to N sequential matvecs when the set is not a uniform
2498    /// q8-family/F32 group or there is no pool.
2499    pub fn matvec_many<const N: usize>(
2500        ts: [&QTensor; N],
2501        x: &[f32],
2502        mut outs: [&mut [f32]; N],
2503        pool: Option<&Pool>,
2504    ) {
2505        let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
2506        if ts.iter().any(|t| t.has_prism_contract()) {
2507            // The fused range kernels have no transform descriptor.  Let
2508            // each tensor's ordinary matvec dispatch perform the explicit
2509            // signed FWHT (and retain CPU fallback for mixed q2tp/q4tp).
2510            for (t, o) in ts.iter().zip(outs.iter_mut()) {
2511                t.matvec(x, o, pool);
2512            }
2513            return;
2514        }
2515        let uniform_q8 = ts.iter().all(|t| {
2516            matches!(
2517                t,
2518                Self::Mapped {
2519                    dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f,
2520                    ..
2521                }
2522            )
2523        });
2524        let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
2525        if uniform_f32 && crate::f32_backend::active() {
2526            for (t, o) in ts.iter().zip(outs.iter_mut()) {
2527                t.matvec(x, o, pool);
2528            }
2529            return;
2530        }
2531        let uniform_q4 = ts.iter().all(|t| {
2532            matches!(
2533                t,
2534                Self::Mapped {
2535                    dtype: TensorDtype::Q4Block,
2536                    ..
2537                }
2538            )
2539        });
2540        let uniform_vbit = ts.iter().all(|t| {
2541            matches!(
2542                t,
2543                Self::Mapped {
2544                    dtype: TensorDtype::Vbit | TensorDtype::VbitRo,
2545                    ..
2546                }
2547            )
2548        });
2549        let uniform_q1 = ts.iter().all(|t| {
2550            matches!(
2551                t,
2552                Self::Mapped {
2553                    dtype: TensorDtype::Q1,
2554                    ..
2555                }
2556            )
2557        });
2558        let uniform_q1t = ts.iter().all(|t| {
2559            matches!(
2560                t,
2561                Self::Mapped {
2562                    dtype: TensorDtype::Q1T,
2563                    ..
2564                }
2565            )
2566        });
2567        // q4tp is the skeleton dtype of the big MoE files, and without an arm
2568        // here every projection that shares an input paid its own pool
2569        // barrier: DeepSeek-V4's attention step alone hands this function
2570        // wq_a, wkv and both compressors' pairs off the same hidden state.
2571        let uniform_q4tp = ts.iter().all(|t| {
2572            matches!(
2573                t,
2574                Self::Mapped {
2575                    dtype: TensorDtype::Q4TiledP,
2576                    ..
2577                }
2578            )
2579        }) && ts
2580            .iter()
2581            .all(|t| t.cols() == ts[0].cols() && t.cols() % GROUP_SIZE == 0);
2582        let Some(pool) = pool else {
2583            for (t, o) in ts.iter().zip(outs.iter_mut()) {
2584                t.matvec(x, o, None);
2585            }
2586            return;
2587        };
2588        if total_rows < 256
2589            || !(uniform_q8
2590                || uniform_f32
2591                || uniform_q4
2592                || uniform_vbit
2593                || uniform_q1
2594                || uniform_q1t
2595                || uniform_q4tp)
2596        {
2597            for (t, o) in ts.iter().zip(outs.iter_mut()) {
2598                t.matvec(x, o, Some(pool));
2599            }
2600            return;
2601        }
2602
2603        if uniform_q4tp {
2604            // Every tensor's rows laid end to end in one virtual row space,
2605            // so the whole set is ONE dispatch. The per-row body is the
2606            // `q4tp_matvec` arm verbatim — same activation split, same
2607            // accumulation order — so the outputs are bit-identical to the
2608            // sequential calls this replaces.
2609            let cols = ts[0].cols();
2610            let gpr = cols / GROUP_SIZE;
2611            let views: [Q4tpView; N] =
2612                std::array::from_fn(|i| Q4tpView::new(ts[i].quant_bytes(), ts[i].rows(), cols));
2613            let rows_of: [usize; N] = std::array::from_fn(|i| ts[i].rows());
2614            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2615            // flat index -> (which tensor, which of its rows)
2616            let locate = |flat: usize| -> (usize, usize) {
2617                let mut acc = 0;
2618                for (i, &r) in rows_of.iter().enumerate() {
2619                    if flat < acc + r {
2620                        return (i, flat - acc);
2621                    }
2622                    acc += r;
2623                }
2624                (rows_of.len() - 1, 0)
2625            };
2626            let (views, outs_addr) = (&views, &outs_addr);
2627            if a8w8_enabled() {
2628                let act = split_act(x);
2629                let act = &act;
2630                let run = |start: usize, end: usize| {
2631                    with_krow(gpr, |sc| {
2632                        for flat in start..end {
2633                            let (t, r) = locate(flat);
2634                            let v = &views[t];
2635                            v.scales_into(r, gpr, sc);
2636                            let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, sc) * act.sx;
2637                            for &(j, xv) in &act.outliers {
2638                                let (w, s) = q4tp_outlier(v.nib, r, gpr, j, sc);
2639                                acc += w * s * xv;
2640                            }
2641                            // SAFETY: one worker owns each (tensor, row) pair.
2642                            unsafe { *outs_addr[t].at(r) = acc };
2643                        }
2644                    });
2645                };
2646                pool.run_rows(total_rows, &run);
2647            } else {
2648                let run = |start: usize, end: usize| {
2649                    with_krow(gpr, |sc| {
2650                        for flat in start..end {
2651                            let (t, r) = locate(flat);
2652                            let v = &views[t];
2653                            v.scales_into(r, gpr, sc);
2654                            // SAFETY: one worker owns each (tensor, row) pair.
2655                            unsafe { *outs_addr[t].at(r) = q4tp_row_exact(v.nib, r, gpr, x, sc) };
2656                        }
2657                    });
2658                };
2659                pool.run_rows(total_rows, &run);
2660            }
2661            return;
2662        }
2663
2664        if uniform_q1 {
2665            // One shared activation split + group sums (q1 has no col
2666            // field; the same input feeds every tensor).
2667            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2668            if a8w8_enabled() {
2669                let act = split_act(x);
2670                let gsum = q1_group_sums(&act.xq, ts[0].cols() / GROUP_SIZE);
2671                let (act, gsum) = (&act, &gsum);
2672                let closures: [_; N] = std::array::from_fn(|i| {
2673                    let (bytes, gpr, out) =
2674                        (ts[i].quant_bytes(), ts[i].cols() / GROUP_SIZE, outs_addr[i]);
2675                    move |s: usize, e: usize| q1_range_a8w8(bytes, gpr, act, gsum, out, s, e)
2676                });
2677                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2678                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2679                pool.run_many(&parts);
2680            } else {
2681                let closures: [_; N] = std::array::from_fn(|i| {
2682                    let (bytes, gpr, out) =
2683                        (ts[i].quant_bytes(), ts[i].cols() / GROUP_SIZE, outs_addr[i]);
2684                    move |s: usize, e: usize| q1_range_f32(bytes, gpr, x, out, s, e)
2685                });
2686                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2687                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2688                pool.run_many(&parts);
2689            }
2690            return;
2691        }
2692
2693        if uniform_q1t {
2694            // Q1T batched: one shared activation split + overlay decode,
2695            // all tensors' rows in ONE pool dispatch (saves N−1 dispatches
2696            // and N−1 redundant split_act calls per layer).
2697            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2698            const TILE: usize = cortiq_core::quant::Q1T_TILE;
2699            if a8w8_enabled() {
2700                let act = split_act(x);
2701                let act = &act;
2702                let x_ref = x;
2703                let closures: [_; N] = std::array::from_fn(|i| {
2704                    let bytes = ts[i].quant_bytes();
2705                    let (rows, cols) = (ts[i].rows(), ts[i].cols());
2706                    let gpr = cols / GROUP_SIZE;
2707                    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
2708                    let out = outs_addr[i];
2709                    move |s: usize, e: usize| {
2710                        q1t_range_a8w8(bytes, gpr, rp_off, ent_off, has_ov, act, x_ref, out, s, e)
2711                    }
2712                });
2713                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2714                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2715                pool.run_many(&parts);
2716            } else {
2717                let x_ref = x;
2718                let closures: [_; N] = std::array::from_fn(|i| {
2719                    let bytes = ts[i].quant_bytes();
2720                    let (rows, cols) = (ts[i].rows(), ts[i].cols());
2721                    let gpr = cols / GROUP_SIZE;
2722                    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
2723                    let out = outs_addr[i];
2724                    move |s: usize, e: usize| {
2725                        q1t_range_f32_batch(bytes, gpr, rp_off, ent_off, has_ov, x_ref, out, s, e)
2726                    }
2727                });
2728                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2729                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2730                pool.run_many(&parts);
2731            }
2732            return;
2733        }
2734
2735        if uniform_q4 || uniform_vbit {
2736            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2737            // q4/vbit share one activation split — no per-tensor col field.
2738            if a8w8_enabled() {
2739                let act = split_act(x);
2740                let act = &act;
2741                if uniform_q4 {
2742                    let closures: [_; N] = std::array::from_fn(|i| {
2743                        let (packed, scales) =
2744                            q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
2745                        let (gpr, cols, out) =
2746                            (ts[i].cols() / GROUP_SIZE, ts[i].cols(), outs_addr[i]);
2747                        move |s: usize, e: usize| {
2748                            q4_range_a8w8(packed, scales, gpr, cols, act, out, s, e)
2749                        }
2750                    });
2751                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2752                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2753                    pool.run_many(&parts);
2754                } else {
2755                    let closures: [_; N] = std::array::from_fn(|i| {
2756                        let Self::Mapped { vbit_offsets, .. } = ts[i] else {
2757                            unreachable!()
2758                        };
2759                        let (bytes, rows, cols, out) = (
2760                            ts[i].quant_bytes(),
2761                            ts[i].rows(),
2762                            ts[i].cols(),
2763                            outs_addr[i],
2764                        );
2765                        move |s: usize, e: usize| {
2766                            vbit_range_a8w8(bytes, vbit_offsets, x, act, rows, cols, out, s, e)
2767                        }
2768                    });
2769                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2770                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2771                    pool.run_many(&parts);
2772                }
2773                return;
2774            }
2775            if uniform_q4 {
2776                let closures: [_; N] = std::array::from_fn(|i| {
2777                    let (packed, scales) =
2778                        q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
2779                    let (gpr, out) = (ts[i].cols() / GROUP_SIZE, outs_addr[i]);
2780                    move |s: usize, e: usize| q4_range_f32(packed, scales, gpr, x, out, s, e)
2781                });
2782                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2783                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2784                pool.run_many(&parts);
2785            } else {
2786                let closures: [_; N] = std::array::from_fn(|i| {
2787                    let Self::Mapped { vbit_offsets, .. } = ts[i] else {
2788                        unreachable!()
2789                    };
2790                    let (bytes, rows, cols, out) = (
2791                        ts[i].quant_bytes(),
2792                        ts[i].rows(),
2793                        ts[i].cols(),
2794                        outs_addr[i],
2795                    );
2796                    move |s: usize, e: usize| {
2797                        vbit_range_f32(bytes, vbit_offsets, x, rows, cols, out, s, e)
2798                    }
2799                });
2800                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2801                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2802                pool.run_many(&parts);
2803            }
2804            return;
2805        }
2806
2807        if uniform_f32 {
2808            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2809            let closures: [_; N] = std::array::from_fn(|i| {
2810                let Self::F32 { data, cols, .. } = ts[i] else {
2811                    unreachable!()
2812                };
2813                let out = outs_addr[i];
2814                move |start: usize, end: usize| {
2815                    for o in start..end {
2816                        let row = &data[o * cols..(o + 1) * cols];
2817                        let mut sum = 0.0f32;
2818                        for j in 0..*cols {
2819                            sum += row[j] * x[j];
2820                        }
2821                        // SAFETY: disjoint (tensor, row) cells per worker.
2822                        unsafe { *out.at(o) = sum };
2823                    }
2824                }
2825            });
2826            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2827                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2828            pool.run_many(&parts);
2829            return;
2830        }
2831
2832        // Uniform q8-family: per-tensor prescale (q8_2f col fields
2833        // differ per tensor) + the shared range kernels.
2834        struct Ctx<'a> {
2835            bytes: &'a [u8],
2836            #[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
2837            rep: &'a [u8],
2838            row_scale: &'a [f32],
2839            cols: usize,
2840            xs: std::borrow::Cow<'a, [f32]>,
2841        }
2842        let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
2843            let Self::Mapped {
2844                dtype,
2845                cols,
2846                row_scale,
2847                col_field,
2848                repack,
2849                ..
2850            } = ts[i]
2851            else {
2852                unreachable!()
2853            };
2854            Ctx {
2855                bytes: ts[i].quant_bytes(),
2856                rep: repack,
2857                row_scale,
2858                cols: *cols,
2859                xs: prescale(x, col_field, *dtype),
2860            }
2861        });
2862        let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2863        #[cfg(target_arch = "aarch64")]
2864        if sdot_enabled() {
2865            let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
2866            let closures: [_; N] = std::array::from_fn(|i| {
2867                let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
2868                move |start: usize, end: usize| {
2869                    q8_range_sdot(c.bytes, c.rep, c.row_scale, act, c.cols, out, start, end)
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            return;
2876        }
2877        #[cfg(target_arch = "x86_64")]
2878        if avx2_a8w8_enabled() {
2879            let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
2880            let closures: [_; N] = std::array::from_fn(|i| {
2881                let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
2882                move |start: usize, end: usize| {
2883                    q8_range_avx2(c.bytes, c.row_scale, act, c.cols, out, start, end)
2884                }
2885            });
2886            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2887                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2888            pool.run_many(&parts);
2889            return;
2890        }
2891        let closures: [_; N] = std::array::from_fn(|i| {
2892            let (c, out) = (&ctxs[i], outs_addr[i]);
2893            move |start: usize, end: usize| {
2894                q8_range_f32(c.bytes, c.row_scale, &c.xs, c.cols, out, start, end)
2895            }
2896        });
2897        let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2898            std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2899        pool.run_many(&parts);
2900    }
2901}
2902
2903impl QTensor {
2904    /// Pair-input multi-matrix job: N tensors × 2 shared inputs under a
2905    /// single pool dispatch — the MTP/pair decode path publishes one job
2906    /// for Q/K/V (and one for gate+up) instead of one per tensor.
2907    /// Per-row math is exactly `matvec2`'s kernels; bit-identical.
2908    #[allow(clippy::needless_range_loop)]
2909    pub fn matvec2_many<const N: usize>(
2910        ts: [&QTensor; N],
2911        x1: &[f32],
2912        x2: &[f32],
2913        mut o1s: [&mut [f32]; N],
2914        mut o2s: [&mut [f32]; N],
2915        pool: Option<&Pool>,
2916    ) {
2917        let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
2918        if ts.iter().any(|t| t.has_prism_contract()) {
2919            for i in 0..N {
2920                ts[i].matvec2(x1, x2, o1s[i], o2s[i], pool);
2921            }
2922            return;
2923        }
2924        let uniform_q8 = ts.iter().all(|t| {
2925            matches!(
2926                t,
2927                Self::Mapped {
2928                    dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f,
2929                    ..
2930                }
2931            )
2932        });
2933        let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
2934        let uniform_q4 = ts.iter().all(|t| {
2935            matches!(
2936                t,
2937                Self::Mapped {
2938                    dtype: TensorDtype::Q4Block,
2939                    ..
2940                }
2941            )
2942        });
2943        let uniform_vbit = ts.iter().all(|t| {
2944            matches!(
2945                t,
2946                Self::Mapped {
2947                    dtype: TensorDtype::Vbit | TensorDtype::VbitRo,
2948                    ..
2949                }
2950            )
2951        });
2952        let fusable = pool.is_some()
2953            && total_rows >= 256
2954            && (uniform_q8 || uniform_f32 || uniform_q4 || uniform_vbit);
2955        if !fusable {
2956            for i in 0..N {
2957                ts[i].matvec2(x1, x2, o1s[i], o2s[i], pool);
2958            }
2959            return;
2960        }
2961        let pool = pool.unwrap();
2962
2963        if uniform_q4 || uniform_vbit {
2964            let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
2965            let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
2966            // q4/vbit share activation splits — no per-tensor col field.
2967            if a8w8_enabled() {
2968                let a1 = split_act(x1);
2969                let a2 = split_act(x2);
2970                let (a1, a2) = (&a1, &a2);
2971                if uniform_q4 {
2972                    let closures: [_; N] = std::array::from_fn(|i| {
2973                        let (packed, scales) =
2974                            q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
2975                        let (gpr, cols, o1, o2) =
2976                            (ts[i].cols() / GROUP_SIZE, ts[i].cols(), p1[i], p2[i]);
2977                        move |s: usize, e: usize| {
2978                            q4_range2_a8w8(packed, scales, gpr, cols, a1, a2, o1, o2, s, e)
2979                        }
2980                    });
2981                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2982                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2983                    pool.run_many(&parts);
2984                } else {
2985                    let closures: [_; N] = std::array::from_fn(|i| {
2986                        let Self::Mapped { vbit_offsets, .. } = ts[i] else {
2987                            unreachable!()
2988                        };
2989                        let (bytes, rows, cols, o1, o2) = (
2990                            ts[i].quant_bytes(),
2991                            ts[i].rows(),
2992                            ts[i].cols(),
2993                            p1[i],
2994                            p2[i],
2995                        );
2996                        move |s: usize, e: usize| {
2997                            vbit_range2_a8w8(
2998                                bytes,
2999                                vbit_offsets,
3000                                x1,
3001                                x2,
3002                                a1,
3003                                a2,
3004                                rows,
3005                                cols,
3006                                o1,
3007                                o2,
3008                                s,
3009                                e,
3010                            )
3011                        }
3012                    });
3013                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3014                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3015                    pool.run_many(&parts);
3016                }
3017                return;
3018            }
3019            if uniform_q4 {
3020                let closures: [_; N] = std::array::from_fn(|i| {
3021                    let (packed, scales) =
3022                        q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
3023                    let (gpr, o1, o2) = (ts[i].cols() / GROUP_SIZE, p1[i], p2[i]);
3024                    move |s: usize, e: usize| {
3025                        q4_range2_f32(packed, scales, gpr, x1, x2, o1, o2, s, e)
3026                    }
3027                });
3028                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3029                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3030                pool.run_many(&parts);
3031            } else {
3032                let closures: [_; N] = std::array::from_fn(|i| {
3033                    let Self::Mapped { vbit_offsets, .. } = ts[i] else {
3034                        unreachable!()
3035                    };
3036                    let (bytes, rows, cols, o1, o2) = (
3037                        ts[i].quant_bytes(),
3038                        ts[i].rows(),
3039                        ts[i].cols(),
3040                        p1[i],
3041                        p2[i],
3042                    );
3043                    move |s: usize, e: usize| {
3044                        vbit_range2_f32(bytes, vbit_offsets, x1, x2, rows, cols, o1, o2, s, e)
3045                    }
3046                });
3047                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3048                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3049                pool.run_many(&parts);
3050            }
3051            return;
3052        }
3053
3054        if uniform_f32 {
3055            let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
3056            let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
3057            let closures: [_; N] = std::array::from_fn(|i| {
3058                let Self::F32 { data, cols, .. } = ts[i] else {
3059                    unreachable!()
3060                };
3061                let (o1, o2) = (p1[i], p2[i]);
3062                move |start: usize, end: usize| {
3063                    for o in start..end {
3064                        let row = &data[o * cols..(o + 1) * cols];
3065                        let (mut s1, mut s2) = (0.0f32, 0.0f32);
3066                        for j in 0..*cols {
3067                            s1 += row[j] * x1[j];
3068                            s2 += row[j] * x2[j];
3069                        }
3070                        // SAFETY: disjoint (tensor, row) cells per worker.
3071                        unsafe {
3072                            *o1.at(o) = s1;
3073                            *o2.at(o) = s2;
3074                        }
3075                    }
3076                }
3077            });
3078            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3079                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3080            pool.run_many(&parts);
3081            return;
3082        }
3083
3084        struct Ctx<'a> {
3085            bytes: &'a [u8],
3086            row_scale: &'a [f32],
3087            cols: usize,
3088            xs1: std::borrow::Cow<'a, [f32]>,
3089            xs2: std::borrow::Cow<'a, [f32]>,
3090        }
3091        let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
3092            let Self::Mapped {
3093                dtype,
3094                cols,
3095                row_scale,
3096                col_field,
3097                ..
3098            } = ts[i]
3099            else {
3100                unreachable!()
3101            };
3102            Ctx {
3103                bytes: ts[i].quant_bytes(),
3104                row_scale,
3105                cols: *cols,
3106                xs1: prescale(x1, col_field, *dtype),
3107                xs2: prescale(x2, col_field, *dtype),
3108            }
3109        });
3110        let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
3111        let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
3112        #[cfg(target_arch = "aarch64")]
3113        if sdot_enabled() {
3114            let acts: [(SplitAct, SplitAct); N] =
3115                std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
3116            let closures: [_; N] = std::array::from_fn(|i| {
3117                let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
3118                move |start: usize, end: usize| {
3119                    q8_range2_sdot(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
3120                }
3121            });
3122            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3123                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3124            pool.run_many(&parts);
3125            return;
3126        }
3127        #[cfg(target_arch = "x86_64")]
3128        if avx2_a8w8_enabled() {
3129            let acts: [(SplitAct, SplitAct); N] =
3130                std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
3131            let closures: [_; N] = std::array::from_fn(|i| {
3132                let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
3133                move |start: usize, end: usize| {
3134                    q8_range2_avx2(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
3135                }
3136            });
3137            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3138                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3139            pool.run_many(&parts);
3140            return;
3141        }
3142        let closures: [_; N] = std::array::from_fn(|i| {
3143            let (c, o1, o2) = (&ctxs[i], p1[i], p2[i]);
3144            move |start: usize, end: usize| {
3145                q8_range2_f32(
3146                    c.bytes,
3147                    c.row_scale,
3148                    &c.xs1,
3149                    &c.xs2,
3150                    c.cols,
3151                    o1,
3152                    o2,
3153                    start,
3154                    end,
3155                )
3156            }
3157        });
3158        let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3159            std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3160        pool.run_many(&parts);
3161    }
3162
3163    /// Fused gate+up matvec with SiLU·mul: for each row r, computes
3164    /// `silu(gate·x) * (up·x)` and writes to `out[r]`. ONE pool dispatch,
3165    /// no intermediate g/u buffers, no separate silu pass. Falls back
3166    /// (returns false) for unsupported dtype combos.
3167    pub fn matvec_silu_mul(
3168        gate: &QTensor,
3169        up: &QTensor,
3170        x: &[f32],
3171        out: &mut [f32],
3172        pool: Option<&Pool>,
3173    ) -> bool {
3174        Self::matvec_silu_mul_limited(gate, up, x, out, 0.0, pool)
3175    }
3176
3177    /// Fused gate+up+SiLU with the GLM asymmetrical clamp.  `limit == 0`
3178    /// preserves the historical unclamped helper; a positive limit clamps
3179    /// `up` to both sides and `gate` only from above, matching the GLM
3180    /// SwiGLU reference.  Keeping the limit in the row kernel avoids the two
3181    /// intermediate vectors and the extra combine pass on the Q2TP experts.
3182    pub fn matvec_silu_mul_limited(
3183        gate: &QTensor,
3184        up: &QTensor,
3185        x: &[f32],
3186        out: &mut [f32],
3187        limit: f32,
3188        pool: Option<&Pool>,
3189    ) -> bool {
3190        if gate.has_prism_contract() || up.has_prism_contract() {
3191            // The fused gate/up kernels consume x directly.  Prism requires
3192            // a per-matrix signed FWHT, so the caller must use two ordinary
3193            // descriptor-aware matvecs instead of an unrotated fast path.
3194            return false;
3195        }
3196        let inter = gate.rows();
3197        debug_assert_eq!(up.rows(), inter);
3198        debug_assert_eq!(out.len(), inter);
3199        debug_assert_eq!(gate.cols(), up.cols());
3200        if !a8w8_enabled() {
3201            return false;
3202        }
3203        let act = split_act(x);
3204        let act = &act;
3205        let x_ref = x;
3206        let out_addr = SendMut(out.as_mut_ptr());
3207
3208        match (gate, up) {
3209            // Q4Block gate + Q4Block up (most common mobile q4 models)
3210            (
3211                Self::Mapped {
3212                    dtype: TensorDtype::Q4Block,
3213                    ..
3214                },
3215                Self::Mapped {
3216                    dtype: TensorDtype::Q4Block,
3217                    ..
3218                },
3219            ) => {
3220                let (gp, gs) = q4_split(gate.quant_bytes(), gate.rows(), gate.cols());
3221                let (up_p, up_s) = q4_split(up.quant_bytes(), up.rows(), up.cols());
3222                let gpr = gate.cols() / GROUP_SIZE;
3223                let cols = gate.cols();
3224                let run = move |start: usize, end: usize| {
3225                    for r in start..end {
3226                        let mut gv = dot_q4_row_i8(gp, gs, r * gpr, gpr, &act.xq) * act.sx;
3227                        let mut uv = dot_q4_row_i8(up_p, up_s, r * gpr, gpr, &act.xq) * act.sx;
3228                        for &(j, xv) in &act.outliers {
3229                            let flat = r * cols + j;
3230                            let gb = gp[flat / 2];
3231                            let gn = if flat & 1 == 0 { gb & 0x0F } else { gb >> 4 };
3232                            let gsc = f16_to_f32(u16::from_le_bytes([
3233                                gs[(flat / GROUP_SIZE) * 2],
3234                                gs[(flat / GROUP_SIZE) * 2 + 1],
3235                            ]));
3236                            gv += ((gn as i32 - 8) as f32) * gsc * xv;
3237                            let ub = up_p[flat / 2];
3238                            let un = if flat & 1 == 0 { ub & 0x0F } else { ub >> 4 };
3239                            let usc = f16_to_f32(u16::from_le_bytes([
3240                                up_s[(flat / GROUP_SIZE) * 2],
3241                                up_s[(flat / GROUP_SIZE) * 2 + 1],
3242                            ]));
3243                            uv += ((un as i32 - 8) as f32) * usc * xv;
3244                        }
3245                        // SAFETY: disjoint row ranges per worker.
3246                        unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3247                    }
3248                };
3249                dispatch_rows(pool, inter, &run);
3250                true
3251            }
3252            // Q4Tiled gate + Q4Tiled up — one row pass, both tile
3253            // streams sequential, silu·mul fused (same per-row math as
3254            // `q4t_matvec`).
3255            (
3256                Self::Mapped {
3257                    dtype: TensorDtype::Q4Tiled,
3258                    ..
3259                },
3260                Self::Mapped {
3261                    dtype: TensorDtype::Q4Tiled,
3262                    ..
3263                },
3264            ) => {
3265                let g_bytes = gate.quant_bytes();
3266                let u_bytes = up.quant_bytes();
3267                let gpr = gate.cols() / GROUP_SIZE;
3268                let run = move |start: usize, end: usize| {
3269                    for r in start..end {
3270                        let mut gv = dot_q4t_row_i8(g_bytes, r, gpr, &act.xq) * act.sx;
3271                        let mut uv = dot_q4t_row_i8(u_bytes, r, gpr, &act.xq) * act.sx;
3272                        for &(j, xv) in &act.outliers {
3273                            let (w, s) = q4t_outlier(g_bytes, r, gpr, j);
3274                            gv += w * s * xv;
3275                            let (w, s) = q4t_outlier(u_bytes, r, gpr, j);
3276                            uv += w * s * xv;
3277                        }
3278                        // SAFETY: disjoint row ranges per worker.
3279                        unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3280                    }
3281                };
3282                dispatch_rows(pool, inter, &run);
3283                true
3284            }
3285            // Q4TiledP gate + Q4TiledP up — the same fused row pass, with
3286            // each row's two ladders built once and spent on both streams.
3287            (
3288                Self::Mapped {
3289                    dtype: TensorDtype::Q4TiledP,
3290                    ..
3291                },
3292                Self::Mapped {
3293                    dtype: TensorDtype::Q4TiledP,
3294                    ..
3295                },
3296            ) => {
3297                let cols = gate.cols();
3298                let gpr = cols / GROUP_SIZE;
3299                let gv_view = Q4tpView::new(gate.quant_bytes(), inter, cols);
3300                let uv_view = Q4tpView::new(up.quant_bytes(), inter, cols);
3301                let run = |start: usize, end: usize| {
3302                    with_krows(gpr, |gsc, usc| {
3303                        for r in start..end {
3304                            gv_view.scales_into(r, gpr, gsc);
3305                            uv_view.scales_into(r, gpr, usc);
3306                            let mut gv =
3307                                dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsc) * act.sx;
3308                            let mut uv =
3309                                dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, usc) * act.sx;
3310                            for &(j, xv) in &act.outliers {
3311                                let (w, s) = q4tp_outlier(gv_view.nib, r, gpr, j, gsc);
3312                                gv += w * s * xv;
3313                                let (w, s) = q4tp_outlier(uv_view.nib, r, gpr, j, usc);
3314                                uv += w * s * xv;
3315                            }
3316                            // SAFETY: disjoint row ranges per worker.
3317                            unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3318                        }
3319                    });
3320                };
3321                dispatch_rows(pool, inter, &run);
3322                true
3323            }
3324            // Q1 gate + Q1 up — one row pass over both sign streams,
3325            // silu·mul fused (the per-row math of `q1_range_a8w8`); the
3326            // activation group sums are shared by both streams. Without
3327            // this arm a q1 dense FFN paid two dispatches + a combine
3328            // loop — the exact barrier this function exists to remove.
3329            (
3330                Self::Mapped {
3331                    dtype: TensorDtype::Q1,
3332                    ..
3333                },
3334                Self::Mapped {
3335                    dtype: TensorDtype::Q1,
3336                    ..
3337                },
3338            ) => {
3339                let g_bytes = gate.quant_bytes();
3340                let u_bytes = up.quant_bytes();
3341                let gpr = gate.cols() / GROUP_SIZE;
3342                let gsum = q1_group_sums(&act.xq, gpr);
3343                let gsum = &gsum;
3344                let run = move |start: usize, end: usize| {
3345                    for r in start..end {
3346                        let mut gv = dot_q1_row_i8(g_bytes, r, gpr, &act.xq, gsum) * act.sx;
3347                        let mut uv = dot_q1_row_i8(u_bytes, r, gpr, &act.xq, gsum) * act.sx;
3348                        for &(j, xv) in &act.outliers {
3349                            let (w, s) = q1_outlier(g_bytes, r, gpr, j);
3350                            gv += w * s * xv;
3351                            let (w, s) = q1_outlier(u_bytes, r, gpr, j);
3352                            uv += w * s * xv;
3353                        }
3354                        // SAFETY: disjoint row ranges per worker.
3355                        unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3356                    }
3357                };
3358                dispatch_rows(pool, inter, &run);
3359                true
3360            }
3361            // Q2TiledP gate + Q2TiledP up — the 2-bit expert pair (MoE
3362            // FFNs of the W2 class): one row pass, both ladders built
3363            // once, integer code dots with shared group sums.
3364            (
3365                Self::Mapped {
3366                    dtype: TensorDtype::Q2TiledP,
3367                    ..
3368                },
3369                Self::Mapped {
3370                    dtype: TensorDtype::Q2TiledP,
3371                    ..
3372                },
3373            ) => {
3374                let cols = gate.cols();
3375                let gpr = cols / GROUP_SIZE;
3376                let gv_view = Q4tpView::new_q2(gate.quant_bytes(), inter, cols);
3377                let uv_view = Q4tpView::new_q2(up.quant_bytes(), inter, cols);
3378                let gsum = q1_group_sums(&act.xq, gpr);
3379                let gsum = &gsum;
3380                let run = move |start: usize, end: usize| {
3381                    with_krows(gpr, |gsc, usc| {
3382                        for r in start..end {
3383                            gv_view.scales_into(r, gpr, gsc);
3384                            uv_view.scales_into(r, gpr, usc);
3385                            let mut gv =
3386                                dot_q2tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsum, gsc) * act.sx;
3387                            let mut uv =
3388                                dot_q2tp_row_i8(uv_view.nib, r, gpr, &act.xq, gsum, usc) * act.sx;
3389                            for &(j, xv) in &act.outliers {
3390                                let (w, s) = q2tp_outlier(gv_view.nib, r, gpr, j, gsc);
3391                                gv += w * s * xv;
3392                                let (w, s) = q2tp_outlier(uv_view.nib, r, gpr, j, usc);
3393                                uv += w * s * xv;
3394                            }
3395                            // SAFETY: disjoint row ranges per worker.
3396                            unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3397                        }
3398                    });
3399                };
3400                dispatch_rows(pool, inter, &run);
3401                true
3402            }
3403            // Q8Row gate + Q8Row up — one row pass over both i8 streams.
3404            // Q8_2f stays out on purpose: its column field prescales the
3405            // activations PER TENSOR, which breaks this fn's shared
3406            // split_act contract — it keeps the two-dispatch path.
3407            (
3408                Self::Mapped {
3409                    dtype: TensorDtype::Q8Row,
3410                    row_scale: g_rs,
3411                    ..
3412                },
3413                Self::Mapped {
3414                    dtype: TensorDtype::Q8Row,
3415                    row_scale: u_rs,
3416                    ..
3417                },
3418            ) => {
3419                let g_bytes = gate.quant_bytes();
3420                let u_bytes = up.quant_bytes();
3421                let cols = gate.cols();
3422                let run = move |start: usize, end: usize| {
3423                    for r in start..end {
3424                        let gv = q8_row_dot(&g_bytes[r * cols..(r + 1) * cols], act) * g_rs[r];
3425                        let uv = q8_row_dot(&u_bytes[r * cols..(r + 1) * cols], act) * u_rs[r];
3426                        // SAFETY: disjoint row ranges per worker.
3427                        unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3428                    }
3429                };
3430                dispatch_rows(pool, inter, &run);
3431                true
3432            }
3433            // Q1T gate + Q1T up
3434            (
3435                Self::Mapped {
3436                    dtype: TensorDtype::Q1T,
3437                    ..
3438                },
3439                Self::Mapped {
3440                    dtype: TensorDtype::Q1T,
3441                    ..
3442                },
3443            ) => {
3444                const TILE: usize = cortiq_core::quant::Q1T_TILE;
3445                let g_bytes = gate.quant_bytes();
3446                let u_bytes = up.quant_bytes();
3447                let gpr = gate.cols() / GROUP_SIZE;
3448                let (g_rp, g_ent, g_ov) = q1t_overlay(g_bytes, inter * gpr * TILE, inter);
3449                let (u_rp, u_ent, u_ov) = q1t_overlay(u_bytes, inter * gpr * TILE, inter);
3450                let run = move |start: usize, end: usize| {
3451                    for r in start..end {
3452                        let mut gv = q1t_dot_row_i8(g_bytes, r, gpr, &act.xq) * act.sx;
3453                        let mut uv = q1t_dot_row_i8(u_bytes, r, gpr, &act.xq) * act.sx;
3454                        for &(j, xv) in &act.outliers {
3455                            gv += q1t_base_weight(g_bytes, r, gpr, j) * xv;
3456                            uv += q1t_base_weight(u_bytes, r, gpr, j) * xv;
3457                        }
3458                        gv += q1t_row_outlier_correction(g_bytes, r, g_rp, g_ent, g_ov, x_ref);
3459                        uv += q1t_row_outlier_correction(u_bytes, r, u_rp, u_ent, u_ov, x_ref);
3460                        // SAFETY: disjoint row ranges per worker.
3461                        unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3462                    }
3463                };
3464                dispatch_rows(pool, inter, &run);
3465                true
3466            }
3467            _ => false,
3468        }
3469    }
3470
3471    /// Every routed expert's fused gate/up/SiLU under ONE pool dispatch.
3472    ///
3473    /// The per-expert path pays a pool barrier per expert per stage: at 9
3474    /// experts over 40 layers that is ~720 barriers a token, and a decode
3475    /// profile of Qwen3.6-35B-A3B showed the pool parked in
3476    /// `psynch_cvwait` about twice as long as it spent computing. Laying
3477    /// every expert's rows end-to-end in one virtual row space collapses
3478    /// the stage to a single dispatch. The per-row body is the
3479    /// single-expert q4tp arm verbatim, so outputs are bit-identical.
3480    ///
3481    /// `false` = something is outside the fused q4tp kernel (dtype, shape,
3482    /// or a transformed tensor); the caller walks the ordinary per-expert
3483    /// path. Float activations use the same exact scalar rows, still fused
3484    /// under one pool dispatch.
3485    pub fn moe_gate_up_many(
3486        pairs: &[(&QTensor, &QTensor)],
3487        x: &[f32],
3488        outs: &mut [Vec<f32>],
3489        pool: Option<&Pool>,
3490    ) -> bool {
3491        Self::moe_gate_up_many_limited(pairs, x, outs, 0.0, pool)
3492    }
3493
3494    /// Batched gate/up/SiLU with the optional GLM clamp.  The public legacy
3495    /// helper above keeps its historical unclamped semantics; callers that
3496    /// implement a reference with a positive SwiGLU limit use this variant.
3497    pub fn moe_gate_up_many_limited(
3498        pairs: &[(&QTensor, &QTensor)],
3499        x: &[f32],
3500        outs: &mut [Vec<f32>],
3501        limit: f32,
3502        pool: Option<&Pool>,
3503    ) -> bool {
3504        if pairs.is_empty() || pairs.len() != outs.len() {
3505            return false;
3506        }
3507        if !a8w8_enabled() {
3508            if limit > 0.0 {
3509                // The exact-row fallback applies no SwiGLU clamp; the
3510                // caller's per-expert path carries it instead.
3511                return false;
3512            }
3513            let groups = vec![vec![0]; pairs.len()];
3514            return Self::moe_gate_up_rows(pairs, &groups, x, outs, pool);
3515        }
3516        let inter = pairs[0].0.rows();
3517        let cols = pairs[0].0.cols();
3518        if cols % GROUP_SIZE != 0 {
3519            return false;
3520        }
3521        let gpr = cols / GROUP_SIZE;
3522        // Uniform layout across every routed pair: q4tp, or the 2-bit
3523        // profile's q2tp gate/up (the W2 class). Mixed sets refuse.
3524        let q2 = matches!(
3525            pairs[0].0,
3526            Self::Mapped {
3527                dtype: TensorDtype::Q2TiledP,
3528                ..
3529            }
3530        );
3531        let want = if q2 {
3532            TensorDtype::Q2TiledP
3533        } else {
3534            TensorDtype::Q4TiledP
3535        };
3536        let mut views = Vec::with_capacity(pairs.len() * 2);
3537        for ((g, u), o) in pairs.iter().zip(outs.iter()) {
3538            let both = matches!(g, Self::Mapped { dtype, .. } if *dtype == want)
3539                && matches!(u, Self::Mapped { dtype, .. } if *dtype == want);
3540            if !both
3541                || g.rows() != inter
3542                || u.rows() != inter
3543                || g.cols() != cols
3544                || u.cols() != cols
3545                || o.len() != inter
3546            {
3547                return false;
3548            }
3549            let mk = if q2 { Q4tpView::new_q2 } else { Q4tpView::new };
3550            views.push(mk(g.quant_bytes(), inter, cols));
3551            views.push(mk(u.quant_bytes(), inter, cols));
3552        }
3553        let act = split_act(x);
3554        let gsum = if q2 {
3555            q1_group_sums(&act.xq, gpr)
3556        } else {
3557            Vec::new()
3558        };
3559        let (act, gsum) = (&act, &gsum);
3560        let ptrs: Vec<SendMut> = outs.iter_mut().map(|o| SendMut(o.as_mut_ptr())).collect();
3561        let (views, ptrs) = (&views, &ptrs);
3562        let run = |start: usize, end: usize| {
3563            with_krows(gpr, |gsc, usc| {
3564                for flat in start..end {
3565                    let (e, r) = (flat / inter, flat % inter);
3566                    let gv_view = &views[e * 2];
3567                    let uv_view = &views[e * 2 + 1];
3568                    gv_view.scales_into(r, gpr, gsc);
3569                    uv_view.scales_into(r, gpr, usc);
3570                    let (mut gv, mut uv) = if q2 {
3571                        (
3572                            dot_q2tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsum, gsc) * act.sx,
3573                            dot_q2tp_row_i8(uv_view.nib, r, gpr, &act.xq, gsum, usc) * act.sx,
3574                        )
3575                    } else {
3576                        (
3577                            dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsc) * act.sx,
3578                            dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, usc) * act.sx,
3579                        )
3580                    };
3581                    for &(j, xv) in &act.outliers {
3582                        let (og, ou) = if q2 {
3583                            (
3584                                q2tp_outlier(gv_view.nib, r, gpr, j, gsc),
3585                                q2tp_outlier(uv_view.nib, r, gpr, j, usc),
3586                            )
3587                        } else {
3588                            (
3589                                q4tp_outlier(gv_view.nib, r, gpr, j, gsc),
3590                                q4tp_outlier(uv_view.nib, r, gpr, j, usc),
3591                            )
3592                        };
3593                        gv += og.0 * og.1 * xv;
3594                        uv += ou.0 * ou.1 * xv;
3595                    }
3596                    // SAFETY: one worker owns each (expert, row) pair.
3597                    unsafe { *ptrs[e].at(r) = silu_mul_limited(gv, uv, limit) };
3598                }
3599            });
3600        };
3601        dispatch_rows(pool, pairs.len() * inter, &run);
3602        true
3603    }
3604
3605    /// Every routed expert's down projection, weighted and summed into
3606    /// `out`, under ONE pool dispatch.
3607    ///
3608    /// Partitioned by OUTPUT row rather than by expert: each row is owned
3609    /// by a single worker, so the experts are summed in the caller's order
3610    /// — the same sequence of f32 adds the serial `out[i] += w·eo[i]` loop
3611    /// performs, hence bit-identical. Partitioning by expert instead would
3612    /// race on the shared accumulator.
3613    pub fn moe_down_many(
3614        downs: &[&QTensor],
3615        gs: &[Vec<f32>],
3616        weights: &[f32],
3617        out: &mut [f32],
3618        pool: Option<&Pool>,
3619    ) -> bool {
3620        if downs.is_empty() || downs.len() != gs.len() || downs.len() != weights.len() {
3621            return false;
3622        }
3623        if !a8w8_enabled() {
3624            let mut terms = vec![vec![0.0; out.len()]; downs.len()];
3625            if !Self::moe_down_rows(downs, &vec![1; downs.len()], gs, &mut terms, pool) {
3626                return false;
3627            }
3628            out.fill(0.0);
3629            for (row, &w) in terms.iter().zip(weights) {
3630                for (o, &v) in out.iter_mut().zip(row) {
3631                    *o += w * v;
3632                }
3633            }
3634            return true;
3635        }
3636        let rows = out.len();
3637        let cols = downs[0].cols();
3638        if cols % GROUP_SIZE != 0 {
3639            return false;
3640        }
3641        let gpr = cols / GROUP_SIZE;
3642        let mut views = Vec::with_capacity(downs.len());
3643        for (d, g) in downs.iter().zip(gs.iter()) {
3644            if !matches!(
3645                d,
3646                Self::Mapped {
3647                    dtype: TensorDtype::Q4TiledP,
3648                    ..
3649                }
3650            ) || d.rows() != rows
3651                || d.cols() != cols
3652                || g.len() != cols
3653            {
3654                return false;
3655            }
3656            views.push(Q4tpView::new(d.quant_bytes(), rows, cols));
3657        }
3658        // One int8 split per expert — the activation vectors differ.
3659        let acts: Vec<SplitAct> = gs.iter().map(|g| split_act(g)).collect();
3660        // Partitioned by OUTPUT row, with the experts folded inside: each
3661        // row is owned by one worker, so they are summed in the caller's
3662        // order — the same f32 sequence the serial `out[i] += w·eo[i]`
3663        // loop produces. Partitioning by expert instead would either race
3664        // on the accumulator or need a scratch plane and a second pass;
3665        // measured, that variant was a wash, so this keeps the simpler
3666        // shape.
3667        let out_addr = SendMut(out.as_mut_ptr());
3668        let (views, acts, weights) = (&views, &acts, &weights);
3669        let run = |start: usize, end: usize| {
3670            with_krow(gpr, |sc| {
3671                for r in start..end {
3672                    let mut acc = 0f32;
3673                    for (e, v) in views.iter().enumerate() {
3674                        v.scales_into(r, gpr, sc);
3675                        let a = &acts[e];
3676                        let mut d = dot_q4tp_row_i8(v.nib, r, gpr, &a.xq, sc) * a.sx;
3677                        for &(j, xv) in &a.outliers {
3678                            let (w, s) = q4tp_outlier(v.nib, r, gpr, j, sc);
3679                            d += w * s * xv;
3680                        }
3681                        acc += weights[e] * d;
3682                    }
3683                    // SAFETY: disjoint row ranges per worker.
3684                    unsafe { *out_addr.at(r) = acc };
3685                }
3686            });
3687        };
3688        dispatch_rows(pool, rows, &run);
3689        true
3690    }
3691
3692    /// `moe_gate_up_many` for SEVERAL tokens at once, decode-exact: expert
3693    /// `e` (`pairs[e]`, q4tp) serves the tokens `groups[e]` (row indices
3694    /// into `xs`, each `cols` wide). Every (expert, token) output is
3695    /// bit-identical to `moe_gate_up_many` run on that token alone — the
3696    /// same int8 activation split, VNNI dots, outlier terms and inline
3697    /// SiLU — while each weight row is read once for all the tokens routed
3698    /// to its expert (the speculative verify's expert sharing). `outs` is
3699    /// flat in (expert, token-of-group) order. False = not covered (not
3700    /// q4tp): the caller takes the per-token path. With float activations,
3701    /// the exact scalar row kernel replaces the int8 dot without changing
3702    /// the shared dispatch or route-order reduction.
3703    pub fn moe_gate_up_rows(
3704        pairs: &[(&QTensor, &QTensor)],
3705        groups: &[Vec<usize>],
3706        xs: &[f32],
3707        outs: &mut [Vec<f32>],
3708        pool: Option<&Pool>,
3709    ) -> bool {
3710        if pairs.is_empty() || pairs.len() != groups.len() {
3711            return false;
3712        }
3713        let inter = pairs[0].0.rows();
3714        let cols = pairs[0].0.cols();
3715        let n_pairs: usize = groups.iter().map(|g| g.len()).sum();
3716        if cols == 0 || cols % GROUP_SIZE != 0 || outs.len() != n_pairs || xs.len() % cols != 0 {
3717            return false;
3718        }
3719        let b = xs.len() / cols;
3720        let gpr = cols / GROUP_SIZE;
3721        let mut views = Vec::with_capacity(pairs.len() * 2);
3722        for (g, u) in pairs {
3723            let q4tp = |t: &QTensor| {
3724                matches!(
3725                    t,
3726                    Self::Mapped {
3727                        dtype: TensorDtype::Q4TiledP,
3728                        ..
3729                    }
3730                )
3731            };
3732            if g.has_prism_contract()
3733                || u.has_prism_contract()
3734                || !q4tp(g)
3735                || !q4tp(u)
3736                || g.rows() != inter
3737                || u.rows() != inter
3738                || g.cols() != cols
3739                || u.cols() != cols
3740            {
3741                return false;
3742            }
3743            views.push(Q4tpView::new(g.quant_bytes(), inter, cols));
3744            views.push(Q4tpView::new(u.quant_bytes(), inter, cols));
3745        }
3746        if outs.iter().any(|o| o.len() != inter) || groups.iter().flatten().any(|&t| t >= b) {
3747            return false;
3748        }
3749        let quantized = a8w8_enabled();
3750        let acts: Vec<SplitAct> = if quantized {
3751            (0..b)
3752                .map(|t| split_act(&xs[t * cols..(t + 1) * cols]))
3753                .collect()
3754        } else {
3755            Vec::new()
3756        };
3757        let mut offs = Vec::with_capacity(groups.len());
3758        let mut o = 0usize;
3759        for g in groups {
3760            offs.push(o);
3761            o += g.len();
3762        }
3763        let ptrs: Vec<SendMut> = outs.iter_mut().map(|o| SendMut(o.as_mut_ptr())).collect();
3764        let (views, ptrs, acts, offs) = (&views, &ptrs, &acts, &offs);
3765        let run = |start: usize, end: usize| {
3766            let (mut gsc, mut usc) = (vec![0f32; gpr], vec![0f32; gpr]);
3767            for flat in start..end {
3768                let (e, r) = (flat / inter, flat % inter);
3769                let (gv_view, uv_view) = (&views[e * 2], &views[e * 2 + 1]);
3770                gv_view.scales_into(r, gpr, &mut gsc);
3771                uv_view.scales_into(r, gpr, &mut usc);
3772                for (k, &t) in groups[e].iter().enumerate() {
3773                    if !quantized {
3774                        let x = &xs[t * cols..(t + 1) * cols];
3775                        let gv = q4tp_row_exact(gv_view.nib, r, gpr, x, &gsc);
3776                        let uv = q4tp_row_exact(uv_view.nib, r, gpr, x, &usc);
3777                        unsafe { *ptrs[offs[e] + k].at(r) = (gv / (1.0 + (-gv).exp())) * uv };
3778                        continue;
3779                    }
3780                    let act = &acts[t];
3781                    let mut gv = dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, &gsc) * act.sx;
3782                    let mut uv = dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, &usc) * act.sx;
3783                    for &(j, xv) in &act.outliers {
3784                        let og = q4tp_outlier(gv_view.nib, r, gpr, j, &gsc);
3785                        let ou = q4tp_outlier(uv_view.nib, r, gpr, j, &usc);
3786                        gv += og.0 * og.1 * xv;
3787                        uv += ou.0 * ou.1 * xv;
3788                    }
3789                    let silu_g = gv / (1.0 + (-gv).exp());
3790                    // SAFETY: one worker owns each (expert, row) cell of
3791                    // every output of the expert's group.
3792                    unsafe { *ptrs[offs[e] + k].at(r) = silu_g * uv };
3793                }
3794            }
3795        };
3796        dispatch_rows(pool, pairs.len() * inter, &run);
3797        true
3798    }
3799
3800    /// The per-(expert, token) down terms `moe_down_many` weights and sums,
3801    /// for SEVERAL tokens: `outs[p][o] = down_e[o] · gs[p]` (int8 split of
3802    /// `gs[p]`, VNNI dot, outlier terms — bit-identical to that kernel's
3803    /// `d`), each down row read once for its expert's whole group. The
3804    /// caller sums `w·d` per token in its route order, which reproduces
3805    /// `moe_down_many`'s f32 sequence exactly. Layout as `moe_gate_up_rows`.
3806    pub fn moe_down_rows(
3807        downs: &[&QTensor],
3808        group_lens: &[usize],
3809        gs: &[Vec<f32>],
3810        outs: &mut [Vec<f32>],
3811        pool: Option<&Pool>,
3812    ) -> bool {
3813        if downs.is_empty() || downs.len() != group_lens.len() {
3814            return false;
3815        }
3816        let rows = downs[0].rows();
3817        let cols = downs[0].cols();
3818        let n_pairs: usize = group_lens.iter().sum();
3819        if cols == 0 || cols % GROUP_SIZE != 0 || gs.len() != n_pairs || outs.len() != n_pairs {
3820            return false;
3821        }
3822        let gpr = cols / GROUP_SIZE;
3823        let mut views = Vec::with_capacity(downs.len());
3824        for d in downs {
3825            if d.has_prism_contract()
3826                || !matches!(
3827                    d,
3828                    Self::Mapped {
3829                        dtype: TensorDtype::Q4TiledP,
3830                        ..
3831                    }
3832                )
3833                || d.rows() != rows
3834                || d.cols() != cols
3835            {
3836                return false;
3837            }
3838            views.push(Q4tpView::new(d.quant_bytes(), rows, cols));
3839        }
3840        if gs.iter().any(|g| g.len() != cols) || outs.iter().any(|o| o.len() != rows) {
3841            return false;
3842        }
3843        let quantized = a8w8_enabled();
3844        let acts: Vec<SplitAct> = if quantized {
3845            gs.iter().map(|g| split_act(g)).collect()
3846        } else {
3847            Vec::new()
3848        };
3849        let mut offs = Vec::with_capacity(group_lens.len());
3850        let mut o = 0usize;
3851        for &l in group_lens {
3852            offs.push(o);
3853            o += l;
3854        }
3855        let ptrs: Vec<SendMut> = outs.iter_mut().map(|o| SendMut(o.as_mut_ptr())).collect();
3856        let (views, ptrs, acts, offs) = (&views, &ptrs, &acts, &offs);
3857        let run = |start: usize, end: usize| {
3858            let mut sc = vec![0f32; gpr];
3859            for flat in start..end {
3860                let (e, r) = (flat / rows, flat % rows);
3861                let v = &views[e];
3862                v.scales_into(r, gpr, &mut sc);
3863                for k in 0..group_lens[e] {
3864                    if !quantized {
3865                        let d = q4tp_row_exact(v.nib, r, gpr, &gs[offs[e] + k], &sc);
3866                        unsafe { *ptrs[offs[e] + k].at(r) = d };
3867                        continue;
3868                    }
3869                    let a = &acts[offs[e] + k];
3870                    let mut d = dot_q4tp_row_i8(v.nib, r, gpr, &a.xq, &sc) * a.sx;
3871                    for &(j, xv) in &a.outliers {
3872                        let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
3873                        d += w * s * xv;
3874                    }
3875                    // SAFETY: one worker owns each (expert, row) cell.
3876                    unsafe { *ptrs[offs[e] + k].at(r) = d };
3877                }
3878            }
3879        };
3880        dispatch_rows(pool, downs.len() * rows, &run);
3881        true
3882    }
3883}
3884
3885/// Batched q8 kernel: same math as qmatvec, the row makes a single
3886/// pass from memory for the whole batch.
3887/// Accelerate CBLAS — the Apple AMX matrix units, the same engine
3888/// llama.cpp's `-ngl 0` prefill rides via ggml-blas.
3889#[cfg(target_os = "macos")]
3890mod accel_blas {
3891    #[link(name = "Accelerate", kind = "framework")]
3892    unsafe extern "C" {
3893        pub fn cblas_sgemm(
3894            order: i32,
3895            trans_a: i32,
3896            trans_b: i32,
3897            m: i32,
3898            n: i32,
3899            k: i32,
3900            alpha: f32,
3901            a: *const f32,
3902            lda: i32,
3903            b: *const f32,
3904            ldb: i32,
3905            beta: f32,
3906            c: *mut f32,
3907            ldc: i32,
3908        );
3909    }
3910}
3911
3912#[cfg(target_os = "macos")]
3913pub(crate) fn accel_gemm_enabled() -> bool {
3914    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
3915    *ON.get_or_init(|| std::env::var("CMF_ACCEL").map(|v| v != "0").unwrap_or(true))
3916}
3917
3918/// Off macOS the "accel" GEMM is the portable NEON micro-kernel below —
3919/// same entry point, so the batched-attention path opens on mobile.
3920#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
3921pub(crate) fn accel_gemm_enabled() -> bool {
3922    true
3923}
3924
3925/// Portable NEON f32 GEMM (row-major, optional Bᵀ): a 4×8 fmla
3926/// micro-kernel with A broadcast against B panels — the mobile stand-in
3927/// for Accelerate in the batched causal attention (QKᵀ and P·V). Not a
3928/// BLAS: shapes here are the attention panels (m ≤ heads·chunk,
3929/// k = head_dim or context), and the goal is removing the per-position
3930/// quadratic wall, not peak GEMM.
3931#[cfg(target_arch = "aarch64")]
3932#[allow(clippy::too_many_arguments)]
3933pub(crate) fn neon_gemm_rm(
3934    m: usize,
3935    n: usize,
3936    k: usize,
3937    alpha: f32,
3938    a: &[f32],
3939    lda: usize,
3940    b_mat: &[f32],
3941    ldb: usize,
3942    b_rows_are_n: bool,
3943    c: &mut [f32],
3944    ldc: usize,
3945) {
3946    debug_assert!(a.len() >= (m - 1) * lda + k);
3947    debug_assert!(c.len() >= (m - 1) * ldc + n);
3948    // SAFETY: bounds asserted above; NEON is baseline on aarch64.
3949    unsafe {
3950        use core::arch::aarch64::*;
3951        let mut i = 0usize;
3952        while i < m {
3953            let mi = (m - i).min(4);
3954            let mut j = 0usize;
3955            while j < n {
3956                let nj = (n - j).min(8);
3957                if mi == 4 && nj == 8 {
3958                    let (mut c0a, mut c0b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
3959                    let (mut c1a, mut c1b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
3960                    let (mut c2a, mut c2b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
3961                    let (mut c3a, mut c3b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
3962                    for p in 0..k {
3963                        let (b0, b1) = if b_rows_are_n {
3964                            // B is [n, k]: column p of Bᵀ = element p of
3965                            // eight consecutive B rows — gathered.
3966                            let base = b_mat.as_ptr().add(j * ldb + p);
3967                            let g = |o: usize| *base.add(o * ldb);
3968                            ([g(0), g(1), g(2), g(3)], [g(4), g(5), g(6), g(7)])
3969                        } else {
3970                            let base = b_mat.as_ptr().add(p * ldb + j);
3971                            (
3972                                [*base, *base.add(1), *base.add(2), *base.add(3)],
3973                                [*base.add(4), *base.add(5), *base.add(6), *base.add(7)],
3974                            )
3975                        };
3976                        let bv0 = vld1q_f32(b0.as_ptr());
3977                        let bv1 = vld1q_f32(b1.as_ptr());
3978                        let a0 = vdupq_n_f32(*a.as_ptr().add(i * lda + p));
3979                        let a1 = vdupq_n_f32(*a.as_ptr().add((i + 1) * lda + p));
3980                        let a2 = vdupq_n_f32(*a.as_ptr().add((i + 2) * lda + p));
3981                        let a3 = vdupq_n_f32(*a.as_ptr().add((i + 3) * lda + p));
3982                        c0a = vfmaq_f32(c0a, a0, bv0);
3983                        c0b = vfmaq_f32(c0b, a0, bv1);
3984                        c1a = vfmaq_f32(c1a, a1, bv0);
3985                        c1b = vfmaq_f32(c1b, a1, bv1);
3986                        c2a = vfmaq_f32(c2a, a2, bv0);
3987                        c2b = vfmaq_f32(c2b, a2, bv1);
3988                        c3a = vfmaq_f32(c3a, a3, bv0);
3989                        c3b = vfmaq_f32(c3b, a3, bv1);
3990                    }
3991                    let al = vdupq_n_f32(alpha);
3992                    for (r, (ca, cb)) in [(c0a, c0b), (c1a, c1b), (c2a, c2b), (c3a, c3b)]
3993                        .iter()
3994                        .enumerate()
3995                    {
3996                        let dst = c.as_mut_ptr().add((i + r) * ldc + j);
3997                        vst1q_f32(dst, vmulq_f32(*ca, al));
3998                        vst1q_f32(dst.add(4), vmulq_f32(*cb, al));
3999                    }
4000                } else {
4001                    for r in 0..mi {
4002                        for q in 0..nj {
4003                            let mut acc = 0f32;
4004                            for p in 0..k {
4005                                let bv = if b_rows_are_n {
4006                                    b_mat[(j + q) * ldb + p]
4007                                } else {
4008                                    b_mat[p * ldb + j + q]
4009                                };
4010                                acc += a[(i + r) * lda + p] * bv;
4011                            }
4012                            c[(i + r) * ldc + j + q] = acc * alpha;
4013                        }
4014                    }
4015                }
4016                j += nj;
4017            }
4018            i += mi;
4019        }
4020    }
4021}
4022
4023/// Off-macOS aarch64: the batched attention rides the NEON micro-GEMM.
4024#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
4025#[allow(clippy::too_many_arguments)]
4026pub(crate) fn sgemm_rm(
4027    m: usize,
4028    n: usize,
4029    k: usize,
4030    alpha: f32,
4031    a: &[f32],
4032    lda: usize,
4033    b_mat: &[f32],
4034    ldb: usize,
4035    b_rows_are_n: bool,
4036    c: &mut [f32],
4037    ldc: usize,
4038) {
4039    neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
4040}
4041
4042/// Row-major f32 GEMM, exposed for offline tools (the AWNP pass builds a
4043/// per-layer projection and applies it to every expert; a naive triple loop
4044/// would turn a two-minute job into half an hour).
4045#[allow(clippy::too_many_arguments)]
4046pub fn sgemm_public(
4047    m: usize,
4048    n: usize,
4049    k: usize,
4050    alpha: f32,
4051    a: &[f32],
4052    lda: usize,
4053    b_mat: &[f32],
4054    ldb: usize,
4055    b_rows_are_n: bool,
4056    c: &mut [f32],
4057    ldc: usize,
4058) {
4059    #[cfg(any(target_os = "macos", target_arch = "aarch64"))]
4060    {
4061        sgemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
4062    }
4063    // x86 without Accelerate has no sgemm_rm: the specialized paths there are
4064    // quantized kernels, not an f32 GEMM. Only the offline AWNP pass reaches
4065    // this, so correctness matters and throughput does not — a triple loop is
4066    // the honest fallback rather than a reason to make the tool macOS-only.
4067    #[cfg(not(any(target_os = "macos", target_arch = "aarch64")))]
4068    {
4069        for i in 0..m {
4070            for j in 0..n {
4071                let mut acc = 0f32;
4072                for p in 0..k {
4073                    let bv = if b_rows_are_n {
4074                        b_mat[j * ldb + p]
4075                    } else {
4076                        b_mat[p * ldb + j]
4077                    };
4078                    acc += a[i * lda + p] * bv;
4079                }
4080                c[i * ldc + j] = alpha * acc;
4081            }
4082        }
4083    }
4084}
4085
4086/// Row-major f32 GEMM on Accelerate: C[m,n] = alpha·A[m,k] × B(ᵀ).
4087/// `b_rows_are_n` = true multiplies by Bᵀ where B is stored [n, k].
4088#[cfg(target_os = "macos")]
4089#[allow(clippy::too_many_arguments)]
4090pub(crate) fn sgemm_rm(
4091    m: usize,
4092    n: usize,
4093    k: usize,
4094    alpha: f32,
4095    a: &[f32],
4096    lda: usize,
4097    b_mat: &[f32],
4098    ldb: usize,
4099    b_rows_are_n: bool,
4100    c: &mut [f32],
4101    ldc: usize,
4102) {
4103    debug_assert!(a.len() >= (m - 1) * lda + k);
4104    debug_assert!(c.len() >= (m - 1) * ldc + n);
4105    // Test hook: route the attention GEMMs through the portable NEON
4106    // micro-kernel ON APPLE SILICON — how the mobile batched attend is
4107    // measured without a phone in the loop. (Intel macOS has no NEON —
4108    // the hook is a no-op there, Accelerate continues below.)
4109    #[cfg(target_arch = "aarch64")]
4110    if std::env::var("CMF_FORCE_NEON_GEMM")
4111        .map(|v| v == "1")
4112        .unwrap_or(false)
4113    {
4114        return neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
4115    }
4116    unsafe {
4117        accel_blas::cblas_sgemm(
4118            101, // RowMajor
4119            111, // NoTrans A
4120            if b_rows_are_n { 112 } else { 111 },
4121            m as i32,
4122            n as i32,
4123            k as i32,
4124            alpha,
4125            a.as_ptr(),
4126            lda as i32,
4127            b_mat.as_ptr(),
4128            ldb as i32,
4129            0.0,
4130            c.as_mut_ptr(),
4131            ldc as i32,
4132        );
4133    }
4134}
4135
4136/// Prefill GEMM through Accelerate (macOS): dequantize q8 rows into
4137/// f32 tiles (scale folded in, pool-parallel) and multiply each tile
4138/// on the AMX with one row-major sgemm. Tiles live in cache, weights
4139/// stream once. Numerics are f32-GEMM (not the int8 dot): prefill
4140/// logits shift within f32 rounding — tolerance-class, like every
4141/// reduction-order change; decode (M=1) never takes this path.
4142#[cfg(target_os = "macos")]
4143fn qmatmat_accel(
4144    q: &[u8],
4145    row_scale: &[f32],
4146    pre: &[std::borrow::Cow<'_, [f32]>],
4147    rows: usize,
4148    cols: usize,
4149    out: &mut [f32],
4150    pool: Option<&Pool>,
4151) {
4152    // NOTE: double-buffering the dequant against the sgemm (a scoped
4153    // thread driving the pool on tile k+1 while the caller multiplies
4154    // tile k) was tried and LOST ~6%: Accelerate's sgemm is itself
4155    // multithreaded, and the dequant workers just steal its cores.
4156    const TR: usize = 2048;
4157    let b = pre.len();
4158    thread_local! {
4159        static XPANEL: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
4160        static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
4161    }
4162    XPANEL.with(|xp| {
4163        WTILE.with(|wt| {
4164            let mut xpanel = xp.borrow_mut();
4165            xpanel.clear();
4166            for x in pre {
4167                xpanel.extend_from_slice(x);
4168            }
4169            let mut wtile = wt.borrow_mut();
4170            wtile.resize(TR * cols, 0.0);
4171            let mut r0 = 0usize;
4172            while r0 < rows {
4173                let tr = TR.min(rows - r0);
4174                // Dequant the tile (scale folded) — pool-parallel.
4175                let wt_addr = SendMut(wtile.as_mut_ptr());
4176                let run = |start: usize, end: usize| {
4177                    for r in start..end {
4178                        let row = &q[(r0 + r) * cols..(r0 + r + 1) * cols];
4179                        let s = row_scale[r0 + r];
4180                        // SAFETY: workers cover disjoint r ranges.
4181                        let dst =
4182                            unsafe { std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols) };
4183                        for (d, &v) in dst.iter_mut().zip(row) {
4184                            *d = (v as i8) as f32 * s;
4185                        }
4186                    }
4187                };
4188                dispatch_rows(pool, tr, &run);
4189                // C[b, tr] (at column r0 of out[b, rows]) = X · Wtileᵀ
4190                unsafe {
4191                    accel_blas::cblas_sgemm(
4192                        101, // RowMajor
4193                        111, // NoTrans A
4194                        112, // Trans B
4195                        b as i32,
4196                        tr as i32,
4197                        cols as i32,
4198                        1.0,
4199                        xpanel.as_ptr(),
4200                        cols as i32,
4201                        wtile.as_ptr(),
4202                        cols as i32,
4203                        0.0,
4204                        out.as_mut_ptr().add(r0),
4205                        rows as i32,
4206                    );
4207                }
4208                r0 += tr;
4209            }
4210        })
4211    });
4212}
4213
4214fn qmatmat(
4215    q: &[u8],
4216    row_scale: &[f32],
4217    pre: &[std::borrow::Cow<'_, [f32]>],
4218    rows: usize,
4219    cols: usize,
4220    out: &mut [f32],
4221    pool: Option<&Pool>,
4222) {
4223    let b = pre.len();
4224    debug_assert_eq!(out.len(), b * rows);
4225    // Big prefill batches ride the AMX (roadmap PR3): the row×batch
4226    // SDOT loop below peaks near the CPU's dot throughput, an order
4227    // below the matrix units. Small tensors and tiny test models stay
4228    // on the exact integer path.
4229    #[cfg(target_os = "macos")]
4230    if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
4231        qmatmat_accel(q, row_scale, pre, rows, cols, out, pool);
4232        return;
4233    }
4234    #[cfg(target_arch = "aarch64")]
4235    if sdot_enabled() {
4236        let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
4237        let out_addr = SendMut(out.as_mut_ptr());
4238        // Blocked 2×4 (mobile prefill: no AMX to fall back on — this
4239        // path IS the ARM prefill GEMM off Apple silicon).
4240        let blocked_ok = blocked_enabled();
4241        let use_i8mm = i8mm_enabled();
4242        if blocked_ok {
4243            let run = |start: usize, end: usize| {
4244                let mut o = start;
4245                while o < end {
4246                    if o + 2 <= end {
4247                        let r0 = &q[o * cols..(o + 1) * cols];
4248                        let r1 = &q[(o + 1) * cols..(o + 2) * cols];
4249                        let mut bi = 0usize;
4250                        while bi + 4 <= acts.len() {
4251                            let xs = [
4252                                acts[bi].xq.as_slice(),
4253                                acts[bi + 1].xq.as_slice(),
4254                                acts[bi + 2].xq.as_slice(),
4255                                acts[bi + 3].xq.as_slice(),
4256                            ];
4257                            let d = if use_i8mm {
4258                                unsafe { dot_i8_smmla_2x4(r0, r1, xs) }
4259                            } else {
4260                                unsafe { dot_i8_sdot_2x4(r0, r1, xs) }
4261                            };
4262                            for (r, row) in [r0, r1].into_iter().enumerate() {
4263                                for k in 0..4 {
4264                                    let act = &acts[bi + k];
4265                                    let mut v = d[r][k] as f32 * act.sx;
4266                                    for &(j, xv) in &act.outliers {
4267                                        v += (row[j] as i8) as f32 * xv;
4268                                    }
4269                                    unsafe {
4270                                        *out_addr.at((bi + k) * rows + o + r) = v * row_scale[o + r]
4271                                    };
4272                                }
4273                            }
4274                            bi += 4;
4275                        }
4276                        while bi < acts.len() {
4277                            for (r, row) in [r0, r1].into_iter().enumerate() {
4278                                let v = row_dot_sdot(row, &acts[bi]) * row_scale[o + r];
4279                                unsafe { *out_addr.at(bi * rows + o + r) = v };
4280                            }
4281                            bi += 1;
4282                        }
4283                        o += 2;
4284                    } else {
4285                        let row = &q[o * cols..(o + 1) * cols];
4286                        for (bi, act) in acts.iter().enumerate() {
4287                            let v = row_dot_sdot(row, act) * row_scale[o];
4288                            unsafe { *out_addr.at(bi * rows + o) = v };
4289                        }
4290                        o += 1;
4291                    }
4292                }
4293            };
4294            dispatch_rows(pool, rows, &run);
4295            return;
4296        }
4297        let run = |start: usize, end: usize| {
4298            for o in start..end {
4299                let row = &q[o * cols..(o + 1) * cols];
4300                for (bi, act) in acts.iter().enumerate() {
4301                    let v = row_dot_sdot(row, act) * row_scale[o];
4302                    unsafe { *out_addr.at(bi * rows + o) = v };
4303                }
4304            }
4305        };
4306        dispatch_rows(pool, rows, &run);
4307        return;
4308    }
4309    // x86 A8W8 batch. Non-VNNI parts take the BLOCKED 2×4 kernel
4310    // (roadmap P0: two weight rows' abs() stay in registers across four
4311    // activation streams); VNNI machines keep the per-row bias-trick
4312    // dot, which is already throughput-bound there.
4313    #[cfg(target_arch = "x86_64")]
4314    if avx2_a8w8_enabled() {
4315        let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
4316        let out_addr = SendMut(out.as_mut_ptr());
4317        // CMF_X86_BLOCKED=0 forces the per-row path (paired in-process
4318        // A/B on noisy shared-vCPU hosts).
4319        let blocked_ok = blocked_enabled();
4320        if !avx512vnni_enabled() && blocked_ok && !row_exact() {
4321            let run = |start: usize, end: usize| {
4322                let mut o = start;
4323                while o < end {
4324                    if o + 2 <= end {
4325                        let r0 = &q[o * cols..(o + 1) * cols];
4326                        let r1 = &q[(o + 1) * cols..(o + 2) * cols];
4327                        let mut bi = 0usize;
4328                        while bi + 4 <= acts.len() {
4329                            let xs = [
4330                                acts[bi].xq.as_slice(),
4331                                acts[bi + 1].xq.as_slice(),
4332                                acts[bi + 2].xq.as_slice(),
4333                                acts[bi + 3].xq.as_slice(),
4334                            ];
4335                            let d = unsafe { dot_i8_i8_avx2_2x4(r0, r1, xs) };
4336                            for (r, row) in [r0, r1].into_iter().enumerate() {
4337                                for k in 0..4 {
4338                                    let act = &acts[bi + k];
4339                                    let mut v = d[r][k] as f32 * act.sx;
4340                                    for &(j, xv) in &act.outliers {
4341                                        v += (row[j] as i8) as f32 * xv;
4342                                    }
4343                                    unsafe {
4344                                        *out_addr.at((bi + k) * rows + o + r) = v * row_scale[o + r]
4345                                    };
4346                                }
4347                            }
4348                            bi += 4;
4349                        }
4350                        while bi < acts.len() {
4351                            for (r, row) in [r0, r1].into_iter().enumerate() {
4352                                let v = row_dot_avx2(row, &acts[bi]) * row_scale[o + r];
4353                                unsafe { *out_addr.at(bi * rows + o + r) = v };
4354                            }
4355                            bi += 1;
4356                        }
4357                        o += 2;
4358                    } else {
4359                        let row = &q[o * cols..(o + 1) * cols];
4360                        for (bi, act) in acts.iter().enumerate() {
4361                            let v = row_dot_avx2(row, act) * row_scale[o];
4362                            unsafe { *out_addr.at(bi * rows + o) = v };
4363                        }
4364                        o += 1;
4365                    }
4366                }
4367            };
4368            dispatch_rows(pool, rows, &run);
4369            return;
4370        }
4371        let run = |start: usize, end: usize| {
4372            for o in start..end {
4373                let row = &q[o * cols..(o + 1) * cols];
4374                for (bi, act) in acts.iter().enumerate() {
4375                    let v = row_dot_avx2(row, act) * row_scale[o];
4376                    unsafe { *out_addr.at(bi * rows + o) = v };
4377                }
4378            }
4379        };
4380        dispatch_rows(pool, rows, &run);
4381        return;
4382    }
4383    let out_addr = SendMut(out.as_mut_ptr());
4384    let run = |start: usize, end: usize| {
4385        for o in start..end {
4386            let row = &q[o * cols..(o + 1) * cols];
4387            for (bi, x) in pre.iter().enumerate() {
4388                let mut acc = 0f32;
4389                for j in 0..cols {
4390                    acc += (row[j] as i8) as f32 * x[j];
4391                }
4392                unsafe { *out_addr.at(bi * rows + o) = acc * row_scale[o] };
4393            }
4394        }
4395    };
4396    dispatch_rows(pool, rows, &run);
4397}
4398
4399/// Split rows across pool workers (shared qmatvec pattern). Self-balancing
4400/// — see `Pool::run_rows` for why a static 1/n split is wrong here.
4401fn dispatch_rows(pool: Option<&Pool>, rows: usize, run: &(dyn Fn(usize, usize) + Sync)) {
4402    match pool {
4403        Some(pool) if rows >= 256 => pool.run_rows(rows, run),
4404        _ => run(0, rows),
4405    }
4406}
4407
4408/// Split a q4_block blob into (packed nibbles, f16 group scales).
4409fn q4_split(bytes: &[u8], rows: usize, cols: usize) -> (&[u8], &[u8]) {
4410    let groups = rows * cols / GROUP_SIZE;
4411    bytes.split_at(groups * 16)
4412}
4413
4414/// SIMD unpack for the dominant vbit width B=4 (94% of rows on the
4415/// log2-shape calibration): 16 packed bytes -> 32 centered i8 values.
4416/// vbit packs MSB-first, so the HIGH nibble is the even element
4417/// (opposite of q4_block's lo-first interleave). Centering is u-7.
4418#[inline]
4419fn vbit_fill4(data: &[u8], buf: &mut [u8]) {
4420    #[cfg(target_arch = "aarch64")]
4421    unsafe {
4422        return vbit_fill4_neon(data, buf);
4423    }
4424    #[cfg(target_arch = "x86_64")]
4425    if avx2_enabled() {
4426        return unsafe { vbit_fill4_avx2(data, buf) };
4427    }
4428    #[allow(unreachable_code)]
4429    for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
4430        let u = unpack8::<4>(&data[blk * 4..]);
4431        for k in 0..8 {
4432            chunk[k] = (u[k] - 7) as i8 as u8;
4433        }
4434    }
4435}
4436
4437#[cfg(target_arch = "aarch64")]
4438#[target_feature(enable = "neon")]
4439unsafe fn vbit_fill4_neon(data: &[u8], buf: &mut [u8]) {
4440    // SAFETY: buf.len() is a multiple of GROUP_SIZE=32; data holds
4441    // buf.len()/2 packed bytes (validated at load).
4442    unsafe {
4443        use core::arch::aarch64::*;
4444        let n = buf.len();
4445        let mask = vdupq_n_u8(0x0F);
4446        let seven = vdupq_n_s8(7);
4447        let mut g = 0usize;
4448        while g * 32 + 32 <= n {
4449            let b = vld1q_u8(data.as_ptr().add(g * 16));
4450            let hi = vshrq_n_u8::<4>(b);
4451            let lo = vandq_u8(b, mask);
4452            let z0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(hi, lo)), seven);
4453            let z1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(hi, lo)), seven);
4454            vst1q_u8(buf.as_mut_ptr().add(g * 32), vreinterpretq_u8_s8(z0));
4455            vst1q_u8(buf.as_mut_ptr().add(g * 32 + 16), vreinterpretq_u8_s8(z1));
4456            g += 1;
4457        }
4458    }
4459}
4460
4461#[cfg(target_arch = "x86_64")]
4462#[target_feature(enable = "avx2")]
4463unsafe fn vbit_fill4_avx2(data: &[u8], buf: &mut [u8]) {
4464    // SAFETY: see vbit_fill4_neon.
4465    unsafe {
4466        use core::arch::x86_64::*;
4467        let n = buf.len();
4468        let mask = _mm_set1_epi8(0x0F);
4469        let seven = _mm256_set1_epi8(7);
4470        let mut g = 0usize;
4471        while g * 32 + 32 <= n {
4472            let b = _mm_loadu_si128(data.as_ptr().add(g * 16) as *const __m128i);
4473            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), mask);
4474            let lo = _mm_and_si128(b, mask);
4475            let z = _mm256_sub_epi8(
4476                _mm256_set_m128i(_mm_unpackhi_epi8(hi, lo), _mm_unpacklo_epi8(hi, lo)),
4477                seven,
4478            );
4479            _mm256_storeu_si256(buf.as_mut_ptr().add(g * 32) as *mut __m256i, z);
4480            g += 1;
4481        }
4482    }
4483}
4484
4485/// Unpack 8 MSB-first B-bit values from exactly B bytes (fixed shifts —
4486/// no serial bit-buffer, auto-vectorizable). Every 32-value group starts
4487/// byte-aligned (32·B/8 is integral for B∈3..8), so groups decompose
4488/// into 4 such blocks.
4489#[inline(always)]
4490fn unpack8<const B: usize>(data: &[u8]) -> [i32; 8] {
4491    let mut acc = 0u64;
4492    for i in 0..B {
4493        acc = (acc << 8) | data[i] as u64;
4494    }
4495    let mask = (1u64 << B) - 1;
4496    let mut out = [0i32; 8];
4497    for (k, o) in out.iter_mut().enumerate() {
4498        *o = ((acc >> ((7 - k) * B)) & mask) as i32;
4499    }
4500    out
4501}
4502
4503/// Fused vbit matvec straight from the mapped bytes (spec §3, P13
4504/// FIG.3): [u8 bits: rows][f16 scales: rows·cols/32][bit-packed rows,
4505/// MSB-first, byte-padded]. Row data offsets are precomputed at load
4506/// (`vbit_row_offsets`) — the per-call prefix scan was O(rows) pure
4507/// overhead on every matvec.
4508#[allow(clippy::too_many_arguments)]
4509fn vbitmatvec(
4510    bytes: &[u8],
4511    offsets: &[usize],
4512    x: &[f32],
4513    rows: usize,
4514    cols: usize,
4515    out: &mut [f32],
4516    pool: Option<&Pool>,
4517) {
4518    debug_assert_eq!(out.len(), rows);
4519    debug_assert_eq!(offsets.len(), rows + 1);
4520
4521    // SDOT path: unpack the row to centered i8 once, then per-group
4522    // int8 dot against the quantized activations — same A8W8 contract
4523    // as q8 (bounded noise; CMF_SDOT=0 keeps the exact scalar path).
4524    if a8w8_enabled() {
4525        let act = split_act(x);
4526        let out_addr = SendMut(out.as_mut_ptr());
4527        let run = move |start: usize, end: usize| {
4528            vbit_range_a8w8(bytes, offsets, x, &act, rows, cols, out_addr, start, end)
4529        };
4530        dispatch_rows(pool, rows, &run);
4531        return;
4532    }
4533
4534    let out_addr = SendMut(out.as_mut_ptr());
4535    let run = move |start: usize, end: usize| {
4536        vbit_range_f32(bytes, offsets, x, rows, cols, out_addr, start, end)
4537    };
4538    dispatch_rows(pool, rows, &run);
4539}
4540
4541/// One vbit row range via the A8W8 int8 path — kernel body of
4542/// `vbitmatvec`, extracted so multi-matrix jobs can drive it for
4543/// several tensors in one dispatch (b=8 rows go exact f32).
4544#[allow(clippy::too_many_arguments)]
4545fn vbit_range_a8w8(
4546    bytes: &[u8],
4547    offsets: &[usize],
4548    x: &[f32],
4549    act: &SplitAct,
4550    rows: usize,
4551    cols: usize,
4552    out: SendMut,
4553    start: usize,
4554    end: usize,
4555) {
4556    let ng = cols / GROUP_SIZE;
4557    let bits = &bytes[..rows];
4558    let sc_off = rows;
4559    let row_dot = |r: usize| -> f32 {
4560        let b = bits[r] as usize;
4561        let l = (1i32 << (b - 1)) - 1;
4562        let mask = (1u64 << b) - 1;
4563        let data = &bytes[offsets[r]..offsets[r + 1]];
4564        if b == 8 {
4565            // u−L reaches 128 → does not fit i8; exact f32 path.
4566            let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
4567            let mut dot = 0f32;
4568            for g in 0..ng {
4569                let so = (r * ng + g) * 2;
4570                let sgf = f16_to_f32(u16::from_le_bytes([
4571                    bytes[sc_off + so],
4572                    bytes[sc_off + so + 1],
4573                ]));
4574                let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
4575                let mut gd = 0f32;
4576                for &xv in xg.iter() {
4577                    if nbits < 8 {
4578                        acc = (acc << 8) | data[idx] as u64;
4579                        idx += 1;
4580                        nbits += 8;
4581                    }
4582                    let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
4583                    nbits -= 8;
4584                    gd += (u - l) as f32 * xv;
4585                }
4586                dot += gd * sgf;
4587            }
4588            return dot;
4589        }
4590        // Per-worker scratch: this closure runs for every row of the
4591        // tensor (lm_head ≈ 150k rows/token) — a heap allocation per
4592        // row was measurable pure overhead.
4593        thread_local! {
4594            static VBIT_SCRATCH: std::cell::RefCell<Vec<u8>> =
4595                const { std::cell::RefCell::new(Vec::new()) };
4596        }
4597        #[inline(always)]
4598        fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
4599            for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
4600                let u = unpack8::<B>(&data[blk * B..]);
4601                for k in 0..8 {
4602                    chunk[k] = (u[k] - l) as i8 as u8;
4603                }
4604            }
4605        }
4606        let _ = mask;
4607        VBIT_SCRATCH.with(|scratch| {
4608            let mut buf = scratch.borrow_mut();
4609            buf.resize(cols, 0);
4610            match b {
4611                3 => fill::<3>(data, l, &mut buf),
4612                4 => vbit_fill4(data, &mut buf),
4613                5 => fill::<5>(data, l, &mut buf),
4614                6 => fill::<6>(data, l, &mut buf),
4615                _ => unreachable!(),
4616            }
4617            let mut dot = 0f32;
4618            for g in 0..ng {
4619                let so = (r * ng + g) * 2;
4620                let s = f16_to_f32(u16::from_le_bytes([
4621                    bytes[sc_off + so],
4622                    bytes[sc_off + so + 1],
4623                ]));
4624                let d = dot_i8_i8(
4625                    &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
4626                    &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
4627                ) as f32
4628                    * act.sx;
4629                dot += d * s;
4630            }
4631            for &(j, xv) in &act.outliers {
4632                let so = (r * ng + j / GROUP_SIZE) * 2;
4633                let s = f16_to_f32(u16::from_le_bytes([
4634                    bytes[sc_off + so],
4635                    bytes[sc_off + so + 1],
4636                ]));
4637                // xq is zeroed at outlier slots — add the exact term.
4638                dot += (buf[j] as i8) as f32 * s * xv;
4639            }
4640            dot
4641        })
4642    };
4643    for r in start..end {
4644        // SAFETY: disjoint row ranges per worker.
4645        unsafe { *out.at(r) = row_dot(r) };
4646    }
4647}
4648
4649/// Exact scalar vbit row range (same extraction, non-SDOT path).
4650#[allow(clippy::too_many_arguments)]
4651fn vbit_range_f32(
4652    bytes: &[u8],
4653    offsets: &[usize],
4654    x: &[f32],
4655    rows: usize,
4656    cols: usize,
4657    out: SendMut,
4658    start: usize,
4659    end: usize,
4660) {
4661    let ng = cols / GROUP_SIZE;
4662    let bits = &bytes[..rows];
4663    let sc_off = rows;
4664    // Per-bit-width specialized inner loops: the compiler unrolls the
4665    // constant shifts (the generic bit-buffer loop was branch-bound —
4666    // 5.6 vs 13.2 tok/s q4 on the 0.8B).
4667    #[inline(always)]
4668    fn dot_row<const B: usize>(
4669        data: &[u8],
4670        bytes: &[u8],
4671        sc_off: usize,
4672        r: usize,
4673        ng: usize,
4674        x: &[f32],
4675    ) -> f32 {
4676        let l = ((1i32 << (B - 1)) - 1) as f32;
4677        let gbytes = GROUP_SIZE * B / 8;
4678        let mut dot = 0f32;
4679        for g in 0..ng {
4680            let so = (r * ng + g) * 2;
4681            let s = f16_to_f32(u16::from_le_bytes([
4682                bytes[sc_off + so],
4683                bytes[sc_off + so + 1],
4684            ]));
4685            let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
4686            let gd0 = &data[g * gbytes..(g + 1) * gbytes];
4687            let mut gd = 0f32;
4688            for blk in 0..GROUP_SIZE / 8 {
4689                let u = unpack8::<B>(&gd0[blk * B..]);
4690                let xb = &xg[blk * 8..blk * 8 + 8];
4691                for k in 0..8 {
4692                    gd += (u[k] as f32 - l) * xb[k];
4693                }
4694            }
4695            dot += gd * s;
4696        }
4697        dot
4698    }
4699    for r in start..end {
4700        let data = &bytes[offsets[r]..offsets[r + 1]];
4701        let v = match bits[r] {
4702            3 => dot_row::<3>(data, bytes, sc_off, r, ng, x),
4703            4 => dot_row::<4>(data, bytes, sc_off, r, ng, x),
4704            5 => dot_row::<5>(data, bytes, sc_off, r, ng, x),
4705            6 => dot_row::<6>(data, bytes, sc_off, r, ng, x),
4706            8 => dot_row::<8>(data, bytes, sc_off, r, ng, x),
4707            b => unreachable!("vbit bit-width {b} (validated at load)"),
4708        };
4709        // SAFETY: disjoint row ranges per worker.
4710        unsafe { *out.at(r) = v };
4711    }
4712}
4713
4714/// Fused two-input vbit matvec: each row is unpacked from the mmap ONCE
4715/// and dotted against BOTH activations (MTP verify / pair prefill used
4716/// to run two full matvecs — double weight traffic and double unpack).
4717/// Per-input math is identical to `vbitmatvec` → same accuracy contract.
4718#[allow(clippy::too_many_arguments)]
4719fn vbitmatvec2(
4720    bytes: &[u8],
4721    offsets: &[usize],
4722    x1: &[f32],
4723    x2: &[f32],
4724    rows: usize,
4725    cols: usize,
4726    o1: &mut [f32],
4727    o2: &mut [f32],
4728    pool: Option<&Pool>,
4729) {
4730    debug_assert_eq!(o1.len(), rows);
4731    debug_assert_eq!(o2.len(), rows);
4732
4733    if a8w8_enabled() {
4734        let a1 = split_act(x1);
4735        let a2 = split_act(x2);
4736        let p1 = SendMut(o1.as_mut_ptr());
4737        let p2 = SendMut(o2.as_mut_ptr());
4738        let run = move |start: usize, end: usize| {
4739            vbit_range2_a8w8(
4740                bytes, offsets, x1, x2, &a1, &a2, rows, cols, p1, p2, start, end,
4741            )
4742        };
4743        dispatch_rows(pool, rows, &run);
4744        return;
4745    }
4746
4747    let p1 = SendMut(o1.as_mut_ptr());
4748    let p2 = SendMut(o2.as_mut_ptr());
4749    let run = move |start: usize, end: usize| {
4750        vbit_range2_f32(bytes, offsets, x1, x2, rows, cols, p1, p2, start, end)
4751    };
4752    dispatch_rows(pool, rows, &run);
4753}
4754
4755/// Two-input vbit row range via the A8W8 int8 path — kernel body of
4756/// `vbitmatvec2`, extracted for pair multi-matrix jobs (b=8 rows go
4757/// exact f32 for both lanes, bits streamed once).
4758#[allow(clippy::too_many_arguments)]
4759fn vbit_range2_a8w8(
4760    bytes: &[u8],
4761    offsets: &[usize],
4762    x1: &[f32],
4763    x2: &[f32],
4764    a1: &SplitAct,
4765    a2: &SplitAct,
4766    rows: usize,
4767    cols: usize,
4768    p1: SendMut,
4769    p2: SendMut,
4770    start: usize,
4771    end: usize,
4772) {
4773    let ng = cols / GROUP_SIZE;
4774    let bits = &bytes[..rows];
4775    let sc_off = rows;
4776    let row_dots = |r: usize| -> (f32, f32) {
4777        let b = bits[r] as usize;
4778        let l = (1i32 << (b - 1)) - 1;
4779        let data = &bytes[offsets[r]..offsets[r + 1]];
4780        if b == 8 {
4781            // u−L reaches 128 → does not fit i8; exact f32 path,
4782            // bits still streamed once for both lanes.
4783            let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
4784            let (mut d1, mut d2) = (0f32, 0f32);
4785            for g in 0..ng {
4786                let so = (r * ng + g) * 2;
4787                let sgf = f16_to_f32(u16::from_le_bytes([
4788                    bytes[sc_off + so],
4789                    bytes[sc_off + so + 1],
4790                ]));
4791                let (mut g1, mut g2) = (0f32, 0f32);
4792                for k in 0..GROUP_SIZE {
4793                    if nbits < 8 {
4794                        acc = (acc << 8) | data[idx] as u64;
4795                        idx += 1;
4796                        nbits += 8;
4797                    }
4798                    let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
4799                    nbits -= 8;
4800                    let w = (u - l) as f32;
4801                    g1 += w * x1[g * GROUP_SIZE + k];
4802                    g2 += w * x2[g * GROUP_SIZE + k];
4803                }
4804                d1 += g1 * sgf;
4805                d2 += g2 * sgf;
4806            }
4807            return (d1, d2);
4808        }
4809        thread_local! {
4810            static VBIT_SCRATCH2: std::cell::RefCell<Vec<u8>> =
4811                const { std::cell::RefCell::new(Vec::new()) };
4812        }
4813        #[inline(always)]
4814        fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
4815            for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
4816                let u = unpack8::<B>(&data[blk * B..]);
4817                for k in 0..8 {
4818                    chunk[k] = (u[k] - l) as i8 as u8;
4819                }
4820            }
4821        }
4822        VBIT_SCRATCH2.with(|scratch| {
4823            let mut buf = scratch.borrow_mut();
4824            buf.resize(cols, 0);
4825            match b {
4826                3 => fill::<3>(data, l, &mut buf),
4827                4 => vbit_fill4(data, &mut buf),
4828                5 => fill::<5>(data, l, &mut buf),
4829                6 => fill::<6>(data, l, &mut buf),
4830                _ => unreachable!(),
4831            }
4832            let (mut d1, mut d2) = (0f32, 0f32);
4833            for g in 0..ng {
4834                let so = (r * ng + g) * 2;
4835                let s = f16_to_f32(u16::from_le_bytes([
4836                    bytes[sc_off + so],
4837                    bytes[sc_off + so + 1],
4838                ]));
4839                let wg = &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
4840                let v1 = dot_i8_i8(wg, &a1.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a1.sx;
4841                let v2 = dot_i8_i8(wg, &a2.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a2.sx;
4842                d1 += v1 * s;
4843                d2 += v2 * s;
4844            }
4845            for &(j, xv) in &a1.outliers {
4846                let so = (r * ng + j / GROUP_SIZE) * 2;
4847                let s = f16_to_f32(u16::from_le_bytes([
4848                    bytes[sc_off + so],
4849                    bytes[sc_off + so + 1],
4850                ]));
4851                d1 += (buf[j] as i8) as f32 * s * xv;
4852            }
4853            for &(j, xv) in &a2.outliers {
4854                let so = (r * ng + j / GROUP_SIZE) * 2;
4855                let s = f16_to_f32(u16::from_le_bytes([
4856                    bytes[sc_off + so],
4857                    bytes[sc_off + so + 1],
4858                ]));
4859                d2 += (buf[j] as i8) as f32 * s * xv;
4860            }
4861            (d1, d2)
4862        })
4863    };
4864    for r in start..end {
4865        let (v1, v2) = row_dots(r);
4866        // SAFETY: disjoint row ranges per worker.
4867        unsafe {
4868            *p1.at(r) = v1;
4869            *p2.at(r) = v2;
4870        }
4871    }
4872}
4873
4874/// Two-input exact scalar vbit row range (same extraction) —
4875/// per-bit-width specialized, two accumulators per row; per-lane
4876/// accumulation order matches `vbitmatvec` exactly.
4877#[allow(clippy::too_many_arguments)]
4878fn vbit_range2_f32(
4879    bytes: &[u8],
4880    offsets: &[usize],
4881    x1: &[f32],
4882    x2: &[f32],
4883    rows: usize,
4884    cols: usize,
4885    p1: SendMut,
4886    p2: SendMut,
4887    start: usize,
4888    end: usize,
4889) {
4890    let ng = cols / GROUP_SIZE;
4891    let bits = &bytes[..rows];
4892    let sc_off = rows;
4893    #[inline(always)]
4894    #[allow(clippy::too_many_arguments)]
4895    fn dot_row2<const B: usize>(
4896        data: &[u8],
4897        bytes: &[u8],
4898        sc_off: usize,
4899        r: usize,
4900        ng: usize,
4901        x1: &[f32],
4902        x2: &[f32],
4903    ) -> (f32, f32) {
4904        let l = ((1i32 << (B - 1)) - 1) as f32;
4905        let gbytes = GROUP_SIZE * B / 8;
4906        let (mut d1, mut d2) = (0f32, 0f32);
4907        for g in 0..ng {
4908            let so = (r * ng + g) * 2;
4909            let s = f16_to_f32(u16::from_le_bytes([
4910                bytes[sc_off + so],
4911                bytes[sc_off + so + 1],
4912            ]));
4913            let x1g = &x1[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
4914            let x2g = &x2[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
4915            let gd0 = &data[g * gbytes..(g + 1) * gbytes];
4916            let (mut g1, mut g2) = (0f32, 0f32);
4917            for blk in 0..GROUP_SIZE / 8 {
4918                let u = unpack8::<B>(&gd0[blk * B..]);
4919                for k in 0..8 {
4920                    let w = u[k] as f32 - l;
4921                    g1 += w * x1g[blk * 8 + k];
4922                    g2 += w * x2g[blk * 8 + k];
4923                }
4924            }
4925            d1 += g1 * s;
4926            d2 += g2 * s;
4927        }
4928        (d1, d2)
4929    }
4930    for r in start..end {
4931        let data = &bytes[offsets[r]..offsets[r + 1]];
4932        let (v1, v2) = match bits[r] {
4933            3 => dot_row2::<3>(data, bytes, sc_off, r, ng, x1, x2),
4934            4 => dot_row2::<4>(data, bytes, sc_off, r, ng, x1, x2),
4935            5 => dot_row2::<5>(data, bytes, sc_off, r, ng, x1, x2),
4936            6 => dot_row2::<6>(data, bytes, sc_off, r, ng, x1, x2),
4937            8 => dot_row2::<8>(data, bytes, sc_off, r, ng, x1, x2),
4938            b => unreachable!("vbit bit-width {b} (validated at load)"),
4939        };
4940        // SAFETY: disjoint row ranges per worker.
4941        unsafe {
4942            *p1.at(r) = v1;
4943            *p2.at(r) = v2;
4944        }
4945    }
4946}
4947
4948// ───────────────────── q4_tiled kernels (§4.3) ─────────────────────
4949
4950/// One q4_tiled row dot on the A8W8 int8 path: per 32-group the tile
4951/// is ONE sequential read — [f16 scale][16B nibbles] — versus the two
4952/// distant streams of the split layout. Values/order identical to the
4953/// split kernels.
4954#[inline]
4955#[allow(unreachable_code)]
4956fn dot_q4t_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
4957    #[cfg(target_arch = "aarch64")]
4958    unsafe {
4959        return dot_q4t_row_sdot(bytes, r, gpr, xq);
4960    }
4961    #[cfg(target_arch = "x86_64")]
4962    unsafe {
4963        if vnni_tiles_enabled() {
4964            return dot_q4t_row_vnni(bytes, r, gpr, xq);
4965        }
4966        return dot_q4t_row_avx2(bytes, r, gpr, xq);
4967    }
4968    let mut acc = 0f32;
4969    for gi in 0..gpr {
4970        let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
4971        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
4972        let mut d = 0i32;
4973        for (k, &b) in tile[2..].iter().enumerate() {
4974            d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
4975                + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
4976        }
4977        acc += d as f32 * s;
4978    }
4979    acc
4980}
4981
4982#[cfg(target_arch = "aarch64")]
4983#[target_feature(enable = "neon,dotprod")]
4984unsafe fn dot_q4t_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
4985    // SAFETY: callers uphold slice-length contracts (18B tile per group,
4986    // xq.len() == gpr·GROUP_SIZE).
4987    unsafe {
4988        use core::arch::aarch64::*;
4989        use core::arch::asm;
4990        let lomask = vdupq_n_u8(0x0F);
4991        let eight = vdupq_n_s8(8);
4992        let mut acc = 0f32;
4993        for gi in 0..gpr {
4994            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4995            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4996            let b = vld1q_u8(t.add(2));
4997            let lo = vandq_u8(b, lomask);
4998            let hi = vshrq_n_u8::<4>(b);
4999            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
5000            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
5001            let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
5002            let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
5003            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5004            asm!(
5005                "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
5006                "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
5007                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5008                e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
5009                options(pure, nomem, nostack),
5010            );
5011            acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
5012        }
5013        acc
5014    }
5015}
5016
5017#[cfg(target_arch = "x86_64")]
5018#[target_feature(enable = "avx2")]
5019unsafe fn dot_q4t_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
5020    // SAFETY: see dot_q4t_row_sdot.
5021    unsafe {
5022        use core::arch::x86_64::*;
5023        let lomask = _mm_set1_epi8(0x0F);
5024        let eight = _mm256_set1_epi8(8);
5025        let ones = _mm256_set1_epi16(1);
5026        let mut acc = 0f32;
5027        for gi in 0..gpr {
5028            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
5029            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5030            let b = _mm_loadu_si128(t.add(2) as *const __m128i);
5031            let lo = _mm_and_si128(b, lomask);
5032            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
5033            let w = _mm256_sub_epi8(
5034                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5035                eight,
5036            );
5037            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5038            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
5039            let d = _mm256_madd_epi16(p16, ones);
5040            let hi128 = _mm256_extracti128_si256::<1>(d);
5041            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
5042            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
5043            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
5044            acc += _mm_cvtsi128_si32(s32) as f32 * s;
5045        }
5046        acc
5047    }
5048}
5049
5050/// VNNI twin of `dot_q4t_row_avx2`: same unpack, `vpdpbusd` replaces
5051/// the maddubs+madd pair (see `dpbusd_hsum` — sums are bit-identical).
5052/// 256-bit VL encoding, so the VEX `vpsignb` stays usable.
5053#[cfg(target_arch = "x86_64")]
5054#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
5055unsafe fn dot_q4t_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
5056    // SAFETY: see dot_q4t_row_sdot.
5057    unsafe {
5058        use core::arch::x86_64::*;
5059        let lomask = _mm_set1_epi8(0x0F);
5060        let eight = _mm256_set1_epi8(8);
5061        let mut acc = 0f32;
5062        for gi in 0..gpr {
5063            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
5064            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5065            let b = _mm_loadu_si128(t.add(2) as *const __m128i);
5066            let lo = _mm_and_si128(b, lomask);
5067            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
5068            let w = _mm256_sub_epi8(
5069                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5070                eight,
5071            );
5072            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5073            let d = dpbusd_hsum(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
5074            acc += d as f32 * s;
5075        }
5076        acc
5077    }
5078}
5079
5080/// One q4_tiled row against FOUR activation streams: the nibble unpack
5081/// and abs() happen once per group instead of once per (group,
5082/// activation) — the unpack is the dominant per-element cost of the
5083/// tiled format (roadmap P0 portable blocking, q4t leg).
5084#[cfg(target_arch = "x86_64")]
5085// `fma` is NOT implied by `avx2`: without it LLVM lowers _mm256_fmadd_ps
5086// to a libm call per lane — measured 2x slower than the reduction this
5087// kernel replaces. The runtime gate (`avx2_enabled`) already requires
5088// both features, so declaring it here is safe.
5089#[target_feature(enable = "avx2,fma")]
5090unsafe fn dot_q4t_row_1x4_avx2(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
5091    // SAFETY: callers uphold the 18B-tile and xq-length contracts.
5092    unsafe {
5093        use core::arch::x86_64::*;
5094        let lomask = _mm_set1_epi8(0x0F);
5095        let eight = _mm256_set1_epi8(8);
5096        let ones = _mm256_set1_epi16(1);
5097        // One f32 accumulator VECTOR per activation, reduced once at the
5098        // end. Folding each group's i32 lanes to a scalar inside the loop
5099        // costs an extracti128 + three shift/add + a movd — a cross-lane
5100        // dependency chain per (group, activation), 288 of them per row at
5101        // cols=2304. The per-group scale is what forces a float
5102        // accumulator; it does not force a horizontal sum.
5103        //
5104        // The four accumulators are NAMED, not an array: as `[__m256; 4]`
5105        // indexed by a loop variable LLVM keeps them in memory and every
5106        // group pays four 32-byte loads and stores. That alone made this
5107        // kernel 2x SLOWER than the per-group reduction it replaces
5108        // (measured on the EPYC box: 150 s vs 71 s for two 256² steps).
5109        let mut f0 = _mm256_setzero_ps();
5110        let mut f1 = _mm256_setzero_ps();
5111        let mut f2 = _mm256_setzero_ps();
5112        let mut f3 = _mm256_setzero_ps();
5113        for gi in 0..gpr {
5114            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
5115            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5116            let sv = _mm256_set1_ps(s);
5117            let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
5118            let lo = _mm_and_si128(bb, lomask);
5119            let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
5120            let w = _mm256_sub_epi8(
5121                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5122                eight,
5123            );
5124            let aw = _mm256_abs_epi8(w);
5125            let off = gi * GROUP_SIZE;
5126            let dot = |xq: &[i8]| {
5127                let x = _mm256_loadu_si256(xq.as_ptr().add(off) as *const __m256i);
5128                let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
5129                _mm256_cvtepi32_ps(_mm256_madd_epi16(p16, ones))
5130            };
5131            f0 = _mm256_fmadd_ps(dot(xs[0]), sv, f0);
5132            f1 = _mm256_fmadd_ps(dot(xs[1]), sv, f1);
5133            f2 = _mm256_fmadd_ps(dot(xs[2]), sv, f2);
5134            f3 = _mm256_fmadd_ps(dot(xs[3]), sv, f3);
5135        }
5136        [
5137            hsum256_ps(f0),
5138            hsum256_ps(f1),
5139            hsum256_ps(f2),
5140            hsum256_ps(f3),
5141        ]
5142    }
5143}
5144
5145/// Horizontal sum of eight f32 lanes — the one cross-lane reduction the
5146/// blocked kernels pay, once per row instead of once per group.
5147#[cfg(target_arch = "x86_64")]
5148#[target_feature(enable = "avx2")]
5149#[inline]
5150unsafe fn hsum256_ps(v: core::arch::x86_64::__m256) -> f32 {
5151    // SAFETY: pure register arithmetic on the caller's vector.
5152    unsafe {
5153        use core::arch::x86_64::*;
5154        let hi = _mm256_extractf128_ps::<1>(v);
5155        let s = _mm_add_ps(_mm256_castps256_ps128(v), hi);
5156        let s = _mm_add_ps(s, _mm_movehl_ps(s, s));
5157        let s = _mm_add_ss(s, _mm_shuffle_ps::<0x55>(s, s));
5158        _mm_cvtss_f32(s)
5159    }
5160}
5161
5162/// VNNI twin of `dot_q4t_row_1x4_avx2` (see `dpbusd_hsum`).
5163#[cfg(target_arch = "x86_64")]
5164#[target_feature(enable = "avx2,fma,avx512f,avx512bw,avx512vl,avx512vnni")]
5165unsafe fn dot_q4t_row_1x4_vnni(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
5166    // SAFETY: callers uphold the 18B-tile and xq-length contracts.
5167    unsafe {
5168        use core::arch::x86_64::*;
5169        let lomask = _mm_set1_epi8(0x0F);
5170        let eight = _mm256_set1_epi8(8);
5171        // Same shape as the AVX2 twin: accumulate in f32 vectors and pay
5172        // one cross-lane reduction per row, not per (group, activation).
5173        let mut f0 = _mm256_setzero_ps();
5174        let mut f1 = _mm256_setzero_ps();
5175        let mut f2 = _mm256_setzero_ps();
5176        let mut f3 = _mm256_setzero_ps();
5177        for gi in 0..gpr {
5178            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
5179            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5180            let sv = _mm256_set1_ps(s);
5181            let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
5182            let lo = _mm_and_si128(bb, lomask);
5183            let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
5184            let w = _mm256_sub_epi8(
5185                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5186                eight,
5187            );
5188            let aw = _mm256_abs_epi8(w);
5189            let off = gi * GROUP_SIZE;
5190            let dot = |xq: &[i8]| {
5191                let x = _mm256_loadu_si256(xq.as_ptr().add(off) as *const __m256i);
5192                _mm256_cvtepi32_ps(_mm256_dpbusd_epi32(
5193                    _mm256_setzero_si256(),
5194                    aw,
5195                    _mm256_sign_epi8(x, w),
5196                ))
5197            };
5198            f0 = _mm256_fmadd_ps(dot(xs[0]), sv, f0);
5199            f1 = _mm256_fmadd_ps(dot(xs[1]), sv, f1);
5200            f2 = _mm256_fmadd_ps(dot(xs[2]), sv, f2);
5201            f3 = _mm256_fmadd_ps(dot(xs[3]), sv, f3);
5202        }
5203        let acc = [
5204            hsum256_ps(f0),
5205            hsum256_ps(f1),
5206            hsum256_ps(f2),
5207            hsum256_ps(f3),
5208        ];
5209        acc
5210    }
5211}
5212
5213/// ARM twin of `dot_q4t_row_1x4_avx2`: one nibble unpack per group
5214/// serves FOUR activation streams. Per stream the group order and f32
5215/// accumulation match `dot_q4t_row_sdot` exactly — batch == matvec
5216/// bit-for-bit.
5217#[cfg(target_arch = "aarch64")]
5218#[target_feature(enable = "neon,dotprod")]
5219unsafe fn dot_q4t_row_1x4_sdot(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
5220    // SAFETY: callers uphold the 18B-tile and xq-length contracts.
5221    unsafe {
5222        use core::arch::aarch64::*;
5223        use core::arch::asm;
5224        let lomask = vdupq_n_u8(0x0F);
5225        let eight = vdupq_n_s8(8);
5226        let mut acc = [0f32; 4];
5227        for gi in 0..gpr {
5228            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
5229            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5230            let b = vld1q_u8(t.add(2));
5231            let lo = vandq_u8(b, lomask);
5232            let hi = vshrq_n_u8::<4>(b);
5233            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
5234            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
5235            for (k, xq) in xs.iter().enumerate() {
5236                let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
5237                let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
5238                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5239                asm!(
5240                    "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
5241                    "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
5242                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5243                    e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
5244                    options(pure, nomem, nostack),
5245                );
5246                acc[k] += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
5247            }
5248        }
5249        acc
5250    }
5251}
5252
5253/// Exact-term correction for A8W8 outliers on a tiled row.
5254#[inline]
5255fn q4t_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
5256    let gi = j / GROUP_SIZE;
5257    let k = j % GROUP_SIZE;
5258    let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
5259    let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
5260    let byte = tile[2 + k / 2];
5261    let nib = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
5262    ((nib as i32 - 8) as f32, s)
5263}
5264
5265/// Exact scalar q4_tiled row (CMF_SDOT=0 contract) — same pairwise
5266/// accumulation shape as `q4_range_f32`.
5267#[inline]
5268fn q4t_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
5269    let mut acc = 0f32;
5270    for gi in 0..gpr {
5271        let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
5272        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
5273        let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
5274        let mut ga = 0f32;
5275        for (k, &b) in tile[2..].iter().enumerate() {
5276            ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
5277                + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
5278        }
5279        acc += ga * s;
5280    }
5281    acc
5282}
5283
5284/// Split view of a `q4tp` payload. The three planes are resolved once per
5285/// matvec instead of per row — `q4tp_sections` is cheap, but doing it inside
5286/// the row loop would put a division on the hot path for nothing.
5287struct Q4tpView<'a> {
5288    nib: &'a [u8],
5289    params: &'a [u8],
5290    codes: &'a [u8],
5291    stride: usize,
5292    /// q2tp reads the ladder with rung 0 = exact zero.
5293    zero_rung: bool,
5294}
5295
5296impl<'a> Q4tpView<'a> {
5297    fn new(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
5298        let (params_off, codes_off, stride) = q4tp_sections(rows, cols);
5299        Self {
5300            nib: &bytes[..params_off],
5301            params: &bytes[params_off..codes_off],
5302            codes: &bytes[codes_off..],
5303            stride,
5304            zero_rung: false,
5305        }
5306    }
5307
5308    /// The q2tp view: identical params/codes planes, 8 B weight chunks.
5309    fn new_q2(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
5310        let (params_off, codes_off, stride) = q2tp_sections(rows, cols);
5311        Self {
5312            nib: &bytes[..params_off],
5313            params: &bytes[params_off..codes_off],
5314            codes: &bytes[codes_off..],
5315            stride,
5316            zero_rung: true,
5317        }
5318    }
5319
5320    /// Expand row `r`'s per-tile scales into `out` (length `gpr`).
5321    ///
5322    /// Doing this once per row — rather than decoding a 5-bit code inside the
5323    /// tile loop — is what makes the format free at runtime. Random access to
5324    /// a packed 5-bit field costs a division, two bounds checks and a branch;
5325    /// the tile's actual work is two `sdot`s, so per-tile decoding dominated
5326    /// the kernel and cost 5x (measured: 1.4 vs 6.9 tok/s on Nanbeige-3B).
5327    /// Walking the plane sequentially with a bit accumulator is ~3 ops.
5328    /// Eight 5-bit codes are exactly five bytes, so a whole group of
5329    /// eight decodes from one little-endian word at fixed shifts. The
5330    /// bit-accumulator this replaces carried a data-dependent `while
5331    /// have < 5` refill whose branch sat in the innermost loop of every
5332    /// q4tp row; a decode profile put this function above the dot
5333    /// products it feeds. Same bitstream, same codes — just no branch
5334    /// and eight independent extractions.
5335    #[inline]
5336    fn scales_into(&self, r: usize, gpr: usize, out: &mut [f32]) {
5337        let tab = if self.zero_rung {
5338            q2tp_ladder(self.params, r)
5339        } else {
5340            q4tp_ladder(self.params, r)
5341        };
5342        let codes = &self.codes[r * self.stride..(r + 1) * self.stride];
5343        let out = &mut out[..gpr];
5344        let mut chunks = out.chunks_exact_mut(8);
5345        let mut ci = 0usize;
5346        for c in &mut chunks {
5347            let w = u64::from(codes[ci])
5348                | u64::from(codes[ci + 1]) << 8
5349                | u64::from(codes[ci + 2]) << 16
5350                | u64::from(codes[ci + 3]) << 24
5351                | u64::from(codes[ci + 4]) << 32;
5352            for (k, o) in c.iter_mut().enumerate() {
5353                *o = tab[((w >> (5 * k)) & 31) as usize];
5354            }
5355            ci += 5;
5356        }
5357        // Fewer than eight codes left: the shared total accessor, which
5358        // tolerates a 5-bit field whose spill byte is past the stride.
5359        let tail = &codes[ci..];
5360        for (k, o) in chunks.into_remainder().iter_mut().enumerate() {
5361            *o = tab[q4tp_code(tail, k)];
5362        }
5363    }
5364}
5365
5366#[inline]
5367fn dot_q4tp_row_i8(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
5368    #[cfg(target_arch = "aarch64")]
5369    unsafe {
5370        return dot_q4tp_row_sdot(nib, r, gpr, xq, scales);
5371    }
5372    #[cfg(target_arch = "x86_64")]
5373    unsafe {
5374        if vnni_tiles_enabled() {
5375            return dot_q4tp_row_vnni(nib, r, gpr, xq, scales);
5376        }
5377        return dot_q4tp_row_avx2(nib, r, gpr, xq, scales);
5378    }
5379    #[allow(unreachable_code)]
5380    {
5381        let mut acc = 0f32;
5382        for gi in 0..gpr {
5383            let tile = &nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
5384            let s = scales[gi];
5385            let mut d = 0i32;
5386            for (k, &b) in tile.iter().enumerate() {
5387                d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
5388                    + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
5389            }
5390            acc += d as f32 * s;
5391        }
5392        acc
5393    }
5394}
5395
5396/// q4tp twin of `dot_q4t_row_sdot`: identical nibble math, but the tile
5397/// stride is 16 B (no inline scale) and the scale is a ladder lookup.
5398#[cfg(target_arch = "aarch64")]
5399#[target_feature(enable = "neon,dotprod")]
5400unsafe fn dot_q4tp_row_sdot(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
5401    // SAFETY: callers uphold slice-length contracts (16B tile per group,
5402    // xq.len() == gpr·GROUP_SIZE, codes covering gpr 5-bit fields).
5403    unsafe {
5404        use core::arch::aarch64::*;
5405        use core::arch::asm;
5406        let lomask = vdupq_n_u8(0x0F);
5407        let eight = vdupq_n_s8(8);
5408        let mut acc = 0f32;
5409        for gi in 0..gpr {
5410            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5411            let s = *scales.get_unchecked(gi);
5412            let b = vld1q_u8(t);
5413            let lo = vandq_u8(b, lomask);
5414            let hi = vshrq_n_u8::<4>(b);
5415            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
5416            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
5417            let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
5418            let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
5419            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5420            asm!(
5421                "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
5422                "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
5423                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5424                e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
5425                options(pure, nomem, nostack),
5426            );
5427            acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
5428        }
5429        acc
5430    }
5431}
5432
5433#[cfg(target_arch = "x86_64")]
5434#[target_feature(enable = "avx2")]
5435unsafe fn dot_q4tp_row_avx2(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
5436    // SAFETY: see dot_q4tp_row_sdot.
5437    unsafe {
5438        use core::arch::x86_64::*;
5439        let lomask = _mm_set1_epi8(0x0F);
5440        let eight = _mm256_set1_epi8(8);
5441        let ones = _mm256_set1_epi16(1);
5442        let mut acc = 0f32;
5443        for gi in 0..gpr {
5444            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5445            let s = *scales.get_unchecked(gi);
5446            let b = _mm_loadu_si128(t as *const __m128i);
5447            let lo = _mm_and_si128(b, lomask);
5448            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
5449            let w = _mm256_sub_epi8(
5450                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5451                eight,
5452            );
5453            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5454            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
5455            let d = _mm256_madd_epi16(p16, ones);
5456            let hi128 = _mm256_extracti128_si256::<1>(d);
5457            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
5458            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
5459            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
5460            acc += _mm_cvtsi128_si32(s32) as f32 * s;
5461        }
5462        acc
5463    }
5464}
5465
5466/// VNNI twin of `dot_q4tp_row_avx2` (see `dot_q4t_row_vnni` for why the
5467/// 256-bit VL encoding is the one to use here).
5468#[cfg(target_arch = "x86_64")]
5469#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
5470unsafe fn dot_q4tp_row_vnni(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
5471    // SAFETY: see dot_q4tp_row_sdot.
5472    unsafe {
5473        use core::arch::x86_64::*;
5474        let lomask = _mm_set1_epi8(0x0F);
5475        let eight = _mm256_set1_epi8(8);
5476        let mut acc = 0f32;
5477        for gi in 0..gpr {
5478            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5479            let s = *scales.get_unchecked(gi);
5480            let b = _mm_loadu_si128(t as *const __m128i);
5481            let lo = _mm_and_si128(b, lomask);
5482            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
5483            let w = _mm256_sub_epi8(
5484                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5485                eight,
5486            );
5487            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5488            acc += dpbusd_hsum(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w)) as f32 * s;
5489        }
5490        acc
5491    }
5492}
5493
5494/// Exact scalar q4tp row — the `CMF_SDOT=0` contract, same pairwise
5495/// accumulation shape as `q4t_row_exact`.
5496#[inline]
5497fn q4tp_row_exact(nib: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
5498    #[cfg(target_arch = "x86_64")]
5499    if avx2_enabled() {
5500        // Keep the scalar pair/group reduction order, not a vector sum.
5501        return unsafe { q4tp_row_float_avx2(nib, r, gpr, x, scales) };
5502    }
5503    q4tp_row_float_scalar(nib, r, gpr, x, scales)
5504}
5505
5506#[inline]
5507fn q4tp_row_float_scalar(nib: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
5508    let mut acc = 0f32;
5509    for gi in 0..gpr {
5510        let tile = &nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
5511        let s = scales[gi];
5512        let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
5513        let mut ga = 0f32;
5514        for (k, &b) in tile.iter().enumerate() {
5515            ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
5516                + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
5517        }
5518        acc += ga * s;
5519    }
5520    acc
5521}
5522
5523/// Vectorize unpack, conversion and multiplication, but preserve every
5524/// pair addition and the scalar accumulation order. No activation rounding
5525/// or FMA: bit-identical to the float scalar row, including its group scale.
5526#[cfg(target_arch = "x86_64")]
5527#[target_feature(enable = "avx2")]
5528unsafe fn q4tp_row_float_avx2(nib: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
5529    // SAFETY: caller checks AVX2 and provides the same complete 32-element
5530    // groups as the scalar row. Loads/stores are explicitly unaligned.
5531    unsafe {
5532        use core::arch::x86_64::*;
5533        let mask = _mm_set1_epi8(15);
5534        let eight = _mm_set1_epi8(8);
5535        let order = _mm256_setr_epi32(0, 1, 4, 5, 2, 3, 6, 7);
5536        let mut acc = 0.0f32;
5537        for gi in 0..gpr {
5538            let packed = _mm_loadu_si128(nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB).cast());
5539            let lo = _mm_and_si128(packed, mask);
5540            let hi = _mm_and_si128(_mm_srli_epi16::<4>(packed), mask);
5541            let w0 = _mm_sub_epi8(_mm_unpacklo_epi8(lo, hi), eight);
5542            let w1 = _mm_sub_epi8(_mm_unpackhi_epi8(lo, hi), eight);
5543            let xp = x.as_ptr().add(gi * GROUP_SIZE);
5544            let a = _mm256_mul_ps(
5545                _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(w0)),
5546                _mm256_loadu_ps(xp),
5547            );
5548            let b = _mm256_mul_ps(
5549                _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(w0))),
5550                _mm256_loadu_ps(xp.add(8)),
5551            );
5552            let c = _mm256_mul_ps(
5553                _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(w1)),
5554                _mm256_loadu_ps(xp.add(16)),
5555            );
5556            let d = _mm256_mul_ps(
5557                _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(w1))),
5558                _mm256_loadu_ps(xp.add(24)),
5559            );
5560            let mut pairs = [0.0f32; 16];
5561            _mm256_storeu_ps(
5562                pairs.as_mut_ptr(),
5563                _mm256_permutevar8x32_ps(_mm256_hadd_ps(a, b), order),
5564            );
5565            _mm256_storeu_ps(
5566                pairs.as_mut_ptr().add(8),
5567                _mm256_permutevar8x32_ps(_mm256_hadd_ps(c, d), order),
5568            );
5569            let mut ga = 0.0f32;
5570            for v in pairs {
5571                ga += v;
5572            }
5573            acc += ga * scales[gi];
5574        }
5575        acc
5576    }
5577}
5578
5579/// Single weight of a q4tp tensor — the a8w8 outlier path, which restores
5580/// activation outliers at full precision after the int8 pass.
5581#[inline]
5582fn q4tp_outlier(nib: &[u8], r: usize, gpr: usize, j: usize, scales: &[f32]) -> (f32, f32) {
5583    let (gi, k) = (j / GROUP_SIZE, j % GROUP_SIZE);
5584    let byte = nib[(r * gpr + gi) * Q4TP_NIB + k / 2];
5585    let n = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
5586    ((n as i32 - 8) as f32, scales[gi])
5587}
5588
5589/// Fused q4tp matvec (dispatch mirrors `q4t_matvec`).
5590fn q4tp_matvec(
5591    bytes: &[u8],
5592    x: &[f32],
5593    rows: usize,
5594    cols: usize,
5595    out: &mut [f32],
5596    pool: Option<&Pool>,
5597) {
5598    debug_assert_eq!(out.len(), rows);
5599    let gpr = cols / GROUP_SIZE;
5600    let v = Q4tpView::new(bytes, rows, cols);
5601    let out_addr = SendMut(out.as_mut_ptr());
5602    if a8w8_enabled() {
5603        let act = split_act(x);
5604        let run = |start: usize, end: usize| {
5605            // One scratch row of scales per worker — borrowed, not minted.
5606            with_krow(gpr, |sc| {
5607                for r in start..end {
5608                    v.scales_into(r, gpr, sc);
5609                    let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, sc) * act.sx;
5610                    for &(j, xv) in &act.outliers {
5611                        let (w, s) = q4tp_outlier(v.nib, r, gpr, j, sc);
5612                        acc += w * s * xv;
5613                    }
5614                    // SAFETY: disjoint row ranges per worker.
5615                    unsafe { *out_addr.at(r) = acc };
5616                }
5617            })
5618        };
5619        dispatch_rows(pool, rows, &run);
5620        return;
5621    }
5622    let run = |start: usize, end: usize| {
5623        with_krow(gpr, |sc| {
5624            for r in start..end {
5625                v.scales_into(r, gpr, sc);
5626                // SAFETY: disjoint row ranges per worker.
5627                unsafe { *out_addr.at(r) = q4tp_row_exact(v.nib, r, gpr, x, sc) };
5628            }
5629        })
5630    };
5631    dispatch_rows(pool, rows, &run);
5632}
5633
5634/// Fused two-input q4tp matvec — the SwiGLU gate/up pair. Weights and the
5635/// row ladder are read once and spent on both activation streams.
5636#[allow(clippy::too_many_arguments)]
5637fn q4tp_matvec2(
5638    bytes: &[u8],
5639    x1: &[f32],
5640    x2: &[f32],
5641    rows: usize,
5642    cols: usize,
5643    o1: &mut [f32],
5644    o2: &mut [f32],
5645    pool: Option<&Pool>,
5646) {
5647    let gpr = cols / GROUP_SIZE;
5648    let v = Q4tpView::new(bytes, rows, cols);
5649    let (p1, p2) = (SendMut(o1.as_mut_ptr()), SendMut(o2.as_mut_ptr()));
5650    let run = |start: usize, end: usize| {
5651        let mut sc = vec![0f32; gpr];
5652        for r in start..end {
5653            v.scales_into(r, gpr, &mut sc);
5654            // SAFETY: disjoint row ranges per worker.
5655            unsafe {
5656                *p1.at(r) = q4tp_row_exact(v.nib, r, gpr, x1, &sc);
5657                *p2.at(r) = q4tp_row_exact(v.nib, r, gpr, x2, &sc);
5658            }
5659        }
5660    };
5661    dispatch_rows(pool, rows, &run);
5662}
5663
5664/// One q2tp outlier weight at column `j` of row `r`: the 2-bit code and
5665/// its group scale, mirrored on `q4tp_outlier`.
5666#[inline]
5667fn q2tp_outlier(chunks: &[u8], r: usize, gpr: usize, j: usize, scales: &[f32]) -> (f32, f32) {
5668    let (gi, k) = (j / GROUP_SIZE, j % GROUP_SIZE);
5669    let byte = chunks[(r * gpr + gi) * Q2TP_CHUNK + k / 4];
5670    let c = (byte >> (2 * (k % 4))) & 3;
5671    (c as f32 - 1.5, scales[gi])
5672}
5673
5674#[cfg(target_arch = "x86_64")]
5675const Q2TP_DECODE_U32: [u32; 256] = {
5676    let mut tab = [0u32; 256];
5677    let mut b = 0usize;
5678    while b < 256 {
5679        tab[b] = ((b as u32) & 3)
5680            | ((((b as u32) >> 2) & 3) << 8)
5681            | ((((b as u32) >> 4) & 3) << 16)
5682            | ((((b as u32) >> 6) & 3) << 24);
5683        b += 1;
5684    }
5685    tab
5686};
5687
5688/// Eight packed q2tp bytes against 32 signed activation bytes. `maddubs`
5689/// exactly computes unsigned 2-bit code × signed i8; its pair sums cannot
5690/// saturate (2 × 3 × 127 < i16::MAX), and the second madd widens to i32.
5691#[cfg(target_arch = "x86_64")]
5692#[target_feature(enable = "avx2")]
5693unsafe fn q2tp_code_dot_avx2(ch: &[u8], x: &[i8]) -> i32 {
5694    use core::arch::x86_64::*;
5695    debug_assert!(ch.len() >= Q2TP_CHUNK && x.len() >= GROUP_SIZE);
5696    let codes = _mm256_setr_epi32(
5697        Q2TP_DECODE_U32[ch[0] as usize] as i32,
5698        Q2TP_DECODE_U32[ch[1] as usize] as i32,
5699        Q2TP_DECODE_U32[ch[2] as usize] as i32,
5700        Q2TP_DECODE_U32[ch[3] as usize] as i32,
5701        Q2TP_DECODE_U32[ch[4] as usize] as i32,
5702        Q2TP_DECODE_U32[ch[5] as usize] as i32,
5703        Q2TP_DECODE_U32[ch[6] as usize] as i32,
5704        Q2TP_DECODE_U32[ch[7] as usize] as i32,
5705    );
5706    let xv = unsafe { _mm256_loadu_si256(x.as_ptr().cast()) };
5707    let pair = _mm256_maddubs_epi16(codes, xv);
5708    let quad = _mm256_madd_epi16(pair, _mm256_set1_epi16(1));
5709    let sum128 = _mm_add_epi32(
5710        _mm256_castsi256_si128(quad),
5711        _mm256_extracti128_si256(quad, 1),
5712    );
5713    let sum64 = _mm_hadd_epi32(sum128, sum128);
5714    _mm_cvtsi128_si32(_mm_hadd_epi32(sum64, sum64))
5715}
5716
5717#[cfg(target_arch = "x86_64")]
5718#[inline]
5719fn q2tp_avx2_enabled() -> bool {
5720    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
5721    *ON.get_or_init(|| std::arch::is_x86_feature_detected!("avx2"))
5722}
5723
5724/// Integer dot of one q2tp row against pre-quantized activations:
5725/// Σ_g s_g · (Σ c·xq − 1.5·Σ xq). The half-integer grid (c − 1.5)
5726/// becomes exact integer math through the group sums — the same trick
5727/// every a8w8 kernel in this file rides. The codes decode into a
5728/// 32-byte scratch in natural order and the dot itself is the shared
5729/// SDOT primitive; elsewhere a scalar integer loop.
5730#[inline]
5731fn dot_q2tp_row_i8(
5732    chunks: &[u8],
5733    r: usize,
5734    gpr: usize,
5735    xq: &[i8],
5736    gsum: &[i32],
5737    scales: &[f32],
5738) -> f32 {
5739    let mut acc = 0f32;
5740    let base = r * gpr * Q2TP_CHUNK;
5741    #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
5742    let mut codes = [0i8; GROUP_SIZE];
5743    #[cfg(target_arch = "x86_64")]
5744    // Cache the process-wide AVX2 decision; probing it for every 32-weight
5745    // group adds a branch to the hottest Q2TP row loop while retaining the
5746    // established table-load AVX2 arithmetic (and without coupling this path
5747    // to the separate FMA-gated A8W8 switch).
5748    let avx2 = q2tp_avx2_enabled();
5749    for gi in 0..gpr {
5750        let ch = &chunks[base + gi * Q2TP_CHUNK..base + (gi + 1) * Q2TP_CHUNK];
5751        let xg = &xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
5752        #[cfg(target_arch = "aarch64")]
5753        // NEON: the byte's four 2-bit fields land in four lane vectors
5754        // (shift+mask), vld4 de-interleaves xq to match (xj[k] =
5755        // xq[4k+j]), widening MACs accumulate exactly in i32. A scalar
5756        // decode here cost as much as the dot it fed — the profile put
5757        // it at the top of the whole W2 decode.
5758        let dot = unsafe {
5759            use core::arch::aarch64::*;
5760            let b = vld1_u8(ch.as_ptr());
5761            let three = vdup_n_u8(3);
5762            let c0 = vreinterpret_s8_u8(vand_u8(b, three));
5763            let c1 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 2), three));
5764            let c2 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 4), three));
5765            let c3 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 6), three));
5766            let x4 = vld4_s8(xg.as_ptr());
5767            let mut acc4 = vdupq_n_s32(0);
5768            acc4 = vpadalq_s16(acc4, vmull_s8(c0, x4.0));
5769            acc4 = vpadalq_s16(acc4, vmull_s8(c1, x4.1));
5770            acc4 = vpadalq_s16(acc4, vmull_s8(c2, x4.2));
5771            acc4 = vpadalq_s16(acc4, vmull_s8(c3, x4.3));
5772            vaddvq_s32(acc4)
5773        };
5774        #[cfg(target_arch = "x86_64")]
5775        let dot: i32 = if avx2 {
5776            // SAFETY: the runtime feature check gates the target-feature body;
5777            // the group slices above are exactly 8 and 32 bytes long.
5778            unsafe { q2tp_code_dot_avx2(ch, xg) }
5779        } else {
5780            ch.iter()
5781                .enumerate()
5782                .map(|(k, &b)| {
5783                    ((b & 3) as i32) * xg[k * 4] as i32
5784                        + (((b >> 2) & 3) as i32) * xg[k * 4 + 1] as i32
5785                        + (((b >> 4) & 3) as i32) * xg[k * 4 + 2] as i32
5786                        + (((b >> 6) & 3) as i32) * xg[k * 4 + 3] as i32
5787                })
5788                .sum()
5789        };
5790        #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
5791        let dot: i32 = {
5792            for (k, &b) in ch.iter().enumerate() {
5793                codes[k * 4] = (b & 3) as i8;
5794                codes[k * 4 + 1] = ((b >> 2) & 3) as i8;
5795                codes[k * 4 + 2] = ((b >> 4) & 3) as i8;
5796                codes[k * 4 + 3] = ((b >> 6) & 3) as i8;
5797            }
5798            codes
5799                .iter()
5800                .zip(xg)
5801                .map(|(&c, &x)| c as i32 * x as i32)
5802                .sum()
5803        };
5804        acc += scales[gi] * (dot as f32 - 1.5 * gsum[gi] as f32);
5805    }
5806    acc
5807}
5808
5809/// Exact f32 dot of one q2tp row: 2-bit fields LSB-first, (c − 1.5)·s.
5810/// Scalar on purpose — the 2-bit class targets the GPU graph; the CPU
5811/// path exists for parity gates and small-machine fallback.
5812fn q2tp_row_exact(chunks: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
5813    q2tp_row_exact_center(chunks, r, gpr, x, scales, 1.5)
5814}
5815
5816/// Fused Prism affine row: the derived correction is applied inside the
5817/// decoded code, avoiding a second accumulated dot and avoiding cancellation
5818/// between `B=(c-1.5)s` and `+.5s` for long 5120/17408 rows.
5819#[inline]
5820fn q2tp_affine_row_exact(chunks: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
5821    q2tp_row_exact_center(chunks, r, gpr, x, scales, 1.0)
5822}
5823
5824#[inline]
5825fn q2tp_row_exact_center(
5826    chunks: &[u8],
5827    r: usize,
5828    gpr: usize,
5829    x: &[f32],
5830    scales: &[f32],
5831    center: f32,
5832) -> f32 {
5833    let mut acc = 0f32;
5834    for gi in 0..gpr {
5835        let ch = &chunks[(r * gpr + gi) * Q2TP_CHUNK..(r * gpr + gi + 1) * Q2TP_CHUNK];
5836        let s = scales[gi];
5837        let xb = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
5838        let mut g = 0f32;
5839        for (k, &b) in ch.iter().enumerate() {
5840            g += ((b & 3) as f32 - center) * xb[k * 4]
5841                + (((b >> 2) & 3) as f32 - center) * xb[k * 4 + 1]
5842                + (((b >> 4) & 3) as f32 - center) * xb[k * 4 + 2]
5843                + (((b >> 6) & 3) as f32 - center) * xb[k * 4 + 3];
5844        }
5845        acc += s * g;
5846    }
5847    acc
5848}
5849
5850fn q2tp_matvec(
5851    bytes: &[u8],
5852    x: &[f32],
5853    rows: usize,
5854    cols: usize,
5855    out: &mut [f32],
5856    pool: Option<&Pool>,
5857) {
5858    q2tp_matvec_mode(bytes, x, rows, cols, out, pool, false);
5859}
5860
5861fn q2tp_affine_matvec(
5862    bytes: &[u8],
5863    x: &[f32],
5864    rows: usize,
5865    cols: usize,
5866    out: &mut [f32],
5867    pool: Option<&Pool>,
5868) {
5869    q2tp_matvec_mode(bytes, x, rows, cols, out, pool, true);
5870}
5871
5872fn q2tp_matvec_mode(
5873    bytes: &[u8],
5874    x: &[f32],
5875    rows: usize,
5876    cols: usize,
5877    out: &mut [f32],
5878    pool: Option<&Pool>,
5879    affine: bool,
5880) {
5881    debug_assert_eq!(out.len(), rows);
5882    let gpr = cols / GROUP_SIZE;
5883    let v = Q4tpView::new_q2(bytes, rows, cols);
5884    let out_addr = SendMut(out.as_mut_ptr());
5885    // a8w8 fast path (CMF_SDOT=0 keeps the exact scalar walk): integer
5886    // code dots + group sums, exact outlier correction — the same
5887    // contract as every sibling kernel; measured 2-bit rows were the
5888    // only scalar holdout in the family.
5889    if !affine && a8w8_enabled() {
5890        let act = split_act(x);
5891        let gsum = q1_group_sums(&act.xq, gpr);
5892        let (act, gsum) = (&act, &gsum);
5893        let run = move |start: usize, end: usize| {
5894            with_krow(gpr, |sc| {
5895                for r in start..end {
5896                    v.scales_into(r, gpr, sc);
5897                    let mut acc = dot_q2tp_row_i8(v.nib, r, gpr, &act.xq, gsum, sc) * act.sx;
5898                    for &(j, xv) in &act.outliers {
5899                        let (w, s) = q2tp_outlier(v.nib, r, gpr, j, sc);
5900                        acc += w * s * xv;
5901                    }
5902                    // SAFETY: disjoint row ranges per worker.
5903                    unsafe { *out_addr.at(r) = acc };
5904                }
5905            })
5906        };
5907        dispatch_rows(pool, rows, &run);
5908        return;
5909    }
5910    let run = |start: usize, end: usize| {
5911        with_krow(gpr, |sc| {
5912            for r in start..end {
5913                v.scales_into(r, gpr, sc);
5914                // SAFETY: disjoint row ranges per worker.
5915                unsafe {
5916                    *out_addr.at(r) = if affine {
5917                        q2tp_affine_row_exact(v.nib, r, gpr, x, sc)
5918                    } else {
5919                        q2tp_row_exact(v.nib, r, gpr, x, sc)
5920                    }
5921                };
5922            }
5923        })
5924    };
5925    dispatch_rows(pool, rows, &run);
5926}
5927
5928/// Fused two-input q2tp matvec — the SwiGLU gate/up pair.
5929#[allow(clippy::too_many_arguments)]
5930fn q2tp_matvec2(
5931    bytes: &[u8],
5932    x1: &[f32],
5933    x2: &[f32],
5934    rows: usize,
5935    cols: usize,
5936    o1: &mut [f32],
5937    o2: &mut [f32],
5938    pool: Option<&Pool>,
5939) {
5940    q2tp_matvec2_mode(bytes, x1, x2, rows, cols, o1, o2, pool, false);
5941}
5942
5943#[allow(clippy::too_many_arguments)]
5944fn q2tp_affine_matvec2(
5945    bytes: &[u8],
5946    x1: &[f32],
5947    x2: &[f32],
5948    rows: usize,
5949    cols: usize,
5950    o1: &mut [f32],
5951    o2: &mut [f32],
5952    pool: Option<&Pool>,
5953) {
5954    q2tp_matvec2_mode(bytes, x1, x2, rows, cols, o1, o2, pool, true);
5955}
5956
5957#[allow(clippy::too_many_arguments)]
5958fn q2tp_matvec2_mode(
5959    bytes: &[u8],
5960    x1: &[f32],
5961    x2: &[f32],
5962    rows: usize,
5963    cols: usize,
5964    o1: &mut [f32],
5965    o2: &mut [f32],
5966    pool: Option<&Pool>,
5967    affine: bool,
5968) {
5969    let gpr = cols / GROUP_SIZE;
5970    let v = Q4tpView::new_q2(bytes, rows, cols);
5971    let (p1, p2) = (SendMut(o1.as_mut_ptr()), SendMut(o2.as_mut_ptr()));
5972    let run = |start: usize, end: usize| {
5973        let mut sc = vec![0f32; gpr];
5974        for r in start..end {
5975            v.scales_into(r, gpr, &mut sc);
5976            // SAFETY: disjoint row ranges per worker.
5977            unsafe {
5978                *p1.at(r) = if affine {
5979                    q2tp_affine_row_exact(v.nib, r, gpr, x1, &sc)
5980                } else {
5981                    q2tp_row_exact(v.nib, r, gpr, x1, &sc)
5982                };
5983                *p2.at(r) = if affine {
5984                    q2tp_affine_row_exact(v.nib, r, gpr, x2, &sc)
5985                } else {
5986                    q2tp_row_exact(v.nib, r, gpr, x2, &sc)
5987                };
5988            }
5989        }
5990    };
5991    dispatch_rows(pool, rows, &run);
5992}
5993
5994/// Batched q2tp matmat: scalar row kernel over every batch column. CPU
5995/// prefill only — decode rides the graph, so plain and correct beats
5996/// clever here.
5997/// Test doors into the host 2-bit kernels: the stand's heap corruption
5998/// pointed at down-shaped tensors, and the private fns need a way to be
5999/// held to a reference without a model file around them.
6000pub fn q2tp_matvec_for_test(bytes: &[u8], x: &[f32], rows: usize, cols: usize, out: &mut [f32]) {
6001    // The facade IS the reference: encoder oracles hold requant output
6002    // to the exact scalar walk. The production dispatch may take the i8
6003    // fast path, whose error scale is the ACTIVATIONS' — a different
6004    // claim than the encoder correctness these tests pin.
6005    let gpr = cols / GROUP_SIZE;
6006    let v = Q4tpView::new_q2(bytes, rows, cols);
6007    with_krow(gpr, |sc| {
6008        for r in 0..rows {
6009            v.scales_into(r, gpr, sc);
6010            out[r] = q2tp_row_exact(v.nib, r, gpr, x, sc);
6011        }
6012    });
6013}
6014
6015/// Test door for the descriptor-specific fused affine decode.  Production
6016/// callers select this through a validated Prism header, never by dtype alone.
6017pub fn q2tp_affine_matvec_for_test(
6018    bytes: &[u8],
6019    x: &[f32],
6020    rows: usize,
6021    cols: usize,
6022    out: &mut [f32],
6023) {
6024    q2tp_affine_matvec(bytes, x, rows, cols, out, None);
6025}
6026
6027pub fn q2tp_matmat_for_test(
6028    bytes: &[u8],
6029    xs_all: &[f32],
6030    b: usize,
6031    rows: usize,
6032    cols: usize,
6033    out: &mut [f32],
6034) {
6035    q2tp_matmat(bytes, xs_all, b, rows, cols, out, None);
6036}
6037
6038fn q2tp_matmat(
6039    bytes: &[u8],
6040    xs_all: &[f32],
6041    b: usize,
6042    rows: usize,
6043    cols: usize,
6044    out: &mut [f32],
6045    pool: Option<&Pool>,
6046) {
6047    q2tp_matmat_mode(bytes, xs_all, b, rows, cols, out, pool, false);
6048}
6049
6050fn q2tp_affine_matmat(
6051    bytes: &[u8],
6052    xs_all: &[f32],
6053    b: usize,
6054    rows: usize,
6055    cols: usize,
6056    out: &mut [f32],
6057    pool: Option<&Pool>,
6058) {
6059    q2tp_matmat_mode(bytes, xs_all, b, rows, cols, out, pool, true);
6060}
6061
6062fn q2tp_matmat_mode(
6063    bytes: &[u8],
6064    xs_all: &[f32],
6065    b: usize,
6066    rows: usize,
6067    cols: usize,
6068    out: &mut [f32],
6069    pool: Option<&Pool>,
6070    affine: bool,
6071) {
6072    debug_assert_eq!(out.len(), b * rows);
6073    let gpr = cols / GROUP_SIZE;
6074    let v = Q4tpView::new_q2(bytes, rows, cols);
6075    let out_addr = SendMut(out.as_mut_ptr());
6076    let run = |start: usize, end: usize| {
6077        let mut sc = vec![0f32; gpr];
6078        for r in start..end {
6079            v.scales_into(r, gpr, &mut sc);
6080            for bi in 0..b {
6081                let x = &xs_all[bi * cols..(bi + 1) * cols];
6082                // SAFETY: disjoint row ranges per worker.
6083                unsafe {
6084                    *out_addr.at(bi * rows + r) = if affine {
6085                        q2tp_affine_row_exact(v.nib, r, gpr, x, &sc)
6086                    } else {
6087                        q2tp_row_exact(v.nib, r, gpr, x, &sc)
6088                    }
6089                };
6090            }
6091        }
6092    };
6093    dispatch_rows(pool, rows, &run);
6094}
6095
6096/// The pre-vectorised shape, kept for A/B (`CMF_Q4TP_V1=1`): the
6097/// horizontal add lands once per group per column instead of once per
6098/// row. Same weights, same activations — only the reduction differs.
6099/// It is also the row-exact batch kernel: per group and column it forms
6100/// `int dot as f32 * scale` and adds it to a scalar running sum, which is
6101/// `dot_q4tp_row_sdot` step for step (Rust never contracts to an fma),
6102/// so each column is bit-identical to that column's matvec.
6103#[cfg(target_arch = "aarch64")]
6104#[target_feature(enable = "neon,dotprod")]
6105unsafe fn dot_q4tp_row_1x4_sdot_v1(
6106    nib: &[u8],
6107    r: usize,
6108    gpr: usize,
6109    xs: [&[i8]; 4],
6110    scales: &[f32],
6111) -> [f32; 4] {
6112    unsafe {
6113        use core::arch::aarch64::*;
6114        use core::arch::asm;
6115        let lomask = vdupq_n_u8(0x0F);
6116        let eight = vdupq_n_s8(8);
6117        let (mut f0, mut f1, mut f2, mut f3) = (0f32, 0f32, 0f32, 0f32);
6118        for gi in 0..gpr {
6119            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
6120            let s = *scales.get_unchecked(gi);
6121            let bb = vld1q_u8(t);
6122            let lo = vandq_u8(bb, lomask);
6123            let hi = vshrq_n_u8::<4>(bb);
6124            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
6125            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
6126            let mut d = [0f32; 4];
6127            for (k, dk) in d.iter_mut().enumerate() {
6128                let x0 = vld1q_s8(xs[k].as_ptr().add(gi * GROUP_SIZE));
6129                let x1 = vld1q_s8(xs[k].as_ptr().add(gi * GROUP_SIZE + 16));
6130                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6131                asm!(
6132                    "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
6133                    "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
6134                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6135                    e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
6136                    options(pure, nomem, nostack),
6137                );
6138                *dk = vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
6139            }
6140            f0 += d[0];
6141            f1 += d[1];
6142            f2 += d[2];
6143            f3 += d[3];
6144        }
6145        [f0, f1, f2, f3]
6146    }
6147}
6148
6149/// Which q4tp batch kernel to run: 1 = the previous one, 2 = the tuned
6150/// one, 0 = decide from the CPU. An atomic rather than a `OnceLock` so a
6151/// benchmark can alternate the two inside one process, where the machine's
6152/// mood — a shared box drifts ±25% between runs — is the same for both.
6153/// What the two mean is per-architecture: on x86 the blocked AVX-512 path
6154/// against the per-column one, on ARM the two reduction shapes.
6155#[allow(dead_code)]
6156static Q4TP_ALT: std::sync::atomic::AtomicU8 = std::sync::atomic::AtomicU8::new(0);
6157
6158/// Tests that store `Q4TP_ALT` hold this, so one test's kernel pick does
6159/// not leak into another's bit-exact comparison running in parallel.
6160#[cfg(test)]
6161static Q4TP_ALT_TEST_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
6162
6163/// Blocking pays on x86 only with 512-bit VNNI. With AVX2 alone, four
6164/// columns sharing an unpack still measured slower than the per-column
6165/// path (23.2 ms against 19.4 on a 48-thread EPYC), because that path
6166/// already dequantizes the row once — so the blocked kernel bought a
6167/// second unpack-free pass at the price of half the vector width.
6168#[cfg(target_arch = "x86_64")]
6169fn q4tp_blocked_x86() -> bool {
6170    match Q4TP_ALT.load(std::sync::atomic::Ordering::Relaxed) {
6171        1 => false,
6172        // A forced ON still asks the CPU. The switch exists so a bench can
6173        // pick a kernel, not so it can promise instructions the machine
6174        // does not have — CI caught that as a SIGILL on a runner without
6175        // AVX-512, where the parity test had turned the path on by hand.
6176        2 => avx512vnni_enabled(),
6177        // Deliberately not cached back into the switch: both gates below
6178        // hold their own `OnceLock`, and latching their answer here would
6179        // make a test's override outlive the test that set it.
6180        _ => blocked_enabled() && avx512vnni_enabled(),
6181    }
6182}
6183
6184/// `CMF_Q4TP_V1=1` picks the old reduction shape (A/B only).
6185#[cfg(target_arch = "aarch64")]
6186#[allow(dead_code)]
6187fn q4tp_v1() -> bool {
6188    match Q4TP_ALT.load(std::sync::atomic::Ordering::Relaxed) {
6189        1 => true,
6190        2 => false,
6191        _ => {
6192            static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6193            *ON.get_or_init(|| std::env::var("CMF_Q4TP_V1").is_ok_and(|v| v != "0"))
6194        }
6195    }
6196}
6197
6198/// Two weight rows against eight columns. The activation load is the
6199/// same for both rows, so it is paid once for twice the arithmetic, and
6200/// sixteen accumulator chains run where eight did — which is what a kernel
6201/// retiring 0.29 instructions a cycle is short of. Register pressure is
6202/// the limit: sixteen `zmm` accumulators, two weight tiles, one
6203/// activation, of thirty-two.
6204///
6205/// Four rows by four columns spends the same sixteen accumulators the
6206/// other way and measured worse — 1488 GFLOP/s against 1644 — so the
6207/// unpack, which four rows pay twice as often, costs more than the extra
6208/// sharing of one activation load buys.
6209#[cfg(target_arch = "x86_64")]
6210#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
6211unsafe fn dot_q4tp_2x8_avx512(
6212    nib: &[u8],
6213    r0: usize,
6214    gpr: usize,
6215    xs: [&[i8]; 8],
6216    sc0: &[f32],
6217    sc1: &[f32],
6218) -> [[f32; 8]; 2] {
6219    // SAFETY: as dot_q4tp_row_1x8_avx512, two adjacent rows at once; the
6220    // caller guarantees r0 + 1 < rows and the ISA.
6221    unsafe {
6222        use core::arch::x86_64::*;
6223        let lomask = _mm256_set1_epi8(0x0F);
6224        let eight = _mm256_set1_epi8(8);
6225        let zero = _mm512_setzero_si512();
6226        let mut v0 = [_mm512_setzero_ps(); 8];
6227        let mut v1 = [_mm512_setzero_ps(); 8];
6228        let pairs = gpr / 2;
6229        let unpack = |r: usize, gi: usize| -> (__m512i, __mmask64) {
6230            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
6231            let bb = _mm256_loadu_si256(t as *const __m256i);
6232            let lo = _mm256_and_si256(bb, lomask);
6233            let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
6234            let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
6235            let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
6236            let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
6237            let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
6238            (_mm512_abs_epi8(w), _mm512_movepi8_mask(w))
6239        };
6240        for gp in 0..pairs {
6241            let gi = gp * 2;
6242            let (wa0, neg0) = unpack(r0, gi);
6243            let (wa1, neg1) = unpack(r0 + 1, gi);
6244            let off = gi * GROUP_SIZE;
6245            let sv = |sc: &[f32]| {
6246                _mm512_insertf32x8::<1>(
6247                    _mm512_castps256_ps512(_mm256_set1_ps(*sc.get_unchecked(gi))),
6248                    _mm256_set1_ps(*sc.get_unchecked(gi + 1)),
6249                )
6250            };
6251            let s0 = sv(sc0);
6252            let s1 = sv(sc1);
6253            for k in 0..8 {
6254                let xv = _mm512_loadu_si512(xs[k].as_ptr().add(off) as *const __m512i);
6255                let d0 = _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(
6256                    zero,
6257                    wa0,
6258                    _mm512_mask_sub_epi8(xv, neg0, zero, xv),
6259                ));
6260                let d1 = _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(
6261                    zero,
6262                    wa1,
6263                    _mm512_mask_sub_epi8(xv, neg1, zero, xv),
6264                ));
6265                v0[k] = _mm512_fmadd_ps(d0, s0, v0[k]);
6266                v1[k] = _mm512_fmadd_ps(d1, s1, v1[k]);
6267            }
6268        }
6269        let mut acc = [[0f32; 8]; 2];
6270        for k in 0..8 {
6271            acc[0][k] = _mm512_reduce_add_ps(v0[k]);
6272            acc[1][k] = _mm512_reduce_add_ps(v1[k]);
6273        }
6274        if gpr % 2 == 1 {
6275            let off = (gpr - 1) * GROUP_SIZE;
6276            for j in off..off + GROUP_SIZE {
6277                let (w0, sa) = q4tp_outlier(nib, r0, gpr, j, sc0);
6278                let (w1, sb) = q4tp_outlier(nib, r0 + 1, gpr, j, sc1);
6279                for k in 0..8 {
6280                    let x = *xs[k].get_unchecked(j) as f32;
6281                    acc[0][k] += w0 * sa * x;
6282                    acc[1][k] += w1 * sb * x;
6283                }
6284            }
6285        }
6286        acc
6287    }
6288}
6289
6290/// The same, eight columns at a time. One unpack then feeds twice as many
6291/// activation streams, so a wide batch reads the weight tile half as
6292/// often; the price is eight accumulators live at once. Measured 9.0 ->
6293/// 8.3 ms at 9216x2304, b=296 on a 48-thread EPYC 9B45.
6294#[cfg(target_arch = "x86_64")]
6295#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
6296unsafe fn dot_q4tp_row_1x8_avx512(
6297    nib: &[u8],
6298    r: usize,
6299    gpr: usize,
6300    xs: [&[i8]; 8],
6301    scales: &[f32],
6302) -> [f32; 8] {
6303    // SAFETY: as dot_q4tp_row_1x4_avx2; caller guarantees the ISA.
6304    unsafe {
6305        use core::arch::x86_64::*;
6306        let lomask = _mm256_set1_epi8(0x0F);
6307        let eight = _mm256_set1_epi8(8);
6308        let zero = _mm512_setzero_si512();
6309        let (mut v0, mut v1, mut v2, mut v3) = (
6310            _mm512_setzero_ps(),
6311            _mm512_setzero_ps(),
6312            _mm512_setzero_ps(),
6313            _mm512_setzero_ps(),
6314        );
6315        let (mut v4, mut v5, mut v6, mut v7) = (
6316            _mm512_setzero_ps(),
6317            _mm512_setzero_ps(),
6318            _mm512_setzero_ps(),
6319            _mm512_setzero_ps(),
6320        );
6321        let pairs = gpr / 2;
6322        for gp in 0..pairs {
6323            let gi = gp * 2;
6324            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
6325            let bb = _mm256_loadu_si256(t as *const __m256i);
6326            let lo = _mm256_and_si256(bb, lomask);
6327            let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
6328            // `unpack` works per 128-bit lane, so the halves come out as
6329            // [A.lo, B.lo] and [A.hi, B.hi]; the shuffle reorders the four
6330            // 128-bit lanes into the weights' natural order, which is what
6331            // the straight activation load expects.
6332            let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
6333            let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
6334            let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
6335            let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
6336            let wabs = _mm512_abs_epi8(w);
6337            let neg = _mm512_movepi8_mask(w);
6338            let off = gi * GROUP_SIZE;
6339            let sv = _mm512_insertf32x8::<1>(
6340                _mm512_castps256_ps512(_mm256_set1_ps(*scales.get_unchecked(gi))),
6341                _mm256_set1_ps(*scales.get_unchecked(gi + 1)),
6342            );
6343            let dot = |x: &[i8]| -> __m512 {
6344                let xv = _mm512_loadu_si512(x.as_ptr().add(off) as *const __m512i);
6345                let sx = _mm512_mask_sub_epi8(xv, neg, zero, xv);
6346                _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(zero, wabs, sx))
6347            };
6348            v0 = _mm512_fmadd_ps(dot(xs[0]), sv, v0);
6349            v1 = _mm512_fmadd_ps(dot(xs[1]), sv, v1);
6350            v2 = _mm512_fmadd_ps(dot(xs[2]), sv, v2);
6351            v3 = _mm512_fmadd_ps(dot(xs[3]), sv, v3);
6352            v4 = _mm512_fmadd_ps(dot(xs[4]), sv, v4);
6353            v5 = _mm512_fmadd_ps(dot(xs[5]), sv, v5);
6354            v6 = _mm512_fmadd_ps(dot(xs[6]), sv, v6);
6355            v7 = _mm512_fmadd_ps(dot(xs[7]), sv, v7);
6356        }
6357        let mut acc = [
6358            _mm512_reduce_add_ps(v0),
6359            _mm512_reduce_add_ps(v1),
6360            _mm512_reduce_add_ps(v2),
6361            _mm512_reduce_add_ps(v3),
6362            _mm512_reduce_add_ps(v4),
6363            _mm512_reduce_add_ps(v5),
6364            _mm512_reduce_add_ps(v6),
6365            _mm512_reduce_add_ps(v7),
6366        ];
6367        // An odd group count leaves one group over; the narrow kernel
6368        // finishes it rather than the tail being a special case here.
6369        if gpr % 2 == 1 {
6370            let off = (gpr - 1) * GROUP_SIZE;
6371            for j in off..off + GROUP_SIZE {
6372                let (w, s) = q4tp_outlier(nib, r, gpr, j, scales);
6373                let ws = w * s;
6374                for k in 0..8 {
6375                    acc[k] += ws * *xs[k].get_unchecked(j) as f32;
6376                }
6377            }
6378        }
6379        acc
6380    }
6381}
6382
6383/// The same four columns, 512 bits wide. Two groups (64 weights) ride one
6384/// unpack and one `vpdpbusd`, where AVX2 needs two unpacks and four
6385/// `maddubs`/`madd` pairs — about 2.3x fewer instructions for the same
6386/// arithmetic. The two groups carry different scales, so the fma takes a
6387/// vector whose halves hold each group's scale rather than a broadcast.
6388///
6389/// There is no 512-bit `vpsignb`, so the activation's sign is applied by
6390/// negating under a mask taken from the weight's sign bits. That mask is
6391/// per-tile, so it is hoisted out of the column loop and the per-column
6392/// cost stays exactly one instruction, as with `sign_epi8`. Weights of
6393/// zero are not zeroed by the mask trick and do not need to be: their
6394/// magnitude is zero, so the product is.
6395#[cfg(target_arch = "x86_64")]
6396#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
6397unsafe fn dot_q4tp_row_1x4_avx512(
6398    nib: &[u8],
6399    r: usize,
6400    gpr: usize,
6401    xs: [&[i8]; 4],
6402    scales: &[f32],
6403) -> [f32; 4] {
6404    // SAFETY: as dot_q4tp_row_1x4_avx2; caller guarantees the ISA.
6405    unsafe {
6406        use core::arch::x86_64::*;
6407        let lomask = _mm256_set1_epi8(0x0F);
6408        let eight = _mm256_set1_epi8(8);
6409        let zero = _mm512_setzero_si512();
6410        let (mut v0, mut v1, mut v2, mut v3) = (
6411            _mm512_setzero_ps(),
6412            _mm512_setzero_ps(),
6413            _mm512_setzero_ps(),
6414            _mm512_setzero_ps(),
6415        );
6416        let pairs = gpr / 2;
6417        for gp in 0..pairs {
6418            let gi = gp * 2;
6419            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
6420            let bb = _mm256_loadu_si256(t as *const __m256i);
6421            let lo = _mm256_and_si256(bb, lomask);
6422            let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
6423            // `unpack` works per 128-bit lane, so the halves come out as
6424            // [A.lo, B.lo] and [A.hi, B.hi]; the shuffle reorders the four
6425            // 128-bit lanes into the weights' natural order, which is what
6426            // the straight activation load expects.
6427            let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
6428            let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
6429            let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
6430            let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
6431            let wabs = _mm512_abs_epi8(w);
6432            let neg = _mm512_movepi8_mask(w);
6433            let off = gi * GROUP_SIZE;
6434            let sv = _mm512_insertf32x8::<1>(
6435                _mm512_castps256_ps512(_mm256_set1_ps(*scales.get_unchecked(gi))),
6436                _mm256_set1_ps(*scales.get_unchecked(gi + 1)),
6437            );
6438            let dot = |x: &[i8]| -> __m512 {
6439                let xv = _mm512_loadu_si512(x.as_ptr().add(off) as *const __m512i);
6440                let sx = _mm512_mask_sub_epi8(xv, neg, zero, xv);
6441                _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(zero, wabs, sx))
6442            };
6443            v0 = _mm512_fmadd_ps(dot(xs[0]), sv, v0);
6444            v1 = _mm512_fmadd_ps(dot(xs[1]), sv, v1);
6445            v2 = _mm512_fmadd_ps(dot(xs[2]), sv, v2);
6446            v3 = _mm512_fmadd_ps(dot(xs[3]), sv, v3);
6447        }
6448        let mut acc = [
6449            _mm512_reduce_add_ps(v0),
6450            _mm512_reduce_add_ps(v1),
6451            _mm512_reduce_add_ps(v2),
6452            _mm512_reduce_add_ps(v3),
6453        ];
6454        // An odd group count leaves one group over; the narrow kernel
6455        // finishes it rather than the tail being a special case here.
6456        if gpr % 2 == 1 {
6457            let off = (gpr - 1) * GROUP_SIZE;
6458            for j in off..off + GROUP_SIZE {
6459                let (w, s) = q4tp_outlier(nib, r, gpr, j, scales);
6460                let ws = w * s;
6461                for k in 0..4 {
6462                    acc[k] += ws * *xs[k].get_unchecked(j) as f32;
6463                }
6464            }
6465        }
6466        acc
6467    }
6468}
6469
6470/// Four batch columns against one q4tp row: the tile is unpacked ONCE and
6471/// spent on four activation streams, which is where a prefill batch stops
6472/// being weight-bandwidth-bound. Twin of `dot_q4t_row_1x4_sdot`.
6473#[cfg(target_arch = "aarch64")]
6474#[target_feature(enable = "neon,dotprod")]
6475unsafe fn dot_q4tp_row_1x4_sdot(
6476    nib: &[u8],
6477    r: usize,
6478    gpr: usize,
6479    xs: [&[i8]; 4],
6480    scales: &[f32],
6481) -> [f32; 4] {
6482    // SAFETY: see dot_q4tp_row_sdot; every xs[k] is gpr·GROUP_SIZE long.
6483    unsafe {
6484        use core::arch::aarch64::*;
6485        use core::arch::asm;
6486        let lomask = vdupq_n_u8(0x0F);
6487        let eight = vdupq_n_s8(8);
6488        // Named accumulators, NOT an array indexed by a loop variable: the
6489        // latter does not stay in registers (the same defect cost 2x in the
6490        // AVX2 q4t kernel and again in WGSL).
6491        //
6492        // They are VECTORS, and the horizontal add happens once at the end
6493        // instead of once per group per column. `vaddvq` is a cross-lane
6494        // reduction — with 72 groups and four columns the old shape paid
6495        // 288 of them per row, each one a dependency stall the pipeline
6496        // cannot hide, to save four float adds. The group's scale now
6497        // rides an fma into the lane accumulators, so the arithmetic per
6498        // group is one convert and one fma. Summation order changes (the
6499        // lanes carry independent partial sums), which is the same
6500        // round-off class the SDOT path already lives in — the strict
6501        // kernel (`CMF_SDOT=0`, what `cortiq ppl` runs) is unchanged and
6502        // stays the reference.
6503        let (mut v0, mut v1, mut v2, mut v3) = (
6504            vdupq_n_f32(0.0),
6505            vdupq_n_f32(0.0),
6506            vdupq_n_f32(0.0),
6507            vdupq_n_f32(0.0),
6508        );
6509        for gi in 0..gpr {
6510            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
6511            let s = *scales.get_unchecked(gi);
6512            let bb = vld1q_u8(t);
6513            let lo = vandq_u8(bb, lomask);
6514            let hi = vshrq_n_u8::<4>(bb);
6515            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
6516            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
6517            let off = gi * GROUP_SIZE;
6518            let dot4 = |x: &[i8]| -> int32x4_t {
6519                let x0 = vld1q_s8(x.as_ptr().add(off));
6520                let x1 = vld1q_s8(x.as_ptr().add(off + 16));
6521                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6522                asm!(
6523                    "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
6524                    "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
6525                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6526                    e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
6527                    options(pure, nomem, nostack),
6528                );
6529                vaddq_s32(a0, a1)
6530            };
6531            v0 = vfmaq_n_f32(v0, vcvtq_f32_s32(dot4(xs[0])), s);
6532            v1 = vfmaq_n_f32(v1, vcvtq_f32_s32(dot4(xs[1])), s);
6533            v2 = vfmaq_n_f32(v2, vcvtq_f32_s32(dot4(xs[2])), s);
6534            v3 = vfmaq_n_f32(v3, vcvtq_f32_s32(dot4(xs[3])), s);
6535        }
6536        [
6537            vaddvq_f32(v0),
6538            vaddvq_f32(v1),
6539            vaddvq_f32(v2),
6540            vaddvq_f32(v3),
6541        ]
6542    }
6543}
6544
6545/// Fused q4tp matmat — the same three arms `q4t_matmat` has. Shipping only
6546/// the scalar one made Nanbeige-3B decode at 1.2 tok/s against q4t's 5.9:
6547/// the format was fine, the missing arms were the whole regression.
6548///
6549/// Under `row_exact()` every cell is summed in `q4tp_matvec`'s order on
6550/// every architecture; the mode is read once, so one call never mixes
6551/// the two contracts when a concurrent scope opens or closes mid-call.
6552fn q4tp_matmat(
6553    bytes: &[u8],
6554    xs_all: &[f32],
6555    b: usize,
6556    rows: usize,
6557    cols: usize,
6558    out: &mut [f32],
6559    pool: Option<&Pool>,
6560) {
6561    q4tp_matmat_with(bytes, xs_all, b, rows, cols, out, pool, row_exact())
6562}
6563
6564/// `q4tp_matmat` with the row-exact mode passed in rather than read from
6565/// the shared counter: `exact` = each cell equals its token's matvec bit
6566/// for bit, otherwise the fast blocked / AMX arms are free to reorder.
6567#[allow(clippy::too_many_arguments)]
6568fn q4tp_matmat_with(
6569    bytes: &[u8],
6570    xs_all: &[f32],
6571    b: usize,
6572    rows: usize,
6573    cols: usize,
6574    out: &mut [f32],
6575    pool: Option<&Pool>,
6576    exact: bool,
6577) {
6578    debug_assert_eq!(out.len(), b * rows);
6579    let gpr = cols / GROUP_SIZE;
6580    let v = Q4tpView::new(bytes, rows, cols);
6581
6582    // Wide batches ride the AMX through a dequant-tile sgemm, as in q4t.
6583    // An f32 GEMM over dequantized weights is not the matvec's int8 sum,
6584    // so a row-exact batch never takes it.
6585    #[cfg(target_os = "macos")]
6586    if !exact && b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
6587        dequant_matmat_accel(
6588            &|r, dst| {
6589                let mut sc = [0f32; 32];
6590                let mut scv;
6591                let s: &[f32] = if gpr <= 32 {
6592                    v.scales_into(r, gpr, &mut sc);
6593                    &sc[..gpr]
6594                } else {
6595                    scv = vec![0f32; gpr];
6596                    v.scales_into(r, gpr, &mut scv);
6597                    &scv
6598                };
6599                for gi in 0..gpr {
6600                    let tile = &v.nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
6601                    for (k, &bb) in tile.iter().enumerate() {
6602                        dst[gi * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s[gi];
6603                        dst[gi * GROUP_SIZE + k * 2 + 1] =
6604                            (((bb >> 4) & 0x0F) as f32 - 8.0) * s[gi];
6605                    }
6606                }
6607            },
6608            xs_all,
6609            b,
6610            rows,
6611            cols,
6612            out,
6613            pool,
6614        );
6615        return;
6616    }
6617
6618    let out_addr = SendMut(out.as_mut_ptr());
6619    if a8w8_enabled() {
6620        let acts: Vec<SplitAct> = (0..b)
6621            .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
6622            .collect();
6623        let acts = &acts;
6624        // ARM stays blocked under `exact` too: the 1x4 kernel then runs in
6625        // its v1 shape, whose per-group `int dot as f32 * scale` and scalar
6626        // running sum are exactly `dot_q4tp_row_sdot`'s, so the tile is
6627        // still unpacked once for four columns and each column equals its
6628        // matvec. The tuned shape (fma into lane partials, one horizontal
6629        // add per row) is 1-16 ulp off the matvec and is kept for `!exact`.
6630        #[cfg(target_arch = "aarch64")]
6631        let blocked_ok = sdot_enabled() && blocked_enabled();
6632        // x86 gets the same blocking: one tile unpack spent on four
6633        // columns. Without it every column re-decoded the row, which is
6634        // why a 48-core EPYC measured a sixth of an M4's per-core rate.
6635        // The gate is `avx2_enabled`, as in q4t — `sdot_enabled` answers
6636        // for ARM's dotprod and is hard-wired false everywhere else, so
6637        // asking it here left the whole blocked path unreachable on x86.
6638        #[cfg(target_arch = "x86_64")]
6639        let blocked_ok = q4tp_blocked_x86() && !exact;
6640        #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
6641        let blocked_ok = {
6642            let _ = exact;
6643            false
6644        };
6645        // Columns are swept in panels that fit L2. Without this a
6646        // row-pair walks every activation in the batch — 4.8 MB at
6647        // 512x512 — and does it again for the next pair, so the whole
6648        // batch streams out of the shared cache once per row. Measured
6649        // 800 GB/s of it, flat across batch sizes, which is the signature
6650        // of a loop bound by traffic rather than by arithmetic. A panel of
6651        // 256 columns is 590 KB beside 221 KB of this worker's weights:
6652        // both stay resident and the batch crosses L3 once instead of
6653        // once per row.
6654        let panel_cols: usize = std::env::var("CMF_Q4TP_PANEL")
6655            .ok()
6656            .and_then(|v| v.parse().ok())
6657            .filter(|v| *v > 0)
6658            .unwrap_or(256);
6659        let run = |start: usize, end: usize| {
6660            for abase in (0..acts.len()).step_by(panel_cols) {
6661                let alen = (acts.len() - abase).min(panel_cols);
6662                let mut sc = vec![0f32; gpr];
6663                #[cfg(target_arch = "x86_64")]
6664                let mut r_lo = start;
6665                #[cfg(target_arch = "x86_64")]
6666                if blocked_ok && alen >= 8 {
6667                    let mut sc1 = vec![0f32; gpr];
6668                    while r_lo + 2 <= end {
6669                        v.scales_into(r_lo, gpr, &mut sc);
6670                        v.scales_into(r_lo + 1, gpr, &mut sc1);
6671                        let mut bi = 0usize;
6672                        while bi + 8 <= alen {
6673                            let xs = [
6674                                acts[abase + bi].xq.as_slice(),
6675                                acts[abase + bi + 1].xq.as_slice(),
6676                                acts[abase + bi + 2].xq.as_slice(),
6677                                acts[abase + bi + 3].xq.as_slice(),
6678                                acts[abase + bi + 4].xq.as_slice(),
6679                                acts[abase + bi + 5].xq.as_slice(),
6680                                acts[abase + bi + 6].xq.as_slice(),
6681                                acts[abase + bi + 7].xq.as_slice(),
6682                            ];
6683                            let d = unsafe { dot_q4tp_2x8_avx512(v.nib, r_lo, gpr, xs, &sc, &sc1) };
6684                            for (row, dr, scr) in [(r_lo, &d[0], &sc), (r_lo + 1, &d[1], &sc1)] {
6685                                for k in 0..8 {
6686                                    let act = &acts[abase + bi + k];
6687                                    let mut acc = dr[k] * act.sx;
6688                                    for &(j, xv) in &act.outliers {
6689                                        let (w, s) = q4tp_outlier(v.nib, row, gpr, j, scr);
6690                                        acc += w * s * xv;
6691                                    }
6692                                    // SAFETY: disjoint (bi, r) cells per worker.
6693                                    unsafe { *out_addr.at((abase + bi + k) * rows + row) = acc };
6694                                }
6695                            }
6696                            bi += 8;
6697                        }
6698                        // Columns past the last group of eight, both rows —
6699                        // the same single-row kernel the tail below uses.
6700                        for row in [r_lo, r_lo + 1] {
6701                            let scr: &[f32] = if row == r_lo { &sc } else { &sc1 };
6702                            for b2 in bi..alen {
6703                                let act = &acts[abase + b2];
6704                                let xs4 = [
6705                                    act.xq.as_slice(),
6706                                    act.xq.as_slice(),
6707                                    act.xq.as_slice(),
6708                                    act.xq.as_slice(),
6709                                ];
6710                                let d =
6711                                    unsafe { dot_q4tp_row_1x4_avx512(v.nib, row, gpr, xs4, scr) };
6712                                let mut acc = d[0] * act.sx;
6713                                for &(j, xv) in &act.outliers {
6714                                    let (w, s) = q4tp_outlier(v.nib, row, gpr, j, scr);
6715                                    acc += w * s * xv;
6716                                }
6717                                // SAFETY: disjoint (bi, r) cells per worker.
6718                                unsafe { *out_addr.at((abase + b2) * rows + row) = acc };
6719                            }
6720                        }
6721                        r_lo += 2;
6722                    }
6723                }
6724                #[cfg(target_arch = "x86_64")]
6725                let row_start = r_lo;
6726                #[cfg(not(target_arch = "x86_64"))]
6727                let row_start = start;
6728                for r in row_start..end {
6729                    v.scales_into(r, gpr, &mut sc);
6730                    let mut bi = 0usize;
6731                    #[cfg(target_arch = "x86_64")]
6732                    if blocked_ok {
6733                        while bi + 8 <= alen {
6734                            let xs = [
6735                                acts[abase + bi].xq.as_slice(),
6736                                acts[abase + bi + 1].xq.as_slice(),
6737                                acts[abase + bi + 2].xq.as_slice(),
6738                                acts[abase + bi + 3].xq.as_slice(),
6739                                acts[abase + bi + 4].xq.as_slice(),
6740                                acts[abase + bi + 5].xq.as_slice(),
6741                                acts[abase + bi + 6].xq.as_slice(),
6742                                acts[abase + bi + 7].xq.as_slice(),
6743                            ];
6744                            let d = unsafe { dot_q4tp_row_1x8_avx512(v.nib, r, gpr, xs, &sc) };
6745                            for k in 0..8 {
6746                                let act = &acts[abase + bi + k];
6747                                let mut acc = d[k] * act.sx;
6748                                for &(j, xv) in &act.outliers {
6749                                    let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
6750                                    acc += w * s * xv;
6751                                }
6752                                // SAFETY: disjoint (bi, r) cells per worker.
6753                                unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
6754                            }
6755                            bi += 8;
6756                        }
6757                        while bi + 4 <= alen {
6758                            let xs = [
6759                                acts[abase + bi].xq.as_slice(),
6760                                acts[abase + bi + 1].xq.as_slice(),
6761                                acts[abase + bi + 2].xq.as_slice(),
6762                                acts[abase + bi + 3].xq.as_slice(),
6763                            ];
6764                            let d = unsafe { dot_q4tp_row_1x4_avx512(v.nib, r, gpr, xs, &sc) };
6765                            for k in 0..4 {
6766                                let act = &acts[abase + bi + k];
6767                                let mut acc = d[k] * act.sx;
6768                                for &(j, xv) in &act.outliers {
6769                                    let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
6770                                    acc += w * s * xv;
6771                                }
6772                                // SAFETY: disjoint (bi, r) cells per worker.
6773                                unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
6774                            }
6775                            bi += 4;
6776                        }
6777                    }
6778                    #[cfg(target_arch = "aarch64")]
6779                    if blocked_ok {
6780                        while bi + 4 <= alen {
6781                            let xs = [
6782                                acts[abase + bi].xq.as_slice(),
6783                                acts[abase + bi + 1].xq.as_slice(),
6784                                acts[abase + bi + 2].xq.as_slice(),
6785                                acts[abase + bi + 3].xq.as_slice(),
6786                            ];
6787                            let d = unsafe {
6788                                if exact || q4tp_v1() {
6789                                    dot_q4tp_row_1x4_sdot_v1(v.nib, r, gpr, xs, &sc)
6790                                } else {
6791                                    dot_q4tp_row_1x4_sdot(v.nib, r, gpr, xs, &sc)
6792                                }
6793                            };
6794                            for k in 0..4 {
6795                                let act = &acts[abase + bi + k];
6796                                let mut acc = d[k] * act.sx;
6797                                for &(j, xv) in &act.outliers {
6798                                    let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
6799                                    acc += w * s * xv;
6800                                }
6801                                // SAFETY: disjoint (bi, r) cells per worker.
6802                                unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
6803                            }
6804                            bi += 4;
6805                        }
6806                    }
6807                    let _ = blocked_ok;
6808                    while bi < alen {
6809                        let act = &acts[abase + bi];
6810                        let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, &sc) * act.sx;
6811                        for &(j, xv) in &act.outliers {
6812                            let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
6813                            acc += w * s * xv;
6814                        }
6815                        // SAFETY: disjoint (bi, r) cells per worker range.
6816                        unsafe { *out_addr.at((abase + bi) * rows + r) = acc };
6817                        bi += 1;
6818                    }
6819                }
6820            }
6821        };
6822        dispatch_rows(pool, rows, &run);
6823        return;
6824    }
6825
6826    let run = |start: usize, end: usize| {
6827        let mut sc = vec![0f32; gpr];
6828        for r in start..end {
6829            v.scales_into(r, gpr, &mut sc);
6830            for bi in 0..b {
6831                let x = &xs_all[bi * cols..(bi + 1) * cols];
6832                // SAFETY: disjoint (bi, r) cells per worker range.
6833                unsafe { *out_addr.at(bi * rows + r) = q4tp_row_exact(v.nib, r, gpr, x, &sc) };
6834            }
6835        }
6836    };
6837    dispatch_rows(pool, rows, &run);
6838}
6839
6840/// Fused q4_tiled matvec (dispatch mirrors `q4matvec`).
6841fn q4t_matvec(
6842    bytes: &[u8],
6843    x: &[f32],
6844    rows: usize,
6845    cols: usize,
6846    out: &mut [f32],
6847    pool: Option<&Pool>,
6848) {
6849    debug_assert_eq!(out.len(), rows);
6850    let gpr = cols / GROUP_SIZE;
6851    let out_addr = SendMut(out.as_mut_ptr());
6852    if a8w8_enabled() {
6853        let act = split_act(x);
6854        let run = move |start: usize, end: usize| {
6855            for r in start..end {
6856                let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
6857                for &(j, xv) in &act.outliers {
6858                    let (w, s) = q4t_outlier(bytes, r, gpr, j);
6859                    acc += w * s * xv;
6860                }
6861                // SAFETY: disjoint row ranges per worker.
6862                unsafe { *out_addr.at(r) = acc };
6863            }
6864        };
6865        dispatch_rows(pool, rows, &run);
6866        return;
6867    }
6868    let run = move |start: usize, end: usize| {
6869        for r in start..end {
6870            // SAFETY: disjoint row ranges per worker.
6871            unsafe { *out_addr.at(r) = q4t_row_exact(bytes, r, gpr, x) };
6872        }
6873    };
6874    dispatch_rows(pool, rows, &run);
6875}
6876
6877/// Fused two-input q4_tiled matvec (weights read once per pair).
6878#[allow(clippy::too_many_arguments)]
6879fn q4t_matvec2(
6880    bytes: &[u8],
6881    x1: &[f32],
6882    x2: &[f32],
6883    rows: usize,
6884    cols: usize,
6885    o1: &mut [f32],
6886    o2: &mut [f32],
6887    pool: Option<&Pool>,
6888) {
6889    let gpr = cols / GROUP_SIZE;
6890    let p1 = SendMut(o1.as_mut_ptr());
6891    let p2 = SendMut(o2.as_mut_ptr());
6892    if a8w8_enabled() {
6893        let a1 = split_act(x1);
6894        let a2 = split_act(x2);
6895        let run = move |start: usize, end: usize| {
6896            for r in start..end {
6897                let mut v1 = dot_q4t_row_i8(bytes, r, gpr, &a1.xq) * a1.sx;
6898                let mut v2 = dot_q4t_row_i8(bytes, r, gpr, &a2.xq) * a2.sx;
6899                for &(j, xv) in &a1.outliers {
6900                    let (w, s) = q4t_outlier(bytes, r, gpr, j);
6901                    v1 += w * s * xv;
6902                }
6903                for &(j, xv) in &a2.outliers {
6904                    let (w, s) = q4t_outlier(bytes, r, gpr, j);
6905                    v2 += w * s * xv;
6906                }
6907                // SAFETY: disjoint row ranges per worker.
6908                unsafe {
6909                    *p1.at(r) = v1;
6910                    *p2.at(r) = v2;
6911                }
6912            }
6913        };
6914        dispatch_rows(pool, rows, &run);
6915        return;
6916    }
6917    let run = move |start: usize, end: usize| {
6918        for r in start..end {
6919            // SAFETY: disjoint row ranges per worker.
6920            unsafe {
6921                *p1.at(r) = q4t_row_exact(bytes, r, gpr, x1);
6922                *p2.at(r) = q4t_row_exact(bytes, r, gpr, x2);
6923            }
6924        }
6925    };
6926    dispatch_rows(pool, rows, &run);
6927}
6928
6929/// Batched q4_tiled matmat: each row's tiles stream once per microbatch.
6930#[allow(clippy::too_many_arguments)]
6931/// Prefill GEMM through Accelerate for group-quantized codecs: a
6932/// caller-supplied row dequantizer fills f32 tiles (pool-parallel) and
6933/// each tile rides the AMX with one sgemm — the generic sibling of
6934/// `qmatmat_accel` (q8). Numerics are f32-GEMM (tolerance class);
6935/// decode (b=1) never takes this path.
6936#[cfg(target_os = "macos")]
6937fn dequant_matmat_accel(
6938    dequant_row: &(dyn Fn(usize, &mut [f32]) + Sync),
6939    xs_all: &[f32],
6940    b: usize,
6941    rows: usize,
6942    cols: usize,
6943    out: &mut [f32],
6944    pool: Option<&Pool>,
6945) {
6946    const TR: usize = 2048;
6947    thread_local! {
6948        static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
6949    }
6950    WTILE.with(|wt| {
6951        let mut wtile = wt.borrow_mut();
6952        wtile.resize(TR * cols, 0.0);
6953        let mut r0 = 0usize;
6954        while r0 < rows {
6955            let tr = TR.min(rows - r0);
6956            let wt_addr = SendMut(wtile.as_mut_ptr());
6957            let run = |start: usize, end: usize| {
6958                for r in start..end {
6959                    // SAFETY: workers cover disjoint r ranges.
6960                    let dst = unsafe { std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols) };
6961                    dequant_row(r0 + r, dst);
6962                }
6963            };
6964            dispatch_rows(pool, tr, &run);
6965            unsafe {
6966                accel_blas::cblas_sgemm(
6967                    101, // RowMajor
6968                    111, // NoTrans A
6969                    112, // Trans B
6970                    b as i32,
6971                    tr as i32,
6972                    cols as i32,
6973                    1.0,
6974                    xs_all.as_ptr(),
6975                    cols as i32,
6976                    wtile.as_ptr(),
6977                    cols as i32,
6978                    0.0,
6979                    out.as_mut_ptr().add(r0),
6980                    rows as i32,
6981                );
6982            }
6983            r0 += tr;
6984        }
6985    });
6986}
6987
6988fn q4t_matmat(
6989    bytes: &[u8],
6990    xs_all: &[f32],
6991    b: usize,
6992    rows: usize,
6993    cols: usize,
6994    out: &mut [f32],
6995    pool: Option<&Pool>,
6996) {
6997    debug_assert_eq!(out.len(), b * rows);
6998    let gpr = cols / GROUP_SIZE;
6999    // Wide batches ride the AMX like q8's qmatmat: on Apple silicon
7000    // the dequant-tile sgemm is an order above the SDOT row loop for
7001    // prefill shapes (imagegen DiT forwards are exactly this).
7002    #[cfg(target_os = "macos")]
7003    if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
7004        dequant_matmat_accel(
7005            &|r, dst| {
7006                for gi in 0..gpr {
7007                    let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
7008                    let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
7009                    for (k, &bb) in tile[2..].iter().enumerate() {
7010                        dst[gi * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s;
7011                        dst[gi * GROUP_SIZE + k * 2 + 1] = (((bb >> 4) & 0x0F) as f32 - 8.0) * s;
7012                    }
7013                }
7014            },
7015            xs_all,
7016            b,
7017            rows,
7018            cols,
7019            out,
7020            pool,
7021        );
7022        return;
7023    }
7024    let out_addr = SendMut(out.as_mut_ptr());
7025    if a8w8_enabled() {
7026        let acts: Vec<SplitAct> = (0..b)
7027            .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
7028            .collect();
7029        let acts = &acts;
7030        #[cfg(target_arch = "x86_64")]
7031        let blocked_ok = avx2_enabled() && blocked_enabled();
7032        #[cfg(target_arch = "aarch64")]
7033        let blocked_ok = sdot_enabled() && blocked_enabled();
7034        #[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))]
7035        let blocked_ok = false;
7036        let run = move |start: usize, end: usize| {
7037            for r in start..end {
7038                let mut bi = 0usize;
7039                #[cfg(target_arch = "aarch64")]
7040                if blocked_ok {
7041                    while bi + 4 <= acts.len() {
7042                        let xs = [
7043                            acts[bi].xq.as_slice(),
7044                            acts[bi + 1].xq.as_slice(),
7045                            acts[bi + 2].xq.as_slice(),
7046                            acts[bi + 3].xq.as_slice(),
7047                        ];
7048                        let d = unsafe { dot_q4t_row_1x4_sdot(bytes, r, gpr, xs) };
7049                        for k in 0..4 {
7050                            let act = &acts[bi + k];
7051                            let mut acc = d[k] * act.sx;
7052                            for &(j, xv) in &act.outliers {
7053                                let (w, sc) = q4t_outlier(bytes, r, gpr, j);
7054                                acc += w * sc * xv;
7055                            }
7056                            // SAFETY: disjoint (bi, r) cells per worker.
7057                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
7058                        }
7059                        bi += 4;
7060                    }
7061                }
7062                #[cfg(target_arch = "x86_64")]
7063                if blocked_ok {
7064                    while bi + 4 <= acts.len() {
7065                        let xs = [
7066                            acts[bi].xq.as_slice(),
7067                            acts[bi + 1].xq.as_slice(),
7068                            acts[bi + 2].xq.as_slice(),
7069                            acts[bi + 3].xq.as_slice(),
7070                        ];
7071                        let d = unsafe {
7072                            if vnni_tiles_enabled() {
7073                                dot_q4t_row_1x4_vnni(bytes, r, gpr, xs)
7074                            } else {
7075                                dot_q4t_row_1x4_avx2(bytes, r, gpr, xs)
7076                            }
7077                        };
7078                        for k in 0..4 {
7079                            let act = &acts[bi + k];
7080                            let mut acc = d[k] * act.sx;
7081                            for &(j, xv) in &act.outliers {
7082                                let (w, sc) = q4t_outlier(bytes, r, gpr, j);
7083                                acc += w * sc * xv;
7084                            }
7085                            // SAFETY: disjoint (bi, r) cells per worker.
7086                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
7087                        }
7088                        bi += 4;
7089                    }
7090                }
7091                let _ = blocked_ok;
7092                while bi < acts.len() {
7093                    let act = &acts[bi];
7094                    let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
7095                    for &(j, xv) in &act.outliers {
7096                        let (w, s) = q4t_outlier(bytes, r, gpr, j);
7097                        acc += w * s * xv;
7098                    }
7099                    // SAFETY: disjoint (bi, r) cells per worker range.
7100                    unsafe { *out_addr.at(bi * rows + r) = acc };
7101                    bi += 1;
7102                }
7103            }
7104        };
7105        dispatch_rows(pool, rows, &run);
7106        return;
7107    }
7108    let run = move |start: usize, end: usize| {
7109        for r in start..end {
7110            for bi in 0..b {
7111                let x = &xs_all[bi * cols..(bi + 1) * cols];
7112                // SAFETY: disjoint (bi, r) cells per worker range.
7113                unsafe { *out_addr.at(bi * rows + r) = q4t_row_exact(bytes, r, gpr, x) };
7114            }
7115        }
7116    };
7117    dispatch_rows(pool, rows, &run);
7118}
7119
7120// ── q1 (dtype 12): binary weights, [f16 scale][4B sign bits] per
7121// 32-group tile. The kernel family mirrors q4_tiled: one sequential
7122// stream of 6-byte tiles, per-tile integer dot × scale, exact outlier
7123// correction (A8W8 contract), exact scalar path under CMF_SDOT=0. ──
7124
7125/// Per-32-group sums of the quantized activation — the ±1 identity's
7126/// shared half: `dot = −2·sdot(mask, x) − gsum[g]`, computed ONCE per
7127/// matvec and reused by every row.
7128fn q1_group_sums(xq: &[i8], gpr: usize) -> Vec<i32> {
7129    (0..gpr)
7130        .map(|gi| {
7131            xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE]
7132                .iter()
7133                .map(|&v| v as i32)
7134                .sum()
7135        })
7136        .collect()
7137}
7138
7139/// One q1 row via the A8W8 int8 path — mask-SDOT on ARM (no ±1
7140/// expansion at all), scalar bit loop elsewhere (AVX2 queued with the
7141/// x86 pass).
7142#[inline]
7143#[allow(unreachable_code)]
7144/// AVX2 q1 row via the same ±1 identity as the ARM sdot kernel: the
7145/// sign bits expand to a {0, −1} byte mask through shuffle+cmpeq, the
7146/// masked activation sums through maddubs(1, x&mask), and
7147/// `dot = −(2·masked_sum + Σx_group)` — bit-identical integer math.
7148#[cfg(target_arch = "x86_64")]
7149#[target_feature(enable = "avx2")]
7150unsafe fn dot_q1_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
7151    // SAFETY: callers uphold the 6B-tile and xq/gsum length contracts.
7152    unsafe {
7153        use core::arch::x86_64::*;
7154        // Byte j of the mask must replicate bits-byte j/8.
7155        let expand = _mm256_setr_epi8(
7156            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,
7157            3, 3, 3,
7158        );
7159        let bitsel = _mm256_setr_epi8(
7160            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
7161            -128, 1, 2, 4, 8, 16, 32, 64, -128,
7162        );
7163        let ones8 = _mm256_set1_epi8(1);
7164        let ones16 = _mm256_set1_epi16(1);
7165        let mut acc = 0f32;
7166        for gi in 0..gpr {
7167            let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
7168            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
7169            let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
7170            let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
7171            let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
7172            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7173            let sel = _mm256_and_si256(x, mask);
7174            // Σ of selected i8 lanes: maddubs(1u8, sel_i8) pairs → madd.
7175            let p16 = _mm256_maddubs_epi16(ones8, sel);
7176            let d32 = _mm256_madd_epi16(p16, ones16);
7177            let hi128 = _mm256_extracti128_si256::<1>(d32);
7178            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
7179            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7180            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7181            let msum = _mm_cvtsi128_si32(s32);
7182            // The and-select keeps x UN-negated (unlike ARM's −1-mask
7183            // sdot): d = Σ_set − Σ_unset = 2·Σ_set − Σ_all.
7184            let d = 2 * msum - gsum[gi];
7185            acc += d as f32 * s;
7186        }
7187        acc
7188    }
7189}
7190
7191/// VNNI twin of `dot_q1_row_avx2`: the masked-select sum goes through
7192/// one `vpdpbusd(1u8, sel)` (see `dpbusd_hsum` — bit-identical).
7193#[cfg(target_arch = "x86_64")]
7194#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
7195unsafe fn dot_q1_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
7196    // SAFETY: callers uphold the 6B-tile and xq/gsum length contracts.
7197    unsafe {
7198        use core::arch::x86_64::*;
7199        let expand = _mm256_setr_epi8(
7200            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,
7201            3, 3, 3,
7202        );
7203        let bitsel = _mm256_setr_epi8(
7204            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
7205            -128, 1, 2, 4, 8, 16, 32, 64, -128,
7206        );
7207        let ones8 = _mm256_set1_epi8(1);
7208        let mut acc = 0f32;
7209        for gi in 0..gpr {
7210            let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
7211            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
7212            let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
7213            let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
7214            let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
7215            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7216            let msum = dpbusd_hsum(ones8, _mm256_and_si256(x, mask));
7217            let d = 2 * msum - gsum[gi];
7218            acc += d as f32 * s;
7219        }
7220        acc
7221    }
7222}
7223
7224/// VNNI twin of `dot_q1_row_1x4_avx2` (see `dpbusd_hsum`).
7225#[cfg(target_arch = "x86_64")]
7226#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
7227unsafe fn dot_q1_row_1x4_vnni(
7228    bytes: &[u8],
7229    r: usize,
7230    gpr: usize,
7231    xs: [&[i8]; 4],
7232    gsums: [&[i32]; 4],
7233) -> [f32; 4] {
7234    // SAFETY: callers uphold the 6B-tile and xq/gsum length contracts.
7235    unsafe {
7236        use core::arch::x86_64::*;
7237        let expand = _mm256_setr_epi8(
7238            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,
7239            3, 3, 3,
7240        );
7241        let bitsel = _mm256_setr_epi8(
7242            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
7243            -128, 1, 2, 4, 8, 16, 32, 64, -128,
7244        );
7245        let ones8 = _mm256_set1_epi8(1);
7246        let mut acc = [0f32; 4];
7247        for gi in 0..gpr {
7248            let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
7249            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
7250            let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
7251            let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
7252            let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
7253            for (k, xq) in xs.iter().enumerate() {
7254                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7255                let msum = dpbusd_hsum(ones8, _mm256_and_si256(x, mask));
7256                let d = 2 * msum - gsums[k][gi];
7257                acc[k] += d as f32 * s;
7258            }
7259        }
7260        acc
7261    }
7262}
7263
7264/// The blocked 1×4 flavor: the expanded bit mask serves four activation
7265/// streams per group (mask build once, four select+reduce chains).
7266#[cfg(target_arch = "x86_64")]
7267#[target_feature(enable = "avx2")]
7268unsafe fn dot_q1_row_1x4_avx2(
7269    bytes: &[u8],
7270    r: usize,
7271    gpr: usize,
7272    xs: [&[i8]; 4],
7273    gsums: [&[i32]; 4],
7274) -> [f32; 4] {
7275    // SAFETY: callers uphold the 6B-tile and xq/gsum length contracts.
7276    unsafe {
7277        use core::arch::x86_64::*;
7278        let expand = _mm256_setr_epi8(
7279            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,
7280            3, 3, 3,
7281        );
7282        let bitsel = _mm256_setr_epi8(
7283            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
7284            -128, 1, 2, 4, 8, 16, 32, 64, -128,
7285        );
7286        let ones8 = _mm256_set1_epi8(1);
7287        let ones16 = _mm256_set1_epi16(1);
7288        let mut acc = [0f32; 4];
7289        for gi in 0..gpr {
7290            let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
7291            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
7292            let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
7293            let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
7294            let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
7295            for (k, xq) in xs.iter().enumerate() {
7296                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7297                let sel = _mm256_and_si256(x, mask);
7298                let p16 = _mm256_maddubs_epi16(ones8, sel);
7299                let d32 = _mm256_madd_epi16(p16, ones16);
7300                let hi128 = _mm256_extracti128_si256::<1>(d32);
7301                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
7302                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7303                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7304                let msum = _mm_cvtsi128_si32(s32);
7305                let d = 2 * msum - gsums[k][gi];
7306                acc[k] += d as f32 * s;
7307            }
7308        }
7309        acc
7310    }
7311}
7312
7313#[allow(unreachable_code)]
7314fn dot_q1_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
7315    #[cfg(target_arch = "aarch64")]
7316    unsafe {
7317        return dot_q1_row_sdot(bytes, r, gpr, xq, gsum);
7318    }
7319    #[cfg(target_arch = "x86_64")]
7320    if avx2_enabled() {
7321        unsafe {
7322            if vnni_tiles_enabled() {
7323                return dot_q1_row_vnni(bytes, r, gpr, xq, gsum);
7324            }
7325            return dot_q1_row_avx2(bytes, r, gpr, xq, gsum);
7326        }
7327    }
7328    let _ = gsum;
7329    let mut acc = 0f32;
7330    for gi in 0..gpr {
7331        let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
7332        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
7333        let mut d = 0i32;
7334        for (j, &b) in tile[2..].iter().enumerate() {
7335            for k in 0..8 {
7336                let w = ((b >> k) & 1) as i32 * 2 - 1;
7337                d += w * xq[gi * GROUP_SIZE + j * 8 + k] as i32;
7338            }
7339        }
7340        acc += d as f32 * s;
7341    }
7342    acc
7343}
7344
7345/// SDOT q1 row via the ±1 identity: the vtst mask (0xFF where the bit
7346/// is set, i.e. −1 as i8) feeds `sdot` DIRECTLY — no expansion to ±1
7347/// lanes at all — and `dot = −(2·sdot(mask, x) + Σx_group)`, with the
7348/// per-group activation sums shared across every row of the matvec.
7349/// Four tiles (128 weights) per iteration: integer dots reduce through
7350/// a vpaddq tree into ONE i32x4 that meets its four scales in a single
7351/// fused f32 multiply-add. Integer math throughout — bit-identical to
7352/// the scalar ±1 reference.
7353#[cfg(target_arch = "aarch64")]
7354#[target_feature(enable = "neon,dotprod")]
7355unsafe fn dot_q1_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
7356    // SAFETY: callers uphold slice-length contracts (6B tile per group,
7357    // xq.len() == gpr·GROUP_SIZE, gsum.len() == gpr).
7358    unsafe {
7359        use core::arch::aarch64::*;
7360        use core::arch::asm;
7361        const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
7362        let m = vld1q_u8(MASKS.as_ptr());
7363        // One tile's −Σ_set(x) as an UNREDUCED i32x4 (two mask-sdots).
7364        macro_rules! tile_dot {
7365            ($t:expr, $x:expr) => {{
7366                let v0 = vcombine_u8(vdup_n_u8(*$t.add(2)), vdup_n_u8(*$t.add(3)));
7367                let v1 = vcombine_u8(vdup_n_u8(*$t.add(4)), vdup_n_u8(*$t.add(5)));
7368                let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
7369                let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
7370                let x0 = vld1q_s8($x);
7371                let x1 = vld1q_s8($x.add(16));
7372                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
7373                asm!(
7374                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
7375                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
7376                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
7377                    w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
7378                    options(pure, nomem, nostack),
7379                );
7380                vaddq_s32(a0, a1)
7381            }};
7382        }
7383        // TBL unpack over PAIR loads: one vld1q covers two 6B tiles
7384        // ([s s b b b b][s s b b b b] + 4B slack), TBL replicates each
7385        // bit-byte across 8 lanes for vtst, and the four scales gather
7386        // through tbl2 into one fcvtl — the 16 ld1r broadcast loads and
7387        // 4 branchy software f16 conversions per 128 weights (the
7388        // measured load-port wall of this kernel) become 2 vector
7389        // loads + 9 table lookups. Integer math order is unchanged —
7390        // bit-identical results (FCVTL is exact on every f16).
7391        const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
7392        const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
7393        const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
7394        const IW11: [u8; 16] = [
7395            10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
7396        ];
7397        const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
7398        let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
7399        let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
7400        let isc = vld1_u8(ISC.as_ptr());
7401        // One tile's −Σ_set(x) from a TBL-unpacked pair load.
7402        macro_rules! tile_dot_tbl {
7403            ($ld:expr, $i0:expr, $i1:expr, $x:expr) => {{
7404                let w0 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i0), m));
7405                let w1 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i1), m));
7406                let x0 = vld1q_s8($x);
7407                let x1 = vld1q_s8($x.add(16));
7408                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
7409                asm!(
7410                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
7411                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
7412                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
7413                    w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
7414                    options(pure, nomem, nostack),
7415                );
7416                vaddq_s32(a0, a1)
7417            }};
7418        }
7419        let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
7420        let row_base = r * gpr * Q1_TILE;
7421        let abs_end = bytes.len();
7422        let xp = xq.as_ptr();
7423        let gp = gsum.as_ptr();
7424        let mut accv = vdupq_n_f32(0.0);
7425        let mut gi = 0;
7426        // The second pair load reads 4B past tile gi+3 — stay inside
7427        // the payload slice (only the file's final tiles fall back).
7428        while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
7429            let t0 = base.add(gi * Q1_TILE);
7430            let ld_a = vld1q_u8(t0);
7431            let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
7432            let d0 = tile_dot_tbl!(ld_a, iw00, iw01, xp.add(gi * GROUP_SIZE));
7433            let d1 = tile_dot_tbl!(ld_a, iw10, iw11, xp.add((gi + 1) * GROUP_SIZE));
7434            let d2 = tile_dot_tbl!(ld_b, iw00, iw01, xp.add((gi + 2) * GROUP_SIZE));
7435            let d3 = tile_dot_tbl!(ld_b, iw10, iw11, xp.add((gi + 3) * GROUP_SIZE));
7436            // [−Σ0, −Σ1, −Σ2, −Σ3] → dots = −(2·Σset_neg + gsum)
7437            let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
7438            let g = vld1q_s32(gp.add(gi));
7439            let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
7440            let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
7441            let scf: float32x4_t;
7442            asm!(
7443                "fcvtl {o:v}.4s, {i:v}.4h",
7444                o = out(vreg) scf, i = in(vreg) sc16,
7445                options(pure, nomem, nostack),
7446            );
7447            accv = vfmaq_f32(accv, vcvtq_f32_s32(dots), scf);
7448            gi += 4;
7449        }
7450        let mut acc = vaddvq_f32(accv);
7451        while gi < gpr {
7452            let t = base.add(gi * Q1_TILE);
7453            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
7454            let d = vaddvq_s32(tile_dot!(t, xp.add(gi * GROUP_SIZE)));
7455            acc += (-(2 * d + *gp.add(gi))) as f32 * s;
7456            gi += 1;
7457        }
7458        acc
7459    }
7460}
7461
7462/// Blocked q1 1×4: one TBL unpack of the tile pair serves FOUR
7463/// activation streams (prefill amortization — the same idea as the
7464/// AVX2 twin; per stream the group order, fma order and tail match the
7465/// single-row kernel exactly, so batch == matvec bit-for-bit).
7466#[cfg(target_arch = "aarch64")]
7467#[target_feature(enable = "neon,dotprod")]
7468unsafe fn dot_q1_row_1x4_sdot(
7469    bytes: &[u8],
7470    r: usize,
7471    gpr: usize,
7472    xs: [&[i8]; 4],
7473    gs: [&[i32]; 4],
7474) -> [f32; 4] {
7475    // SAFETY: same slice-length contracts as `dot_q1_row_sdot`, ×4.
7476    unsafe {
7477        use core::arch::aarch64::*;
7478        use core::arch::asm;
7479        const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
7480        const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
7481        const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
7482        const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
7483        const IW11: [u8; 16] = [
7484            10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
7485        ];
7486        const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
7487        let m = vld1q_u8(MASKS.as_ptr());
7488        let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
7489        let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
7490        let isc = vld1_u8(ISC.as_ptr());
7491        macro_rules! sdot2 {
7492            ($w0:expr, $w1:expr, $x:expr) => {{
7493                let x0 = vld1q_s8($x);
7494                let x1 = vld1q_s8($x.add(16));
7495                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
7496                asm!(
7497                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
7498                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
7499                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
7500                    w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
7501                    options(pure, nomem, nostack),
7502                );
7503                vaddq_s32(a0, a1)
7504            }};
7505        }
7506        let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
7507        let row_base = r * gpr * Q1_TILE;
7508        let abs_end = bytes.len();
7509        let mut accv = [vdupq_n_f32(0.0); 4];
7510        let mut gi = 0;
7511        while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
7512            let t0 = base.add(gi * Q1_TILE);
7513            let ld_a = vld1q_u8(t0);
7514            let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
7515            // Unpack ONCE — eight ±mask vectors serve all four streams.
7516            let w00 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw00), m));
7517            let w01 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw01), m));
7518            let w10 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw10), m));
7519            let w11 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw11), m));
7520            let w20 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw00), m));
7521            let w21 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw01), m));
7522            let w30 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw10), m));
7523            let w31 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw11), m));
7524            let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
7525            let scf: float32x4_t;
7526            asm!(
7527                "fcvtl {o:v}.4s, {i:v}.4h",
7528                o = out(vreg) scf, i = in(vreg) sc16,
7529                options(pure, nomem, nostack),
7530            );
7531            for k in 0..4 {
7532                let xp = xs[k].as_ptr();
7533                let d0 = sdot2!(w00, w01, xp.add(gi * GROUP_SIZE));
7534                let d1 = sdot2!(w10, w11, xp.add((gi + 1) * GROUP_SIZE));
7535                let d2 = sdot2!(w20, w21, xp.add((gi + 2) * GROUP_SIZE));
7536                let d3 = sdot2!(w30, w31, xp.add((gi + 3) * GROUP_SIZE));
7537                let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
7538                let g = vld1q_s32(gs[k].as_ptr().add(gi));
7539                let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
7540                accv[k] = vfmaq_f32(accv[k], vcvtq_f32_s32(dots), scf);
7541            }
7542            gi += 4;
7543        }
7544        let mut acc = [
7545            vaddvq_f32(accv[0]),
7546            vaddvq_f32(accv[1]),
7547            vaddvq_f32(accv[2]),
7548            vaddvq_f32(accv[3]),
7549        ];
7550        while gi < gpr {
7551            let t = base.add(gi * Q1_TILE);
7552            let sc = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
7553            let v0 = vcombine_u8(vdup_n_u8(*t.add(2)), vdup_n_u8(*t.add(3)));
7554            let v1 = vcombine_u8(vdup_n_u8(*t.add(4)), vdup_n_u8(*t.add(5)));
7555            let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
7556            let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
7557            for k in 0..4 {
7558                let d = vaddvq_s32(sdot2!(w0, w1, xs[k].as_ptr().add(gi * GROUP_SIZE)));
7559                acc[k] += (-(2 * d + *gs[k].as_ptr().add(gi))) as f32 * sc;
7560            }
7561            gi += 1;
7562        }
7563        acc
7564    }
7565}
7566
7567/// (weight ±1, scale) of one q1 element — the exact outlier term.
7568#[inline]
7569fn q1_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
7570    let gi = j / GROUP_SIZE;
7571    let k = j % GROUP_SIZE;
7572    let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
7573    let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
7574    let bit = (tile[2 + k / 8] >> (k % 8)) & 1;
7575    ((bit as i32 * 2 - 1) as f32, s)
7576}
7577
7578/// Exact scalar q1 row (CMF_SDOT=0 contract).
7579#[inline]
7580fn q1_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
7581    let mut acc = 0f32;
7582    for gi in 0..gpr {
7583        let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
7584        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
7585        let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
7586        let mut ga = 0f32;
7587        for (j, &b) in tile[2..].iter().enumerate() {
7588            for k in 0..8 {
7589                ga += (((b >> k) & 1) as f32 * 2.0 - 1.0) * xg[j * 8 + k];
7590            }
7591        }
7592        acc += ga * s;
7593    }
7594    acc
7595}
7596
7597/// One q1 row range via A8W8 (the body of `q1_matvec`'s hot loop,
7598/// extracted so multi-matrix jobs drive the same kernel).
7599#[allow(clippy::too_many_arguments)]
7600fn q1_range_a8w8(
7601    bytes: &[u8],
7602    gpr: usize,
7603    act: &SplitAct,
7604    gsum: &[i32],
7605    out: SendMut,
7606    start: usize,
7607    end: usize,
7608) {
7609    for r in start..end {
7610        let mut acc = dot_q1_row_i8(bytes, r, gpr, &act.xq, gsum) * act.sx;
7611        for &(j, xv) in &act.outliers {
7612            let (w, s) = q1_outlier(bytes, r, gpr, j);
7613            acc += w * s * xv;
7614        }
7615        // SAFETY: disjoint row ranges per worker.
7616        unsafe { *out.at(r) = acc };
7617    }
7618}
7619
7620/// Exact-scalar q1 row range (CMF_SDOT=0 contract).
7621fn q1_range_f32(bytes: &[u8], gpr: usize, x: &[f32], out: SendMut, start: usize, end: usize) {
7622    for r in start..end {
7623        // SAFETY: disjoint row ranges per worker.
7624        unsafe { *out.at(r) = q1_row_exact(bytes, r, gpr, x) };
7625    }
7626}
7627
7628/// q1t per-row overlay locator. After the base (`base_len`) come
7629/// `[u32 row_ptr[rows+1]]` then `[(u16 col, f16 val)]` grouped by row (row
7630/// `r`'s entries are `[row_ptr[r], row_ptr[r+1])`). Returns
7631/// `(row_ptr offset, entries offset, present)`.
7632fn q1t_overlay(bytes: &[u8], base_len: usize, rows: usize) -> (usize, usize, bool) {
7633    let entries = base_len + (rows + 1) * 4;
7634    (base_len, entries, entries <= bytes.len())
7635}
7636
7637/// Read `row_ptr[r]` from the overlay's prefix-sum table.
7638#[inline]
7639fn q1t_rowptr(bytes: &[u8], rp_off: usize, r: usize) -> usize {
7640    let o = rp_off + r * 4;
7641    u32::from_le_bytes([bytes[o], bytes[o + 1], bytes[o + 2], bytes[o + 3]]) as usize
7642}
7643
7644/// Byte → the 5 ternary signs it packs `{−1,0,+1}` as f32, precomputed so
7645/// decoding a q1t code is a table load, not the base-3 divide/modulo per
7646/// weight (division is ~20–40× the cost of a load). Built at compile time.
7647const SIGN5: [[f32; 5]; 256] = {
7648    let mut lut = [[0.0f32; 5]; 256];
7649    let pow3 = [1u16, 3, 9, 27, 81];
7650    let mut byte = 0usize;
7651    while byte < 256 {
7652        let mut i = 0usize;
7653        while i < 5 {
7654            let code = (byte as u16 / pow3[i]) % 3;
7655            lut[byte][i] = if code == 1 {
7656                1.0
7657            } else if code == 2 {
7658                -1.0
7659            } else {
7660                0.0
7661            };
7662            i += 1;
7663        }
7664        byte += 1;
7665    }
7666    lut
7667};
7668
7669/// Same table, as i8 signs — the operand for the int8 SDOT base kernel.
7670const SIGN5_I8: [[i8; 5]; 256] = {
7671    let mut lut = [[0i8; 5]; 256];
7672    let pow3 = [1u16, 3, 9, 27, 81];
7673    let mut byte = 0usize;
7674    while byte < 256 {
7675        let mut i = 0usize;
7676        while i < 5 {
7677            let code = (byte as u16 / pow3[i]) % 3;
7678            lut[byte][i] = if code == 1 {
7679                1
7680            } else if code == 2 {
7681                -1
7682            } else {
7683                0
7684            };
7685            i += 1;
7686        }
7687        byte += 1;
7688    }
7689    lut
7690};
7691
7692/// The same 5 i8 signs packed into a u64 (`[s0 s1 s2 s3 s4 0 0 0]`, LE) so the
7693/// group unpack is 7 unaligned u64 stores at offsets 0,5,10,…,30 instead of
7694/// six 5-byte copies + LUT indexing — each store's trailing zeros are fixed by
7695/// the next store, and the last one runs 6 B past the 32nd weight (the unpack
7696/// buffer is padded to 40). This is the decode/prefill hot inner op.
7697const SIGN5_U64: [u64; 256] = {
7698    let mut lut = [0u64; 256];
7699    let pow3 = [1u16, 3, 9, 27, 81];
7700    let mut byte = 0usize;
7701    while byte < 256 {
7702        let mut v = 0u64;
7703        let mut i = 0usize;
7704        while i < 5 {
7705            let code = (byte as u16 / pow3[i]) % 3;
7706            let s: u8 = if code == 1 {
7707                1
7708            } else if code == 2 {
7709                0xFF
7710            } else {
7711                0
7712            };
7713            v |= (s as u64) << (i * 8);
7714            i += 1;
7715        }
7716        lut[byte] = v;
7717        byte += 1;
7718    }
7719    lut
7720};
7721
7722/// Ternary base weight at `(row r, col j)` = `sign(code)·s_group`. Used to add
7723/// back activation-outlier columns, whose `x` was zeroed for the int8 bulk dot
7724/// (`split_act`). At a weight-outlier position the code is 0, so this is 0 and
7725/// the overlay correction owns that column — no double counting.
7726#[inline]
7727fn q1t_base_weight(bytes: &[u8], r: usize, gpr: usize, j: usize) -> f32 {
7728    const TILE: usize = cortiq_core::quant::Q1T_TILE;
7729    let off = (r * gpr + j / GROUP_SIZE) * TILE;
7730    let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
7731    let within = j % GROUP_SIZE;
7732    SIGN5[bytes[off + 2 + within / 5] as usize][within % 5] * s
7733}
7734
7735/// One 32-group int8 dot via two SDOTs. Bit-exact vs the scalar i8 sum
7736/// (integer accumulation is order-independent).
7737#[cfg(target_arch = "aarch64")]
7738#[target_feature(enable = "neon,dotprod")]
7739#[inline]
7740unsafe fn sdot32_i8(w: *const i8, x: *const i8) -> i32 {
7741    // SAFETY: caller guarantees 32 readable i8 at each pointer.
7742    unsafe {
7743        use core::arch::aarch64::*;
7744        use core::arch::asm;
7745        let w0 = vld1q_s8(w);
7746        let w1 = vld1q_s8(w.add(16));
7747        let x0 = vld1q_s8(x);
7748        let x1 = vld1q_s8(x.add(16));
7749        let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
7750        asm!(
7751            "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
7752            "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
7753            a0 = inout(vreg) a0, a1 = inout(vreg) a1,
7754            w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
7755            options(pure, nomem, nostack),
7756        );
7757        vaddvq_s32(vaddq_s32(a0, a1))
7758    }
7759}
7760
7761/// One 32-group int8 dot via AVX2: signed·signed as `maddubs(|w|, sign(x,w))`
7762/// then `madd` and a horizontal reduce (the same idiom as `dot_q4t_row_avx2`).
7763#[cfg(target_arch = "x86_64")]
7764#[target_feature(enable = "avx2")]
7765#[inline]
7766unsafe fn i8dot32_avx2(w: *const i8, x: *const i8) -> i32 {
7767    // SAFETY: caller guarantees 32 readable i8 at each pointer.
7768    unsafe {
7769        use core::arch::x86_64::*;
7770        let wv = _mm256_loadu_si256(w as *const __m256i);
7771        let xv = _mm256_loadu_si256(x as *const __m256i);
7772        let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
7773        let d = _mm256_madd_epi16(p16, _mm256_set1_epi16(1));
7774        let hi128 = _mm256_extracti128_si256::<1>(d);
7775        let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
7776        let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7777        let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7778        _mm_cvtsi128_si32(s32)
7779    }
7780}
7781
7782/// Unpack one q1t group's base-3 codes into 32 i8 signs via 7 unaligned u64
7783/// stores (see `SIGN5_U64`). `dst` MUST have ≥ 40 bytes: the 7th store writes
7784/// `dst[30..38]`. Stores go in order so each one's trailing zeros are
7785/// overwritten by the next; the final 6 padding bytes are unused by the dot.
7786#[inline]
7787fn q1t_unpack_group_i8(codes: *const u8, dst: &mut [i8]) {
7788    debug_assert!(dst.len() >= 40);
7789    // SAFETY: codes points at 7 readable bytes; dst has ≥ 40 bytes so every
7790    // 8-byte store at offset bi*5 (bi ≤ 6 → ≤ 30) stays in bounds.
7791    unsafe {
7792        let p = dst.as_mut_ptr();
7793        for bi in 0..7 {
7794            core::ptr::write_unaligned(
7795                p.add(bi * 5) as *mut u64,
7796                SIGN5_U64[*codes.add(bi) as usize],
7797            );
7798        }
7799    }
7800}
7801
7802/// One 32-group int8 dot, arch-dispatched (the matmat inner loop, where the
7803/// row's signs are unpacked once and dotted against every batch input).
7804/// Callers are gated by `a8w8_enabled()`, so the target-feature arms are
7805/// reachable; the scalar arm is a non-SIMD-arch fallback.
7806#[inline]
7807fn q1t_i8dot32(w: *const i8, x: *const i8) -> i32 {
7808    #[cfg(target_arch = "aarch64")]
7809    unsafe {
7810        return sdot32_i8(w, x);
7811    }
7812    #[cfg(target_arch = "x86_64")]
7813    unsafe {
7814        return i8dot32_avx2(w, x);
7815    }
7816    #[allow(unreachable_code)]
7817    unsafe {
7818        let mut s = 0i32;
7819        for k in 0..GROUP_SIZE {
7820            s += *w.add(k) as i32 * *x.add(k) as i32;
7821        }
7822        s
7823    }
7824}
7825
7826#[inline]
7827unsafe fn q1t_unpack_reg_u64s(codes: *const u8) -> (u64, u64, u64, u64) {
7828    let (s0, s1, s2, s3, s4, s5, s6) = unsafe {
7829        (
7830            SIGN5_U64[*codes as usize],
7831            SIGN5_U64[*codes.add(1) as usize],
7832            SIGN5_U64[*codes.add(2) as usize],
7833            SIGN5_U64[*codes.add(3) as usize],
7834            SIGN5_U64[*codes.add(4) as usize],
7835            SIGN5_U64[*codes.add(5) as usize],
7836            SIGN5_U64[*codes.add(6) as usize],
7837        )
7838    };
7839
7840    let u0 = s0 | (s1 << 40);
7841    let u1 = (s1 >> 24) | (s2 << 16) | (s3 << 56);
7842    let u2 = (s3 >> 8) | (s4 << 32);
7843    let u3 = (s4 >> 32) | (s5 << 8) | (s6 << 48);
7844
7845    (u0, u1, u2, u3)
7846}
7847
7848/// One q1t row's int8 base dot: `Σ_group s·dot(signs, xq)` (before the shared
7849/// `sx`). Direct register unpacking (zero stack stores/loads, no STLF stalls).
7850/// ARM SDOT.
7851#[cfg(target_arch = "aarch64")]
7852#[target_feature(enable = "neon,dotprod")]
7853unsafe fn q1t_dot_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
7854    use core::arch::aarch64::*;
7855    use core::arch::asm;
7856    unsafe {
7857        const TILE: usize = cortiq_core::quant::Q1T_TILE;
7858        let mut acc = 0f32;
7859        let bytes_ptr = bytes.as_ptr();
7860        let xq_ptr = xq.as_ptr();
7861        let row_off = r * gpr * TILE;
7862
7863        let gpr2 = gpr & !1;
7864        let mut gi = 0;
7865        while gi < gpr2 {
7866            let off0 = row_off + gi * TILE;
7867            let off1 = off0 + TILE;
7868            let s0 = f16_to_f32(u16::from_le_bytes([
7869                *bytes_ptr.add(off0),
7870                *bytes_ptr.add(off0 + 1),
7871            ]));
7872            let s1 = f16_to_f32(u16::from_le_bytes([
7873                *bytes_ptr.add(off1),
7874                *bytes_ptr.add(off1 + 1),
7875            ]));
7876
7877            let (u0_0, u1_0, u2_0, u3_0) = q1t_unpack_reg_u64s(bytes_ptr.add(off0 + 2));
7878            let (u0_1, u1_1, u2_1, u3_1) = q1t_unpack_reg_u64s(bytes_ptr.add(off1 + 2));
7879
7880            let w0_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_0), vcreate_u64(u1_0)));
7881            let w1_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_0), vcreate_u64(u3_0)));
7882            let w0_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_1), vcreate_u64(u1_1)));
7883            let w1_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_1), vcreate_u64(u3_1)));
7884
7885            let x0_0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE));
7886            let x1_0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE + 16));
7887            let x0_1 = vld1q_s8(xq_ptr.add((gi + 1) * GROUP_SIZE));
7888            let x1_1 = vld1q_s8(xq_ptr.add((gi + 1) * GROUP_SIZE + 16));
7889
7890            let (mut a0_0, mut a1_0) = (vdupq_n_s32(0), vdupq_n_s32(0));
7891            let (mut a0_1, mut a1_1) = (vdupq_n_s32(0), vdupq_n_s32(0));
7892            asm!(
7893                "sdot {a0_0:v}.4s, {w0_0:v}.16b, {x0_0:v}.16b",
7894                "sdot {a1_0:v}.4s, {w1_0:v}.16b, {x1_0:v}.16b",
7895                "sdot {a0_1:v}.4s, {w0_1:v}.16b, {x0_1:v}.16b",
7896                "sdot {a1_1:v}.4s, {w1_1:v}.16b, {x1_1:v}.16b",
7897                a0_0 = inout(vreg) a0_0, a1_0 = inout(vreg) a1_0,
7898                a0_1 = inout(vreg) a0_1, a1_1 = inout(vreg) a1_1,
7899                w0_0 = in(vreg) w0_0, x0_0 = in(vreg) x0_0, w1_0 = in(vreg) w1_0, x1_0 = in(vreg) x1_0,
7900                w0_1 = in(vreg) w0_1, x0_1 = in(vreg) x0_1, w1_1 = in(vreg) w1_1, x1_1 = in(vreg) x1_1,
7901                options(pure, nomem, nostack),
7902            );
7903            let d0 = vaddvq_s32(vaddq_s32(a0_0, a1_0));
7904            let d1 = vaddvq_s32(vaddq_s32(a0_1, a1_1));
7905            acc += d0 as f32 * s0 + d1 as f32 * s1;
7906            gi += 2;
7907        }
7908
7909        if gi < gpr {
7910            let off = row_off + gi * TILE;
7911            let s = f16_to_f32(u16::from_le_bytes([
7912                *bytes_ptr.add(off),
7913                *bytes_ptr.add(off + 1),
7914            ]));
7915            let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
7916            let w0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0), vcreate_u64(u1)));
7917            let w1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2), vcreate_u64(u3)));
7918            let x0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE));
7919            let x1 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE + 16));
7920            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
7921            asm!(
7922                "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
7923                "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
7924                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
7925                w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
7926                options(pure, nomem, nostack),
7927            );
7928            let d = vaddvq_s32(vaddq_s32(a0, a1));
7929            acc += d as f32 * s;
7930        }
7931        acc
7932    }
7933}
7934
7935/// x86 AVX2 mirror of `q1t_dot_row_sdot` (maddubs int8 dot per group).
7936#[cfg(target_arch = "x86_64")]
7937#[target_feature(enable = "avx2")]
7938unsafe fn q1t_dot_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
7939    use core::arch::x86_64::*;
7940    unsafe {
7941        const TILE: usize = cortiq_core::quant::Q1T_TILE;
7942        let mut acc = 0f32;
7943        let bytes_ptr = bytes.as_ptr();
7944        let xq_ptr = xq.as_ptr();
7945        let row_off = r * gpr * TILE;
7946
7947        let ones = _mm256_set1_epi16(1);
7948        for gi in 0..gpr {
7949            let off = row_off + gi * TILE;
7950            let s = f16_to_f32(u16::from_le_bytes([
7951                *bytes_ptr.add(off),
7952                *bytes_ptr.add(off + 1),
7953            ]));
7954            let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
7955            let wv = _mm256_set_epi64x(u3 as i64, u2 as i64, u1 as i64, u0 as i64);
7956            let xv = _mm256_loadu_si256(xq_ptr.add(gi * GROUP_SIZE) as *const __m256i);
7957            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
7958            let d256 = _mm256_madd_epi16(p16, ones);
7959            let d128 = _mm_add_epi32(
7960                _mm256_castsi256_si128(d256),
7961                _mm256_extracti128_si256(d256, 1),
7962            );
7963            let d64 = _mm_add_epi32(d128, _mm_shuffle_epi32(d128, 0xee));
7964            let d32 = _mm_cvtsi128_si32(_mm_add_epi32(d64, _mm_shuffle_epi32(d64, 0x55)));
7965            acc += d32 as f32 * s;
7966        }
7967        acc
7968    }
7969}
7970
7971/// VNNI twin of `q1t_dot_row_avx2` (see `dpbusd_hsum`).
7972#[cfg(target_arch = "x86_64")]
7973#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
7974unsafe fn q1t_dot_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
7975    use core::arch::x86_64::*;
7976    // SAFETY: same tile/xq contracts as `q1t_dot_row_avx2`.
7977    unsafe {
7978        const TILE: usize = cortiq_core::quant::Q1T_TILE;
7979        let mut acc = 0f32;
7980        let bytes_ptr = bytes.as_ptr();
7981        let xq_ptr = xq.as_ptr();
7982        let row_off = r * gpr * TILE;
7983        for gi in 0..gpr {
7984            let off = row_off + gi * TILE;
7985            let s = f16_to_f32(u16::from_le_bytes([
7986                *bytes_ptr.add(off),
7987                *bytes_ptr.add(off + 1),
7988            ]));
7989            let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
7990            let wv = _mm256_set_epi64x(u3 as i64, u2 as i64, u1 as i64, u0 as i64);
7991            let xv = _mm256_loadu_si256(xq_ptr.add(gi * GROUP_SIZE) as *const __m256i);
7992            let d = dpbusd_hsum(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
7993            acc += d as f32 * s;
7994        }
7995        acc
7996    }
7997}
7998
7999/// Per-row int8 base dot, dispatched once per row (matvec decode hot path).
8000/// Callers are gated by `a8w8_enabled()`, so the target-feature kernels are
8001/// reachable.
8002#[inline]
8003fn q1t_dot_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
8004    #[cfg(target_arch = "aarch64")]
8005    unsafe {
8006        return q1t_dot_row_sdot(bytes, r, gpr, xq);
8007    }
8008    #[cfg(target_arch = "x86_64")]
8009    unsafe {
8010        if vnni_tiles_enabled() {
8011            return q1t_dot_row_vnni(bytes, r, gpr, xq);
8012        }
8013        return q1t_dot_row_avx2(bytes, r, gpr, xq);
8014    }
8015    #[allow(unreachable_code)]
8016    {
8017        const TILE: usize = cortiq_core::quant::Q1T_TILE;
8018        let mut acc = 0f32;
8019        let mut sg = [0i8; GROUP_SIZE + 8]; // +8 slack for the u64-store unpack
8020        for gi in 0..gpr {
8021            let off = (r * gpr + gi) * TILE;
8022            let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
8023            q1t_unpack_group_i8(bytes.as_ptr().wrapping_add(off + 2), &mut sg);
8024            let mut d = 0i32;
8025            for k in 0..GROUP_SIZE {
8026                d += sg[k] as i32 * xq[gi * GROUP_SIZE + k] as i32;
8027            }
8028            acc += d as f32 * s;
8029        }
8030        acc
8031    }
8032}
8033
8034/// Σ over a row's outliers of `value·x[col]` — the correction that adds the
8035/// overlay's exact weights on top of the base dot. INVARIANT: the encoder
8036/// writes ternary code 0 at every outlier position (`quantize_q1t`), so the
8037/// base contributes nothing there and this is a plain `value·x`, not
8038/// `(value − base)·x` — no scattered per-outlier scale read. Row `r`'s entries
8039/// are the contiguous slice `[row_ptr[r], row_ptr[r+1])`, so no binary search.
8040fn q1t_row_outlier_correction(
8041    bytes: &[u8],
8042    r: usize,
8043    rp_off: usize,
8044    entries_off: usize,
8045    has_ov: bool,
8046    x: &[f32],
8047) -> f32 {
8048    if !has_ov {
8049        return 0.0;
8050    }
8051    let (c0, c1) = (
8052        q1t_rowptr(bytes, rp_off, r),
8053        q1t_rowptr(bytes, rp_off, r + 1),
8054    );
8055    let mut corr = 0f32;
8056    for p in c0..c1 {
8057        let e = entries_off + p * 4;
8058        let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
8059        let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
8060        corr += val * x[col];
8061    }
8062    corr
8063}
8064
8065/// Dequantize one q1t row into `buf[..cols]` via the sign LUT (no division),
8066/// then apply the row's outliers (its `[row_ptr[r], row_ptr[r+1])` slice).
8067/// Used by the batched (prefill) path where the decode amortizes over the batch.
8068fn q1t_dequant_row(
8069    bytes: &[u8],
8070    r: usize,
8071    gpr: usize,
8072    rp_off: usize,
8073    entries_off: usize,
8074    has_ov: bool,
8075    buf: &mut [f32],
8076) {
8077    const TILE: usize = cortiq_core::quant::Q1T_TILE;
8078    for g in 0..gpr {
8079        let off = (r * gpr + g) * TILE;
8080        let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
8081        let codes = &bytes[off + 2..off + TILE];
8082        let bc = g * GROUP_SIZE;
8083        // 6 full bytes (30 codes) + a 7th byte holding the last 2.
8084        for bi in 0..6 {
8085            let lut = &SIGN5[codes[bi] as usize];
8086            let d = &mut buf[bc + bi * 5..bc + bi * 5 + 5];
8087            for i in 0..5 {
8088                d[i] = lut[i] * s;
8089            }
8090        }
8091        let lut = &SIGN5[codes[6] as usize];
8092        buf[bc + 30] = lut[0] * s;
8093        buf[bc + 31] = lut[1] * s;
8094    }
8095    if !has_ov {
8096        return;
8097    }
8098    let (c0, c1) = (
8099        q1t_rowptr(bytes, rp_off, r),
8100        q1t_rowptr(bytes, rp_off, r + 1),
8101    );
8102    for p in c0..c1 {
8103        let e = entries_off + p * 4;
8104        let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
8105        buf[col] = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
8106    }
8107}
8108
8109/// Add the sparse outlier overlay onto a base dot already in `out` (the GPU
8110/// computes the ternary base; the overlay stays on the CPU — its entries are
8111/// few and its per-row gather doesn't vectorize on the GPU). Row-parallel.
8112fn q1t_add_overlay(
8113    bytes: &[u8],
8114    x: &[f32],
8115    rows: usize,
8116    cols: usize,
8117    out: &mut [f32],
8118    pool: Option<&Pool>,
8119) {
8120    const TILE: usize = cortiq_core::quant::Q1T_TILE;
8121    let gpr = cols / GROUP_SIZE;
8122    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
8123    if !has_ov {
8124        return;
8125    }
8126    let out_addr = SendMut(out.as_mut_ptr());
8127    let run = move |start: usize, end: usize| {
8128        for r in start..end {
8129            let corr = q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
8130            // SAFETY: disjoint rows; add onto the base the GPU already wrote.
8131            unsafe { *out_addr.at(r) += corr };
8132        }
8133    };
8134    dispatch_rows(pool, rows, &run);
8135}
8136
8137/// Q1T row range via the A8W8 int8 path — shared activation split,
8138/// per-row: base SDOT dot + outlier correction + overlay.
8139#[allow(clippy::too_many_arguments)]
8140fn q1t_range_a8w8(
8141    bytes: &[u8],
8142    gpr: usize,
8143    rp_off: usize,
8144    ent_off: usize,
8145    has_ov: bool,
8146    act: &SplitAct,
8147    x: &[f32],
8148    out: SendMut,
8149    start: usize,
8150    end: usize,
8151) {
8152    for r in start..end {
8153        let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
8154        for &(j, xv) in &act.outliers {
8155            acc += q1t_base_weight(bytes, r, gpr, j) * xv;
8156        }
8157        acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
8158        // SAFETY: disjoint row ranges per worker.
8159        unsafe { *out.at(r) = acc };
8160    }
8161}
8162
8163/// Q1T row range via the f32 path (no SDOT) — for matvec_many batched
8164/// dispatch when a8w8 is unavailable.
8165#[allow(clippy::too_many_arguments)]
8166fn q1t_range_f32_batch(
8167    bytes: &[u8],
8168    gpr: usize,
8169    rp_off: usize,
8170    ent_off: usize,
8171    has_ov: bool,
8172    x: &[f32],
8173    out: SendMut,
8174    start: usize,
8175    end: usize,
8176) {
8177    const TILE: usize = cortiq_core::quant::Q1T_TILE;
8178    let mut sg = [0f32; GROUP_SIZE];
8179    for r in start..end {
8180        let mut acc = 0f32;
8181        for g in 0..gpr {
8182            let off = (r * gpr + g) * TILE;
8183            let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
8184            let codes = &bytes[off + 2..off + TILE];
8185            let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
8186            for bi in 0..6 {
8187                sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
8188            }
8189            let lut = &SIGN5[codes[6] as usize];
8190            sg[30] = lut[0];
8191            sg[31] = lut[1];
8192            let mut gsum = 0f32;
8193            for k in 0..GROUP_SIZE {
8194                gsum += sg[k] * xg[k];
8195            }
8196            acc += s * gsum;
8197        }
8198        acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
8199        // SAFETY: disjoint row ranges per worker.
8200        unsafe { *out.at(r) = acc };
8201    }
8202}
8203
8204/// Ternary (q1t) matvec — decode+dot straight from mmap, one group at a time:
8205/// no per-ROW buffer, no division (the sign LUT), and a tiny per-group sign
8206/// buffer so the 32-wide dot vectorizes. This is the decode hot path.
8207fn q1t_matvec(
8208    bytes: &[u8],
8209    x: &[f32],
8210    rows: usize,
8211    cols: usize,
8212    out: &mut [f32],
8213    pool: Option<&Pool>,
8214) {
8215    debug_assert_eq!(out.len(), rows);
8216    const TILE: usize = cortiq_core::quant::Q1T_TILE;
8217    let gpr = cols / GROUP_SIZE;
8218    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
8219    let out_addr = SendMut(out.as_mut_ptr());
8220    // int8 SDOT base dot (ARM dotprod): ~4× the f32 arithmetic. x → i8 once
8221    // (`split_act`), activation outliers added back exactly in f32, weight
8222    // overlay on top. ARM SDOT / x86 AVX2; CMF_SDOT=0 keeps the exact f32 path.
8223    if a8w8_enabled() {
8224        let act = split_act(x);
8225        let act = &act;
8226        let run = move |start: usize, end: usize| {
8227            for r in start..end {
8228                let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
8229                for &(j, xv) in &act.outliers {
8230                    acc += q1t_base_weight(bytes, r, gpr, j) * xv;
8231                }
8232                acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
8233                // SAFETY: disjoint row ranges per worker.
8234                unsafe { *out_addr.at(r) = acc };
8235            }
8236        };
8237        dispatch_rows(pool, rows, &run);
8238        return;
8239    }
8240    let run = move |start: usize, end: usize| {
8241        // Per-group signs, unpacked contiguously so the dot below is a clean
8242        // 32-wide reduction the autovectorizer turns into f32x4 FMAs — the
8243        // 5-values-per-byte base-3 layout won't SIMD in place.
8244        let mut sg = [0f32; GROUP_SIZE];
8245        for r in start..end {
8246            let mut acc = 0f32;
8247            for g in 0..gpr {
8248                let off = (r * gpr + g) * TILE;
8249                let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
8250                let codes = &bytes[off + 2..off + TILE];
8251                let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
8252                for bi in 0..6 {
8253                    sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
8254                }
8255                let lut = &SIGN5[codes[6] as usize];
8256                sg[30] = lut[0];
8257                sg[31] = lut[1];
8258                let mut gsum = 0f32;
8259                for k in 0..GROUP_SIZE {
8260                    gsum += sg[k] * xg[k];
8261                }
8262                acc += s * gsum;
8263            }
8264            acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
8265            unsafe { *out_addr.at(r) = acc };
8266        }
8267    };
8268    dispatch_rows(pool, rows, &run);
8269}
8270
8271/// Fused-pair twin of `q1t_dot_row_sdot`: ONE register unpack of the
8272/// ternary codes serves BOTH activation streams (the unpack chain is
8273/// the dominant per-row cost — MTP verify pairs paid it twice). Per
8274/// stream the group order and f32 accumulation match the single-row
8275/// kernel exactly, so pair == 2×matvec bit-for-bit.
8276#[cfg(target_arch = "aarch64")]
8277#[target_feature(enable = "neon,dotprod")]
8278unsafe fn q1t_dot_row_sdot2(bytes: &[u8], r: usize, gpr: usize, xa: &[i8], xb: &[i8]) -> [f32; 2] {
8279    use core::arch::aarch64::*;
8280    use core::arch::asm;
8281    // SAFETY: same slice-length contracts as `q1t_dot_row_sdot`, ×2.
8282    unsafe {
8283        const TILE: usize = cortiq_core::quant::Q1T_TILE;
8284        let bytes_ptr = bytes.as_ptr();
8285        let row_off = r * gpr * TILE;
8286        let xp = [xa.as_ptr(), xb.as_ptr()];
8287        let mut acc = [0f32; 2];
8288        macro_rules! sdot2 {
8289            ($w0:expr, $w1:expr, $x:expr) => {{
8290                let x0 = vld1q_s8($x);
8291                let x1 = vld1q_s8($x.add(16));
8292                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
8293                asm!(
8294                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
8295                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
8296                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
8297                    w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
8298                    options(pure, nomem, nostack),
8299                );
8300                vaddvq_s32(vaddq_s32(a0, a1))
8301            }};
8302        }
8303        let gpr2 = gpr & !1;
8304        let mut gi = 0;
8305        while gi < gpr2 {
8306            let off0 = row_off + gi * TILE;
8307            let off1 = off0 + TILE;
8308            let s0 = f16_to_f32(u16::from_le_bytes([
8309                *bytes_ptr.add(off0),
8310                *bytes_ptr.add(off0 + 1),
8311            ]));
8312            let s1 = f16_to_f32(u16::from_le_bytes([
8313                *bytes_ptr.add(off1),
8314                *bytes_ptr.add(off1 + 1),
8315            ]));
8316            let (u0_0, u1_0, u2_0, u3_0) = q1t_unpack_reg_u64s(bytes_ptr.add(off0 + 2));
8317            let (u0_1, u1_1, u2_1, u3_1) = q1t_unpack_reg_u64s(bytes_ptr.add(off1 + 2));
8318            let w0_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_0), vcreate_u64(u1_0)));
8319            let w1_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_0), vcreate_u64(u3_0)));
8320            let w0_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_1), vcreate_u64(u1_1)));
8321            let w1_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_1), vcreate_u64(u3_1)));
8322            for k in 0..2 {
8323                let d0 = sdot2!(w0_0, w1_0, xp[k].add(gi * GROUP_SIZE));
8324                let d1 = sdot2!(w0_1, w1_1, xp[k].add((gi + 1) * GROUP_SIZE));
8325                acc[k] += d0 as f32 * s0 + d1 as f32 * s1;
8326            }
8327            gi += 2;
8328        }
8329        if gi < gpr {
8330            let off = row_off + gi * TILE;
8331            let s = f16_to_f32(u16::from_le_bytes([
8332                *bytes_ptr.add(off),
8333                *bytes_ptr.add(off + 1),
8334            ]));
8335            let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
8336            let w0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0), vcreate_u64(u1)));
8337            let w1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2), vcreate_u64(u3)));
8338            for k in 0..2 {
8339                let d = sdot2!(w0, w1, xp[k].add(gi * GROUP_SIZE));
8340                acc[k] += d as f32 * s;
8341            }
8342        }
8343        acc
8344    }
8345}
8346
8347/// Fused Q1T pair matvec: ONE pass over the rows serves both
8348/// activation streams — on ARM the ternary register unpack happens
8349/// once per tile pair (`q1t_dot_row_sdot2`); elsewhere the second dot
8350/// rides the row's L1-warm tile bytes. Per stream the math matches
8351/// `q1t_matvec` exactly.
8352fn q1t_matvec2(
8353    bytes: &[u8],
8354    x1: &[f32],
8355    x2: &[f32],
8356    rows: usize,
8357    cols: usize,
8358    o1: &mut [f32],
8359    o2: &mut [f32],
8360    pool: Option<&Pool>,
8361) {
8362    debug_assert_eq!(o1.len(), rows);
8363    debug_assert_eq!(o2.len(), rows);
8364    const TILE: usize = cortiq_core::quant::Q1T_TILE;
8365    let gpr = cols / GROUP_SIZE;
8366    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
8367    let out1 = SendMut(o1.as_mut_ptr());
8368    let out2 = SendMut(o2.as_mut_ptr());
8369    if a8w8_enabled() {
8370        let a1 = split_act(x1);
8371        let a2 = split_act(x2);
8372        let (a1, a2) = (&a1, &a2);
8373        let run = move |start: usize, end: usize| {
8374            for r in start..end {
8375                #[cfg(target_arch = "aarch64")]
8376                // a8w8 on aarch64 ⇔ sdot_enabled(), so the kernel's
8377                // target features are present.
8378                let ds = unsafe { q1t_dot_row_sdot2(bytes, r, gpr, &a1.xq, &a2.xq) };
8379                #[cfg(not(target_arch = "aarch64"))]
8380                let ds = [
8381                    q1t_dot_row_i8(bytes, r, gpr, &a1.xq),
8382                    q1t_dot_row_i8(bytes, r, gpr, &a2.xq),
8383                ];
8384                let mut acc1 = ds[0] * a1.sx;
8385                for &(j, xv) in &a1.outliers {
8386                    acc1 += q1t_base_weight(bytes, r, gpr, j) * xv;
8387                }
8388                acc1 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x1);
8389                let mut acc2 = ds[1] * a2.sx;
8390                for &(j, xv) in &a2.outliers {
8391                    acc2 += q1t_base_weight(bytes, r, gpr, j) * xv;
8392                }
8393                acc2 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x2);
8394                // SAFETY: disjoint row ranges per worker.
8395                unsafe {
8396                    *out1.at(r) = acc1;
8397                    *out2.at(r) = acc2;
8398                }
8399            }
8400        };
8401        dispatch_rows(pool, rows, &run);
8402        return;
8403    }
8404    let run = move |start: usize, end: usize| {
8405        // Exact path (CMF_SDOT=0): unpack the sign LUT once per group,
8406        // dot both streams — same op order per stream as `q1t_matvec`.
8407        let mut sg = [0f32; GROUP_SIZE];
8408        for r in start..end {
8409            let mut acc1 = 0f32;
8410            let mut acc2 = 0f32;
8411            for g in 0..gpr {
8412                let off = (r * gpr + g) * TILE;
8413                let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
8414                let codes = &bytes[off + 2..off + TILE];
8415                for bi in 0..6 {
8416                    sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
8417                }
8418                let lut = &SIGN5[codes[6] as usize];
8419                sg[30] = lut[0];
8420                sg[31] = lut[1];
8421                let xg1 = &x1[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
8422                let xg2 = &x2[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
8423                let mut gsum1 = 0f32;
8424                for k in 0..GROUP_SIZE {
8425                    gsum1 += sg[k] * xg1[k];
8426                }
8427                acc1 += s * gsum1;
8428                let mut gsum2 = 0f32;
8429                for k in 0..GROUP_SIZE {
8430                    gsum2 += sg[k] * xg2[k];
8431                }
8432                acc2 += s * gsum2;
8433            }
8434            acc1 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x1);
8435            acc2 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x2);
8436            // SAFETY: disjoint row ranges per worker.
8437            unsafe {
8438                *out1.at(r) = acc1;
8439                *out2.at(r) = acc2;
8440            }
8441        }
8442    };
8443    dispatch_rows(pool, rows, &run);
8444}
8445
8446/// Ternary (q1t) matmat (prefill) — dequant each row once, dot the whole
8447/// batch against it (amortizes the per-row decode).
8448fn q1t_matmat(
8449    bytes: &[u8],
8450    xs: &[f32],
8451    b: usize,
8452    rows: usize,
8453    cols: usize,
8454    out: &mut [f32],
8455    pool: Option<&Pool>,
8456) {
8457    debug_assert_eq!(out.len(), b * rows);
8458    const TILE: usize = cortiq_core::quant::Q1T_TILE;
8459    let gpr = cols / GROUP_SIZE;
8460    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
8461    let out_addr = SendMut(out.as_mut_ptr());
8462    // int8 prefill (ARM SDOT / x86 AVX2): quantize the B inputs once, unpack
8463    // each weight row's signs to i8 ONCE, then int8-dot against every input —
8464    // the row sign-decode amortizes over the whole batch. CMF_SDOT=0 → f32.
8465    if a8w8_enabled() {
8466        let acts: Vec<SplitAct> = (0..b)
8467            .map(|bi| split_act(&xs[bi * cols..(bi + 1) * cols]))
8468            .collect();
8469        let acts = &acts;
8470        let run = move |start: usize, end: usize| {
8471            let mut sg = vec![0i8; cols + 8]; // row signs, i8 (+8 unpack slack)
8472            let mut sc = vec![0f32; gpr]; // per-group scales
8473            let mut accs = vec![0f32; b]; // per-batch accumulators, reused per row
8474            for r in start..end {
8475                for g in 0..gpr {
8476                    let off = (r * gpr + g) * TILE;
8477                    sc[g] = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
8478                    q1t_unpack_group_i8(
8479                        bytes.as_ptr().wrapping_add(off + 2),
8480                        &mut sg[g * GROUP_SIZE..],
8481                    );
8482                }
8483                for bi in 0..b {
8484                    let act = &acts[bi];
8485                    let mut isum = 0f32;
8486                    for g in 0..gpr {
8487                        let d = q1t_i8dot32(
8488                            sg.as_ptr().wrapping_add(g * GROUP_SIZE),
8489                            act.xq.as_ptr().wrapping_add(g * GROUP_SIZE),
8490                        );
8491                        isum += d as f32 * sc[g];
8492                    }
8493                    let mut acc = isum * act.sx;
8494                    for &(j, xv) in &act.outliers {
8495                        acc += q1t_base_weight(bytes, r, gpr, j) * xv;
8496                    }
8497                    accs[bi] = acc;
8498                }
8499                // Overlay ONCE per row for the whole batch: read each (col, val)
8500                // from mmap a single time (was b× — the re-read dominated prefill)
8501                // and fan it out over the batch via the cached inputs.
8502                if has_ov {
8503                    let (c0, c1) = (
8504                        q1t_rowptr(bytes, rp_off, r),
8505                        q1t_rowptr(bytes, rp_off, r + 1),
8506                    );
8507                    for p in c0..c1 {
8508                        let e = ent_off + p * 4;
8509                        let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
8510                        let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
8511                        for bi in 0..b {
8512                            accs[bi] += val * xs[bi * cols + col];
8513                        }
8514                    }
8515                }
8516                for bi in 0..b {
8517                    unsafe { *out_addr.at(bi * rows + r) = accs[bi] };
8518                }
8519            }
8520        };
8521        dispatch_rows(pool, rows, &run);
8522        return;
8523    }
8524    let run = move |start: usize, end: usize| {
8525        let mut buf = vec![0f32; cols];
8526        for r in start..end {
8527            q1t_dequant_row(bytes, r, gpr, rp_off, ent_off, has_ov, &mut buf);
8528            for bi in 0..b {
8529                let xr = &xs[bi * cols..(bi + 1) * cols];
8530                let mut acc = 0f32;
8531                for j in 0..cols {
8532                    acc += buf[j] * xr[j];
8533                }
8534                unsafe { *out_addr.at(bi * rows + r) = acc };
8535            }
8536        }
8537    };
8538    dispatch_rows(pool, rows, &run);
8539}
8540
8541fn q1_matvec(
8542    bytes: &[u8],
8543    x: &[f32],
8544    rows: usize,
8545    cols: usize,
8546    out: &mut [f32],
8547    pool: Option<&Pool>,
8548) {
8549    debug_assert_eq!(out.len(), rows);
8550    let gpr = cols / GROUP_SIZE;
8551    let out_addr = SendMut(out.as_mut_ptr());
8552    if a8w8_enabled() {
8553        let act = split_act(x);
8554        let gsum = q1_group_sums(&act.xq, gpr);
8555        let (act, gsum) = (&act, &gsum);
8556        let run = move |start: usize, end: usize| {
8557            q1_range_a8w8(bytes, gpr, act, gsum, out_addr, start, end)
8558        };
8559        dispatch_rows(pool, rows, &run);
8560        return;
8561    }
8562    let run = move |start: usize, end: usize| q1_range_f32(bytes, gpr, x, out_addr, start, end);
8563    dispatch_rows(pool, rows, &run);
8564}
8565
8566/// Fused two-input q1 matvec (weights read once per pair).
8567#[allow(clippy::too_many_arguments)]
8568fn q1_matvec2(
8569    bytes: &[u8],
8570    x1: &[f32],
8571    x2: &[f32],
8572    rows: usize,
8573    cols: usize,
8574    o1: &mut [f32],
8575    o2: &mut [f32],
8576    pool: Option<&Pool>,
8577) {
8578    let gpr = cols / GROUP_SIZE;
8579    let p1 = SendMut(o1.as_mut_ptr());
8580    let p2 = SendMut(o2.as_mut_ptr());
8581    if a8w8_enabled() {
8582        let a1 = split_act(x1);
8583        let a2 = split_act(x2);
8584        let g1 = q1_group_sums(&a1.xq, gpr);
8585        let g2 = q1_group_sums(&a2.xq, gpr);
8586        let (a1, a2, g1, g2) = (&a1, &a2, &g1, &g2);
8587        let run = move |start: usize, end: usize| {
8588            for r in start..end {
8589                let mut v1 = dot_q1_row_i8(bytes, r, gpr, &a1.xq, g1) * a1.sx;
8590                let mut v2 = dot_q1_row_i8(bytes, r, gpr, &a2.xq, g2) * a2.sx;
8591                for &(j, xv) in &a1.outliers {
8592                    let (w, s) = q1_outlier(bytes, r, gpr, j);
8593                    v1 += w * s * xv;
8594                }
8595                for &(j, xv) in &a2.outliers {
8596                    let (w, s) = q1_outlier(bytes, r, gpr, j);
8597                    v2 += w * s * xv;
8598                }
8599                // SAFETY: disjoint row ranges per worker.
8600                unsafe {
8601                    *p1.at(r) = v1;
8602                    *p2.at(r) = v2;
8603                }
8604            }
8605        };
8606        dispatch_rows(pool, rows, &run);
8607        return;
8608    }
8609    let run = move |start: usize, end: usize| {
8610        for r in start..end {
8611            // SAFETY: disjoint row ranges per worker.
8612            unsafe {
8613                *p1.at(r) = q1_row_exact(bytes, r, gpr, x1);
8614                *p2.at(r) = q1_row_exact(bytes, r, gpr, x2);
8615            }
8616        }
8617    };
8618    dispatch_rows(pool, rows, &run);
8619}
8620
8621/// Batched q1 matmat: each row's tiles stream once per microbatch.
8622#[allow(clippy::too_many_arguments)]
8623fn q1_matmat(
8624    bytes: &[u8],
8625    xs_all: &[f32],
8626    b: usize,
8627    rows: usize,
8628    cols: usize,
8629    out: &mut [f32],
8630    pool: Option<&Pool>,
8631) {
8632    debug_assert_eq!(out.len(), b * rows);
8633    let gpr = cols / GROUP_SIZE;
8634    let out_addr = SendMut(out.as_mut_ptr());
8635    if a8w8_enabled() {
8636        let acts: Vec<(SplitAct, Vec<i32>)> = (0..b)
8637            .map(|bi| {
8638                let act = split_act(&xs_all[bi * cols..(bi + 1) * cols]);
8639                let gsum = q1_group_sums(&act.xq, gpr);
8640                (act, gsum)
8641            })
8642            .collect();
8643        let acts = &acts;
8644        #[cfg(target_arch = "x86_64")]
8645        let blocked_ok = avx2_enabled() && blocked_enabled();
8646        #[cfg(target_arch = "aarch64")]
8647        let blocked_ok = sdot_enabled() && blocked_enabled();
8648        let run = move |start: usize, end: usize| {
8649            for r in start..end {
8650                let mut bi = 0usize;
8651                // Blocked 1×4: the unpacked bit mask serves four
8652                // activation streams per group.
8653                #[cfg(target_arch = "aarch64")]
8654                if blocked_ok {
8655                    while bi + 4 <= acts.len() {
8656                        let xs = [
8657                            acts[bi].0.xq.as_slice(),
8658                            acts[bi + 1].0.xq.as_slice(),
8659                            acts[bi + 2].0.xq.as_slice(),
8660                            acts[bi + 3].0.xq.as_slice(),
8661                        ];
8662                        let gs = [
8663                            acts[bi].1.as_slice(),
8664                            acts[bi + 1].1.as_slice(),
8665                            acts[bi + 2].1.as_slice(),
8666                            acts[bi + 3].1.as_slice(),
8667                        ];
8668                        let d = unsafe { dot_q1_row_1x4_sdot(bytes, r, gpr, xs, gs) };
8669                        for k in 0..4 {
8670                            let (act, _) = &acts[bi + k];
8671                            let mut acc = d[k] * act.sx;
8672                            for &(j, xv) in &act.outliers {
8673                                let (w, sc) = q1_outlier(bytes, r, gpr, j);
8674                                acc += w * sc * xv;
8675                            }
8676                            // SAFETY: disjoint (bi, r) cells per worker.
8677                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
8678                        }
8679                        bi += 4;
8680                    }
8681                }
8682                #[cfg(target_arch = "x86_64")]
8683                if blocked_ok {
8684                    while bi + 4 <= acts.len() {
8685                        let xs = [
8686                            acts[bi].0.xq.as_slice(),
8687                            acts[bi + 1].0.xq.as_slice(),
8688                            acts[bi + 2].0.xq.as_slice(),
8689                            acts[bi + 3].0.xq.as_slice(),
8690                        ];
8691                        let gs = [
8692                            acts[bi].1.as_slice(),
8693                            acts[bi + 1].1.as_slice(),
8694                            acts[bi + 2].1.as_slice(),
8695                            acts[bi + 3].1.as_slice(),
8696                        ];
8697                        let d = unsafe {
8698                            if vnni_tiles_enabled() {
8699                                dot_q1_row_1x4_vnni(bytes, r, gpr, xs, gs)
8700                            } else {
8701                                dot_q1_row_1x4_avx2(bytes, r, gpr, xs, gs)
8702                            }
8703                        };
8704                        for k in 0..4 {
8705                            let (act, _) = &acts[bi + k];
8706                            let mut acc = d[k] * act.sx;
8707                            for &(j, xv) in &act.outliers {
8708                                let (w, sc) = q1_outlier(bytes, r, gpr, j);
8709                                acc += w * sc * xv;
8710                            }
8711                            // SAFETY: disjoint (bi, r) cells per worker.
8712                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
8713                        }
8714                        bi += 4;
8715                    }
8716                }
8717                while bi < acts.len() {
8718                    let (act, gsum) = &acts[bi];
8719                    let mut acc = dot_q1_row_i8(bytes, r, gpr, &act.xq, gsum) * act.sx;
8720                    for &(j, xv) in &act.outliers {
8721                        let (w, s) = q1_outlier(bytes, r, gpr, j);
8722                        acc += w * s * xv;
8723                    }
8724                    // SAFETY: disjoint (bi, r) cells per worker range.
8725                    unsafe { *out_addr.at(bi * rows + r) = acc };
8726                    bi += 1;
8727                }
8728            }
8729        };
8730        dispatch_rows(pool, rows, &run);
8731        return;
8732    }
8733    let run = move |start: usize, end: usize| {
8734        for r in start..end {
8735            for bi in 0..b {
8736                let x = &xs_all[bi * cols..(bi + 1) * cols];
8737                // SAFETY: disjoint (bi, r) cells per worker range.
8738                unsafe { *out_addr.at(bi * rows + r) = q1_row_exact(bytes, r, gpr, x) };
8739            }
8740        }
8741    };
8742    dispatch_rows(pool, rows, &run);
8743}
8744
8745/// Fused q4_block matvec straight from the mapped bytes. SDOT path when
8746/// dotprod is available (port of vmfcore `dot_q4_block_sdot`, measured
8747/// +23% on q4 decode): nibbles → centered i8, int8×int8 `sdot` per
8748/// 32-group, exact outlier correction — the same A8W8 contract as q8.
8749/// `CMF_SDOT=0` keeps the exact scalar path.
8750fn q4matvec(
8751    bytes: &[u8],
8752    x: &[f32],
8753    rows: usize,
8754    cols: usize,
8755    out: &mut [f32],
8756    pool: Option<&Pool>,
8757) {
8758    debug_assert_eq!(out.len(), rows);
8759    let (packed, scales) = q4_split(bytes, rows, cols);
8760    let gpr = cols / GROUP_SIZE;
8761    let out_addr = SendMut(out.as_mut_ptr());
8762
8763    if a8w8_enabled() {
8764        let act = split_act(x);
8765        let run = move |start: usize, end: usize| {
8766            q4_range_a8w8(packed, scales, gpr, cols, &act, out_addr, start, end)
8767        };
8768        dispatch_rows(pool, rows, &run);
8769        return;
8770    }
8771
8772    let run =
8773        move |start: usize, end: usize| q4_range_f32(packed, scales, gpr, x, out_addr, start, end);
8774    dispatch_rows(pool, rows, &run);
8775}
8776
8777/// One q4 row via the A8W8 int8 path — SDOT on ARM, AVX2 maddubs on
8778/// x86 (scalar fallback is unreachable: callers gate on a8w8_enabled).
8779#[inline]
8780#[allow(unreachable_code)]
8781/// One UNPACKED q4 row (centered i8 in `buf`) against four activation
8782/// streams: the 32-byte weight chunk and its abs() load once per group,
8783/// the per-group f16 scale decodes once — four maddubs+reduce chains
8784/// instead of four full (load, abs, dot) rounds.
8785#[cfg(target_arch = "x86_64")]
8786#[target_feature(enable = "avx2")]
8787unsafe fn dot_q4b_row_1x4_avx2(
8788    buf: &[u8],
8789    scales: &[u8],
8790    g0: usize,
8791    gpr: usize,
8792    xs: [&[i8]; 4],
8793) -> [f32; 4] {
8794    // SAFETY: callers uphold buffer contracts (buf.len() == gpr·32).
8795    unsafe {
8796        use core::arch::x86_64::*;
8797        let ones = _mm256_set1_epi16(1);
8798        let mut acc = [0f32; 4];
8799        for gi in 0..gpr {
8800            let s = f16_to_f32(u16::from_le_bytes([
8801                scales[(g0 + gi) * 2],
8802                scales[(g0 + gi) * 2 + 1],
8803            ]));
8804            let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8805            let aw = _mm256_abs_epi8(w);
8806            for (k, xq) in xs.iter().enumerate() {
8807                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8808                let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
8809                let d = _mm256_madd_epi16(p16, ones);
8810                let hi128 = _mm256_extracti128_si256::<1>(d);
8811                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
8812                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8813                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8814                acc[k] += _mm_cvtsi128_si32(s32) as f32 * s;
8815            }
8816        }
8817        acc
8818    }
8819}
8820
8821/// VNNI twin of `dot_q4b_row_1x4_avx2` (see `dpbusd_hsum`).
8822#[cfg(target_arch = "x86_64")]
8823#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
8824unsafe fn dot_q4b_row_1x4_vnni(
8825    buf: &[u8],
8826    scales: &[u8],
8827    g0: usize,
8828    gpr: usize,
8829    xs: [&[i8]; 4],
8830) -> [f32; 4] {
8831    // SAFETY: callers uphold buffer contracts (buf.len() == gpr·32).
8832    unsafe {
8833        use core::arch::x86_64::*;
8834        let mut acc = [0f32; 4];
8835        for gi in 0..gpr {
8836            let s = f16_to_f32(u16::from_le_bytes([
8837                scales[(g0 + gi) * 2],
8838                scales[(g0 + gi) * 2 + 1],
8839            ]));
8840            let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8841            let aw = _mm256_abs_epi8(w);
8842            for (k, xq) in xs.iter().enumerate() {
8843                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8844                let d = dpbusd_hsum(aw, _mm256_sign_epi8(x, w));
8845                acc[k] += d as f32 * s;
8846            }
8847        }
8848        acc
8849    }
8850}
8851
8852/// The vbit flavor of the blocked 1×4: the per-activation A8W8 scale
8853/// folds in PER GROUP as `(d·sx)·s` — bit-matching the single-matvec
8854/// accumulation order (the q4_block flavor applies sx once at the end,
8855/// matching ITS single path; the two conventions are historical and
8856/// each blocked leg must mirror its own).
8857#[cfg(target_arch = "x86_64")]
8858#[target_feature(enable = "avx2")]
8859unsafe fn dot_q4b_row_1x4_sx_avx2(
8860    buf: &[u8],
8861    scales: &[u8],
8862    g0: usize,
8863    gpr: usize,
8864    xs: [&[i8]; 4],
8865    sxs: [f32; 4],
8866) -> [f32; 4] {
8867    // SAFETY: callers uphold buffer contracts (buf.len() == gpr·32).
8868    unsafe {
8869        use core::arch::x86_64::*;
8870        let ones = _mm256_set1_epi16(1);
8871        let mut acc = [0f32; 4];
8872        for gi in 0..gpr {
8873            let s = f16_to_f32(u16::from_le_bytes([
8874                scales[(g0 + gi) * 2],
8875                scales[(g0 + gi) * 2 + 1],
8876            ]));
8877            let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8878            let aw = _mm256_abs_epi8(w);
8879            for (k, xq) in xs.iter().enumerate() {
8880                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8881                let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
8882                let d = _mm256_madd_epi16(p16, ones);
8883                let hi128 = _mm256_extracti128_si256::<1>(d);
8884                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
8885                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8886                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8887                acc[k] += (_mm_cvtsi128_si32(s32) as f32 * sxs[k]) * s;
8888            }
8889        }
8890        acc
8891    }
8892}
8893
8894/// VNNI twin of `dot_q4b_row_1x4_sx_avx2` (see `dpbusd_hsum`; the
8895/// per-group `(d·sx)·s` fold mirrors the vbit single path).
8896#[cfg(target_arch = "x86_64")]
8897#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
8898unsafe fn dot_q4b_row_1x4_sx_vnni(
8899    buf: &[u8],
8900    scales: &[u8],
8901    g0: usize,
8902    gpr: usize,
8903    xs: [&[i8]; 4],
8904    sxs: [f32; 4],
8905) -> [f32; 4] {
8906    // SAFETY: callers uphold buffer contracts (buf.len() == gpr·32).
8907    unsafe {
8908        use core::arch::x86_64::*;
8909        let mut acc = [0f32; 4];
8910        for gi in 0..gpr {
8911            let s = f16_to_f32(u16::from_le_bytes([
8912                scales[(g0 + gi) * 2],
8913                scales[(g0 + gi) * 2 + 1],
8914            ]));
8915            let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8916            let aw = _mm256_abs_epi8(w);
8917            for (k, xq) in xs.iter().enumerate() {
8918                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8919                let d = dpbusd_hsum(aw, _mm256_sign_epi8(x, w));
8920                acc[k] += (d as f32 * sxs[k]) * s;
8921            }
8922        }
8923        acc
8924    }
8925}
8926
8927#[allow(unreachable_code)]
8928fn dot_q4_row_i8(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
8929    #[cfg(target_arch = "aarch64")]
8930    unsafe {
8931        return dot_q4_row_sdot(packed, scales, g0, gpr, xq);
8932    }
8933    #[cfg(target_arch = "x86_64")]
8934    unsafe {
8935        return dot_q4_row_avx2(packed, scales, g0, gpr, xq);
8936    }
8937    let mut acc = 0f32;
8938    for gi in 0..gpr {
8939        let g = g0 + gi;
8940        let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
8941        let mut d = 0i32;
8942        for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
8943            d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
8944                + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
8945        }
8946        acc += d as f32 * s;
8947    }
8948    acc
8949}
8950
8951/// Two-activation q4 row via the A8W8 int8 path (see `dot_q4_row_i8`).
8952#[inline]
8953#[allow(unreachable_code)]
8954fn dot_q4_row_i8_2(
8955    packed: &[u8],
8956    scales: &[u8],
8957    g0: usize,
8958    gpr: usize,
8959    xq1: &[i8],
8960    xq2: &[i8],
8961) -> (f32, f32) {
8962    #[cfg(target_arch = "aarch64")]
8963    unsafe {
8964        return dot_q4_row_sdot2(packed, scales, g0, gpr, xq1, xq2);
8965    }
8966    #[cfg(target_arch = "x86_64")]
8967    unsafe {
8968        return dot_q4_row_avx2_2(packed, scales, g0, gpr, xq1, xq2);
8969    }
8970    (
8971        dot_q4_row_i8(packed, scales, g0, gpr, xq1),
8972        dot_q4_row_i8(packed, scales, g0, gpr, xq2),
8973    )
8974}
8975
8976/// One q4 row range via SDOT (kernel body of `q4matvec`, extracted so
8977/// multi-matrix jobs can drive it for several tensors in one dispatch).
8978#[allow(clippy::too_many_arguments)]
8979fn q4_range_a8w8(
8980    packed: &[u8],
8981    scales: &[u8],
8982    gpr: usize,
8983    cols: usize,
8984    act: &SplitAct,
8985    out: SendMut,
8986    start: usize,
8987    end: usize,
8988) {
8989    for r in start..end {
8990        let mut acc = dot_q4_row_i8(packed, scales, r * gpr, gpr, &act.xq) * act.sx;
8991        // xq is zeroed at outlier slots — add the exact terms.
8992        for &(j, xv) in &act.outliers {
8993            let flat = r * cols + j;
8994            let byte = packed[flat / 2];
8995            let nib = if flat & 1 == 0 {
8996                byte & 0x0F
8997            } else {
8998                byte >> 4
8999            };
9000            let s = f16_to_f32(u16::from_le_bytes([
9001                scales[(flat / GROUP_SIZE) * 2],
9002                scales[(flat / GROUP_SIZE) * 2 + 1],
9003            ]));
9004            acc += ((nib as i32 - 8) as f32) * s * xv;
9005        }
9006        // SAFETY: disjoint row ranges per worker.
9007        unsafe { *out.at(r) = acc };
9008    }
9009}
9010
9011/// Two-input q4 row range via the A8W8 int8 path — kernel body of
9012/// `q4matvec2`, extracted for pair multi-matrix jobs.
9013#[allow(clippy::too_many_arguments)]
9014fn q4_range2_a8w8(
9015    packed: &[u8],
9016    scales: &[u8],
9017    gpr: usize,
9018    cols: usize,
9019    a1: &SplitAct,
9020    a2: &SplitAct,
9021    p1: SendMut,
9022    p2: SendMut,
9023    start: usize,
9024    end: usize,
9025) {
9026    for r in start..end {
9027        let (s1, s2) = dot_q4_row_i8_2(packed, scales, r * gpr, gpr, &a1.xq, &a2.xq);
9028        let mut acc1 = s1 * a1.sx;
9029        let mut acc2 = s2 * a2.sx;
9030        // xq is zeroed at outlier slots — add the exact terms.
9031        let fix = |outliers: &[(usize, f32)], acc: &mut f32| {
9032            for &(j, xv) in outliers {
9033                let flat = r * cols + j;
9034                let byte = packed[flat / 2];
9035                let nib = if flat & 1 == 0 {
9036                    byte & 0x0F
9037                } else {
9038                    byte >> 4
9039                };
9040                let s = f16_to_f32(u16::from_le_bytes([
9041                    scales[(flat / GROUP_SIZE) * 2],
9042                    scales[(flat / GROUP_SIZE) * 2 + 1],
9043                ]));
9044                *acc += ((nib as i32 - 8) as f32) * s * xv;
9045            }
9046        };
9047        fix(&a1.outliers, &mut acc1);
9048        fix(&a2.outliers, &mut acc2);
9049        // SAFETY: disjoint row ranges per worker.
9050        unsafe {
9051            *p1.at(r) = acc1;
9052            *p2.at(r) = acc2;
9053        }
9054    }
9055}
9056
9057/// Exact scalar q4 row range (same extraction, non-SDOT path).
9058fn q4_range_f32(
9059    packed: &[u8],
9060    scales: &[u8],
9061    gpr: usize,
9062    x: &[f32],
9063    out: SendMut,
9064    start: usize,
9065    end: usize,
9066) {
9067    for r in start..end {
9068        let mut acc = 0f32;
9069        for gi in 0..gpr {
9070            let g = r * gpr + gi;
9071            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
9072            let pk = &packed[g * 16..(g + 1) * 16];
9073            let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
9074            let mut ga = 0f32;
9075            for (k, &b) in pk.iter().enumerate() {
9076                ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
9077                    + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
9078            }
9079            acc += ga * s;
9080        }
9081        // SAFETY: disjoint row ranges per worker.
9082        unsafe { *out.at(r) = acc };
9083    }
9084}
9085
9086/// Fused two-input q4 matvec: nibbles are unpacked ONCE per group and
9087/// dotted against both activations (was: two full matvecs — double
9088/// weight traffic). Per-lane math matches `q4matvec` exactly.
9089#[allow(clippy::too_many_arguments)]
9090fn q4matvec2(
9091    bytes: &[u8],
9092    x1: &[f32],
9093    x2: &[f32],
9094    rows: usize,
9095    cols: usize,
9096    o1: &mut [f32],
9097    o2: &mut [f32],
9098    pool: Option<&Pool>,
9099) {
9100    debug_assert_eq!(o1.len(), rows);
9101    debug_assert_eq!(o2.len(), rows);
9102    let (packed, scales) = q4_split(bytes, rows, cols);
9103    let gpr = cols / GROUP_SIZE;
9104
9105    if a8w8_enabled() {
9106        let a1 = split_act(x1);
9107        let a2 = split_act(x2);
9108        let p1 = SendMut(o1.as_mut_ptr());
9109        let p2 = SendMut(o2.as_mut_ptr());
9110        let run = move |start: usize, end: usize| {
9111            q4_range2_a8w8(packed, scales, gpr, cols, &a1, &a2, p1, p2, start, end)
9112        };
9113        dispatch_rows(pool, rows, &run);
9114        return;
9115    }
9116
9117    let p1 = SendMut(o1.as_mut_ptr());
9118    let p2 = SendMut(o2.as_mut_ptr());
9119    let run = move |start: usize, end: usize| {
9120        q4_range2_f32(packed, scales, gpr, x1, x2, p1, p2, start, end)
9121    };
9122    dispatch_rows(pool, rows, &run);
9123}
9124
9125/// Two-input exact scalar q4 row range (same extraction).
9126#[allow(clippy::too_many_arguments)]
9127fn q4_range2_f32(
9128    packed: &[u8],
9129    scales: &[u8],
9130    gpr: usize,
9131    x1: &[f32],
9132    x2: &[f32],
9133    p1: SendMut,
9134    p2: SendMut,
9135    start: usize,
9136    end: usize,
9137) {
9138    for r in start..end {
9139        let (mut acc1, mut acc2) = (0f32, 0f32);
9140        for gi in 0..gpr {
9141            let g = r * gpr + gi;
9142            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
9143            let pk = &packed[g * 16..(g + 1) * 16];
9144            let x1g = &x1[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
9145            let x2g = &x2[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
9146            let (mut g1, mut g2) = (0f32, 0f32);
9147            for (k, &b) in pk.iter().enumerate() {
9148                let wl = (b & 0x0F) as f32 - 8.0;
9149                let wh = ((b >> 4) & 0x0F) as f32 - 8.0;
9150                g1 += wl * x1g[k * 2] + wh * x1g[k * 2 + 1];
9151                g2 += wl * x2g[k * 2] + wh * x2g[k * 2 + 1];
9152            }
9153            acc1 += g1 * s;
9154            acc2 += g2 * s;
9155        }
9156        // SAFETY: disjoint row ranges per worker.
9157        unsafe {
9158            *p1.at(r) = acc1;
9159            *p2.at(r) = acc2;
9160        }
9161    }
9162}
9163
9164thread_local! {
9165    /// Per-worker decoded-row scratch for the batched q4/vbit kernels
9166    /// (centered i8 for SDOT, f32 for the exact/scalar paths).
9167    static ROW_I8: std::cell::RefCell<Vec<u8>> = const { std::cell::RefCell::new(Vec::new()) };
9168    static ROW_F32: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
9169}
9170
9171/// Batched q4 matmat: each weight row is unpacked from the mmap ONCE
9172/// and dotted against ALL b activations (prefill used to fall back to b
9173/// full matvecs — b× weight traffic and b× nibble decode). Per-position
9174/// math matches `q4matvec` exactly: same group order, same accumulation.
9175/// `out` is row-major [b, rows] like `qmatmat`.
9176#[allow(clippy::too_many_arguments)]
9177fn q4matmat(
9178    bytes: &[u8],
9179    xs_all: &[f32],
9180    b: usize,
9181    rows: usize,
9182    cols: usize,
9183    out: &mut [f32],
9184    pool: Option<&Pool>,
9185) {
9186    debug_assert_eq!(xs_all.len(), b * cols);
9187    debug_assert_eq!(out.len(), b * rows);
9188    let (packed, scales) = q4_split(bytes, rows, cols);
9189    let gpr = cols / GROUP_SIZE;
9190    let gscale = |g: usize| f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
9191
9192    if a8w8_enabled() {
9193        let acts: Vec<SplitAct> = (0..b)
9194            .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
9195            .collect();
9196        let acts = &acts;
9197        let out_addr = SendMut(out.as_mut_ptr());
9198        let run = move |start: usize, end: usize| {
9199            ROW_I8.with(|rb| {
9200                let mut buf = rb.borrow_mut();
9201                buf.resize(cols, 0);
9202                for r in start..end {
9203                    // Unpack the row's nibbles to centered i8 once
9204                    // (element 2k = low nibble, 2k+1 = high — flat order,
9205                    // same as dot_q4_row_sdot's zip).
9206                    for gi in 0..gpr {
9207                        let g = r * gpr + gi;
9208                        for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
9209                            buf[gi * GROUP_SIZE + k * 2] = ((bt & 0x0F) as i32 - 8) as i8 as u8;
9210                            buf[gi * GROUP_SIZE + k * 2 + 1] =
9211                                (((bt >> 4) & 0x0F) as i32 - 8) as i8 as u8;
9212                        }
9213                    }
9214                    let mut bi = 0usize;
9215                    #[cfg(target_arch = "x86_64")]
9216                    if avx2_enabled() && blocked_enabled() {
9217                        while bi + 4 <= acts.len() {
9218                            let xs = [
9219                                acts[bi].xq.as_slice(),
9220                                acts[bi + 1].xq.as_slice(),
9221                                acts[bi + 2].xq.as_slice(),
9222                                acts[bi + 3].xq.as_slice(),
9223                            ];
9224                            let d = unsafe {
9225                                if vnni_tiles_enabled() {
9226                                    dot_q4b_row_1x4_vnni(&buf, scales, r * gpr, gpr, xs)
9227                                } else {
9228                                    dot_q4b_row_1x4_avx2(&buf, scales, r * gpr, gpr, xs)
9229                                }
9230                            };
9231                            for k in 0..4 {
9232                                let act = &acts[bi + k];
9233                                let mut acc = d[k] * act.sx;
9234                                for &(j, xv) in &act.outliers {
9235                                    acc += (buf[j] as i8) as f32
9236                                        * gscale((r * cols + j) / GROUP_SIZE)
9237                                        * xv;
9238                                }
9239                                // SAFETY: disjoint (bi, r) cells per worker.
9240                                unsafe { *out_addr.at((bi + k) * rows + r) = acc };
9241                            }
9242                            bi += 4;
9243                        }
9244                    }
9245                    while bi < acts.len() {
9246                        let act = &acts[bi];
9247                        let mut acc = 0f32;
9248                        for gi in 0..gpr {
9249                            let d = dot_i8_i8(
9250                                &buf[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
9251                                &act.xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
9252                            );
9253                            acc += d as f32 * gscale(r * gpr + gi);
9254                        }
9255                        acc *= act.sx;
9256                        // xq is zeroed at outlier slots — exact terms.
9257                        for &(j, xv) in &act.outliers {
9258                            acc += (buf[j] as i8) as f32 * gscale((r * cols + j) / GROUP_SIZE) * xv;
9259                        }
9260                        // SAFETY: disjoint (bi, r) cells per worker row range.
9261                        unsafe { *out_addr.at(bi * rows + r) = acc };
9262                        bi += 1;
9263                    }
9264                }
9265            })
9266        };
9267        dispatch_rows(pool, rows, &run);
9268        return;
9269    }
9270
9271    let out_addr = SendMut(out.as_mut_ptr());
9272    let run = move |start: usize, end: usize| {
9273        ROW_F32.with(|rb| {
9274            let mut buf = rb.borrow_mut();
9275            buf.resize(cols, 0.0);
9276            for r in start..end {
9277                // Decode raw (nib − 8) values once; scales stay per-group
9278                // so the accumulation order matches q4matvec bit-for-bit.
9279                for gi in 0..gpr {
9280                    let g = r * gpr + gi;
9281                    for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
9282                        buf[gi * GROUP_SIZE + k * 2] = (bt & 0x0F) as f32 - 8.0;
9283                        buf[gi * GROUP_SIZE + k * 2 + 1] = ((bt >> 4) & 0x0F) as f32 - 8.0;
9284                    }
9285                }
9286                for bi in 0..b {
9287                    let x = &xs_all[bi * cols..(bi + 1) * cols];
9288                    let mut acc = 0f32;
9289                    for gi in 0..gpr {
9290                        let mut ga = 0f32;
9291                        // Pairwise (lo + hi) addition, matching
9292                        // q4matvec's `ga += lo·x + hi·x` shape exactly —
9293                        // a flat one-per-element loop rounds differently
9294                        // and broke bit-parity on the scalar (x86) path.
9295                        for k in 0..GROUP_SIZE / 2 {
9296                            let e = gi * GROUP_SIZE + k * 2;
9297                            ga += buf[e] * x[e] + buf[e + 1] * x[e + 1];
9298                        }
9299                        acc += ga * gscale(r * gpr + gi);
9300                    }
9301                    // SAFETY: disjoint (bi, r) cells per worker row range.
9302                    unsafe { *out_addr.at(bi * rows + r) = acc };
9303                }
9304            }
9305        })
9306    };
9307    dispatch_rows(pool, rows, &run);
9308}
9309
9310/// Batched vbit matmat: each variable-bit row is decoded from the mmap
9311/// ONCE for the whole microbatch. Same per-position math as
9312/// `vbitmatvec` (SDOT A8W8 with exact outliers / exact f32 for b=8 rows
9313/// and the scalar path).
9314#[allow(clippy::too_many_arguments)]
9315fn vbitmatmat(
9316    bytes: &[u8],
9317    offsets: &[usize],
9318    xs_all: &[f32],
9319    b: usize,
9320    rows: usize,
9321    cols: usize,
9322    out: &mut [f32],
9323    pool: Option<&Pool>,
9324) {
9325    debug_assert_eq!(xs_all.len(), b * cols);
9326    debug_assert_eq!(out.len(), b * rows);
9327    debug_assert_eq!(offsets.len(), rows + 1);
9328    let ng = cols / GROUP_SIZE;
9329    let bits = &bytes[..rows];
9330    let sc_off = rows;
9331    let gscale = |r: usize, g: usize| {
9332        let so = (r * ng + g) * 2;
9333        f16_to_f32(u16::from_le_bytes([
9334            bytes[sc_off + so],
9335            bytes[sc_off + so + 1],
9336        ]))
9337    };
9338
9339    // Decode row r's raw (u − L) values into `dst` (f32, unscaled).
9340    let decode_f32 = |r: usize, dst: &mut [f32]| {
9341        let bw = bits[r] as usize;
9342        let l = ((1i32 << (bw - 1)) - 1) as f32;
9343        let data = &bytes[offsets[r]..offsets[r + 1]];
9344        let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
9345        for d in dst.iter_mut() {
9346            while nbits < bw {
9347                acc = (acc << 8) | data[idx] as u64;
9348                idx += 1;
9349                nbits += 8;
9350            }
9351            let u = ((acc >> (nbits - bw)) & ((1u64 << bw) - 1)) as f32;
9352            nbits -= bw;
9353            *d = u - l;
9354        }
9355    };
9356
9357    if a8w8_enabled() {
9358        let acts: Vec<SplitAct> = (0..b)
9359            .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
9360            .collect();
9361        let acts = &acts;
9362        let out_addr = SendMut(out.as_mut_ptr());
9363        let run = move |start: usize, end: usize| {
9364            for r in start..end {
9365                let bw = bits[r] as usize;
9366                if bw == 8 {
9367                    // u−L reaches 128 → no i8 path; decode once, exact
9368                    // f32 dots for every position (same as vbitmatvec).
9369                    ROW_F32.with(|rb| {
9370                        let mut buf = rb.borrow_mut();
9371                        buf.resize(cols, 0.0);
9372                        decode_f32(r, &mut buf);
9373                        for bi in 0..b {
9374                            let x = &xs_all[bi * cols..(bi + 1) * cols];
9375                            let mut dot = 0f32;
9376                            for g in 0..ng {
9377                                let mut gd = 0f32;
9378                                for k in 0..GROUP_SIZE {
9379                                    gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
9380                                }
9381                                dot += gd * gscale(r, g);
9382                            }
9383                            // SAFETY: disjoint (bi, r) cells per worker range.
9384                            unsafe { *out_addr.at(bi * rows + r) = dot };
9385                        }
9386                    });
9387                    continue;
9388                }
9389                let l = (1i32 << (bw - 1)) - 1;
9390                let data = &bytes[offsets[r]..offsets[r + 1]];
9391                ROW_I8.with(|rb| {
9392                    let mut buf = rb.borrow_mut();
9393                    buf.resize(cols, 0);
9394                    #[inline(always)]
9395                    fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
9396                        for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
9397                            let u = unpack8::<B>(&data[blk * B..]);
9398                            for k in 0..8 {
9399                                chunk[k] = (u[k] - l) as i8 as u8;
9400                            }
9401                        }
9402                    }
9403                    match bw {
9404                        3 => fill::<3>(data, l, &mut buf),
9405                        4 => vbit_fill4(data, &mut buf),
9406                        5 => fill::<5>(data, l, &mut buf),
9407                        6 => fill::<6>(data, l, &mut buf),
9408                        _ => unreachable!("vbit bit-width {bw} (validated at load)"),
9409                    }
9410                    let mut bi = 0usize;
9411                    // The vbit scale table shares q4_block's layout
9412                    // (contiguous f16 per (row·ng + g)), so the same
9413                    // blocked 1×4 kernel serves the decoded row.
9414                    #[cfg(target_arch = "x86_64")]
9415                    if avx2_enabled() && blocked_enabled() {
9416                        while bi + 4 <= acts.len() {
9417                            let xs = [
9418                                acts[bi].xq.as_slice(),
9419                                acts[bi + 1].xq.as_slice(),
9420                                acts[bi + 2].xq.as_slice(),
9421                                acts[bi + 3].xq.as_slice(),
9422                            ];
9423                            let sxs = [
9424                                acts[bi].sx,
9425                                acts[bi + 1].sx,
9426                                acts[bi + 2].sx,
9427                                acts[bi + 3].sx,
9428                            ];
9429                            let d = unsafe {
9430                                if vnni_tiles_enabled() {
9431                                    dot_q4b_row_1x4_sx_vnni(
9432                                        &buf,
9433                                        &bytes[sc_off..],
9434                                        r * ng,
9435                                        ng,
9436                                        xs,
9437                                        sxs,
9438                                    )
9439                                } else {
9440                                    dot_q4b_row_1x4_sx_avx2(
9441                                        &buf,
9442                                        &bytes[sc_off..],
9443                                        r * ng,
9444                                        ng,
9445                                        xs,
9446                                        sxs,
9447                                    )
9448                                }
9449                            };
9450                            for k in 0..4 {
9451                                let act = &acts[bi + k];
9452                                let mut dot = d[k];
9453                                for &(j, xv) in &act.outliers {
9454                                    dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
9455                                }
9456                                // SAFETY: disjoint (bi, r) cells per worker.
9457                                unsafe { *out_addr.at((bi + k) * rows + r) = dot };
9458                            }
9459                            bi += 4;
9460                        }
9461                    }
9462                    while bi < acts.len() {
9463                        let act = &acts[bi];
9464                        let mut dot = 0f32;
9465                        for g in 0..ng {
9466                            let d = dot_i8_i8(
9467                                &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
9468                                &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
9469                            ) as f32
9470                                * act.sx;
9471                            dot += d * gscale(r, g);
9472                        }
9473                        for &(j, xv) in &act.outliers {
9474                            dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
9475                        }
9476                        // SAFETY: disjoint (bi, r) cells per worker range.
9477                        unsafe { *out_addr.at(bi * rows + r) = dot };
9478                        bi += 1;
9479                    }
9480                });
9481            }
9482        };
9483        dispatch_rows(pool, rows, &run);
9484        return;
9485    }
9486
9487    let out_addr = SendMut(out.as_mut_ptr());
9488    let run = move |start: usize, end: usize| {
9489        ROW_F32.with(|rb| {
9490            let mut buf = rb.borrow_mut();
9491            buf.resize(cols, 0.0);
9492            for r in start..end {
9493                decode_f32(r, &mut buf);
9494                for bi in 0..b {
9495                    let x = &xs_all[bi * cols..(bi + 1) * cols];
9496                    let mut dot = 0f32;
9497                    for g in 0..ng {
9498                        let mut gd = 0f32;
9499                        for k in 0..GROUP_SIZE {
9500                            gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
9501                        }
9502                        dot += gd * gscale(r, g);
9503                    }
9504                    // SAFETY: disjoint (bi, r) cells per worker range.
9505                    unsafe { *out_addr.at(bi * rows + r) = dot };
9506                }
9507            }
9508        })
9509    };
9510    dispatch_rows(pool, rows, &run);
9511}
9512
9513/// Build a GPU batch job for a q8-family mapped tensor (primary
9514/// shard): prescaled input + directory coordinates. None → not
9515/// GPU-eligible, caller stays on the CPU.
9516pub(crate) fn gpu_batch_job<'a>(
9517    t: &'a QTensor,
9518    x: &[f32],
9519) -> Option<(std::sync::Arc<CmfModel>, crate::gpu::BatchJob<'a>)> {
9520    match t {
9521        QTensor::Mapped {
9522            model,
9523            idx,
9524            dtype: dt @ (TensorDtype::Q8Row | TensorDtype::Q8_2f),
9525            rows,
9526            cols,
9527            row_scale,
9528            col_field,
9529            ..
9530        } => Some((
9531            model.clone(),
9532            crate::gpu::BatchJob {
9533                idx: *idx,
9534                rows: *rows,
9535                cols: *cols,
9536                row_scale,
9537                xs: prescale(x, col_field, *dt).into_owned(),
9538                layout: crate::gpu::BatchLayout::Q8,
9539            },
9540        )),
9541        // q1: raw f32 activations, tile-embedded scales.
9542        QTensor::Mapped {
9543            model,
9544            idx,
9545            dtype: TensorDtype::Q1,
9546            rows,
9547            cols,
9548            ..
9549        } => Some((
9550            model.clone(),
9551            crate::gpu::BatchJob {
9552                idx: *idx,
9553                rows: *rows,
9554                cols: *cols,
9555                row_scale: &[],
9556                xs: x.to_vec(),
9557                layout: crate::gpu::BatchLayout::Q1,
9558            },
9559        )),
9560        // q4_tiled / q4tp: raw f32 activations; the scales live in the
9561        // payload (inline tiles / row ladder), so row_scale stays empty.
9562        // The GDN projection batch already runs these layouts on Metal —
9563        // this arm lets the attention QKV batch reach the same kernels.
9564        QTensor::Mapped {
9565            model,
9566            idx,
9567            dtype: dt @ (TensorDtype::Q4Tiled | TensorDtype::Q4TiledP),
9568            rows,
9569            cols,
9570            ..
9571        } => Some((
9572            model.clone(),
9573            crate::gpu::BatchJob {
9574                idx: *idx,
9575                rows: *rows,
9576                cols: *cols,
9577                row_scale: &[],
9578                xs: x.to_vec(),
9579                layout: if *dt == TensorDtype::Q4Tiled {
9580                    crate::gpu::BatchLayout::Q4t
9581                } else {
9582                    crate::gpu::BatchLayout::Q4tp
9583                },
9584            },
9585        )),
9586        _ => None,
9587    }
9588}
9589
9590thread_local! {
9591    static PRESCALE_BUF1: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
9592    static PRESCALE_BUF2: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
9593}
9594
9595pub(crate) fn prescale<'a>(
9596    x: &'a [f32],
9597    col_field: &[f32],
9598    dtype: TensorDtype,
9599) -> std::borrow::Cow<'a, [f32]> {
9600    if dtype == TensorDtype::Q8_2f {
9601        x.iter().zip(col_field).map(|(a, c)| a * c).collect()
9602    } else {
9603        std::borrow::Cow::Borrowed(x)
9604    }
9605}
9606
9607/// θ col-field fold for q8_2f activations. Borrowed pass-through for
9608/// every other dtype, using thread-local buffers to eliminate per-matvec allocations.
9609pub(crate) fn prescale_with<R, F: FnOnce(&[f32]) -> R>(
9610    x: &[f32],
9611    col_field: &[f32],
9612    dtype: TensorDtype,
9613    buf_id: u8,
9614    f: F,
9615) -> R {
9616    if dtype == TensorDtype::Q8_2f {
9617        if buf_id == 1 {
9618            PRESCALE_BUF1.with(|b| {
9619                let mut buf = b.borrow_mut();
9620                buf.clear();
9621                buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
9622                f(&buf)
9623            })
9624        } else {
9625            PRESCALE_BUF2.with(|b| {
9626                let mut buf = b.borrow_mut();
9627                buf.clear();
9628                buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
9629                f(&buf)
9630            })
9631        }
9632    } else {
9633        f(x)
9634    }
9635}
9636
9637// ───────────────────── x86-64 AVX2 kernels (roadmap этап 2) ─────────────────────
9638
9639/// AVX2+FMA available? Default ON when the CPU supports both;
9640/// `CMF_AVX2=0` disables (falls back to the autovectorized loops).
9641#[cfg(target_arch = "x86_64")]
9642pub(crate) fn avx2_enabled() -> bool {
9643    use std::sync::OnceLock;
9644    static ON: OnceLock<bool> = OnceLock::new();
9645    *ON.get_or_init(|| {
9646        std::env::var("CMF_AVX2").map(|v| v != "0").unwrap_or(true)
9647            && std::arch::is_x86_feature_detected!("avx2")
9648            && std::arch::is_x86_feature_detected!("fma")
9649    })
9650}
9651
9652/// AVX2 A8W8 allowed? The quantized-activation contract is switched by
9653/// the SAME env as the ARM SDOT path: `CMF_SDOT=0` keeps exact kernels
9654/// (the golden-parity exact gate relies on it) — AVX2 f32 kernels stay
9655/// active either way, they are exact (regrouped sums only).
9656#[cfg(target_arch = "x86_64")]
9657fn avx2_a8w8_enabled() -> bool {
9658    if FLOAT_ACTIVATIONS.get() {
9659        return false;
9660    }
9661    use std::sync::OnceLock;
9662    static ON: OnceLock<bool> = OnceLock::new();
9663    *ON.get_or_init(|| {
9664        avx2_enabled() && std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true)
9665    })
9666}
9667
9668thread_local! {
9669    static FULL_GPU_Q8: std::cell::Cell<bool> = const { std::cell::Cell::new(false) };
9670}
9671
9672/// Match the graph's full-device q8 projection precision on MiMo's host
9673/// tail. Does not enable the GPU or bypass a CPU-only/device-refusal gate.
9674pub(crate) fn enter_full_gpu_q8_scope() -> impl Drop {
9675    struct Restore(bool, std::marker::PhantomData<std::rc::Rc<()>>);
9676    impl Drop for Restore {
9677        fn drop(&mut self) {
9678            FULL_GPU_Q8.set(self.0);
9679        }
9680    }
9681    Restore(FULL_GPU_Q8.replace(true), std::marker::PhantomData)
9682}
9683
9684// Dynamic MiMo experts must not change activation precision when a cache
9685// fill moves them from CPU to GPU. Thread-local: only the cold-expert
9686// dispatch selects float kernels; concurrent pipelines keep their policy.
9687thread_local! {
9688    static FLOAT_ACTIVATIONS: std::cell::Cell<bool> = const { std::cell::Cell::new(false) };
9689}
9690
9691pub(crate) fn float_activations_scope<R>(f: impl FnOnce() -> R) -> R {
9692    struct Restore(bool);
9693    impl Drop for Restore {
9694        fn drop(&mut self) {
9695            FLOAT_ACTIVATIONS.set(self.0);
9696        }
9697    }
9698    let _restore = Restore(FLOAT_ACTIVATIONS.replace(true));
9699    f()
9700}
9701
9702/// Row-exact batching: while set, the x86 batched kernels (`qmatmat`,
9703/// `q4tp_matmat`) compute every (weight row, token) cell with the
9704/// single-token kernel instead of the blocked 2×4 / 1×4 / 1×8 tiles, so a
9705/// token's result does not depend on the batch it rides in and equals its
9706/// matvec. On ARM `q4tp_matmat` keeps its 1×4 tile but in the matvec's
9707/// reduction order, and no q4tp batch takes the AMX or device GEMM. The
9708/// MiMo speculative verify holds it (`row_exact_scope`) — its accepted
9709/// rows must be the rows plain decode would have produced.
9710// Shared with pool workers, so overlapping requests must keep the mode
9711// enabled until the LAST scope leaves. Saving/restoring a global bool is
9712// incorrect when two threads enter and leave in a non-LIFO order.
9713static ROW_EXACT: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
9714
9715pub(crate) fn row_exact() -> bool {
9716    ROW_EXACT.load(std::sync::atomic::Ordering::Acquire) != 0
9717}
9718
9719fn counted_row_exact_scope<R>(active: &std::sync::atomic::AtomicUsize, f: impl FnOnce() -> R) -> R {
9720    struct Restore<'a>(&'a std::sync::atomic::AtomicUsize);
9721    impl Drop for Restore<'_> {
9722        fn drop(&mut self) {
9723            self.0.fetch_sub(1, std::sync::atomic::Ordering::AcqRel);
9724        }
9725    }
9726    active.fetch_add(1, std::sync::atomic::Ordering::AcqRel);
9727    let _restore = Restore(active);
9728    f()
9729}
9730
9731/// Run `f` with row-exact batching on (also released on unwind).
9732pub(crate) fn row_exact_scope<R>(f: impl FnOnce() -> R) -> R {
9733    counted_row_exact_scope(&ROW_EXACT, f)
9734}
9735
9736/// A8W8 quantized-activation path available on THIS machine? One
9737/// switch across architectures: ARM dotprod (CMF_SDOT) or x86 AVX2
9738/// (CMF_AVX2 + the same CMF_SDOT exact-contract override).
9739#[inline]
9740pub(crate) fn a8w8_enabled() -> bool {
9741    #[cfg(target_arch = "aarch64")]
9742    {
9743        sdot_enabled()
9744    }
9745    #[cfg(target_arch = "x86_64")]
9746    {
9747        avx2_a8w8_enabled()
9748    }
9749    #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
9750    {
9751        false
9752    }
9753}
9754
9755/// int8·int8 dot dispatch: SDOT on ARM; AVX-512 VNNI (vpdpbusd) or AVX2
9756/// maddubs on x86. Callers are gated by `a8w8_enabled()`.
9757#[inline]
9758#[allow(unreachable_code)]
9759fn dot_i8_i8(w: &[u8], xq: &[i8]) -> i32 {
9760    #[cfg(target_arch = "aarch64")]
9761    unsafe {
9762        return dot_i8_sdot(w, xq);
9763    }
9764    #[cfg(target_arch = "x86_64")]
9765    unsafe {
9766        if avx512vnni_enabled() {
9767            return dot_i8_i8_vnni(w, xq);
9768        }
9769        return dot_i8_i8_avx2(w, xq);
9770    }
9771    w.iter()
9772        .zip(xq)
9773        .map(|(&a, &b)| (a as i8) as i32 * b as i32)
9774        .sum()
9775}
9776
9777/// AVX-512 VNNI available? (F+BW+VL+VNNI; `CMF_AVX512=0` falls back to
9778/// AVX2.) VL matters: short 32-byte groups (q4/vbit) ride the 256-bit
9779/// `vpdpbusd` encoding.
9780#[cfg(target_arch = "x86_64")]
9781fn avx512vnni_enabled() -> bool {
9782    use std::sync::OnceLock;
9783    static ON: OnceLock<bool> = OnceLock::new();
9784    *ON.get_or_init(|| {
9785        std::env::var("CMF_AVX512")
9786            .map(|v| v != "0")
9787            .unwrap_or(true)
9788            && std::arch::is_x86_feature_detected!("avx512f")
9789            && std::arch::is_x86_feature_detected!("avx512bw")
9790            && std::arch::is_x86_feature_detected!("avx512vl")
9791            && std::arch::is_x86_feature_detected!("avx512vnni")
9792    })
9793}
9794
9795/// Grouped-codec VNNI arms (the q4t/q4b/q1/q1t tile kernels): default
9796/// ON where AVX-512 VNNI exists (`CMF_VNNI_TILES=0` opt-out). Measured
9797/// on Ryzen 7950X (Zen4, 3 alternating process pairs, blocked GEMM
9798/// 4864×896 b=256): q4t 63→68 GF/s (+8%), q1 53→56 (+6%), q4b 72→75
9799/// (+4%) — consistent, no leg regressed. The tile kernels keep a
9800/// horizontal reduce per 32-weight group, so the `vpdpbusd` saving is
9801/// smaller than the long-dot q8 win (+13%), but it is real and free.
9802#[cfg(target_arch = "x86_64")]
9803fn vnni_tiles_enabled() -> bool {
9804    use std::sync::OnceLock;
9805    static ON: OnceLock<bool> = OnceLock::new();
9806    *ON.get_or_init(|| {
9807        std::env::var("CMF_VNNI_TILES")
9808            .map(|v| v != "0")
9809            .unwrap_or(true)
9810            && avx512vnni_enabled()
9811    })
9812}
9813
9814/// One 256-bit u8×i8 dot → i32 via `vpdpbusd` into a fresh accumulator
9815/// plus the same horizontal reduce the AVX2 kernels use. Products are
9816/// bounded (|w| ≤ 8 or ≤ 1), so maddubs never saturated — the i32 sum
9817/// is bit-identical to the maddubs+madd pair it replaces.
9818#[cfg(target_arch = "x86_64")]
9819#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
9820#[inline]
9821unsafe fn dpbusd_hsum(aw: core::arch::x86_64::__m256i, xs: core::arch::x86_64::__m256i) -> i32 {
9822    // SAFETY: pure register math.
9823    unsafe {
9824        use core::arch::x86_64::*;
9825        let d = _mm256_dpbusd_epi32(_mm256_setzero_si256(), aw, xs);
9826        let hi128 = _mm256_extracti128_si256::<1>(d);
9827        let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
9828        let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
9829        let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
9830        _mm_cvtsi128_si32(s32)
9831    }
9832}
9833
9834/// int8·int8 via AVX-512 VNNI: `vpdpbusd` fuses the maddubs+madd+add
9835/// triple into one u8×i8 dot-accumulate. AVX-512 has no vpsignb, so the
9836/// |w|·sign(x,w) trick becomes |w| × (x negated where w<0) via a mask
9837/// subtract — w==0 lanes contribute 0 through |w|=0 either way.
9838#[cfg(target_arch = "x86_64")]
9839#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
9840unsafe fn dot_i8_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
9841    // SAFETY: callers uphold slice-length contracts (see call sites).
9842    unsafe {
9843        use core::arch::x86_64::*;
9844        let n = w.len();
9845        let mut j = 0usize;
9846        let mut total: i32;
9847        // 4 independent accumulators: vpdpbusd is its own loop-carried
9848        // dependency (~5-cycle latency) — a single-acc loop runs
9849        // latency-bound and LOSES to the AVX2 maddubs kernel, measured
9850        // on Granite Rapids.
9851        {
9852            #[inline(always)]
9853            unsafe fn step(
9854                w: *const u8,
9855                x: *const i8,
9856                acc: core::arch::x86_64::__m512i,
9857            ) -> core::arch::x86_64::__m512i {
9858                unsafe {
9859                    use core::arch::x86_64::*;
9860                    let wv = _mm512_loadu_si512(w as *const _);
9861                    let xv = _mm512_loadu_si512(x as *const _);
9862                    let aw = _mm512_abs_epi8(wv);
9863                    let neg = _mm512_movepi8_mask(wv);
9864                    let sx = _mm512_mask_sub_epi8(xv, neg, _mm512_setzero_si512(), xv);
9865                    _mm512_dpbusd_epi32(acc, aw, sx)
9866                }
9867            }
9868            let (mut a0, mut a1, mut a2, mut a3) = (
9869                _mm512_setzero_si512(),
9870                _mm512_setzero_si512(),
9871                _mm512_setzero_si512(),
9872                _mm512_setzero_si512(),
9873            );
9874            while j + 256 <= n {
9875                a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
9876                a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), a1);
9877                a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), a2);
9878                a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), a3);
9879                j += 256;
9880            }
9881            while j + 64 <= n {
9882                a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
9883                j += 64;
9884            }
9885            let s01 = _mm512_add_epi32(a0, a1);
9886            let s23 = _mm512_add_epi32(a2, a3);
9887            total = _mm512_reduce_add_epi32(_mm512_add_epi32(s01, s23));
9888        }
9889        // 32-wide (q4/vbit groups are exactly 32 bytes).
9890        if j + 32 <= n {
9891            let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
9892            let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
9893            let d = _mm256_dpbusd_epi32(
9894                _mm256_setzero_si256(),
9895                _mm256_abs_epi8(wv),
9896                _mm256_sign_epi8(xv, wv),
9897            );
9898            let hi128 = _mm256_extracti128_si256::<1>(d);
9899            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
9900            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
9901            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
9902            total += _mm_cvtsi128_si32(s32);
9903            j += 32;
9904        }
9905        while j < n {
9906            total += (w[j] as i8) as i32 * xq[j] as i32;
9907            j += 1;
9908        }
9909        total
9910    }
9911}
9912
9913/// i8 row · f32 x via AVX2/FMA (x86 mirror of `dot_i8_f32_neon`).
9914#[cfg(target_arch = "x86_64")]
9915#[target_feature(enable = "avx2,fma")]
9916unsafe fn dot_i8_f32_avx2(w: &[u8], x: &[f32]) -> f32 {
9917    // SAFETY: callers uphold slice-length contracts (see call sites).
9918    unsafe {
9919        use core::arch::x86_64::*;
9920        let n = x.len();
9921        let wp = w.as_ptr();
9922        let xp = x.as_ptr();
9923        let (mut a0, mut a1) = (_mm256_setzero_ps(), _mm256_setzero_ps());
9924        let mut j = 0usize;
9925        while j + 16 <= n {
9926            let wb = _mm_loadu_si128(wp.add(j) as *const __m128i);
9927            let lo = _mm256_cvtepi8_epi32(wb);
9928            let hi = _mm256_cvtepi8_epi32(_mm_srli_si128::<8>(wb));
9929            a0 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(lo), _mm256_loadu_ps(xp.add(j)), a0);
9930            a1 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(hi), _mm256_loadu_ps(xp.add(j + 8)), a1);
9931            j += 16;
9932        }
9933        let acc = _mm256_add_ps(a0, a1);
9934        let hi128 = _mm256_extractf128_ps::<1>(acc);
9935        let s128 = _mm_add_ps(_mm256_castps256_ps128(acc), hi128);
9936        let s64 = _mm_add_ps(s128, _mm_movehl_ps(s128, s128));
9937        let s32 = _mm_add_ss(s64, _mm_shuffle_ps::<1>(s64, s64));
9938        let mut sum = _mm_cvtss_f32(s32);
9939        while j < n {
9940            sum += (*wp.add(j) as i8) as f32 * *xp.add(j);
9941            j += 1;
9942        }
9943        sum
9944    }
9945}
9946
9947/// int8(weight)·int8(activation) → i32 via AVX2 maddubs — the x86
9948/// analogue of the SDOT A8W8 path. `maddubs` takes u8×i8, so the
9949/// standard sign trick applies: |w| × sign(x, w) ≡ w × x per lane.
9950/// Pair saturation is safe: |w|≤128, |x|≤127 → 2·128·127 < 32767.
9951#[cfg(target_arch = "x86_64")]
9952#[target_feature(enable = "avx2")]
9953unsafe fn dot_i8_i8_avx2(w: &[u8], xq: &[i8]) -> i32 {
9954    // SAFETY: callers uphold slice-length contracts (see call sites).
9955    unsafe {
9956        use core::arch::x86_64::*;
9957        let n = w.len();
9958        let ones = _mm256_set1_epi16(1);
9959        let mut acc = _mm256_setzero_si256();
9960        let mut j = 0usize;
9961        while j + 32 <= n {
9962            let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
9963            let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
9964            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
9965            acc = _mm256_add_epi32(acc, _mm256_madd_epi16(p16, ones));
9966            j += 32;
9967        }
9968        let hi128 = _mm256_extracti128_si256::<1>(acc);
9969        let s128 = _mm_add_epi32(_mm256_castsi256_si128(acc), hi128);
9970        let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
9971        let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
9972        let mut s = _mm_cvtsi128_si32(s32);
9973        while j < n {
9974            s += (w[j] as i8) as i32 * xq[j] as i32;
9975            j += 1;
9976        }
9977        s
9978    }
9979}
9980
9981/// smmla 2×4: one instruction covers a 2-row × 2-activation × 8-deep
9982/// tile (32 MACs vs sdot's 16) — the weight pair loads once per 8-k
9983/// slice as a combined 2×8 register and meets two activation pairs.
9984#[cfg(target_arch = "aarch64")]
9985#[target_feature(enable = "neon,i8mm")]
9986unsafe fn dot_i8_smmla_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
9987    // SAFETY: callers uphold slice-length contracts.
9988    unsafe {
9989        use core::arch::aarch64::*;
9990        use core::arch::asm;
9991        let n = w0.len();
9992        let w0p = w0.as_ptr() as *const i8;
9993        let w1p = w1.as_ptr() as *const i8;
9994        // acc01 holds [c(r0,x0) c(r0,x1) c(r1,x0) c(r1,x1)]; acc23 the
9995        // same for x2/x3.
9996        let mut acc01 = vdupq_n_s32(0);
9997        let mut acc23 = vdupq_n_s32(0);
9998        let mut i = 0usize;
9999        while i + 8 <= n {
10000            let wa = vcombine_s8(vld1_s8(w0p.add(i)), vld1_s8(w1p.add(i)));
10001            let xb01 = vcombine_s8(
10002                vld1_s8(xs[0].as_ptr().add(i)),
10003                vld1_s8(xs[1].as_ptr().add(i)),
10004            );
10005            let xb23 = vcombine_s8(
10006                vld1_s8(xs[2].as_ptr().add(i)),
10007                vld1_s8(xs[3].as_ptr().add(i)),
10008            );
10009            asm!(
10010                "smmla {a01:v}.4s, {w:v}.16b, {x01:v}.16b",
10011                "smmla {a23:v}.4s, {w:v}.16b, {x23:v}.16b",
10012                a01 = inout(vreg) acc01, a23 = inout(vreg) acc23,
10013                w = in(vreg) wa, x01 = in(vreg) xb01, x23 = in(vreg) xb23,
10014                options(pure, nomem, nostack),
10015            );
10016            i += 8;
10017        }
10018        let mut out = [[0i32; 4]; 2];
10019        let a01: [i32; 4] = core::mem::transmute(acc01);
10020        let a23: [i32; 4] = core::mem::transmute(acc23);
10021        out[0][0] = a01[0];
10022        out[0][1] = a01[1];
10023        out[1][0] = a01[2];
10024        out[1][1] = a01[3];
10025        out[0][2] = a23[0];
10026        out[0][3] = a23[1];
10027        out[1][2] = a23[2];
10028        out[1][3] = a23[3];
10029        if i < n {
10030            for (k, x) in xs.iter().enumerate() {
10031                for j in i..n {
10032                    out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
10033                    out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
10034                }
10035            }
10036        }
10037        out
10038    }
10039}
10040
10041/// ARM twin of the x86 blocked prefill GEMM: two weight rows stay in
10042/// registers across four activation streams, eight sdot accumulators.
10043/// (The per-row form re-read each W row once per activation.)
10044#[cfg(target_arch = "aarch64")]
10045#[target_feature(enable = "neon,dotprod")]
10046unsafe fn dot_i8_sdot_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
10047    // SAFETY: callers uphold slice-length contracts.
10048    unsafe {
10049        use core::arch::aarch64::*;
10050        use core::arch::asm;
10051        let n = w0.len();
10052        let w0p = w0.as_ptr() as *const i8;
10053        let w1p = w1.as_ptr() as *const i8;
10054        let mut acc = [[vdupq_n_s32(0); 4]; 2];
10055        let mut i = 0usize;
10056        while i + 16 <= n {
10057            let wv0 = vld1q_s8(w0p.add(i));
10058            let wv1 = vld1q_s8(w1p.add(i));
10059            for (k, x) in xs.iter().enumerate() {
10060                let xv = vld1q_s8(x.as_ptr().add(i));
10061                let (mut a0, mut a1) = (acc[0][k], acc[1][k]);
10062                asm!(
10063                    "sdot {a0:v}.4s, {w0:v}.16b, {x:v}.16b",
10064                    "sdot {a1:v}.4s, {w1:v}.16b, {x:v}.16b",
10065                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
10066                    w0 = in(vreg) wv0, w1 = in(vreg) wv1, x = in(vreg) xv,
10067                    options(pure, nomem, nostack),
10068                );
10069                acc[0][k] = a0;
10070                acc[1][k] = a1;
10071            }
10072            i += 16;
10073        }
10074        let mut out = [[0i32; 4]; 2];
10075        for r in 0..2 {
10076            for k in 0..4 {
10077                out[r][k] = vaddvq_s32(acc[r][k]);
10078            }
10079        }
10080        if i < n {
10081            for (k, x) in xs.iter().enumerate() {
10082                for j in i..n {
10083                    out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
10084                    out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
10085                }
10086            }
10087        }
10088        out
10089    }
10090}
10091
10092/// Blocked 2 weight rows × 4 activations for the prefill GEMM
10093/// (roadmap P0: packed panels + multi-row accumulators). The two rows'
10094/// abs() live in registers across all four activation streams; the
10095/// sign-fixup is recomputed per pair (the price of the maddubs trick).
10096/// Returns raw i8·i8 dots; the caller applies scales and outliers.
10097#[cfg(target_arch = "x86_64")]
10098#[target_feature(enable = "avx2")]
10099unsafe fn dot_i8_i8_avx2_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
10100    // SAFETY: callers uphold slice-length contracts.
10101    unsafe {
10102        use core::arch::x86_64::*;
10103        let n = w0.len();
10104        let ones = _mm256_set1_epi16(1);
10105        let mut acc = [[_mm256_setzero_si256(); 4]; 2];
10106        let mut j = 0usize;
10107        while j + 32 <= n {
10108            let wv0 = _mm256_loadu_si256(w0.as_ptr().add(j) as *const __m256i);
10109            let wv1 = _mm256_loadu_si256(w1.as_ptr().add(j) as *const __m256i);
10110            let aw0 = _mm256_abs_epi8(wv0);
10111            let aw1 = _mm256_abs_epi8(wv1);
10112            for (k, x) in xs.iter().enumerate() {
10113                let xv = _mm256_loadu_si256(x.as_ptr().add(j) as *const __m256i);
10114                let p0 = _mm256_maddubs_epi16(aw0, _mm256_sign_epi8(xv, wv0));
10115                acc[0][k] = _mm256_add_epi32(acc[0][k], _mm256_madd_epi16(p0, ones));
10116                let p1 = _mm256_maddubs_epi16(aw1, _mm256_sign_epi8(xv, wv1));
10117                acc[1][k] = _mm256_add_epi32(acc[1][k], _mm256_madd_epi16(p1, ones));
10118            }
10119            j += 32;
10120        }
10121        let mut out = [[0i32; 4]; 2];
10122        for r in 0..2 {
10123            for k in 0..4 {
10124                let a = acc[r][k];
10125                let hi128 = _mm256_extracti128_si256::<1>(a);
10126                let s128 = _mm_add_epi32(_mm256_castsi256_si128(a), hi128);
10127                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
10128                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
10129                out[r][k] = _mm_cvtsi128_si32(s32);
10130            }
10131        }
10132        if j < n {
10133            for (k, x) in xs.iter().enumerate() {
10134                for i in j..n {
10135                    out[0][k] += (w0[i] as i8) as i32 * x[i] as i32;
10136                    out[1][k] += (w1[i] as i8) as i32 * x[i] as i32;
10137                }
10138            }
10139        }
10140        out
10141    }
10142}
10143
10144/// AVX2/VNNI q8 row dot with exact outlier correction (x86 mirror of
10145/// `row_dot_sdot` — same A8W8 contract). With AVX-512 VNNI the row goes
10146/// through the bias trick: Σ(w+128)·x via pure `vpdpbusd` (no per-lane
10147/// sign fixups), corrected by −128·Σx with Σx precomputed per split.
10148#[cfg(target_arch = "x86_64")]
10149#[inline]
10150fn row_dot_avx2(row: &[u8], act: &SplitAct) -> f32 {
10151    let dot = if avx512vnni_enabled() && row.len() >= 64 {
10152        (unsafe { dot_u8p128_i8_vnni(row, &act.xq) }) - 128 * act.xsum
10153    } else {
10154        unsafe { dot_i8_i8_avx2(row, &act.xq) }
10155    };
10156    let mut acc = dot as f32 * act.sx;
10157    for &(j, xv) in &act.outliers {
10158        acc += (row[j] as i8) as f32 * xv;
10159    }
10160    acc
10161}
10162
10163/// Σ (w[i]+128)·x[i] via pure `vpdpbusd` — the caller subtracts
10164/// 128·Σx. Four independent accumulators (dpbusd is ~5-cycle latency;
10165/// a single-acc loop runs latency-bound, measured on Granite Rapids).
10166#[cfg(target_arch = "x86_64")]
10167#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
10168unsafe fn dot_u8p128_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
10169    // SAFETY: callers uphold slice-length contracts (see call sites).
10170    unsafe {
10171        use core::arch::x86_64::*;
10172        let n = w.len();
10173        let flip = _mm512_set1_epi8(-128); // XOR 0x80: i8 w → u8 (w+128)
10174        #[inline(always)]
10175        unsafe fn step(
10176            w: *const u8,
10177            x: *const i8,
10178            flip: core::arch::x86_64::__m512i,
10179            acc: core::arch::x86_64::__m512i,
10180        ) -> core::arch::x86_64::__m512i {
10181            unsafe {
10182                use core::arch::x86_64::*;
10183                let wv = _mm512_xor_si512(_mm512_loadu_si512(w as *const _), flip);
10184                _mm512_dpbusd_epi32(acc, wv, _mm512_loadu_si512(x as *const _))
10185            }
10186        }
10187        let (mut a0, mut a1, mut a2, mut a3) = (
10188            _mm512_setzero_si512(),
10189            _mm512_setzero_si512(),
10190            _mm512_setzero_si512(),
10191            _mm512_setzero_si512(),
10192        );
10193        let mut j = 0usize;
10194        while j + 256 <= n {
10195            a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
10196            a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), flip, a1);
10197            a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), flip, a2);
10198            a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), flip, a3);
10199            j += 256;
10200        }
10201        while j + 64 <= n {
10202            a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
10203            j += 64;
10204        }
10205        let mut total = _mm512_reduce_add_epi32(_mm512_add_epi32(
10206            _mm512_add_epi32(a0, a1),
10207            _mm512_add_epi32(a2, a3),
10208        ));
10209        // Scalar tail: (w as i8) + 128 ≡ (w as u8) ^ 0x80.
10210        while j < n {
10211            total += ((w[j] ^ 0x80) as i32) * xq[j] as i32;
10212            j += 1;
10213        }
10214        total
10215    }
10216}
10217
10218/// One q4 row via AVX2: nibbles → centered i8 (unpacklo/hi restores the
10219/// writer's flat order, same as the NEON vzip pair), maddubs against
10220/// the pre-quantized activation group, × the group's f16 scale. Pair
10221/// saturation safe: |w|≤8, |x|≤127 → 2·8·127 ≪ 32767. Mirror of
10222/// `dot_q4_row_sdot`.
10223#[cfg(target_arch = "x86_64")]
10224#[target_feature(enable = "avx2")]
10225unsafe fn dot_q4_row_avx2(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
10226    // SAFETY: callers uphold slice-length contracts (16 packed bytes and
10227    // 2 scale bytes per group; xq.len() == gpr·GROUP_SIZE).
10228    unsafe {
10229        use core::arch::x86_64::*;
10230        let lomask = _mm_set1_epi8(0x0F);
10231        let eight = _mm256_set1_epi8(8);
10232        let ones = _mm256_set1_epi16(1);
10233        let mut acc = 0f32;
10234        for gi in 0..gpr {
10235            let g = g0 + gi;
10236            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
10237            let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
10238            let lo = _mm_and_si128(b, lomask);
10239            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
10240            let w = _mm256_sub_epi8(
10241                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
10242                eight,
10243            );
10244            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
10245            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
10246            let d = _mm256_madd_epi16(p16, ones);
10247            let hi128 = _mm256_extracti128_si256::<1>(d);
10248            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
10249            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
10250            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
10251            acc += _mm_cvtsi128_si32(s32) as f32 * s;
10252        }
10253        acc
10254    }
10255}
10256
10257/// Two-activation q4 row via AVX2: nibbles unpacked ONCE per group,
10258/// both activations dotted against the same centered i8 register.
10259#[cfg(target_arch = "x86_64")]
10260#[target_feature(enable = "avx2")]
10261unsafe fn dot_q4_row_avx2_2(
10262    packed: &[u8],
10263    scales: &[u8],
10264    g0: usize,
10265    gpr: usize,
10266    xq1: &[i8],
10267    xq2: &[i8],
10268) -> (f32, f32) {
10269    // SAFETY: callers uphold slice-length contracts (see dot_q4_row_avx2).
10270    unsafe {
10271        use core::arch::x86_64::*;
10272        let lomask = _mm_set1_epi8(0x0F);
10273        let eight = _mm256_set1_epi8(8);
10274        let ones = _mm256_set1_epi16(1);
10275        let (mut acc1, mut acc2) = (0f32, 0f32);
10276        #[inline(always)]
10277        unsafe fn hsum(d: core::arch::x86_64::__m256i) -> i32 {
10278            unsafe {
10279                use core::arch::x86_64::*;
10280                let hi128 = _mm256_extracti128_si256::<1>(d);
10281                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
10282                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
10283                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
10284                _mm_cvtsi128_si32(s32)
10285            }
10286        }
10287        for gi in 0..gpr {
10288            let g = g0 + gi;
10289            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
10290            let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
10291            let lo = _mm_and_si128(b, lomask);
10292            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
10293            let w = _mm256_sub_epi8(
10294                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
10295                eight,
10296            );
10297            let aw = _mm256_abs_epi8(w);
10298            let x1 = _mm256_loadu_si256(xq1.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
10299            let x2 = _mm256_loadu_si256(xq2.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
10300            let d1 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x1, w)), ones);
10301            let d2 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x2, w)), ones);
10302            acc1 += hsum(d1) as f32 * s;
10303            acc2 += hsum(d2) as f32 * s;
10304        }
10305        (acc1, acc2)
10306    }
10307}
10308
10309/// One q8 row range via AVX2 (x86 mirror of `q8_range_sdot`).
10310#[cfg(target_arch = "x86_64")]
10311fn q8_range_avx2(
10312    q: &[u8],
10313    row_scale: &[f32],
10314    act: &SplitAct,
10315    cols: usize,
10316    out_addr: SendMut,
10317    start: usize,
10318    end: usize,
10319) {
10320    for o in start..end {
10321        let v = row_dot_avx2(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
10322        // SAFETY: disjoint row ranges per worker.
10323        unsafe { *out_addr.at(o) = v };
10324    }
10325}
10326
10327/// Two-input q8 row range via AVX2 (x86 mirror of `q8_range2_sdot`).
10328#[cfg(target_arch = "x86_64")]
10329#[allow(clippy::too_many_arguments)]
10330fn q8_range2_avx2(
10331    q: &[u8],
10332    row_scale: &[f32],
10333    a1: &SplitAct,
10334    a2: &SplitAct,
10335    cols: usize,
10336    p1: SendMut,
10337    p2: SendMut,
10338    start: usize,
10339    end: usize,
10340) {
10341    for o in start..end {
10342        let row = &q[o * cols..(o + 1) * cols];
10343        // SAFETY: disjoint row ranges per worker.
10344        unsafe {
10345            *p1.at(o) = row_dot_avx2(row, a1) * row_scale[o];
10346            *p2.at(o) = row_dot_avx2(row, a2) * row_scale[o];
10347        }
10348    }
10349}
10350
10351// ───────────────────── A8W8 SDOT path (port of vmfcore, ×1.78 decode) ─────────────────────
10352
10353/// ARMv8.6 i8mm (smmla): 32 int8 MACs per instruction vs sdot's 16 —
10354/// yet MEASURED 2.4× SLOWER than the blocked sdot on Apple silicon
10355/// (108 vs 264 GF/s): the on-the-fly vcombine packing and the two-
10356/// accumulator dependency chain swamp the MAC advantage, and Apple's
10357/// four SIMD pipes already keep sdot fed. OPT-IN (CMF_I8MM=1) for
10358/// field trials on Cortex-A710/X-class parts with two pipes, where the
10359/// balance may differ; a pre-interleaved weight layout (repack infra)
10360/// is the known path if it ever earns its keep.
10361#[cfg(target_arch = "aarch64")]
10362fn i8mm_enabled() -> bool {
10363    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
10364    *ON.get_or_init(|| {
10365        std::env::var("CMF_I8MM").map(|v| v == "1").unwrap_or(false)
10366            && std::arch::is_aarch64_feature_detected!("i8mm")
10367    })
10368}
10369
10370/// SDOT enabled? Default ON when the CPU has ARMv8.2 dotprod;
10371/// `CMF_SDOT=0` disables (falls back to i8×f32 NEON).
10372/// (On non-ARM release builds only the test tolerance switch calls it.)
10373#[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
10374fn sdot_enabled() -> bool {
10375    if FLOAT_ACTIVATIONS.get() {
10376        return false;
10377    }
10378    use std::sync::OnceLock;
10379    static ON: OnceLock<bool> = OnceLock::new();
10380    *ON.get_or_init(|| {
10381        let want = std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true);
10382        if !want {
10383            return false;
10384        }
10385
10386        #[cfg(target_arch = "aarch64")]
10387        {
10388            if std::arch::is_aarch64_feature_detected!("dotprod") {
10389                return true;
10390            }
10391            #[cfg(target_os = "android")]
10392            {
10393                if let Ok(cpuinfo) = std::fs::read_to_string("/proc/cpuinfo") {
10394                    if cpuinfo.lines().any(|l| {
10395                        (l.starts_with("Features") || l.starts_with("features"))
10396                            && l.contains("asimddp")
10397                    }) {
10398                        return true;
10399                    }
10400                }
10401            }
10402            false
10403        }
10404        #[cfg(not(target_arch = "aarch64"))]
10405        {
10406            false
10407        }
10408    })
10409}
10410
10411/// Two-field activation split (≡ vmfcore `q8_split_prep`): outlier
10412/// channels (>8·rms) are computed exactly in f32; the bulk (outliers
10413/// zeroed → clean absmax) goes through int8 SDOT. Computed ONCE per
10414/// matvec, shared by all rows/workers.
10415struct SplitAct {
10416    xq: Vec<i8>,
10417    sx: f32,
10418    outliers: Vec<(usize, f32)>,
10419    /// Σ xq — the VNNI bias-trick correction (`(w+128)·x` sums need
10420    /// `−128·Σx`); one i32 per split, computed once per matvec.
10421    #[cfg_attr(not(target_arch = "x86_64"), allow(dead_code))]
10422    xsum: i32,
10423}
10424
10425thread_local! {
10426    /// Recycled xq buffers: split_act runs for every matvec (~200/token)
10427    /// and its hidden-size allocation was steady-state heap churn.
10428    static XQ_FREE: std::cell::RefCell<Vec<Vec<i8>>> =
10429        const { std::cell::RefCell::new(Vec::new()) };
10430}
10431
10432impl Drop for SplitAct {
10433    fn drop(&mut self) {
10434        let buf = std::mem::take(&mut self.xq);
10435        if buf.capacity() > 0 {
10436            XQ_FREE.with(|f| {
10437                let mut f = f.borrow_mut();
10438                if f.len() < 16 {
10439                    f.push(buf);
10440                }
10441            });
10442        }
10443    }
10444}
10445
10446thread_local! {
10447    /// One scratch row per WORKER, kept for the life of the thread.
10448    ///
10449    /// The kernels take a row of group scales per dispatch, and a fresh
10450    /// `vec![0f32; gpr]` inside the closure is one allocation per worker per
10451    /// dispatch — on the release checkpoint about six thousand a token, a
10452    /// quarter of everything the benchmark counts.
10453    static KROW: UnsafeCell<Vec<f32>> = const { UnsafeCell::new(Vec::new()) };
10454    /// Two scratch rows for a gate/up pair.  Q2TP and Q4TP fused SwiGLU
10455    /// decode each projection's ladder independently, but both ladders have
10456    /// the same lifetime; retaining them per worker removes two heap trips
10457    /// from every dispatched expert pair.
10458    static KROWS: UnsafeCell<[Vec<f32>; 2]> =
10459        const { UnsafeCell::new([Vec::new(), Vec::new()]) };
10460}
10461
10462/// Borrow `n` floats of the calling worker's scratch. Nothing inside a
10463/// kernel body borrows it again, which is what keeps the RefCell honest.
10464#[inline]
10465fn with_krow<R>(n: usize, f: impl FnOnce(&mut [f32]) -> R) -> R {
10466    KROW.with(|s| {
10467        // SAFETY: `KROW` is thread-local and this function does not recurse;
10468        // each worker has exclusive access to its own scratch row.
10469        let b = unsafe { &mut *s.get() };
10470        if b.len() < n {
10471            b.resize(n, 0.0);
10472        }
10473        f(&mut b[..n])
10474    })
10475}
10476
10477/// `t.round().clamp(-127.0, 127.0) as i8`, bit for bit, without the libm
10478/// call. On baseline x86-64 (no SSE4.1 `roundps`) `f32::round` is a
10479/// function call per element, and split_act runs it over every hidden
10480/// state before every matvec: measured 27 us a call on a 2048-wide
10481/// activation on an EPYC 7763 — 5.4 ms of a 55 ms decode token, all of
10482/// it on the caller's thread while thirty workers wait. Clamping first is
10483/// equivalent (round is monotonic and ±127 are integers), and after the
10484/// clamp `t - trunc(t)` is exact, so the half-away-from-zero decision is
10485/// the one `round` makes. NaN clamps to NaN and converts to 0, as before.
10486/// The loop vectorizes (cvttps2dq + compare/select).
10487#[inline(always)]
10488fn q8_round(t: f32) -> i8 {
10489    let t = t.clamp(-127.0, 127.0);
10490    let i = t as i32;
10491    let f = t - i as f32;
10492    let r = if f >= 0.5 {
10493        i + 1
10494    } else if f <= -0.5 {
10495        i - 1
10496    } else {
10497        i
10498    };
10499    r as i8
10500}
10501
10502#[inline]
10503fn with_krows<R>(n: usize, f: impl FnOnce(&mut [f32], &mut [f32]) -> R) -> R {
10504    KROWS.with(|s| {
10505        // SAFETY: KROWS is thread-local and no kernel body recursively calls
10506        // this helper; each participant owns its two vectors exclusively.
10507        let b = unsafe { &mut *s.get() };
10508        for row in b.iter_mut() {
10509            if row.len() < n {
10510                row.resize(n, 0.0);
10511            }
10512        }
10513        let (a, b) = b.split_at_mut(1);
10514        f(&mut a[0][..n], &mut b[0][..n])
10515    })
10516}
10517
10518#[inline]
10519fn silu_mul_limited(mut gate: f32, mut up: f32, limit: f32) -> f32 {
10520    if limit > 0.0 {
10521        up = up.clamp(-limit, limit);
10522        gate = gate.min(limit);
10523    }
10524    gate / (1.0 + (-gate).exp()) * up
10525}
10526
10527fn split_act(x: &[f32]) -> SplitAct {
10528    let _prof = crate::cpuprof::time(crate::cpuprof::Slot::SplitAct);
10529    let n = x.len();
10530    let rms = (x.iter().map(|&v| (v * v) as f64).sum::<f64>() / n.max(1) as f64).sqrt() as f32;
10531    let thr = 8.0 * rms;
10532    // One pass: collect outliers and the bulk absmax (outliers excluded —
10533    // identical to the old zero-then-fold over a copied buffer, minus the
10534    // full-vector copy).
10535    let mut outliers: Vec<(usize, f32)> = Vec::new();
10536    let mut amax = 0f32;
10537    for (j, &v) in x.iter().enumerate() {
10538        let a = v.abs();
10539        if a > thr {
10540            outliers.push((j, v));
10541        } else if a > amax {
10542            amax = a;
10543        }
10544    }
10545    let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
10546    let inv = 1.0 / sx;
10547    let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
10548    xq.clear();
10549    xq.reserve(n);
10550    if outliers.is_empty() {
10551        xq.extend(
10552            x.iter()
10553                .map(|&v| q8_round(v * inv)),
10554        );
10555    } else {
10556        // Outlier slots quantize to 0 (their exact term is added later).
10557        xq.extend(x.iter().map(|&v| {
10558            if v.abs() > thr {
10559                0
10560            } else {
10561                q8_round(v * inv)
10562            }
10563        }));
10564    }
10565    let xsum = xq.iter().map(|&v| v as i32).sum();
10566    SplitAct {
10567        xq,
10568        sx,
10569        outliers,
10570        xsum,
10571    }
10572}
10573
10574fn split_act_q8_2f(x: &[f32], col: &[f32]) -> SplitAct {
10575    let _prof = crate::cpuprof::time(crate::cpuprof::Slot::SplitAct);
10576    let n = x.len();
10577    let rms = (x
10578        .iter()
10579        .zip(col)
10580        .map(|(&a, &c)| {
10581            let v = a * c;
10582            (v * v) as f64
10583        })
10584        .sum::<f64>()
10585        / n.max(1) as f64)
10586        .sqrt() as f32;
10587    let thr = 8.0 * rms;
10588
10589    let mut outliers = Vec::new();
10590    let mut amax = 0f32;
10591    for (j, (&a, &c)) in x.iter().zip(col).enumerate() {
10592        let v = a * c;
10593        let s = v.abs();
10594        if s > thr {
10595            outliers.push((j, v));
10596        } else if s > amax {
10597            amax = s;
10598        }
10599    }
10600
10601    let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
10602    let inv = 1.0 / sx;
10603    let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
10604    xq.clear();
10605    xq.reserve(n);
10606    if outliers.is_empty() {
10607        xq.extend(
10608            x.iter()
10609                .zip(col)
10610                .map(|(&a, &c)| q8_round((a * c) * inv)),
10611        );
10612    } else {
10613        xq.extend(x.iter().zip(col).map(|(&a, &c)| {
10614            let v = a * c;
10615            if v.abs() > thr {
10616                0
10617            } else {
10618                q8_round(v * inv)
10619            }
10620        }));
10621    }
10622    let xsum = xq.iter().map(|&v| v as i32).sum();
10623    SplitAct {
10624        xq,
10625        sx,
10626        outliers,
10627        xsum,
10628    }
10629}
10630
10631/// int8(weight)·int8(activation) → i32 via `sdot` (inline asm — the
10632/// vdotq intrinsic is unstable; port of vmfcore `dot_i8_sdot`).
10633#[cfg(target_arch = "aarch64")]
10634#[target_feature(enable = "neon,dotprod")]
10635unsafe fn dot_i8_sdot(w: &[u8], xq: &[i8]) -> i32 {
10636    // SAFETY: callers uphold slice-length contracts (see call sites).
10637    unsafe {
10638        use core::arch::aarch64::*;
10639        use core::arch::asm;
10640        let wp = w.as_ptr() as *const i8;
10641        let n = w.len();
10642        let (mut a0, mut a1, mut a2, mut a3) = (
10643            vdupq_n_s32(0),
10644            vdupq_n_s32(0),
10645            vdupq_n_s32(0),
10646            vdupq_n_s32(0),
10647        );
10648        let mut i = 0;
10649        while i + 64 <= n {
10650            let (w0, x0) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
10651            let (w1, x1) = (vld1q_s8(wp.add(i + 16)), vld1q_s8(xq.as_ptr().add(i + 16)));
10652            let (w2, x2) = (vld1q_s8(wp.add(i + 32)), vld1q_s8(xq.as_ptr().add(i + 32)));
10653            let (w3, x3) = (vld1q_s8(wp.add(i + 48)), vld1q_s8(xq.as_ptr().add(i + 48)));
10654            asm!(
10655                "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
10656                "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
10657                "sdot {a2:v}.4s, {w2:v}.16b, {x2:v}.16b",
10658                "sdot {a3:v}.4s, {w3:v}.16b, {x3:v}.16b",
10659                a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
10660                w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
10661                w2 = in(vreg) w2, x2 = in(vreg) x2, w3 = in(vreg) w3, x3 = in(vreg) x3,
10662                options(pure, nomem, nostack),
10663            );
10664            i += 64;
10665        }
10666        while i + 16 <= n {
10667            let (wv, xv) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
10668            asm!("sdot {a:v}.4s, {w:v}.16b, {x:v}.16b",
10669                 a = inout(vreg) a0, w = in(vreg) wv, x = in(vreg) xv, options(pure, nomem, nostack));
10670            i += 16;
10671        }
10672        let mut s = vaddvq_s32(vaddq_s32(vaddq_s32(a0, a1), vaddq_s32(a2, a3)));
10673        while i < n {
10674            s += (*wp.add(i)) as i32 * xq[i] as i32;
10675            i += 1;
10676        }
10677        s
10678    }
10679}
10680
10681/// Row-blocked SDOT: 4 output rows per pass — the activation chunk is
10682/// loaded once and reused, 4 independent accumulators hide sdot latency
10683/// (port of vmfcore `dot_i8_sdot_4rows`).
10684#[cfg(target_arch = "aarch64")]
10685#[target_feature(enable = "neon,dotprod")]
10686unsafe fn dot_i8_sdot_4rows(w0: &[u8], w1: &[u8], w2: &[u8], w3: &[u8], xq: &[i8]) -> [i32; 4] {
10687    // SAFETY: callers uphold slice-length contracts (see call sites).
10688    unsafe {
10689        use core::arch::aarch64::*;
10690        use core::arch::asm;
10691        let n = xq.len();
10692        let px = xq.as_ptr();
10693        let (p0, p1, p2, p3) = (
10694            w0.as_ptr() as *const i8,
10695            w1.as_ptr() as *const i8,
10696            w2.as_ptr() as *const i8,
10697            w3.as_ptr() as *const i8,
10698        );
10699        let (mut a0, mut a1, mut a2, mut a3) = (
10700            vdupq_n_s32(0),
10701            vdupq_n_s32(0),
10702            vdupq_n_s32(0),
10703            vdupq_n_s32(0),
10704        );
10705        let mut i = 0;
10706        while i + 16 <= n {
10707            let x = vld1q_s8(px.add(i));
10708            let v0 = vld1q_s8(p0.add(i));
10709            let v1 = vld1q_s8(p1.add(i));
10710            let v2 = vld1q_s8(p2.add(i));
10711            let v3 = vld1q_s8(p3.add(i));
10712            asm!(
10713                "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
10714                "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
10715                "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
10716                "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
10717                a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
10718                v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
10719                options(pure, nomem, nostack),
10720            );
10721            i += 16;
10722        }
10723        let mut r = [
10724            vaddvq_s32(a0),
10725            vaddvq_s32(a1),
10726            vaddvq_s32(a2),
10727            vaddvq_s32(a3),
10728        ];
10729        while i < n {
10730            let xi = *px.add(i) as i32;
10731            r[0] += (*p0.add(i)) as i32 * xi;
10732            r[1] += (*p1.add(i)) as i32 * xi;
10733            r[2] += (*p2.add(i)) as i32 * xi;
10734            r[3] += (*p3.add(i)) as i32 * xi;
10735            i += 1;
10736        }
10737        r
10738    }
10739}
10740
10741/// 4 interleaved rows in one pass: the repacked group is [r0[c], r1[c],
10742/// r2[c], r3[c]] per 16-byte chunk, so each iteration reads ONE 64-byte
10743/// line plus the shared activation chunk — a single sequential weight
10744/// stream per worker. Per-row accumulation is the same one-accumulator
10745/// scheme as `dot_i8_sdot_4rows`; integer sums are exact, so outputs
10746/// are bit-identical to the mmap-layout kernel.
10747#[cfg(target_arch = "aarch64")]
10748#[target_feature(enable = "neon,dotprod")]
10749unsafe fn dot_i8_sdot_4rows_il(g: &[u8], xq: &[i8]) -> [i32; 4] {
10750    // SAFETY: callers uphold slice-length contracts (g.len() == 4·n,
10751    // n % 16 == 0 — guaranteed by the repack gate).
10752    unsafe {
10753        use core::arch::aarch64::*;
10754        use core::arch::asm;
10755        let n = xq.len();
10756        let px = xq.as_ptr();
10757        let pg = g.as_ptr() as *const i8;
10758        let (mut a0, mut a1, mut a2, mut a3) = (
10759            vdupq_n_s32(0),
10760            vdupq_n_s32(0),
10761            vdupq_n_s32(0),
10762            vdupq_n_s32(0),
10763        );
10764        let mut i = 0;
10765        while i + 16 <= n {
10766            let x = vld1q_s8(px.add(i));
10767            let base = pg.add(4 * i);
10768            let v0 = vld1q_s8(base);
10769            let v1 = vld1q_s8(base.add(16));
10770            let v2 = vld1q_s8(base.add(32));
10771            let v3 = vld1q_s8(base.add(48));
10772            asm!(
10773                "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
10774                "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
10775                "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
10776                "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
10777                a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
10778                v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
10779                options(pure, nomem, nostack),
10780            );
10781            i += 16;
10782        }
10783        [
10784            vaddvq_s32(a0),
10785            vaddvq_s32(a1),
10786            vaddvq_s32(a2),
10787            vaddvq_s32(a3),
10788        ]
10789    }
10790}
10791
10792/// One q8 row range via SDOT (4-row blocks + tail) — the body of
10793/// `qmatvec`'s hot loop, extracted so multi-matrix jobs can drive the
10794/// SAME kernel for several tensors under one pool dispatch. `rep` — the
10795/// load-time interleaved repack (empty = mmap layout only); rows outside
10796/// full 4-row groups always come from the mmap layout.
10797#[cfg(target_arch = "aarch64")]
10798fn q8_range_sdot(
10799    q: &[u8],
10800    rep: &[u8],
10801    row_scale: &[f32],
10802    act: &SplitAct,
10803    cols: usize,
10804    out_addr: SendMut,
10805    start: usize,
10806    end: usize,
10807) {
10808    let mut o = start;
10809    // Leading rows to the group boundary (repack path only): the pool
10810    // splits row ranges arbitrarily, groups are absolute.
10811    if !rep.is_empty() {
10812        while o < end && o % 4 != 0 {
10813            let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
10814            unsafe { *out_addr.at(o) = v };
10815            o += 1;
10816        }
10817    }
10818    while o + 4 <= end {
10819        let r = if rep.is_empty() {
10820            unsafe {
10821                dot_i8_sdot_4rows(
10822                    &q[o * cols..(o + 1) * cols],
10823                    &q[(o + 1) * cols..(o + 2) * cols],
10824                    &q[(o + 2) * cols..(o + 3) * cols],
10825                    &q[(o + 3) * cols..(o + 4) * cols],
10826                    &act.xq,
10827                )
10828            }
10829        } else {
10830            unsafe { dot_i8_sdot_4rows_il(&rep[o * cols..(o + 4) * cols], &act.xq) }
10831        };
10832        for k in 0..4 {
10833            let mut acc = r[k] as f32 * act.sx;
10834            for &(j, xv) in &act.outliers {
10835                acc += (q[(o + k) * cols + j] as i8) as f32 * xv;
10836            }
10837            // SAFETY: disjoint row ranges per worker.
10838            unsafe { *out_addr.at(o + k) = acc * row_scale[o + k] };
10839        }
10840        o += 4;
10841    }
10842    while o < end {
10843        let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
10844        unsafe { *out_addr.at(o) = v };
10845        o += 1;
10846    }
10847}
10848
10849/// Two-input q8 row range via SDOT — `qmatvec2`'s hot loop, extracted
10850/// for the fused pair multi-matrix job (`matvec2_many`).
10851#[cfg(target_arch = "aarch64")]
10852#[allow(clippy::too_many_arguments)]
10853fn q8_range2_sdot(
10854    q: &[u8],
10855    row_scale: &[f32],
10856    a1: &SplitAct,
10857    a2: &SplitAct,
10858    cols: usize,
10859    p1: SendMut,
10860    p2: SendMut,
10861    start: usize,
10862    end: usize,
10863) {
10864    for o in start..end {
10865        let row = &q[o * cols..(o + 1) * cols];
10866        // SAFETY: disjoint row ranges per worker.
10867        unsafe {
10868            *p1.at(o) = row_dot_sdot(row, a1) * row_scale[o];
10869            *p2.at(o) = row_dot_sdot(row, a2) * row_scale[o];
10870        }
10871    }
10872}
10873
10874/// Two-input q8 row range, f32 kernel (non-SDOT) — same extraction.
10875#[allow(clippy::too_many_arguments)]
10876fn q8_range2_f32(
10877    q: &[u8],
10878    row_scale: &[f32],
10879    x1: &[f32],
10880    x2: &[f32],
10881    cols: usize,
10882    p1: SendMut,
10883    p2: SendMut,
10884    start: usize,
10885    end: usize,
10886) {
10887    for o in start..end {
10888        let row = &q[o * cols..(o + 1) * cols];
10889        // SAFETY: disjoint row ranges per worker.
10890        unsafe {
10891            *p1.at(o) = dot_i8_f32(row, x1) * row_scale[o];
10892            *p2.at(o) = dot_i8_f32(row, x2) * row_scale[o];
10893        }
10894    }
10895}
10896
10897/// Scalar/NEON-f32 q8 row range (non-SDOT platforms) — same extraction.
10898fn q8_range_f32(
10899    q: &[u8],
10900    row_scale: &[f32],
10901    xs: &[f32],
10902    cols: usize,
10903    out_addr: SendMut,
10904    start: usize,
10905    end: usize,
10906) {
10907    for o in start..end {
10908        let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
10909        // SAFETY: disjoint row ranges per worker.
10910        unsafe { *out_addr.at(o) = v };
10911    }
10912}
10913
10914/// One q8 row against a split activation, portable: the per-arch fast
10915/// dots where they exist, the exact scalar loop elsewhere. The scalar
10916/// arm is also the test oracle for both fast arms.
10917#[inline]
10918fn q8_row_dot(row: &[u8], act: &SplitAct) -> f32 {
10919    #[cfg(target_arch = "aarch64")]
10920    return row_dot_sdot(row, act);
10921    #[cfg(target_arch = "x86_64")]
10922    return row_dot_avx2(row, act);
10923    #[allow(unreachable_code)]
10924    q8_row_dot_scalar(row, act)
10925}
10926
10927#[allow(dead_code)]
10928fn q8_row_dot_scalar(row: &[u8], act: &SplitAct) -> f32 {
10929    let mut acc = 0i32;
10930    for (k, &b) in row.iter().enumerate() {
10931        acc += (b as i8) as i32 * act.xq[k] as i32;
10932    }
10933    let mut acc = acc as f32 * act.sx;
10934    for &(j, xv) in &act.outliers {
10935        acc += (row[j] as i8) as f32 * xv;
10936    }
10937    acc
10938}
10939
10940/// SDOT row dot with exact outlier correction:
10941/// `dot = sdot(w, xq)·sx + Σ_outl w[j]·x[j]` (then × row_scale by caller).
10942#[cfg(target_arch = "aarch64")]
10943#[inline]
10944fn row_dot_sdot(row: &[u8], act: &SplitAct) -> f32 {
10945    let mut acc = unsafe { dot_i8_sdot(row, &act.xq) } as f32 * act.sx;
10946    for &(j, xv) in &act.outliers {
10947        acc += (row[j] as i8) as f32 * xv;
10948    }
10949    acc
10950}
10951
10952/// One q4 row via SDOT: each 32-group's nibbles unpack to centered i8
10953/// (nib−8 ∈ [−8,7]), int8×int8 `sdot` against the pre-quantized
10954/// activation group, × the group's f16 scale. Returns Σ_g dot_g·s_g;
10955/// the caller multiplies by the activation scale and adds the exact
10956/// outlier terms (port of vmfcore `dot_q4_block_sdot`, +23% measured).
10957/// Nibble order matches the writer: element 2k = low nibble, 2k+1 = high
10958/// → zip(lo,hi) restores flat order.
10959#[cfg(target_arch = "aarch64")]
10960#[target_feature(enable = "neon,dotprod")]
10961unsafe fn dot_q4_row_sdot(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
10962    // SAFETY: callers uphold slice-length contracts (16 packed bytes and
10963    // 2 scale bytes per group; xq.len() == gpr·GROUP_SIZE).
10964    unsafe {
10965        use core::arch::aarch64::*;
10966        use core::arch::asm;
10967        let lomask = vdupq_n_u8(0x0F);
10968        let eight = vdupq_n_s8(8);
10969        let mut acc = 0f32;
10970        for gi in 0..gpr {
10971            let g = g0 + gi;
10972            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
10973            let b = vld1q_u8(packed.as_ptr().add(g * 16));
10974            let lo = vandq_u8(b, lomask);
10975            let hi = vshrq_n_u8::<4>(b);
10976            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
10977            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
10978            let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
10979            let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
10980            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
10981            asm!(
10982                "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
10983                "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
10984                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
10985                e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
10986                options(pure, nomem, nostack),
10987            );
10988            acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
10989        }
10990        acc
10991    }
10992}
10993
10994/// Two-activation q4 row via SDOT: the nibble unpack (the expensive
10995/// part) happens ONCE per group; both pre-quantized activations are
10996/// dotted against the same centered i8 registers. Per-lane math matches
10997/// `dot_q4_row_sdot` exactly.
10998#[cfg(target_arch = "aarch64")]
10999#[target_feature(enable = "neon,dotprod")]
11000unsafe fn dot_q4_row_sdot2(
11001    packed: &[u8],
11002    scales: &[u8],
11003    g0: usize,
11004    gpr: usize,
11005    xq1: &[i8],
11006    xq2: &[i8],
11007) -> (f32, f32) {
11008    // SAFETY: callers uphold slice-length contracts (16 packed bytes and
11009    // 2 scale bytes per group; xq*.len() == gpr·GROUP_SIZE).
11010    unsafe {
11011        use core::arch::aarch64::*;
11012        use core::arch::asm;
11013        let lomask = vdupq_n_u8(0x0F);
11014        let eight = vdupq_n_s8(8);
11015        let (mut acc1, mut acc2) = (0f32, 0f32);
11016        for gi in 0..gpr {
11017            let g = g0 + gi;
11018            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
11019            let b = vld1q_u8(packed.as_ptr().add(g * 16));
11020            let lo = vandq_u8(b, lomask);
11021            let hi = vshrq_n_u8::<4>(b);
11022            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
11023            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
11024            let x10 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE));
11025            let x11 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE + 16));
11026            let x20 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE));
11027            let x21 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE + 16));
11028            let (mut a0, mut a1, mut b0, mut b1) = (
11029                vdupq_n_s32(0),
11030                vdupq_n_s32(0),
11031                vdupq_n_s32(0),
11032                vdupq_n_s32(0),
11033            );
11034            asm!(
11035                "sdot {a0:v}.4s, {e0:v}.16b, {x10:v}.16b",
11036                "sdot {a1:v}.4s, {e1:v}.16b, {x11:v}.16b",
11037                "sdot {b0:v}.4s, {e0:v}.16b, {x20:v}.16b",
11038                "sdot {b1:v}.4s, {e1:v}.16b, {x21:v}.16b",
11039                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
11040                b0 = inout(vreg) b0, b1 = inout(vreg) b1,
11041                e0 = in(vreg) e0, e1 = in(vreg) e1,
11042                x10 = in(vreg) x10, x11 = in(vreg) x11,
11043                x20 = in(vreg) x20, x21 = in(vreg) x21,
11044                options(pure, nomem, nostack),
11045            );
11046            acc1 += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
11047            acc2 += vaddvq_s32(vaddq_s32(b0, b1)) as f32 * s;
11048        }
11049        (acc1, acc2)
11050    }
11051}
11052
11053// ───────────────────── fused int8 kernels ─────────────────────
11054
11055/// `acc += w · row` where the row is centered i8 — NEON widen+fma on
11056/// aarch64, scalar elsewhere. The KV-cache q8 value path rides on this.
11057#[inline]
11058pub(crate) fn axpy_i8_f32(acc: &mut [f32], row: &[i8], w: f32) {
11059    #[cfg(target_arch = "aarch64")]
11060    unsafe {
11061        return axpy_i8_f32_neon(acc, row, w);
11062    }
11063    #[cfg(target_arch = "x86_64")]
11064    if avx2_enabled() {
11065        return unsafe { axpy_i8_f32_avx2(acc, row, w) };
11066    }
11067    #[allow(unreachable_code)]
11068    {
11069        for (a, &b) in acc.iter_mut().zip(row) {
11070            *a += w * b as f32;
11071        }
11072    }
11073}
11074
11075/// i8→f32 axpy via AVX2/FMA (x86 mirror of `axpy_i8_f32_neon`).
11076#[cfg(target_arch = "x86_64")]
11077#[target_feature(enable = "avx2,fma")]
11078unsafe fn axpy_i8_f32_avx2(acc: &mut [f32], row: &[i8], w: f32) {
11079    // SAFETY: callers uphold slice-length contracts (see call sites).
11080    unsafe {
11081        use core::arch::x86_64::*;
11082        let n = acc.len().min(row.len());
11083        let ap = acc.as_mut_ptr();
11084        let rp = row.as_ptr();
11085        let wv = _mm256_set1_ps(w);
11086        let mut j = 0usize;
11087        while j + 16 <= n {
11088            let rb = _mm_loadu_si128(rp.add(j) as *const __m128i);
11089            let lo = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(rb));
11090            let hi = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(rb)));
11091            let v0 = _mm256_fmadd_ps(wv, lo, _mm256_loadu_ps(ap.add(j)));
11092            let v1 = _mm256_fmadd_ps(wv, hi, _mm256_loadu_ps(ap.add(j + 8)));
11093            _mm256_storeu_ps(ap.add(j), v0);
11094            _mm256_storeu_ps(ap.add(j + 8), v1);
11095            j += 16;
11096        }
11097        while j < n {
11098            *ap.add(j) += w * (*rp.add(j)) as f32;
11099            j += 1;
11100        }
11101    }
11102}
11103
11104#[cfg(target_arch = "aarch64")]
11105#[target_feature(enable = "neon")]
11106unsafe fn axpy_i8_f32_neon(acc: &mut [f32], row: &[i8], w: f32) {
11107    // SAFETY: callers uphold slice-length contracts (see call sites).
11108    unsafe {
11109        use core::arch::aarch64::*;
11110        let n = acc.len().min(row.len());
11111        let ap = acc.as_mut_ptr();
11112        let rp = row.as_ptr();
11113        let wv = vdupq_n_f32(w);
11114        let mut j = 0usize;
11115        while j + 16 <= n {
11116            let rb = vld1q_s8(rp.add(j));
11117            let lo = vmovl_s8(vget_low_s8(rb));
11118            let hi = vmovl_s8(vget_high_s8(rb));
11119            for (off, half) in [(0, lo), (8, hi)] {
11120                let f0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(half)));
11121                let f1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(half)));
11122                let o = j + off;
11123                vst1q_f32(ap.add(o), vfmaq_f32(vld1q_f32(ap.add(o)), wv, f0));
11124                vst1q_f32(ap.add(o + 4), vfmaq_f32(vld1q_f32(ap.add(o + 4)), wv, f1));
11125            }
11126            j += 16;
11127        }
11128        while j < n {
11129            *ap.add(j) += w * (*rp.add(j)) as f32;
11130            j += 1;
11131        }
11132    }
11133}
11134
11135/// i8 row · f32 x. NEON on aarch64 (ported from vmfcore `dot_i8_f32_neon`,
11136/// ≈9× scalar), scalar elsewhere.
11137#[inline]
11138pub(crate) fn dot_i8_f32(w: &[u8], x: &[f32]) -> f32 {
11139    #[cfg(target_arch = "aarch64")]
11140    unsafe {
11141        return dot_i8_f32_neon(w, x);
11142    }
11143    #[cfg(target_arch = "x86_64")]
11144    if avx2_enabled() {
11145        return unsafe { dot_i8_f32_avx2(w, x) };
11146    }
11147    #[allow(unreachable_code)]
11148    {
11149        let mut sum = 0.0f32;
11150        for (j, &b) in w.iter().enumerate() {
11151            sum += (b as i8) as f32 * x[j];
11152        }
11153        sum
11154    }
11155}
11156
11157/// i8 row · (x ⊙ col_field) — the q8_2f row dot with the θ col-field
11158/// folded into the product (no prescaled copy of x). NEON on aarch64,
11159/// scalar elsewhere. Used by the active-neuron path `row_dot`.
11160#[inline]
11161fn dot_i8_col_f32(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
11162    #[cfg(target_arch = "aarch64")]
11163    unsafe {
11164        return dot_i8_col_f32_neon(w, x, col);
11165    }
11166    #[allow(unreachable_code)]
11167    {
11168        let mut sum = 0.0f32;
11169        for (j, &b) in w.iter().enumerate() {
11170            sum += (b as i8) as f32 * x[j] * col[j];
11171        }
11172        sum
11173    }
11174}
11175
11176#[cfg(target_arch = "aarch64")]
11177#[target_feature(enable = "neon")]
11178unsafe fn dot_i8_col_f32_neon(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
11179    // SAFETY: callers uphold slice-length contracts (see call sites).
11180    unsafe {
11181        use core::arch::aarch64::*;
11182        let n = x.len();
11183        let wp = w.as_ptr() as *const i8;
11184        let xp = x.as_ptr();
11185        let cp = col.as_ptr();
11186        let (mut a0, mut a1, mut a2, mut a3) = (
11187            vdupq_n_f32(0.0),
11188            vdupq_n_f32(0.0),
11189            vdupq_n_f32(0.0),
11190            vdupq_n_f32(0.0),
11191        );
11192        let mut j = 0usize;
11193        while j + 16 <= n {
11194            let wb = vld1q_s8(wp.add(j));
11195            let lo = vmovl_s8(vget_low_s8(wb));
11196            let hi = vmovl_s8(vget_high_s8(wb));
11197            let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
11198            let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
11199            let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
11200            let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
11201            a0 = vfmaq_f32(
11202                a0,
11203                w0,
11204                vmulq_f32(vld1q_f32(xp.add(j)), vld1q_f32(cp.add(j))),
11205            );
11206            a1 = vfmaq_f32(
11207                a1,
11208                w1,
11209                vmulq_f32(vld1q_f32(xp.add(j + 4)), vld1q_f32(cp.add(j + 4))),
11210            );
11211            a2 = vfmaq_f32(
11212                a2,
11213                w2,
11214                vmulq_f32(vld1q_f32(xp.add(j + 8)), vld1q_f32(cp.add(j + 8))),
11215            );
11216            a3 = vfmaq_f32(
11217                a3,
11218                w3,
11219                vmulq_f32(vld1q_f32(xp.add(j + 12)), vld1q_f32(cp.add(j + 12))),
11220            );
11221            j += 16;
11222        }
11223        let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
11224        while j < n {
11225            sum += (*wp.add(j)) as f32 * *xp.add(j) * *cp.add(j);
11226            j += 1;
11227        }
11228        sum
11229    }
11230}
11231
11232#[cfg(target_arch = "aarch64")]
11233#[target_feature(enable = "neon")]
11234unsafe fn dot_i8_f32_neon(w: &[u8], x: &[f32]) -> f32 {
11235    // SAFETY: callers uphold slice-length contracts (see call sites).
11236    unsafe {
11237        use core::arch::aarch64::*;
11238        let n = x.len();
11239        let wp = w.as_ptr() as *const i8;
11240        let xp = x.as_ptr();
11241        let (mut a0, mut a1, mut a2, mut a3) = (
11242            vdupq_n_f32(0.0),
11243            vdupq_n_f32(0.0),
11244            vdupq_n_f32(0.0),
11245            vdupq_n_f32(0.0),
11246        );
11247        let mut j = 0usize;
11248        while j + 16 <= n {
11249            let wb = vld1q_s8(wp.add(j));
11250            let lo = vmovl_s8(vget_low_s8(wb));
11251            let hi = vmovl_s8(vget_high_s8(wb));
11252            let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
11253            let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
11254            let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
11255            let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
11256            a0 = vfmaq_f32(a0, w0, vld1q_f32(xp.add(j)));
11257            a1 = vfmaq_f32(a1, w1, vld1q_f32(xp.add(j + 4)));
11258            a2 = vfmaq_f32(a2, w2, vld1q_f32(xp.add(j + 8)));
11259            a3 = vfmaq_f32(a3, w3, vld1q_f32(xp.add(j + 12)));
11260            j += 16;
11261        }
11262        let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
11263        while j < n {
11264            sum += (*wp.add(j)) as f32 * *xp.add(j);
11265            j += 1;
11266        }
11267        sum
11268    }
11269}
11270
11271#[allow(clippy::too_many_arguments)]
11272fn qmatvec(
11273    q: &[u8],
11274    rep: &[u8],
11275    row_scale: &[f32],
11276    x: &[f32],
11277    col_field: &[f32],
11278    dtype: TensorDtype,
11279    rows: usize,
11280    cols: usize,
11281    out: &mut [f32],
11282    pool: Option<&Pool>,
11283) {
11284    debug_assert_eq!(out.len(), rows);
11285    #[cfg(not(target_arch = "aarch64"))]
11286    let _ = rep;
11287
11288    #[cfg(target_arch = "aarch64")]
11289    if sdot_enabled() {
11290        let act = if dtype == TensorDtype::Q8_2f {
11291            split_act_q8_2f(x, col_field)
11292        } else {
11293            split_act(x)
11294        };
11295        let out_addr = SendMut(out.as_mut_ptr());
11296        let run_range = |start: usize, end: usize| {
11297            q8_range_sdot(q, rep, row_scale, &act, cols, out_addr, start, end)
11298        };
11299        match pool {
11300            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
11301            _ => run_range(0, rows),
11302        }
11303        return;
11304    }
11305    // x86 A8W8 via AVX2 maddubs — same quantized-activation contract as
11306    // the SDOT path (CMF_AVX2=0 keeps the exact i8×f32 loop).
11307    #[cfg(target_arch = "x86_64")]
11308    if avx2_a8w8_enabled() {
11309        let act = if dtype == TensorDtype::Q8_2f {
11310            split_act_q8_2f(x, col_field)
11311        } else {
11312            split_act(x)
11313        };
11314        let out_addr = SendMut(out.as_mut_ptr());
11315        let run_range = |start: usize, end: usize| {
11316            q8_range_avx2(q, row_scale, &act, cols, out_addr, start, end)
11317        };
11318        match pool {
11319            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
11320            _ => run_range(0, rows),
11321        }
11322        return;
11323    }
11324
11325    prescale_with(x, col_field, dtype, 1, |xs| {
11326        let out_addr = SendMut(out.as_mut_ptr());
11327        let run_range = move |start: usize, end: usize| {
11328            for o in start..end {
11329                let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
11330                // SAFETY: disjoint row ranges per worker.
11331                unsafe { *out_addr.at(o) = v };
11332            }
11333        };
11334        match pool {
11335            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
11336            _ => run_range(0, rows),
11337        }
11338    });
11339}
11340
11341#[allow(clippy::too_many_arguments)]
11342fn qmatvec2(
11343    q: &[u8],
11344    row_scale: &[f32],
11345    x1: &[f32],
11346    x2: &[f32],
11347    col_field: &[f32],
11348    dtype: TensorDtype,
11349    rows: usize,
11350    cols: usize,
11351    o1: &mut [f32],
11352    o2: &mut [f32],
11353    pool: Option<&Pool>,
11354) {
11355    #[cfg(target_arch = "aarch64")]
11356    if sdot_enabled() {
11357        let a1s = if dtype == TensorDtype::Q8_2f {
11358            split_act_q8_2f(x1, col_field)
11359        } else {
11360            split_act(x1)
11361        };
11362        let a2s = if dtype == TensorDtype::Q8_2f {
11363            split_act_q8_2f(x2, col_field)
11364        } else {
11365            split_act(x2)
11366        };
11367        let p1 = SendMut(o1.as_mut_ptr());
11368        let p2 = SendMut(o2.as_mut_ptr());
11369        let run_range = |start: usize, end: usize| {
11370            q8_range2_sdot(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
11371        };
11372        match pool {
11373            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
11374            _ => run_range(0, rows),
11375        }
11376        return;
11377    }
11378    #[cfg(target_arch = "x86_64")]
11379    if avx2_a8w8_enabled() {
11380        let a1s = if dtype == TensorDtype::Q8_2f {
11381            split_act_q8_2f(x1, col_field)
11382        } else {
11383            split_act(x1)
11384        };
11385        let a2s = if dtype == TensorDtype::Q8_2f {
11386            split_act_q8_2f(x2, col_field)
11387        } else {
11388            split_act(x2)
11389        };
11390        let p1 = SendMut(o1.as_mut_ptr());
11391        let p2 = SendMut(o2.as_mut_ptr());
11392        let run_range = |start: usize, end: usize| {
11393            q8_range2_avx2(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
11394        };
11395        match pool {
11396            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
11397            _ => run_range(0, rows),
11398        }
11399        return;
11400    }
11401
11402    prescale_with(x1, col_field, dtype, 1, |x1s| {
11403        prescale_with(x2, col_field, dtype, 2, |x2s| {
11404            let p1 = SendMut(o1.as_mut_ptr());
11405            let p2 = SendMut(o2.as_mut_ptr());
11406            let run_range = move |start: usize, end: usize| {
11407                for o in start..end {
11408                    let row = &q[o * cols..(o + 1) * cols];
11409                    let s1 = dot_i8_f32(row, x1s) * row_scale[o];
11410                    let s2 = dot_i8_f32(row, x2s) * row_scale[o];
11411                    // SAFETY: disjoint row ranges per worker.
11412                    unsafe {
11413                        *p1.at(o) = s1;
11414                        *p2.at(o) = s2;
11415                    }
11416                }
11417            };
11418            match pool {
11419                Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
11420                _ => run_range(0, rows),
11421            }
11422        });
11423    });
11424}
11425
11426#[derive(Clone, Copy)]
11427struct SendMut(*mut f32);
11428unsafe impl Send for SendMut {}
11429unsafe impl Sync for SendMut {}
11430
11431impl SendMut {
11432    #[inline]
11433    fn at(self, i: usize) -> *mut f32 {
11434        unsafe { self.0.add(i) }
11435    }
11436}
11437
11438#[cfg(test)]
11439mod tests {
11440    /// `q8_round` must be `round().clamp(±127) as i8` bit for bit: every
11441    /// half-integer, their neighbours one ulp either side, the clamp
11442    /// boundary, huge values, infinities and NaN, plus a dense sweep.
11443    #[test]
11444    fn q8_round_is_round_clamp() {
11445        let reference = |t: f32| t.round().clamp(-127.0, 127.0) as i8;
11446        let mut probe = vec![
11447            0.0f32,
11448            -0.0,
11449            f32::NAN,
11450            f32::INFINITY,
11451            f32::NEG_INFINITY,
11452            f32::MAX,
11453            f32::MIN,
11454            1e30,
11455            -1e30,
11456            f32::MIN_POSITIVE,
11457            -f32::MIN_POSITIVE,
11458        ];
11459        for k in -300i32..=300 {
11460            let h = k as f32 * 0.5;
11461            let up = f32::from_bits(h.to_bits() + 1);
11462            let down = f32::from_bits(h.to_bits().wrapping_sub(1));
11463            for t in [h, up, down] {
11464                probe.push(t);
11465                probe.push(-t);
11466            }
11467        }
11468        let mut t = -140.0f32;
11469        while t < 140.0 {
11470            probe.push(t);
11471            t += 0.000_731;
11472        }
11473        for t in probe {
11474            assert_eq!(q8_round(t), reference(t), "t = {t:e} ({:#x})", t.to_bits());
11475        }
11476    }
11477
11478    use super::*;
11479
11480    /// The f32 `matmat` walks the batch four positions a row at a time; every
11481    /// output must still be the scalar-order sum `matvec` gives, bit for bit
11482    /// (a 7-row batch covers the four-wide block and the remainder).
11483    #[test]
11484    fn f32_matmat_equals_matvec_bitwise() {
11485        let (rows, cols, b) = (5usize, 37usize, 7usize);
11486        let w: Vec<f32> = (0..rows * cols).map(|i| ((i * 7919 % 113) as f32 - 56.0) / 37.0).collect();
11487        let xs: Vec<f32> = (0..b * cols).map(|i| ((i * 104_729 % 97) as f32 - 48.0) / 29.0).collect();
11488        let t = QTensor::from_f32(w, rows, cols);
11489        let mut out = vec![0f32; b * rows];
11490        t.matmat(&xs, b, &mut out, None);
11491        for bi in 0..b {
11492            let mut one = vec![0f32; rows];
11493            t.matvec(&xs[bi * cols..(bi + 1) * cols], &mut one, None);
11494            for o in 0..rows {
11495                assert_eq!(out[bi * rows + o].to_bits(), one[o].to_bits(), "bi {bi} row {o}");
11496            }
11497        }
11498    }
11499
11500    #[test]
11501    fn q2tp_i8_dot_matches_exact_on_grid() {
11502        // On-grid activations (±1 → sx=1/127, xq=±127 dequantizes
11503        // exactly, no outliers) must make the integer path agree with
11504        // the exact scalar walk to f32 rounding.
11505        let (rows, cols) = (5, 64);
11506        let gpr = cols / GROUP_SIZE;
11507        // Synthetic codes plane + a flat ladder: scales_into is not under
11508        // test here, so drive dot_q2tp_row_i8 / q2tp_row_exact directly
11509        // with hand-made scales.
11510        let chunks: Vec<u8> = (0..rows * gpr * Q2TP_CHUNK)
11511            .map(|i| (i as u32).wrapping_mul(2654435761) as u8)
11512            .collect();
11513        let scales: Vec<f32> = (0..gpr).map(|g| 0.5 + g as f32 * 0.25).collect();
11514        let x: Vec<f32> = (0..cols)
11515            .map(|i| if i % 3 == 0 { -1.0 } else { 1.0 })
11516            .collect();
11517        let act = split_act(&x);
11518        assert!(
11519            act.outliers.is_empty(),
11520            "on-grid input must have no outliers"
11521        );
11522        let gsum = q1_group_sums(&act.xq, gpr);
11523        for r in 0..rows {
11524            let exact = q2tp_row_exact(&chunks, r, gpr, &x, &scales);
11525            let fast = dot_q2tp_row_i8(&chunks, r, gpr, &act.xq, &gsum, &scales) * act.sx;
11526            assert!(
11527                (exact - fast).abs() <= exact.abs() * 1e-5 + 1e-5,
11528                "row {r}: exact {exact} vs i8 {fast}"
11529            );
11530        }
11531    }
11532
11533    #[cfg(target_arch = "x86_64")]
11534    #[test]
11535    fn q2tp_avx2_dot_matches_scalar_for_random_patterns() {
11536        // Compare the release AVX2 integer dot against the scalar oracle over
11537        // arbitrary packed bytes/activation signs.  This guards the exact
11538        // table-load path used after rejecting a faster-looking decoder whose
11539        // full-checkpoint greedy output drifted.
11540        if !std::arch::is_x86_feature_detected!("avx2") {
11541            return;
11542        }
11543        let mut seed = 0x9e3779b9u32;
11544        let mut next = || {
11545            seed = seed.wrapping_mul(1664525).wrapping_add(1013904223);
11546            seed
11547        };
11548        for _ in 0..20_000 {
11549            let mut ch = [0u8; Q2TP_CHUNK];
11550            let mut x = [0i8; GROUP_SIZE];
11551            for b in &mut ch {
11552                *b = next() as u8;
11553            }
11554            for v in &mut x {
11555                *v = (next() >> 24) as i8;
11556            }
11557            let mut reference = 0i32;
11558            for (k, &b) in ch.iter().enumerate() {
11559                reference += (b & 3) as i32 * x[k * 4] as i32;
11560                reference += ((b >> 2) & 3) as i32 * x[k * 4 + 1] as i32;
11561                reference += ((b >> 4) & 3) as i32 * x[k * 4 + 2] as i32;
11562                reference += ((b >> 6) & 3) as i32 * x[k * 4 + 3] as i32;
11563            }
11564            // SAFETY: guarded by the runtime AVX2 feature check and fixed
11565            // 8-byte/32-byte slice lengths above.
11566            let got = unsafe { q2tp_code_dot_avx2(&ch, &x) };
11567            assert_eq!(got, reference, "packed q2 lane mismatch");
11568        }
11569    }
11570
11571    #[test]
11572    fn q2tp_affine_fuses_half_scale_correction_without_changing_raw_decode() {
11573        let (rows, cols) = (1usize, GROUP_SIZE);
11574        let mut bytes = vec![0u8; Q2TP_CHUNK + 4 + 1];
11575        // Repeating symbols 0,1,2,0 at unit scale.  q2tp's raw B is
11576        // (c-1.5), while the affine Prism operator is (c-1.0).
11577        bytes[..Q2TP_CHUNK].fill(0x24); // codes 0,1,2,0 in LSB-first order
11578        bytes[Q2TP_CHUNK..Q2TP_CHUNK + 2].copy_from_slice(&0u16.to_le_bytes());
11579        bytes[Q2TP_CHUNK + 2..Q2TP_CHUNK + 4].copy_from_slice(&0u16.to_le_bytes());
11580        bytes[Q2TP_CHUNK + 4] = 1; // dtype16 rung 1 = 1.0
11581        let x = vec![1.0f32; cols];
11582        let mut raw = vec![0.0f32; rows];
11583        let mut affine = vec![0.0f32; rows];
11584        q2tp_matvec_for_test(&bytes, &x, rows, cols, &mut raw);
11585        q2tp_affine_matvec_for_test(&bytes, &x, rows, cols, &mut affine);
11586        assert_eq!(raw, vec![-24.0]);
11587        assert_eq!(affine, vec![-8.0]);
11588        assert!((affine[0] - (raw[0] + 0.5 * cols as f32)).abs() < 1e-6);
11589    }
11590
11591    #[test]
11592    fn q8_row_dot_fast_matches_scalar() {
11593        // The per-arch fast dot must agree with the exact scalar oracle
11594        // (same contract the fused q8 FFN arm rides on).
11595        let cols = 96;
11596        let row: Vec<u8> = (0..cols)
11597            .map(|i| ((i * 37 % 251) - 125) as i8 as u8)
11598            .collect();
11599        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.13).sin()).collect();
11600        let act = split_act(&x);
11601        let fast = q8_row_dot(&row, &act);
11602        let scalar = q8_row_dot_scalar(&row, &act);
11603        assert!(
11604            (fast - scalar).abs() <= scalar.abs() * 1e-5 + 1e-5,
11605            "fast {fast} vs scalar {scalar}"
11606        );
11607    }
11608
11609    #[test]
11610    fn f32_matvec_matches_matvec_rows_bitexact() {
11611        let (rows, cols) = (300, 40);
11612        let w: Vec<f32> = (0..rows * cols).map(|i| (i as f32 * 0.017).sin()).collect();
11613        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.05).cos()).collect();
11614        let qt = QTensor::from_f32(w.clone(), rows, cols);
11615
11616        let mut a = vec![0.0f32; rows];
11617        matvec_rows(None, &w, &x, &mut a);
11618        let mut b = vec![0.0f32; rows];
11619        qt.matvec(&x, &mut b, None);
11620        assert_eq!(a, b);
11621    }
11622
11623    #[test]
11624    fn sdot_kernel_exact_on_grid() {
11625        // Activations already on the i8 grid (±1 with amax=1 → sx=1/127,
11626        // xq=±127 dequantizes EXACTLY) → the SDOT path must match the
11627        // exact f32 dot to float rounding. This isolates kernel
11628        // correctness from quantization noise.
11629        eprintln!("sdot_enabled = {}", sdot_enabled());
11630        let (rows, cols) = (9, 80); // odd rows → exercises 4-row + tail
11631        let w: Vec<u8> = (0..rows * cols)
11632            .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
11633            .collect();
11634        let scales: Vec<f32> = (0..rows).map(|o| 0.005 + o as f32 * 0.001).collect();
11635        let x: Vec<f32> = (0..cols)
11636            .map(|i| match i % 3 {
11637                0 => 1.0,
11638                1 => -1.0,
11639                _ => 0.0,
11640            })
11641            .collect();
11642        let mut a = vec![0.0f32; rows];
11643        qmatvec(
11644            &w,
11645            &[],
11646            &scales,
11647            &x,
11648            &[],
11649            TensorDtype::Q8Row,
11650            rows,
11651            cols,
11652            &mut a,
11653            None,
11654        );
11655        for o in 0..rows {
11656            let mut acc = 0.0f32;
11657            for j in 0..cols {
11658                acc += (w[o * cols + j] as i8) as f32 * x[j];
11659            }
11660            let expect = acc * scales[o];
11661            assert!(
11662                (a[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
11663                "row {o}: {} vs {expect}",
11664                a[o]
11665            );
11666        }
11667    }
11668
11669    #[test]
11670    fn q1_tbl_fast_path_matches_reference() {
11671        // gpr = 8 exercises the TBL pair-load fast loop, and the LAST
11672        // row's final 4-tile window trips the 4B-overread guard (the
11673        // payload ends exactly at the last tile) — both paths must
11674        // agree with the dequant reference.
11675        let (rows, cols) = (5, 256);
11676        let gpr = cols / GROUP_SIZE;
11677        let mut bytes = Vec::new();
11678        for t in 0..rows * gpr {
11679            let s = 0.007 + (t % 11) as f32 * 0.004;
11680            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
11681            for j in 0..4 {
11682                bytes.push(((t * 53 + j * 89 + 7) % 249) as u8);
11683            }
11684        }
11685        let x: Vec<f32> = (0..cols)
11686            .map(|i| if (i * 5) % 7 < 3 { 1.0 } else { -1.0 })
11687            .collect();
11688        let mut w = vec![0.0f32; rows * cols];
11689        cortiq_core::quant::dequant_q1(&bytes, &mut w);
11690        let mut got = vec![0.0f32; rows];
11691        q1_matvec(&bytes, &x, rows, cols, &mut got, None);
11692        for o in 0..rows {
11693            let expect: f32 = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
11694            assert!(
11695                (got[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
11696                "row {o}: {} vs {expect}",
11697                got[o]
11698            );
11699        }
11700        // Blocked 1×4 batch (b=5: one quad + remainder) must equal the
11701        // single-matvec path bit-for-bit.
11702        let b = 5usize;
11703        let mut xs_all = Vec::new();
11704        for bi in 0..b {
11705            xs_all.extend(x.iter().map(|v| if bi % 2 == 0 { *v } else { -*v }));
11706        }
11707        let mut mm = vec![0.0f32; b * rows];
11708        q1_matmat(&bytes, &xs_all, b, rows, cols, &mut mm, None);
11709        for bi in 0..b {
11710            let mut single = vec![0.0f32; rows];
11711            q1_matvec(
11712                &bytes,
11713                &xs_all[bi * cols..(bi + 1) * cols],
11714                rows,
11715                cols,
11716                &mut single,
11717                None,
11718            );
11719            assert_eq!(&mm[bi * rows..(bi + 1) * rows], &single[..], "stream {bi}");
11720        }
11721    }
11722
11723    #[test]
11724    fn q1_kernels_match_exact_reference() {
11725        // Synthetic q1 payload: 6-byte tiles [f16 scale][4B bits].
11726        let (rows, cols) = (7, 96);
11727        let gpr = cols / GROUP_SIZE;
11728        let mut bytes = Vec::new();
11729        for t in 0..rows * gpr {
11730            let s = 0.01 + (t % 13) as f32 * 0.003;
11731            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
11732            for j in 0..4 {
11733                bytes.push(((t * 31 + j * 97) % 251) as u8);
11734            }
11735        }
11736        // On-grid activations (±1, amax 1) → the SDOT path is exact.
11737        let x: Vec<f32> = (0..cols)
11738            .map(|i| if i % 3 == 0 { 1.0 } else { -1.0 })
11739            .collect();
11740        // Reference through the core dequant.
11741        let mut w = vec![0.0f32; rows * cols];
11742        cortiq_core::quant::dequant_q1(&bytes, &mut w);
11743        let mut expect = vec![0.0f32; rows];
11744        for o in 0..rows {
11745            expect[o] = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
11746        }
11747        let mut got = vec![0.0f32; rows];
11748        q1_matvec(&bytes, &x, rows, cols, &mut got, None);
11749        for o in 0..rows {
11750            assert!(
11751                (got[o] - expect[o]).abs() < 1e-3 * expect[o].abs().max(1e-3),
11752                "row {o}: {} vs {}",
11753                got[o],
11754                expect[o]
11755            );
11756        }
11757        // Pair and batch paths agree with the single path.
11758        let x2: Vec<f32> = x.iter().map(|v| -v).collect();
11759        let (mut a1, mut a2) = (vec![0.0f32; rows], vec![0.0f32; rows]);
11760        q1_matvec2(&bytes, &x, &x2, rows, cols, &mut a1, &mut a2, None);
11761        assert_eq!(a1, got);
11762        let mut xs = x.clone();
11763        xs.extend_from_slice(&x2);
11764        let mut mm = vec![0.0f32; 2 * rows];
11765        q1_matmat(&bytes, &xs, 2, rows, cols, &mut mm, None);
11766        assert_eq!(&mm[..rows], got.as_slice());
11767        assert_eq!(&mm[rows..], a2.as_slice());
11768    }
11769
11770    #[test]
11771    fn repack_is_bit_identical() {
11772        // The interleaved-repack kernel must produce EXACTLY the same
11773        // bits as the mmap-layout kernel: integer accumulation is order-
11774        // exact, the f32 epilogue is identical. Odd rows exercise the
11775        // tail; direct range calls exercise unaligned pool splits.
11776        let (rows, cols) = (267, 96); // 66 groups + 3 tail rows, cols % 16 == 0
11777        let w: Vec<u8> = (0..rows * cols)
11778            .map(|i| (((i * 89) % 253) as i32 - 126) as i8 as u8)
11779            .collect();
11780        let scales: Vec<f32> = (0..rows).map(|o| 0.003 + o as f32 * 0.0007).collect();
11781        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.37).sin() * 2.0).collect();
11782        let rep = q8_repack_layout(&w, rows, cols);
11783        // Group interleave round-trips.
11784        for g in 0..rows / 4 {
11785            for c in 0..cols / 16 {
11786                for lane in 0..4 {
11787                    assert_eq!(
11788                        &rep[g * 4 * cols + c * 64 + lane * 16
11789                            ..g * 4 * cols + c * 64 + lane * 16 + 16],
11790                        &w[(g * 4 + lane) * cols + c * 16..(g * 4 + lane) * cols + c * 16 + 16],
11791                    );
11792                }
11793            }
11794        }
11795        let mut a = vec![0.0f32; rows];
11796        qmatvec(
11797            &w,
11798            &[],
11799            &scales,
11800            &x,
11801            &[],
11802            TensorDtype::Q8Row,
11803            rows,
11804            cols,
11805            &mut a,
11806            None,
11807        );
11808        let mut b = vec![0.0f32; rows];
11809        qmatvec(
11810            &w,
11811            &rep,
11812            &scales,
11813            &x,
11814            &[],
11815            TensorDtype::Q8Row,
11816            rows,
11817            cols,
11818            &mut b,
11819            None,
11820        );
11821        assert_eq!(a, b, "full-range repack output diverged");
11822
11823        #[cfg(target_arch = "aarch64")]
11824        if sdot_enabled() {
11825            // Unaligned range split (pool workers get arbitrary bounds).
11826            let act = split_act(&x);
11827            let mut c1 = vec![0.0f32; rows];
11828            let mut c2 = vec![0.0f32; rows];
11829            q8_range_sdot(
11830                &w,
11831                &[],
11832                &scales,
11833                &act,
11834                cols,
11835                SendMut(c1.as_mut_ptr()),
11836                3,
11837                rows - 2,
11838            );
11839            q8_range_sdot(
11840                &w,
11841                &rep,
11842                &scales,
11843                &act,
11844                cols,
11845                SendMut(c2.as_mut_ptr()),
11846                3,
11847                rows - 2,
11848            );
11849            assert_eq!(c1, c2, "unaligned-range repack output diverged");
11850        }
11851    }
11852
11853    #[test]
11854    fn sdot_a8w8_noise_is_bounded() {
11855        // Off-grid activations: A8 quantization noise must stay small in
11856        // relative L2 over the whole output (realistic accuracy contract;
11857        // vmfcore measured argmax-identical decode on real models).
11858        let (rows, cols) = (16, 512);
11859        let w: Vec<u8> = (0..rows * cols)
11860            .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
11861            .collect();
11862        let scales = vec![0.01f32; rows];
11863        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.21).sin()).collect();
11864        let mut a = vec![0.0f32; rows];
11865        qmatvec(
11866            &w,
11867            &[],
11868            &scales,
11869            &x,
11870            &[],
11871            TensorDtype::Q8Row,
11872            rows,
11873            cols,
11874            &mut a,
11875            None,
11876        );
11877        let (mut num, mut den) = (0f64, 0f64);
11878        for o in 0..rows {
11879            let mut acc = 0.0f32;
11880            for j in 0..cols {
11881                acc += (w[o * cols + j] as i8) as f32 * x[j];
11882            }
11883            let expect = acc * scales[o];
11884            num += ((a[o] - expect) as f64).powi(2);
11885            den += (expect as f64).powi(2);
11886        }
11887        let rel = (num / den.max(1e-12)).sqrt();
11888        assert!(rel < 0.05, "A8W8 relative L2 error too high: {rel}");
11889    }
11890
11891    #[test]
11892    fn i8_dot_neon_matches_scalar() {
11893        let n = 100;
11894        let w: Vec<u8> = (0..n).map(|i| ((i * 37 + 11) % 251) as u8).collect();
11895        let x: Vec<f32> = (0..n).map(|i| (i as f32 * 0.13).sin()).collect();
11896        let mut scalar = 0.0f32;
11897        for j in 0..n {
11898            scalar += (w[j] as i8) as f32 * x[j];
11899        }
11900        let fast = dot_i8_f32(&w, &x);
11901        assert!((scalar - fast).abs() < 1e-3 * scalar.abs().max(1.0));
11902    }
11903
11904    /// Fused vbit matvec must match full dequant_vbit + dense matvec.
11905    #[test]
11906    fn vbitmatvec_matches_full_dequant() {
11907        let (rows, cols) = (6, 64);
11908        let ng = cols / GROUP_SIZE;
11909        // Hand-craft: bits per row, f16 scales, packed rows.
11910        let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
11911        let mut bytes = bits.clone();
11912        for g in 0..rows * ng {
11913            let s = 0.02 + 0.001 * g as f32;
11914            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
11915        }
11916        for r in 0..rows {
11917            let b = bits[r] as usize;
11918            let (mut acc, mut nb) = (0u64, 0usize);
11919            let mut rowbytes = Vec::new();
11920            for i in 0..cols {
11921                let v = ((i * 7 + r * 13) % (1 << b)) as u64;
11922                acc = (acc << b) | v;
11923                nb += b;
11924                while nb >= 8 {
11925                    nb -= 8;
11926                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
11927                }
11928            }
11929            if nb > 0 {
11930                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
11931            }
11932            bytes.extend_from_slice(&rowbytes);
11933        }
11934        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
11935
11936        let mut reference = vec![0f32; rows * cols];
11937        cortiq_core::quant::dequant_vbit(&bytes, rows, cols, &mut reference).unwrap();
11938        let mut expect = vec![0f32; rows];
11939        for r in 0..rows {
11940            expect[r] = reference[r * cols..(r + 1) * cols]
11941                .iter()
11942                .zip(&x)
11943                .map(|(w, xv)| w * xv)
11944                .sum();
11945        }
11946        let mut got = vec![0f32; rows];
11947        let offsets = vbit_row_offsets(&bytes, rows, cols);
11948        vbitmatvec(&bytes, &offsets, &x, rows, cols, &mut got, None);
11949        // SDOT path quantizes activations to i8 (A8W8): bounded noise,
11950        // same contract as q8 (exact path is pinned by CMF_SDOT=0 in
11951        // the golden-parity gate).
11952        let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
11953        let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
11954        for r in 0..rows {
11955            assert!(
11956                (got[r] - expect[r]).abs() < tol * scale,
11957                "row {r}: {} vs {}",
11958                got[r],
11959                expect[r]
11960            );
11961        }
11962    }
11963
11964    /// Fused q4 matvec must match the reference full-dequant + dense
11965    /// matvec bit-for-bit in structure (same f32 math, group order).
11966    /// vbit matmat: the blocked 1×4 leg must match the per-row path
11967    /// (paired env toggle; larger shape so both code paths engage).
11968    #[test]
11969    #[cfg(target_arch = "x86_64")]
11970    fn vbit_matmat_blocked_matches_per_row() {
11971        let (rows, cols, b) = (64usize, 128usize, 9usize);
11972        let ng = cols / GROUP_SIZE;
11973        let bits: Vec<u8> = (0..rows).map(|r| [3u8, 4, 5, 6][r % 4]).collect();
11974        let mut bytes = bits.clone();
11975        for g in 0..rows * ng {
11976            let sc = 0.02 + 0.0005 * g as f32;
11977            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
11978        }
11979        for r in 0..rows {
11980            let bw = bits[r] as usize;
11981            let (mut acc, mut nb) = (0u64, 0usize);
11982            let mut rowbytes = Vec::new();
11983            for i in 0..cols {
11984                let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
11985                acc = (acc << bw) | v;
11986                nb += bw;
11987                while nb >= 8 {
11988                    nb -= 8;
11989                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
11990                }
11991            }
11992            if nb > 0 {
11993                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
11994            }
11995            bytes.extend_from_slice(&rowbytes);
11996        }
11997        let x: Vec<f32> = (0..b * cols)
11998            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
11999            .collect();
12000        let offsets = vbit_row_offsets(&bytes, rows, cols);
12001        let mut y_a = vec![0f32; b * rows];
12002        let mut y_b = vec![0f32; b * rows];
12003        unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
12004        vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_a, None);
12005        unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
12006        vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_b, None);
12007        unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
12008        let max_d = y_a
12009            .iter()
12010            .zip(&y_b)
12011            .map(|(p, q)| (p - q).abs())
12012            .fold(0.0f32, f32::max);
12013        assert!(max_d < 1e-4, "vbit blocked ≠ per-row: max|Δ| = {max_d}");
12014    }
12015
12016    /// q4t blocked 1×4 (SDOT on ARM, AVX2 on x86) must equal the
12017    /// per-row path exactly: same nibble unpack, same group order,
12018    /// same f32 accumulation — batch == matvec bit-for-bit. b=9 covers
12019    /// two full 1×4 blocks plus a remainder through the single-row
12020    /// kernel. (Both paths produce identical output, so the shared
12021    /// CMF_X86_BLOCKED env var racing with other tests cannot flip
12022    /// the verdict — worst case both sides take the same path.)
12023    #[test]
12024    fn q4t_matmat_blocked_matches_per_row() {
12025        let (rows, cols, b) = (16usize, 64usize, 9usize);
12026        let gpr = cols / GROUP_SIZE;
12027        let mut bytes = vec![0u8; rows * gpr * Q4_TILE];
12028        for r in 0..rows {
12029            for g in 0..gpr {
12030                let t = (r * gpr + g) * Q4_TILE;
12031                let sc = 0.02 + 0.001 * (r * gpr + g) as f32;
12032                bytes[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
12033                for k in 0..16 {
12034                    bytes[t + 2 + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
12035                }
12036            }
12037        }
12038        let x: Vec<f32> = (0..b * cols)
12039            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
12040            .collect();
12041        let mut y_blk = vec![0f32; b * rows];
12042        let mut y_row = vec![0f32; b * rows];
12043        unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
12044        q4t_matmat(&bytes, &x, b, rows, cols, &mut y_blk, None);
12045        unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
12046        q4t_matmat(&bytes, &x, b, rows, cols, &mut y_row, None);
12047        unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
12048        assert_eq!(y_blk, y_row, "q4t blocked 1x4 ≠ per-row");
12049    }
12050
12051    /// The wide-batch Accelerate arm of q4t_matmat vs a brute-force
12052    /// f32 dequant matmul: both are f32 GEMMs, so only reduction
12053    /// order differs — tight tolerance.
12054    /// A synthetic q4tp payload: random nibbles plus a per-row ladder whose
12055    /// span varies row to row, so the codes actually exercise the full 0..31
12056    /// range rather than clustering on one rung.
12057    fn synth_q4tp(rows: usize, cols: usize) -> Vec<u8> {
12058        use cortiq_core::quant::{f32_to_f16, q4tp_code_stride, q4tp_put_code};
12059        let gpr = cols / GROUP_SIZE;
12060        let stride = q4tp_code_stride(gpr);
12061        let (params_off, codes_off, _) = q4tp_sections(rows, cols);
12062        let mut b = vec![0u8; codes_off + rows * stride];
12063        for r in 0..rows {
12064            for g in 0..gpr {
12065                let t = (r * gpr + g) * Q4TP_NIB;
12066                for k in 0..16 {
12067                    b[t + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
12068                }
12069            }
12070            let lo = -6.0 - 0.03 * (r % 17) as f32;
12071            let step = 0.01 + 0.004 * (r % 11) as f32;
12072            let p = params_off + r * 4;
12073            b[p..p + 2].copy_from_slice(&f32_to_f16(lo).to_le_bytes());
12074            b[p + 2..p + 4].copy_from_slice(&f32_to_f16(step).to_le_bytes());
12075            let crow = &mut b[codes_off + r * stride..codes_off + (r + 1) * stride];
12076            for g in 0..gpr {
12077                q4tp_put_code(crow, g, (r * 5 + g * 3) % 32);
12078            }
12079        }
12080        b
12081    }
12082
12083    /// The same weights re-expressed as q4_tiled, so the proven kernel can
12084    /// be the reference: each tile stores the ladder scale its code selects.
12085    /// Only the f16 rounding of that scale separates the two payloads.
12086    fn q4tp_as_q4t(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
12087        let gpr = cols / GROUP_SIZE;
12088        let v = Q4tpView::new(bytes, rows, cols);
12089        let mut out = vec![0u8; rows * gpr * Q4_TILE];
12090        let mut sc = vec![0f32; gpr];
12091        for r in 0..rows {
12092            v.scales_into(r, gpr, &mut sc);
12093            for g in 0..gpr {
12094                let t = (r * gpr + g) * Q4_TILE;
12095                let s = sc[g];
12096                out[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
12097                let src = (r * gpr + g) * Q4TP_NIB;
12098                out[t + 2..t + Q4_TILE].copy_from_slice(&v.nib[src..src + Q4TP_NIB]);
12099            }
12100        }
12101        out
12102    }
12103
12104    /// The exact (`CMF_SDOT=0`) path must reproduce `dequant_q4tp` to f32
12105    /// rounding — that scalar routine is the format's definition, and the
12106    /// kernels re-derive the scale from the ladder independently. Call the
12107    /// row kernel directly: `matmat` picks the int8 arm when a8w8 is on,
12108    /// so routing through it would test the other path by accident.
12109    #[test]
12110    fn q4tp_exact_path_matches_dequant_reference() {
12111        let (rows, cols) = (256usize, 512usize);
12112        let gpr = cols / GROUP_SIZE;
12113        let bytes = synth_q4tp(rows, cols);
12114        let mut w = vec![0f32; rows * cols];
12115        cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
12116
12117        let x: Vec<f32> = (0..cols)
12118            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
12119            .collect();
12120        let v = Q4tpView::new(&bytes, rows, cols);
12121        let mut sc = vec![0f32; gpr];
12122        for r in 0..rows {
12123            v.scales_into(r, gpr, &mut sc);
12124            let got = q4tp_row_exact(v.nib, r, gpr, &x, &sc);
12125            let want: f32 = (0..cols).map(|c| w[r * cols + c] * x[c]).sum();
12126            // These dot products cancel down to ~1e-3 from terms of ~5e-2, so
12127            // the meaningful yardstick is the summed magnitude, not the result:
12128            // against the result any reordering of a 512-term f32 sum "fails".
12129            let mag: f32 = (0..cols).map(|c| (w[r * cols + c] * x[c]).abs()).sum();
12130            assert!(
12131                (got - want).abs() <= 1e-5 * mag,
12132                "row {r}: kernel {got} vs dequant {want}"
12133            );
12134        }
12135    }
12136
12137    /// The int8 (a8w8) path can't be checked against an f32 reference — the
12138    /// activation quantization dominates. Check it against the q4t kernel it
12139    /// was ported from instead, on payloads holding the same weights: that
12140    /// isolates exactly what the port could break (16 B stride, ladder
12141    /// lookup, nibble unpack) from what it deliberately shares.
12142    #[test]
12143    fn q4tp_matvec_matches_the_q4t_kernel_it_was_ported_from() {
12144        let (rows, cols) = (256usize, 512usize);
12145        let bytes = synth_q4tp(rows, cols);
12146        let twin = q4tp_as_q4t(&bytes, rows, cols);
12147        let x: Vec<f32> = (0..cols)
12148            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
12149            .collect();
12150
12151        let mut got = vec![0f32; rows];
12152        q4tp_matvec(&bytes, &x, rows, cols, &mut got, None);
12153        let mut want = vec![0f32; rows];
12154        q4t_matvec(&twin, &x, rows, cols, &mut want, None);
12155
12156        // Scale is f16 in the twin and f32 here, so allow that rounding on
12157        // top of the summed magnitude (same cancellation argument as above).
12158        let mut w = vec![0f32; rows * cols];
12159        cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
12160        for r in 0..rows {
12161            let mag: f32 = (0..cols).map(|c| (w[r * cols + c] * x[c]).abs()).sum();
12162            assert!(
12163                (got[r] - want[r]).abs() <= 1e-3 * mag,
12164                "row {r}: q4tp {} vs q4t {}",
12165                got[r],
12166                want[r]
12167            );
12168        }
12169    }
12170
12171    /// `matmat` carries three arms (Accelerate, blocked int8 1x4, scalar).
12172    /// Batch 5 crosses the blocked kernel's stride, so this exercises the
12173    /// 1x4 path AND its scalar tail in one run — the blocked kernel is new
12174    /// code and its four accumulators are exactly what tends to go wrong.
12175    #[test]
12176    fn q4tp_matmat_matches_the_q4t_kernel_it_was_ported_from() {
12177        let (rows, cols, b) = (256usize, 512usize, 5usize);
12178        let bytes = synth_q4tp(rows, cols);
12179        let twin = q4tp_as_q4t(&bytes, rows, cols);
12180        let xs: Vec<f32> = (0..b * cols)
12181            .map(|i| ((i * 29 + 11) % 89) as f32 / 89.0 - 0.5)
12182            .collect();
12183
12184        let mut got = vec![0f32; b * rows];
12185        q4tp_matmat(&bytes, &xs, b, rows, cols, &mut got, None);
12186        let mut want = vec![0f32; b * rows];
12187        q4t_matmat(&twin, &xs, b, rows, cols, &mut want, None);
12188
12189        let mut w = vec![0f32; rows * cols];
12190        cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
12191        for t in 0..b {
12192            for r in 0..rows {
12193                let mag: f32 = (0..cols)
12194                    .map(|c| (w[r * cols + c] * xs[t * cols + c]).abs())
12195                    .sum();
12196                let (g, wa) = (got[t * rows + r], want[t * rows + r]);
12197                assert!(
12198                    (g - wa).abs() <= 1e-3 * mag,
12199                    "batch {t} row {r}: q4tp {g} vs q4t {wa}"
12200                );
12201            }
12202        }
12203    }
12204
12205    #[test]
12206    fn q4tp_matvec2_matches_the_single_stream_kernel() {
12207        let (rows, cols) = (128usize, 256usize);
12208        let gpr = cols / GROUP_SIZE;
12209        let bytes = synth_q4tp(rows, cols);
12210        let xs: Vec<f32> = (0..2 * cols)
12211            .map(|i| ((i * 29 + 11) % 89) as f32 / 89.0 - 0.5)
12212            .collect();
12213
12214        let (mut o1, mut o2) = (vec![0f32; rows], vec![0f32; rows]);
12215        q4tp_matvec2(
12216            &bytes,
12217            &xs[..cols],
12218            &xs[cols..],
12219            rows,
12220            cols,
12221            &mut o1,
12222            &mut o2,
12223            None,
12224        );
12225
12226        // matvec2 takes the exact path for both streams, so the single-row
12227        // kernel is an exact reference — no tolerance for path differences.
12228        let v = Q4tpView::new(&bytes, rows, cols);
12229        let mut sc = vec![0f32; gpr];
12230        for r in 0..rows {
12231            v.scales_into(r, gpr, &mut sc);
12232            assert_eq!(o1[r], q4tp_row_exact(v.nib, r, gpr, &xs[..cols], &sc));
12233            assert_eq!(o2[r], q4tp_row_exact(v.nib, r, gpr, &xs[cols..], &sc));
12234        }
12235    }
12236
12237    /// q4tp must not COST speed — it exists to save bytes, and a format that
12238    /// trades 7% of a file for a slower model is a bad trade. This guard is
12239    /// here because correctness tests happily passed while `q4tp_matmat` was
12240    /// missing its int8 and Accelerate arms and the model ran 5x slower.
12241    /// Measured on M-series: 0.97-1.04x, i.e. parity (16 B tiles are better
12242    /// aligned than q4t's 18 B, which pays for the scale indirection).
12243    #[test]
12244    fn q4tp_matvec_keeps_pace_with_q4t() {
12245        let (rows, cols) = (4096usize, 3072usize);
12246        let bytes = synth_q4tp(rows, cols);
12247        let twin = q4tp_as_q4t(&bytes, rows, cols);
12248        let x: Vec<f32> = (0..cols).map(|i| (i % 97) as f32 / 97.0 - 0.5).collect();
12249        let mut o = vec![0f32; rows];
12250        let n = 12;
12251        let mut best = (f64::MAX, f64::MAX);
12252        // Interleaved A/B, minimum statistic: this machine throttles, and a
12253        // mean over a thermal ramp reliably indicts whichever ran second.
12254        for _ in 0..3 {
12255            let t0 = std::time::Instant::now();
12256            for _ in 0..n {
12257                q4t_matvec(&twin, &x, rows, cols, &mut o, None);
12258            }
12259            best.0 = best.0.min(t0.elapsed().as_secs_f64());
12260            let t0 = std::time::Instant::now();
12261            for _ in 0..n {
12262                q4tp_matvec(&bytes, &x, rows, cols, &mut o, None);
12263            }
12264            best.1 = best.1.min(t0.elapsed().as_secs_f64());
12265        }
12266        let ratio = best.1 / best.0;
12267        println!(
12268            "q4t {:.3} ms | q4tp {:.3} ms | {ratio:.2}x",
12269            best.0 * 1e3 / n as f64,
12270            best.1 * 1e3 / n as f64
12271        );
12272        assert!(ratio < 2.0, "q4tp matvec {ratio:.2}x slower than q4t");
12273    }
12274
12275    #[cfg(target_os = "macos")]
12276    #[test]
12277    fn q4t_matmat_accel_matches_dequant_reference() {
12278        if !accel_gemm_enabled() {
12279            return; // CMF_ACCEL=0
12280        }
12281        let (rows, cols, b) = (512usize, 1024usize, 8usize); // ≥500K → accel arm
12282        let gpr = cols / GROUP_SIZE;
12283        let mut bytes = vec![0u8; rows * gpr * Q4_TILE];
12284        for r in 0..rows {
12285            for g in 0..gpr {
12286                let t = (r * gpr + g) * Q4_TILE;
12287                let sc = 0.02 + 0.0005 * ((r * gpr + g) % 64) as f32;
12288                bytes[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
12289                for k in 0..16 {
12290                    bytes[t + 2 + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
12291                }
12292            }
12293        }
12294        let x: Vec<f32> = (0..b * cols)
12295            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
12296            .collect();
12297        let mut got = vec![0f32; b * rows];
12298        q4t_matmat(&bytes, &x, b, rows, cols, &mut got, None);
12299        // Brute-force reference off the same tiles.
12300        let mut w = vec![0f32; rows * cols];
12301        for r in 0..rows {
12302            for g in 0..gpr {
12303                let t = (r * gpr + g) * Q4_TILE;
12304                let s = f16_to_f32(u16::from_le_bytes([bytes[t], bytes[t + 1]]));
12305                for (k, &bb) in bytes[t + 2..t + Q4_TILE].iter().enumerate() {
12306                    w[r * cols + g * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s;
12307                    w[r * cols + g * GROUP_SIZE + k * 2 + 1] =
12308                        (((bb >> 4) & 0x0F) as f32 - 8.0) * s;
12309                }
12310            }
12311        }
12312        for bi in 0..b {
12313            for r in 0..rows {
12314                let want: f32 = (0..cols).map(|j| x[bi * cols + j] * w[r * cols + j]).sum();
12315                let d = (got[bi * rows + r] - want).abs();
12316                assert!(
12317                    d <= want.abs().max(1.0) * 1e-4,
12318                    "accel q4t GEMM diverged at ({bi},{r}): {} vs {want}",
12319                    got[bi * rows + r]
12320                );
12321            }
12322        }
12323    }
12324
12325    #[test]
12326    fn q4matvec_matches_full_dequant() {
12327        let (rows, cols) = (8, 64);
12328        let groups = rows * cols / GROUP_SIZE;
12329        // Hand-craft a q4_block blob: nibbles then f16 scales.
12330        let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
12331        for i in 0..groups * 16 {
12332            bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
12333        }
12334        for g in 0..groups {
12335            let s = 0.01 + 0.003 * g as f32;
12336            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
12337        }
12338        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
12339
12340        let mut reference = vec![0.0f32; rows * cols];
12341        cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
12342        let mut expect = vec![0.0f32; rows];
12343        for r in 0..rows {
12344            expect[r] = reference[r * cols..(r + 1) * cols]
12345                .iter()
12346                .zip(&x)
12347                .map(|(w, xv)| w * xv)
12348                .sum();
12349        }
12350
12351        let mut got = vec![0.0f32; rows];
12352        q4matvec(&bytes, &x, rows, cols, &mut got, None);
12353        // SDOT path quantizes activations to i8 (A8W8): bounded noise,
12354        // same contract as q8/vbit (exact path is pinned by CMF_SDOT=0
12355        // in the golden-parity gate).
12356        let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
12357        let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
12358        for r in 0..rows {
12359            assert!(
12360                (got[r] - expect[r]).abs() < tol * scale,
12361                "row {r}: {} vs {}",
12362                got[r],
12363                expect[r]
12364            );
12365        }
12366    }
12367
12368    /// Fused two-input vbit matvec must equal two single matvecs exactly
12369    /// (same per-lane accumulation order on both scalar and SDOT paths).
12370    #[test]
12371    fn vbitmatvec2_equals_two_singles() {
12372        let (rows, cols) = (6, 64);
12373        let ng = cols / GROUP_SIZE;
12374        let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
12375        let mut bytes = bits.clone();
12376        for g in 0..rows * ng {
12377            let s = 0.02 + 0.001 * g as f32;
12378            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
12379        }
12380        for r in 0..rows {
12381            let b = bits[r] as usize;
12382            let (mut acc, mut nb) = (0u64, 0usize);
12383            let mut rowbytes = Vec::new();
12384            for i in 0..cols {
12385                let v = ((i * 7 + r * 13) % (1 << b)) as u64;
12386                acc = (acc << b) | v;
12387                nb += b;
12388                while nb >= 8 {
12389                    nb -= 8;
12390                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
12391                }
12392            }
12393            if nb > 0 {
12394                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
12395            }
12396            bytes.extend_from_slice(&rowbytes);
12397        }
12398        let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
12399        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).cos()).collect();
12400        let offsets = vbit_row_offsets(&bytes, rows, cols);
12401
12402        let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
12403        vbitmatvec(&bytes, &offsets, &x1, rows, cols, &mut a1, None);
12404        vbitmatvec(&bytes, &offsets, &x2, rows, cols, &mut a2, None);
12405        let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
12406        vbitmatvec2(
12407            &bytes, &offsets, &x1, &x2, rows, cols, &mut b1, &mut b2, None,
12408        );
12409        assert_eq!(a1, b1, "fused vbit lane 1 must be bit-identical");
12410        assert_eq!(a2, b2, "fused vbit lane 2 must be bit-identical");
12411    }
12412
12413    /// Fused two-input q4 matvec must equal two single matvecs exactly.
12414    #[test]
12415    fn q4matvec2_equals_two_singles() {
12416        let (rows, cols) = (8, 128);
12417        let groups = rows * cols / GROUP_SIZE;
12418        let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
12419        for i in 0..groups * 16 {
12420            bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
12421        }
12422        for g in 0..groups {
12423            let s = 0.01 + 0.003 * g as f32;
12424            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
12425        }
12426        // Include an outlier channel so the SDOT correction path is
12427        // exercised in the pair kernel too.
12428        let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
12429        x1[9] = 250.0;
12430        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
12431
12432        let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
12433        q4matvec(&bytes, &x1, rows, cols, &mut a1, None);
12434        q4matvec(&bytes, &x2, rows, cols, &mut a2, None);
12435        let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
12436        q4matvec2(&bytes, &x1, &x2, rows, cols, &mut b1, &mut b2, None);
12437        assert_eq!(a1, b1, "fused q4 lane 1 must be bit-identical");
12438        assert_eq!(a2, b2, "fused q4 lane 2 must be bit-identical");
12439    }
12440
12441    /// Multi-matrix job must equal separate matvecs exactly — same
12442    /// kernels, only the dispatch is fused.
12443    #[test]
12444    fn matvec_many_equals_separate_matvecs() {
12445        use crate::pool::Pool;
12446        let (r1, r2, cols) = (300, 200, 64);
12447        let mk = |salt: usize, rows: usize| {
12448            QTensor::from_f32(
12449                (0..rows * cols)
12450                    .map(|i| ((i * 7 + salt) % 97) as f32 / 97.0 - 0.5)
12451                    .collect(),
12452                rows,
12453                cols,
12454            )
12455        };
12456        let (a, b) = (mk(1, r1), mk(5, r2));
12457        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).sin()).collect();
12458        let pool = Pool::new(3);
12459
12460        let (mut ea, mut eb) = (vec![0f32; r1], vec![0f32; r2]);
12461        a.matvec(&x, &mut ea, Some(&pool));
12462        b.matvec(&x, &mut eb, Some(&pool));
12463        let (mut ga, mut gb) = (vec![0f32; r1], vec![0f32; r2]);
12464        QTensor::matvec_many([&a, &b], &x, [&mut ga, &mut gb], Some(&pool));
12465        assert_eq!(ea, ga, "fused multi-matrix lane 1 must be bit-identical");
12466        assert_eq!(eb, gb, "fused multi-matrix lane 2 must be bit-identical");
12467    }
12468
12469    /// The public Q4TP operator must take the real mapped matvec_many arm,
12470    /// rather than the F32 fallback above.  Build a tiny valid CMF so both
12471    /// handles retain their mmap payloads, then compare the fused dispatch
12472    /// with two ordinary mapped matvec calls bit-for-bit.
12473    #[test]
12474    fn q4tp_matvec_many_equals_separate_matvecs() {
12475        use crate::pool::Pool;
12476        use cortiq_core::{CMF_VERSION, CmfHeader, CmfModel, QuantType, TensorSpec};
12477
12478        let (r1, r2, cols) = (300usize, 200usize, 64usize);
12479        let arch: cortiq_core::ModelArch = serde_json::from_value(serde_json::json!({
12480            "arch_name": "tiny-q4tp",
12481            "hidden_size": cols,
12482            "intermediate_size": cols * 2,
12483            "num_layers": 1,
12484            "num_attention_heads": 2,
12485            "num_kv_heads": 1,
12486            "head_dim": 32,
12487            "vocab_size": r1,
12488            "layer_types": ["FullAttention"],
12489            "rms_norm_eps": 1e-6,
12490            "max_position_embeddings": 8,
12491            "linear_conv_kernel_dim": 0,
12492            "linear_num_key_heads": 0,
12493            "linear_num_value_heads": 0
12494        }))
12495        .unwrap();
12496        let header = CmfHeader {
12497            format: "cmf".into(),
12498            version: CMF_VERSION,
12499            arch,
12500            quant_type: QuantType::Q4Block,
12501            provenance: None,
12502            tokenizer_config: None,
12503            section_hashes: None,
12504            skills: Vec::new(),
12505            shard: None,
12506            calibration: None,
12507            routing: None,
12508            genome: None,
12509            lineage: Vec::new(),
12510            router: None,
12511            segments: Vec::new(),
12512        };
12513        let specs = [
12514            TensorSpec {
12515                name: "q".into(),
12516                dtype: TensorDtype::Q4TiledP,
12517                shape: vec![r1, cols],
12518                data: synth_q4tp(r1, cols),
12519            },
12520            TensorSpec {
12521                name: "kv".into(),
12522                dtype: TensorDtype::Q4TiledP,
12523                shape: vec![r2, cols],
12524                data: synth_q4tp(r2, cols),
12525            },
12526        ];
12527        let dir = std::env::temp_dir().join(format!("cmf-q4tp-many-{}", std::process::id()));
12528        std::fs::create_dir_all(&dir).unwrap();
12529        let path = dir.join("m.cmf");
12530        CmfModel::write(&path, &header, &specs, None, None).unwrap();
12531        let model = Arc::new(CmfModel::open(&path).unwrap());
12532        let (a, b) = (
12533            QTensor::from_model(&model, "q").unwrap(),
12534            QTensor::from_model(&model, "kv").unwrap(),
12535        );
12536        assert_eq!(a.model_dtype(), Some(TensorDtype::Q4TiledP));
12537        assert_eq!(b.model_dtype(), Some(TensorDtype::Q4TiledP));
12538        let x: Vec<f32> = (0..cols)
12539            .map(|i| ((i * 17 + 3) % 97) as f32 / 97.0 - 0.5)
12540            .collect();
12541        let pool = Pool::new(3);
12542        let (mut ea, mut eb) = (vec![0.0f32; r1], vec![0.0f32; r2]);
12543        a.matvec(&x, &mut ea, Some(&pool));
12544        b.matvec(&x, &mut eb, Some(&pool));
12545        let (mut ga, mut gb) = (vec![0.0f32; r1], vec![0.0f32; r2]);
12546        QTensor::matvec_many([&a, &b], &x, [&mut ga, &mut gb], Some(&pool));
12547        assert_eq!(ea, ga, "Q4TP fused lane 1 must be bit-identical");
12548        assert_eq!(eb, gb, "Q4TP fused lane 2 must be bit-identical");
12549        let _ = std::fs::remove_dir_all(&dir);
12550    }
12551
12552    /// The MiMo speculative verify's kernels: several tokens' MoE through
12553    /// `moe_gate_up_rows` / `moe_down_rows` (+ the caller's route-order sum)
12554    /// is bit-identical to each token alone through `moe_gate_up_many` /
12555    /// `moe_down_many` (decode), and a row-exact `q4tp_matmat` of five
12556    /// tokens (wide enough for the blocked tiles) equals five matvecs.
12557    #[test]
12558    fn multi_token_moe_rows_equal_single_token_decode() {
12559        use crate::pool::Pool;
12560        use cortiq_core::{CMF_VERSION, CmfHeader, CmfModel, QuantType, TensorSpec};
12561
12562        let (h, inter, ne) = (64usize, 128usize, 3usize);
12563        let arch: cortiq_core::ModelArch = serde_json::from_value(serde_json::json!({
12564            "arch_name": "tiny-q4tp-moe",
12565            "hidden_size": h,
12566            "intermediate_size": inter,
12567            "num_layers": 1,
12568            "num_attention_heads": 2,
12569            "num_kv_heads": 1,
12570            "head_dim": 32,
12571            "vocab_size": 8,
12572            "layer_types": ["FullAttention"],
12573            "rms_norm_eps": 1e-6,
12574            "max_position_embeddings": 8,
12575            "linear_conv_kernel_dim": 0,
12576            "linear_num_key_heads": 0,
12577            "linear_num_value_heads": 0
12578        }))
12579        .unwrap();
12580        let header = CmfHeader {
12581            format: "cmf".into(),
12582            version: CMF_VERSION,
12583            arch,
12584            quant_type: QuantType::Q4Block,
12585            provenance: None,
12586            tokenizer_config: None,
12587            section_hashes: None,
12588            skills: Vec::new(),
12589            shard: None,
12590            calibration: None,
12591            routing: None,
12592            genome: None,
12593            lineage: Vec::new(),
12594            router: None,
12595            segments: Vec::new(),
12596        };
12597        let mut specs = Vec::new();
12598        for e in 0..ne {
12599            for (k, (n, r, c)) in [("g", inter, h), ("u", inter, h), ("d", h, inter)]
12600                .into_iter()
12601                .enumerate()
12602            {
12603                // Distinct experts: perturb only the nibble plane (any byte
12604                // is a valid pair of codes; the ladder stays intact).
12605                let mut data = synth_q4tp(r, c);
12606                for (i, byte) in data[..r * (c / GROUP_SIZE) * Q4TP_NIB]
12607                    .iter_mut()
12608                    .enumerate()
12609                {
12610                    *byte ^= ((i * (e * 3 + k + 1)) % 251) as u8;
12611                }
12612                specs.push(TensorSpec {
12613                    name: format!("{n}{e}"),
12614                    dtype: TensorDtype::Q4TiledP,
12615                    shape: vec![r, c],
12616                    data,
12617                });
12618            }
12619        }
12620        let dir = std::env::temp_dir().join(format!(
12621            "cmf-moe-rows-{}-{}",
12622            std::process::id(),
12623            FLOAT_ACTIVATIONS.get()
12624        ));
12625        std::fs::create_dir_all(&dir).unwrap();
12626        let path = dir.join("m.cmf");
12627        CmfModel::write(&path, &header, &specs, None, None).unwrap();
12628        let model = Arc::new(CmfModel::open(&path).unwrap());
12629        let t = |n: String| QTensor::from_model(&model, &n).unwrap();
12630        let g: Vec<QTensor> = (0..ne).map(|e| t(format!("g{e}"))).collect();
12631        let u: Vec<QTensor> = (0..ne).map(|e| t(format!("u{e}"))).collect();
12632        let d: Vec<QTensor> = (0..ne).map(|e| t(format!("d{e}"))).collect();
12633        let b = 4usize;
12634        let mut xs: Vec<f32> = (0..b * h)
12635            .map(|i| ((i * 31 + 7) % 89) as f32 / 89.0 - 0.5)
12636            .collect();
12637        xs[5] = 9.0; // an activation outlier on token 0
12638        // Token -> (experts in route order, weights).
12639        let routes: Vec<(Vec<usize>, Vec<f32>)> = vec![
12640            (vec![2, 0], vec![0.6, 0.4]),
12641            (vec![0, 1, 2], vec![0.2, 0.5, 0.3]),
12642            (vec![1], vec![1.0]),
12643            (vec![2, 1, 0], vec![0.25, 0.25, 0.5]),
12644        ];
12645        let pool = Pool::new(3);
12646        // Decode reference, token by token.
12647        let mut want = vec![0f32; b * h];
12648        for (tk, (idx, w)) in routes.iter().enumerate() {
12649            let x = &xs[tk * h..(tk + 1) * h];
12650            let pairs: Vec<(&QTensor, &QTensor)> = idx.iter().map(|&e| (&g[e], &u[e])).collect();
12651            let mut gs: Vec<Vec<f32>> = idx.iter().map(|_| vec![0f32; inter]).collect();
12652            assert!(QTensor::moe_gate_up_many(&pairs, x, &mut gs, Some(&pool)));
12653            if FLOAT_ACTIVATIONS.get() {
12654                for (slot, &e) in idx.iter().enumerate() {
12655                    let (mut gate, mut up) = (vec![0.0; inter], vec![0.0; inter]);
12656                    g[e].matvec(x, &mut gate, Some(&pool));
12657                    u[e].matvec(x, &mut up, Some(&pool));
12658                    for (v, u) in gate.iter_mut().zip(up) {
12659                        *v = (*v / (1.0 + (-*v).exp())) * u;
12660                    }
12661                    assert_eq!(gs[slot], gate, "float gate/up must equal ordinary matvecs");
12662                }
12663            }
12664            let downs: Vec<&QTensor> = idx.iter().map(|&e| &d[e]).collect();
12665            assert!(QTensor::moe_down_many(
12666                &downs,
12667                &gs,
12668                w,
12669                &mut want[tk * h..(tk + 1) * h],
12670                Some(&pool)
12671            ));
12672        }
12673        if FLOAT_ACTIVATIONS.get() {
12674            for (tk, (idx, w)) in routes.iter().enumerate() {
12675                let mut scalar = vec![0.0; h];
12676                for (&e, &weight) in idx.iter().zip(w) {
12677                    let (mut gate, mut up, mut down) =
12678                        (vec![0.0; inter], vec![0.0; inter], vec![0.0; h]);
12679                    g[e].matvec(&xs[tk * h..(tk + 1) * h], &mut gate, Some(&pool));
12680                    u[e].matvec(&xs[tk * h..(tk + 1) * h], &mut up, Some(&pool));
12681                    for (v, u) in gate.iter_mut().zip(up) {
12682                        *v = (*v / (1.0 + (-*v).exp())) * u;
12683                    }
12684                    d[e].matvec(&gate, &mut down, Some(&pool));
12685                    for (v, d) in scalar.iter_mut().zip(down) {
12686                        *v += weight * d;
12687                    }
12688                }
12689                assert_eq!(
12690                    &want[tk * h..(tk + 1) * h],
12691                    scalar,
12692                    "float many equals scalar experts"
12693                );
12694            }
12695        }
12696        // All four tokens at once, grouped by expert.
12697        let mut experts: Vec<usize> = Vec::new();
12698        let mut groups: Vec<Vec<usize>> = Vec::new();
12699        for (tk, (idx, _)) in routes.iter().enumerate() {
12700            for &e in idx {
12701                match experts.iter().position(|&x| x == e) {
12702                    Some(k) => groups[k].push(tk),
12703                    None => {
12704                        experts.push(e);
12705                        groups.push(vec![tk]);
12706                    }
12707                }
12708            }
12709        }
12710        let n_pairs: usize = groups.iter().map(|g| g.len()).sum();
12711        let pairs: Vec<(&QTensor, &QTensor)> = experts.iter().map(|&e| (&g[e], &u[e])).collect();
12712        let mut gs: Vec<Vec<f32>> = (0..n_pairs).map(|_| vec![0f32; inter]).collect();
12713        assert!(QTensor::moe_gate_up_rows(
12714            &pairs,
12715            &groups,
12716            &xs,
12717            &mut gs,
12718            Some(&pool)
12719        ));
12720        let downs: Vec<&QTensor> = experts.iter().map(|&e| &d[e]).collect();
12721        let lens: Vec<usize> = groups.iter().map(|g| g.len()).collect();
12722        let mut ds: Vec<Vec<f32>> = (0..n_pairs).map(|_| vec![0f32; h]).collect();
12723        assert!(QTensor::moe_down_rows(
12724            &downs,
12725            &lens,
12726            &gs,
12727            &mut ds,
12728            Some(&pool)
12729        ));
12730        let slot = |tk: usize, e: usize| {
12731            let k = experts.iter().position(|&x| x == e).unwrap();
12732            groups[..k].iter().map(|g| g.len()).sum::<usize>()
12733                + groups[k].iter().position(|&x| x == tk).unwrap()
12734        };
12735        let mut got = vec![0f32; b * h];
12736        for (tk, (idx, w)) in routes.iter().enumerate() {
12737            for i in 0..h {
12738                let mut acc = 0f32;
12739                for (&e, &we) in idx.iter().zip(w) {
12740                    acc += we * ds[slot(tk, e)][i];
12741                }
12742                got[tk * h + i] = acc;
12743            }
12744        }
12745        assert!(want.iter().any(|v| *v != 0.0));
12746        assert_eq!(
12747            want.iter().map(|v| v.to_bits()).collect::<Vec<_>>(),
12748            got.iter().map(|v| v.to_bits()).collect::<Vec<_>>(),
12749            "multi-token MoE must equal decode bit for bit"
12750        );
12751
12752        // Row-exact q4tp_matmat: five tokens (a blocked 1x4 tile + tail
12753        // otherwise) equal five matvecs.
12754        let b5 = 5usize;
12755        let x5: Vec<f32> = (0..b5 * h)
12756            .map(|i| ((i * 13 + 5) % 71) as f32 / 71.0 - 0.5)
12757            .collect();
12758        let mut mm = vec![0f32; b5 * inter];
12759        row_exact_scope(|| g[1].matmat(&x5, b5, &mut mm, Some(&pool)));
12760        for tk in 0..b5 {
12761            let mut mv = vec![0f32; inter];
12762            g[1].matvec(&x5[tk * h..(tk + 1) * h], &mut mv, Some(&pool));
12763            assert_eq!(
12764                mv.iter().map(|v| v.to_bits()).collect::<Vec<_>>(),
12765                mm[tk * inter..(tk + 1) * inter]
12766                    .iter()
12767                    .map(|v| v.to_bits())
12768                    .collect::<Vec<_>>(),
12769                "row-exact matmat token {tk}"
12770            );
12771        }
12772        // Other concurrent tests/requests may still hold the shared mode.
12773        // Nested, overlapping and unwind restoration is checked separately.
12774        let _ = std::fs::remove_dir_all(&dir);
12775    }
12776
12777    /// The row-exact fix must not touch the fast path. Outside the scope
12778    /// `q4tp_matmat` has to produce exactly what it did before, and on ARM
12779    /// "before" is spelled out below: the tuned 1x4 SDOT tile for every
12780    /// four columns and the single-row kernel for the tail. Inside the
12781    /// scope every column equals its token's matvec. `q4tp_matmat_with`
12782    /// takes the mode as an argument, so a concurrent test holding the
12783    /// shared scope cannot flip it under this one.
12784    #[test]
12785    fn q4tp_matmat_fast_path_unchanged_outside_row_exact() {
12786        use crate::pool::Pool;
12787        use std::sync::atomic::Ordering::Relaxed;
12788        let _alt = Q4TP_ALT_TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner());
12789        // The tuned ARM shape, which is also what an unset switch picks
12790        // unless CMF_Q4TP_V1 is exported.
12791        Q4TP_ALT.store(2, Relaxed);
12792        let pool = Pool::new(3);
12793        // Under 500k cells, so macOS keeps the matmat off the AMX; the
12794        // second shape runs across pool workers (rows >= 256) with 32
12795        // groups of accumulation and a tail after two 1x4 tiles.
12796        for &(rows, cols, b) in &[(64usize, 256usize, 7usize), (320, 1024, 9)] {
12797            let bytes = synth_q4tp(rows, cols);
12798            let mut xs: Vec<f32> = (0..b * cols)
12799                .map(|i| ((i * 29 + 11) % 83) as f32 / 83.0 - 0.5)
12800                .collect();
12801            xs[3] = 7.5; // an activation outlier on token 0
12802            let run = |exact: bool| {
12803                let mut out = vec![0f32; b * rows];
12804                q4tp_matmat_with(&bytes, &xs, b, rows, cols, &mut out, Some(&pool), exact);
12805                out
12806            };
12807            let (fast, exact) = (run(false), run(true));
12808            #[cfg(not(target_arch = "aarch64"))]
12809            let _ = fast;
12810            let bits = |v: &[f32]| v.iter().map(|x| x.to_bits()).collect::<Vec<_>>();
12811            let mut matvecs = vec![0f32; b * rows];
12812            for (bi, o) in matvecs.chunks_mut(rows).enumerate() {
12813                q4tp_matvec(
12814                    &bytes,
12815                    &xs[bi * cols..(bi + 1) * cols],
12816                    rows,
12817                    cols,
12818                    o,
12819                    Some(&pool),
12820                );
12821            }
12822            assert!(matvecs.iter().any(|v| *v != 0.0));
12823            assert_eq!(
12824                bits(&exact),
12825                bits(&matvecs),
12826                "{rows}x{cols} b={b}: row-exact matmat must equal per-token matvecs"
12827            );
12828            #[cfg(target_arch = "aarch64")]
12829            {
12830                // The pre-fix ARM loop, cell for cell.
12831                let gpr = cols / GROUP_SIZE;
12832                let v = Q4tpView::new(&bytes, rows, cols);
12833                let mut old = vec![0f32; b * rows];
12834                let mut sc = vec![0f32; gpr];
12835                let a8w8 = a8w8_enabled();
12836                let blocked = sdot_enabled() && blocked_enabled();
12837                let acts: Vec<SplitAct> = (0..b)
12838                    .map(|bi| split_act(&xs[bi * cols..(bi + 1) * cols]))
12839                    .collect();
12840                for r in 0..rows {
12841                    v.scales_into(r, gpr, &mut sc);
12842                    if !a8w8 {
12843                        for bi in 0..b {
12844                            let x = &xs[bi * cols..(bi + 1) * cols];
12845                            old[bi * rows + r] = q4tp_row_exact(v.nib, r, gpr, x, &sc);
12846                        }
12847                        continue;
12848                    }
12849                    let finish = |d: f32, act: &SplitAct| {
12850                        let mut acc = d * act.sx;
12851                        for &(j, xv) in &act.outliers {
12852                            let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
12853                            acc += w * s * xv;
12854                        }
12855                        acc
12856                    };
12857                    let mut bi = 0usize;
12858                    while blocked && bi + 4 <= b {
12859                        let xs4 = [
12860                            acts[bi].xq.as_slice(),
12861                            acts[bi + 1].xq.as_slice(),
12862                            acts[bi + 2].xq.as_slice(),
12863                            acts[bi + 3].xq.as_slice(),
12864                        ];
12865                        let d = unsafe { dot_q4tp_row_1x4_sdot(v.nib, r, gpr, xs4, &sc) };
12866                        for k in 0..4 {
12867                            old[(bi + k) * rows + r] = finish(d[k], &acts[bi + k]);
12868                        }
12869                        bi += 4;
12870                    }
12871                    for (bi, act) in acts.iter().enumerate().skip(bi) {
12872                        let d = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, &sc);
12873                        old[bi * rows + r] = finish(d, act);
12874                    }
12875                }
12876                assert_eq!(
12877                    bits(&fast),
12878                    bits(&old),
12879                    "{rows}x{cols} b={b}: the fast path changed outside row_exact"
12880                );
12881                // And it is still the fast tile that runs: its lane-parallel
12882                // fma sum rounds differently from the matvec somewhere.
12883                if blocked {
12884                    assert_ne!(
12885                        bits(&fast),
12886                        bits(&matvecs),
12887                        "{rows}x{cols} b={b}: the tuned tile no longer runs outside row_exact"
12888                    );
12889                }
12890            }
12891        }
12892        Q4TP_ALT.store(0, Relaxed);
12893    }
12894
12895    #[test]
12896    fn row_exact_scopes_survive_overlap_nesting_and_unwind() {
12897        use std::sync::{Barrier, atomic::{AtomicUsize, Ordering}};
12898        // A private counter makes this restoration test independent of
12899        // numerical tests concurrently using the production counter.
12900        let active = AtomicUsize::new(0);
12901        counted_row_exact_scope(&active, || {
12902            assert_eq!(active.load(Ordering::Acquire), 1);
12903            counted_row_exact_scope(&active, || {
12904                assert_eq!(active.load(Ordering::Acquire), 2);
12905            });
12906            assert_eq!(active.load(Ordering::Acquire), 1);
12907        });
12908        assert_eq!(active.load(Ordering::Acquire), 0);
12909
12910        let both_entered = Barrier::new(2);
12911        let release_last = Barrier::new(2);
12912        std::thread::scope(|s| {
12913            let first = s.spawn(|| counted_row_exact_scope(&active, || {
12914                both_entered.wait();
12915            }));
12916            let last = s.spawn(|| counted_row_exact_scope(&active, || {
12917                both_entered.wait();
12918                release_last.wait();
12919            }));
12920            first.join().unwrap();
12921            let after_first = active.load(Ordering::Acquire);
12922            release_last.wait();
12923            last.join().unwrap();
12924            assert_eq!(after_first, 1, "second request must remain exact");
12925        });
12926        assert_eq!(active.load(Ordering::Acquire), 0);
12927        let panic = std::panic::catch_unwind(|| {
12928            counted_row_exact_scope(&active, || panic!("scope unwind"));
12929        });
12930        assert!(panic.is_err());
12931        assert_eq!(active.load(Ordering::Acquire), 0);
12932    }
12933
12934    #[test]
12935    #[cfg(target_arch = "x86_64")]
12936    fn q4tp_float_avx2_is_bitwise_scalar() {
12937        if !avx2_enabled() {
12938            return;
12939        }
12940        for cols in [32, 64, 96, 2048, 4096] {
12941            let rows = 9;
12942            let bytes = synth_q4tp(rows, cols);
12943            let v = Q4tpView::new(&bytes, rows, cols);
12944            let gpr = cols / GROUP_SIZE;
12945            let mut sc = vec![0.0; gpr];
12946            for seed in 1..=5 {
12947                let xs: Vec<f32> = (0..cols)
12948                    .map(|i| (((i * 104729 + seed * 8191) % 100003) as f32 - 50001.0) / 7919.0)
12949                    .collect();
12950                for r in 0..rows {
12951                    v.scales_into(r, gpr, &mut sc);
12952                    let scalar = q4tp_row_float_scalar(v.nib, r, gpr, &xs, &sc);
12953                    let vector = unsafe { q4tp_row_float_avx2(v.nib, r, gpr, &xs, &sc) };
12954                    assert_eq!(
12955                        scalar.to_bits(),
12956                        vector.to_bits(),
12957                        "cols={cols} row={r} seed={seed}"
12958                    );
12959                }
12960            }
12961        }
12962    }
12963
12964    #[test]
12965    fn multi_token_moe_rows_float_equal_single_token_decode() {
12966        float_activations_scope(multi_token_moe_rows_equal_single_token_decode);
12967    }
12968
12969    #[test]
12970    fn full_gpu_q8_scope_is_nested_and_thread_local() {
12971        assert!(!FULL_GPU_Q8.get());
12972        let before = gpu_split_frac();
12973        {
12974            let _guard = enter_full_gpu_q8_scope();
12975            assert_eq!(gpu_split_frac(), 1.0);
12976            {
12977                let _nested = enter_full_gpu_q8_scope();
12978            }
12979            assert_eq!(gpu_split_frac(), 1.0);
12980            std::thread::spawn(|| assert!(!FULL_GPU_Q8.get())).join().unwrap();
12981        }
12982        assert!(!FULL_GPU_Q8.get());
12983        assert_eq!(gpu_split_frac(), before);
12984    }
12985
12986    #[test]
12987    fn float_activation_scope_is_nested_thread_local_and_unwind_safe() {
12988        assert!(!FLOAT_ACTIVATIONS.get());
12989        let before = a8w8_enabled();
12990        float_activations_scope(|| {
12991            assert!(!a8w8_enabled());
12992            float_activations_scope(|| assert!(!a8w8_enabled()));
12993            assert!(FLOAT_ACTIVATIONS.get());
12994            std::thread::spawn(|| assert!(!FLOAT_ACTIVATIONS.get()))
12995                .join()
12996                .unwrap();
12997        });
12998        assert!(!FLOAT_ACTIVATIONS.get());
12999        assert_eq!(a8w8_enabled(), before);
13000        let _ = std::panic::catch_unwind(|| float_activations_scope(|| panic!("test unwind")));
13001        assert!(!FLOAT_ACTIVATIONS.get());
13002    }
13003
13004    /// Batched q4/vbit matmat must equal per-position matvec calls
13005    /// exactly (the fallback it replaced) — same kernels, same order.
13006    #[test]
13007    fn batched_matmat_equals_per_position_matvec() {
13008        let (rows, cols, b) = (8, 64, 5);
13009        // q4 blob.
13010        let groups = rows * cols / GROUP_SIZE;
13011        let mut q4 = Vec::new();
13012        for i in 0..groups * 16 {
13013            q4.push((((i * 7 + 3) % 256) & 0xFF) as u8);
13014        }
13015        for g in 0..groups {
13016            q4.extend_from_slice(
13017                &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
13018            );
13019        }
13020        // vbit blob (mixed widths incl. 8).
13021        let ng = cols / GROUP_SIZE;
13022        let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4, 5, 3];
13023        let mut vb = bits.clone();
13024        for g in 0..rows * ng {
13025            vb.extend_from_slice(
13026                &cortiq_core::quant::f32_to_f16(0.02 + 0.001 * g as f32).to_le_bytes(),
13027            );
13028        }
13029        for r in 0..rows {
13030            let bw = bits[r] as usize;
13031            let (mut acc, mut nb) = (0u64, 0usize);
13032            let mut rowbytes = Vec::new();
13033            for i in 0..cols {
13034                let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
13035                acc = (acc << bw) | v;
13036                nb += bw;
13037                while nb >= 8 {
13038                    nb -= 8;
13039                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
13040                }
13041            }
13042            if nb > 0 {
13043                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
13044            }
13045            vb.extend_from_slice(&rowbytes);
13046        }
13047        let offsets = vbit_row_offsets(&vb, rows, cols);
13048
13049        let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
13050
13051        // q4: batch vs singles.
13052        let mut got = vec![0f32; b * rows];
13053        q4matmat(&q4, &xs, b, rows, cols, &mut got, None);
13054        for bi in 0..b {
13055            let mut expect = vec![0f32; rows];
13056            q4matvec(
13057                &q4,
13058                &xs[bi * cols..(bi + 1) * cols],
13059                rows,
13060                cols,
13061                &mut expect,
13062                None,
13063            );
13064            assert_eq!(
13065                &got[bi * rows..(bi + 1) * rows],
13066                &expect[..],
13067                "q4 batch pos {bi}"
13068            );
13069        }
13070
13071        // vbit: batch vs singles.
13072        let mut got = vec![0f32; b * rows];
13073        vbitmatmat(&vb, &offsets, &xs, b, rows, cols, &mut got, None);
13074        for bi in 0..b {
13075            let mut expect = vec![0f32; rows];
13076            vbitmatvec(
13077                &vb,
13078                &offsets,
13079                &xs[bi * cols..(bi + 1) * cols],
13080                rows,
13081                cols,
13082                &mut expect,
13083                None,
13084            );
13085            assert_eq!(
13086                &got[bi * rows..(bi + 1) * rows],
13087                &expect[..],
13088                "vbit batch pos {bi}"
13089            );
13090        }
13091    }
13092
13093    /// q4_tiled kernels must produce BIT-identical outputs to the q4
13094    /// split kernels on the same values (same ints, same order — only
13095    /// the byte placement differs).
13096    #[test]
13097    fn q4_tiled_matches_q4_block_bitexact() {
13098        let (rows, cols, b) = (8usize, 128usize, 3usize);
13099        let groups = rows * cols / GROUP_SIZE;
13100        let mut split = Vec::with_capacity(groups * 18);
13101        for i in 0..groups * 16 {
13102            split.push((((i * 7 + 3) % 256) & 0xFF) as u8);
13103        }
13104        for g in 0..groups {
13105            split.extend_from_slice(
13106                &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
13107            );
13108        }
13109        // Re-tile: [scale][nibbles] per group.
13110        let (packed, scales) = split.split_at(groups * 16);
13111        let mut tiled = Vec::with_capacity(groups * Q4_TILE);
13112        for g in 0..groups {
13113            tiled.extend_from_slice(&scales[g * 2..g * 2 + 2]);
13114            tiled.extend_from_slice(&packed[g * 16..(g + 1) * 16]);
13115        }
13116
13117        let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
13118        x1[9] = 250.0; // exercise the outlier path
13119        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
13120
13121        let (mut a, mut t) = (vec![0f32; rows], vec![0f32; rows]);
13122        q4matvec(&split, &x1, rows, cols, &mut a, None);
13123        q4t_matvec(&tiled, &x1, rows, cols, &mut t, None);
13124        assert_eq!(a, t, "q4t matvec must match q4 bit-for-bit");
13125
13126        let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
13127        let (mut t1, mut t2) = (vec![0f32; rows], vec![0f32; rows]);
13128        q4matvec2(&split, &x1, &x2, rows, cols, &mut a1, &mut a2, None);
13129        q4t_matvec2(&tiled, &x1, &x2, rows, cols, &mut t1, &mut t2, None);
13130        assert_eq!(a1, t1);
13131        assert_eq!(a2, t2);
13132
13133        let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
13134        let (mut am, mut tm) = (vec![0f32; b * rows], vec![0f32; b * rows]);
13135        q4matmat(&split, &xs, b, rows, cols, &mut am, None);
13136        q4t_matmat(&tiled, &xs, b, rows, cols, &mut tm, None);
13137        assert_eq!(am, tm, "q4t matmat must match q4 bit-for-bit");
13138    }
13139
13140    /// q4 SDOT outlier correction: a single huge activation channel
13141    /// (>8·rms → outlier, zeroed in xq) must still contribute its EXACT
13142    /// term. On-grid bulk (±1/0 → xq dequantizes exactly) isolates the
13143    /// correction from A8W8 noise. cols must exceed 64: at n=64 the
13144    /// 8·rms threshold equals sqrt(v²+rest) ≥ v, so a single outlier
13145    /// can never qualify (8² = n).
13146    #[test]
13147    fn q4matvec_sdot_outlier_exact() {
13148        let (rows, cols) = (4, 128);
13149        let groups = rows * cols / GROUP_SIZE;
13150        let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
13151        for i in 0..groups * 16 {
13152            bytes.push(((i * 11 + 5) % 256) as u8);
13153        }
13154        for g in 0..groups {
13155            let s = 0.02 + 0.002 * g as f32;
13156            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
13157        }
13158        let mut x: Vec<f32> = (0..cols)
13159            .map(|i| match i % 3 {
13160                0 => 1.0,
13161                1 => -1.0,
13162                _ => 0.0,
13163            })
13164            .collect();
13165        x[17] = 300.0; // ≫ 8·rms → outlier channel
13166
13167        let mut reference = vec![0.0f32; rows * cols];
13168        cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
13169        let mut expect = vec![0.0f32; rows];
13170        for r in 0..rows {
13171            expect[r] = reference[r * cols..(r + 1) * cols]
13172                .iter()
13173                .zip(&x)
13174                .map(|(w, xv)| w * xv)
13175                .sum();
13176        }
13177        let mut got = vec![0.0f32; rows];
13178        q4matvec(&bytes, &x, rows, cols, &mut got, None);
13179        let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
13180        for r in 0..rows {
13181            assert!(
13182                (got[r] - expect[r]).abs() < 2e-3 * scale,
13183                "row {r}: {} vs {} (outlier term must be exact)",
13184                got[r],
13185                expect[r]
13186            );
13187        }
13188    }
13189
13190    /// The fused q1t matvec must equal the reference (dequant_q1t → dot),
13191    /// including the ternary zero level and the binary-searched outlier
13192    /// overlay. Guards the mmap kernel that makes a 12B q1t runnable.
13193    #[test]
13194    fn q1t_matvec_matches_reference() {
13195        use cortiq_core::quant::{dequant_q1t, f32_to_f16};
13196        let (rows, cols) = (3usize, 64usize); // gpr = 2
13197        let gpr = cols / GROUP_SIZE;
13198        let scales = [0.5f32, 0.3, 0.7, 0.2, 0.6, 0.15];
13199        // Overlay (must be sorted by flat index): a few spikes across rows.
13200        let outliers: [(u32, f32); 3] = [(5, 9.0), (70, -4.5), (150, 3.25)];
13201        let is_out = |flat: usize| outliers.iter().any(|&(i, _)| i as usize == flat);
13202        let mut bytes = Vec::new();
13203        for r in 0..rows {
13204            for g in 0..gpr {
13205                bytes.extend_from_slice(&f32_to_f16(scales[r * gpr + g]).to_le_bytes());
13206                let mut c = [0u8; 7];
13207                for k in 0..GROUP_SIZE {
13208                    // Encoder invariant: code 0 at outlier positions.
13209                    let code = if is_out(r * cols + g * GROUP_SIZE + k) {
13210                        0
13211                    } else {
13212                        ((k + r * 3 + g) % 3) as u8 // 0,1,2
13213                    };
13214                    cortiq_core::quant::q1t_pack(&mut c, k, code);
13215                }
13216                bytes.extend_from_slice(&c);
13217            }
13218        }
13219        // Per-row overlay: [u32 row_ptr[rows+1]] then [(u16 col, f16 val)] by
13220        // row (outliers are sorted by flat index → already grouped by row).
13221        let mut row_ptr = vec![0u32; rows + 1];
13222        for &(idx, _) in &outliers {
13223            row_ptr[idx as usize / cols + 1] += 1;
13224        }
13225        for r in 0..rows {
13226            row_ptr[r + 1] += row_ptr[r];
13227        }
13228        for &p in &row_ptr {
13229            bytes.extend_from_slice(&p.to_le_bytes());
13230        }
13231        for &(idx, v) in &outliers {
13232            bytes.extend_from_slice(&((idx as usize % cols) as u16).to_le_bytes());
13233            bytes.extend_from_slice(&f32_to_f16(v).to_le_bytes());
13234        }
13235
13236        let mut refw = vec![0f32; rows * cols];
13237        dequant_q1t(&bytes, rows, cols, &mut refw);
13238        // On-grid activations (±1, amax 1) so the int8 SDOT path reconstructs
13239        // x exactly and matches the f32 reference (same trick as the q1 test).
13240        let x: Vec<f32> = (0..cols)
13241            .map(|j| if j % 3 == 0 { 1.0 } else { -1.0 })
13242            .collect();
13243        let mut expect = vec![0f32; rows];
13244        for r in 0..rows {
13245            let mut a = 0.0f32;
13246            for j in 0..cols {
13247                a += refw[r * cols + j] * x[j];
13248            }
13249            expect[r] = a;
13250        }
13251        let tol = |e: f32| 1e-3 * e.abs().max(1e-3);
13252        let mut got = vec![0f32; rows];
13253        q1t_matvec(&bytes, &x, rows, cols, &mut got, None);
13254        for r in 0..rows {
13255            assert!(
13256                (got[r] - expect[r]).abs() < tol(expect[r]),
13257                "row {r}: {} vs {}",
13258                got[r],
13259                expect[r]
13260            );
13261        }
13262        // matmat (b=2, f32 decode path) must agree too.
13263        let x2: Vec<f32> = x.iter().chain(x.iter()).copied().collect();
13264        let mut gm = vec![0f32; 2 * rows];
13265        q1t_matmat(&bytes, &x2, 2, rows, cols, &mut gm, None);
13266        for r in 0..rows {
13267            assert!((gm[r] - expect[r]).abs() < tol(expect[r]));
13268            assert!((gm[rows + r] - expect[r]).abs() < tol(expect[r]));
13269        }
13270        // Fused pair (q1t_matvec2) must equal two single matvecs
13271        // bit-for-bit: same unpack, same group order, same f32
13272        // accumulation per stream. Distinct x2 exercises both lanes.
13273        let xb: Vec<f32> = (0..cols)
13274            .map(|j| if j % 5 == 0 { -1.0 } else { 1.0 })
13275            .collect();
13276        let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
13277        q1t_matvec(&bytes, &x, rows, cols, &mut s1, None);
13278        q1t_matvec(&bytes, &xb, rows, cols, &mut s2, None);
13279        let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
13280        q1t_matvec2(&bytes, &x, &xb, rows, cols, &mut p1, &mut p2, None);
13281        assert_eq!(p1, s1, "q1t pair lane 1 ≠ single matvec");
13282        assert_eq!(p2, s2, "q1t pair lane 2 ≠ single matvec");
13283    }
13284
13285    /// Pair == 2×matvec with an ODD group count (the kernel's tail
13286    /// group) and no overlay section.
13287    #[test]
13288    fn q1t_matvec2_odd_gpr_matches_singles() {
13289        use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_pack};
13290        let (rows, cols) = (5usize, 96usize); // gpr = 3 → paired + tail
13291        let gpr = cols / GROUP_SIZE;
13292        let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE);
13293        for r in 0..rows {
13294            for g in 0..gpr {
13295                bytes.extend_from_slice(&f32_to_f16(0.1 + 0.05 * (r + g) as f32).to_le_bytes());
13296                let mut c = [0u8; 7];
13297                for k in 0..GROUP_SIZE {
13298                    q1t_pack(&mut c, k, ((k * 7 + r * 5 + g * 3) % 3) as u8);
13299                }
13300                bytes.extend_from_slice(&c);
13301            }
13302        }
13303        let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.31).sin()).collect();
13304        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).cos()).collect();
13305        let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
13306        q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
13307        q1t_matvec(&bytes, &x2, rows, cols, &mut s2, None);
13308        let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
13309        q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
13310        assert_eq!(p1, s1, "odd-gpr pair lane 1 ≠ single");
13311        assert_eq!(p2, s2, "odd-gpr pair lane 2 ≠ single");
13312    }
13313
13314    // Speed A/B: fused pair (one unpack, two streams) vs two single
13315    // matvecs. Single-threaded, FFN-sized, min-of paired in-process.
13316    //   cargo test -p cortiq-engine --release q1t_matvec2_speed -- --ignored --nocapture
13317    #[test]
13318    #[ignore]
13319    fn q1t_matvec2_speed() {
13320        use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_pack};
13321        use std::time::Instant;
13322        let (rows, cols) = (8192usize, 4096usize);
13323        let gpr = cols / GROUP_SIZE;
13324        let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE);
13325        for r in 0..rows {
13326            for g in 0..gpr {
13327                let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
13328                bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
13329                let mut c = [0u8; 7];
13330                for k in 0..GROUP_SIZE {
13331                    q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
13332                }
13333                bytes.extend_from_slice(&c);
13334            }
13335        }
13336        let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.31).sin()).collect();
13337        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).cos()).collect();
13338        let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
13339        let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
13340        // Warm both paths once.
13341        q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
13342        q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
13343        let (mut t_pair, mut t_two) = (f64::MAX, f64::MAX);
13344        for _ in 0..8 {
13345            let t0 = Instant::now();
13346            q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
13347            t_pair = t_pair.min(t0.elapsed().as_secs_f64() * 1000.0);
13348            let t1 = Instant::now();
13349            q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
13350            q1t_matvec(&bytes, &x2, rows, cols, &mut s2, None);
13351            t_two = t_two.min(t1.elapsed().as_secs_f64() * 1000.0);
13352        }
13353        assert_eq!(p1, s1);
13354        assert_eq!(p2, s2);
13355        println!("q1t pair {rows}x{cols}: fused {t_pair:.2} ms | two singles {t_two:.2} ms");
13356    }
13357
13358    // Speed A/B: the base-3-division decode (what the packing commit left in
13359    // place) vs the fused sign-LUT matvec. Both single-threaded, same bytes.
13360    //   cargo test -p cortiq-engine q1t_matvec_speed -- --ignored --nocapture
13361    #[test]
13362    #[ignore]
13363    fn q1t_matvec_speed() {
13364        use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_code, q1t_pack};
13365        use std::time::Instant;
13366        let (rows, cols) = (8192usize, 4096usize); // FFN-sized
13367        let gpr = cols / GROUP_SIZE;
13368        let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE + 16);
13369        for r in 0..rows {
13370            for g in 0..gpr {
13371                let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
13372                bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
13373                let mut c = [0u8; 7];
13374                for k in 0..GROUP_SIZE {
13375                    q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
13376                }
13377                bytes.extend_from_slice(&c);
13378            }
13379        }
13380        let (n, stride) = (rows * cols, 40usize); // ~2.5% outliers, per-row overlay
13381        let mut row_ptr = vec![0u32; rows + 1];
13382        let mut idx = 0usize;
13383        while idx < n {
13384            row_ptr[idx / cols + 1] += 1;
13385            idx += stride;
13386        }
13387        for r in 0..rows {
13388            row_ptr[r + 1] += row_ptr[r];
13389        }
13390        for &p in &row_ptr {
13391            bytes.extend_from_slice(&p.to_le_bytes());
13392        }
13393        let mut idx = 0usize;
13394        while idx < n {
13395            bytes.extend_from_slice(&((idx % cols) as u16).to_le_bytes());
13396            bytes.extend_from_slice(&f32_to_f16((idx % 13) as f32 * 0.1 - 0.6).to_le_bytes());
13397            idx += stride;
13398        }
13399        // On-grid ±1 so the fast path's int8 SDOT is exact vs the f32 "slow"
13400        // reference (the A/B is a timing check; values must still agree).
13401        let x: Vec<f32> = (0..cols)
13402            .map(|j| if j % 3 == 0 { 1.0 } else { -1.0 })
13403            .collect();
13404        let (rp_off, ent_off, has_ov) = q1t_overlay(&bytes, rows * gpr * Q1T_TILE, rows);
13405
13406        // "before": base-3 division decode into a buffer, then dot.
13407        let slow = |out: &mut [f32]| {
13408            let mut buf = vec![0f32; cols];
13409            for r in 0..rows {
13410                for g in 0..gpr {
13411                    let off = (r * gpr + g) * Q1T_TILE;
13412                    let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
13413                    let codes = &bytes[off + 2..off + Q1T_TILE];
13414                    for k in 0..GROUP_SIZE {
13415                        buf[g * GROUP_SIZE + k] = match q1t_code(codes, k) {
13416                            1 => s,
13417                            2 => -s,
13418                            _ => 0.0,
13419                        };
13420                    }
13421                }
13422                out[r] = q1t_row_outlier_correction(&bytes, r, rp_off, ent_off, has_ov, &x)
13423                    + (0..cols).map(|j| buf[j] * x[j]).sum::<f32>();
13424            }
13425        };
13426        let iters = 5;
13427        let mut a = vec![0f32; rows];
13428        slow(&mut a); // warm
13429        let t = Instant::now();
13430        for _ in 0..iters {
13431            slow(&mut a);
13432        }
13433        let slow_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
13434
13435        let mut b = vec![0f32; rows];
13436        q1t_matvec(&bytes, &x, rows, cols, &mut b, None); // warm
13437        let t = Instant::now();
13438        for _ in 0..iters {
13439            q1t_matvec(&bytes, &x, rows, cols, &mut b, None);
13440        }
13441        let fast_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
13442
13443        for r in 0..rows {
13444            assert!((a[r] - b[r]).abs() < 1e-2, "mismatch row {r}");
13445        }
13446        println!(
13447            "q1t matvec {rows}x{cols} (1 thread): div-decode {slow_ms:.2} ms  fused-LUT {fast_ms:.2} ms  => {:.2}x",
13448            slow_ms / fast_ms
13449        );
13450    }
13451}
13452
13453#[cfg(test)]
13454mod gemm_bench {
13455    /// `cargo test -p cortiq-engine --release q4tp_matmat_throughput -- --ignored --nocapture`
13456    /// Times the batched q4tp GEMM at the shapes the image DiT runs
13457    /// (b=296 tokens, 2304 -> 9216), on synthetic bytes: no model, no
13458    /// mmap, no thermal drift over minutes — a kernel change shows up
13459    /// here in seconds where a full render hides it in noise.
13460    ///
13461    /// On macOS add `CMF_ACCEL=0`: this shape is over the 500k-cell mark
13462    /// where the matmat hands off to Accelerate's dequant sgemm, and
13463    /// without the opt-out both rows below measure the AMX, not the
13464    /// kernel under test.
13465    #[test]
13466    #[ignore]
13467    fn q4tp_matmat_throughput() {
13468        let _alt = super::Q4TP_ALT_TEST_LOCK
13469            .lock()
13470            .unwrap_or_else(|e| e.into_inner());
13471        // 296 is a prompt-encode batch; the image DiT runs 2085 at
13472        // 512x512, where the activation panel stops fitting L2 and the
13473        // loop's shape starts to matter more than its instructions.
13474        let b: usize = std::env::var("CMF_BENCH_B")
13475            .ok()
13476            .and_then(|v| v.parse().ok())
13477            .unwrap_or(296);
13478        let (rows, cols) = (9216usize, 2304usize);
13479        let (_, _, _) = (rows, cols, b);
13480        let total =
13481            cortiq_core::quant::expected_nbytes(cortiq_core::TensorDtype::Q4TiledP, &[rows, cols])
13482                .unwrap();
13483        // Random nibbles are fine, but the row params are f16 (lo, step)
13484        // of a geometric ladder: garbage there gives exp2 of a huge
13485        // exponent, the scales come back inf, and the whole bench times
13486        // NaN arithmetic instead of the kernel.
13487        let (params_off, codes_off, _) = cortiq_core::quant::q4tp_sections(rows, cols);
13488        let mut bytes: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
13489        let lo = cortiq_core::quant::f32_to_f16(-4.0);
13490        let step = cortiq_core::quant::f32_to_f16(0.1);
13491        for r in 0..rows {
13492            let o = params_off + r * 4;
13493            bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
13494            bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
13495        }
13496        let _ = codes_off;
13497        let xs: Vec<f32> = (0..b * cols)
13498            .map(|i| ((i % 97) as f32 - 48.0) / 48.0)
13499            .collect();
13500        let mut out = vec![0f32; b * rows];
13501        let pool = crate::pool::Pool::from_env();
13502        // A shared 48-core stand drifts ±25% run to run, which is wider
13503        // than any kernel change worth making. So: alternate the two
13504        // kernels inside one process and keep the BEST time for
13505        // each. Interleaving makes both see the same interference, and a
13506        // minimum is the one statistic another tenant cannot inflate.
13507        super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
13508        let reps: usize = std::env::var("CMF_BENCH_REPS")
13509            .ok()
13510            .and_then(|v| v.parse().ok())
13511            .unwrap_or(10);
13512        let mut best = [f64::MAX; 2];
13513        let mut sums = [0f32; 2];
13514        for _ in 0..reps {
13515            for (k, w) in [(0usize, 1u8), (1usize, 2u8)] {
13516                super::Q4TP_ALT.store(w, std::sync::atomic::Ordering::Relaxed);
13517                let t = std::time::Instant::now();
13518                super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
13519                best[k] = best[k].min(t.elapsed().as_secs_f64());
13520                sums[k] = out.iter().take(64).sum::<f32>();
13521            }
13522        }
13523        let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
13524        for (k, name) in ["previous", "tuned   "].iter().enumerate() {
13525            println!(
13526                "q4tp matmat {rows}x{cols} b={b} {name}: {:.1} ms  {:.1} GFLOP/s  (checksum {:.3})",
13527                best[k] * 1e3,
13528                flops / best[k] / 1e9,
13529                sums[k]
13530            );
13531        }
13532        assert!(
13533            (sums[0] - sums[1]).abs() < 1e-2,
13534            "the tuned kernel changed the result: {} vs {}",
13535            sums[0],
13536            sums[1]
13537        );
13538    }
13539
13540    /// The blocked kernel must agree with the per-column path exactly —
13541    /// same weights, same activation split, only a different instruction
13542    /// mix. Shapes are chosen to hit the awkward cases: a column count
13543    /// that leaves an odd group (the 512-bit kernel does two at a time),
13544    /// and a batch that does not divide by four.
13545    #[test]
13546    fn q4tp_matmat_blocked_matches_scalar() {
13547        use std::sync::atomic::Ordering::Relaxed;
13548        let _alt = super::Q4TP_ALT_TEST_LOCK
13549            .lock()
13550            .unwrap_or_else(|e| e.into_inner());
13551        // The last shape carries the image DiT's column count — 2304, so
13552        // 72 groups of accumulation, which is where a reordered sum can
13553        // actually drift — and runs through the thread pool, since the
13554        // blocked path splits rows across workers. Its row count stays
13555        // under 500k cells on purpose: above that, macOS diverts the whole
13556        // matmat to the Accelerate/AMX dequant sgemm and neither kernel
13557        // here would run.
13558        for &(rows, cols, b) in &[
13559            (64usize, 128usize, 7usize),
13560            (33, 96, 4),
13561            (16, 256, 9),
13562            (192, 2304, 37),
13563        ] {
13564            let total = cortiq_core::quant::expected_nbytes(
13565                cortiq_core::TensorDtype::Q4TiledP,
13566                &[rows, cols],
13567            )
13568            .unwrap();
13569            let (params_off, _, _) = cortiq_core::quant::q4tp_sections(rows, cols);
13570            let mut bytes: Vec<u8> = (0..total).map(|i| (i * 61 % 251) as u8).collect();
13571            let lo = cortiq_core::quant::f32_to_f16(-4.0);
13572            let step = cortiq_core::quant::f32_to_f16(0.1);
13573            for r in 0..rows {
13574                let o = params_off + r * 4;
13575                bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
13576                bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
13577            }
13578            let xs: Vec<f32> = (0..b * cols)
13579                .map(|i| ((i % 89) as f32 - 44.0) / 44.0)
13580                .collect();
13581            let mut got = vec![0f32; b * rows];
13582            let mut want = vec![0f32; b * rows];
13583            let gpr = cols / 32;
13584            let view = super::Q4tpView::new(&bytes, rows, cols);
13585            let pool = crate::pool::Pool::from_env();
13586            super::Q4TP_ALT.store(2, Relaxed);
13587            super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut got, pool.as_deref());
13588            super::Q4TP_ALT.store(1, Relaxed);
13589            super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut want, pool.as_deref());
13590            super::Q4TP_ALT.store(0, Relaxed);
13591            // Measured against the output's scale, not cell by cell: a
13592            // dot product of 2304 terms lands near zero wherever the row
13593            // and the activation nearly cancel, and there a per-cell
13594            // ratio reports 1e-3 for an absolute error of 5e-6 — f32's
13595            // own rounding, reordered. What must stay small is the error
13596            // relative to what the layer actually outputs.
13597            let scale = want.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
13598            let (mut worst, mut at) = (0f32, 0usize);
13599            for (i, (g, w)) in got.iter().zip(&want).enumerate() {
13600                if (g - w).abs() > worst {
13601                    worst = (g - w).abs();
13602                    at = i;
13603                }
13604            }
13605            assert!(
13606                worst <= 1e-4 * scale,
13607                "{rows}x{cols} b={b}: blocked and scalar disagree by {worst:.3e} \
13608                 (scale {scale:.3e}) at cell {at}: {} vs {}",
13609                got[at],
13610                want[at]
13611            );
13612
13613            // "Same speed, no quality loss" is a claim about which answer
13614            // is RIGHT, not about which two agree. Both paths sum the same
13615            // 2304 products in different orders, so f64 decides: the
13616            // blocked kernel keeps sixteen partial sums and folds them at
13617            // the end, which is a shallower addition tree than the
13618            // per-column path's running scalar, and it must not be worse.
13619            let (mut e_blocked, mut e_scalar) = (0f64, 0f64);
13620            for bi in 0..b {
13621                let act = super::split_act(&xs[bi * cols..(bi + 1) * cols]);
13622                for r in 0..rows {
13623                    let mut sc = vec![0f32; gpr];
13624                    view.scales_into(r, gpr, &mut sc);
13625                    let mut exact = 0f64;
13626                    for j in 0..cols {
13627                        let (w, sq) = super::q4tp_outlier(view.nib, r, gpr, j, &sc);
13628                        exact += w as f64 * sq as f64 * act.xq[j] as f64;
13629                    }
13630                    exact *= act.sx as f64;
13631                    for &(j, xv) in &act.outliers {
13632                        let (w, sq) = super::q4tp_outlier(view.nib, r, gpr, j, &sc);
13633                        exact += w as f64 * sq as f64 * xv as f64;
13634                    }
13635                    let i = bi * rows + r;
13636                    e_blocked = e_blocked.max((got[i] as f64 - exact).abs());
13637                    e_scalar = e_scalar.max((want[i] as f64 - exact).abs());
13638                }
13639            }
13640            println!(
13641                "{rows}x{cols} b={b}: worst error vs f64 — blocked {e_blocked:.3e}, \
13642                 per-column {e_scalar:.3e}"
13643            );
13644            // An absolute bar, not a race between the two: at these
13645            // magnitudes both sit in f32's last bits, and on a small shape
13646            // whichever one happens to round the unluckiest cell "wins" by
13647            // a factor the next seed reverses.
13648            assert!(
13649                e_blocked <= 1e-5 * scale as f64 && e_scalar <= 1e-5 * scale as f64,
13650                "{rows}x{cols} b={b}: error against f64 too large — blocked \
13651                 {e_blocked:.3e}, per-column {e_scalar:.3e}, scale {scale:.3e}"
13652            );
13653        }
13654    }
13655
13656    /// What the row-exact contract costs a speculative-verify panel on the
13657    /// host: the same batch through the fast arms, through the row-exact
13658    /// arms, and as one matvec per token (the other way to be exact).
13659    /// Arms alternate inside one process and keep their best time.
13660    /// `CMF_GPU=0 cargo test -p cortiq-engine --release q4tp_matmat_row_exact_cost -- --ignored --nocapture`
13661    #[test]
13662    #[ignore]
13663    fn q4tp_matmat_row_exact_cost() {
13664        let _alt = super::Q4TP_ALT_TEST_LOCK
13665            .lock()
13666            .unwrap_or_else(|e| e.into_inner());
13667        let pool = crate::pool::Pool::from_env();
13668        let reps: usize = std::env::var("CMF_BENCH_REPS")
13669            .ok()
13670            .and_then(|v| v.parse().ok())
13671            .unwrap_or(30);
13672        for &(rows, cols) in &[(2048usize, 4096usize), (4096, 2048), (4096, 4096)] {
13673            let total = cortiq_core::quant::expected_nbytes(
13674                cortiq_core::TensorDtype::Q4TiledP,
13675                &[rows, cols],
13676            )
13677            .unwrap();
13678            let (params_off, _, _) = cortiq_core::quant::q4tp_sections(rows, cols);
13679            let mut bytes: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
13680            let lo = cortiq_core::quant::f32_to_f16(-4.0);
13681            let step = cortiq_core::quant::f32_to_f16(0.1);
13682            for r in 0..rows {
13683                let o = params_off + r * 4;
13684                bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
13685                bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
13686            }
13687            for &b in &[2usize, 4, 5, 8] {
13688                let xs: Vec<f32> = (0..b * cols)
13689                    .map(|i| ((i % 97) as f32 - 48.0) / 48.0)
13690                    .collect();
13691                let mut out = vec![0f32; b * rows];
13692                let mut best = [f64::MAX; 3];
13693                for _ in 0..reps {
13694                    for (k, best_k) in best.iter_mut().enumerate() {
13695                        let t = std::time::Instant::now();
13696                        match k {
13697                            0 | 1 => super::q4tp_matmat_with(
13698                                &bytes,
13699                                &xs,
13700                                b,
13701                                rows,
13702                                cols,
13703                                &mut out,
13704                                pool.as_deref(),
13705                                k == 1,
13706                            ),
13707                            _ => {
13708                                for (bi, o) in out.chunks_mut(rows).enumerate() {
13709                                    super::q4tp_matvec(
13710                                        &bytes,
13711                                        &xs[bi * cols..(bi + 1) * cols],
13712                                        rows,
13713                                        cols,
13714                                        o,
13715                                        pool.as_deref(),
13716                                    );
13717                                }
13718                            }
13719                        }
13720                        *best_k = best_k.min(t.elapsed().as_secs_f64());
13721                    }
13722                }
13723                println!(
13724                    "q4tp {rows}x{cols} b={b}: fast {:.3} ms, row-exact {:.3} ms, \
13725                     {b} matvecs {:.3} ms",
13726                    best[0] * 1e3,
13727                    best[1] * 1e3,
13728                    best[2] * 1e3
13729                );
13730            }
13731        }
13732    }
13733
13734    /// The q4t twin of the throughput bench, same shape and rules, so the
13735    /// two quantisations' batch kernels can be read against each other.
13736    /// `cargo test -p cortiq-engine --release q4t_matmat_throughput -- --ignored --nocapture`
13737    #[test]
13738    #[ignore]
13739    fn q4t_matmat_throughput() {
13740        let (rows, cols, b) = (9216usize, 2304usize, 296usize);
13741        let total =
13742            cortiq_core::quant::expected_nbytes(cortiq_core::TensorDtype::Q4Tiled, &[rows, cols])
13743                .unwrap();
13744        // q4t carries a per-group f16 scale in the tile's first two bytes;
13745        // random bytes there decode to inf and the bench would time NaNs.
13746        let mut bytes: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
13747        let sc = cortiq_core::quant::f32_to_f16(0.02);
13748        for t in bytes.chunks_mut(super::Q4_TILE) {
13749            t[..2].copy_from_slice(&sc.to_le_bytes());
13750        }
13751        let xs: Vec<f32> = (0..b * cols)
13752            .map(|i| ((i % 97) as f32 - 48.0) / 48.0)
13753            .collect();
13754        let mut out = vec![0f32; b * rows];
13755        let pool = crate::pool::Pool::from_env();
13756        super::q4t_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
13757        let reps: usize = std::env::var("CMF_BENCH_REPS")
13758            .ok()
13759            .and_then(|v| v.parse().ok())
13760            .unwrap_or(10);
13761        let mut best = f64::MAX;
13762        for _ in 0..reps {
13763            let t = std::time::Instant::now();
13764            super::q4t_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
13765            best = best.min(t.elapsed().as_secs_f64());
13766        }
13767        let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
13768        println!(
13769            "q4t matmat {rows}x{cols} b={b}: {:.1} ms  {:.1} GFLOP/s  (checksum {:.3})",
13770            best * 1e3,
13771            flops / best / 1e9,
13772            out.iter().take(64).sum::<f32>()
13773        );
13774    }
13775}