Skip to main content

cortiq_engine/
qtensor.rs

1//! QTensor — weight tensor with pluggable storage.
2//!
3//! Two backings, one interface:
4//! - `F32`   — owned dense floats (small models, tests). Every operation
5//!   is bit-identical to the historical `&[f32]` code paths.
6//! - `Mapped` — quantized bytes zero-copy from the CMF mmap (`q8_row` /
7//!   `q8_2f`). The matvec is fused: int8 rows × f32 activations, the
8//!   q8_2f column field folds into a pre-scale of the input
9//!   (`x'[i] = col[i]·x[i]`), so the inner loop is the same i8 dot as
10//!   q8_row. This is what lets a 15B file run in a few GB of RSS.
11//!
12//! Extension point: new dtypes = new match arm here, nothing else moves.
13
14use crate::pool::{Pool, matvec_rows, matvec_rows2};
15use cortiq_core::quant::{
16    GROUP_SIZE, Q1_TILE, Q2TP_CHUNK, Q4_TILE, Q4TP_NIB, f16_to_f32, q2tp_ladder, q2tp_sections,
17    q4tp_code, q4tp_ladder, q4tp_sections,
18};
19use cortiq_core::{CmfModel, TensorDtype};
20use std::cell::UnsafeCell;
21use std::sync::Arc;
22
23/// `mxfp4` (OCP MX FP4) CPU kernels.
24mod mxfp4;
25
26pub enum QTensor {
27    F32 {
28        data: Vec<f32>,
29        rows: usize,
30        cols: usize,
31    },
32    Mapped {
33        model: Arc<CmfModel>,
34        /// Index into the model's tensor directory.
35        idx: usize,
36        dtype: TensorDtype,
37        rows: usize,
38        cols: usize,
39        /// Per-row scales, dequantized to f32 up front (tiny).
40        row_scale: Vec<f32>,
41        /// q8_2f column field (θ), dequantized up front; empty for q8_row.
42        col_field: Vec<f32>,
43        /// Vbit only: byte offset of each row's packed data within the
44        /// tensor blob (`[rows + 1]`, computed once at load — the per-
45        /// matvec prefix scan over row bit-widths was O(rows) each call).
46        vbit_offsets: Vec<usize>,
47        /// q8-family decode repack (load-time, optional): rows in groups
48        /// of 4, interleaved in 16-byte units — one 64-byte line per
49        /// iteration feeds all 4 sdot lanes, ONE sequential weight
50        /// stream per worker instead of four (this is where llama.cpp's
51        /// repacked Q8 kernels get their bandwidth). Empty = off
52        /// (CMF_REPACK=0, non-SDOT arch, or an ineligible shape). Trades
53        /// an anonymous copy of the quants for mmap pages that go cold.
54        repack: Vec<u8>,
55    },
56    /// Binary low-rank (`blr`): ±1 paths, signs from the mmap, scales as
57    /// f32 (docs/BLR_NVG.md §3). Every device-graph accessor answers `None`
58    /// for it (whole-token graphs decline to per-op/host paths); the CPU
59    /// kernels live in `crate::blr`, the per-op wgpu kernels in
60    /// `gpu_wgpu::blr`.
61    Blr(Box<crate::blr::BlrTensor>),
62}
63
64/// Load-time q8 repack gate (see `Mapped::repack`). OPT-IN
65/// (`CMF_REPACK=1`): the single-stream hypothesis LOST on Apple Silicon
66/// (M4, interleaved A/B: decode 101 vs 94 tok/s — four adjacent row
67/// streams per worker feed the prefetcher MORE memory-level parallelism
68/// than one); kept as an experiment flag for x86, where the tradeoff
69/// may land differently.
70fn repack_enabled() -> bool {
71    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
72    *ON.get_or_init(|| {
73        std::env::var("CMF_REPACK")
74            .map(|v| v == "1")
75            .unwrap_or(cfg!(target_os = "android"))
76    })
77}
78
79/// Interleave q8 rows for the decode kernel: group g holds rows
80/// 4g..4g+4 as [r0[c], r1[c], r2[c], r3[c]] per 16-byte chunk c. Only
81/// full groups are packed — tail rows keep reading the mmap layout.
82fn q8_repack(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
83    #[cfg(target_arch = "aarch64")]
84    let arch_ok = sdot_enabled();
85    #[cfg(not(target_arch = "aarch64"))]
86    let arch_ok = false;
87    if !arch_ok || !repack_enabled() || rows < 256 || cols % 16 != 0 {
88        return Vec::new();
89    }
90    q8_repack_layout(bytes, rows, cols)
91}
92
93/// The pure layout transform behind `q8_repack` (tested directly —
94/// the gate depends on arch and env).
95fn q8_repack_layout(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
96    let groups = rows / 4;
97    let mut rep = vec![0u8; groups * 4 * cols];
98    for g in 0..groups {
99        let dst = &mut rep[g * 4 * cols..(g + 1) * 4 * cols];
100        for c in 0..cols / 16 {
101            for lane in 0..4 {
102                let src = (g * 4 + lane) * cols + c * 16;
103                dst[c * 64 + lane * 16..c * 64 + lane * 16 + 16]
104                    .copy_from_slice(&bytes[src..src + 16]);
105            }
106        }
107    }
108    rep
109}
110
111/// Prefix-sum of vbit row payload offsets (absolute within the tensor
112/// bytes). `offsets[r]..offsets[r+1]` is row r's packed data.
113fn vbit_row_offsets(bytes: &[u8], rows: usize, cols: usize) -> Vec<usize> {
114    let ng = cols / GROUP_SIZE;
115    let bits = &bytes[..rows];
116    let mut offsets = Vec::with_capacity(rows + 1);
117    let mut off = rows + rows * ng * 2;
118    for r in 0..rows {
119        offsets.push(off);
120        off += (cols * bits[r] as usize).div_ceil(8);
121    }
122    offsets.push(off);
123    offsets
124}
125
126/// `CMF_X86_BLOCKED` / `CMF_GPU_LMHEAD` / `CMF_GPU_SPLIT`, read once. They
127/// used to be read from the environment on every large matvec and on every
128/// matmat in six places — microseconds each, but also a knob that could
129/// change under a running process, which is not a thing a kernel choice
130/// should be able to do mid-sequence.
131fn blocked_enabled() -> bool {
132    use std::sync::atomic::Ordering::Relaxed;
133    match BLOCKED_OVERRIDE.load(Relaxed) {
134        1 => false,
135        2 => true,
136        _ => {
137            static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
138            *ON.get_or_init(|| {
139                std::env::var("CMF_X86_BLOCKED")
140                    .map(|v| v != "0")
141                    .unwrap_or(true)
142            })
143        }
144    }
145}
146
147static BLOCKED_OVERRIDE: std::sync::atomic::AtomicU8 = std::sync::atomic::AtomicU8::new(0);
148
149/// Force the blocked GEMM on or off, ignoring the environment; `None`
150/// restores it. For tests that need to run BOTH paths and compare them:
151/// `blocked_enabled` caches its answer for the life of the process, which
152/// is right when the environment is the only input, but leaves a test that
153/// flips `CMF_X86_BLOCKED` between two calls comparing a path against
154/// itself — or against whatever a test running in parallel latched first.
155pub fn set_blocked_override(on: Option<bool>) {
156    let v = match on {
157        None => 0,
158        Some(false) => 1,
159        Some(true) => 2,
160    };
161    BLOCKED_OVERRIDE.store(v, std::sync::atomic::Ordering::Relaxed);
162}
163
164fn gpu_lmhead_enabled() -> bool {
165    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
166    *ON.get_or_init(|| {
167        std::env::var("CMF_GPU_LMHEAD")
168            .map(|v| v != "0")
169            .unwrap_or(true)
170    })
171}
172
173fn gpu_split_frac() -> f32 {
174    // MiMo's banked verification projects every row on the device. Plain
175    // decode must not quantize half of its O/head activation on the CPU.
176    if FULL_GPU_Q8.get() {
177        return 1.0;
178    }
179    static F: std::sync::OnceLock<f32> = std::sync::OnceLock::new();
180    *F.get_or_init(|| {
181        std::env::var("CMF_GPU_SPLIT")
182            .ok()
183            .and_then(|v| v.parse::<f32>().ok())
184            .unwrap_or(0.5)
185            .clamp(0.0, 1.0)
186    })
187}
188
189impl QTensor {
190    pub fn from_f32(data: Vec<f32>, rows: usize, cols: usize) -> Self {
191        debug_assert_eq!(data.len(), rows * cols);
192        Self::F32 { data, rows, cols }
193    }
194
195    /// Wrap a directory tensor without dequantizing the payload.
196    /// Falls back to dequantized f32 for dtypes without a fused kernel.
197    pub fn from_model(model: &Arc<CmfModel>, name: &str) -> Result<Self, String> {
198        // Indexed lookup: the linear directory scan made pipeline build
199        // O(N²) on MoE/skills files with thousands of tensors.
200        let idx = model
201            .tensor_index(name)
202            .ok_or_else(|| format!("tensor '{name}' not found in CMF directory"))?;
203        let entry = &model.tensors[idx];
204        if entry.shape.len() != 2 {
205            return Err(format!("QTensor::from_model needs 2-D, got '{name}'"));
206        }
207        let (rows, cols) = (entry.shape[0], entry.shape[1]);
208        let bytes = model.entry_bytes(entry);
209
210        match entry.dtype {
211            TensorDtype::Q8Row | TensorDtype::Q8_2f => {
212                let n = rows * cols;
213                let scales_off = n;
214                let row_scale: Vec<f32> = (0..rows)
215                    .map(|o| {
216                        f16_to_f32(u16::from_le_bytes([
217                            bytes[scales_off + o * 2],
218                            bytes[scales_off + o * 2 + 1],
219                        ]))
220                    })
221                    .collect();
222                let col_field: Vec<f32> = if entry.dtype == TensorDtype::Q8_2f {
223                    let col_off = n + rows * 2;
224                    (0..cols)
225                        .map(|i| {
226                            f16_to_f32(u16::from_le_bytes([
227                                bytes[col_off + i * 2],
228                                bytes[col_off + i * 2 + 1],
229                            ]))
230                        })
231                        .collect()
232                } else {
233                    Vec::new()
234                };
235                Ok(Self::Mapped {
236                    model: model.clone(),
237                    idx,
238                    dtype: entry.dtype,
239                    rows,
240                    cols,
241                    row_scale,
242                    col_field,
243                    vbit_offsets: Vec::new(),
244                    repack: q8_repack(bytes, rows, cols),
245                })
246            }
247            // vbit: fused kernel unpacks variable-bit rows from mmap.
248            TensorDtype::Vbit if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
249                model: model.clone(),
250                idx,
251                dtype: entry.dtype,
252                rows,
253                cols,
254                row_scale: Vec::new(),
255                col_field: Vec::new(),
256                vbit_offsets: vbit_row_offsets(bytes, rows, cols),
257                repack: Vec::new(),
258            }),
259            // vbit_ro (§4.2): the offset table comes straight from the
260            // file — no load-time prefix scan; kernels are shared with
261            // legacy vbit (they consume absolute offsets either way).
262            TensorDtype::VbitRo if cols % GROUP_SIZE == 0 => {
263                let (_, off_off, packed_off) = cortiq_core::quant::vbit_ro_sections(rows, cols);
264                let offsets: Vec<usize> = (0..=rows)
265                    .map(|r| packed_off + cortiq_core::quant::vbit_ro_offset(bytes, off_off, r))
266                    .collect();
267                Ok(Self::Mapped {
268                    model: model.clone(),
269                    idx,
270                    dtype: entry.dtype,
271                    rows,
272                    cols,
273                    row_scale: Vec::new(),
274                    col_field: Vec::new(),
275                    vbit_offsets: offsets,
276                    repack: Vec::new(),
277                })
278            }
279            // q4_block: fused kernel reads nibbles straight from mmap —
280            // a 14B q4 file no longer explodes into ×8 f32 RAM.
281            // q4_tiled (§4.3): interleaved [scale][nibbles] tiles — one
282            // sequential memory stream (measured ×1.66 ARM / ×1.13 AVX2
283            // at kernel level over the split layout).
284            TensorDtype::Q4Tiled if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
285                model: model.clone(),
286                idx,
287                dtype: entry.dtype,
288                rows,
289                cols,
290                row_scale: Vec::new(),
291                col_field: Vec::new(),
292                vbit_offsets: Vec::new(),
293                repack: Vec::new(),
294            }),
295            // q4tp (§4.10): nibbles from mmap, scale from the row ladder —
296            // 7.3% less file than q4t at the same 4-bit grid.
297            TensorDtype::Q4TiledP if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
298                model: model.clone(),
299                idx,
300                dtype: entry.dtype,
301                rows,
302                cols,
303                row_scale: Vec::new(),
304                col_field: Vec::new(),
305                vbit_offsets: Vec::new(),
306                repack: Vec::new(),
307            }),
308            // q2tp: 2-bit chunks from mmap, scale from the same row ladder.
309            TensorDtype::Q2TiledP if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
310                model: model.clone(),
311                idx,
312                dtype: entry.dtype,
313                rows,
314                cols,
315                row_scale: Vec::new(),
316                col_field: Vec::new(),
317                vbit_offsets: Vec::new(),
318                repack: Vec::new(),
319            }),
320            TensorDtype::Q4Block if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
321                model: model.clone(),
322                idx,
323                dtype: entry.dtype,
324                rows,
325                cols,
326                row_scale: Vec::new(),
327                col_field: Vec::new(),
328                vbit_offsets: Vec::new(),
329                repack: Vec::new(),
330            }),
331            // q1: binary sign-bit tiles from mmap (1-bit-trained models).
332            // q1g: the same signs as planes with one scale per 128-group.
333            TensorDtype::Q1 if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
334                model: model.clone(),
335                idx,
336                dtype: entry.dtype,
337                rows,
338                cols,
339                row_scale: Vec::new(),
340                col_field: Vec::new(),
341                vbit_offsets: Vec::new(),
342                repack: Vec::new(),
343            }),
344            TensorDtype::Q1G if cols % cortiq_core::quant::Q1G_GROUP == 0 => Ok(Self::Mapped {
345                model: model.clone(),
346                idx,
347                dtype: entry.dtype,
348                rows,
349                cols,
350                row_scale: Vec::new(),
351                col_field: Vec::new(),
352                vbit_offsets: Vec::new(),
353                repack: Vec::new(),
354            }),
355            // mxfp4: E2M1 codes + E8M0 scales straight from the mmap (the
356            // native FP4 expert codec; `qtensor::mxfp4`).
357            TensorDtype::Mxfp4 if cols % cortiq_core::mxfp4::GROUP == 0 => Ok(Self::Mapped {
358                model: model.clone(),
359                idx,
360                dtype: entry.dtype,
361                rows,
362                cols,
363                row_scale: Vec::new(),
364                col_field: Vec::new(),
365                vbit_offsets: Vec::new(),
366                repack: Vec::new(),
367            }),
368            // q1t (ternary + outlier overlay): fused per-row dequant kernel
369            // reads straight from mmap — a 12B q1t stays ~its file size in
370            // RAM instead of dequantizing to ~48 GB of f32.
371            TensorDtype::Q1T if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
372                model: model.clone(),
373                idx,
374                dtype: entry.dtype,
375                rows,
376                cols,
377                row_scale: Vec::new(),
378                col_field: Vec::new(),
379                vbit_offsets: Vec::new(),
380                repack: Vec::new(),
381            }),
382            // blr: factored kernels over the packed signs. CMF_BLR_DENSE=1
383            // dequantizes to f32 instead (A/B parity runs).
384            TensorDtype::BinaryLowRank if std::env::var("CMF_BLR_DENSE").as_deref() != Ok("1") => {
385                Ok(Self::Blr(Box::new(crate::blr::BlrTensor::from_model(
386                    model, idx,
387                )?)))
388            }
389            // No fused kernel yet → dequantize once (correct, more RAM).
390            _ => {
391                let mut data = vec![0.0f32; rows * cols];
392                cortiq_core::quant::dequant_tensor(entry, bytes, &mut data)?;
393                Ok(Self::from_f32(data, rows, cols))
394            }
395        }
396    }
397
398    /// Binary low-rank tensor? Fused multi-tensor kernels refuse it and fall
399    /// back to per-tensor calls.
400    pub fn is_blr(&self) -> bool {
401        matches!(self, Self::Blr(_))
402    }
403
404    /// Memory-mapped `mxfp4` tensor? Like `blr`, it has no row-range kernel
405    /// in the fused multi-tensor jobs: those run it per tensor.
406    pub fn is_mxfp4(&self) -> bool {
407        matches!(
408            self,
409            Self::Mapped {
410                dtype: TensorDtype::Mxfp4,
411                ..
412            }
413        )
414    }
415
416    /// Codecs the fused multi-tensor jobs (`matvec_many`, the SwiGLU and
417    /// MoE row kernels) do not carry: they fall back to per-tensor calls.
418    fn per_tensor_only(&self) -> bool {
419        self.is_blr() || self.is_mxfp4()
420    }
421
422    /// (model, idx) of a `blr` tensor, for the wgpu graphs only (their
423    /// `blr` stage kernels, kind [`crate::gpu::GRAPH_KIND_BLR`]). Kept apart
424    /// from [`graph_weight`](Self::graph_weight), which stays None for
425    /// `blr`: every other graph consumer keys its kernels on that answer.
426    /// A file with a Prism contract keeps the per-op path.
427    pub fn graph_blr(&self) -> Option<(&Arc<CmfModel>, usize)> {
428        match self {
429            Self::Blr(t) if !crate::prism::has_contract(&t.model) => Some((&t.model, t.idx)),
430            _ => None,
431        }
432    }
433
434    /// [`graph_weight`](Self::graph_weight), then the Prism descriptor
435    /// form, then a `blr` tensor as kind [`crate::gpu::GRAPH_KIND_BLR`] —
436    /// the wgpu token and batch graph builders' lookup.
437    pub(crate) fn graph_weight_or_blr(
438        &self,
439    ) -> Option<(&std::sync::Arc<CmfModel>, usize, u8, &[f32])> {
440        self.graph_weight()
441            .or_else(|| self.graph_weight_descriptor())
442            .or_else(|| {
443                self.graph_blr()
444                    .map(|(m, i)| (m, i, crate::gpu::GRAPH_KIND_BLR, &[][..]))
445            })
446    }
447
448    /// Binary-weight-mapped tensor, q1 or q1g? (GPU gates: the binary CPU
449    /// kernels are compute-bound, so offload pays at much smaller shapes
450    /// than q8.)
451    pub(crate) fn is_q1(&self) -> bool {
452        matches!(
453            self,
454            Self::Mapped {
455                dtype: TensorDtype::Q1 | TensorDtype::Q1G,
456                ..
457            }
458        )
459    }
460
461    /// The binary-weight view (q1 tiles or q1g planes) the shared CPU row
462    /// drivers take; None for every other storage.
463    fn bin_w(&self) -> Option<BinW<'_>> {
464        match self {
465            Self::Mapped {
466                dtype: TensorDtype::Q1,
467                ..
468            } => Some(BinW::Q1(self.quant_bytes())),
469            Self::Mapped {
470                dtype: TensorDtype::Q1G,
471                rows,
472                cols,
473                ..
474            } => Some(BinW::q1g(self.quant_bytes(), *rows, *cols)),
475            _ => None,
476        }
477    }
478
479    /// Owned-f32 view (data, rows, cols) — the GDN a/b gate projections
480    /// arrive dequantized (force-f16 in the converter → F32 in RAM).
481    pub(crate) fn f32_parts(&self) -> Option<(&[f32], usize, usize)> {
482        match self {
483            Self::F32 { data, rows, cols } => Some((data, *rows, *cols)),
484            _ => None,
485        }
486    }
487
488    /// (directory idx, rows, cols) of a q1-mapped tensor — the
489    /// whole-block GPU path resolves offsets itself.
490    /// (idx, rows, cols) of a mapped tensor the whole-token GPU graph can drive
491    /// — Q1, Q1T or Q4-block (it resolves the offset and picks the kernel by
492    /// dtype). Q4-block lets a precise down_proj/lm_head stay on-device.
493    /// Named `q1_parts` for historical reasons.
494    pub(crate) fn q1_parts(&self) -> Option<(usize, usize, usize)> {
495        if self.has_prism_contract() {
496            return None;
497        }
498        match self {
499            #[cfg(target_os = "macos")]
500            Self::Mapped {
501                dtype: TensorDtype::Q1T,
502                ..
503            } if !crate::gpu::metal_q1t_enabled() => None,
504            Self::Mapped {
505                idx,
506                dtype:
507                    TensorDtype::Q1
508                    | TensorDtype::Q1G
509                    | TensorDtype::Q1T
510                    | TensorDtype::Q4Block
511                    | TensorDtype::Q4Tiled
512                    // Q2TiledP deliberately absent: the Metal graph has no
513                    // q2tp kernel, and advertising it here made the block
514                    // plan truncate mid-run at the first q2tp layer.
515                    | TensorDtype::Q4TiledP
516                    | TensorDtype::Q8Row
517                    | TensorDtype::Q8_2f,
518                rows,
519                cols,
520                ..
521            } => Some((*idx, *rows, *cols)),
522            _ => None,
523        }
524    }
525
526    /// `(directory idx, rows, cols)` for the native Metal token graph.  The
527    /// historical q1 graph gate intentionally refuses every Prism tensor so
528    /// an untransformed q2 payload cannot slip into the resident path.  The
529    /// Metal2 graph is descriptor-aware and admits only the production
530    /// q2tp-affine forward targets; ordinary q1/q4 callers retain the old
531    /// `q1_parts` behaviour.
532    #[cfg(target_os = "macos")]
533    pub(crate) fn metal_graph_parts(&self) -> Option<(usize, usize, usize)> {
534        if let Some((model, idx, kind, _)) = self.graph_weight_descriptor() {
535            let name = &model.tensors[idx].name;
536            let forward = kind == 9 && crate::prism::is_forward_weight(model, name);
537            let affine = kind == 9 && crate::prism::is_affine_target(model, name);
538            if forward && affine {
539                let e = model.tensors.get(idx)?;
540                return Some((idx, *e.shape.first()?, *e.shape.get(1)?));
541            }
542        }
543        self.q1_parts()
544    }
545
546    /// (directory idx, rows, cols) of a q4_tiled mapped tensor. The
547    /// chunk-prefill graph takes it in the same 4-tuple slot as
548    /// `q8_row_parts` with an EMPTY row_scale — q4t carries its scales
549    /// inside the 18-byte tiles, and the empty slice is what tells the
550    /// encoder to reach for the q4t kernels.
551    pub(crate) fn q4t_parts(&self) -> Option<(usize, usize, usize)> {
552        if self.has_prism_contract() {
553            return None;
554        }
555        match self {
556            Self::Mapped {
557                idx,
558                dtype: TensorDtype::Q4Tiled,
559                rows,
560                cols,
561                ..
562            } => Some((*idx, *rows, *cols)),
563            _ => None,
564        }
565    }
566
567    /// (directory idx, rows, cols) of a q4tp mapped tensor. Same empty-scale
568    /// slot as `q4t_parts` in the chunk graph — the encoder tells the two
569    /// apart by the tensor's dtype, not by the slot.
570    pub(crate) fn q4tp_parts(&self) -> Option<(usize, usize, usize)> {
571        if self.has_prism_contract() {
572            return None;
573        }
574        match self {
575            Self::Mapped {
576                idx,
577                dtype: TensorDtype::Q4TiledP,
578                rows,
579                cols,
580                ..
581            } => Some((*idx, *rows, *cols)),
582            _ => None,
583        }
584    }
585
586    /// (directory idx, rows, cols, row_scale) of a plain q8_row mapped
587    /// tensor — the chunk-prefill GPU graph resolves offsets itself.
588    /// q8_2f is excluded on purpose: its column field would need a
589    /// prescale stage on the device.
590    pub(crate) fn q8_row_parts(&self) -> Option<(usize, usize, usize, &[f32])> {
591        if self.has_prism_contract() {
592            return None;
593        }
594        match self {
595            Self::Mapped {
596                idx,
597                dtype: TensorDtype::Q8Row,
598                rows,
599                cols,
600                row_scale,
601                col_field,
602                ..
603            } if col_field.is_empty() => Some((*idx, *rows, *cols, row_scale)),
604            _ => None,
605        }
606    }
607
608    /// (directory idx, rows, cols, row_scale, col_field) of a q8_2f mapped
609    /// tensor — the chunk-prefill graph stages x·col on the device first.
610    pub(crate) fn q8_2f_parts(&self) -> Option<(usize, usize, usize, &[f32], &[f32])> {
611        if self.has_prism_contract() {
612            return None;
613        }
614        match self {
615            Self::Mapped {
616                idx,
617                dtype: TensorDtype::Q8_2f,
618                rows,
619                cols,
620                row_scale,
621                col_field,
622                ..
623            } if col_field.len() == *cols => Some((*idx, *rows, *cols, row_scale, col_field)),
624            _ => None,
625        }
626    }
627
628    /// The layout this tensor is stored in, when it is mapped from a model.
629    /// The frames branch on it — a q2tp gate against a q4tp down is a real
630    /// combination in the 2-bit profile and needs a different kernel.
631    pub fn model_dtype(&self) -> Option<cortiq_core::TensorDtype> {
632        match self {
633            Self::Mapped { dtype, .. } => Some(*dtype),
634            _ => None,
635        }
636    }
637
638    /// The tensor's index in the model directory, when it is mapped from one.
639    /// The GPU frames bind by index rather than by name — a name lookup per
640    /// layer per token is not free, and the index is what the device cache is
641    /// keyed on anyway.
642    pub fn model_idx(&self) -> Option<usize> {
643        match self {
644            Self::Mapped { idx, .. } => Some(*idx),
645            _ => None,
646        }
647    }
648
649    /// The model this tensor is mapped from, when it is mapped at all. The
650    /// GPU frames need the container to reach the bytes; a QTensor already
651    /// holds it, and threading a second handle down every call site to say
652    /// the same thing invites the two to disagree.
653    pub fn model_arc(&self) -> Option<std::sync::Arc<cortiq_core::CmfModel>> {
654        match self {
655            Self::Mapped { model, .. } => Some(model.clone()),
656            _ => None,
657        }
658    }
659
660    /// Whether this mapped tensor belongs to the Prism/Bonsai transform
661    /// contract.  Device graphs do not carry the descriptor, so callers use
662    /// this conservative predicate to stay on the descriptor-aware CPU path
663    /// instead of silently executing an unrotated matrix.
664    pub(crate) fn has_prism_contract(&self) -> bool {
665        matches!(self, Self::Mapped { model, .. } if crate::prism::has_contract(model))
666    }
667
668    pub fn rows(&self) -> usize {
669        match self {
670            Self::F32 { rows, .. } | Self::Mapped { rows, .. } => *rows,
671            Self::Blr(t) => t.rows,
672        }
673    }
674
675    /// Mapped q4t handle (model + directory index) — the fused GPU FFN
676    /// needs the raw file coordinates of its three projections.
677    pub(crate) fn mapped_q4t(&self) -> Option<(&Arc<CmfModel>, usize)> {
678        if self.has_prism_contract() {
679            return None;
680        }
681        match self {
682            Self::Mapped {
683                model,
684                idx,
685                dtype: TensorDtype::Q4Tiled,
686                ..
687            } => Some((model, *idx)),
688            _ => None,
689        }
690    }
691
692    /// Same slot as `mapped_q4t` for a q4tp tensor — the fused DiT FFN picks
693    /// its kernels by which of the two answers.
694    pub fn mapped_q4tp(&self) -> Option<(&Arc<CmfModel>, usize)> {
695        if self.has_prism_contract() {
696            return None;
697        }
698        match self {
699            Self::Mapped {
700                model,
701                idx,
702                dtype: TensorDtype::Q4TiledP,
703                ..
704            } => Some((model, *idx)),
705            _ => None,
706        }
707    }
708
709    /// (model, tensor idx) for a mapped weight in ANY codec the fused device
710    /// paths can run — four-bit tiled or either int8 layout.
711    ///
712    /// The fused DiT chains asked for `mapped_q4tp` by name, so an eight-bit
713    /// container never reached them and rendered through per-op GEMMs even
714    /// after those kernels learned its codec. The gate is what the codec has
715    /// a device GEMM for, not which codec it is.
716    pub fn mapped_device_gemm(&self) -> Option<(&Arc<CmfModel>, usize)> {
717        if self.has_prism_contract() {
718            return None;
719        }
720        match self {
721            Self::Mapped {
722                model,
723                idx,
724                dtype: TensorDtype::Q4TiledP | TensorDtype::Q8Row | TensorDtype::Q8_2f,
725                ..
726            } => Some((model, *idx)),
727            _ => None,
728        }
729    }
730
731    /// (model, tensor idx) for a q8_2f mapped weight — the eight-bit slot of
732    /// the MoE graph ladder (int8 body, f16 row scales, f16 column field).
733    pub fn mapped_q8_2f(&self) -> Option<(&Arc<CmfModel>, usize)> {
734        if self.has_prism_contract() {
735            return None;
736        }
737        match self {
738            Self::Mapped {
739                model,
740                idx,
741                dtype: TensorDtype::Q8_2f,
742                ..
743            } => Some((model, *idx)),
744            _ => None,
745        }
746    }
747
748    /// (model, tensor idx) for a q2tp mapped weight — the 2-bit twin of
749    /// `mapped_q4tp`, used by the mixed MoE profile.
750    pub fn mapped_q2tp(&self) -> Option<(&Arc<CmfModel>, usize)> {
751        if self.has_prism_contract() {
752            return None;
753        }
754        match self {
755            Self::Mapped {
756                model,
757                idx,
758                dtype: TensorDtype::Q2TiledP,
759                ..
760            } => Some((model, *idx)),
761            _ => None,
762        }
763    }
764
765    pub fn cols(&self) -> usize {
766        match self {
767            Self::F32 { cols, .. } | Self::Mapped { cols, .. } => *cols,
768            Self::Blr(t) => t.cols,
769        }
770    }
771
772    /// (model, tensor idx) for a q1 mapped weight — the wgpu token graph
773    /// keys its resident VRAM cache by idx. None for any other dtype/kind.
774    pub fn mapped_q1(&self) -> Option<(&std::sync::Arc<CmfModel>, usize)> {
775        if self.has_prism_contract() {
776            return None;
777        }
778        match self {
779            Self::Mapped {
780                model,
781                idx,
782                dtype: TensorDtype::Q1,
783                ..
784            } => Some((model, *idx)),
785            _ => None,
786        }
787    }
788
789    /// (model, idx, kind, row_scale) for a graph-capable mapped weight.
790    /// kind: 0=q8_row (per-row scales), 1=q1, 2=q4_block, 3=q1t
791    /// (tile-embedded, no rs), 5=q4_tiled, 6=q4tp, 7=q8_2f (both scale
792    /// planes live inside the tensor), 9=q2tp, 10=q1g (sign + scale
793    /// planes). None only for `vbit`.
794    ///
795    /// The old comment here claimed q4_block was unhandled while the arm
796    /// right below mapped it, and it named q8_2f as unhandled after that
797    /// stopped being true — a stale comment on this function is how a
798    /// model silently loses the graph, so it is worth keeping honest.
799    pub fn graph_weight(&self) -> Option<(&std::sync::Arc<CmfModel>, usize, u8, &[f32])> {
800        if self.has_prism_contract() {
801            return None;
802        }
803        self.graph_weight_descriptor()
804    }
805
806    /// Descriptor-aware graph handle used only by the Prism token graph.
807    /// Ordinary graph callers continue to use [`graph_weight`] and therefore
808    /// remain fail-closed until they provide the same explicit transform
809    /// contract.
810    pub(crate) fn graph_weight_descriptor(
811        &self,
812    ) -> Option<(&std::sync::Arc<CmfModel>, usize, u8, &[f32])> {
813        match self {
814            Self::Mapped {
815                model,
816                idx,
817                dtype: TensorDtype::Q8Row,
818                row_scale,
819                ..
820            } => Some((model, *idx, 0, row_scale.as_slice())),
821            Self::Mapped {
822                model,
823                idx,
824                dtype: TensorDtype::Q1,
825                ..
826            } => Some((model, *idx, 1, &[])),
827            // Kind 10, never 1: q1g's sign and scale planes are a different
828            // walk from q1's 6-byte tiles (the kind-2/5 mixup taught that).
829            Self::Mapped {
830                model,
831                idx,
832                dtype: TensorDtype::Q1G,
833                ..
834            } => Some((model, *idx, 10, &[])),
835            // Q4Tiled is kind 5, NOT 2: both carried 2 historically, and
836            // the wgpu token graph fed 18B interleaved tiles to the
837            // split-layout q4b kernel — garbage output on q4t models
838            // (caught by an end-to-end answer check on real Vulkan).
839            Self::Mapped {
840                model,
841                idx,
842                dtype: TensorDtype::Q4Tiled,
843                ..
844            } => Some((model, *idx, 5, &[])),
845            // Kind 6, not 5: q4tp's nibble stride and scale planes differ,
846            // and feeding them to the q4t kernel is exactly the mistake that
847            // produced garbage when Q4Tiled shared kind 2 with Q4Block.
848            Self::Mapped {
849                model,
850                idx,
851                dtype: TensorDtype::Q4TiledP,
852                ..
853            } => Some((model, *idx, 6, &[])),
854            Self::Mapped {
855                model,
856                idx,
857                dtype: TensorDtype::Q4Block,
858                ..
859            } => Some((model, *idx, 2, &[])),
860            // q8_2f carries BOTH scale planes after the int8 body (rows
861            // f16, then cols f16), so the graph takes the whole tensor
862            // and the kernel reads them where they lie — no host-side
863            // prescale, which is what the per-op path does instead.
864            Self::Mapped {
865                model,
866                idx,
867                dtype: TensorDtype::Q8_2f,
868                ..
869            } => Some((model, *idx, 7, &[])),
870            Self::Mapped {
871                model,
872                idx,
873                dtype: TensorDtype::Q1T,
874                ..
875            } => Some((model, *idx, 3, &[])),
876            // Kind 9: the 2-bit plane on the q4tp ladder (dense FFN gate/up
877            // of the q2tp profile). Its own kernel — 8 bytes a group where
878            // q4tp has 16, and rung 0 is the exact zero.
879            Self::Mapped {
880                model,
881                idx,
882                dtype: TensorDtype::Q2TiledP,
883                ..
884            } => Some((model, *idx, 9, &[])),
885            _ => None,
886        }
887    }
888
889    /// Dense f32 view — only for owned tensors. Masked/sparse execution
890    /// paths require it; quantized weights don't support masks yet.
891    pub fn as_f32(&self) -> Option<&[f32]> {
892        match self {
893            Self::F32 { data, .. } => Some(data),
894            Self::Mapped { .. } | Self::Blr(_) => None,
895        }
896    }
897
898    fn quant_bytes(&self) -> &[u8] {
899        match self {
900            Self::Mapped { model, idx, .. } => model.entry_bytes(&model.tensors[*idx]),
901            Self::Blr(t) => t.bytes(),
902            Self::F32 { .. } => unreachable!("quant_bytes on F32"),
903        }
904    }
905
906    /// Dequantize one row into `dst` (embedding lookup).
907    pub fn row_f32(&self, r: usize, dst: &mut [f32]) {
908        let cols = self.cols();
909        debug_assert_eq!(dst.len(), cols);
910        match self {
911            Self::F32 { data, .. } => dst.copy_from_slice(&data[r * cols..(r + 1) * cols]),
912            Self::Blr(t) => t.row_f32(r, dst),
913            Self::Mapped {
914                model,
915                idx,
916                dtype,
917                row_scale,
918                col_field,
919                vbit_offsets,
920                ..
921            } => {
922                if *dtype == TensorDtype::Q4Tiled {
923                    let bytes = self.quant_bytes();
924                    let gpr = cols / GROUP_SIZE;
925                    for gi in 0..gpr {
926                        let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
927                        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
928                        for (k, &b) in tile[2..].iter().enumerate() {
929                            dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
930                            dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
931                        }
932                    }
933                    if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
934                        crate::prism::inverse_embedding(model, dst);
935                    }
936                    return;
937                }
938                if *dtype == TensorDtype::Q4TiledP {
939                    let bytes = self.quant_bytes();
940                    let gpr = cols / GROUP_SIZE;
941                    let v = Q4tpView::new(bytes, self.rows(), cols);
942                    let mut sc = vec![0f32; gpr];
943                    v.scales_into(r, gpr, &mut sc);
944                    for gi in 0..gpr {
945                        let tile = &v.nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
946                        let s = sc[gi];
947                        for (k, &b) in tile.iter().enumerate() {
948                            dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
949                            dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
950                        }
951                    }
952                    if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
953                        crate::prism::inverse_embedding(model, dst);
954                    }
955                    return;
956                }
957                if *dtype == TensorDtype::Q2TiledP {
958                    let bytes = self.quant_bytes();
959                    let gpr = cols / GROUP_SIZE;
960                    let v = Q4tpView::new_q2(bytes, self.rows(), cols);
961                    let mut sc = vec![0f32; gpr];
962                    v.scales_into(r, gpr, &mut sc);
963                    for gi in 0..gpr {
964                        let ch =
965                            &v.nib[(r * gpr + gi) * Q2TP_CHUNK..(r * gpr + gi + 1) * Q2TP_CHUNK];
966                        let s = sc[gi];
967                        for (k, &b) in ch.iter().enumerate() {
968                            for j in 0..4 {
969                                let center = if crate::prism::is_affine_target(
970                                    model,
971                                    &model.tensors[*idx].name,
972                                ) {
973                                    1.0
974                                } else {
975                                    1.5
976                                };
977                                dst[gi * GROUP_SIZE + k * 4 + j] =
978                                    (((b >> (2 * j)) & 3) as f32 - center) * s;
979                            }
980                        }
981                    }
982                    if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
983                        crate::prism::inverse_embedding(model, dst);
984                    }
985                    return;
986                }
987                if *dtype == TensorDtype::Q4Block {
988                    let (packed, scales) = q4_split(self.quant_bytes(), self.rows(), cols);
989                    let gpr = cols / GROUP_SIZE;
990                    for gi in 0..gpr {
991                        let g = r * gpr + gi;
992                        let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
993                        for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
994                            dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
995                            dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
996                        }
997                    }
998                    if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
999                        crate::prism::inverse_embedding(model, dst);
1000                    }
1001                    return;
1002                }
1003                if *dtype == TensorDtype::Mxfp4 {
1004                    mxfp4::Mx::new(self.quant_bytes(), self.rows(), cols).row_f32(r, dst);
1005                    if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
1006                        crate::prism::inverse_embedding(model, dst);
1007                    }
1008                    return;
1009                }
1010                if *dtype == TensorDtype::Q1 {
1011                    let bytes = self.quant_bytes();
1012                    let gpr = cols / GROUP_SIZE;
1013                    for gi in 0..gpr {
1014                        let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
1015                        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
1016                        for (j, &b) in tile[2..].iter().enumerate() {
1017                            for k in 0..8 {
1018                                dst[gi * GROUP_SIZE + j * 8 + k] =
1019                                    (((b >> k) & 1) as f32 * 2.0 - 1.0) * s;
1020                            }
1021                        }
1022                    }
1023                    if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
1024                        crate::prism::inverse_embedding(model, dst);
1025                    }
1026                    return;
1027                }
1028                if *dtype == TensorDtype::Q1G {
1029                    // The tied Bonsai embedding is q1g too: one row = cols/8
1030                    // sign bytes + cols/128 scales, same ±s values as q1.
1031                    let bytes = self.quant_bytes();
1032                    let sc_off = cortiq_core::quant::q1g_scale_offset(self.rows(), cols);
1033                    let gpr = cols / GROUP_SIZE;
1034                    for gi in 0..gpr {
1035                        let t = r * gpr + gi;
1036                        let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, t));
1037                        for (j, &b) in bytes[t * 4..t * 4 + 4].iter().enumerate() {
1038                            for k in 0..8 {
1039                                dst[gi * GROUP_SIZE + j * 8 + k] =
1040                                    (((b >> k) & 1) as f32 * 2.0 - 1.0) * s;
1041                            }
1042                        }
1043                    }
1044                    if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
1045                        crate::prism::inverse_embedding(model, dst);
1046                    }
1047                    return;
1048                }
1049                if *dtype == TensorDtype::Q1T {
1050                    let bytes = self.quant_bytes();
1051                    let gpr = cols / GROUP_SIZE;
1052                    let base_len = self.rows() * gpr * cortiq_core::quant::Q1T_TILE;
1053                    for gi in 0..gpr {
1054                        let off = (r * gpr + gi) * cortiq_core::quant::Q1T_TILE;
1055                        let s = cortiq_core::quant::f16_to_f32(u16::from_le_bytes([
1056                            bytes[off],
1057                            bytes[off + 1],
1058                        ]));
1059                        let codes = &bytes[off + 2..off + cortiq_core::quant::Q1T_TILE];
1060                        for k in 0..GROUP_SIZE {
1061                            dst[gi * GROUP_SIZE + k] = match cortiq_core::quant::q1t_code(codes, k)
1062                            {
1063                                1 => s,
1064                                2 => -s,
1065                                _ => 0.0,
1066                            };
1067                        }
1068                    }
1069                    // Overlay
1070                    let rows = self.rows();
1071                    let entries = base_len + (rows + 1) * 4;
1072                    if entries <= bytes.len() {
1073                        let ptrs = &bytes[base_len..base_len + (rows + 1) * 4];
1074                        let r0 = u32::from_le_bytes([
1075                            ptrs[r * 4],
1076                            ptrs[r * 4 + 1],
1077                            ptrs[r * 4 + 2],
1078                            ptrs[r * 4 + 3],
1079                        ]) as usize;
1080                        let r1 = u32::from_le_bytes([
1081                            ptrs[(r + 1) * 4],
1082                            ptrs[(r + 1) * 4 + 1],
1083                            ptrs[(r + 1) * 4 + 2],
1084                            ptrs[(r + 1) * 4 + 3],
1085                        ]) as usize;
1086                        let off = entries + r0 * 4;
1087                        for i in 0..r1 - r0 {
1088                            let item = &bytes[off + i * 4..off + i * 4 + 4];
1089                            let c = u16::from_le_bytes([item[0], item[1]]) as usize;
1090                            let v = cortiq_core::quant::f16_to_f32(u16::from_le_bytes([
1091                                item[2], item[3],
1092                            ]));
1093                            if c < cols {
1094                                dst[c] = v;
1095                            }
1096                        }
1097                    }
1098                    if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
1099                        crate::prism::inverse_embedding(model, dst);
1100                    }
1101                    return;
1102                }
1103                if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1104                    let bytes = self.quant_bytes();
1105                    let rows = self.rows();
1106                    let ng = cols / GROUP_SIZE;
1107                    let bits = &bytes[..rows];
1108                    let sc_off = rows;
1109                    // Precomputed at load — embedding lookup used to scan
1110                    // the bit-widths of every preceding row (O(token_id)).
1111                    let off = vbit_offsets[r];
1112                    let b = bits[r] as usize;
1113                    let l = ((1usize << (b - 1)) - 1) as f32;
1114                    let data = &bytes[off..];
1115                    let (mut acc, mut nbits, mut byte_idx) = (0u64, 0usize, 0usize);
1116                    for (i, d) in dst.iter_mut().enumerate() {
1117                        while nbits < b {
1118                            acc = (acc << 8) | data[byte_idx] as u64;
1119                            byte_idx += 1;
1120                            nbits += 8;
1121                        }
1122                        let u = ((acc >> (nbits - b)) & ((1u64 << b) - 1)) as f32;
1123                        nbits -= b;
1124                        let so = (r * ng + i / GROUP_SIZE) * 2;
1125                        let sv = f16_to_f32(u16::from_le_bytes([
1126                            bytes[sc_off + so],
1127                            bytes[sc_off + so + 1],
1128                        ]));
1129                        *d = (u - l) * sv;
1130                    }
1131                    if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
1132                        crate::prism::inverse_embedding(model, dst);
1133                    }
1134                    return;
1135                }
1136                let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
1137                let s = row_scale[r];
1138                match dtype {
1139                    TensorDtype::Q8Row => {
1140                        for (d, &b) in dst.iter_mut().zip(q) {
1141                            *d = (b as i8) as f32 * s;
1142                        }
1143                    }
1144                    TensorDtype::Q8_2f => {
1145                        for (i, (d, &b)) in dst.iter_mut().zip(q).enumerate() {
1146                            *d = (b as i8) as f32 * s * col_field[i];
1147                        }
1148                    }
1149                    _ => unreachable!(),
1150                }
1151                if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
1152                    crate::prism::inverse_embedding(model, dst);
1153                }
1154            }
1155        }
1156    }
1157
1158    /// Can this tensor's columns be read cheaply (for sparse down_proj)?
1159    /// True for F32/Q8Row/Q8_2f (per-row scale, direct strided access);
1160    /// false for group-packed q4/vbit (column access would unpack whole
1161    /// groups — sparse execution falls back to f32 for those).
1162    pub fn sparse_col_ok(&self) -> bool {
1163        match self {
1164            Self::F32 { .. } => true,
1165            Self::Blr(_) => false,
1166            Self::Mapped { dtype, .. } => {
1167                matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
1168            }
1169        }
1170    }
1171
1172    /// down_proj [hidden, inter]: accumulate `w · col(c)` into `out`
1173    /// [hidden] — reads ONLY column `c` (one neuron) from the mmap,
1174    /// no full-matrix dequant. `out[k] += w · down[k, c]`.
1175    pub fn add_col_scaled(&self, c: usize, w: f32, out: &mut [f32]) {
1176        let inter = self.cols();
1177        let hidden = self.rows();
1178        debug_assert_eq!(out.len(), hidden);
1179        match self {
1180            Self::F32 { data, .. } => {
1181                for (k, o) in out.iter_mut().enumerate() {
1182                    *o += w * data[k * inter + c];
1183                }
1184            }
1185            Self::Blr(t) => t.add_col_scaled(c, w, out),
1186            Self::Mapped {
1187                dtype,
1188                row_scale,
1189                col_field,
1190                ..
1191            } => {
1192                let q = self.quant_bytes();
1193                let colf = if *dtype == TensorDtype::Q8_2f {
1194                    col_field[c]
1195                } else {
1196                    1.0
1197                };
1198                let wc = w * colf;
1199                for (k, o) in out.iter_mut().enumerate() {
1200                    let b = q[k * inter + c] as i8 as f32;
1201                    *o += wc * b * row_scale[k];
1202                }
1203            }
1204        }
1205    }
1206
1207    /// Touch the head of row `r` so the DRAM latency of the next
1208    /// neuron's weights overlaps the current one's arithmetic.
1209    ///
1210    /// Scattered rows are what per-token sparsity reads, and a 2 KB
1211    /// stride is past what the hardware prefetcher follows: without this
1212    /// every row starts with a cold miss that nothing hides. One touch
1213    /// per 512 bytes is enough — the rest of the row is a sequential run
1214    /// the prefetcher does pick up.
1215    #[inline]
1216    pub fn prefetch_row(&self, r: usize) {
1217        let Self::Mapped { dtype, .. } = self else {
1218            return;
1219        };
1220        if !matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f) {
1221            return;
1222        }
1223        let cols = self.cols();
1224        let q = self.quant_bytes();
1225        let (a, b) = (r * cols, (r + 1) * cols);
1226        if b > q.len() {
1227            return;
1228        }
1229        let mut j = a;
1230        while j < b {
1231            unsafe { std::ptr::read_volatile(q.as_ptr().add(j)) };
1232            j += 512;
1233        }
1234    }
1235
1236    /// `out += w · row(r)` — the transposed twin of `add_col_scaled`.
1237    ///
1238    /// A neuron's `down` weights are a COLUMN of `[hidden, inter]`, and a
1239    /// column is strided: reading one costs a cache line per element, so
1240    /// per-neuron dynamic sparsity saves arithmetic and no bytes. Stored
1241    /// transposed (`down_proj.t.weight`, `[inter, hidden]`) the same
1242    /// weights are a contiguous ROW, and this accumulate reads exactly
1243    /// the neurons the token asked for.
1244    pub fn add_row_scaled(&self, r: usize, w: f32, out: &mut [f32], scratch: &mut [f32]) {
1245        let cols = self.cols();
1246        debug_assert_eq!(out.len(), cols);
1247        match self {
1248            Self::F32 { data, .. } => {
1249                let row = &data[r * cols..(r + 1) * cols];
1250                for (o, v) in out.iter_mut().zip(row) {
1251                    *o += w * v;
1252                }
1253            }
1254            Self::Blr(t) => {
1255                let tmp = &mut scratch[..cols];
1256                t.row_f32(r, tmp);
1257                for (o, v) in out.iter_mut().zip(tmp.iter()) {
1258                    *o += w * v;
1259                }
1260            }
1261            Self::Mapped {
1262                dtype,
1263                row_scale,
1264                col_field,
1265                ..
1266            } => match dtype {
1267                TensorDtype::Q8Row => {
1268                    let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
1269                    let ws = w * row_scale[r];
1270                    let row: &[i8] =
1271                        unsafe { std::slice::from_raw_parts(q.as_ptr() as *const i8, q.len()) };
1272                    axpy_i8_f32(out, row, ws);
1273                }
1274                TensorDtype::Q8_2f => {
1275                    let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
1276                    let ws = w * row_scale[r];
1277                    for ((o, b), c) in out.iter_mut().zip(q).zip(col_field) {
1278                        *o += ws * c * (*b as i8 as f32);
1279                    }
1280                }
1281                _ => {
1282                    self.row_f32(r, scratch);
1283                    for (o, v) in out.iter_mut().zip(scratch.iter()) {
1284                        *o += w * v;
1285                    }
1286                }
1287            },
1288        }
1289    }
1290
1291    /// Dot of row `r` with `x` (gate/up active-neuron path). Reads only
1292    /// row `r` from the mmap — no full dequant. q4/vbit dequant the row
1293    /// into `scratch` first (rare for active-FFN weights).
1294    pub fn row_dot(&self, r: usize, x: &[f32], scratch: &mut [f32]) -> f32 {
1295        let cols = self.cols();
1296        match self {
1297            Self::F32 { data, .. } => {
1298                let row = &data[r * cols..(r + 1) * cols];
1299                row.iter().zip(x).map(|(w, v)| w * v).sum()
1300            }
1301            Self::Blr(t) => {
1302                let tmp = &mut scratch[..cols];
1303                t.row_f32(r, tmp);
1304                tmp.iter().zip(x).map(|(w, v)| w * v).sum()
1305            }
1306            Self::Mapped {
1307                model,
1308                idx,
1309                dtype,
1310                row_scale,
1311                col_field,
1312                ..
1313            } => {
1314                let prism_forward =
1315                    crate::prism::is_forward_weight(model, &model.tensors[*idx].name);
1316                if prism_forward {
1317                    let transformed = crate::prism::forward(model, &x[..cols]);
1318                    let gpr = cols / GROUP_SIZE;
1319                    match dtype {
1320                        TensorDtype::Q2TiledP => {
1321                            let v = Q4tpView::new_q2(self.quant_bytes(), self.rows(), cols);
1322                            let mut sc = vec![0f32; gpr];
1323                            v.scales_into(r, gpr, &mut sc);
1324                            if crate::prism::is_affine_target(model, &model.tensors[*idx].name) {
1325                                return q2tp_affine_row_exact(v.nib, r, gpr, &transformed, &sc);
1326                            }
1327                            return q2tp_row_exact(v.nib, r, gpr, &transformed, &sc);
1328                        }
1329                        _ => {
1330                            self.row_f32(r, scratch);
1331                            return scratch.iter().zip(&transformed).map(|(w, v)| w * v).sum();
1332                        }
1333                    }
1334                }
1335                match dtype {
1336                    TensorDtype::Q8Row => {
1337                        let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
1338                        dot_i8_f32(q, x) * row_scale[r]
1339                    }
1340                    TensorDtype::Q8_2f => {
1341                        let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
1342                        dot_i8_col_f32(q, x, col_field) * row_scale[r]
1343                    }
1344                    _ => {
1345                        self.row_f32(r, scratch);
1346                        scratch.iter().zip(x).map(|(w, v)| w * v).sum()
1347                    }
1348                }
1349            }
1350        }
1351    }
1352
1353    /// `out = W · x` (row-major). F32 delegates to the historical
1354    /// bit-exact path; Mapped runs the fused int8 kernel.
1355    pub fn matvec(&self, x: &[f32], out: &mut [f32], pool: Option<&Pool>) {
1356        if self.flint_exact_cpu_default() {
1357            return float_activations_scope(|| self.matvec(x, out, pool));
1358        }
1359        match self {
1360            // NOTE: `out.len()` DRIVES this arm — it computes that many rows,
1361            // and `x.len()` is the stride. A short `out` is legitimate here,
1362            // which is why the check below lives in the Mapped arm only.
1363            Self::F32 { data, .. } => {
1364                if !crate::f32_backend::matvec(data, x, out) {
1365                    matvec_rows(pool, data, x, out);
1366                }
1367            }
1368            Self::Blr(t) => {
1369                if crate::gpu::enabled_here()
1370                    && crate::blr::gpu_wanted()
1371                    && crate::gpu::blr_matmat(t, &x[..t.cols], 1, &mut out[..t.rows])
1372                {
1373                    return;
1374                }
1375                t.matvec(x, out, pool)
1376            }
1377            Self::Mapped {
1378                model,
1379                idx,
1380                dtype,
1381                rows,
1382                cols,
1383                row_scale,
1384                col_field,
1385                vbit_offsets,
1386                repack,
1387            } => {
1388                let _ = (model, idx);
1389                // Every kernel below writes `rows` entries through a raw
1390                // pointer, so a short `out` is an out-of-bounds WRITE, not a
1391                // wrong answer: it scribbles on the allocator's metadata and
1392                // the process aborts much later, somewhere innocent
1393                // (`double free or corruption`, `corrupted double-linked
1394                // list`). The debug_assert two of the kernels carried is
1395                // compiled out of the release — exactly the build where it
1396                // matters. Fail here instead, while the caller is still on
1397                // the stack to be named.
1398                assert!(
1399                    out.len() >= *rows && x.len() >= *cols,
1400                    "matvec {rows}x{cols}: out {} (need {rows}), x {} (need {cols})",
1401                    out.len(),
1402                    x.len(),
1403                );
1404                let prism_forward =
1405                    crate::prism::is_forward_weight(model, &model.tensors[*idx].name);
1406                if *dtype == TensorDtype::Q2TiledP
1407                    && std::env::var("CMF_Q2TP_TRACE").as_deref() == Ok("1")
1408                {
1409                    use std::sync::atomic::{AtomicUsize, Ordering};
1410                    static N: AtomicUsize = AtomicUsize::new(0);
1411                    let n = N.fetch_add(1, Ordering::Relaxed);
1412                    if n < 128 {
1413                        eprintln!(
1414                            "q2tp-dispatch #{n} name={} prism={} rows={} cols={} gpu={} optin={} layer={}",
1415                            model.tensors[*idx].name,
1416                            prism_forward,
1417                            rows,
1418                            cols,
1419                            crate::gpu::enabled_here(),
1420                            crate::gpu::q2tp_gpu_opt_in(),
1421                            crate::gpu::cur_layer(),
1422                        );
1423                    }
1424                }
1425                // Prism stores every manifest-listed forward matrix in the
1426                // signed-Hadamard basis.  The q2tp WGSL path receives that
1427                // transformed vector and an explicit affine bit; codecs
1428                // without a descriptor-aware kernel remain on CPU below.
1429                if prism_forward {
1430                    let transformed = crate::prism::forward(model, &x[..*cols]);
1431                    match dtype {
1432                        TensorDtype::Q4Block => {
1433                            q4matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1434                        }
1435                        TensorDtype::Q4Tiled => {
1436                            q4t_matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1437                        }
1438                        TensorDtype::Q4TiledP => {
1439                            q4tp_matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1440                        }
1441                        TensorDtype::Q2TiledP => {
1442                            let affine =
1443                                crate::prism::is_affine_target(model, &model.tensors[*idx].name);
1444                            if *rows * *cols >= 8_388_608
1445                                && crate::gpu::enabled_here()
1446                                && crate::gpu::q2tp_gpu_opt_in()
1447                            {
1448                                let gpu_ok = if affine {
1449                                    crate::gpu::q2tp_affine_matvec(
1450                                        model,
1451                                        *idx,
1452                                        &transformed,
1453                                        *rows,
1454                                        *cols,
1455                                        out,
1456                                    )
1457                                } else {
1458                                    crate::gpu::q2tp_matvec(
1459                                        model,
1460                                        *idx,
1461                                        &transformed,
1462                                        *rows,
1463                                        *cols,
1464                                        out,
1465                                    )
1466                                };
1467                                if gpu_ok {
1468                                    return;
1469                                }
1470                            }
1471                            if affine {
1472                                q2tp_affine_matvec(
1473                                    self.quant_bytes(),
1474                                    &transformed,
1475                                    *rows,
1476                                    *cols,
1477                                    out,
1478                                    pool,
1479                                )
1480                            } else {
1481                                q2tp_matvec(
1482                                    self.quant_bytes(),
1483                                    &transformed,
1484                                    *rows,
1485                                    *cols,
1486                                    out,
1487                                    pool,
1488                                )
1489                            }
1490                        }
1491                        TensorDtype::Q1 | TensorDtype::Q1G => bin_matvec(
1492                            self.bin_w().expect("binary dtype"),
1493                            &transformed,
1494                            *rows,
1495                            *cols,
1496                            out,
1497                            pool,
1498                        ),
1499                        TensorDtype::Q1T => {
1500                            q1t_matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1501                        }
1502                        TensorDtype::Mxfp4 => {
1503                            mxfp4::matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1504                        }
1505                        TensorDtype::Vbit | TensorDtype::VbitRo => vbitmatvec(
1506                            self.quant_bytes(),
1507                            vbit_offsets,
1508                            &transformed,
1509                            *rows,
1510                            *cols,
1511                            out,
1512                            pool,
1513                        ),
1514                        TensorDtype::Q8Row | TensorDtype::Q8_2f => qmatvec(
1515                            self.quant_bytes(),
1516                            repack,
1517                            row_scale,
1518                            &transformed,
1519                            col_field,
1520                            *dtype,
1521                            *rows,
1522                            *cols,
1523                            out,
1524                            pool,
1525                        ),
1526                        _ => unreachable!("unsupported mapped Prism dtype {dtype:?}"),
1527                    }
1528                    return;
1529                }
1530                if *dtype == TensorDtype::Q4Block {
1531                    // GPU route (wgpu q4b kernel) for large q4_block matvecs —
1532                    // gives NVIDIA/AMD/Intel q4 models a GPU path. Probe keeps
1533                    // the winner; Metal returns false → the CPU kernel below.
1534                    if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1535                        let t0 = std::time::Instant::now();
1536                        match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
1537                            crate::gpu::ProbeArm::Gpu => {
1538                                if crate::gpu::q4b_matvec(model, *idx, x, *rows, *cols, out) {
1539                                    crate::gpu::probe_record(
1540                                        crate::gpu::OpClass::Matvec,
1541                                        true,
1542                                        t0.elapsed(),
1543                                    );
1544                                    return;
1545                                }
1546                            }
1547                            crate::gpu::ProbeArm::CpuTimed => {
1548                                q4matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1549                                crate::gpu::probe_record(
1550                                    crate::gpu::OpClass::Matvec,
1551                                    false,
1552                                    t0.elapsed(),
1553                                );
1554                                return;
1555                            }
1556                            crate::gpu::ProbeArm::Cpu => {}
1557                        }
1558                    }
1559                    q4matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1560                    return;
1561                }
1562                if *dtype == TensorDtype::Q4Tiled {
1563                    // GPU route for large q4t matvecs — the lm_head class,
1564                    // same shape as the q4tp arm below. The probe keeps the
1565                    // winner; a backend without the kernel refuses and the
1566                    // CPU path stays.
1567                    if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1568                        let t0 = std::time::Instant::now();
1569                        let cls = crate::gpu::matvec_class(*rows, *cols);
1570                        match crate::gpu::probe_arm(cls) {
1571                            crate::gpu::ProbeArm::Gpu => {
1572                                if crate::gpu::q4t_matvec(model, *idx, x, *rows, *cols, out) {
1573                                    crate::gpu::probe_record(cls, true, t0.elapsed());
1574                                    return;
1575                                }
1576                            }
1577                            crate::gpu::ProbeArm::CpuTimed => {
1578                                q4t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1579                                crate::gpu::probe_record(cls, false, t0.elapsed());
1580                                return;
1581                            }
1582                            crate::gpu::ProbeArm::Cpu => {}
1583                        }
1584                    }
1585                    q4t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1586                    return;
1587                }
1588                if *dtype == TensorDtype::Q4TiledP {
1589                    // GPU route for large q4tp matvecs — the lm_head class.
1590                    // On a q4tp checkpoint the head is the biggest single
1591                    // host matvec left in the decode step, and the batched
1592                    // kernel at b=1 already exists on both backends. Probe
1593                    // keeps the winner, same as q4_block above.
1594                    if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1595                        let t0 = std::time::Instant::now();
1596                        let cls = crate::gpu::matvec_class(*rows, *cols);
1597                        match crate::gpu::probe_arm(cls) {
1598                            crate::gpu::ProbeArm::Gpu => {
1599                                if crate::gpu::q4tp_matvec(model, *idx, x, *rows, *cols, out) {
1600                                    crate::gpu::probe_record(cls, true, t0.elapsed());
1601                                    return;
1602                                }
1603                            }
1604                            crate::gpu::ProbeArm::CpuTimed => {
1605                                q4tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1606                                crate::gpu::probe_record(cls, false, t0.elapsed());
1607                                return;
1608                            }
1609                            crate::gpu::ProbeArm::Cpu => {}
1610                        }
1611                    }
1612                    q4tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1613                    return;
1614                }
1615                if *dtype == TensorDtype::Q2TiledP {
1616                    q2tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1617                    return;
1618                }
1619                if *dtype == TensorDtype::Mxfp4 {
1620                    self.mxfp4_matvec(x, out, pool);
1621                    return;
1622                }
1623                if *dtype == TensorDtype::Q1 || *dtype == TensorDtype::Q1G {
1624                    // GPU route for large q1 matvecs (out_proj / lm_head
1625                    // class): the CPU q1 kernel is load-port-bound at
1626                    // ~4 GB/s/core, the GPU one is bandwidth-bound — the
1627                    // probe measures both arms and keeps the winner.
1628                    let w = self.bin_w().expect("binary dtype");
1629                    if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1630                        let t0 = std::time::Instant::now();
1631                        let arm = if crate::gpu::q1_force() {
1632                            crate::gpu::ProbeArm::Gpu
1633                        } else {
1634                            crate::gpu::probe_arm(crate::gpu::OpClass::Matvec)
1635                        };
1636                        match arm {
1637                            crate::gpu::ProbeArm::Gpu => {
1638                                let ok = if *dtype == TensorDtype::Q1G {
1639                                    crate::gpu::q1g_matvec(model, *idx, x, *rows, *cols, out)
1640                                } else {
1641                                    crate::gpu::q1_matvec(model, *idx, x, *rows, *cols, out)
1642                                };
1643                                if ok {
1644                                    crate::gpu::probe_record(
1645                                        crate::gpu::OpClass::Matvec,
1646                                        true,
1647                                        t0.elapsed(),
1648                                    );
1649                                    return;
1650                                }
1651                            }
1652                            crate::gpu::ProbeArm::CpuTimed => {
1653                                bin_matvec(w, x, *rows, *cols, out, pool);
1654                                crate::gpu::probe_record(
1655                                    crate::gpu::OpClass::Matvec,
1656                                    false,
1657                                    t0.elapsed(),
1658                                );
1659                                return;
1660                            }
1661                            crate::gpu::ProbeArm::Cpu => {}
1662                        }
1663                    }
1664                    bin_matvec(w, x, *rows, *cols, out, pool);
1665                    return;
1666                }
1667                if *dtype == TensorDtype::Q1T {
1668                    // GPU route for large q1t matvecs: the ternary BASE dot runs
1669                    // on the GPU (load-port-bound on CPU, like q1), then the
1670                    // sparse overlay is added on the CPU. Probe keeps the winner.
1671                    if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1672                        let t0 = std::time::Instant::now();
1673                        match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
1674                            crate::gpu::ProbeArm::Gpu => {
1675                                if crate::gpu::q1t_matvec(model, *idx, x, *rows, *cols, out) {
1676                                    q1t_add_overlay(self.quant_bytes(), x, *rows, *cols, out, pool);
1677                                    crate::gpu::probe_record(
1678                                        crate::gpu::OpClass::Matvec,
1679                                        true,
1680                                        t0.elapsed(),
1681                                    );
1682                                    return;
1683                                }
1684                            }
1685                            crate::gpu::ProbeArm::CpuTimed => {
1686                                q1t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1687                                crate::gpu::probe_record(
1688                                    crate::gpu::OpClass::Matvec,
1689                                    false,
1690                                    t0.elapsed(),
1691                                );
1692                                return;
1693                            }
1694                            crate::gpu::ProbeArm::Cpu => {}
1695                        }
1696                    }
1697                    q1t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1698                    return;
1699                }
1700                if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1701                    vbitmatvec(self.quant_bytes(), vbit_offsets, x, *rows, *cols, out, pool);
1702                    return;
1703                }
1704                let xs = prescale(x, col_field, *dtype);
1705                // D5: large q8 matrices (lm_head-class) — hybrid
1706                // CPU∥GPU: split the rows, both sides compute
1707                // SIMULTANEOUSLY (same math, shared prescale).
1708                // GPU share: CMF_GPU_SPLIT (0..1, default 0.5).
1709                if *rows >= crate::gpu::min_rows()
1710                    && matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
1711                    && gpu_lmhead_enabled()
1712                    && crate::gpu::enabled_here()
1713                {
1714                    // Runtime probe: alternate the hybrid against the
1715                    // pure-CPU matvec, keep whichever is faster HERE.
1716                    let t0 = std::time::Instant::now();
1717                    match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
1718                        crate::gpu::ProbeArm::Gpu => {}
1719                        crate::gpu::ProbeArm::CpuTimed => {
1720                            qmatvec(
1721                                self.quant_bytes(),
1722                                repack,
1723                                row_scale,
1724                                x,
1725                                col_field,
1726                                *dtype,
1727                                *rows,
1728                                *cols,
1729                                out,
1730                                pool,
1731                            );
1732                            crate::gpu::probe_record(
1733                                crate::gpu::OpClass::Matvec,
1734                                false,
1735                                t0.elapsed(),
1736                            );
1737                            return;
1738                        }
1739                        crate::gpu::ProbeArm::Cpu => {
1740                            qmatvec(
1741                                self.quant_bytes(),
1742                                repack,
1743                                row_scale,
1744                                x,
1745                                col_field,
1746                                *dtype,
1747                                *rows,
1748                                *cols,
1749                                out,
1750                                pool,
1751                            );
1752                            return;
1753                        }
1754                    }
1755                    let frac = gpu_split_frac();
1756                    let cpu_rows = ((*rows as f32) * (1.0 - frac)) as usize;
1757                    let (out_cpu, out_gpu) = out.split_at_mut(cpu_rows);
1758                    let bytes = self.quant_bytes();
1759                    let ok = std::thread::scope(|sc| {
1760                        let g = sc.spawn(|| {
1761                            crate::gpu::q8_matvec_range(
1762                                model,
1763                                *idx,
1764                                cpu_rows,
1765                                &row_scale[cpu_rows..],
1766                                &xs,
1767                                *rows - cpu_rows,
1768                                *cols,
1769                                out_gpu,
1770                            )
1771                        });
1772                        if cpu_rows > 0 {
1773                            // Repack prefix covers the full groups of the
1774                            // CPU half (the split starts at row 0).
1775                            let rep_cpu = if repack.is_empty() {
1776                                &[][..]
1777                            } else {
1778                                &repack[..(cpu_rows / 4) * 4 * *cols]
1779                            };
1780                            qmatvec(
1781                                &bytes[..cpu_rows * *cols],
1782                                rep_cpu,
1783                                &row_scale[..cpu_rows],
1784                                x,
1785                                col_field,
1786                                *dtype,
1787                                cpu_rows,
1788                                *cols,
1789                                out_cpu,
1790                                pool,
1791                            );
1792                        }
1793                        g.join().unwrap_or(false)
1794                    });
1795                    if ok {
1796                        crate::gpu::probe_record(crate::gpu::OpClass::Matvec, true, t0.elapsed());
1797                        return;
1798                    }
1799                    // GPU failed — CPU finishes its half (rows rebased —
1800                    // group offsets don't line up, mmap layout only).
1801                    qmatvec(
1802                        &bytes[cpu_rows * *cols..(*rows) * *cols],
1803                        &[],
1804                        &row_scale[cpu_rows..],
1805                        x,
1806                        col_field,
1807                        *dtype,
1808                        *rows - cpu_rows,
1809                        *cols,
1810                        out_gpu,
1811                        pool,
1812                    );
1813                    return;
1814                }
1815                qmatvec(
1816                    self.quant_bytes(),
1817                    repack,
1818                    row_scale,
1819                    x,
1820                    col_field,
1821                    *dtype,
1822                    *rows,
1823                    *cols,
1824                    out,
1825                    pool,
1826                );
1827            }
1828        }
1829    }
1830
1831    /// Fused two-input matvec (MTP verify pair): weights streamed once.
1832    pub fn matvec2(
1833        &self,
1834        x1: &[f32],
1835        x2: &[f32],
1836        o1: &mut [f32],
1837        o2: &mut [f32],
1838        pool: Option<&Pool>,
1839    ) {
1840        if self.flint_exact_cpu_default() {
1841            return float_activations_scope(|| self.matvec2(x1, x2, o1, o2, pool));
1842        }
1843        match self {
1844            Self::F32 { data, .. } => matvec_rows2(pool, data, x1, x2, o1, o2),
1845            Self::Blr(_) => {
1846                self.matvec(x1, o1, pool);
1847                self.matvec(x2, o2, pool);
1848            }
1849            Self::Mapped {
1850                model,
1851                idx,
1852                dtype,
1853                rows,
1854                cols,
1855                row_scale,
1856                col_field,
1857                vbit_offsets,
1858                ..
1859            } => {
1860                if crate::prism::is_forward_weight(model, &model.tensors[*idx].name) {
1861                    let tx1 = crate::prism::forward(model, &x1[..*cols]);
1862                    let tx2 = crate::prism::forward(model, &x2[..*cols]);
1863                    match dtype {
1864                        TensorDtype::Q4Block => {
1865                            q4matvec2(self.quant_bytes(), &tx1, &tx2, *rows, *cols, o1, o2, pool)
1866                        }
1867                        TensorDtype::Q4Tiled => {
1868                            q4t_matvec2(self.quant_bytes(), &tx1, &tx2, *rows, *cols, o1, o2, pool)
1869                        }
1870                        TensorDtype::Q4TiledP => {
1871                            q4tp_matvec2(self.quant_bytes(), &tx1, &tx2, *rows, *cols, o1, o2, pool)
1872                        }
1873                        TensorDtype::Q2TiledP => {
1874                            if crate::prism::is_affine_target(model, &model.tensors[*idx].name) {
1875                                q2tp_affine_matvec2(
1876                                    self.quant_bytes(),
1877                                    &tx1,
1878                                    &tx2,
1879                                    *rows,
1880                                    *cols,
1881                                    o1,
1882                                    o2,
1883                                    pool,
1884                                )
1885                            } else {
1886                                q2tp_matvec2(
1887                                    self.quant_bytes(),
1888                                    &tx1,
1889                                    &tx2,
1890                                    *rows,
1891                                    *cols,
1892                                    o1,
1893                                    o2,
1894                                    pool,
1895                                )
1896                            }
1897                        }
1898                        TensorDtype::Q1 | TensorDtype::Q1G => bin_matvec2(
1899                            self.bin_w().expect("binary dtype"),
1900                            &tx1,
1901                            &tx2,
1902                            *rows,
1903                            *cols,
1904                            o1,
1905                            o2,
1906                            pool,
1907                        ),
1908                        TensorDtype::Q1T => {
1909                            q1t_matvec2(self.quant_bytes(), &tx1, &tx2, *rows, *cols, o1, o2, pool)
1910                        }
1911                        TensorDtype::Mxfp4 => mxfp4::matvec2(
1912                            self.quant_bytes(),
1913                            &tx1,
1914                            &tx2,
1915                            *rows,
1916                            *cols,
1917                            o1,
1918                            o2,
1919                            pool,
1920                        ),
1921                        TensorDtype::Vbit | TensorDtype::VbitRo => vbitmatvec2(
1922                            self.quant_bytes(),
1923                            vbit_offsets,
1924                            &tx1,
1925                            &tx2,
1926                            *rows,
1927                            *cols,
1928                            o1,
1929                            o2,
1930                            pool,
1931                        ),
1932                        TensorDtype::Q8Row | TensorDtype::Q8_2f => qmatvec2(
1933                            self.quant_bytes(),
1934                            row_scale,
1935                            &tx1,
1936                            &tx2,
1937                            col_field,
1938                            *dtype,
1939                            *rows,
1940                            *cols,
1941                            o1,
1942                            o2,
1943                            pool,
1944                        ),
1945                        _ => unreachable!("unsupported mapped Prism dtype {dtype:?}"),
1946                    }
1947                    return;
1948                }
1949                if *dtype == TensorDtype::Q4Block {
1950                    q4matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1951                    return;
1952                }
1953                if *dtype == TensorDtype::Q4Tiled {
1954                    q4t_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1955                    return;
1956                }
1957                if *dtype == TensorDtype::Q4TiledP {
1958                    q4tp_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1959                    return;
1960                }
1961                if *dtype == TensorDtype::Q2TiledP {
1962                    q2tp_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1963                    return;
1964                }
1965                if *dtype == TensorDtype::Mxfp4 {
1966                    mxfp4::matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1967                    return;
1968                }
1969                if let Some(w) = self.bin_w() {
1970                    bin_matvec2(w, x1, x2, *rows, *cols, o1, o2, pool);
1971                    return;
1972                }
1973                if *dtype == TensorDtype::Q1T {
1974                    // Fused ternary pair: one row pass, the register
1975                    // unpack shared across both streams on ARM. (Q1T
1976                    // lacks a row_scale array — scales live inline in
1977                    // the tiles — so it must not fall through to the
1978                    // q8 qmatvec2 below.)
1979                    q1t_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1980                    return;
1981                }
1982                if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1983                    vbitmatvec2(
1984                        self.quant_bytes(),
1985                        vbit_offsets,
1986                        x1,
1987                        x2,
1988                        *rows,
1989                        *cols,
1990                        o1,
1991                        o2,
1992                        pool,
1993                    );
1994                    return;
1995                }
1996                qmatvec2(
1997                    self.quant_bytes(),
1998                    row_scale,
1999                    x1,
2000                    x2,
2001                    col_field,
2002                    *dtype,
2003                    *rows,
2004                    *cols,
2005                    o1,
2006                    o2,
2007                    pool,
2008                );
2009            }
2010        }
2011    }
2012}
2013
2014impl QTensor {
2015    /// Batched matvec (prefill-GEMM): xs — row-major [b, cols],
2016    /// out — row-major [b, rows]. Element-wise semantics are IDENTICAL
2017    /// to b matvec calls (same dot kernels in the same order); the win —
2018    /// the weight row streams from DRAM once per batch, not b times.
2019    /// `(model, index)` when this is a memory-mapped q4tp tensor — the
2020    /// identity a device-resident chain needs to hand `tp_matmat` the
2021    /// weight without going through this struct's own dispatch.
2022    pub fn q4tp_mapped(&self) -> Option<(&std::sync::Arc<CmfModel>, usize)> {
2023        if self.has_prism_contract() {
2024            return None;
2025        }
2026        match self {
2027            Self::Mapped {
2028                model, idx, dtype, ..
2029            } if *dtype == TensorDtype::Q4TiledP => Some((model, *idx)),
2030            _ => None,
2031        }
2032    }
2033
2034    pub fn matmat(&self, xs_all: &[f32], b: usize, out: &mut [f32], pool: Option<&Pool>) {
2035        if self.flint_exact_cpu_default() {
2036            return float_activations_scope(|| self.matmat(xs_all, b, out, pool));
2037        }
2038        let cols = self.cols();
2039        let rows = self.rows();
2040        debug_assert_eq!(xs_all.len(), b * cols);
2041        debug_assert_eq!(out.len(), b * rows);
2042        let _prof = crate::cpuprof::time(crate::cpuprof::Slot::Matmat);
2043        // GPTQ calibration: fold this layer's inputs into its Hessian. Only
2044        // Mapped tensors carry a directory name; the check is a relaxed
2045        // atomic load, free when not calibrating.
2046        if crate::gptq_capture::capturing() {
2047            if let Self::Mapped { model, idx, .. } = self {
2048                crate::gptq_capture::accumulate(&model.tensors[*idx].name, xs_all, b, cols);
2049            }
2050            if let Self::Blr(t) = self {
2051                crate::gptq_capture::accumulate(t.name(), xs_all, b, cols);
2052            }
2053        }
2054        match self {
2055            Self::Blr(t) => {
2056                if crate::gpu::enabled_here()
2057                    && crate::blr::gpu_wanted()
2058                    && crate::gpu::blr_matmat(t, &xs_all[..b * cols], b, &mut out[..b * rows])
2059                {
2060                    return;
2061                }
2062                t.matmat(xs_all, b, out, pool)
2063            }
2064            Self::F32 { data, .. } => {
2065                if crate::f32_backend::matmat(data, xs_all, b, rows, cols, out) {
2066                    return;
2067                }
2068                let out_addr = SendMut(out.as_mut_ptr());
2069                // Over the batch, four positions a row at a time: four
2070                // independent chains, each still summing j = 0..cols in the
2071                // scalar order (bit-identical to the one-at-a-time loop and
2072                // to `matvec_rows`). The batch is the wide axis of a prefill;
2073                // the rows may be few (Spark-X2.5's 16-row g_proj spent
2074                // 17 ms a layer here split over its rows).
2075                let run = |start: usize, end: usize| {
2076                    let mut bi = start;
2077                    while bi < end {
2078                        let n = (end - bi).min(4);
2079                        for o in 0..rows {
2080                            let row = &data[o * cols..(o + 1) * cols];
2081                            if n == 4 {
2082                                let x = |k: usize| &xs_all[(bi + k) * cols..(bi + k + 1) * cols];
2083                                let (x0, x1, x2, x3) = (x(0), x(1), x(2), x(3));
2084                                let (mut a0, mut a1, mut a2, mut a3) = (0f32, 0f32, 0f32, 0f32);
2085                                for j in 0..cols {
2086                                    let w = row[j];
2087                                    a0 += w * x0[j];
2088                                    a1 += w * x1[j];
2089                                    a2 += w * x2[j];
2090                                    a3 += w * x3[j];
2091                                }
2092                                for (k, acc) in [a0, a1, a2, a3].into_iter().enumerate() {
2093                                    unsafe { *out_addr.at((bi + k) * rows + o) = acc };
2094                                }
2095                            } else {
2096                                for k in 0..n {
2097                                    let x = &xs_all[(bi + k) * cols..(bi + k + 1) * cols];
2098                                    let mut acc = 0f32;
2099                                    for j in 0..cols {
2100                                        acc += row[j] * x[j];
2101                                    }
2102                                    unsafe { *out_addr.at((bi + k) * rows + o) = acc };
2103                                }
2104                            }
2105                        }
2106                        bi += n;
2107                    }
2108                };
2109                dispatch_rows(pool, b, &run);
2110            }
2111            Self::Mapped {
2112                model,
2113                idx,
2114                dtype,
2115                row_scale,
2116                col_field,
2117                vbit_offsets,
2118                ..
2119            } => {
2120                if crate::prism::is_forward_weight(model, &model.tensors[*idx].name) {
2121                    let mut transformed = Vec::with_capacity(xs_all.len());
2122                    for bi in 0..b {
2123                        transformed.extend_from_slice(&crate::prism::forward(
2124                            model,
2125                            &xs_all[bi * cols..(bi + 1) * cols],
2126                        ));
2127                    }
2128                    match dtype {
2129                        TensorDtype::Q4Block => {
2130                            q4matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
2131                        }
2132                        TensorDtype::Q4Tiled => {
2133                            q4t_matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
2134                        }
2135                        TensorDtype::Q4TiledP => {
2136                            q4tp_matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
2137                        }
2138                        TensorDtype::Q2TiledP => {
2139                            let affine =
2140                                crate::prism::is_affine_target(model, &model.tensors[*idx].name);
2141                            // Affine Prism Q2TP has a descriptor-aware GPU
2142                            // kernel for short/tail batches too.  Unlike the
2143                            // ordinary Q2TP path, don't force b<32 back to a
2144                            // scalar CPU matmat: prefill chunks and the final
2145                            // tail both need to stay on the tested GPU arm.
2146                            let gpu_batch_ok = if affine {
2147                                b >= 2
2148                            } else {
2149                                b >= 32 && b * rows * cols >= 128_000_000
2150                            };
2151                            if gpu_batch_ok
2152                                && cols % 32 == 0
2153                                && crate::gpu::enabled_here()
2154                                && crate::gpu::q2tp_gpu_opt_in()
2155                            {
2156                                let gpu_ok = if affine {
2157                                    crate::gpu::q2tp_affine_matmat(
2158                                        model,
2159                                        *idx,
2160                                        &transformed,
2161                                        b,
2162                                        rows,
2163                                        cols,
2164                                        out,
2165                                    )
2166                                } else {
2167                                    crate::gpu::q2tp_matmat(
2168                                        model,
2169                                        *idx,
2170                                        &transformed,
2171                                        b,
2172                                        rows,
2173                                        cols,
2174                                        out,
2175                                    )
2176                                };
2177                                if gpu_ok {
2178                                    return;
2179                                }
2180                            }
2181                            // A one-token Prism decode is the other short
2182                            // case.  Use the descriptor-aware matvec kernel
2183                            // before falling back to the exact CPU path.
2184                            if affine
2185                                && b == 1
2186                                && cols % 32 == 0
2187                                && crate::gpu::enabled_here()
2188                                && crate::gpu::q2tp_gpu_opt_in()
2189                                && crate::gpu::q2tp_affine_matvec(
2190                                    model,
2191                                    *idx,
2192                                    &transformed[..cols],
2193                                    rows,
2194                                    cols,
2195                                    &mut out[..rows],
2196                                )
2197                            {
2198                                return;
2199                            }
2200                            if affine {
2201                                q2tp_affine_matmat(
2202                                    self.quant_bytes(),
2203                                    &transformed,
2204                                    b,
2205                                    rows,
2206                                    cols,
2207                                    out,
2208                                    pool,
2209                                )
2210                            } else {
2211                                q2tp_matmat(
2212                                    self.quant_bytes(),
2213                                    &transformed,
2214                                    b,
2215                                    rows,
2216                                    cols,
2217                                    out,
2218                                    pool,
2219                                )
2220                            }
2221                        }
2222                        TensorDtype::Q1 | TensorDtype::Q1G => bin_matmat(
2223                            self.bin_w().expect("binary dtype"),
2224                            &transformed,
2225                            b,
2226                            rows,
2227                            cols,
2228                            out,
2229                            pool,
2230                        ),
2231                        TensorDtype::Q1T => {
2232                            q1t_matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
2233                        }
2234                        TensorDtype::Mxfp4 => {
2235                            mxfp4::matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
2236                        }
2237                        TensorDtype::Vbit | TensorDtype::VbitRo => vbitmatmat(
2238                            self.quant_bytes(),
2239                            vbit_offsets,
2240                            &transformed,
2241                            b,
2242                            rows,
2243                            cols,
2244                            out,
2245                            pool,
2246                        ),
2247                        TensorDtype::Q8Row | TensorDtype::Q8_2f => {
2248                            let pre: Vec<std::borrow::Cow<'_, [f32]>> = (0..b)
2249                                .map(|bi| {
2250                                    prescale(
2251                                        &transformed[bi * cols..(bi + 1) * cols],
2252                                        col_field,
2253                                        *dtype,
2254                                    )
2255                                })
2256                                .collect();
2257                            qmatmat(self.quant_bytes(), row_scale, &pre, rows, cols, out, pool)
2258                        }
2259                        _ => unreachable!("unsupported mapped Prism dtype {dtype:?}"),
2260                    }
2261                    return;
2262                }
2263                if *dtype == TensorDtype::Q4Block {
2264                    q4matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2265                    return;
2266                }
2267                if *dtype == TensorDtype::Mxfp4 {
2268                    self.mxfp4_matmat(xs_all, b, out, pool);
2269                    return;
2270                }
2271                if *dtype == TensorDtype::Q4TiledP {
2272                    // GPU batched q4tp GEMM (dequant + f32nt mul_mm on the
2273                    // device); the probe keeps whichever beats the CPU arm.
2274                    // Narrow (prompt-encode) and wide (DiT) batches probe
2275                    // as separate classes — the regimes have opposite
2276                    // winners and one shared verdict locked the wrong arm.
2277                    // Kill switch (gpu::mm_kill): one grossly slow GPU op
2278                    // (a fair-condition op is ≤~100 ms even at 1024px)
2279                    // means the device is contended by another process
2280                    // (e.g. a simulator) — verdicts are per-process, so
2281                    // without the bail the whole render crawls behind
2282                    // someone else's queue.
2283                    // Row-exact batches stay on the host: the device GEMM
2284                    // is an f32 dequant-sgemm, not the host matvec's sum.
2285                    if b >= 32
2286                        && b * rows * cols >= 128_000_000
2287                        && cols % 32 == 0
2288                        && !row_exact()
2289                        && !crate::gpu::mm_killed()
2290                        && crate::gpu::enabled_here()
2291                    {
2292                        let class = if b >= 128 {
2293                            crate::gpu::OpClass::MatmatWide
2294                        } else {
2295                            crate::gpu::OpClass::Matmat
2296                        };
2297                        if let Self::Mapped { model, idx, .. } = self {
2298                            // In-process A/B (`CMF_MM_AB=1`). Three
2299                            // wall-clock A/Bs on a shared stand disagreed
2300                            // with each other by 25% on the same change,
2301                            // because the machine drifts between processes
2302                            // and interleaving whole renders does not fix
2303                            // that. Here both arms run back to back on the
2304                            // SAME data inside one call, so whatever the
2305                            // machine is doing, it does to both — and the
2306                            // disagreement between their outputs falls out
2307                            // for free. Doubles the work; a diagnostic,
2308                            // not a mode.
2309                            if crate::mm_ab::on() {
2310                                let mut g = vec![0f32; b * rows];
2311                                let t = std::time::Instant::now();
2312                                let took = crate::gpu::q4tp_matmat(
2313                                    model, *idx, xs_all, b, rows, cols, &mut g,
2314                                );
2315                                let dg = t.elapsed();
2316                                let t = std::time::Instant::now();
2317                                q4tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2318                                let dc = t.elapsed();
2319                                crate::mm_ab::record(b, rows, cols, took, dg, dc, &g, out);
2320                                return;
2321                            }
2322                            let t0 = std::time::Instant::now();
2323                            // A cold call takes the device arm: its sample
2324                            // is discarded either way, and the upload is
2325                            // what the next step needs.
2326                            let resident = crate::gpu::weight_is_resident(model, *idx);
2327                            match crate::gpu::probe_arm_cold_prefers_gpu(class, resident) {
2328                                crate::gpu::ProbeArm::Gpu => {
2329                                    if crate::gpu::q4tp_matmat(
2330                                        model, *idx, xs_all, b, rows, cols, out,
2331                                    ) {
2332                                        let el = t0.elapsed();
2333                                        // Work-proportional budget: ~8× the
2334                                        // fair-device estimate (+20 ms slack).
2335                                        // An absolute cap missed the worst
2336                                        // case — contended ops sit at
2337                                        // 100–240 ms each and still bury a
2338                                        // render whose fair op is 3–9 ms.
2339                                        // Cold ops (first PSO build, buffer
2340                                        // alloc) are exempt: a one-off
2341                                        // ~50 ms compile is not contention.
2342                                        let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
2343                                        let budget = std::time::Duration::from_secs_f64(
2344                                            flops / 1.5e12 * 8.0 + 0.020,
2345                                        );
2346                                        crate::gpu::mm_budget_check(
2347                                            "q4tp matmat",
2348                                            el,
2349                                            budget,
2350                                            crate::gpu::probe_was_cold() || !resident,
2351                                        );
2352                                        crate::gpu::probe_record(class, true, el);
2353                                        return;
2354                                    }
2355                                }
2356                                crate::gpu::ProbeArm::CpuTimed => {
2357                                    q4tp_matmat(
2358                                        self.quant_bytes(),
2359                                        xs_all,
2360                                        b,
2361                                        rows,
2362                                        cols,
2363                                        out,
2364                                        pool,
2365                                    );
2366                                    crate::gpu::probe_record(class, false, t0.elapsed());
2367                                    return;
2368                                }
2369                                crate::gpu::ProbeArm::Cpu => {}
2370                            }
2371                        }
2372                    }
2373                    q4tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2374                    return;
2375                }
2376                if *dtype == TensorDtype::Q2TiledP {
2377                    // Same device arm as q4tp, behind the same probe:
2378                    // the planes differ, the dispatch does not. Without
2379                    // this a q2tp file ran its widest projections on the
2380                    // host while the 4-bit one had the card, which is a
2381                    // codec paying for its size twice.
2382                    if b >= 32
2383                        && b * rows * cols >= 128_000_000
2384                        && cols % 32 == 0
2385                        && !crate::gpu::mm_killed()
2386                        && crate::gpu::enabled_here()
2387                    {
2388                        let class = if b >= 128 {
2389                            crate::gpu::OpClass::MatmatWide
2390                        } else {
2391                            crate::gpu::OpClass::Matmat
2392                        };
2393                        if let Self::Mapped { model, idx, .. } = self {
2394                            let t0 = std::time::Instant::now();
2395                            match crate::gpu::probe_arm(class) {
2396                                crate::gpu::ProbeArm::Gpu => {
2397                                    if crate::gpu::q2tp_matmat(
2398                                        model, *idx, xs_all, b, rows, cols, out,
2399                                    ) {
2400                                        crate::gpu::probe_record(class, true, t0.elapsed());
2401                                        return;
2402                                    }
2403                                }
2404                                crate::gpu::ProbeArm::CpuTimed => {
2405                                    q2tp_matmat(
2406                                        self.quant_bytes(),
2407                                        xs_all,
2408                                        b,
2409                                        rows,
2410                                        cols,
2411                                        out,
2412                                        pool,
2413                                    );
2414                                    crate::gpu::probe_record(class, false, t0.elapsed());
2415                                    return;
2416                                }
2417                                crate::gpu::ProbeArm::Cpu => {}
2418                            }
2419                        }
2420                    }
2421                    // Without a host arm a q2tp tensor falls through to
2422                    // the q8 fallback, which reads it at one BYTE per
2423                    // weight — a 2x overrun that killed pool workers
2424                    // mid-prefill while the dispatcher waited forever.
2425                    q2tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2426                    return;
2427                }
2428                if *dtype == TensorDtype::Q4Tiled {
2429                    // GPU batched q4t GEMM (dequant + f32nt mul_mm on the
2430                    // device); the probe keeps whichever beats the CPU arm.
2431                    // Narrow (prompt-encode) and wide (DiT) batches probe
2432                    // as separate classes — the regimes have opposite
2433                    // winners and one shared verdict locked the wrong arm.
2434                    // Kill switch (gpu::mm_kill): one grossly slow GPU op
2435                    // (a fair-condition op is ≤~100 ms even at 1024px)
2436                    // means the device is contended by another process
2437                    // (e.g. a simulator) — verdicts are per-process, so
2438                    // without the bail the whole render crawls behind
2439                    // someone else's queue.
2440                    if b >= 32
2441                        && b * rows * cols >= 128_000_000
2442                        && cols % 32 == 0
2443                        && !crate::gpu::mm_killed()
2444                        && crate::gpu::enabled_here()
2445                    {
2446                        let class = if b >= 128 {
2447                            crate::gpu::OpClass::MatmatWide
2448                        } else {
2449                            crate::gpu::OpClass::Matmat
2450                        };
2451                        if let Self::Mapped { model, idx, .. } = self {
2452                            let t0 = std::time::Instant::now();
2453                            match crate::gpu::probe_arm(class) {
2454                                crate::gpu::ProbeArm::Gpu => {
2455                                    if crate::gpu::q4t_matmat(
2456                                        model, *idx, xs_all, b, rows, cols, out,
2457                                    ) {
2458                                        let el = t0.elapsed();
2459                                        // Work-proportional budget: ~8× the
2460                                        // fair-device estimate (+20 ms slack).
2461                                        // An absolute cap missed the worst
2462                                        // case — contended ops sit at
2463                                        // 100–240 ms each and still bury a
2464                                        // render whose fair op is 3–9 ms.
2465                                        // Cold ops (first PSO build, buffer
2466                                        // alloc) are exempt: a one-off
2467                                        // ~50 ms compile is not contention.
2468                                        let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
2469                                        let budget = std::time::Duration::from_secs_f64(
2470                                            flops / 1.5e12 * 8.0 + 0.020,
2471                                        );
2472                                        crate::gpu::mm_budget_check(
2473                                            "q4t matmat",
2474                                            el,
2475                                            budget,
2476                                            crate::gpu::probe_was_cold(),
2477                                        );
2478                                        crate::gpu::probe_record(class, true, el);
2479                                        return;
2480                                    }
2481                                }
2482                                crate::gpu::ProbeArm::CpuTimed => {
2483                                    q4t_matmat(
2484                                        self.quant_bytes(),
2485                                        xs_all,
2486                                        b,
2487                                        rows,
2488                                        cols,
2489                                        out,
2490                                        pool,
2491                                    );
2492                                    crate::gpu::probe_record(class, false, t0.elapsed());
2493                                    return;
2494                                }
2495                                crate::gpu::ProbeArm::Cpu => {}
2496                            }
2497                        }
2498                    }
2499                    q4t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2500                    return;
2501                }
2502                if *dtype == TensorDtype::Q1 || *dtype == TensorDtype::Q1G {
2503                    // GPU batched q1 GEMM for wide prefill (q1_mul_mm on the
2504                    // device); the probe keeps whichever beats the CPU matmat.
2505                    // q1g has the twin `q1g_mul_mm` (cols % 128 by format).
2506                    let w = self.bin_w().expect("binary dtype");
2507                    if b >= 32
2508                        && b * rows * cols >= 128_000_000
2509                        && cols % 64 == 0
2510                        && crate::gpu::enabled_here()
2511                    {
2512                        if let Self::Mapped { model, idx, .. } = self {
2513                            let t0 = std::time::Instant::now();
2514                            match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
2515                                crate::gpu::ProbeArm::Gpu => {
2516                                    let ok = if *dtype == TensorDtype::Q1G {
2517                                        crate::gpu::q1g_matmat(
2518                                            model, *idx, xs_all, b, rows, cols, out,
2519                                        )
2520                                    } else {
2521                                        crate::gpu::q1_matmat(
2522                                            model, *idx, xs_all, b, rows, cols, out,
2523                                        )
2524                                    };
2525                                    if ok {
2526                                        crate::gpu::probe_record(
2527                                            crate::gpu::OpClass::Matmat,
2528                                            true,
2529                                            t0.elapsed(),
2530                                        );
2531                                        return;
2532                                    }
2533                                }
2534                                crate::gpu::ProbeArm::CpuTimed => {
2535                                    bin_matmat(w, xs_all, b, rows, cols, out, pool);
2536                                    crate::gpu::probe_record(
2537                                        crate::gpu::OpClass::Matmat,
2538                                        false,
2539                                        t0.elapsed(),
2540                                    );
2541                                    return;
2542                                }
2543                                crate::gpu::ProbeArm::Cpu => {}
2544                            }
2545                        }
2546                    }
2547                    bin_matmat(w, xs_all, b, rows, cols, out, pool);
2548                    return;
2549                }
2550                if *dtype == TensorDtype::Q1T {
2551                    // GPU batched GEMM for wide prefill (base + overlay on the
2552                    // device); probe keeps the winner vs the CPU matmat.
2553                    if b >= 32 && b * rows * cols >= 128_000_000 && crate::gpu::enabled_here() {
2554                        if let Self::Mapped { model, idx, .. } = self {
2555                            let t0 = std::time::Instant::now();
2556                            match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
2557                                crate::gpu::ProbeArm::Gpu => {
2558                                    if crate::gpu::q1t_matmat(
2559                                        model, *idx, xs_all, b, rows, cols, out,
2560                                    ) {
2561                                        crate::gpu::probe_record(
2562                                            crate::gpu::OpClass::Matmat,
2563                                            true,
2564                                            t0.elapsed(),
2565                                        );
2566                                        return;
2567                                    }
2568                                }
2569                                crate::gpu::ProbeArm::CpuTimed => {
2570                                    q1t_matmat(
2571                                        self.quant_bytes(),
2572                                        xs_all,
2573                                        b,
2574                                        rows,
2575                                        cols,
2576                                        out,
2577                                        pool,
2578                                    );
2579                                    crate::gpu::probe_record(
2580                                        crate::gpu::OpClass::Matmat,
2581                                        false,
2582                                        t0.elapsed(),
2583                                    );
2584                                    return;
2585                                }
2586                                crate::gpu::ProbeArm::Cpu => {}
2587                            }
2588                        }
2589                    }
2590                    q1t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2591                    return;
2592                }
2593                if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
2594                    vbitmatmat(
2595                        self.quant_bytes(),
2596                        vbit_offsets,
2597                        xs_all,
2598                        b,
2599                        rows,
2600                        cols,
2601                        out,
2602                        pool,
2603                    );
2604                    return;
2605                }
2606                let pre: Vec<std::borrow::Cow<'_, [f32]>> = (0..b)
2607                    .map(|bi| prescale(&xs_all[bi * cols..(bi + 1) * cols], col_field, *dtype))
2608                    .collect();
2609                // MiMo verification is a 2–4 row decode panel, not a wide
2610                // prompt GEMM. Keep q8 projections on the same device as
2611                // decode; the generic b>=8 gate otherwise silently moves
2612                // every projection back to CPU. The short wgpu matmat uses
2613                // the same 64-lane reduction as its single-token matvec.
2614                if row_exact()
2615                    && (1..=4).contains(&b)
2616                    && matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
2617                    && crate::gpu::enabled_here()
2618                    && crate::gpu::wgpu_active()
2619                {
2620                    let flat: Vec<f32> = pre.iter().flat_map(|v| v.iter().copied()).collect();
2621                    if crate::gpu::q8_matmat(model, *idx, row_scale, &flat, b, rows, cols, out) {
2622                        return;
2623                    }
2624                }
2625                // D5: large prefill-batch GEMMs — on the GPU (threshold by
2626                // work volume: submission carries b×rows×cols MACs).
2627                // Runtime probe: the naive GEMM shader + sync readback
2628                // lose to the CPU GEMM on slow driver stacks — alternate
2629                // both arms and keep the winner.
2630                if b >= 8 && b * rows * cols >= 128_000_000 && crate::gpu::enabled_here() {
2631                    if let Self::Mapped { model, idx, .. } = self {
2632                        let t0 = std::time::Instant::now();
2633                        match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
2634                            crate::gpu::ProbeArm::Gpu
2635                                if crate::gpu::probe_deciding(crate::gpu::OpClass::Matmat)
2636                                    && !crate::gpu::q8_resident_or_upload(model, *idx) =>
2637                            {
2638                                // Cold weights during probing: the upload
2639                                // has started, the count runs on the CPU —
2640                                // the GPU arm samples on the next touch.
2641                                let q = self.quant_bytes();
2642                                qmatmat(q, row_scale, &pre, rows, cols, out, pool);
2643                                return;
2644                            }
2645                            crate::gpu::ProbeArm::Gpu => {
2646                                // The two-field codec's own entry first, as
2647                                // `device_matmat` takes it: the column field
2648                                // folds into the weight plane and the GEMM
2649                                // runs on the matrix units. The plain int8
2650                                // entry below is the scalar f32 GEMM — on a
2651                                // Spark-X2.5 4B prefill it ran q|k|v|o at
2652                                // host speed.
2653                                if *dtype == TensorDtype::Q8_2f
2654                                    && std::env::var("CMF_Q8_2F_DEV").as_deref() != Ok("0")
2655                                    && crate::gpu::q8_matmat_2f(
2656                                        model, *idx, row_scale, col_field, xs_all, b, rows,
2657                                        cols, out,
2658                                    )
2659                                {
2660                                    crate::gpu::probe_record(
2661                                        crate::gpu::OpClass::Matmat,
2662                                        true,
2663                                        t0.elapsed(),
2664                                    );
2665                                    return;
2666                                }
2667                                let flat: Vec<f32> =
2668                                    pre.iter().flat_map(|v| v.iter().copied()).collect();
2669                                if crate::gpu::q8_matmat(
2670                                    model, *idx, row_scale, &flat, b, rows, cols, out,
2671                                ) {
2672                                    crate::gpu::probe_record(
2673                                        crate::gpu::OpClass::Matmat,
2674                                        true,
2675                                        t0.elapsed(),
2676                                    );
2677                                    return;
2678                                }
2679                            }
2680                            crate::gpu::ProbeArm::CpuTimed => {
2681                                let q = self.quant_bytes();
2682                                qmatmat(q, row_scale, &pre, rows, cols, out, pool);
2683                                crate::gpu::probe_record(
2684                                    crate::gpu::OpClass::Matmat,
2685                                    false,
2686                                    t0.elapsed(),
2687                                );
2688                                return;
2689                            }
2690                            crate::gpu::ProbeArm::Cpu => {}
2691                        }
2692                    }
2693                }
2694                let q = self.quant_bytes();
2695                qmatmat(q, row_scale, &pre, rows, cols, out, pool);
2696            }
2697        }
2698    }
2699}
2700
2701impl QTensor {
2702    /// `mxfp4` matvec: the wgpu kernel for lm_head/expert-class shapes when
2703    /// the probe prefers the device, else the CPU kernels of
2704    /// `qtensor::mxfp4` (always on native Metal, which has no mxfp4 kernel).
2705    fn mxfp4_matvec(&self, x: &[f32], out: &mut [f32], pool: Option<&Pool>) {
2706        let (rows, cols) = (self.rows(), self.cols());
2707        if let Self::Mapped { model, idx, .. } = self {
2708            if rows * cols >= 8_388_608 && crate::gpu::enabled_here() {
2709                let t0 = std::time::Instant::now();
2710                let cls = crate::gpu::matvec_class(rows, cols);
2711                match crate::gpu::probe_arm(cls) {
2712                    crate::gpu::ProbeArm::Gpu => {
2713                        if crate::gpu::mxfp4_matvec(model, *idx, x, rows, cols, out) {
2714                            crate::gpu::probe_record(cls, true, t0.elapsed());
2715                            return;
2716                        }
2717                    }
2718                    crate::gpu::ProbeArm::CpuTimed => {
2719                        mxfp4::matvec(self.quant_bytes(), x, rows, cols, out, pool);
2720                        crate::gpu::probe_record(cls, false, t0.elapsed());
2721                        return;
2722                    }
2723                    crate::gpu::ProbeArm::Cpu => {}
2724                }
2725            }
2726        }
2727        mxfp4::matvec(self.quant_bytes(), x, rows, cols, out, pool);
2728    }
2729
2730    /// `mxfp4` batched matvec (`xs` `[b, cols]` → `out` `[b, rows]`): the
2731    /// device's batched kernel behind the same probe as the other codecs'
2732    /// GEMMs (never for row-exact batches — the device sums in another
2733    /// order than the host matvec), else the CPU kernels.
2734    fn mxfp4_matmat(&self, xs: &[f32], b: usize, out: &mut [f32], pool: Option<&Pool>) {
2735        let (rows, cols) = (self.rows(), self.cols());
2736        if let Self::Mapped { model, idx, .. } = self {
2737            if b >= 2
2738                && b * rows * cols >= 33_554_432
2739                && !row_exact()
2740                && !crate::gpu::mm_killed()
2741                && crate::gpu::enabled_here()
2742            {
2743                let class = if b >= 128 {
2744                    crate::gpu::OpClass::MatmatWide
2745                } else {
2746                    crate::gpu::OpClass::Matmat
2747                };
2748                let t0 = std::time::Instant::now();
2749                match crate::gpu::probe_arm(class) {
2750                    crate::gpu::ProbeArm::Gpu => {
2751                        if crate::gpu::mxfp4_matmat(model, *idx, xs, b, rows, cols, out) {
2752                            crate::gpu::probe_record(class, true, t0.elapsed());
2753                            return;
2754                        }
2755                    }
2756                    crate::gpu::ProbeArm::CpuTimed => {
2757                        mxfp4::matmat(self.quant_bytes(), xs, b, rows, cols, out, pool);
2758                        crate::gpu::probe_record(class, false, t0.elapsed());
2759                        return;
2760                    }
2761                    crate::gpu::ProbeArm::Cpu => {}
2762                }
2763            }
2764        }
2765        mxfp4::matmat(self.quant_bytes(), xs, b, rows, cols, out, pool);
2766    }
2767}
2768
2769impl QTensor {
2770    /// The device GEMM this tensor would take, run once on the caller's
2771    /// data — the startup parity probe's arm, and the one place that knows
2772    /// which entry point each codec has.
2773    ///
2774    /// It exists because the probe used to look for a `q4tp` weight by
2775    /// name AND dtype, and a container packed any other way was declared
2776    /// "host path" for the whole render even though its codec had a device
2777    /// GEMM of its own. A gate that only recognizes one codec is a gate
2778    /// that silently downgrades every other one.
2779    pub fn device_matmat(&self, xs: &[f32], b: usize, out: &mut [f32]) -> bool {
2780        let (rows, cols) = (self.rows(), self.cols());
2781        let Self::Mapped {
2782            model,
2783            idx,
2784            dtype,
2785            row_scale,
2786            col_field,
2787            ..
2788        } = self
2789        else {
2790            return false;
2791        };
2792        if crate::prism::has_contract(model) {
2793            return false;
2794        }
2795        match *dtype {
2796            TensorDtype::Q4TiledP => crate::gpu::q4tp_matmat(model, *idx, xs, b, rows, cols, out),
2797            // The two-field codec folds its column field into the
2798            // activation, which leaves a plain per-row int8 GEMM — the
2799            // same kernel `q8_row` uses, on both backends.
2800            TensorDtype::Q8Row | TensorDtype::Q8_2f => {
2801                // The field belongs to the weight; only a backend that cannot
2802                // apply it there makes a scaled copy of the activation.
2803                if *dtype == TensorDtype::Q8_2f
2804                    && std::env::var("CMF_Q8_2F_DEV").as_deref() != Ok("0")
2805                    && crate::gpu::q8_matmat_2f(
2806                        model, *idx, row_scale, col_field, xs, b, rows, cols, out,
2807                    )
2808                {
2809                    return true;
2810                }
2811                let flat: Vec<f32> = (0..b)
2812                    .flat_map(|bi| {
2813                        prescale(&xs[bi * cols..(bi + 1) * cols], col_field, *dtype).into_owned()
2814                    })
2815                    .collect();
2816                crate::gpu::q8_matmat(model, *idx, row_scale, &flat, b, rows, cols, out)
2817            }
2818            _ => false,
2819        }
2820    }
2821
2822    /// Multi-matrix job (roadmap §3 P0): N tensors sharing one input
2823    /// run under a SINGLE pool dispatch — QKV or gate+up cost one
2824    /// barrier instead of N. Per-row math is the exact same kernel as
2825    /// `matvec` (bit-identical outputs); only the dispatch is fused.
2826    /// Falls back to N sequential matvecs when the set is not a uniform
2827    /// q8-family/F32 group or there is no pool.
2828    pub fn matvec_many<const N: usize>(
2829        ts: [&QTensor; N],
2830        x: &[f32],
2831        mut outs: [&mut [f32]; N],
2832        pool: Option<&Pool>,
2833    ) {
2834        if Self::any_flint_exact_cpu_default(&ts) {
2835            return float_activations_scope(|| Self::matvec_many(ts, x, outs, pool));
2836        }
2837        let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
2838        if ts.iter().any(|t| t.has_prism_contract() || t.per_tensor_only()) {
2839            // The fused range kernels have no transform descriptor.  Let
2840            // each tensor's ordinary matvec dispatch perform the explicit
2841            // signed FWHT (and retain CPU fallback for mixed q2tp/q4tp).
2842            // blr has no row-range kernel either: per-tensor calls.
2843            for (t, o) in ts.iter().zip(outs.iter_mut()) {
2844                t.matvec(x, o, pool);
2845            }
2846            return;
2847        }
2848        let uniform_q8 = ts.iter().all(|t| {
2849            matches!(
2850                t,
2851                Self::Mapped {
2852                    dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f,
2853                    ..
2854                }
2855            )
2856        });
2857        let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
2858        if uniform_f32 && crate::f32_backend::active() {
2859            for (t, o) in ts.iter().zip(outs.iter_mut()) {
2860                t.matvec(x, o, pool);
2861            }
2862            return;
2863        }
2864        let uniform_q4 = ts.iter().all(|t| {
2865            matches!(
2866                t,
2867                Self::Mapped {
2868                    dtype: TensorDtype::Q4Block,
2869                    ..
2870                }
2871            )
2872        });
2873        let uniform_vbit = ts.iter().all(|t| {
2874            matches!(
2875                t,
2876                Self::Mapped {
2877                    dtype: TensorDtype::Vbit | TensorDtype::VbitRo,
2878                    ..
2879                }
2880            )
2881        });
2882        // q1 and q1g may mix freely: the shared activation split and group
2883        // sums are layout-independent, only the row kernel differs.
2884        let uniform_q1 =
2885            ts.iter().all(|t| t.bin_w().is_some()) && ts.iter().all(|t| t.cols() == ts[0].cols());
2886        let uniform_q1t = ts.iter().all(|t| {
2887            matches!(
2888                t,
2889                Self::Mapped {
2890                    dtype: TensorDtype::Q1T,
2891                    ..
2892                }
2893            )
2894        });
2895        // q4tp is the skeleton dtype of the big MoE files, and without an arm
2896        // here every projection that shares an input paid its own pool
2897        // barrier: DeepSeek-V4's attention step alone hands this function
2898        // wq_a, wkv and both compressors' pairs off the same hidden state.
2899        let uniform_q4tp = ts.iter().all(|t| {
2900            matches!(
2901                t,
2902                Self::Mapped {
2903                    dtype: TensorDtype::Q4TiledP,
2904                    ..
2905                }
2906            )
2907        }) && ts
2908            .iter()
2909            .all(|t| t.cols() == ts[0].cols() && t.cols() % GROUP_SIZE == 0);
2910        let Some(pool) = pool else {
2911            for (t, o) in ts.iter().zip(outs.iter_mut()) {
2912                t.matvec(x, o, None);
2913            }
2914            return;
2915        };
2916        if total_rows < 256
2917            || !(uniform_q8
2918                || uniform_f32
2919                || uniform_q4
2920                || uniform_vbit
2921                || uniform_q1
2922                || uniform_q1t
2923                || uniform_q4tp)
2924        {
2925            for (t, o) in ts.iter().zip(outs.iter_mut()) {
2926                t.matvec(x, o, Some(pool));
2927            }
2928            return;
2929        }
2930
2931        if uniform_q4tp {
2932            // Every tensor's rows laid end to end in one virtual row space,
2933            // so the whole set is ONE dispatch. The per-row body is the
2934            // `q4tp_matvec` arm verbatim — same activation split, same
2935            // accumulation order — so the outputs are bit-identical to the
2936            // sequential calls this replaces.
2937            let cols = ts[0].cols();
2938            let gpr = cols / GROUP_SIZE;
2939            let views: [Q4tpView; N] =
2940                std::array::from_fn(|i| Q4tpView::new(ts[i].quant_bytes(), ts[i].rows(), cols));
2941            let rows_of: [usize; N] = std::array::from_fn(|i| ts[i].rows());
2942            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2943            // flat index -> (which tensor, which of its rows)
2944            let locate = |flat: usize| -> (usize, usize) {
2945                let mut acc = 0;
2946                for (i, &r) in rows_of.iter().enumerate() {
2947                    if flat < acc + r {
2948                        return (i, flat - acc);
2949                    }
2950                    acc += r;
2951                }
2952                (rows_of.len() - 1, 0)
2953            };
2954            let (views, outs_addr) = (&views, &outs_addr);
2955            if a8w8_enabled() {
2956                let act = split_act(x);
2957                let act = &act;
2958                let run = |start: usize, end: usize| {
2959                    with_krow(gpr, |sc| {
2960                        for flat in start..end {
2961                            let (t, r) = locate(flat);
2962                            let v = &views[t];
2963                            v.scales_into(r, gpr, sc);
2964                            let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, sc) * act.sx;
2965                            for &(j, xv) in &act.outliers {
2966                                let (w, s) = q4tp_outlier(v.nib, r, gpr, j, sc);
2967                                acc += w * s * xv;
2968                            }
2969                            // SAFETY: one worker owns each (tensor, row) pair.
2970                            unsafe { *outs_addr[t].at(r) = acc };
2971                        }
2972                    });
2973                };
2974                pool.run_rows(total_rows, &run);
2975            } else {
2976                let run = |start: usize, end: usize| {
2977                    with_krow(gpr, |sc| {
2978                        for flat in start..end {
2979                            let (t, r) = locate(flat);
2980                            let v = &views[t];
2981                            v.scales_into(r, gpr, sc);
2982                            // SAFETY: one worker owns each (tensor, row) pair.
2983                            unsafe { *outs_addr[t].at(r) = q4tp_row_exact(v.nib, r, gpr, x, sc) };
2984                        }
2985                    });
2986                };
2987                pool.run_rows(total_rows, &run);
2988            }
2989            return;
2990        }
2991
2992        if uniform_q1 {
2993            // One shared activation split + group sums (q1 has no col
2994            // field; the same input feeds every tensor).
2995            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2996            if a8w8_enabled() {
2997                let act = split_act(x);
2998                let gsum = q1_group_sums(&act.xq, ts[0].cols() / GROUP_SIZE);
2999                let (act, gsum) = (&act, &gsum);
3000                let closures: [_; N] = std::array::from_fn(|i| {
3001                    let (w, gpr, out) = (
3002                        ts[i].bin_w().expect("uniform_q1 checked"),
3003                        ts[i].cols() / GROUP_SIZE,
3004                        outs_addr[i],
3005                    );
3006                    move |s: usize, e: usize| bin_range_a8w8(w, gpr, act, gsum, out, s, e)
3007                });
3008                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3009                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3010                pool.run_many(&parts);
3011            } else {
3012                let closures: [_; N] = std::array::from_fn(|i| {
3013                    let (w, gpr, out) = (
3014                        ts[i].bin_w().expect("uniform_q1 checked"),
3015                        ts[i].cols() / GROUP_SIZE,
3016                        outs_addr[i],
3017                    );
3018                    move |s: usize, e: usize| bin_range_f32(w, gpr, x, out, s, e)
3019                });
3020                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3021                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3022                pool.run_many(&parts);
3023            }
3024            return;
3025        }
3026
3027        if uniform_q1t {
3028            // Q1T batched: one shared activation split + overlay decode,
3029            // all tensors' rows in ONE pool dispatch (saves N−1 dispatches
3030            // and N−1 redundant split_act calls per layer).
3031            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
3032            const TILE: usize = cortiq_core::quant::Q1T_TILE;
3033            if a8w8_enabled() {
3034                let act = split_act(x);
3035                let act = &act;
3036                let x_ref = x;
3037                let closures: [_; N] = std::array::from_fn(|i| {
3038                    let bytes = ts[i].quant_bytes();
3039                    let (rows, cols) = (ts[i].rows(), ts[i].cols());
3040                    let gpr = cols / GROUP_SIZE;
3041                    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
3042                    let out = outs_addr[i];
3043                    move |s: usize, e: usize| {
3044                        q1t_range_a8w8(bytes, gpr, rp_off, ent_off, has_ov, act, x_ref, out, 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            } else {
3051                let x_ref = x;
3052                let closures: [_; N] = std::array::from_fn(|i| {
3053                    let bytes = ts[i].quant_bytes();
3054                    let (rows, cols) = (ts[i].rows(), ts[i].cols());
3055                    let gpr = cols / GROUP_SIZE;
3056                    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
3057                    let out = outs_addr[i];
3058                    move |s: usize, e: usize| {
3059                        q1t_range_f32_batch(bytes, gpr, rp_off, ent_off, has_ov, x_ref, out, s, e)
3060                    }
3061                });
3062                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3063                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3064                pool.run_many(&parts);
3065            }
3066            return;
3067        }
3068
3069        if uniform_q4 || uniform_vbit {
3070            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
3071            // q4/vbit share one activation split — no per-tensor col field.
3072            if a8w8_enabled() {
3073                let act = split_act(x);
3074                let act = &act;
3075                if uniform_q4 {
3076                    let closures: [_; N] = std::array::from_fn(|i| {
3077                        let (packed, scales) =
3078                            q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
3079                        let (gpr, cols, out) =
3080                            (ts[i].cols() / GROUP_SIZE, ts[i].cols(), outs_addr[i]);
3081                        move |s: usize, e: usize| {
3082                            q4_range_a8w8(packed, scales, gpr, cols, act, out, s, e)
3083                        }
3084                    });
3085                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3086                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3087                    pool.run_many(&parts);
3088                } else {
3089                    let closures: [_; N] = std::array::from_fn(|i| {
3090                        let Self::Mapped { vbit_offsets, .. } = ts[i] else {
3091                            unreachable!()
3092                        };
3093                        let (bytes, rows, cols, out) = (
3094                            ts[i].quant_bytes(),
3095                            ts[i].rows(),
3096                            ts[i].cols(),
3097                            outs_addr[i],
3098                        );
3099                        move |s: usize, e: usize| {
3100                            vbit_range_a8w8(bytes, vbit_offsets, x, act, rows, cols, out, s, e)
3101                        }
3102                    });
3103                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3104                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3105                    pool.run_many(&parts);
3106                }
3107                return;
3108            }
3109            if uniform_q4 {
3110                let closures: [_; N] = std::array::from_fn(|i| {
3111                    let (packed, scales) =
3112                        q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
3113                    let (gpr, out) = (ts[i].cols() / GROUP_SIZE, outs_addr[i]);
3114                    move |s: usize, e: usize| q4_range_f32(packed, scales, gpr, x, out, s, e)
3115                });
3116                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3117                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3118                pool.run_many(&parts);
3119            } else {
3120                let closures: [_; N] = std::array::from_fn(|i| {
3121                    let Self::Mapped { vbit_offsets, .. } = ts[i] else {
3122                        unreachable!()
3123                    };
3124                    let (bytes, rows, cols, out) = (
3125                        ts[i].quant_bytes(),
3126                        ts[i].rows(),
3127                        ts[i].cols(),
3128                        outs_addr[i],
3129                    );
3130                    move |s: usize, e: usize| {
3131                        vbit_range_f32(bytes, vbit_offsets, x, rows, cols, out, s, e)
3132                    }
3133                });
3134                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3135                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3136                pool.run_many(&parts);
3137            }
3138            return;
3139        }
3140
3141        if uniform_f32 {
3142            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
3143            let closures: [_; N] = std::array::from_fn(|i| {
3144                let Self::F32 { data, cols, .. } = ts[i] else {
3145                    unreachable!()
3146                };
3147                let out = outs_addr[i];
3148                move |start: usize, end: usize| {
3149                    for o in start..end {
3150                        let row = &data[o * cols..(o + 1) * cols];
3151                        let mut sum = 0.0f32;
3152                        for j in 0..*cols {
3153                            sum += row[j] * x[j];
3154                        }
3155                        // SAFETY: disjoint (tensor, row) cells per worker.
3156                        unsafe { *out.at(o) = sum };
3157                    }
3158                }
3159            });
3160            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3161                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3162            pool.run_many(&parts);
3163            return;
3164        }
3165
3166        // Uniform q8-family: per-tensor prescale (q8_2f col fields
3167        // differ per tensor) + the shared range kernels.
3168        struct Ctx<'a> {
3169            bytes: &'a [u8],
3170            #[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
3171            rep: &'a [u8],
3172            row_scale: &'a [f32],
3173            cols: usize,
3174            xs: std::borrow::Cow<'a, [f32]>,
3175        }
3176        let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
3177            let Self::Mapped {
3178                dtype,
3179                cols,
3180                row_scale,
3181                col_field,
3182                repack,
3183                ..
3184            } = ts[i]
3185            else {
3186                unreachable!()
3187            };
3188            Ctx {
3189                bytes: ts[i].quant_bytes(),
3190                rep: repack,
3191                row_scale,
3192                cols: *cols,
3193                xs: prescale(x, col_field, *dtype),
3194            }
3195        });
3196        let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
3197        #[cfg(target_arch = "aarch64")]
3198        if sdot_enabled() {
3199            let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
3200            let closures: [_; N] = std::array::from_fn(|i| {
3201                let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
3202                move |start: usize, end: usize| {
3203                    q8_range_sdot(c.bytes, c.rep, c.row_scale, act, c.cols, out, start, end)
3204                }
3205            });
3206            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3207                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3208            pool.run_many(&parts);
3209            return;
3210        }
3211        #[cfg(target_arch = "x86_64")]
3212        if avx2_a8w8_enabled() {
3213            let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
3214            let closures: [_; N] = std::array::from_fn(|i| {
3215                let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
3216                move |start: usize, end: usize| {
3217                    q8_range_avx2(c.bytes, c.row_scale, act, c.cols, out, start, end)
3218                }
3219            });
3220            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3221                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3222            pool.run_many(&parts);
3223            return;
3224        }
3225        let closures: [_; N] = std::array::from_fn(|i| {
3226            let (c, out) = (&ctxs[i], outs_addr[i]);
3227            move |start: usize, end: usize| {
3228                q8_range_f32(c.bytes, c.row_scale, &c.xs, c.cols, out, start, end)
3229            }
3230        });
3231        let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3232            std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3233        pool.run_many(&parts);
3234    }
3235}
3236
3237impl QTensor {
3238    /// Pair-input multi-matrix job: N tensors × 2 shared inputs under a
3239    /// single pool dispatch — the MTP/pair decode path publishes one job
3240    /// for Q/K/V (and one for gate+up) instead of one per tensor.
3241    /// Per-row math is exactly `matvec2`'s kernels; bit-identical.
3242    #[allow(clippy::needless_range_loop)]
3243    pub fn matvec2_many<const N: usize>(
3244        ts: [&QTensor; N],
3245        x1: &[f32],
3246        x2: &[f32],
3247        mut o1s: [&mut [f32]; N],
3248        mut o2s: [&mut [f32]; N],
3249        pool: Option<&Pool>,
3250    ) {
3251        if Self::any_flint_exact_cpu_default(&ts) {
3252            return float_activations_scope(|| {
3253                Self::matvec2_many(ts, x1, x2, o1s, o2s, pool)
3254            });
3255        }
3256        let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
3257        if ts.iter().any(|t| t.has_prism_contract() || t.per_tensor_only()) {
3258            for i in 0..N {
3259                ts[i].matvec2(x1, x2, o1s[i], o2s[i], pool);
3260            }
3261            return;
3262        }
3263        let uniform_q8 = ts.iter().all(|t| {
3264            matches!(
3265                t,
3266                Self::Mapped {
3267                    dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f,
3268                    ..
3269                }
3270            )
3271        });
3272        let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
3273        let uniform_q4 = ts.iter().all(|t| {
3274            matches!(
3275                t,
3276                Self::Mapped {
3277                    dtype: TensorDtype::Q4Block,
3278                    ..
3279                }
3280            )
3281        });
3282        let uniform_vbit = ts.iter().all(|t| {
3283            matches!(
3284                t,
3285                Self::Mapped {
3286                    dtype: TensorDtype::Vbit | TensorDtype::VbitRo,
3287                    ..
3288                }
3289            )
3290        });
3291        let fusable = pool.is_some()
3292            && total_rows >= 256
3293            && (uniform_q8 || uniform_f32 || uniform_q4 || uniform_vbit);
3294        if !fusable {
3295            for i in 0..N {
3296                ts[i].matvec2(x1, x2, o1s[i], o2s[i], pool);
3297            }
3298            return;
3299        }
3300        let pool = pool.unwrap();
3301
3302        if uniform_q4 || uniform_vbit {
3303            let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
3304            let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
3305            // q4/vbit share activation splits — no per-tensor col field.
3306            if a8w8_enabled() {
3307                let a1 = split_act(x1);
3308                let a2 = split_act(x2);
3309                let (a1, a2) = (&a1, &a2);
3310                if uniform_q4 {
3311                    let closures: [_; N] = std::array::from_fn(|i| {
3312                        let (packed, scales) =
3313                            q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
3314                        let (gpr, cols, o1, o2) =
3315                            (ts[i].cols() / GROUP_SIZE, ts[i].cols(), p1[i], p2[i]);
3316                        move |s: usize, e: usize| {
3317                            q4_range2_a8w8(packed, scales, gpr, cols, a1, a2, o1, o2, s, e)
3318                        }
3319                    });
3320                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3321                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3322                    pool.run_many(&parts);
3323                } else {
3324                    let closures: [_; N] = std::array::from_fn(|i| {
3325                        let Self::Mapped { vbit_offsets, .. } = ts[i] else {
3326                            unreachable!()
3327                        };
3328                        let (bytes, rows, cols, o1, o2) = (
3329                            ts[i].quant_bytes(),
3330                            ts[i].rows(),
3331                            ts[i].cols(),
3332                            p1[i],
3333                            p2[i],
3334                        );
3335                        move |s: usize, e: usize| {
3336                            vbit_range2_a8w8(
3337                                bytes,
3338                                vbit_offsets,
3339                                x1,
3340                                x2,
3341                                a1,
3342                                a2,
3343                                rows,
3344                                cols,
3345                                o1,
3346                                o2,
3347                                s,
3348                                e,
3349                            )
3350                        }
3351                    });
3352                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3353                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3354                    pool.run_many(&parts);
3355                }
3356                return;
3357            }
3358            if uniform_q4 {
3359                let closures: [_; N] = std::array::from_fn(|i| {
3360                    let (packed, scales) =
3361                        q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
3362                    let (gpr, o1, o2) = (ts[i].cols() / GROUP_SIZE, p1[i], p2[i]);
3363                    move |s: usize, e: usize| {
3364                        q4_range2_f32(packed, scales, gpr, x1, x2, o1, o2, s, e)
3365                    }
3366                });
3367                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3368                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3369                pool.run_many(&parts);
3370            } else {
3371                let closures: [_; N] = std::array::from_fn(|i| {
3372                    let Self::Mapped { vbit_offsets, .. } = ts[i] else {
3373                        unreachable!()
3374                    };
3375                    let (bytes, rows, cols, o1, o2) = (
3376                        ts[i].quant_bytes(),
3377                        ts[i].rows(),
3378                        ts[i].cols(),
3379                        p1[i],
3380                        p2[i],
3381                    );
3382                    move |s: usize, e: usize| {
3383                        vbit_range2_f32(bytes, vbit_offsets, x1, x2, rows, cols, o1, o2, s, e)
3384                    }
3385                });
3386                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3387                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3388                pool.run_many(&parts);
3389            }
3390            return;
3391        }
3392
3393        if uniform_f32 {
3394            let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
3395            let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
3396            let closures: [_; N] = std::array::from_fn(|i| {
3397                let Self::F32 { data, cols, .. } = ts[i] else {
3398                    unreachable!()
3399                };
3400                let (o1, o2) = (p1[i], p2[i]);
3401                move |start: usize, end: usize| {
3402                    for o in start..end {
3403                        let row = &data[o * cols..(o + 1) * cols];
3404                        let (mut s1, mut s2) = (0.0f32, 0.0f32);
3405                        for j in 0..*cols {
3406                            s1 += row[j] * x1[j];
3407                            s2 += row[j] * x2[j];
3408                        }
3409                        // SAFETY: disjoint (tensor, row) cells per worker.
3410                        unsafe {
3411                            *o1.at(o) = s1;
3412                            *o2.at(o) = s2;
3413                        }
3414                    }
3415                }
3416            });
3417            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3418                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3419            pool.run_many(&parts);
3420            return;
3421        }
3422
3423        struct Ctx<'a> {
3424            bytes: &'a [u8],
3425            row_scale: &'a [f32],
3426            cols: usize,
3427            xs1: std::borrow::Cow<'a, [f32]>,
3428            xs2: std::borrow::Cow<'a, [f32]>,
3429        }
3430        let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
3431            let Self::Mapped {
3432                dtype,
3433                cols,
3434                row_scale,
3435                col_field,
3436                ..
3437            } = ts[i]
3438            else {
3439                unreachable!()
3440            };
3441            Ctx {
3442                bytes: ts[i].quant_bytes(),
3443                row_scale,
3444                cols: *cols,
3445                xs1: prescale(x1, col_field, *dtype),
3446                xs2: prescale(x2, col_field, *dtype),
3447            }
3448        });
3449        let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
3450        let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
3451        #[cfg(target_arch = "aarch64")]
3452        if sdot_enabled() {
3453            let acts: [(SplitAct, SplitAct); N] =
3454                std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
3455            let closures: [_; N] = std::array::from_fn(|i| {
3456                let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
3457                move |start: usize, end: usize| {
3458                    q8_range2_sdot(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
3459                }
3460            });
3461            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3462                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3463            pool.run_many(&parts);
3464            return;
3465        }
3466        #[cfg(target_arch = "x86_64")]
3467        if avx2_a8w8_enabled() {
3468            let acts: [(SplitAct, SplitAct); N] =
3469                std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
3470            let closures: [_; N] = std::array::from_fn(|i| {
3471                let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
3472                move |start: usize, end: usize| {
3473                    q8_range2_avx2(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
3474                }
3475            });
3476            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3477                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3478            pool.run_many(&parts);
3479            return;
3480        }
3481        let closures: [_; N] = std::array::from_fn(|i| {
3482            let (c, o1, o2) = (&ctxs[i], p1[i], p2[i]);
3483            move |start: usize, end: usize| {
3484                q8_range2_f32(
3485                    c.bytes,
3486                    c.row_scale,
3487                    &c.xs1,
3488                    &c.xs2,
3489                    c.cols,
3490                    o1,
3491                    o2,
3492                    start,
3493                    end,
3494                )
3495            }
3496        });
3497        let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3498            std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3499        pool.run_many(&parts);
3500    }
3501
3502    /// Fused gate+up matvec with SiLU·mul: for each row r, computes
3503    /// `silu(gate·x) * (up·x)` and writes to `out[r]`. ONE pool dispatch,
3504    /// no intermediate g/u buffers, no separate silu pass. Falls back
3505    /// (returns false) for unsupported dtype combos.
3506    pub fn matvec_silu_mul(
3507        gate: &QTensor,
3508        up: &QTensor,
3509        x: &[f32],
3510        out: &mut [f32],
3511        pool: Option<&Pool>,
3512    ) -> bool {
3513        Self::matvec_silu_mul_limited(gate, up, x, out, 0.0, pool)
3514    }
3515
3516    /// Fused gate+up+SiLU with the GLM asymmetrical clamp.  `limit == 0`
3517    /// preserves the historical unclamped helper; a positive limit clamps
3518    /// `up` to both sides and `gate` only from above, matching the GLM
3519    /// SwiGLU reference.  Keeping the limit in the row kernel avoids the two
3520    /// intermediate vectors and the extra combine pass on the Q2TP experts.
3521    pub fn matvec_silu_mul_limited(
3522        gate: &QTensor,
3523        up: &QTensor,
3524        x: &[f32],
3525        out: &mut [f32],
3526        limit: f32,
3527        pool: Option<&Pool>,
3528    ) -> bool {
3529        if gate.flint_exact_cpu_default() || up.flint_exact_cpu_default() {
3530            return float_activations_scope(|| {
3531                Self::matvec_silu_mul_limited(gate, up, x, out, limit, pool)
3532            });
3533        }
3534        if gate.has_prism_contract()
3535            || up.has_prism_contract()
3536            || gate.per_tensor_only()
3537            || up.per_tensor_only()
3538        {
3539            // The fused gate/up kernels consume x directly.  Prism requires
3540            // a per-matrix signed FWHT, so the caller must use two ordinary
3541            // descriptor-aware matvecs instead of an unrotated fast path.
3542            return false;
3543        }
3544        let inter = gate.rows();
3545        debug_assert_eq!(up.rows(), inter);
3546        debug_assert_eq!(out.len(), inter);
3547        debug_assert_eq!(gate.cols(), up.cols());
3548        if !a8w8_enabled() {
3549            return false;
3550        }
3551        let act = split_act(x);
3552        let act = &act;
3553        let x_ref = x;
3554        let out_addr = SendMut(out.as_mut_ptr());
3555
3556        match (gate, up) {
3557            // Q4Block gate + Q4Block up (most common mobile q4 models)
3558            (
3559                Self::Mapped {
3560                    dtype: TensorDtype::Q4Block,
3561                    ..
3562                },
3563                Self::Mapped {
3564                    dtype: TensorDtype::Q4Block,
3565                    ..
3566                },
3567            ) => {
3568                let (gp, gs) = q4_split(gate.quant_bytes(), gate.rows(), gate.cols());
3569                let (up_p, up_s) = q4_split(up.quant_bytes(), up.rows(), up.cols());
3570                let gpr = gate.cols() / GROUP_SIZE;
3571                let cols = gate.cols();
3572                let run = move |start: usize, end: usize| {
3573                    for r in start..end {
3574                        let mut gv = dot_q4_row_i8(gp, gs, r * gpr, gpr, &act.xq) * act.sx;
3575                        let mut uv = dot_q4_row_i8(up_p, up_s, r * gpr, gpr, &act.xq) * act.sx;
3576                        for &(j, xv) in &act.outliers {
3577                            let flat = r * cols + j;
3578                            let gb = gp[flat / 2];
3579                            let gn = if flat & 1 == 0 { gb & 0x0F } else { gb >> 4 };
3580                            let gsc = f16_to_f32(u16::from_le_bytes([
3581                                gs[(flat / GROUP_SIZE) * 2],
3582                                gs[(flat / GROUP_SIZE) * 2 + 1],
3583                            ]));
3584                            gv += ((gn as i32 - 8) as f32) * gsc * xv;
3585                            let ub = up_p[flat / 2];
3586                            let un = if flat & 1 == 0 { ub & 0x0F } else { ub >> 4 };
3587                            let usc = f16_to_f32(u16::from_le_bytes([
3588                                up_s[(flat / GROUP_SIZE) * 2],
3589                                up_s[(flat / GROUP_SIZE) * 2 + 1],
3590                            ]));
3591                            uv += ((un as i32 - 8) as f32) * usc * xv;
3592                        }
3593                        // SAFETY: disjoint row ranges per worker.
3594                        unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3595                    }
3596                };
3597                dispatch_rows(pool, inter, &run);
3598                true
3599            }
3600            // Q4Tiled gate + Q4Tiled up — one row pass, both tile
3601            // streams sequential, silu·mul fused (same per-row math as
3602            // `q4t_matvec`).
3603            (
3604                Self::Mapped {
3605                    dtype: TensorDtype::Q4Tiled,
3606                    ..
3607                },
3608                Self::Mapped {
3609                    dtype: TensorDtype::Q4Tiled,
3610                    ..
3611                },
3612            ) => {
3613                let g_bytes = gate.quant_bytes();
3614                let u_bytes = up.quant_bytes();
3615                let gpr = gate.cols() / GROUP_SIZE;
3616                let run = move |start: usize, end: usize| {
3617                    for r in start..end {
3618                        let mut gv = dot_q4t_row_i8(g_bytes, r, gpr, &act.xq) * act.sx;
3619                        let mut uv = dot_q4t_row_i8(u_bytes, r, gpr, &act.xq) * act.sx;
3620                        for &(j, xv) in &act.outliers {
3621                            let (w, s) = q4t_outlier(g_bytes, r, gpr, j);
3622                            gv += w * s * xv;
3623                            let (w, s) = q4t_outlier(u_bytes, r, gpr, j);
3624                            uv += w * s * xv;
3625                        }
3626                        // SAFETY: disjoint row ranges per worker.
3627                        unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3628                    }
3629                };
3630                dispatch_rows(pool, inter, &run);
3631                true
3632            }
3633            // Q4TiledP gate + Q4TiledP up — the same fused row pass, with
3634            // each row's two ladders built once and spent on both streams.
3635            (
3636                Self::Mapped {
3637                    dtype: TensorDtype::Q4TiledP,
3638                    ..
3639                },
3640                Self::Mapped {
3641                    dtype: TensorDtype::Q4TiledP,
3642                    ..
3643                },
3644            ) => {
3645                let cols = gate.cols();
3646                let gpr = cols / GROUP_SIZE;
3647                let gv_view = Q4tpView::new(gate.quant_bytes(), inter, cols);
3648                let uv_view = Q4tpView::new(up.quant_bytes(), inter, cols);
3649                let run = |start: usize, end: usize| {
3650                    with_krows(gpr, |gsc, usc| {
3651                        for r in start..end {
3652                            gv_view.scales_into(r, gpr, gsc);
3653                            uv_view.scales_into(r, gpr, usc);
3654                            let mut gv =
3655                                dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsc) * act.sx;
3656                            let mut uv =
3657                                dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, usc) * act.sx;
3658                            for &(j, xv) in &act.outliers {
3659                                let (w, s) = q4tp_outlier(gv_view.nib, r, gpr, j, gsc);
3660                                gv += w * s * xv;
3661                                let (w, s) = q4tp_outlier(uv_view.nib, r, gpr, j, usc);
3662                                uv += w * s * xv;
3663                            }
3664                            // SAFETY: disjoint row ranges per worker.
3665                            unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3666                        }
3667                    });
3668                };
3669                dispatch_rows(pool, inter, &run);
3670                true
3671            }
3672            // Q1 gate + Q1 up — one row pass over both sign streams,
3673            // silu·mul fused (the per-row math of `q1_range_a8w8`); the
3674            // activation group sums are shared by both streams. Without
3675            // this arm a q1 dense FFN paid two dispatches + a combine
3676            // loop — the exact barrier this function exists to remove.
3677            // q1g (and a q1/q1g mix) rides the same arm through `BinW`.
3678            (
3679                Self::Mapped {
3680                    dtype: TensorDtype::Q1 | TensorDtype::Q1G,
3681                    ..
3682                },
3683                Self::Mapped {
3684                    dtype: TensorDtype::Q1 | TensorDtype::Q1G,
3685                    ..
3686                },
3687            ) => {
3688                let gw = gate.bin_w().expect("binary dtype matched");
3689                let uw = up.bin_w().expect("binary dtype matched");
3690                let gpr = gate.cols() / GROUP_SIZE;
3691                let gsum = q1_group_sums(&act.xq, gpr);
3692                let gsum = &gsum;
3693                let run = move |start: usize, end: usize| {
3694                    for r in start..end {
3695                        let mut gv = gw.dot_i8(r, gpr, &act.xq, gsum) * act.sx;
3696                        let mut uv = uw.dot_i8(r, gpr, &act.xq, gsum) * act.sx;
3697                        for &(j, xv) in &act.outliers {
3698                            let (w, s) = gw.outlier(r, gpr, j);
3699                            gv += w * s * xv;
3700                            let (w, s) = uw.outlier(r, gpr, j);
3701                            uv += w * s * xv;
3702                        }
3703                        // SAFETY: disjoint row ranges per worker.
3704                        unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3705                    }
3706                };
3707                dispatch_rows(pool, inter, &run);
3708                true
3709            }
3710            // Q2TiledP gate + Q2TiledP up — the 2-bit expert pair (MoE
3711            // FFNs of the W2 class): one row pass, both ladders built
3712            // once, integer code dots with shared group sums.
3713            (
3714                Self::Mapped {
3715                    dtype: TensorDtype::Q2TiledP,
3716                    ..
3717                },
3718                Self::Mapped {
3719                    dtype: TensorDtype::Q2TiledP,
3720                    ..
3721                },
3722            ) => {
3723                let cols = gate.cols();
3724                let gpr = cols / GROUP_SIZE;
3725                let gv_view = Q4tpView::new_q2(gate.quant_bytes(), inter, cols);
3726                let uv_view = Q4tpView::new_q2(up.quant_bytes(), inter, cols);
3727                let gsum = q1_group_sums(&act.xq, gpr);
3728                let gsum = &gsum;
3729                let run = move |start: usize, end: usize| {
3730                    with_krows(gpr, |gsc, usc| {
3731                        for r in start..end {
3732                            gv_view.scales_into(r, gpr, gsc);
3733                            uv_view.scales_into(r, gpr, usc);
3734                            let mut gv =
3735                                dot_q2tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsum, gsc) * act.sx;
3736                            let mut uv =
3737                                dot_q2tp_row_i8(uv_view.nib, r, gpr, &act.xq, gsum, usc) * act.sx;
3738                            for &(j, xv) in &act.outliers {
3739                                let (w, s) = q2tp_outlier(gv_view.nib, r, gpr, j, gsc);
3740                                gv += w * s * xv;
3741                                let (w, s) = q2tp_outlier(uv_view.nib, r, gpr, j, usc);
3742                                uv += w * s * xv;
3743                            }
3744                            // SAFETY: disjoint row ranges per worker.
3745                            unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3746                        }
3747                    });
3748                };
3749                dispatch_rows(pool, inter, &run);
3750                true
3751            }
3752            // Q8Row gate + Q8Row up — one row pass over both i8 streams.
3753            // Q8_2f stays out on purpose: its column field prescales the
3754            // activations PER TENSOR, which breaks this fn's shared
3755            // split_act contract — it keeps the two-dispatch path.
3756            (
3757                Self::Mapped {
3758                    dtype: TensorDtype::Q8Row,
3759                    row_scale: g_rs,
3760                    ..
3761                },
3762                Self::Mapped {
3763                    dtype: TensorDtype::Q8Row,
3764                    row_scale: u_rs,
3765                    ..
3766                },
3767            ) => {
3768                let g_bytes = gate.quant_bytes();
3769                let u_bytes = up.quant_bytes();
3770                let cols = gate.cols();
3771                let run = move |start: usize, end: usize| {
3772                    for r in start..end {
3773                        let gv = q8_row_dot(&g_bytes[r * cols..(r + 1) * cols], act) * g_rs[r];
3774                        let uv = q8_row_dot(&u_bytes[r * cols..(r + 1) * cols], act) * u_rs[r];
3775                        // SAFETY: disjoint row ranges per worker.
3776                        unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3777                    }
3778                };
3779                dispatch_rows(pool, inter, &run);
3780                true
3781            }
3782            // Q1T gate + Q1T up
3783            (
3784                Self::Mapped {
3785                    dtype: TensorDtype::Q1T,
3786                    ..
3787                },
3788                Self::Mapped {
3789                    dtype: TensorDtype::Q1T,
3790                    ..
3791                },
3792            ) => {
3793                const TILE: usize = cortiq_core::quant::Q1T_TILE;
3794                let g_bytes = gate.quant_bytes();
3795                let u_bytes = up.quant_bytes();
3796                let gpr = gate.cols() / GROUP_SIZE;
3797                let (g_rp, g_ent, g_ov) = q1t_overlay(g_bytes, inter * gpr * TILE, inter);
3798                let (u_rp, u_ent, u_ov) = q1t_overlay(u_bytes, inter * gpr * TILE, inter);
3799                let run = move |start: usize, end: usize| {
3800                    for r in start..end {
3801                        let mut gv = q1t_dot_row_i8(g_bytes, r, gpr, &act.xq) * act.sx;
3802                        let mut uv = q1t_dot_row_i8(u_bytes, r, gpr, &act.xq) * act.sx;
3803                        for &(j, xv) in &act.outliers {
3804                            gv += q1t_base_weight(g_bytes, r, gpr, j) * xv;
3805                            uv += q1t_base_weight(u_bytes, r, gpr, j) * xv;
3806                        }
3807                        gv += q1t_row_outlier_correction(g_bytes, r, g_rp, g_ent, g_ov, x_ref);
3808                        uv += q1t_row_outlier_correction(u_bytes, r, u_rp, u_ent, u_ov, x_ref);
3809                        // SAFETY: disjoint row ranges per worker.
3810                        unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3811                    }
3812                };
3813                dispatch_rows(pool, inter, &run);
3814                true
3815            }
3816            _ => false,
3817        }
3818    }
3819
3820    /// Every routed expert's fused gate/up/SiLU under ONE pool dispatch.
3821    ///
3822    /// The per-expert path pays a pool barrier per expert per stage: at 9
3823    /// experts over 40 layers that is ~720 barriers a token, and a decode
3824    /// profile of Qwen3.6-35B-A3B showed the pool parked in
3825    /// `psynch_cvwait` about twice as long as it spent computing. Laying
3826    /// every expert's rows end-to-end in one virtual row space collapses
3827    /// the stage to a single dispatch. The per-row body is the
3828    /// single-expert q4tp arm verbatim, so outputs are bit-identical.
3829    ///
3830    /// `false` = something is outside the fused q4tp kernel (dtype, shape,
3831    /// or a transformed tensor); the caller walks the ordinary per-expert
3832    /// path. Float activations use the same exact scalar rows, still fused
3833    /// under one pool dispatch.
3834    pub fn moe_gate_up_many(
3835        pairs: &[(&QTensor, &QTensor)],
3836        x: &[f32],
3837        outs: &mut [Vec<f32>],
3838        pool: Option<&Pool>,
3839    ) -> bool {
3840        Self::moe_gate_up_many_limited(pairs, x, outs, 0.0, pool)
3841    }
3842
3843    /// Batched gate/up/SiLU with the optional GLM clamp.  The public legacy
3844    /// helper above keeps its historical unclamped semantics; callers that
3845    /// implement a reference with a positive SwiGLU limit use this variant.
3846    pub fn moe_gate_up_many_limited(
3847        pairs: &[(&QTensor, &QTensor)],
3848        x: &[f32],
3849        outs: &mut [Vec<f32>],
3850        limit: f32,
3851        pool: Option<&Pool>,
3852    ) -> bool {
3853        if Self::any_flint_exact_cpu_default_pairs(pairs) {
3854            return float_activations_scope(|| {
3855                Self::moe_gate_up_many_limited(pairs, x, outs, limit, pool)
3856            });
3857        }
3858        if pairs.is_empty() || pairs.len() != outs.len() {
3859            return false;
3860        }
3861        if pairs.iter().any(|(g, u)| g.per_tensor_only() || u.per_tensor_only()) {
3862            return false;
3863        }
3864        if !a8w8_enabled() {
3865            if limit > 0.0 {
3866                // The exact-row fallback applies no SwiGLU clamp; the
3867                // caller's per-expert path carries it instead.
3868                return false;
3869            }
3870            let groups = vec![vec![0]; pairs.len()];
3871            return Self::moe_gate_up_rows(pairs, &groups, x, outs, pool);
3872        }
3873        let inter = pairs[0].0.rows();
3874        let cols = pairs[0].0.cols();
3875        if cols % GROUP_SIZE != 0 {
3876            return false;
3877        }
3878        let gpr = cols / GROUP_SIZE;
3879        // Uniform layout across every routed pair: q4tp, or the 2-bit
3880        // profile's q2tp gate/up (the W2 class). Mixed sets refuse.
3881        let q2 = matches!(
3882            pairs[0].0,
3883            Self::Mapped {
3884                dtype: TensorDtype::Q2TiledP,
3885                ..
3886            }
3887        );
3888        let want = if q2 {
3889            TensorDtype::Q2TiledP
3890        } else {
3891            TensorDtype::Q4TiledP
3892        };
3893        let mut views = Vec::with_capacity(pairs.len() * 2);
3894        for ((g, u), o) in pairs.iter().zip(outs.iter()) {
3895            let both = matches!(g, Self::Mapped { dtype, .. } if *dtype == want)
3896                && matches!(u, Self::Mapped { dtype, .. } if *dtype == want);
3897            if !both
3898                || g.rows() != inter
3899                || u.rows() != inter
3900                || g.cols() != cols
3901                || u.cols() != cols
3902                || o.len() != inter
3903            {
3904                return false;
3905            }
3906            let mk = if q2 { Q4tpView::new_q2 } else { Q4tpView::new };
3907            views.push(mk(g.quant_bytes(), inter, cols));
3908            views.push(mk(u.quant_bytes(), inter, cols));
3909        }
3910        let act = split_act(x);
3911        let gsum = if q2 {
3912            q1_group_sums(&act.xq, gpr)
3913        } else {
3914            Vec::new()
3915        };
3916        let (act, gsum) = (&act, &gsum);
3917        let ptrs: Vec<SendMut> = outs.iter_mut().map(|o| SendMut(o.as_mut_ptr())).collect();
3918        let (views, ptrs) = (&views, &ptrs);
3919        let run = |start: usize, end: usize| {
3920            with_krows(gpr, |gsc, usc| {
3921                for flat in start..end {
3922                    let (e, r) = (flat / inter, flat % inter);
3923                    let gv_view = &views[e * 2];
3924                    let uv_view = &views[e * 2 + 1];
3925                    gv_view.scales_into(r, gpr, gsc);
3926                    uv_view.scales_into(r, gpr, usc);
3927                    let (mut gv, mut uv) = if q2 {
3928                        (
3929                            dot_q2tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsum, gsc) * act.sx,
3930                            dot_q2tp_row_i8(uv_view.nib, r, gpr, &act.xq, gsum, usc) * act.sx,
3931                        )
3932                    } else {
3933                        (
3934                            dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsc) * act.sx,
3935                            dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, usc) * act.sx,
3936                        )
3937                    };
3938                    for &(j, xv) in &act.outliers {
3939                        let (og, ou) = if q2 {
3940                            (
3941                                q2tp_outlier(gv_view.nib, r, gpr, j, gsc),
3942                                q2tp_outlier(uv_view.nib, r, gpr, j, usc),
3943                            )
3944                        } else {
3945                            (
3946                                q4tp_outlier(gv_view.nib, r, gpr, j, gsc),
3947                                q4tp_outlier(uv_view.nib, r, gpr, j, usc),
3948                            )
3949                        };
3950                        gv += og.0 * og.1 * xv;
3951                        uv += ou.0 * ou.1 * xv;
3952                    }
3953                    // SAFETY: one worker owns each (expert, row) pair.
3954                    unsafe { *ptrs[e].at(r) = silu_mul_limited(gv, uv, limit) };
3955                }
3956            });
3957        };
3958        dispatch_rows(pool, pairs.len() * inter, &run);
3959        true
3960    }
3961
3962    /// Every routed expert's down projection, weighted and summed into
3963    /// `out`, under ONE pool dispatch.
3964    ///
3965    /// Partitioned by OUTPUT row rather than by expert: each row is owned
3966    /// by a single worker, so the experts are summed in the caller's order
3967    /// — the same sequence of f32 adds the serial `out[i] += w·eo[i]` loop
3968    /// performs, hence bit-identical. Partitioning by expert instead would
3969    /// race on the shared accumulator.
3970    pub fn moe_down_many(
3971        downs: &[&QTensor],
3972        gs: &[Vec<f32>],
3973        weights: &[f32],
3974        out: &mut [f32],
3975        pool: Option<&Pool>,
3976    ) -> bool {
3977        if Self::any_flint_exact_cpu_default(downs) {
3978            return float_activations_scope(|| {
3979                Self::moe_down_many(downs, gs, weights, out, pool)
3980            });
3981        }
3982        if downs.is_empty() || downs.len() != gs.len() || downs.len() != weights.len() {
3983            return false;
3984        }
3985        if downs.iter().any(|d| d.per_tensor_only()) {
3986            return false;
3987        }
3988        if !a8w8_enabled() {
3989            let mut terms = vec![vec![0.0; out.len()]; downs.len()];
3990            if !Self::moe_down_rows(downs, &vec![1; downs.len()], gs, &mut terms, pool) {
3991                return false;
3992            }
3993            out.fill(0.0);
3994            for (row, &w) in terms.iter().zip(weights) {
3995                for (o, &v) in out.iter_mut().zip(row) {
3996                    *o += w * v;
3997                }
3998            }
3999            return true;
4000        }
4001        let rows = out.len();
4002        let cols = downs[0].cols();
4003        if cols % GROUP_SIZE != 0 {
4004            return false;
4005        }
4006        let gpr = cols / GROUP_SIZE;
4007        let mut views = Vec::with_capacity(downs.len());
4008        for (d, g) in downs.iter().zip(gs.iter()) {
4009            if !matches!(
4010                d,
4011                Self::Mapped {
4012                    dtype: TensorDtype::Q4TiledP,
4013                    ..
4014                }
4015            ) || d.rows() != rows
4016                || d.cols() != cols
4017                || g.len() != cols
4018            {
4019                return false;
4020            }
4021            views.push(Q4tpView::new(d.quant_bytes(), rows, cols));
4022        }
4023        // One int8 split per expert — the activation vectors differ.
4024        let acts: Vec<SplitAct> = gs.iter().map(|g| split_act(g)).collect();
4025        // Partitioned by OUTPUT row, with the experts folded inside: each
4026        // row is owned by one worker, so they are summed in the caller's
4027        // order — the same f32 sequence the serial `out[i] += w·eo[i]`
4028        // loop produces. Partitioning by expert instead would either race
4029        // on the accumulator or need a scratch plane and a second pass;
4030        // measured, that variant was a wash, so this keeps the simpler
4031        // shape.
4032        let out_addr = SendMut(out.as_mut_ptr());
4033        let (views, acts, weights) = (&views, &acts, &weights);
4034        let run = |start: usize, end: usize| {
4035            with_krow(gpr, |sc| {
4036                for r in start..end {
4037                    let mut acc = 0f32;
4038                    for (e, v) in views.iter().enumerate() {
4039                        v.scales_into(r, gpr, sc);
4040                        let a = &acts[e];
4041                        let mut d = dot_q4tp_row_i8(v.nib, r, gpr, &a.xq, sc) * a.sx;
4042                        for &(j, xv) in &a.outliers {
4043                            let (w, s) = q4tp_outlier(v.nib, r, gpr, j, sc);
4044                            d += w * s * xv;
4045                        }
4046                        acc += weights[e] * d;
4047                    }
4048                    // SAFETY: disjoint row ranges per worker.
4049                    unsafe { *out_addr.at(r) = acc };
4050                }
4051            });
4052        };
4053        dispatch_rows(pool, rows, &run);
4054        true
4055    }
4056
4057    /// `moe_gate_up_many` for SEVERAL tokens at once, decode-exact: expert
4058    /// `e` (`pairs[e]`, q4tp) serves the tokens `groups[e]` (row indices
4059    /// into `xs`, each `cols` wide). Every (expert, token) output is
4060    /// bit-identical to `moe_gate_up_many` run on that token alone — the
4061    /// same int8 activation split, VNNI dots, outlier terms and inline
4062    /// SiLU — while each weight row is read once for all the tokens routed
4063    /// to its expert (the speculative verify's expert sharing). `outs` is
4064    /// flat in (expert, token-of-group) order. False = not covered (not
4065    /// q4tp): the caller takes the per-token path. With float activations,
4066    /// the exact scalar row kernel replaces the int8 dot without changing
4067    /// the shared dispatch or route-order reduction.
4068    pub fn moe_gate_up_rows(
4069        pairs: &[(&QTensor, &QTensor)],
4070        groups: &[Vec<usize>],
4071        xs: &[f32],
4072        outs: &mut [Vec<f32>],
4073        pool: Option<&Pool>,
4074    ) -> bool {
4075        if Self::any_flint_exact_cpu_default_pairs(pairs) {
4076            return float_activations_scope(|| {
4077                Self::moe_gate_up_rows(pairs, groups, xs, outs, pool)
4078            });
4079        }
4080        if pairs.is_empty() || pairs.len() != groups.len() {
4081            return false;
4082        }
4083        if pairs.iter().any(|(g, u)| g.per_tensor_only() || u.per_tensor_only()) {
4084            return false;
4085        }
4086        let inter = pairs[0].0.rows();
4087        let cols = pairs[0].0.cols();
4088        let n_pairs: usize = groups.iter().map(|g| g.len()).sum();
4089        if cols == 0 || cols % GROUP_SIZE != 0 || outs.len() != n_pairs || xs.len() % cols != 0 {
4090            return false;
4091        }
4092        let b = xs.len() / cols;
4093        let gpr = cols / GROUP_SIZE;
4094        let mut views = Vec::with_capacity(pairs.len() * 2);
4095        for (g, u) in pairs {
4096            let q4tp = |t: &QTensor| {
4097                matches!(
4098                    t,
4099                    Self::Mapped {
4100                        dtype: TensorDtype::Q4TiledP,
4101                        ..
4102                    }
4103                )
4104            };
4105            if g.has_prism_contract()
4106                || u.has_prism_contract()
4107                || !q4tp(g)
4108                || !q4tp(u)
4109                || g.rows() != inter
4110                || u.rows() != inter
4111                || g.cols() != cols
4112                || u.cols() != cols
4113            {
4114                return false;
4115            }
4116            views.push(Q4tpView::new(g.quant_bytes(), inter, cols));
4117            views.push(Q4tpView::new(u.quant_bytes(), inter, cols));
4118        }
4119        if outs.iter().any(|o| o.len() != inter) || groups.iter().flatten().any(|&t| t >= b) {
4120            return false;
4121        }
4122        let quantized = a8w8_enabled();
4123        let acts: Vec<SplitAct> = if quantized {
4124            (0..b)
4125                .map(|t| split_act(&xs[t * cols..(t + 1) * cols]))
4126                .collect()
4127        } else {
4128            Vec::new()
4129        };
4130        let mut offs = Vec::with_capacity(groups.len());
4131        let mut o = 0usize;
4132        for g in groups {
4133            offs.push(o);
4134            o += g.len();
4135        }
4136        let ptrs: Vec<SendMut> = outs.iter_mut().map(|o| SendMut(o.as_mut_ptr())).collect();
4137        let (views, ptrs, acts, offs) = (&views, &ptrs, &acts, &offs);
4138        let run = |start: usize, end: usize| {
4139            let (mut gsc, mut usc) = (vec![0f32; gpr], vec![0f32; gpr]);
4140            for flat in start..end {
4141                let (e, r) = (flat / inter, flat % inter);
4142                let (gv_view, uv_view) = (&views[e * 2], &views[e * 2 + 1]);
4143                gv_view.scales_into(r, gpr, &mut gsc);
4144                uv_view.scales_into(r, gpr, &mut usc);
4145                for (k, &t) in groups[e].iter().enumerate() {
4146                    if !quantized {
4147                        let x = &xs[t * cols..(t + 1) * cols];
4148                        let gv = q4tp_row_exact(gv_view.nib, r, gpr, x, &gsc);
4149                        let uv = q4tp_row_exact(uv_view.nib, r, gpr, x, &usc);
4150                        unsafe { *ptrs[offs[e] + k].at(r) = (gv / (1.0 + (-gv).exp())) * uv };
4151                        continue;
4152                    }
4153                    let act = &acts[t];
4154                    let mut gv = dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, &gsc) * act.sx;
4155                    let mut uv = dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, &usc) * act.sx;
4156                    for &(j, xv) in &act.outliers {
4157                        let og = q4tp_outlier(gv_view.nib, r, gpr, j, &gsc);
4158                        let ou = q4tp_outlier(uv_view.nib, r, gpr, j, &usc);
4159                        gv += og.0 * og.1 * xv;
4160                        uv += ou.0 * ou.1 * xv;
4161                    }
4162                    let silu_g = gv / (1.0 + (-gv).exp());
4163                    // SAFETY: one worker owns each (expert, row) cell of
4164                    // every output of the expert's group.
4165                    unsafe { *ptrs[offs[e] + k].at(r) = silu_g * uv };
4166                }
4167            }
4168        };
4169        dispatch_rows(pool, pairs.len() * inter, &run);
4170        true
4171    }
4172
4173    /// The per-(expert, token) down terms `moe_down_many` weights and sums,
4174    /// for SEVERAL tokens: `outs[p][o] = down_e[o] · gs[p]` (int8 split of
4175    /// `gs[p]`, VNNI dot, outlier terms — bit-identical to that kernel's
4176    /// `d`), each down row read once for its expert's whole group. The
4177    /// caller sums `w·d` per token in its route order, which reproduces
4178    /// `moe_down_many`'s f32 sequence exactly. Layout as `moe_gate_up_rows`.
4179    pub fn moe_down_rows(
4180        downs: &[&QTensor],
4181        group_lens: &[usize],
4182        gs: &[Vec<f32>],
4183        outs: &mut [Vec<f32>],
4184        pool: Option<&Pool>,
4185    ) -> bool {
4186        if Self::any_flint_exact_cpu_default(downs) {
4187            return float_activations_scope(|| {
4188                Self::moe_down_rows(downs, group_lens, gs, outs, pool)
4189            });
4190        }
4191        if downs.is_empty() || downs.len() != group_lens.len() {
4192            return false;
4193        }
4194        if downs.iter().any(|d| d.per_tensor_only()) {
4195            return false;
4196        }
4197        let rows = downs[0].rows();
4198        let cols = downs[0].cols();
4199        let n_pairs: usize = group_lens.iter().sum();
4200        if cols == 0 || cols % GROUP_SIZE != 0 || gs.len() != n_pairs || outs.len() != n_pairs {
4201            return false;
4202        }
4203        let gpr = cols / GROUP_SIZE;
4204        let mut views = Vec::with_capacity(downs.len());
4205        for d in downs {
4206            if d.has_prism_contract()
4207                || !matches!(
4208                    d,
4209                    Self::Mapped {
4210                        dtype: TensorDtype::Q4TiledP,
4211                        ..
4212                    }
4213                )
4214                || d.rows() != rows
4215                || d.cols() != cols
4216            {
4217                return false;
4218            }
4219            views.push(Q4tpView::new(d.quant_bytes(), rows, cols));
4220        }
4221        if gs.iter().any(|g| g.len() != cols) || outs.iter().any(|o| o.len() != rows) {
4222            return false;
4223        }
4224        let quantized = a8w8_enabled();
4225        let acts: Vec<SplitAct> = if quantized {
4226            gs.iter().map(|g| split_act(g)).collect()
4227        } else {
4228            Vec::new()
4229        };
4230        let mut offs = Vec::with_capacity(group_lens.len());
4231        let mut o = 0usize;
4232        for &l in group_lens {
4233            offs.push(o);
4234            o += l;
4235        }
4236        let ptrs: Vec<SendMut> = outs.iter_mut().map(|o| SendMut(o.as_mut_ptr())).collect();
4237        let (views, ptrs, acts, offs) = (&views, &ptrs, &acts, &offs);
4238        let run = |start: usize, end: usize| {
4239            let mut sc = vec![0f32; gpr];
4240            for flat in start..end {
4241                let (e, r) = (flat / rows, flat % rows);
4242                let v = &views[e];
4243                v.scales_into(r, gpr, &mut sc);
4244                for k in 0..group_lens[e] {
4245                    if !quantized {
4246                        let d = q4tp_row_exact(v.nib, r, gpr, &gs[offs[e] + k], &sc);
4247                        unsafe { *ptrs[offs[e] + k].at(r) = d };
4248                        continue;
4249                    }
4250                    let a = &acts[offs[e] + k];
4251                    let mut d = dot_q4tp_row_i8(v.nib, r, gpr, &a.xq, &sc) * a.sx;
4252                    for &(j, xv) in &a.outliers {
4253                        let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
4254                        d += w * s * xv;
4255                    }
4256                    // SAFETY: one worker owns each (expert, row) cell.
4257                    unsafe { *ptrs[offs[e] + k].at(r) = d };
4258                }
4259            }
4260        };
4261        dispatch_rows(pool, downs.len() * rows, &run);
4262        true
4263    }
4264}
4265
4266/// Batched q8 kernel: same math as qmatvec, the row makes a single
4267/// pass from memory for the whole batch.
4268/// Accelerate CBLAS — the Apple AMX matrix units, the same engine
4269/// llama.cpp's `-ngl 0` prefill rides via ggml-blas.
4270#[cfg(target_os = "macos")]
4271mod accel_blas {
4272    #[link(name = "Accelerate", kind = "framework")]
4273    unsafe extern "C" {
4274        pub fn cblas_sgemm(
4275            order: i32,
4276            trans_a: i32,
4277            trans_b: i32,
4278            m: i32,
4279            n: i32,
4280            k: i32,
4281            alpha: f32,
4282            a: *const f32,
4283            lda: i32,
4284            b: *const f32,
4285            ldb: i32,
4286            beta: f32,
4287            c: *mut f32,
4288            ldc: i32,
4289        );
4290    }
4291}
4292
4293#[cfg(target_os = "macos")]
4294pub(crate) fn accel_gemm_enabled() -> bool {
4295    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4296    *ON.get_or_init(|| std::env::var("CMF_ACCEL").map(|v| v != "0").unwrap_or(true))
4297}
4298
4299/// Off macOS the "accel" GEMM is the portable NEON micro-kernel below —
4300/// same entry point, so the batched-attention path opens on mobile.
4301#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
4302pub(crate) fn accel_gemm_enabled() -> bool {
4303    true
4304}
4305
4306/// Portable NEON f32 GEMM (row-major, optional Bᵀ): a 4×8 fmla
4307/// micro-kernel with A broadcast against B panels — the mobile stand-in
4308/// for Accelerate in the batched causal attention (QKᵀ and P·V). Not a
4309/// BLAS: shapes here are the attention panels (m ≤ heads·chunk,
4310/// k = head_dim or context), and the goal is removing the per-position
4311/// quadratic wall, not peak GEMM.
4312#[cfg(target_arch = "aarch64")]
4313#[allow(clippy::too_many_arguments)]
4314pub(crate) fn neon_gemm_rm(
4315    m: usize,
4316    n: usize,
4317    k: usize,
4318    alpha: f32,
4319    a: &[f32],
4320    lda: usize,
4321    b_mat: &[f32],
4322    ldb: usize,
4323    b_rows_are_n: bool,
4324    c: &mut [f32],
4325    ldc: usize,
4326) {
4327    debug_assert!(a.len() >= (m - 1) * lda + k);
4328    debug_assert!(c.len() >= (m - 1) * ldc + n);
4329    // SAFETY: bounds asserted above; NEON is baseline on aarch64.
4330    unsafe {
4331        use core::arch::aarch64::*;
4332        let mut i = 0usize;
4333        while i < m {
4334            let mi = (m - i).min(4);
4335            let mut j = 0usize;
4336            while j < n {
4337                let nj = (n - j).min(8);
4338                if mi == 4 && nj == 8 {
4339                    let (mut c0a, mut c0b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
4340                    let (mut c1a, mut c1b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
4341                    let (mut c2a, mut c2b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
4342                    let (mut c3a, mut c3b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
4343                    for p in 0..k {
4344                        let (b0, b1) = if b_rows_are_n {
4345                            // B is [n, k]: column p of Bᵀ = element p of
4346                            // eight consecutive B rows — gathered.
4347                            let base = b_mat.as_ptr().add(j * ldb + p);
4348                            let g = |o: usize| *base.add(o * ldb);
4349                            ([g(0), g(1), g(2), g(3)], [g(4), g(5), g(6), g(7)])
4350                        } else {
4351                            let base = b_mat.as_ptr().add(p * ldb + j);
4352                            (
4353                                [*base, *base.add(1), *base.add(2), *base.add(3)],
4354                                [*base.add(4), *base.add(5), *base.add(6), *base.add(7)],
4355                            )
4356                        };
4357                        let bv0 = vld1q_f32(b0.as_ptr());
4358                        let bv1 = vld1q_f32(b1.as_ptr());
4359                        let a0 = vdupq_n_f32(*a.as_ptr().add(i * lda + p));
4360                        let a1 = vdupq_n_f32(*a.as_ptr().add((i + 1) * lda + p));
4361                        let a2 = vdupq_n_f32(*a.as_ptr().add((i + 2) * lda + p));
4362                        let a3 = vdupq_n_f32(*a.as_ptr().add((i + 3) * lda + p));
4363                        c0a = vfmaq_f32(c0a, a0, bv0);
4364                        c0b = vfmaq_f32(c0b, a0, bv1);
4365                        c1a = vfmaq_f32(c1a, a1, bv0);
4366                        c1b = vfmaq_f32(c1b, a1, bv1);
4367                        c2a = vfmaq_f32(c2a, a2, bv0);
4368                        c2b = vfmaq_f32(c2b, a2, bv1);
4369                        c3a = vfmaq_f32(c3a, a3, bv0);
4370                        c3b = vfmaq_f32(c3b, a3, bv1);
4371                    }
4372                    let al = vdupq_n_f32(alpha);
4373                    for (r, (ca, cb)) in [(c0a, c0b), (c1a, c1b), (c2a, c2b), (c3a, c3b)]
4374                        .iter()
4375                        .enumerate()
4376                    {
4377                        let dst = c.as_mut_ptr().add((i + r) * ldc + j);
4378                        vst1q_f32(dst, vmulq_f32(*ca, al));
4379                        vst1q_f32(dst.add(4), vmulq_f32(*cb, al));
4380                    }
4381                } else {
4382                    for r in 0..mi {
4383                        for q in 0..nj {
4384                            let mut acc = 0f32;
4385                            for p in 0..k {
4386                                let bv = if b_rows_are_n {
4387                                    b_mat[(j + q) * ldb + p]
4388                                } else {
4389                                    b_mat[p * ldb + j + q]
4390                                };
4391                                acc += a[(i + r) * lda + p] * bv;
4392                            }
4393                            c[(i + r) * ldc + j + q] = acc * alpha;
4394                        }
4395                    }
4396                }
4397                j += nj;
4398            }
4399            i += mi;
4400        }
4401    }
4402}
4403
4404/// Off-macOS aarch64: the batched attention rides the NEON micro-GEMM.
4405#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
4406#[allow(clippy::too_many_arguments)]
4407pub(crate) fn sgemm_rm(
4408    m: usize,
4409    n: usize,
4410    k: usize,
4411    alpha: f32,
4412    a: &[f32],
4413    lda: usize,
4414    b_mat: &[f32],
4415    ldb: usize,
4416    b_rows_are_n: bool,
4417    c: &mut [f32],
4418    ldc: usize,
4419) {
4420    neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
4421}
4422
4423/// Row-major f32 GEMM, exposed for offline tools (the AWNP pass builds a
4424/// per-layer projection and applies it to every expert; a naive triple loop
4425/// would turn a two-minute job into half an hour).
4426#[allow(clippy::too_many_arguments)]
4427pub fn sgemm_public(
4428    m: usize,
4429    n: usize,
4430    k: usize,
4431    alpha: f32,
4432    a: &[f32],
4433    lda: usize,
4434    b_mat: &[f32],
4435    ldb: usize,
4436    b_rows_are_n: bool,
4437    c: &mut [f32],
4438    ldc: usize,
4439) {
4440    #[cfg(any(target_os = "macos", target_arch = "aarch64"))]
4441    {
4442        sgemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
4443    }
4444    // x86 without Accelerate has no sgemm_rm (the specialized paths there are
4445    // quantized kernels): the packed AVX/FMA/AVX-512 kernels of
4446    // matrixmultiply, single-threaded — offline callers split rows across
4447    // their own threads.
4448    #[cfg(not(any(target_os = "macos", target_arch = "aarch64")))]
4449    {
4450        if m == 0 || n == 0 {
4451            return;
4452        }
4453        assert!(a.len() >= (m - 1) * lda + k && c.len() >= (m - 1) * ldc + n);
4454        let (rsb, csb) = if b_rows_are_n { (1, ldb as isize) } else { (ldb as isize, 1) };
4455        assert!(k == 0 || b_mat.len() > if b_rows_are_n { (n - 1) * ldb + k - 1 } else { (k - 1) * ldb + n - 1 });
4456        // SAFETY: the asserts above bound every element the strides reach.
4457        unsafe {
4458            matrixmultiply::sgemm(
4459                m,
4460                k,
4461                n,
4462                alpha,
4463                a.as_ptr(),
4464                lda as isize,
4465                1,
4466                b_mat.as_ptr(),
4467                rsb,
4468                csb,
4469                0.0,
4470                c.as_mut_ptr(),
4471                ldc as isize,
4472                1,
4473            );
4474        }
4475    }
4476}
4477
4478/// Row-major f32 GEMM on Accelerate: C[m,n] = alpha·A[m,k] × B(ᵀ).
4479/// `b_rows_are_n` = true multiplies by Bᵀ where B is stored [n, k].
4480#[cfg(target_os = "macos")]
4481#[allow(clippy::too_many_arguments)]
4482pub(crate) fn sgemm_rm(
4483    m: usize,
4484    n: usize,
4485    k: usize,
4486    alpha: f32,
4487    a: &[f32],
4488    lda: usize,
4489    b_mat: &[f32],
4490    ldb: usize,
4491    b_rows_are_n: bool,
4492    c: &mut [f32],
4493    ldc: usize,
4494) {
4495    debug_assert!(a.len() >= (m - 1) * lda + k);
4496    debug_assert!(c.len() >= (m - 1) * ldc + n);
4497    // Test hook: route the attention GEMMs through the portable NEON
4498    // micro-kernel ON APPLE SILICON — how the mobile batched attend is
4499    // measured without a phone in the loop. (Intel macOS has no NEON —
4500    // the hook is a no-op there, Accelerate continues below.)
4501    #[cfg(target_arch = "aarch64")]
4502    if std::env::var("CMF_FORCE_NEON_GEMM")
4503        .map(|v| v == "1")
4504        .unwrap_or(false)
4505    {
4506        return neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
4507    }
4508    unsafe {
4509        accel_blas::cblas_sgemm(
4510            101, // RowMajor
4511            111, // NoTrans A
4512            if b_rows_are_n { 112 } else { 111 },
4513            m as i32,
4514            n as i32,
4515            k as i32,
4516            alpha,
4517            a.as_ptr(),
4518            lda as i32,
4519            b_mat.as_ptr(),
4520            ldb as i32,
4521            0.0,
4522            c.as_mut_ptr(),
4523            ldc as i32,
4524        );
4525    }
4526}
4527
4528/// Prefill GEMM through Accelerate (macOS): dequantize q8 rows into
4529/// f32 tiles (scale folded in, pool-parallel) and multiply each tile
4530/// on the AMX with one row-major sgemm. Tiles live in cache, weights
4531/// stream once. Numerics are f32-GEMM (not the int8 dot): prefill
4532/// logits shift within f32 rounding — tolerance-class, like every
4533/// reduction-order change; decode (M=1) never takes this path.
4534#[cfg(target_os = "macos")]
4535fn qmatmat_accel(
4536    q: &[u8],
4537    row_scale: &[f32],
4538    pre: &[std::borrow::Cow<'_, [f32]>],
4539    rows: usize,
4540    cols: usize,
4541    out: &mut [f32],
4542    pool: Option<&Pool>,
4543) {
4544    // NOTE: double-buffering the dequant against the sgemm (a scoped
4545    // thread driving the pool on tile k+1 while the caller multiplies
4546    // tile k) was tried and LOST ~6%: Accelerate's sgemm is itself
4547    // multithreaded, and the dequant workers just steal its cores.
4548    const TR: usize = 2048;
4549    let b = pre.len();
4550    thread_local! {
4551        static XPANEL: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
4552        static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
4553    }
4554    XPANEL.with(|xp| {
4555        WTILE.with(|wt| {
4556            let mut xpanel = xp.borrow_mut();
4557            xpanel.clear();
4558            for x in pre {
4559                xpanel.extend_from_slice(x);
4560            }
4561            let mut wtile = wt.borrow_mut();
4562            wtile.resize(TR * cols, 0.0);
4563            let mut r0 = 0usize;
4564            while r0 < rows {
4565                let tr = TR.min(rows - r0);
4566                // Dequant the tile (scale folded) — pool-parallel.
4567                let wt_addr = SendMut(wtile.as_mut_ptr());
4568                let run = |start: usize, end: usize| {
4569                    for r in start..end {
4570                        let row = &q[(r0 + r) * cols..(r0 + r + 1) * cols];
4571                        let s = row_scale[r0 + r];
4572                        // SAFETY: workers cover disjoint r ranges.
4573                        let dst =
4574                            unsafe { std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols) };
4575                        for (d, &v) in dst.iter_mut().zip(row) {
4576                            *d = (v as i8) as f32 * s;
4577                        }
4578                    }
4579                };
4580                dispatch_rows(pool, tr, &run);
4581                // C[b, tr] (at column r0 of out[b, rows]) = X · Wtileᵀ
4582                unsafe {
4583                    accel_blas::cblas_sgemm(
4584                        101, // RowMajor
4585                        111, // NoTrans A
4586                        112, // Trans B
4587                        b as i32,
4588                        tr as i32,
4589                        cols as i32,
4590                        1.0,
4591                        xpanel.as_ptr(),
4592                        cols as i32,
4593                        wtile.as_ptr(),
4594                        cols as i32,
4595                        0.0,
4596                        out.as_mut_ptr().add(r0),
4597                        rows as i32,
4598                    );
4599                }
4600                r0 += tr;
4601            }
4602        })
4603    });
4604}
4605
4606fn qmatmat(
4607    q: &[u8],
4608    row_scale: &[f32],
4609    pre: &[std::borrow::Cow<'_, [f32]>],
4610    rows: usize,
4611    cols: usize,
4612    out: &mut [f32],
4613    pool: Option<&Pool>,
4614) {
4615    let b = pre.len();
4616    debug_assert_eq!(out.len(), b * rows);
4617    // Big prefill batches ride the AMX (roadmap PR3): the row×batch
4618    // SDOT loop below peaks near the CPU's dot throughput, an order
4619    // below the matrix units. Small tensors and tiny test models stay
4620    // on the exact integer path.
4621    #[cfg(target_os = "macos")]
4622    if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
4623        qmatmat_accel(q, row_scale, pre, rows, cols, out, pool);
4624        return;
4625    }
4626    #[cfg(target_arch = "aarch64")]
4627    if sdot_enabled() {
4628        let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
4629        let out_addr = SendMut(out.as_mut_ptr());
4630        // Blocked 2×4 (mobile prefill: no AMX to fall back on — this
4631        // path IS the ARM prefill GEMM off Apple silicon).
4632        let blocked_ok = blocked_enabled();
4633        let use_i8mm = i8mm_enabled();
4634        if blocked_ok {
4635            let run = |start: usize, end: usize| {
4636                let mut o = start;
4637                while o < end {
4638                    if o + 2 <= end {
4639                        let r0 = &q[o * cols..(o + 1) * cols];
4640                        let r1 = &q[(o + 1) * cols..(o + 2) * cols];
4641                        let mut bi = 0usize;
4642                        while bi + 4 <= acts.len() {
4643                            let xs = [
4644                                acts[bi].xq.as_slice(),
4645                                acts[bi + 1].xq.as_slice(),
4646                                acts[bi + 2].xq.as_slice(),
4647                                acts[bi + 3].xq.as_slice(),
4648                            ];
4649                            let d = if use_i8mm {
4650                                unsafe { dot_i8_smmla_2x4(r0, r1, xs) }
4651                            } else {
4652                                unsafe { dot_i8_sdot_2x4(r0, r1, xs) }
4653                            };
4654                            for (r, row) in [r0, r1].into_iter().enumerate() {
4655                                for k in 0..4 {
4656                                    let act = &acts[bi + k];
4657                                    let mut v = d[r][k] as f32 * act.sx;
4658                                    for &(j, xv) in &act.outliers {
4659                                        v += (row[j] as i8) as f32 * xv;
4660                                    }
4661                                    unsafe {
4662                                        *out_addr.at((bi + k) * rows + o + r) = v * row_scale[o + r]
4663                                    };
4664                                }
4665                            }
4666                            bi += 4;
4667                        }
4668                        while bi < acts.len() {
4669                            for (r, row) in [r0, r1].into_iter().enumerate() {
4670                                let v = row_dot_sdot(row, &acts[bi]) * row_scale[o + r];
4671                                unsafe { *out_addr.at(bi * rows + o + r) = v };
4672                            }
4673                            bi += 1;
4674                        }
4675                        o += 2;
4676                    } else {
4677                        let row = &q[o * cols..(o + 1) * cols];
4678                        for (bi, act) in acts.iter().enumerate() {
4679                            let v = row_dot_sdot(row, act) * row_scale[o];
4680                            unsafe { *out_addr.at(bi * rows + o) = v };
4681                        }
4682                        o += 1;
4683                    }
4684                }
4685            };
4686            dispatch_rows(pool, rows, &run);
4687            return;
4688        }
4689        let run = |start: usize, end: usize| {
4690            for o in start..end {
4691                let row = &q[o * cols..(o + 1) * cols];
4692                for (bi, act) in acts.iter().enumerate() {
4693                    let v = row_dot_sdot(row, act) * row_scale[o];
4694                    unsafe { *out_addr.at(bi * rows + o) = v };
4695                }
4696            }
4697        };
4698        dispatch_rows(pool, rows, &run);
4699        return;
4700    }
4701    // x86 A8W8 batch. Non-VNNI parts take the BLOCKED 2×4 kernel
4702    // (roadmap P0: two weight rows' abs() stay in registers across four
4703    // activation streams); VNNI machines keep the per-row bias-trick
4704    // dot, which is already throughput-bound there.
4705    #[cfg(target_arch = "x86_64")]
4706    if avx2_a8w8_enabled() {
4707        let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
4708        let out_addr = SendMut(out.as_mut_ptr());
4709        // CMF_X86_BLOCKED=0 forces the per-row path (paired in-process
4710        // A/B on noisy shared-vCPU hosts).
4711        let blocked_ok = blocked_enabled();
4712        if !avx512vnni_enabled() && blocked_ok && !row_exact() {
4713            let run = |start: usize, end: usize| {
4714                let mut o = start;
4715                while o < end {
4716                    if o + 2 <= end {
4717                        let r0 = &q[o * cols..(o + 1) * cols];
4718                        let r1 = &q[(o + 1) * cols..(o + 2) * cols];
4719                        let mut bi = 0usize;
4720                        while bi + 4 <= acts.len() {
4721                            let xs = [
4722                                acts[bi].xq.as_slice(),
4723                                acts[bi + 1].xq.as_slice(),
4724                                acts[bi + 2].xq.as_slice(),
4725                                acts[bi + 3].xq.as_slice(),
4726                            ];
4727                            let d = unsafe { dot_i8_i8_avx2_2x4(r0, r1, xs) };
4728                            for (r, row) in [r0, r1].into_iter().enumerate() {
4729                                for k in 0..4 {
4730                                    let act = &acts[bi + k];
4731                                    let mut v = d[r][k] as f32 * act.sx;
4732                                    for &(j, xv) in &act.outliers {
4733                                        v += (row[j] as i8) as f32 * xv;
4734                                    }
4735                                    unsafe {
4736                                        *out_addr.at((bi + k) * rows + o + r) = v * row_scale[o + r]
4737                                    };
4738                                }
4739                            }
4740                            bi += 4;
4741                        }
4742                        while bi < acts.len() {
4743                            for (r, row) in [r0, r1].into_iter().enumerate() {
4744                                let v = row_dot_avx2(row, &acts[bi]) * row_scale[o + r];
4745                                unsafe { *out_addr.at(bi * rows + o + r) = v };
4746                            }
4747                            bi += 1;
4748                        }
4749                        o += 2;
4750                    } else {
4751                        let row = &q[o * cols..(o + 1) * cols];
4752                        for (bi, act) in acts.iter().enumerate() {
4753                            let v = row_dot_avx2(row, act) * row_scale[o];
4754                            unsafe { *out_addr.at(bi * rows + o) = v };
4755                        }
4756                        o += 1;
4757                    }
4758                }
4759            };
4760            dispatch_rows(pool, rows, &run);
4761            return;
4762        }
4763        let run = |start: usize, end: usize| {
4764            for o in start..end {
4765                let row = &q[o * cols..(o + 1) * cols];
4766                for (bi, act) in acts.iter().enumerate() {
4767                    let v = row_dot_avx2(row, act) * row_scale[o];
4768                    unsafe { *out_addr.at(bi * rows + o) = v };
4769                }
4770            }
4771        };
4772        dispatch_rows(pool, rows, &run);
4773        return;
4774    }
4775    let out_addr = SendMut(out.as_mut_ptr());
4776    let run = |start: usize, end: usize| {
4777        for o in start..end {
4778            let row = &q[o * cols..(o + 1) * cols];
4779            for (bi, x) in pre.iter().enumerate() {
4780                let mut acc = 0f32;
4781                for j in 0..cols {
4782                    acc += (row[j] as i8) as f32 * x[j];
4783                }
4784                unsafe { *out_addr.at(bi * rows + o) = acc * row_scale[o] };
4785            }
4786        }
4787    };
4788    dispatch_rows(pool, rows, &run);
4789}
4790
4791/// Split rows across pool workers (shared qmatvec pattern). Self-balancing
4792/// — see `Pool::run_rows` for why a static 1/n split is wrong here.
4793fn dispatch_rows(pool: Option<&Pool>, rows: usize, run: &(dyn Fn(usize, usize) + Sync)) {
4794    match pool {
4795        Some(pool) if rows >= 256 => pool.run_rows(rows, run),
4796        _ => run(0, rows),
4797    }
4798}
4799
4800/// Split a q4_block blob into (packed nibbles, f16 group scales).
4801fn q4_split(bytes: &[u8], rows: usize, cols: usize) -> (&[u8], &[u8]) {
4802    let groups = rows * cols / GROUP_SIZE;
4803    bytes.split_at(groups * 16)
4804}
4805
4806/// SIMD unpack for the dominant vbit width B=4 (94% of rows on the
4807/// log2-shape calibration): 16 packed bytes -> 32 centered i8 values.
4808/// vbit packs MSB-first, so the HIGH nibble is the even element
4809/// (opposite of q4_block's lo-first interleave). Centering is u-7.
4810#[inline]
4811fn vbit_fill4(data: &[u8], buf: &mut [u8]) {
4812    #[cfg(target_arch = "aarch64")]
4813    unsafe {
4814        return vbit_fill4_neon(data, buf);
4815    }
4816    #[cfg(target_arch = "x86_64")]
4817    if avx2_enabled() {
4818        return unsafe { vbit_fill4_avx2(data, buf) };
4819    }
4820    #[allow(unreachable_code)]
4821    for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
4822        let u = unpack8::<4>(&data[blk * 4..]);
4823        for k in 0..8 {
4824            chunk[k] = (u[k] - 7) as i8 as u8;
4825        }
4826    }
4827}
4828
4829#[cfg(target_arch = "aarch64")]
4830#[target_feature(enable = "neon")]
4831unsafe fn vbit_fill4_neon(data: &[u8], buf: &mut [u8]) {
4832    // SAFETY: buf.len() is a multiple of GROUP_SIZE=32; data holds
4833    // buf.len()/2 packed bytes (validated at load).
4834    unsafe {
4835        use core::arch::aarch64::*;
4836        let n = buf.len();
4837        let mask = vdupq_n_u8(0x0F);
4838        let seven = vdupq_n_s8(7);
4839        let mut g = 0usize;
4840        while g * 32 + 32 <= n {
4841            let b = vld1q_u8(data.as_ptr().add(g * 16));
4842            let hi = vshrq_n_u8::<4>(b);
4843            let lo = vandq_u8(b, mask);
4844            let z0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(hi, lo)), seven);
4845            let z1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(hi, lo)), seven);
4846            vst1q_u8(buf.as_mut_ptr().add(g * 32), vreinterpretq_u8_s8(z0));
4847            vst1q_u8(buf.as_mut_ptr().add(g * 32 + 16), vreinterpretq_u8_s8(z1));
4848            g += 1;
4849        }
4850    }
4851}
4852
4853#[cfg(target_arch = "x86_64")]
4854#[target_feature(enable = "avx2")]
4855unsafe fn vbit_fill4_avx2(data: &[u8], buf: &mut [u8]) {
4856    // SAFETY: see vbit_fill4_neon.
4857    unsafe {
4858        use core::arch::x86_64::*;
4859        let n = buf.len();
4860        let mask = _mm_set1_epi8(0x0F);
4861        let seven = _mm256_set1_epi8(7);
4862        let mut g = 0usize;
4863        while g * 32 + 32 <= n {
4864            let b = _mm_loadu_si128(data.as_ptr().add(g * 16) as *const __m128i);
4865            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), mask);
4866            let lo = _mm_and_si128(b, mask);
4867            let z = _mm256_sub_epi8(
4868                _mm256_set_m128i(_mm_unpackhi_epi8(hi, lo), _mm_unpacklo_epi8(hi, lo)),
4869                seven,
4870            );
4871            _mm256_storeu_si256(buf.as_mut_ptr().add(g * 32) as *mut __m256i, z);
4872            g += 1;
4873        }
4874    }
4875}
4876
4877/// Unpack 8 MSB-first B-bit values from exactly B bytes (fixed shifts —
4878/// no serial bit-buffer, auto-vectorizable). Every 32-value group starts
4879/// byte-aligned (32·B/8 is integral for B∈3..8), so groups decompose
4880/// into 4 such blocks.
4881#[inline(always)]
4882fn unpack8<const B: usize>(data: &[u8]) -> [i32; 8] {
4883    let mut acc = 0u64;
4884    for i in 0..B {
4885        acc = (acc << 8) | data[i] as u64;
4886    }
4887    let mask = (1u64 << B) - 1;
4888    let mut out = [0i32; 8];
4889    for (k, o) in out.iter_mut().enumerate() {
4890        *o = ((acc >> ((7 - k) * B)) & mask) as i32;
4891    }
4892    out
4893}
4894
4895/// Fused vbit matvec straight from the mapped bytes (spec §3, P13
4896/// FIG.3): [u8 bits: rows][f16 scales: rows·cols/32][bit-packed rows,
4897/// MSB-first, byte-padded]. Row data offsets are precomputed at load
4898/// (`vbit_row_offsets`) — the per-call prefix scan was O(rows) pure
4899/// overhead on every matvec.
4900#[allow(clippy::too_many_arguments)]
4901fn vbitmatvec(
4902    bytes: &[u8],
4903    offsets: &[usize],
4904    x: &[f32],
4905    rows: usize,
4906    cols: usize,
4907    out: &mut [f32],
4908    pool: Option<&Pool>,
4909) {
4910    debug_assert_eq!(out.len(), rows);
4911    debug_assert_eq!(offsets.len(), rows + 1);
4912
4913    // SDOT path: unpack the row to centered i8 once, then per-group
4914    // int8 dot against the quantized activations — same A8W8 contract
4915    // as q8 (bounded noise; CMF_SDOT=0 keeps the exact scalar path).
4916    if a8w8_enabled() {
4917        let act = split_act(x);
4918        let out_addr = SendMut(out.as_mut_ptr());
4919        let run = move |start: usize, end: usize| {
4920            vbit_range_a8w8(bytes, offsets, x, &act, rows, cols, out_addr, start, end)
4921        };
4922        dispatch_rows(pool, rows, &run);
4923        return;
4924    }
4925
4926    let out_addr = SendMut(out.as_mut_ptr());
4927    let run = move |start: usize, end: usize| {
4928        vbit_range_f32(bytes, offsets, x, rows, cols, out_addr, start, end)
4929    };
4930    dispatch_rows(pool, rows, &run);
4931}
4932
4933/// One vbit row range via the A8W8 int8 path — kernel body of
4934/// `vbitmatvec`, extracted so multi-matrix jobs can drive it for
4935/// several tensors in one dispatch (b=8 rows go exact f32).
4936#[allow(clippy::too_many_arguments)]
4937fn vbit_range_a8w8(
4938    bytes: &[u8],
4939    offsets: &[usize],
4940    x: &[f32],
4941    act: &SplitAct,
4942    rows: usize,
4943    cols: usize,
4944    out: SendMut,
4945    start: usize,
4946    end: usize,
4947) {
4948    let ng = cols / GROUP_SIZE;
4949    let bits = &bytes[..rows];
4950    let sc_off = rows;
4951    let row_dot = |r: usize| -> f32 {
4952        let b = bits[r] as usize;
4953        let l = (1i32 << (b - 1)) - 1;
4954        let mask = (1u64 << b) - 1;
4955        let data = &bytes[offsets[r]..offsets[r + 1]];
4956        if b == 8 {
4957            // u−L reaches 128 → does not fit i8; exact f32 path.
4958            let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
4959            let mut dot = 0f32;
4960            for g in 0..ng {
4961                let so = (r * ng + g) * 2;
4962                let sgf = f16_to_f32(u16::from_le_bytes([
4963                    bytes[sc_off + so],
4964                    bytes[sc_off + so + 1],
4965                ]));
4966                let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
4967                let mut gd = 0f32;
4968                for &xv in xg.iter() {
4969                    if nbits < 8 {
4970                        acc = (acc << 8) | data[idx] as u64;
4971                        idx += 1;
4972                        nbits += 8;
4973                    }
4974                    let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
4975                    nbits -= 8;
4976                    gd += (u - l) as f32 * xv;
4977                }
4978                dot += gd * sgf;
4979            }
4980            return dot;
4981        }
4982        // Per-worker scratch: this closure runs for every row of the
4983        // tensor (lm_head ≈ 150k rows/token) — a heap allocation per
4984        // row was measurable pure overhead.
4985        thread_local! {
4986            static VBIT_SCRATCH: std::cell::RefCell<Vec<u8>> =
4987                const { std::cell::RefCell::new(Vec::new()) };
4988        }
4989        #[inline(always)]
4990        fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
4991            for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
4992                let u = unpack8::<B>(&data[blk * B..]);
4993                for k in 0..8 {
4994                    chunk[k] = (u[k] - l) as i8 as u8;
4995                }
4996            }
4997        }
4998        let _ = mask;
4999        VBIT_SCRATCH.with(|scratch| {
5000            let mut buf = scratch.borrow_mut();
5001            buf.resize(cols, 0);
5002            match b {
5003                3 => fill::<3>(data, l, &mut buf),
5004                4 => vbit_fill4(data, &mut buf),
5005                5 => fill::<5>(data, l, &mut buf),
5006                6 => fill::<6>(data, l, &mut buf),
5007                _ => unreachable!(),
5008            }
5009            let mut dot = 0f32;
5010            for g in 0..ng {
5011                let so = (r * ng + g) * 2;
5012                let s = f16_to_f32(u16::from_le_bytes([
5013                    bytes[sc_off + so],
5014                    bytes[sc_off + so + 1],
5015                ]));
5016                let d = dot_i8_i8(
5017                    &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
5018                    &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
5019                ) as f32
5020                    * act.sx;
5021                dot += d * s;
5022            }
5023            for &(j, xv) in &act.outliers {
5024                let so = (r * ng + j / GROUP_SIZE) * 2;
5025                let s = f16_to_f32(u16::from_le_bytes([
5026                    bytes[sc_off + so],
5027                    bytes[sc_off + so + 1],
5028                ]));
5029                // xq is zeroed at outlier slots — add the exact term.
5030                dot += (buf[j] as i8) as f32 * s * xv;
5031            }
5032            dot
5033        })
5034    };
5035    for r in start..end {
5036        // SAFETY: disjoint row ranges per worker.
5037        unsafe { *out.at(r) = row_dot(r) };
5038    }
5039}
5040
5041/// Exact scalar vbit row range (same extraction, non-SDOT path).
5042#[allow(clippy::too_many_arguments)]
5043fn vbit_range_f32(
5044    bytes: &[u8],
5045    offsets: &[usize],
5046    x: &[f32],
5047    rows: usize,
5048    cols: usize,
5049    out: SendMut,
5050    start: usize,
5051    end: usize,
5052) {
5053    let ng = cols / GROUP_SIZE;
5054    let bits = &bytes[..rows];
5055    let sc_off = rows;
5056    // Per-bit-width specialized inner loops: the compiler unrolls the
5057    // constant shifts (the generic bit-buffer loop was branch-bound —
5058    // 5.6 vs 13.2 tok/s q4 on the 0.8B).
5059    #[inline(always)]
5060    fn dot_row<const B: usize>(
5061        data: &[u8],
5062        bytes: &[u8],
5063        sc_off: usize,
5064        r: usize,
5065        ng: usize,
5066        x: &[f32],
5067    ) -> f32 {
5068        let l = ((1i32 << (B - 1)) - 1) as f32;
5069        let gbytes = GROUP_SIZE * B / 8;
5070        let mut dot = 0f32;
5071        for g in 0..ng {
5072            let so = (r * ng + g) * 2;
5073            let s = f16_to_f32(u16::from_le_bytes([
5074                bytes[sc_off + so],
5075                bytes[sc_off + so + 1],
5076            ]));
5077            let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
5078            let gd0 = &data[g * gbytes..(g + 1) * gbytes];
5079            let mut gd = 0f32;
5080            for blk in 0..GROUP_SIZE / 8 {
5081                let u = unpack8::<B>(&gd0[blk * B..]);
5082                let xb = &xg[blk * 8..blk * 8 + 8];
5083                for k in 0..8 {
5084                    gd += (u[k] as f32 - l) * xb[k];
5085                }
5086            }
5087            dot += gd * s;
5088        }
5089        dot
5090    }
5091    for r in start..end {
5092        let data = &bytes[offsets[r]..offsets[r + 1]];
5093        let v = match bits[r] {
5094            3 => dot_row::<3>(data, bytes, sc_off, r, ng, x),
5095            4 => dot_row::<4>(data, bytes, sc_off, r, ng, x),
5096            5 => dot_row::<5>(data, bytes, sc_off, r, ng, x),
5097            6 => dot_row::<6>(data, bytes, sc_off, r, ng, x),
5098            8 => dot_row::<8>(data, bytes, sc_off, r, ng, x),
5099            b => unreachable!("vbit bit-width {b} (validated at load)"),
5100        };
5101        // SAFETY: disjoint row ranges per worker.
5102        unsafe { *out.at(r) = v };
5103    }
5104}
5105
5106/// Fused two-input vbit matvec: each row is unpacked from the mmap ONCE
5107/// and dotted against BOTH activations (MTP verify / pair prefill used
5108/// to run two full matvecs — double weight traffic and double unpack).
5109/// Per-input math is identical to `vbitmatvec` → same accuracy contract.
5110#[allow(clippy::too_many_arguments)]
5111fn vbitmatvec2(
5112    bytes: &[u8],
5113    offsets: &[usize],
5114    x1: &[f32],
5115    x2: &[f32],
5116    rows: usize,
5117    cols: usize,
5118    o1: &mut [f32],
5119    o2: &mut [f32],
5120    pool: Option<&Pool>,
5121) {
5122    debug_assert_eq!(o1.len(), rows);
5123    debug_assert_eq!(o2.len(), rows);
5124
5125    if a8w8_enabled() {
5126        let a1 = split_act(x1);
5127        let a2 = split_act(x2);
5128        let p1 = SendMut(o1.as_mut_ptr());
5129        let p2 = SendMut(o2.as_mut_ptr());
5130        let run = move |start: usize, end: usize| {
5131            vbit_range2_a8w8(
5132                bytes, offsets, x1, x2, &a1, &a2, rows, cols, p1, p2, start, end,
5133            )
5134        };
5135        dispatch_rows(pool, rows, &run);
5136        return;
5137    }
5138
5139    let p1 = SendMut(o1.as_mut_ptr());
5140    let p2 = SendMut(o2.as_mut_ptr());
5141    let run = move |start: usize, end: usize| {
5142        vbit_range2_f32(bytes, offsets, x1, x2, rows, cols, p1, p2, start, end)
5143    };
5144    dispatch_rows(pool, rows, &run);
5145}
5146
5147/// Two-input vbit row range via the A8W8 int8 path — kernel body of
5148/// `vbitmatvec2`, extracted for pair multi-matrix jobs (b=8 rows go
5149/// exact f32 for both lanes, bits streamed once).
5150#[allow(clippy::too_many_arguments)]
5151fn vbit_range2_a8w8(
5152    bytes: &[u8],
5153    offsets: &[usize],
5154    x1: &[f32],
5155    x2: &[f32],
5156    a1: &SplitAct,
5157    a2: &SplitAct,
5158    rows: usize,
5159    cols: usize,
5160    p1: SendMut,
5161    p2: SendMut,
5162    start: usize,
5163    end: usize,
5164) {
5165    let ng = cols / GROUP_SIZE;
5166    let bits = &bytes[..rows];
5167    let sc_off = rows;
5168    let row_dots = |r: usize| -> (f32, f32) {
5169        let b = bits[r] as usize;
5170        let l = (1i32 << (b - 1)) - 1;
5171        let data = &bytes[offsets[r]..offsets[r + 1]];
5172        if b == 8 {
5173            // u−L reaches 128 → does not fit i8; exact f32 path,
5174            // bits still streamed once for both lanes.
5175            let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
5176            let (mut d1, mut d2) = (0f32, 0f32);
5177            for g in 0..ng {
5178                let so = (r * ng + g) * 2;
5179                let sgf = f16_to_f32(u16::from_le_bytes([
5180                    bytes[sc_off + so],
5181                    bytes[sc_off + so + 1],
5182                ]));
5183                let (mut g1, mut g2) = (0f32, 0f32);
5184                for k in 0..GROUP_SIZE {
5185                    if nbits < 8 {
5186                        acc = (acc << 8) | data[idx] as u64;
5187                        idx += 1;
5188                        nbits += 8;
5189                    }
5190                    let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
5191                    nbits -= 8;
5192                    let w = (u - l) as f32;
5193                    g1 += w * x1[g * GROUP_SIZE + k];
5194                    g2 += w * x2[g * GROUP_SIZE + k];
5195                }
5196                d1 += g1 * sgf;
5197                d2 += g2 * sgf;
5198            }
5199            return (d1, d2);
5200        }
5201        thread_local! {
5202            static VBIT_SCRATCH2: std::cell::RefCell<Vec<u8>> =
5203                const { std::cell::RefCell::new(Vec::new()) };
5204        }
5205        #[inline(always)]
5206        fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
5207            for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
5208                let u = unpack8::<B>(&data[blk * B..]);
5209                for k in 0..8 {
5210                    chunk[k] = (u[k] - l) as i8 as u8;
5211                }
5212            }
5213        }
5214        VBIT_SCRATCH2.with(|scratch| {
5215            let mut buf = scratch.borrow_mut();
5216            buf.resize(cols, 0);
5217            match b {
5218                3 => fill::<3>(data, l, &mut buf),
5219                4 => vbit_fill4(data, &mut buf),
5220                5 => fill::<5>(data, l, &mut buf),
5221                6 => fill::<6>(data, l, &mut buf),
5222                _ => unreachable!(),
5223            }
5224            let (mut d1, mut d2) = (0f32, 0f32);
5225            for g in 0..ng {
5226                let so = (r * ng + g) * 2;
5227                let s = f16_to_f32(u16::from_le_bytes([
5228                    bytes[sc_off + so],
5229                    bytes[sc_off + so + 1],
5230                ]));
5231                let wg = &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
5232                let v1 = dot_i8_i8(wg, &a1.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a1.sx;
5233                let v2 = dot_i8_i8(wg, &a2.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a2.sx;
5234                d1 += v1 * s;
5235                d2 += v2 * s;
5236            }
5237            for &(j, xv) in &a1.outliers {
5238                let so = (r * ng + j / GROUP_SIZE) * 2;
5239                let s = f16_to_f32(u16::from_le_bytes([
5240                    bytes[sc_off + so],
5241                    bytes[sc_off + so + 1],
5242                ]));
5243                d1 += (buf[j] as i8) as f32 * s * xv;
5244            }
5245            for &(j, xv) in &a2.outliers {
5246                let so = (r * ng + j / GROUP_SIZE) * 2;
5247                let s = f16_to_f32(u16::from_le_bytes([
5248                    bytes[sc_off + so],
5249                    bytes[sc_off + so + 1],
5250                ]));
5251                d2 += (buf[j] as i8) as f32 * s * xv;
5252            }
5253            (d1, d2)
5254        })
5255    };
5256    for r in start..end {
5257        let (v1, v2) = row_dots(r);
5258        // SAFETY: disjoint row ranges per worker.
5259        unsafe {
5260            *p1.at(r) = v1;
5261            *p2.at(r) = v2;
5262        }
5263    }
5264}
5265
5266/// Two-input exact scalar vbit row range (same extraction) —
5267/// per-bit-width specialized, two accumulators per row; per-lane
5268/// accumulation order matches `vbitmatvec` exactly.
5269#[allow(clippy::too_many_arguments)]
5270fn vbit_range2_f32(
5271    bytes: &[u8],
5272    offsets: &[usize],
5273    x1: &[f32],
5274    x2: &[f32],
5275    rows: usize,
5276    cols: usize,
5277    p1: SendMut,
5278    p2: SendMut,
5279    start: usize,
5280    end: usize,
5281) {
5282    let ng = cols / GROUP_SIZE;
5283    let bits = &bytes[..rows];
5284    let sc_off = rows;
5285    #[inline(always)]
5286    #[allow(clippy::too_many_arguments)]
5287    fn dot_row2<const B: usize>(
5288        data: &[u8],
5289        bytes: &[u8],
5290        sc_off: usize,
5291        r: usize,
5292        ng: usize,
5293        x1: &[f32],
5294        x2: &[f32],
5295    ) -> (f32, f32) {
5296        let l = ((1i32 << (B - 1)) - 1) as f32;
5297        let gbytes = GROUP_SIZE * B / 8;
5298        let (mut d1, mut d2) = (0f32, 0f32);
5299        for g in 0..ng {
5300            let so = (r * ng + g) * 2;
5301            let s = f16_to_f32(u16::from_le_bytes([
5302                bytes[sc_off + so],
5303                bytes[sc_off + so + 1],
5304            ]));
5305            let x1g = &x1[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
5306            let x2g = &x2[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
5307            let gd0 = &data[g * gbytes..(g + 1) * gbytes];
5308            let (mut g1, mut g2) = (0f32, 0f32);
5309            for blk in 0..GROUP_SIZE / 8 {
5310                let u = unpack8::<B>(&gd0[blk * B..]);
5311                for k in 0..8 {
5312                    let w = u[k] as f32 - l;
5313                    g1 += w * x1g[blk * 8 + k];
5314                    g2 += w * x2g[blk * 8 + k];
5315                }
5316            }
5317            d1 += g1 * s;
5318            d2 += g2 * s;
5319        }
5320        (d1, d2)
5321    }
5322    for r in start..end {
5323        let data = &bytes[offsets[r]..offsets[r + 1]];
5324        let (v1, v2) = match bits[r] {
5325            3 => dot_row2::<3>(data, bytes, sc_off, r, ng, x1, x2),
5326            4 => dot_row2::<4>(data, bytes, sc_off, r, ng, x1, x2),
5327            5 => dot_row2::<5>(data, bytes, sc_off, r, ng, x1, x2),
5328            6 => dot_row2::<6>(data, bytes, sc_off, r, ng, x1, x2),
5329            8 => dot_row2::<8>(data, bytes, sc_off, r, ng, x1, x2),
5330            b => unreachable!("vbit bit-width {b} (validated at load)"),
5331        };
5332        // SAFETY: disjoint row ranges per worker.
5333        unsafe {
5334            *p1.at(r) = v1;
5335            *p2.at(r) = v2;
5336        }
5337    }
5338}
5339
5340// ───────────────────── q4_tiled kernels (§4.3) ─────────────────────
5341
5342/// One q4_tiled row dot on the A8W8 int8 path: per 32-group the tile
5343/// is ONE sequential read — [f16 scale][16B nibbles] — versus the two
5344/// distant streams of the split layout. Values/order identical to the
5345/// split kernels.
5346#[inline]
5347#[allow(unreachable_code)]
5348fn dot_q4t_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
5349    #[cfg(target_arch = "aarch64")]
5350    unsafe {
5351        return dot_q4t_row_sdot(bytes, r, gpr, xq);
5352    }
5353    #[cfg(target_arch = "x86_64")]
5354    unsafe {
5355        if vnni_tiles_enabled() {
5356            return dot_q4t_row_vnni(bytes, r, gpr, xq);
5357        }
5358        return dot_q4t_row_avx2(bytes, r, gpr, xq);
5359    }
5360    let mut acc = 0f32;
5361    for gi in 0..gpr {
5362        let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
5363        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
5364        let mut d = 0i32;
5365        for (k, &b) in tile[2..].iter().enumerate() {
5366            d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
5367                + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
5368        }
5369        acc += d as f32 * s;
5370    }
5371    acc
5372}
5373
5374#[cfg(target_arch = "aarch64")]
5375#[target_feature(enable = "neon,dotprod")]
5376unsafe fn dot_q4t_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
5377    // SAFETY: callers uphold slice-length contracts (18B tile per group,
5378    // xq.len() == gpr·GROUP_SIZE).
5379    unsafe {
5380        use core::arch::aarch64::*;
5381        use core::arch::asm;
5382        let lomask = vdupq_n_u8(0x0F);
5383        let eight = vdupq_n_s8(8);
5384        let mut acc = 0f32;
5385        for gi in 0..gpr {
5386            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
5387            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5388            let b = vld1q_u8(t.add(2));
5389            let lo = vandq_u8(b, lomask);
5390            let hi = vshrq_n_u8::<4>(b);
5391            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
5392            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
5393            let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
5394            let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
5395            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5396            asm!(
5397                "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
5398                "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
5399                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5400                e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
5401                options(pure, nomem, nostack),
5402            );
5403            acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
5404        }
5405        acc
5406    }
5407}
5408
5409#[cfg(target_arch = "x86_64")]
5410#[target_feature(enable = "avx2")]
5411unsafe fn dot_q4t_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
5412    // SAFETY: see dot_q4t_row_sdot.
5413    unsafe {
5414        use core::arch::x86_64::*;
5415        let lomask = _mm_set1_epi8(0x0F);
5416        let eight = _mm256_set1_epi8(8);
5417        let ones = _mm256_set1_epi16(1);
5418        let mut acc = 0f32;
5419        for gi in 0..gpr {
5420            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
5421            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5422            let b = _mm_loadu_si128(t.add(2) as *const __m128i);
5423            let lo = _mm_and_si128(b, lomask);
5424            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
5425            let w = _mm256_sub_epi8(
5426                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5427                eight,
5428            );
5429            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5430            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
5431            let d = _mm256_madd_epi16(p16, ones);
5432            let hi128 = _mm256_extracti128_si256::<1>(d);
5433            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
5434            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
5435            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
5436            acc += _mm_cvtsi128_si32(s32) as f32 * s;
5437        }
5438        acc
5439    }
5440}
5441
5442/// VNNI twin of `dot_q4t_row_avx2`: same unpack, `vpdpbusd` replaces
5443/// the maddubs+madd pair (see `dpbusd_hsum` — sums are bit-identical).
5444/// 256-bit VL encoding, so the VEX `vpsignb` stays usable.
5445#[cfg(target_arch = "x86_64")]
5446#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
5447unsafe fn dot_q4t_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
5448    // SAFETY: see dot_q4t_row_sdot.
5449    unsafe {
5450        use core::arch::x86_64::*;
5451        let lomask = _mm_set1_epi8(0x0F);
5452        let eight = _mm256_set1_epi8(8);
5453        let mut acc = 0f32;
5454        for gi in 0..gpr {
5455            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
5456            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5457            let b = _mm_loadu_si128(t.add(2) as *const __m128i);
5458            let lo = _mm_and_si128(b, lomask);
5459            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
5460            let w = _mm256_sub_epi8(
5461                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5462                eight,
5463            );
5464            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5465            let d = dpbusd_hsum(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
5466            acc += d as f32 * s;
5467        }
5468        acc
5469    }
5470}
5471
5472/// One q4_tiled row against FOUR activation streams: the nibble unpack
5473/// and abs() happen once per group instead of once per (group,
5474/// activation) — the unpack is the dominant per-element cost of the
5475/// tiled format (roadmap P0 portable blocking, q4t leg).
5476#[cfg(target_arch = "x86_64")]
5477// `fma` is NOT implied by `avx2`: without it LLVM lowers _mm256_fmadd_ps
5478// to a libm call per lane — measured 2x slower than the reduction this
5479// kernel replaces. The runtime gate (`avx2_enabled`) already requires
5480// both features, so declaring it here is safe.
5481#[target_feature(enable = "avx2,fma")]
5482unsafe fn dot_q4t_row_1x4_avx2(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
5483    // SAFETY: callers uphold the 18B-tile and xq-length contracts.
5484    unsafe {
5485        use core::arch::x86_64::*;
5486        let lomask = _mm_set1_epi8(0x0F);
5487        let eight = _mm256_set1_epi8(8);
5488        let ones = _mm256_set1_epi16(1);
5489        // One f32 accumulator VECTOR per activation, reduced once at the
5490        // end. Folding each group's i32 lanes to a scalar inside the loop
5491        // costs an extracti128 + three shift/add + a movd — a cross-lane
5492        // dependency chain per (group, activation), 288 of them per row at
5493        // cols=2304. The per-group scale is what forces a float
5494        // accumulator; it does not force a horizontal sum.
5495        //
5496        // The four accumulators are NAMED, not an array: as `[__m256; 4]`
5497        // indexed by a loop variable LLVM keeps them in memory and every
5498        // group pays four 32-byte loads and stores. That alone made this
5499        // kernel 2x SLOWER than the per-group reduction it replaces
5500        // (measured on the EPYC box: 150 s vs 71 s for two 256² steps).
5501        let mut f0 = _mm256_setzero_ps();
5502        let mut f1 = _mm256_setzero_ps();
5503        let mut f2 = _mm256_setzero_ps();
5504        let mut f3 = _mm256_setzero_ps();
5505        for gi in 0..gpr {
5506            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
5507            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5508            let sv = _mm256_set1_ps(s);
5509            let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
5510            let lo = _mm_and_si128(bb, lomask);
5511            let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
5512            let w = _mm256_sub_epi8(
5513                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5514                eight,
5515            );
5516            let aw = _mm256_abs_epi8(w);
5517            let off = gi * GROUP_SIZE;
5518            let dot = |xq: &[i8]| {
5519                let x = _mm256_loadu_si256(xq.as_ptr().add(off) as *const __m256i);
5520                let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
5521                _mm256_cvtepi32_ps(_mm256_madd_epi16(p16, ones))
5522            };
5523            f0 = _mm256_fmadd_ps(dot(xs[0]), sv, f0);
5524            f1 = _mm256_fmadd_ps(dot(xs[1]), sv, f1);
5525            f2 = _mm256_fmadd_ps(dot(xs[2]), sv, f2);
5526            f3 = _mm256_fmadd_ps(dot(xs[3]), sv, f3);
5527        }
5528        [
5529            hsum256_ps(f0),
5530            hsum256_ps(f1),
5531            hsum256_ps(f2),
5532            hsum256_ps(f3),
5533        ]
5534    }
5535}
5536
5537/// Horizontal sum of eight f32 lanes — the one cross-lane reduction the
5538/// blocked kernels pay, once per row instead of once per group.
5539#[cfg(target_arch = "x86_64")]
5540#[target_feature(enable = "avx2")]
5541#[inline]
5542unsafe fn hsum256_ps(v: core::arch::x86_64::__m256) -> f32 {
5543    // SAFETY: pure register arithmetic on the caller's vector.
5544    unsafe {
5545        use core::arch::x86_64::*;
5546        let hi = _mm256_extractf128_ps::<1>(v);
5547        let s = _mm_add_ps(_mm256_castps256_ps128(v), hi);
5548        let s = _mm_add_ps(s, _mm_movehl_ps(s, s));
5549        let s = _mm_add_ss(s, _mm_shuffle_ps::<0x55>(s, s));
5550        _mm_cvtss_f32(s)
5551    }
5552}
5553
5554/// VNNI twin of `dot_q4t_row_1x4_avx2` (see `dpbusd_hsum`).
5555#[cfg(target_arch = "x86_64")]
5556#[target_feature(enable = "avx2,fma,avx512f,avx512bw,avx512vl,avx512vnni")]
5557unsafe fn dot_q4t_row_1x4_vnni(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
5558    // SAFETY: callers uphold the 18B-tile and xq-length contracts.
5559    unsafe {
5560        use core::arch::x86_64::*;
5561        let lomask = _mm_set1_epi8(0x0F);
5562        let eight = _mm256_set1_epi8(8);
5563        // Same shape as the AVX2 twin: accumulate in f32 vectors and pay
5564        // one cross-lane reduction per row, not per (group, activation).
5565        let mut f0 = _mm256_setzero_ps();
5566        let mut f1 = _mm256_setzero_ps();
5567        let mut f2 = _mm256_setzero_ps();
5568        let mut f3 = _mm256_setzero_ps();
5569        for gi in 0..gpr {
5570            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
5571            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5572            let sv = _mm256_set1_ps(s);
5573            let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
5574            let lo = _mm_and_si128(bb, lomask);
5575            let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
5576            let w = _mm256_sub_epi8(
5577                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5578                eight,
5579            );
5580            let aw = _mm256_abs_epi8(w);
5581            let off = gi * GROUP_SIZE;
5582            let dot = |xq: &[i8]| {
5583                let x = _mm256_loadu_si256(xq.as_ptr().add(off) as *const __m256i);
5584                _mm256_cvtepi32_ps(_mm256_dpbusd_epi32(
5585                    _mm256_setzero_si256(),
5586                    aw,
5587                    _mm256_sign_epi8(x, w),
5588                ))
5589            };
5590            f0 = _mm256_fmadd_ps(dot(xs[0]), sv, f0);
5591            f1 = _mm256_fmadd_ps(dot(xs[1]), sv, f1);
5592            f2 = _mm256_fmadd_ps(dot(xs[2]), sv, f2);
5593            f3 = _mm256_fmadd_ps(dot(xs[3]), sv, f3);
5594        }
5595        let acc = [
5596            hsum256_ps(f0),
5597            hsum256_ps(f1),
5598            hsum256_ps(f2),
5599            hsum256_ps(f3),
5600        ];
5601        acc
5602    }
5603}
5604
5605/// ARM twin of `dot_q4t_row_1x4_avx2`: one nibble unpack per group
5606/// serves FOUR activation streams. Per stream the group order and f32
5607/// accumulation match `dot_q4t_row_sdot` exactly — batch == matvec
5608/// bit-for-bit.
5609#[cfg(target_arch = "aarch64")]
5610#[target_feature(enable = "neon,dotprod")]
5611unsafe fn dot_q4t_row_1x4_sdot(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
5612    // SAFETY: callers uphold the 18B-tile and xq-length contracts.
5613    unsafe {
5614        use core::arch::aarch64::*;
5615        use core::arch::asm;
5616        let lomask = vdupq_n_u8(0x0F);
5617        let eight = vdupq_n_s8(8);
5618        let mut acc = [0f32; 4];
5619        for gi in 0..gpr {
5620            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
5621            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5622            let b = vld1q_u8(t.add(2));
5623            let lo = vandq_u8(b, lomask);
5624            let hi = vshrq_n_u8::<4>(b);
5625            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
5626            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
5627            for (k, xq) in xs.iter().enumerate() {
5628                let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
5629                let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
5630                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5631                asm!(
5632                    "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
5633                    "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
5634                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5635                    e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
5636                    options(pure, nomem, nostack),
5637                );
5638                acc[k] += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
5639            }
5640        }
5641        acc
5642    }
5643}
5644
5645/// Exact-term correction for A8W8 outliers on a tiled row.
5646#[inline]
5647fn q4t_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
5648    let gi = j / GROUP_SIZE;
5649    let k = j % GROUP_SIZE;
5650    let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
5651    let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
5652    let byte = tile[2 + k / 2];
5653    let nib = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
5654    ((nib as i32 - 8) as f32, s)
5655}
5656
5657/// Exact scalar q4_tiled row (CMF_SDOT=0 contract) — same pairwise
5658/// accumulation shape as `q4_range_f32`.
5659#[inline]
5660fn q4t_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
5661    let mut acc = 0f32;
5662    for gi in 0..gpr {
5663        let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
5664        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
5665        let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
5666        let mut ga = 0f32;
5667        for (k, &b) in tile[2..].iter().enumerate() {
5668            ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
5669                + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
5670        }
5671        acc += ga * s;
5672    }
5673    acc
5674}
5675
5676/// Split view of a `q4tp` payload. The three planes are resolved once per
5677/// matvec instead of per row — `q4tp_sections` is cheap, but doing it inside
5678/// the row loop would put a division on the hot path for nothing.
5679struct Q4tpView<'a> {
5680    nib: &'a [u8],
5681    params: &'a [u8],
5682    codes: &'a [u8],
5683    stride: usize,
5684    /// q2tp reads the ladder with rung 0 = exact zero.
5685    zero_rung: bool,
5686}
5687
5688impl<'a> Q4tpView<'a> {
5689    fn new(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
5690        let (params_off, codes_off, stride) = q4tp_sections(rows, cols);
5691        Self {
5692            nib: &bytes[..params_off],
5693            params: &bytes[params_off..codes_off],
5694            codes: &bytes[codes_off..],
5695            stride,
5696            zero_rung: false,
5697        }
5698    }
5699
5700    /// The q2tp view: identical params/codes planes, 8 B weight chunks.
5701    fn new_q2(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
5702        let (params_off, codes_off, stride) = q2tp_sections(rows, cols);
5703        Self {
5704            nib: &bytes[..params_off],
5705            params: &bytes[params_off..codes_off],
5706            codes: &bytes[codes_off..],
5707            stride,
5708            zero_rung: true,
5709        }
5710    }
5711
5712    /// Expand row `r`'s per-tile scales into `out` (length `gpr`).
5713    ///
5714    /// Doing this once per row — rather than decoding a 5-bit code inside the
5715    /// tile loop — is what makes the format free at runtime. Random access to
5716    /// a packed 5-bit field costs a division, two bounds checks and a branch;
5717    /// the tile's actual work is two `sdot`s, so per-tile decoding dominated
5718    /// the kernel and cost 5x (measured: 1.4 vs 6.9 tok/s on Nanbeige-3B).
5719    /// Walking the plane sequentially with a bit accumulator is ~3 ops.
5720    /// Eight 5-bit codes are exactly five bytes, so a whole group of
5721    /// eight decodes from one little-endian word at fixed shifts. The
5722    /// bit-accumulator this replaces carried a data-dependent `while
5723    /// have < 5` refill whose branch sat in the innermost loop of every
5724    /// q4tp row; a decode profile put this function above the dot
5725    /// products it feeds. Same bitstream, same codes — just no branch
5726    /// and eight independent extractions.
5727    #[inline]
5728    fn scales_into(&self, r: usize, gpr: usize, out: &mut [f32]) {
5729        let tab = if self.zero_rung {
5730            q2tp_ladder(self.params, r)
5731        } else {
5732            q4tp_ladder(self.params, r)
5733        };
5734        let codes = &self.codes[r * self.stride..(r + 1) * self.stride];
5735        let out = &mut out[..gpr];
5736        let mut chunks = out.chunks_exact_mut(8);
5737        let mut ci = 0usize;
5738        for c in &mut chunks {
5739            let w = u64::from(codes[ci])
5740                | u64::from(codes[ci + 1]) << 8
5741                | u64::from(codes[ci + 2]) << 16
5742                | u64::from(codes[ci + 3]) << 24
5743                | u64::from(codes[ci + 4]) << 32;
5744            for (k, o) in c.iter_mut().enumerate() {
5745                *o = tab[((w >> (5 * k)) & 31) as usize];
5746            }
5747            ci += 5;
5748        }
5749        // Fewer than eight codes left: the shared total accessor, which
5750        // tolerates a 5-bit field whose spill byte is past the stride.
5751        let tail = &codes[ci..];
5752        for (k, o) in chunks.into_remainder().iter_mut().enumerate() {
5753            *o = tab[q4tp_code(tail, k)];
5754        }
5755    }
5756}
5757
5758#[inline]
5759fn dot_q4tp_row_i8(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
5760    #[cfg(target_arch = "aarch64")]
5761    unsafe {
5762        return dot_q4tp_row_sdot(nib, r, gpr, xq, scales);
5763    }
5764    #[cfg(target_arch = "x86_64")]
5765    unsafe {
5766        if vnni_tiles_enabled() {
5767            return dot_q4tp_row_vnni(nib, r, gpr, xq, scales);
5768        }
5769        return dot_q4tp_row_avx2(nib, r, gpr, xq, scales);
5770    }
5771    #[allow(unreachable_code)]
5772    {
5773        let mut acc = 0f32;
5774        for gi in 0..gpr {
5775            let tile = &nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
5776            let s = scales[gi];
5777            let mut d = 0i32;
5778            for (k, &b) in tile.iter().enumerate() {
5779                d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
5780                    + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
5781            }
5782            acc += d as f32 * s;
5783        }
5784        acc
5785    }
5786}
5787
5788/// q4tp twin of `dot_q4t_row_sdot`: identical nibble math, but the tile
5789/// stride is 16 B (no inline scale) and the scale is a ladder lookup.
5790#[cfg(target_arch = "aarch64")]
5791#[target_feature(enable = "neon,dotprod")]
5792unsafe fn dot_q4tp_row_sdot(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
5793    // SAFETY: callers uphold slice-length contracts (16B tile per group,
5794    // xq.len() == gpr·GROUP_SIZE, codes covering gpr 5-bit fields).
5795    unsafe {
5796        use core::arch::aarch64::*;
5797        use core::arch::asm;
5798        let lomask = vdupq_n_u8(0x0F);
5799        let eight = vdupq_n_s8(8);
5800        let mut acc = 0f32;
5801        for gi in 0..gpr {
5802            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5803            let s = *scales.get_unchecked(gi);
5804            let b = vld1q_u8(t);
5805            let lo = vandq_u8(b, lomask);
5806            let hi = vshrq_n_u8::<4>(b);
5807            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
5808            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
5809            let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
5810            let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
5811            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5812            asm!(
5813                "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
5814                "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
5815                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5816                e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
5817                options(pure, nomem, nostack),
5818            );
5819            acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
5820        }
5821        acc
5822    }
5823}
5824
5825#[cfg(target_arch = "x86_64")]
5826#[target_feature(enable = "avx2")]
5827unsafe fn dot_q4tp_row_avx2(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
5828    // SAFETY: see dot_q4tp_row_sdot.
5829    unsafe {
5830        use core::arch::x86_64::*;
5831        let lomask = _mm_set1_epi8(0x0F);
5832        let eight = _mm256_set1_epi8(8);
5833        let ones = _mm256_set1_epi16(1);
5834        let mut acc = 0f32;
5835        for gi in 0..gpr {
5836            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5837            let s = *scales.get_unchecked(gi);
5838            let b = _mm_loadu_si128(t as *const __m128i);
5839            let lo = _mm_and_si128(b, lomask);
5840            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
5841            let w = _mm256_sub_epi8(
5842                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5843                eight,
5844            );
5845            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5846            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
5847            let d = _mm256_madd_epi16(p16, ones);
5848            let hi128 = _mm256_extracti128_si256::<1>(d);
5849            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
5850            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
5851            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
5852            acc += _mm_cvtsi128_si32(s32) as f32 * s;
5853        }
5854        acc
5855    }
5856}
5857
5858/// VNNI twin of `dot_q4tp_row_avx2` (see `dot_q4t_row_vnni` for why the
5859/// 256-bit VL encoding is the one to use here).
5860#[cfg(target_arch = "x86_64")]
5861#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
5862unsafe fn dot_q4tp_row_vnni(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
5863    // SAFETY: see dot_q4tp_row_sdot.
5864    unsafe {
5865        use core::arch::x86_64::*;
5866        let lomask = _mm_set1_epi8(0x0F);
5867        let eight = _mm256_set1_epi8(8);
5868        let mut acc = 0f32;
5869        for gi in 0..gpr {
5870            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5871            let s = *scales.get_unchecked(gi);
5872            let b = _mm_loadu_si128(t as *const __m128i);
5873            let lo = _mm_and_si128(b, lomask);
5874            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
5875            let w = _mm256_sub_epi8(
5876                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5877                eight,
5878            );
5879            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5880            acc += dpbusd_hsum(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w)) as f32 * s;
5881        }
5882        acc
5883    }
5884}
5885
5886/// Exact scalar q4tp row — the `CMF_SDOT=0` contract, same pairwise
5887/// accumulation shape as `q4t_row_exact`.
5888#[inline]
5889fn q4tp_row_exact(nib: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
5890    #[cfg(target_arch = "x86_64")]
5891    if avx2_enabled() {
5892        // Keep the scalar pair/group reduction order, not a vector sum.
5893        return unsafe { q4tp_row_float_avx2(nib, r, gpr, x, scales) };
5894    }
5895    q4tp_row_float_scalar(nib, r, gpr, x, scales)
5896}
5897
5898#[inline]
5899fn q4tp_row_float_scalar(nib: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
5900    let mut acc = 0f32;
5901    for gi in 0..gpr {
5902        let tile = &nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
5903        let s = scales[gi];
5904        let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
5905        let mut ga = 0f32;
5906        for (k, &b) in tile.iter().enumerate() {
5907            ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
5908                + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
5909        }
5910        acc += ga * s;
5911    }
5912    acc
5913}
5914
5915/// Vectorize unpack, conversion and multiplication, but preserve every
5916/// pair addition and the scalar accumulation order. No activation rounding
5917/// or FMA: bit-identical to the float scalar row, including its group scale.
5918#[cfg(target_arch = "x86_64")]
5919#[target_feature(enable = "avx2")]
5920unsafe fn q4tp_row_float_avx2(nib: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
5921    // SAFETY: caller checks AVX2 and provides the same complete 32-element
5922    // groups as the scalar row. Loads/stores are explicitly unaligned.
5923    unsafe {
5924        use core::arch::x86_64::*;
5925        let mask = _mm_set1_epi8(15);
5926        let eight = _mm_set1_epi8(8);
5927        let order = _mm256_setr_epi32(0, 1, 4, 5, 2, 3, 6, 7);
5928        let mut acc = 0.0f32;
5929        for gi in 0..gpr {
5930            let packed = _mm_loadu_si128(nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB).cast());
5931            let lo = _mm_and_si128(packed, mask);
5932            let hi = _mm_and_si128(_mm_srli_epi16::<4>(packed), mask);
5933            let w0 = _mm_sub_epi8(_mm_unpacklo_epi8(lo, hi), eight);
5934            let w1 = _mm_sub_epi8(_mm_unpackhi_epi8(lo, hi), eight);
5935            let xp = x.as_ptr().add(gi * GROUP_SIZE);
5936            let a = _mm256_mul_ps(
5937                _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(w0)),
5938                _mm256_loadu_ps(xp),
5939            );
5940            let b = _mm256_mul_ps(
5941                _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(w0))),
5942                _mm256_loadu_ps(xp.add(8)),
5943            );
5944            let c = _mm256_mul_ps(
5945                _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(w1)),
5946                _mm256_loadu_ps(xp.add(16)),
5947            );
5948            let d = _mm256_mul_ps(
5949                _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(w1))),
5950                _mm256_loadu_ps(xp.add(24)),
5951            );
5952            let mut pairs = [0.0f32; 16];
5953            _mm256_storeu_ps(
5954                pairs.as_mut_ptr(),
5955                _mm256_permutevar8x32_ps(_mm256_hadd_ps(a, b), order),
5956            );
5957            _mm256_storeu_ps(
5958                pairs.as_mut_ptr().add(8),
5959                _mm256_permutevar8x32_ps(_mm256_hadd_ps(c, d), order),
5960            );
5961            let mut ga = 0.0f32;
5962            for v in pairs {
5963                ga += v;
5964            }
5965            acc += ga * scales[gi];
5966        }
5967        acc
5968    }
5969}
5970
5971/// Single weight of a q4tp tensor — the a8w8 outlier path, which restores
5972/// activation outliers at full precision after the int8 pass.
5973#[inline]
5974fn q4tp_outlier(nib: &[u8], r: usize, gpr: usize, j: usize, scales: &[f32]) -> (f32, f32) {
5975    let (gi, k) = (j / GROUP_SIZE, j % GROUP_SIZE);
5976    let byte = nib[(r * gpr + gi) * Q4TP_NIB + k / 2];
5977    let n = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
5978    ((n as i32 - 8) as f32, scales[gi])
5979}
5980
5981/// Fused q4tp matvec (dispatch mirrors `q4t_matvec`).
5982fn q4tp_matvec(
5983    bytes: &[u8],
5984    x: &[f32],
5985    rows: usize,
5986    cols: usize,
5987    out: &mut [f32],
5988    pool: Option<&Pool>,
5989) {
5990    debug_assert_eq!(out.len(), rows);
5991    let gpr = cols / GROUP_SIZE;
5992    let v = Q4tpView::new(bytes, rows, cols);
5993    let out_addr = SendMut(out.as_mut_ptr());
5994    if a8w8_enabled() {
5995        let act = split_act(x);
5996        let run = |start: usize, end: usize| {
5997            // One scratch row of scales per worker — borrowed, not minted.
5998            with_krow(gpr, |sc| {
5999                for r in start..end {
6000                    v.scales_into(r, gpr, sc);
6001                    let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, sc) * act.sx;
6002                    for &(j, xv) in &act.outliers {
6003                        let (w, s) = q4tp_outlier(v.nib, r, gpr, j, sc);
6004                        acc += w * s * xv;
6005                    }
6006                    // SAFETY: disjoint row ranges per worker.
6007                    unsafe { *out_addr.at(r) = acc };
6008                }
6009            })
6010        };
6011        dispatch_rows(pool, rows, &run);
6012        return;
6013    }
6014    let run = |start: usize, end: usize| {
6015        with_krow(gpr, |sc| {
6016            for r in start..end {
6017                v.scales_into(r, gpr, sc);
6018                // SAFETY: disjoint row ranges per worker.
6019                unsafe { *out_addr.at(r) = q4tp_row_exact(v.nib, r, gpr, x, sc) };
6020            }
6021        })
6022    };
6023    dispatch_rows(pool, rows, &run);
6024}
6025
6026/// Fused two-input q4tp matvec — the SwiGLU gate/up pair. Weights and the
6027/// row ladder are read once and spent on both activation streams.
6028#[allow(clippy::too_many_arguments)]
6029fn q4tp_matvec2(
6030    bytes: &[u8],
6031    x1: &[f32],
6032    x2: &[f32],
6033    rows: usize,
6034    cols: usize,
6035    o1: &mut [f32],
6036    o2: &mut [f32],
6037    pool: Option<&Pool>,
6038) {
6039    let gpr = cols / GROUP_SIZE;
6040    let v = Q4tpView::new(bytes, rows, cols);
6041    let (p1, p2) = (SendMut(o1.as_mut_ptr()), SendMut(o2.as_mut_ptr()));
6042    let run = |start: usize, end: usize| {
6043        let mut sc = vec![0f32; gpr];
6044        for r in start..end {
6045            v.scales_into(r, gpr, &mut sc);
6046            // SAFETY: disjoint row ranges per worker.
6047            unsafe {
6048                *p1.at(r) = q4tp_row_exact(v.nib, r, gpr, x1, &sc);
6049                *p2.at(r) = q4tp_row_exact(v.nib, r, gpr, x2, &sc);
6050            }
6051        }
6052    };
6053    dispatch_rows(pool, rows, &run);
6054}
6055
6056/// One q2tp outlier weight at column `j` of row `r`: the 2-bit code and
6057/// its group scale, mirrored on `q4tp_outlier`.
6058#[inline]
6059fn q2tp_outlier(chunks: &[u8], r: usize, gpr: usize, j: usize, scales: &[f32]) -> (f32, f32) {
6060    let (gi, k) = (j / GROUP_SIZE, j % GROUP_SIZE);
6061    let byte = chunks[(r * gpr + gi) * Q2TP_CHUNK + k / 4];
6062    let c = (byte >> (2 * (k % 4))) & 3;
6063    (c as f32 - 1.5, scales[gi])
6064}
6065
6066#[cfg(target_arch = "x86_64")]
6067const Q2TP_DECODE_U32: [u32; 256] = {
6068    let mut tab = [0u32; 256];
6069    let mut b = 0usize;
6070    while b < 256 {
6071        tab[b] = ((b as u32) & 3)
6072            | ((((b as u32) >> 2) & 3) << 8)
6073            | ((((b as u32) >> 4) & 3) << 16)
6074            | ((((b as u32) >> 6) & 3) << 24);
6075        b += 1;
6076    }
6077    tab
6078};
6079
6080/// Eight packed q2tp bytes against 32 signed activation bytes. `maddubs`
6081/// exactly computes unsigned 2-bit code × signed i8; its pair sums cannot
6082/// saturate (2 × 3 × 127 < i16::MAX), and the second madd widens to i32.
6083#[cfg(target_arch = "x86_64")]
6084#[target_feature(enable = "avx2")]
6085unsafe fn q2tp_code_dot_avx2(ch: &[u8], x: &[i8]) -> i32 {
6086    use core::arch::x86_64::*;
6087    debug_assert!(ch.len() >= Q2TP_CHUNK && x.len() >= GROUP_SIZE);
6088    let codes = _mm256_setr_epi32(
6089        Q2TP_DECODE_U32[ch[0] as usize] as i32,
6090        Q2TP_DECODE_U32[ch[1] as usize] as i32,
6091        Q2TP_DECODE_U32[ch[2] as usize] as i32,
6092        Q2TP_DECODE_U32[ch[3] as usize] as i32,
6093        Q2TP_DECODE_U32[ch[4] as usize] as i32,
6094        Q2TP_DECODE_U32[ch[5] as usize] as i32,
6095        Q2TP_DECODE_U32[ch[6] as usize] as i32,
6096        Q2TP_DECODE_U32[ch[7] as usize] as i32,
6097    );
6098    let xv = unsafe { _mm256_loadu_si256(x.as_ptr().cast()) };
6099    let pair = _mm256_maddubs_epi16(codes, xv);
6100    let quad = _mm256_madd_epi16(pair, _mm256_set1_epi16(1));
6101    let sum128 = _mm_add_epi32(
6102        _mm256_castsi256_si128(quad),
6103        _mm256_extracti128_si256(quad, 1),
6104    );
6105    let sum64 = _mm_hadd_epi32(sum128, sum128);
6106    _mm_cvtsi128_si32(_mm_hadd_epi32(sum64, sum64))
6107}
6108
6109#[cfg(target_arch = "x86_64")]
6110#[inline]
6111fn q2tp_avx2_enabled() -> bool {
6112    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6113    *ON.get_or_init(|| std::arch::is_x86_feature_detected!("avx2"))
6114}
6115
6116/// Integer dot of one q2tp row against pre-quantized activations:
6117/// Σ_g s_g · (Σ c·xq − 1.5·Σ xq). The half-integer grid (c − 1.5)
6118/// becomes exact integer math through the group sums — the same trick
6119/// every a8w8 kernel in this file rides. The codes decode into a
6120/// 32-byte scratch in natural order and the dot itself is the shared
6121/// SDOT primitive; elsewhere a scalar integer loop.
6122#[inline]
6123fn dot_q2tp_row_i8(
6124    chunks: &[u8],
6125    r: usize,
6126    gpr: usize,
6127    xq: &[i8],
6128    gsum: &[i32],
6129    scales: &[f32],
6130) -> f32 {
6131    let mut acc = 0f32;
6132    let base = r * gpr * Q2TP_CHUNK;
6133    #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
6134    let mut codes = [0i8; GROUP_SIZE];
6135    #[cfg(target_arch = "x86_64")]
6136    // Cache the process-wide AVX2 decision; probing it for every 32-weight
6137    // group adds a branch to the hottest Q2TP row loop while retaining the
6138    // established table-load AVX2 arithmetic (and without coupling this path
6139    // to the separate FMA-gated A8W8 switch).
6140    let avx2 = q2tp_avx2_enabled();
6141    for gi in 0..gpr {
6142        let ch = &chunks[base + gi * Q2TP_CHUNK..base + (gi + 1) * Q2TP_CHUNK];
6143        let xg = &xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
6144        #[cfg(target_arch = "aarch64")]
6145        // NEON: the byte's four 2-bit fields land in four lane vectors
6146        // (shift+mask), vld4 de-interleaves xq to match (xj[k] =
6147        // xq[4k+j]), widening MACs accumulate exactly in i32. A scalar
6148        // decode here cost as much as the dot it fed — the profile put
6149        // it at the top of the whole W2 decode.
6150        let dot = unsafe {
6151            use core::arch::aarch64::*;
6152            let b = vld1_u8(ch.as_ptr());
6153            let three = vdup_n_u8(3);
6154            let c0 = vreinterpret_s8_u8(vand_u8(b, three));
6155            let c1 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 2), three));
6156            let c2 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 4), three));
6157            let c3 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 6), three));
6158            let x4 = vld4_s8(xg.as_ptr());
6159            let mut acc4 = vdupq_n_s32(0);
6160            acc4 = vpadalq_s16(acc4, vmull_s8(c0, x4.0));
6161            acc4 = vpadalq_s16(acc4, vmull_s8(c1, x4.1));
6162            acc4 = vpadalq_s16(acc4, vmull_s8(c2, x4.2));
6163            acc4 = vpadalq_s16(acc4, vmull_s8(c3, x4.3));
6164            vaddvq_s32(acc4)
6165        };
6166        #[cfg(target_arch = "x86_64")]
6167        let dot: i32 = if avx2 {
6168            // SAFETY: the runtime feature check gates the target-feature body;
6169            // the group slices above are exactly 8 and 32 bytes long.
6170            unsafe { q2tp_code_dot_avx2(ch, xg) }
6171        } else {
6172            ch.iter()
6173                .enumerate()
6174                .map(|(k, &b)| {
6175                    ((b & 3) as i32) * xg[k * 4] as i32
6176                        + (((b >> 2) & 3) as i32) * xg[k * 4 + 1] as i32
6177                        + (((b >> 4) & 3) as i32) * xg[k * 4 + 2] as i32
6178                        + (((b >> 6) & 3) as i32) * xg[k * 4 + 3] as i32
6179                })
6180                .sum()
6181        };
6182        #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
6183        let dot: i32 = {
6184            for (k, &b) in ch.iter().enumerate() {
6185                codes[k * 4] = (b & 3) as i8;
6186                codes[k * 4 + 1] = ((b >> 2) & 3) as i8;
6187                codes[k * 4 + 2] = ((b >> 4) & 3) as i8;
6188                codes[k * 4 + 3] = ((b >> 6) & 3) as i8;
6189            }
6190            codes
6191                .iter()
6192                .zip(xg)
6193                .map(|(&c, &x)| c as i32 * x as i32)
6194                .sum()
6195        };
6196        acc += scales[gi] * (dot as f32 - 1.5 * gsum[gi] as f32);
6197    }
6198    acc
6199}
6200
6201/// Exact f32 dot of one q2tp row: 2-bit fields LSB-first, (c − 1.5)·s.
6202/// Scalar on purpose — the 2-bit class targets the GPU graph; the CPU
6203/// path exists for parity gates and small-machine fallback.
6204fn q2tp_row_exact(chunks: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
6205    q2tp_row_exact_center(chunks, r, gpr, x, scales, 1.5)
6206}
6207
6208/// Fused Prism affine row: the derived correction is applied inside the
6209/// decoded code, avoiding a second accumulated dot and avoiding cancellation
6210/// between `B=(c-1.5)s` and `+.5s` for long 5120/17408 rows.
6211#[inline]
6212fn q2tp_affine_row_exact(chunks: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
6213    q2tp_row_exact_center(chunks, r, gpr, x, scales, 1.0)
6214}
6215
6216#[inline]
6217fn q2tp_row_exact_center(
6218    chunks: &[u8],
6219    r: usize,
6220    gpr: usize,
6221    x: &[f32],
6222    scales: &[f32],
6223    center: f32,
6224) -> f32 {
6225    let mut acc = 0f32;
6226    for gi in 0..gpr {
6227        let ch = &chunks[(r * gpr + gi) * Q2TP_CHUNK..(r * gpr + gi + 1) * Q2TP_CHUNK];
6228        let s = scales[gi];
6229        let xb = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
6230        let mut g = 0f32;
6231        for (k, &b) in ch.iter().enumerate() {
6232            g += ((b & 3) as f32 - center) * xb[k * 4]
6233                + (((b >> 2) & 3) as f32 - center) * xb[k * 4 + 1]
6234                + (((b >> 4) & 3) as f32 - center) * xb[k * 4 + 2]
6235                + (((b >> 6) & 3) as f32 - center) * xb[k * 4 + 3];
6236        }
6237        acc += s * g;
6238    }
6239    acc
6240}
6241
6242fn q2tp_matvec(
6243    bytes: &[u8],
6244    x: &[f32],
6245    rows: usize,
6246    cols: usize,
6247    out: &mut [f32],
6248    pool: Option<&Pool>,
6249) {
6250    q2tp_matvec_mode(bytes, x, rows, cols, out, pool, false);
6251}
6252
6253fn q2tp_affine_matvec(
6254    bytes: &[u8],
6255    x: &[f32],
6256    rows: usize,
6257    cols: usize,
6258    out: &mut [f32],
6259    pool: Option<&Pool>,
6260) {
6261    q2tp_matvec_mode(bytes, x, rows, cols, out, pool, true);
6262}
6263
6264fn q2tp_matvec_mode(
6265    bytes: &[u8],
6266    x: &[f32],
6267    rows: usize,
6268    cols: usize,
6269    out: &mut [f32],
6270    pool: Option<&Pool>,
6271    affine: bool,
6272) {
6273    debug_assert_eq!(out.len(), rows);
6274    let gpr = cols / GROUP_SIZE;
6275    let v = Q4tpView::new_q2(bytes, rows, cols);
6276    let out_addr = SendMut(out.as_mut_ptr());
6277    // a8w8 fast path (CMF_SDOT=0 keeps the exact scalar walk): integer
6278    // code dots + group sums, exact outlier correction — the same
6279    // contract as every sibling kernel; measured 2-bit rows were the
6280    // only scalar holdout in the family.
6281    if !affine && a8w8_enabled() {
6282        let act = split_act(x);
6283        let gsum = q1_group_sums(&act.xq, gpr);
6284        let (act, gsum) = (&act, &gsum);
6285        let run = move |start: usize, end: usize| {
6286            with_krow(gpr, |sc| {
6287                for r in start..end {
6288                    v.scales_into(r, gpr, sc);
6289                    let mut acc = dot_q2tp_row_i8(v.nib, r, gpr, &act.xq, gsum, sc) * act.sx;
6290                    for &(j, xv) in &act.outliers {
6291                        let (w, s) = q2tp_outlier(v.nib, r, gpr, j, sc);
6292                        acc += w * s * xv;
6293                    }
6294                    // SAFETY: disjoint row ranges per worker.
6295                    unsafe { *out_addr.at(r) = acc };
6296                }
6297            })
6298        };
6299        dispatch_rows(pool, rows, &run);
6300        return;
6301    }
6302    let run = |start: usize, end: usize| {
6303        with_krow(gpr, |sc| {
6304            for r in start..end {
6305                v.scales_into(r, gpr, sc);
6306                // SAFETY: disjoint row ranges per worker.
6307                unsafe {
6308                    *out_addr.at(r) = if affine {
6309                        q2tp_affine_row_exact(v.nib, r, gpr, x, sc)
6310                    } else {
6311                        q2tp_row_exact(v.nib, r, gpr, x, sc)
6312                    }
6313                };
6314            }
6315        })
6316    };
6317    dispatch_rows(pool, rows, &run);
6318}
6319
6320/// Fused two-input q2tp matvec — the SwiGLU gate/up pair.
6321#[allow(clippy::too_many_arguments)]
6322fn q2tp_matvec2(
6323    bytes: &[u8],
6324    x1: &[f32],
6325    x2: &[f32],
6326    rows: usize,
6327    cols: usize,
6328    o1: &mut [f32],
6329    o2: &mut [f32],
6330    pool: Option<&Pool>,
6331) {
6332    q2tp_matvec2_mode(bytes, x1, x2, rows, cols, o1, o2, pool, false);
6333}
6334
6335#[allow(clippy::too_many_arguments)]
6336fn q2tp_affine_matvec2(
6337    bytes: &[u8],
6338    x1: &[f32],
6339    x2: &[f32],
6340    rows: usize,
6341    cols: usize,
6342    o1: &mut [f32],
6343    o2: &mut [f32],
6344    pool: Option<&Pool>,
6345) {
6346    q2tp_matvec2_mode(bytes, x1, x2, rows, cols, o1, o2, pool, true);
6347}
6348
6349#[allow(clippy::too_many_arguments)]
6350fn q2tp_matvec2_mode(
6351    bytes: &[u8],
6352    x1: &[f32],
6353    x2: &[f32],
6354    rows: usize,
6355    cols: usize,
6356    o1: &mut [f32],
6357    o2: &mut [f32],
6358    pool: Option<&Pool>,
6359    affine: bool,
6360) {
6361    let gpr = cols / GROUP_SIZE;
6362    let v = Q4tpView::new_q2(bytes, rows, cols);
6363    let (p1, p2) = (SendMut(o1.as_mut_ptr()), SendMut(o2.as_mut_ptr()));
6364    let run = |start: usize, end: usize| {
6365        let mut sc = vec![0f32; gpr];
6366        for r in start..end {
6367            v.scales_into(r, gpr, &mut sc);
6368            // SAFETY: disjoint row ranges per worker.
6369            unsafe {
6370                *p1.at(r) = if affine {
6371                    q2tp_affine_row_exact(v.nib, r, gpr, x1, &sc)
6372                } else {
6373                    q2tp_row_exact(v.nib, r, gpr, x1, &sc)
6374                };
6375                *p2.at(r) = if affine {
6376                    q2tp_affine_row_exact(v.nib, r, gpr, x2, &sc)
6377                } else {
6378                    q2tp_row_exact(v.nib, r, gpr, x2, &sc)
6379                };
6380            }
6381        }
6382    };
6383    dispatch_rows(pool, rows, &run);
6384}
6385
6386/// Batched q2tp matmat: scalar row kernel over every batch column. CPU
6387/// prefill only — decode rides the graph, so plain and correct beats
6388/// clever here.
6389/// Test doors into the host 2-bit kernels: the stand's heap corruption
6390/// pointed at down-shaped tensors, and the private fns need a way to be
6391/// held to a reference without a model file around them.
6392pub fn q2tp_matvec_for_test(bytes: &[u8], x: &[f32], rows: usize, cols: usize, out: &mut [f32]) {
6393    // The facade IS the reference: encoder oracles hold requant output
6394    // to the exact scalar walk. The production dispatch may take the i8
6395    // fast path, whose error scale is the ACTIVATIONS' — a different
6396    // claim than the encoder correctness these tests pin.
6397    let gpr = cols / GROUP_SIZE;
6398    let v = Q4tpView::new_q2(bytes, rows, cols);
6399    with_krow(gpr, |sc| {
6400        for r in 0..rows {
6401            v.scales_into(r, gpr, sc);
6402            out[r] = q2tp_row_exact(v.nib, r, gpr, x, sc);
6403        }
6404    });
6405}
6406
6407/// Test door for the descriptor-specific fused affine decode.  Production
6408/// callers select this through a validated Prism header, never by dtype alone.
6409pub fn q2tp_affine_matvec_for_test(
6410    bytes: &[u8],
6411    x: &[f32],
6412    rows: usize,
6413    cols: usize,
6414    out: &mut [f32],
6415) {
6416    q2tp_affine_matvec(bytes, x, rows, cols, out, None);
6417}
6418
6419pub fn q2tp_matmat_for_test(
6420    bytes: &[u8],
6421    xs_all: &[f32],
6422    b: usize,
6423    rows: usize,
6424    cols: usize,
6425    out: &mut [f32],
6426) {
6427    q2tp_matmat(bytes, xs_all, b, rows, cols, out, None);
6428}
6429
6430fn q2tp_matmat(
6431    bytes: &[u8],
6432    xs_all: &[f32],
6433    b: usize,
6434    rows: usize,
6435    cols: usize,
6436    out: &mut [f32],
6437    pool: Option<&Pool>,
6438) {
6439    q2tp_matmat_mode(bytes, xs_all, b, rows, cols, out, pool, false);
6440}
6441
6442fn q2tp_affine_matmat(
6443    bytes: &[u8],
6444    xs_all: &[f32],
6445    b: usize,
6446    rows: usize,
6447    cols: usize,
6448    out: &mut [f32],
6449    pool: Option<&Pool>,
6450) {
6451    q2tp_matmat_mode(bytes, xs_all, b, rows, cols, out, pool, true);
6452}
6453
6454fn q2tp_matmat_mode(
6455    bytes: &[u8],
6456    xs_all: &[f32],
6457    b: usize,
6458    rows: usize,
6459    cols: usize,
6460    out: &mut [f32],
6461    pool: Option<&Pool>,
6462    affine: bool,
6463) {
6464    debug_assert_eq!(out.len(), b * rows);
6465    let gpr = cols / GROUP_SIZE;
6466    let v = Q4tpView::new_q2(bytes, rows, cols);
6467    let out_addr = SendMut(out.as_mut_ptr());
6468    let run = |start: usize, end: usize| {
6469        let mut sc = vec![0f32; gpr];
6470        for r in start..end {
6471            v.scales_into(r, gpr, &mut sc);
6472            for bi in 0..b {
6473                let x = &xs_all[bi * cols..(bi + 1) * cols];
6474                // SAFETY: disjoint row ranges per worker.
6475                unsafe {
6476                    *out_addr.at(bi * rows + r) = if affine {
6477                        q2tp_affine_row_exact(v.nib, r, gpr, x, &sc)
6478                    } else {
6479                        q2tp_row_exact(v.nib, r, gpr, x, &sc)
6480                    }
6481                };
6482            }
6483        }
6484    };
6485    dispatch_rows(pool, rows, &run);
6486}
6487
6488/// The pre-vectorised shape, kept for A/B (`CMF_Q4TP_V1=1`): the
6489/// horizontal add lands once per group per column instead of once per
6490/// row. Same weights, same activations — only the reduction differs.
6491/// It is also the row-exact batch kernel: per group and column it forms
6492/// `int dot as f32 * scale` and adds it to a scalar running sum, which is
6493/// `dot_q4tp_row_sdot` step for step (Rust never contracts to an fma),
6494/// so each column is bit-identical to that column's matvec.
6495#[cfg(target_arch = "aarch64")]
6496#[target_feature(enable = "neon,dotprod")]
6497unsafe fn dot_q4tp_row_1x4_sdot_v1(
6498    nib: &[u8],
6499    r: usize,
6500    gpr: usize,
6501    xs: [&[i8]; 4],
6502    scales: &[f32],
6503) -> [f32; 4] {
6504    unsafe {
6505        use core::arch::aarch64::*;
6506        use core::arch::asm;
6507        let lomask = vdupq_n_u8(0x0F);
6508        let eight = vdupq_n_s8(8);
6509        let (mut f0, mut f1, mut f2, mut f3) = (0f32, 0f32, 0f32, 0f32);
6510        for gi in 0..gpr {
6511            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
6512            let s = *scales.get_unchecked(gi);
6513            let bb = vld1q_u8(t);
6514            let lo = vandq_u8(bb, lomask);
6515            let hi = vshrq_n_u8::<4>(bb);
6516            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
6517            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
6518            let mut d = [0f32; 4];
6519            for (k, dk) in d.iter_mut().enumerate() {
6520                let x0 = vld1q_s8(xs[k].as_ptr().add(gi * GROUP_SIZE));
6521                let x1 = vld1q_s8(xs[k].as_ptr().add(gi * GROUP_SIZE + 16));
6522                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6523                asm!(
6524                    "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
6525                    "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
6526                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6527                    e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
6528                    options(pure, nomem, nostack),
6529                );
6530                *dk = vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
6531            }
6532            f0 += d[0];
6533            f1 += d[1];
6534            f2 += d[2];
6535            f3 += d[3];
6536        }
6537        [f0, f1, f2, f3]
6538    }
6539}
6540
6541/// Which q4tp batch kernel to run: 1 = the previous one, 2 = the tuned
6542/// one, 0 = decide from the CPU. An atomic rather than a `OnceLock` so a
6543/// benchmark can alternate the two inside one process, where the machine's
6544/// mood — a shared box drifts ±25% between runs — is the same for both.
6545/// What the two mean is per-architecture: on x86 the blocked AVX-512 path
6546/// against the per-column one, on ARM the two reduction shapes.
6547#[allow(dead_code)]
6548static Q4TP_ALT: std::sync::atomic::AtomicU8 = std::sync::atomic::AtomicU8::new(0);
6549
6550/// Tests that store `Q4TP_ALT` hold this, so one test's kernel pick does
6551/// not leak into another's bit-exact comparison running in parallel.
6552#[cfg(test)]
6553static Q4TP_ALT_TEST_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
6554
6555/// Blocking pays on x86 only with 512-bit VNNI. With AVX2 alone, four
6556/// columns sharing an unpack still measured slower than the per-column
6557/// path (23.2 ms against 19.4 on a 48-thread EPYC), because that path
6558/// already dequantizes the row once — so the blocked kernel bought a
6559/// second unpack-free pass at the price of half the vector width.
6560#[cfg(target_arch = "x86_64")]
6561fn q4tp_blocked_x86() -> bool {
6562    match Q4TP_ALT.load(std::sync::atomic::Ordering::Relaxed) {
6563        1 => false,
6564        // A forced ON still asks the CPU. The switch exists so a bench can
6565        // pick a kernel, not so it can promise instructions the machine
6566        // does not have — CI caught that as a SIGILL on a runner without
6567        // AVX-512, where the parity test had turned the path on by hand.
6568        2 => avx512vnni_enabled(),
6569        // Deliberately not cached back into the switch: both gates below
6570        // hold their own `OnceLock`, and latching their answer here would
6571        // make a test's override outlive the test that set it.
6572        _ => blocked_enabled() && avx512vnni_enabled(),
6573    }
6574}
6575
6576/// `CMF_Q4TP_V1=1` picks the old reduction shape (A/B only).
6577#[cfg(target_arch = "aarch64")]
6578#[allow(dead_code)]
6579fn q4tp_v1() -> bool {
6580    match Q4TP_ALT.load(std::sync::atomic::Ordering::Relaxed) {
6581        1 => true,
6582        2 => false,
6583        _ => {
6584            static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6585            *ON.get_or_init(|| std::env::var("CMF_Q4TP_V1").is_ok_and(|v| v != "0"))
6586        }
6587    }
6588}
6589
6590/// Two weight rows against eight columns. The activation load is the
6591/// same for both rows, so it is paid once for twice the arithmetic, and
6592/// sixteen accumulator chains run where eight did — which is what a kernel
6593/// retiring 0.29 instructions a cycle is short of. Register pressure is
6594/// the limit: sixteen `zmm` accumulators, two weight tiles, one
6595/// activation, of thirty-two.
6596///
6597/// Four rows by four columns spends the same sixteen accumulators the
6598/// other way and measured worse — 1488 GFLOP/s against 1644 — so the
6599/// unpack, which four rows pay twice as often, costs more than the extra
6600/// sharing of one activation load buys.
6601#[cfg(target_arch = "x86_64")]
6602#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
6603unsafe fn dot_q4tp_2x8_avx512(
6604    nib: &[u8],
6605    r0: usize,
6606    gpr: usize,
6607    xs: [&[i8]; 8],
6608    sc0: &[f32],
6609    sc1: &[f32],
6610) -> [[f32; 8]; 2] {
6611    // SAFETY: as dot_q4tp_row_1x8_avx512, two adjacent rows at once; the
6612    // caller guarantees r0 + 1 < rows and the ISA.
6613    unsafe {
6614        use core::arch::x86_64::*;
6615        let lomask = _mm256_set1_epi8(0x0F);
6616        let eight = _mm256_set1_epi8(8);
6617        let zero = _mm512_setzero_si512();
6618        let mut v0 = [_mm512_setzero_ps(); 8];
6619        let mut v1 = [_mm512_setzero_ps(); 8];
6620        let pairs = gpr / 2;
6621        let unpack = |r: usize, gi: usize| -> (__m512i, __mmask64) {
6622            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
6623            let bb = _mm256_loadu_si256(t as *const __m256i);
6624            let lo = _mm256_and_si256(bb, lomask);
6625            let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
6626            let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
6627            let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
6628            let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
6629            let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
6630            (_mm512_abs_epi8(w), _mm512_movepi8_mask(w))
6631        };
6632        for gp in 0..pairs {
6633            let gi = gp * 2;
6634            let (wa0, neg0) = unpack(r0, gi);
6635            let (wa1, neg1) = unpack(r0 + 1, gi);
6636            let off = gi * GROUP_SIZE;
6637            let sv = |sc: &[f32]| {
6638                _mm512_insertf32x8::<1>(
6639                    _mm512_castps256_ps512(_mm256_set1_ps(*sc.get_unchecked(gi))),
6640                    _mm256_set1_ps(*sc.get_unchecked(gi + 1)),
6641                )
6642            };
6643            let s0 = sv(sc0);
6644            let s1 = sv(sc1);
6645            for k in 0..8 {
6646                let xv = _mm512_loadu_si512(xs[k].as_ptr().add(off) as *const __m512i);
6647                let d0 = _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(
6648                    zero,
6649                    wa0,
6650                    _mm512_mask_sub_epi8(xv, neg0, zero, xv),
6651                ));
6652                let d1 = _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(
6653                    zero,
6654                    wa1,
6655                    _mm512_mask_sub_epi8(xv, neg1, zero, xv),
6656                ));
6657                v0[k] = _mm512_fmadd_ps(d0, s0, v0[k]);
6658                v1[k] = _mm512_fmadd_ps(d1, s1, v1[k]);
6659            }
6660        }
6661        let mut acc = [[0f32; 8]; 2];
6662        for k in 0..8 {
6663            acc[0][k] = _mm512_reduce_add_ps(v0[k]);
6664            acc[1][k] = _mm512_reduce_add_ps(v1[k]);
6665        }
6666        if gpr % 2 == 1 {
6667            let off = (gpr - 1) * GROUP_SIZE;
6668            for j in off..off + GROUP_SIZE {
6669                let (w0, sa) = q4tp_outlier(nib, r0, gpr, j, sc0);
6670                let (w1, sb) = q4tp_outlier(nib, r0 + 1, gpr, j, sc1);
6671                for k in 0..8 {
6672                    let x = *xs[k].get_unchecked(j) as f32;
6673                    acc[0][k] += w0 * sa * x;
6674                    acc[1][k] += w1 * sb * x;
6675                }
6676            }
6677        }
6678        acc
6679    }
6680}
6681
6682/// The same, eight columns at a time. One unpack then feeds twice as many
6683/// activation streams, so a wide batch reads the weight tile half as
6684/// often; the price is eight accumulators live at once. Measured 9.0 ->
6685/// 8.3 ms at 9216x2304, b=296 on a 48-thread EPYC 9B45.
6686#[cfg(target_arch = "x86_64")]
6687#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
6688unsafe fn dot_q4tp_row_1x8_avx512(
6689    nib: &[u8],
6690    r: usize,
6691    gpr: usize,
6692    xs: [&[i8]; 8],
6693    scales: &[f32],
6694) -> [f32; 8] {
6695    // SAFETY: as dot_q4tp_row_1x4_avx2; caller guarantees the ISA.
6696    unsafe {
6697        use core::arch::x86_64::*;
6698        let lomask = _mm256_set1_epi8(0x0F);
6699        let eight = _mm256_set1_epi8(8);
6700        let zero = _mm512_setzero_si512();
6701        let (mut v0, mut v1, mut v2, mut v3) = (
6702            _mm512_setzero_ps(),
6703            _mm512_setzero_ps(),
6704            _mm512_setzero_ps(),
6705            _mm512_setzero_ps(),
6706        );
6707        let (mut v4, mut v5, mut v6, mut v7) = (
6708            _mm512_setzero_ps(),
6709            _mm512_setzero_ps(),
6710            _mm512_setzero_ps(),
6711            _mm512_setzero_ps(),
6712        );
6713        let pairs = gpr / 2;
6714        for gp in 0..pairs {
6715            let gi = gp * 2;
6716            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
6717            let bb = _mm256_loadu_si256(t as *const __m256i);
6718            let lo = _mm256_and_si256(bb, lomask);
6719            let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
6720            // `unpack` works per 128-bit lane, so the halves come out as
6721            // [A.lo, B.lo] and [A.hi, B.hi]; the shuffle reorders the four
6722            // 128-bit lanes into the weights' natural order, which is what
6723            // the straight activation load expects.
6724            let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
6725            let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
6726            let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
6727            let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
6728            let wabs = _mm512_abs_epi8(w);
6729            let neg = _mm512_movepi8_mask(w);
6730            let off = gi * GROUP_SIZE;
6731            let sv = _mm512_insertf32x8::<1>(
6732                _mm512_castps256_ps512(_mm256_set1_ps(*scales.get_unchecked(gi))),
6733                _mm256_set1_ps(*scales.get_unchecked(gi + 1)),
6734            );
6735            let dot = |x: &[i8]| -> __m512 {
6736                let xv = _mm512_loadu_si512(x.as_ptr().add(off) as *const __m512i);
6737                let sx = _mm512_mask_sub_epi8(xv, neg, zero, xv);
6738                _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(zero, wabs, sx))
6739            };
6740            v0 = _mm512_fmadd_ps(dot(xs[0]), sv, v0);
6741            v1 = _mm512_fmadd_ps(dot(xs[1]), sv, v1);
6742            v2 = _mm512_fmadd_ps(dot(xs[2]), sv, v2);
6743            v3 = _mm512_fmadd_ps(dot(xs[3]), sv, v3);
6744            v4 = _mm512_fmadd_ps(dot(xs[4]), sv, v4);
6745            v5 = _mm512_fmadd_ps(dot(xs[5]), sv, v5);
6746            v6 = _mm512_fmadd_ps(dot(xs[6]), sv, v6);
6747            v7 = _mm512_fmadd_ps(dot(xs[7]), sv, v7);
6748        }
6749        let mut acc = [
6750            _mm512_reduce_add_ps(v0),
6751            _mm512_reduce_add_ps(v1),
6752            _mm512_reduce_add_ps(v2),
6753            _mm512_reduce_add_ps(v3),
6754            _mm512_reduce_add_ps(v4),
6755            _mm512_reduce_add_ps(v5),
6756            _mm512_reduce_add_ps(v6),
6757            _mm512_reduce_add_ps(v7),
6758        ];
6759        // An odd group count leaves one group over; the narrow kernel
6760        // finishes it rather than the tail being a special case here.
6761        if gpr % 2 == 1 {
6762            let off = (gpr - 1) * GROUP_SIZE;
6763            for j in off..off + GROUP_SIZE {
6764                let (w, s) = q4tp_outlier(nib, r, gpr, j, scales);
6765                let ws = w * s;
6766                for k in 0..8 {
6767                    acc[k] += ws * *xs[k].get_unchecked(j) as f32;
6768                }
6769            }
6770        }
6771        acc
6772    }
6773}
6774
6775/// The same four columns, 512 bits wide. Two groups (64 weights) ride one
6776/// unpack and one `vpdpbusd`, where AVX2 needs two unpacks and four
6777/// `maddubs`/`madd` pairs — about 2.3x fewer instructions for the same
6778/// arithmetic. The two groups carry different scales, so the fma takes a
6779/// vector whose halves hold each group's scale rather than a broadcast.
6780///
6781/// There is no 512-bit `vpsignb`, so the activation's sign is applied by
6782/// negating under a mask taken from the weight's sign bits. That mask is
6783/// per-tile, so it is hoisted out of the column loop and the per-column
6784/// cost stays exactly one instruction, as with `sign_epi8`. Weights of
6785/// zero are not zeroed by the mask trick and do not need to be: their
6786/// magnitude is zero, so the product is.
6787#[cfg(target_arch = "x86_64")]
6788#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
6789unsafe fn dot_q4tp_row_1x4_avx512(
6790    nib: &[u8],
6791    r: usize,
6792    gpr: usize,
6793    xs: [&[i8]; 4],
6794    scales: &[f32],
6795) -> [f32; 4] {
6796    // SAFETY: as dot_q4tp_row_1x4_avx2; caller guarantees the ISA.
6797    unsafe {
6798        use core::arch::x86_64::*;
6799        let lomask = _mm256_set1_epi8(0x0F);
6800        let eight = _mm256_set1_epi8(8);
6801        let zero = _mm512_setzero_si512();
6802        let (mut v0, mut v1, mut v2, mut v3) = (
6803            _mm512_setzero_ps(),
6804            _mm512_setzero_ps(),
6805            _mm512_setzero_ps(),
6806            _mm512_setzero_ps(),
6807        );
6808        let pairs = gpr / 2;
6809        for gp in 0..pairs {
6810            let gi = gp * 2;
6811            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
6812            let bb = _mm256_loadu_si256(t as *const __m256i);
6813            let lo = _mm256_and_si256(bb, lomask);
6814            let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
6815            // `unpack` works per 128-bit lane, so the halves come out as
6816            // [A.lo, B.lo] and [A.hi, B.hi]; the shuffle reorders the four
6817            // 128-bit lanes into the weights' natural order, which is what
6818            // the straight activation load expects.
6819            let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
6820            let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
6821            let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
6822            let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
6823            let wabs = _mm512_abs_epi8(w);
6824            let neg = _mm512_movepi8_mask(w);
6825            let off = gi * GROUP_SIZE;
6826            let sv = _mm512_insertf32x8::<1>(
6827                _mm512_castps256_ps512(_mm256_set1_ps(*scales.get_unchecked(gi))),
6828                _mm256_set1_ps(*scales.get_unchecked(gi + 1)),
6829            );
6830            let dot = |x: &[i8]| -> __m512 {
6831                let xv = _mm512_loadu_si512(x.as_ptr().add(off) as *const __m512i);
6832                let sx = _mm512_mask_sub_epi8(xv, neg, zero, xv);
6833                _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(zero, wabs, sx))
6834            };
6835            v0 = _mm512_fmadd_ps(dot(xs[0]), sv, v0);
6836            v1 = _mm512_fmadd_ps(dot(xs[1]), sv, v1);
6837            v2 = _mm512_fmadd_ps(dot(xs[2]), sv, v2);
6838            v3 = _mm512_fmadd_ps(dot(xs[3]), sv, v3);
6839        }
6840        let mut acc = [
6841            _mm512_reduce_add_ps(v0),
6842            _mm512_reduce_add_ps(v1),
6843            _mm512_reduce_add_ps(v2),
6844            _mm512_reduce_add_ps(v3),
6845        ];
6846        // An odd group count leaves one group over; the narrow kernel
6847        // finishes it rather than the tail being a special case here.
6848        if gpr % 2 == 1 {
6849            let off = (gpr - 1) * GROUP_SIZE;
6850            for j in off..off + GROUP_SIZE {
6851                let (w, s) = q4tp_outlier(nib, r, gpr, j, scales);
6852                let ws = w * s;
6853                for k in 0..4 {
6854                    acc[k] += ws * *xs[k].get_unchecked(j) as f32;
6855                }
6856            }
6857        }
6858        acc
6859    }
6860}
6861
6862/// Four batch columns against one q4tp row: the tile is unpacked ONCE and
6863/// spent on four activation streams, which is where a prefill batch stops
6864/// being weight-bandwidth-bound. Twin of `dot_q4t_row_1x4_sdot`.
6865#[cfg(target_arch = "aarch64")]
6866#[target_feature(enable = "neon,dotprod")]
6867unsafe fn dot_q4tp_row_1x4_sdot(
6868    nib: &[u8],
6869    r: usize,
6870    gpr: usize,
6871    xs: [&[i8]; 4],
6872    scales: &[f32],
6873) -> [f32; 4] {
6874    // SAFETY: see dot_q4tp_row_sdot; every xs[k] is gpr·GROUP_SIZE long.
6875    unsafe {
6876        use core::arch::aarch64::*;
6877        use core::arch::asm;
6878        let lomask = vdupq_n_u8(0x0F);
6879        let eight = vdupq_n_s8(8);
6880        // Named accumulators, NOT an array indexed by a loop variable: the
6881        // latter does not stay in registers (the same defect cost 2x in the
6882        // AVX2 q4t kernel and again in WGSL).
6883        //
6884        // They are VECTORS, and the horizontal add happens once at the end
6885        // instead of once per group per column. `vaddvq` is a cross-lane
6886        // reduction — with 72 groups and four columns the old shape paid
6887        // 288 of them per row, each one a dependency stall the pipeline
6888        // cannot hide, to save four float adds. The group's scale now
6889        // rides an fma into the lane accumulators, so the arithmetic per
6890        // group is one convert and one fma. Summation order changes (the
6891        // lanes carry independent partial sums), which is the same
6892        // round-off class the SDOT path already lives in — the strict
6893        // kernel (`CMF_SDOT=0`, what `cortiq ppl` runs) is unchanged and
6894        // stays the reference.
6895        let (mut v0, mut v1, mut v2, mut v3) = (
6896            vdupq_n_f32(0.0),
6897            vdupq_n_f32(0.0),
6898            vdupq_n_f32(0.0),
6899            vdupq_n_f32(0.0),
6900        );
6901        for gi in 0..gpr {
6902            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
6903            let s = *scales.get_unchecked(gi);
6904            let bb = vld1q_u8(t);
6905            let lo = vandq_u8(bb, lomask);
6906            let hi = vshrq_n_u8::<4>(bb);
6907            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
6908            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
6909            let off = gi * GROUP_SIZE;
6910            let dot4 = |x: &[i8]| -> int32x4_t {
6911                let x0 = vld1q_s8(x.as_ptr().add(off));
6912                let x1 = vld1q_s8(x.as_ptr().add(off + 16));
6913                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6914                asm!(
6915                    "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
6916                    "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
6917                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6918                    e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
6919                    options(pure, nomem, nostack),
6920                );
6921                vaddq_s32(a0, a1)
6922            };
6923            v0 = vfmaq_n_f32(v0, vcvtq_f32_s32(dot4(xs[0])), s);
6924            v1 = vfmaq_n_f32(v1, vcvtq_f32_s32(dot4(xs[1])), s);
6925            v2 = vfmaq_n_f32(v2, vcvtq_f32_s32(dot4(xs[2])), s);
6926            v3 = vfmaq_n_f32(v3, vcvtq_f32_s32(dot4(xs[3])), s);
6927        }
6928        [
6929            vaddvq_f32(v0),
6930            vaddvq_f32(v1),
6931            vaddvq_f32(v2),
6932            vaddvq_f32(v3),
6933        ]
6934    }
6935}
6936
6937/// Fused q4tp matmat — the same three arms `q4t_matmat` has. Shipping only
6938/// the scalar one made Nanbeige-3B decode at 1.2 tok/s against q4t's 5.9:
6939/// the format was fine, the missing arms were the whole regression.
6940///
6941/// Under `row_exact()` every cell is summed in `q4tp_matvec`'s order on
6942/// every architecture; the mode is read once, so one call never mixes
6943/// the two contracts when a concurrent scope opens or closes mid-call.
6944fn q4tp_matmat(
6945    bytes: &[u8],
6946    xs_all: &[f32],
6947    b: usize,
6948    rows: usize,
6949    cols: usize,
6950    out: &mut [f32],
6951    pool: Option<&Pool>,
6952) {
6953    q4tp_matmat_with(bytes, xs_all, b, rows, cols, out, pool, row_exact())
6954}
6955
6956/// `q4tp_matmat` with the row-exact mode passed in rather than read from
6957/// the shared counter: `exact` = each cell equals its token's matvec bit
6958/// for bit, otherwise the fast blocked / AMX arms are free to reorder.
6959#[allow(clippy::too_many_arguments)]
6960fn q4tp_matmat_with(
6961    bytes: &[u8],
6962    xs_all: &[f32],
6963    b: usize,
6964    rows: usize,
6965    cols: usize,
6966    out: &mut [f32],
6967    pool: Option<&Pool>,
6968    exact: bool,
6969) {
6970    debug_assert_eq!(out.len(), b * rows);
6971    let gpr = cols / GROUP_SIZE;
6972    let v = Q4tpView::new(bytes, rows, cols);
6973
6974    // Wide batches ride the AMX through a dequant-tile sgemm, as in q4t.
6975    // An f32 GEMM over dequantized weights is not the matvec's int8 sum,
6976    // so a row-exact batch never takes it.
6977    #[cfg(target_os = "macos")]
6978    if !exact && b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
6979        dequant_matmat_accel(
6980            &|r, dst| {
6981                let mut sc = [0f32; 32];
6982                let mut scv;
6983                let s: &[f32] = if gpr <= 32 {
6984                    v.scales_into(r, gpr, &mut sc);
6985                    &sc[..gpr]
6986                } else {
6987                    scv = vec![0f32; gpr];
6988                    v.scales_into(r, gpr, &mut scv);
6989                    &scv
6990                };
6991                for gi in 0..gpr {
6992                    let tile = &v.nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
6993                    for (k, &bb) in tile.iter().enumerate() {
6994                        dst[gi * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s[gi];
6995                        dst[gi * GROUP_SIZE + k * 2 + 1] =
6996                            (((bb >> 4) & 0x0F) as f32 - 8.0) * s[gi];
6997                    }
6998                }
6999            },
7000            xs_all,
7001            b,
7002            rows,
7003            cols,
7004            out,
7005            pool,
7006        );
7007        return;
7008    }
7009
7010    let out_addr = SendMut(out.as_mut_ptr());
7011    if a8w8_enabled() {
7012        let acts: Vec<SplitAct> = (0..b)
7013            .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
7014            .collect();
7015        let acts = &acts;
7016        // ARM stays blocked under `exact` too: the 1x4 kernel then runs in
7017        // its v1 shape, whose per-group `int dot as f32 * scale` and scalar
7018        // running sum are exactly `dot_q4tp_row_sdot`'s, so the tile is
7019        // still unpacked once for four columns and each column equals its
7020        // matvec. The tuned shape (fma into lane partials, one horizontal
7021        // add per row) is 1-16 ulp off the matvec and is kept for `!exact`.
7022        #[cfg(target_arch = "aarch64")]
7023        let blocked_ok = sdot_enabled() && blocked_enabled();
7024        // x86 gets the same blocking: one tile unpack spent on four
7025        // columns. Without it every column re-decoded the row, which is
7026        // why a 48-core EPYC measured a sixth of an M4's per-core rate.
7027        // The gate is `avx2_enabled`, as in q4t — `sdot_enabled` answers
7028        // for ARM's dotprod and is hard-wired false everywhere else, so
7029        // asking it here left the whole blocked path unreachable on x86.
7030        #[cfg(target_arch = "x86_64")]
7031        let blocked_ok = q4tp_blocked_x86() && !exact;
7032        #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
7033        let blocked_ok = {
7034            let _ = exact;
7035            false
7036        };
7037        // Columns are swept in panels that fit L2. Without this a
7038        // row-pair walks every activation in the batch — 4.8 MB at
7039        // 512x512 — and does it again for the next pair, so the whole
7040        // batch streams out of the shared cache once per row. Measured
7041        // 800 GB/s of it, flat across batch sizes, which is the signature
7042        // of a loop bound by traffic rather than by arithmetic. A panel of
7043        // 256 columns is 590 KB beside 221 KB of this worker's weights:
7044        // both stay resident and the batch crosses L3 once instead of
7045        // once per row.
7046        let panel_cols: usize = std::env::var("CMF_Q4TP_PANEL")
7047            .ok()
7048            .and_then(|v| v.parse().ok())
7049            .filter(|v| *v > 0)
7050            .unwrap_or(256);
7051        let run = |start: usize, end: usize| {
7052            for abase in (0..acts.len()).step_by(panel_cols) {
7053                let alen = (acts.len() - abase).min(panel_cols);
7054                let mut sc = vec![0f32; gpr];
7055                #[cfg(target_arch = "x86_64")]
7056                let mut r_lo = start;
7057                #[cfg(target_arch = "x86_64")]
7058                if blocked_ok && alen >= 8 {
7059                    let mut sc1 = vec![0f32; gpr];
7060                    while r_lo + 2 <= end {
7061                        v.scales_into(r_lo, gpr, &mut sc);
7062                        v.scales_into(r_lo + 1, gpr, &mut sc1);
7063                        let mut bi = 0usize;
7064                        while bi + 8 <= alen {
7065                            let xs = [
7066                                acts[abase + bi].xq.as_slice(),
7067                                acts[abase + bi + 1].xq.as_slice(),
7068                                acts[abase + bi + 2].xq.as_slice(),
7069                                acts[abase + bi + 3].xq.as_slice(),
7070                                acts[abase + bi + 4].xq.as_slice(),
7071                                acts[abase + bi + 5].xq.as_slice(),
7072                                acts[abase + bi + 6].xq.as_slice(),
7073                                acts[abase + bi + 7].xq.as_slice(),
7074                            ];
7075                            let d = unsafe { dot_q4tp_2x8_avx512(v.nib, r_lo, gpr, xs, &sc, &sc1) };
7076                            for (row, dr, scr) in [(r_lo, &d[0], &sc), (r_lo + 1, &d[1], &sc1)] {
7077                                for k in 0..8 {
7078                                    let act = &acts[abase + bi + k];
7079                                    let mut acc = dr[k] * act.sx;
7080                                    for &(j, xv) in &act.outliers {
7081                                        let (w, s) = q4tp_outlier(v.nib, row, gpr, j, scr);
7082                                        acc += w * s * xv;
7083                                    }
7084                                    // SAFETY: disjoint (bi, r) cells per worker.
7085                                    unsafe { *out_addr.at((abase + bi + k) * rows + row) = acc };
7086                                }
7087                            }
7088                            bi += 8;
7089                        }
7090                        // Columns past the last group of eight, both rows —
7091                        // the same single-row kernel the tail below uses.
7092                        for row in [r_lo, r_lo + 1] {
7093                            let scr: &[f32] = if row == r_lo { &sc } else { &sc1 };
7094                            for b2 in bi..alen {
7095                                let act = &acts[abase + b2];
7096                                let xs4 = [
7097                                    act.xq.as_slice(),
7098                                    act.xq.as_slice(),
7099                                    act.xq.as_slice(),
7100                                    act.xq.as_slice(),
7101                                ];
7102                                let d =
7103                                    unsafe { dot_q4tp_row_1x4_avx512(v.nib, row, gpr, xs4, scr) };
7104                                let mut acc = d[0] * act.sx;
7105                                for &(j, xv) in &act.outliers {
7106                                    let (w, s) = q4tp_outlier(v.nib, row, gpr, j, scr);
7107                                    acc += w * s * xv;
7108                                }
7109                                // SAFETY: disjoint (bi, r) cells per worker.
7110                                unsafe { *out_addr.at((abase + b2) * rows + row) = acc };
7111                            }
7112                        }
7113                        r_lo += 2;
7114                    }
7115                }
7116                #[cfg(target_arch = "x86_64")]
7117                let row_start = r_lo;
7118                #[cfg(not(target_arch = "x86_64"))]
7119                let row_start = start;
7120                for r in row_start..end {
7121                    v.scales_into(r, gpr, &mut sc);
7122                    let mut bi = 0usize;
7123                    #[cfg(target_arch = "x86_64")]
7124                    if blocked_ok {
7125                        while bi + 8 <= alen {
7126                            let xs = [
7127                                acts[abase + bi].xq.as_slice(),
7128                                acts[abase + bi + 1].xq.as_slice(),
7129                                acts[abase + bi + 2].xq.as_slice(),
7130                                acts[abase + bi + 3].xq.as_slice(),
7131                                acts[abase + bi + 4].xq.as_slice(),
7132                                acts[abase + bi + 5].xq.as_slice(),
7133                                acts[abase + bi + 6].xq.as_slice(),
7134                                acts[abase + bi + 7].xq.as_slice(),
7135                            ];
7136                            let d = unsafe { dot_q4tp_row_1x8_avx512(v.nib, r, gpr, xs, &sc) };
7137                            for k in 0..8 {
7138                                let act = &acts[abase + bi + k];
7139                                let mut acc = d[k] * act.sx;
7140                                for &(j, xv) in &act.outliers {
7141                                    let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
7142                                    acc += w * s * xv;
7143                                }
7144                                // SAFETY: disjoint (bi, r) cells per worker.
7145                                unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
7146                            }
7147                            bi += 8;
7148                        }
7149                        while bi + 4 <= alen {
7150                            let xs = [
7151                                acts[abase + bi].xq.as_slice(),
7152                                acts[abase + bi + 1].xq.as_slice(),
7153                                acts[abase + bi + 2].xq.as_slice(),
7154                                acts[abase + bi + 3].xq.as_slice(),
7155                            ];
7156                            let d = unsafe { dot_q4tp_row_1x4_avx512(v.nib, r, gpr, xs, &sc) };
7157                            for k in 0..4 {
7158                                let act = &acts[abase + bi + k];
7159                                let mut acc = d[k] * act.sx;
7160                                for &(j, xv) in &act.outliers {
7161                                    let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
7162                                    acc += w * s * xv;
7163                                }
7164                                // SAFETY: disjoint (bi, r) cells per worker.
7165                                unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
7166                            }
7167                            bi += 4;
7168                        }
7169                    }
7170                    #[cfg(target_arch = "aarch64")]
7171                    if blocked_ok {
7172                        while bi + 4 <= alen {
7173                            let xs = [
7174                                acts[abase + bi].xq.as_slice(),
7175                                acts[abase + bi + 1].xq.as_slice(),
7176                                acts[abase + bi + 2].xq.as_slice(),
7177                                acts[abase + bi + 3].xq.as_slice(),
7178                            ];
7179                            let d = unsafe {
7180                                if exact || q4tp_v1() {
7181                                    dot_q4tp_row_1x4_sdot_v1(v.nib, r, gpr, xs, &sc)
7182                                } else {
7183                                    dot_q4tp_row_1x4_sdot(v.nib, r, gpr, xs, &sc)
7184                                }
7185                            };
7186                            for k in 0..4 {
7187                                let act = &acts[abase + bi + k];
7188                                let mut acc = d[k] * act.sx;
7189                                for &(j, xv) in &act.outliers {
7190                                    let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
7191                                    acc += w * s * xv;
7192                                }
7193                                // SAFETY: disjoint (bi, r) cells per worker.
7194                                unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
7195                            }
7196                            bi += 4;
7197                        }
7198                    }
7199                    let _ = blocked_ok;
7200                    while bi < alen {
7201                        let act = &acts[abase + bi];
7202                        let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, &sc) * act.sx;
7203                        for &(j, xv) in &act.outliers {
7204                            let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
7205                            acc += w * s * xv;
7206                        }
7207                        // SAFETY: disjoint (bi, r) cells per worker range.
7208                        unsafe { *out_addr.at((abase + bi) * rows + r) = acc };
7209                        bi += 1;
7210                    }
7211                }
7212            }
7213        };
7214        dispatch_rows(pool, rows, &run);
7215        return;
7216    }
7217
7218    let run = |start: usize, end: usize| {
7219        let mut sc = vec![0f32; gpr];
7220        for r in start..end {
7221            v.scales_into(r, gpr, &mut sc);
7222            for bi in 0..b {
7223                let x = &xs_all[bi * cols..(bi + 1) * cols];
7224                // SAFETY: disjoint (bi, r) cells per worker range.
7225                unsafe { *out_addr.at(bi * rows + r) = q4tp_row_exact(v.nib, r, gpr, x, &sc) };
7226            }
7227        }
7228    };
7229    dispatch_rows(pool, rows, &run);
7230}
7231
7232/// Fused q4_tiled matvec (dispatch mirrors `q4matvec`).
7233fn q4t_matvec(
7234    bytes: &[u8],
7235    x: &[f32],
7236    rows: usize,
7237    cols: usize,
7238    out: &mut [f32],
7239    pool: Option<&Pool>,
7240) {
7241    debug_assert_eq!(out.len(), rows);
7242    let gpr = cols / GROUP_SIZE;
7243    let out_addr = SendMut(out.as_mut_ptr());
7244    if a8w8_enabled() {
7245        let act = split_act(x);
7246        let run = move |start: usize, end: usize| {
7247            for r in start..end {
7248                let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
7249                for &(j, xv) in &act.outliers {
7250                    let (w, s) = q4t_outlier(bytes, r, gpr, j);
7251                    acc += w * s * xv;
7252                }
7253                // SAFETY: disjoint row ranges per worker.
7254                unsafe { *out_addr.at(r) = acc };
7255            }
7256        };
7257        dispatch_rows(pool, rows, &run);
7258        return;
7259    }
7260    let run = move |start: usize, end: usize| {
7261        for r in start..end {
7262            // SAFETY: disjoint row ranges per worker.
7263            unsafe { *out_addr.at(r) = q4t_row_exact(bytes, r, gpr, x) };
7264        }
7265    };
7266    dispatch_rows(pool, rows, &run);
7267}
7268
7269/// Fused two-input q4_tiled matvec (weights read once per pair).
7270#[allow(clippy::too_many_arguments)]
7271fn q4t_matvec2(
7272    bytes: &[u8],
7273    x1: &[f32],
7274    x2: &[f32],
7275    rows: usize,
7276    cols: usize,
7277    o1: &mut [f32],
7278    o2: &mut [f32],
7279    pool: Option<&Pool>,
7280) {
7281    let gpr = cols / GROUP_SIZE;
7282    let p1 = SendMut(o1.as_mut_ptr());
7283    let p2 = SendMut(o2.as_mut_ptr());
7284    if a8w8_enabled() {
7285        let a1 = split_act(x1);
7286        let a2 = split_act(x2);
7287        let run = move |start: usize, end: usize| {
7288            for r in start..end {
7289                let mut v1 = dot_q4t_row_i8(bytes, r, gpr, &a1.xq) * a1.sx;
7290                let mut v2 = dot_q4t_row_i8(bytes, r, gpr, &a2.xq) * a2.sx;
7291                for &(j, xv) in &a1.outliers {
7292                    let (w, s) = q4t_outlier(bytes, r, gpr, j);
7293                    v1 += w * s * xv;
7294                }
7295                for &(j, xv) in &a2.outliers {
7296                    let (w, s) = q4t_outlier(bytes, r, gpr, j);
7297                    v2 += w * s * xv;
7298                }
7299                // SAFETY: disjoint row ranges per worker.
7300                unsafe {
7301                    *p1.at(r) = v1;
7302                    *p2.at(r) = v2;
7303                }
7304            }
7305        };
7306        dispatch_rows(pool, rows, &run);
7307        return;
7308    }
7309    let run = move |start: usize, end: usize| {
7310        for r in start..end {
7311            // SAFETY: disjoint row ranges per worker.
7312            unsafe {
7313                *p1.at(r) = q4t_row_exact(bytes, r, gpr, x1);
7314                *p2.at(r) = q4t_row_exact(bytes, r, gpr, x2);
7315            }
7316        }
7317    };
7318    dispatch_rows(pool, rows, &run);
7319}
7320
7321/// Batched q4_tiled matmat: each row's tiles stream once per microbatch.
7322#[allow(clippy::too_many_arguments)]
7323/// Prefill GEMM through Accelerate for group-quantized codecs: a
7324/// caller-supplied row dequantizer fills f32 tiles (pool-parallel) and
7325/// each tile rides the AMX with one sgemm — the generic sibling of
7326/// `qmatmat_accel` (q8). Numerics are f32-GEMM (tolerance class);
7327/// decode (b=1) never takes this path.
7328#[cfg(target_os = "macos")]
7329fn dequant_matmat_accel(
7330    dequant_row: &(dyn Fn(usize, &mut [f32]) + Sync),
7331    xs_all: &[f32],
7332    b: usize,
7333    rows: usize,
7334    cols: usize,
7335    out: &mut [f32],
7336    pool: Option<&Pool>,
7337) {
7338    const TR: usize = 2048;
7339    thread_local! {
7340        static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
7341    }
7342    WTILE.with(|wt| {
7343        let mut wtile = wt.borrow_mut();
7344        wtile.resize(TR * cols, 0.0);
7345        let mut r0 = 0usize;
7346        while r0 < rows {
7347            let tr = TR.min(rows - r0);
7348            let wt_addr = SendMut(wtile.as_mut_ptr());
7349            let run = |start: usize, end: usize| {
7350                for r in start..end {
7351                    // SAFETY: workers cover disjoint r ranges.
7352                    let dst = unsafe { std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols) };
7353                    dequant_row(r0 + r, dst);
7354                }
7355            };
7356            dispatch_rows(pool, tr, &run);
7357            unsafe {
7358                accel_blas::cblas_sgemm(
7359                    101, // RowMajor
7360                    111, // NoTrans A
7361                    112, // Trans B
7362                    b as i32,
7363                    tr as i32,
7364                    cols as i32,
7365                    1.0,
7366                    xs_all.as_ptr(),
7367                    cols as i32,
7368                    wtile.as_ptr(),
7369                    cols as i32,
7370                    0.0,
7371                    out.as_mut_ptr().add(r0),
7372                    rows as i32,
7373                );
7374            }
7375            r0 += tr;
7376        }
7377    });
7378}
7379
7380fn q4t_matmat(
7381    bytes: &[u8],
7382    xs_all: &[f32],
7383    b: usize,
7384    rows: usize,
7385    cols: usize,
7386    out: &mut [f32],
7387    pool: Option<&Pool>,
7388) {
7389    debug_assert_eq!(out.len(), b * rows);
7390    let gpr = cols / GROUP_SIZE;
7391    // Wide batches ride the AMX like q8's qmatmat: on Apple silicon
7392    // the dequant-tile sgemm is an order above the SDOT row loop for
7393    // prefill shapes (imagegen DiT forwards are exactly this).
7394    #[cfg(target_os = "macos")]
7395    if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
7396        dequant_matmat_accel(
7397            &|r, dst| {
7398                for gi in 0..gpr {
7399                    let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
7400                    let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
7401                    for (k, &bb) in tile[2..].iter().enumerate() {
7402                        dst[gi * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s;
7403                        dst[gi * GROUP_SIZE + k * 2 + 1] = (((bb >> 4) & 0x0F) as f32 - 8.0) * s;
7404                    }
7405                }
7406            },
7407            xs_all,
7408            b,
7409            rows,
7410            cols,
7411            out,
7412            pool,
7413        );
7414        return;
7415    }
7416    let out_addr = SendMut(out.as_mut_ptr());
7417    if a8w8_enabled() {
7418        let acts: Vec<SplitAct> = (0..b)
7419            .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
7420            .collect();
7421        let acts = &acts;
7422        #[cfg(target_arch = "x86_64")]
7423        let blocked_ok = avx2_enabled() && blocked_enabled();
7424        #[cfg(target_arch = "aarch64")]
7425        let blocked_ok = sdot_enabled() && blocked_enabled();
7426        #[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))]
7427        let blocked_ok = false;
7428        let run = move |start: usize, end: usize| {
7429            for r in start..end {
7430                let mut bi = 0usize;
7431                #[cfg(target_arch = "aarch64")]
7432                if blocked_ok {
7433                    while bi + 4 <= acts.len() {
7434                        let xs = [
7435                            acts[bi].xq.as_slice(),
7436                            acts[bi + 1].xq.as_slice(),
7437                            acts[bi + 2].xq.as_slice(),
7438                            acts[bi + 3].xq.as_slice(),
7439                        ];
7440                        let d = unsafe { dot_q4t_row_1x4_sdot(bytes, r, gpr, xs) };
7441                        for k in 0..4 {
7442                            let act = &acts[bi + k];
7443                            let mut acc = d[k] * act.sx;
7444                            for &(j, xv) in &act.outliers {
7445                                let (w, sc) = q4t_outlier(bytes, r, gpr, j);
7446                                acc += w * sc * xv;
7447                            }
7448                            // SAFETY: disjoint (bi, r) cells per worker.
7449                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
7450                        }
7451                        bi += 4;
7452                    }
7453                }
7454                #[cfg(target_arch = "x86_64")]
7455                if blocked_ok {
7456                    while bi + 4 <= acts.len() {
7457                        let xs = [
7458                            acts[bi].xq.as_slice(),
7459                            acts[bi + 1].xq.as_slice(),
7460                            acts[bi + 2].xq.as_slice(),
7461                            acts[bi + 3].xq.as_slice(),
7462                        ];
7463                        let d = unsafe {
7464                            if vnni_tiles_enabled() {
7465                                dot_q4t_row_1x4_vnni(bytes, r, gpr, xs)
7466                            } else {
7467                                dot_q4t_row_1x4_avx2(bytes, r, gpr, xs)
7468                            }
7469                        };
7470                        for k in 0..4 {
7471                            let act = &acts[bi + k];
7472                            let mut acc = d[k] * act.sx;
7473                            for &(j, xv) in &act.outliers {
7474                                let (w, sc) = q4t_outlier(bytes, r, gpr, j);
7475                                acc += w * sc * xv;
7476                            }
7477                            // SAFETY: disjoint (bi, r) cells per worker.
7478                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
7479                        }
7480                        bi += 4;
7481                    }
7482                }
7483                let _ = blocked_ok;
7484                while bi < acts.len() {
7485                    let act = &acts[bi];
7486                    let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
7487                    for &(j, xv) in &act.outliers {
7488                        let (w, s) = q4t_outlier(bytes, r, gpr, j);
7489                        acc += w * s * xv;
7490                    }
7491                    // SAFETY: disjoint (bi, r) cells per worker range.
7492                    unsafe { *out_addr.at(bi * rows + r) = acc };
7493                    bi += 1;
7494                }
7495            }
7496        };
7497        dispatch_rows(pool, rows, &run);
7498        return;
7499    }
7500    let run = move |start: usize, end: usize| {
7501        for r in start..end {
7502            for bi in 0..b {
7503                let x = &xs_all[bi * cols..(bi + 1) * cols];
7504                // SAFETY: disjoint (bi, r) cells per worker range.
7505                unsafe { *out_addr.at(bi * rows + r) = q4t_row_exact(bytes, r, gpr, x) };
7506            }
7507        }
7508    };
7509    dispatch_rows(pool, rows, &run);
7510}
7511
7512// ── q1 (dtype 12): binary weights, [f16 scale][4B sign bits] per
7513// 32-group tile. The kernel family mirrors q4_tiled: one sequential
7514// stream of 6-byte tiles, per-tile integer dot × scale, exact outlier
7515// correction (A8W8 contract), exact scalar path under CMF_SDOT=0. ──
7516
7517/// Per-32-group sums of the quantized activation — the ±1 identity's
7518/// shared half: `dot = −2·sdot(mask, x) − gsum[g]`, computed ONCE per
7519/// matvec and reused by every row.
7520fn q1_group_sums(xq: &[i8], gpr: usize) -> Vec<i32> {
7521    (0..gpr)
7522        .map(|gi| {
7523            xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE]
7524                .iter()
7525                .map(|&v| v as i32)
7526                .sum()
7527        })
7528        .collect()
7529}
7530
7531/// One q1 row via the A8W8 int8 path — mask-SDOT on ARM (no ±1
7532/// expansion at all), scalar bit loop elsewhere (AVX2 queued with the
7533/// x86 pass).
7534#[inline]
7535#[allow(unreachable_code)]
7536/// AVX2 q1 row via the same ±1 identity as the ARM sdot kernel: the
7537/// sign bits expand to a {0, −1} byte mask through shuffle+cmpeq, the
7538/// masked activation sums through maddubs(1, x&mask), and
7539/// `dot = −(2·masked_sum + Σx_group)` — bit-identical integer math.
7540#[cfg(target_arch = "x86_64")]
7541#[target_feature(enable = "avx2")]
7542unsafe fn dot_q1_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
7543    // SAFETY: callers uphold the 6B-tile and xq/gsum length contracts.
7544    unsafe {
7545        use core::arch::x86_64::*;
7546        // Byte j of the mask must replicate bits-byte j/8.
7547        let expand = _mm256_setr_epi8(
7548            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,
7549            3, 3, 3,
7550        );
7551        let bitsel = _mm256_setr_epi8(
7552            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
7553            -128, 1, 2, 4, 8, 16, 32, 64, -128,
7554        );
7555        let ones8 = _mm256_set1_epi8(1);
7556        let ones16 = _mm256_set1_epi16(1);
7557        let mut acc = 0f32;
7558        for gi in 0..gpr {
7559            let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
7560            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
7561            let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
7562            let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
7563            let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
7564            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7565            let sel = _mm256_and_si256(x, mask);
7566            // Σ of selected i8 lanes: maddubs(1u8, sel_i8) pairs → madd.
7567            let p16 = _mm256_maddubs_epi16(ones8, sel);
7568            let d32 = _mm256_madd_epi16(p16, ones16);
7569            let hi128 = _mm256_extracti128_si256::<1>(d32);
7570            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
7571            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7572            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7573            let msum = _mm_cvtsi128_si32(s32);
7574            // The and-select keeps x UN-negated (unlike ARM's −1-mask
7575            // sdot): d = Σ_set − Σ_unset = 2·Σ_set − Σ_all.
7576            let d = 2 * msum - gsum[gi];
7577            acc += d as f32 * s;
7578        }
7579        acc
7580    }
7581}
7582
7583/// VNNI twin of `dot_q1_row_avx2`: the masked-select sum goes through
7584/// one `vpdpbusd(1u8, sel)` (see `dpbusd_hsum` — bit-identical).
7585#[cfg(target_arch = "x86_64")]
7586#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
7587unsafe fn dot_q1_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
7588    // SAFETY: callers uphold the 6B-tile and xq/gsum length contracts.
7589    unsafe {
7590        use core::arch::x86_64::*;
7591        let expand = _mm256_setr_epi8(
7592            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,
7593            3, 3, 3,
7594        );
7595        let bitsel = _mm256_setr_epi8(
7596            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
7597            -128, 1, 2, 4, 8, 16, 32, 64, -128,
7598        );
7599        let ones8 = _mm256_set1_epi8(1);
7600        let mut acc = 0f32;
7601        for gi in 0..gpr {
7602            let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
7603            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
7604            let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
7605            let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
7606            let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
7607            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7608            let msum = dpbusd_hsum(ones8, _mm256_and_si256(x, mask));
7609            let d = 2 * msum - gsum[gi];
7610            acc += d as f32 * s;
7611        }
7612        acc
7613    }
7614}
7615
7616/// VNNI twin of `dot_q1_row_1x4_avx2` (see `dpbusd_hsum`).
7617#[cfg(target_arch = "x86_64")]
7618#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
7619unsafe fn dot_q1_row_1x4_vnni(
7620    bytes: &[u8],
7621    r: usize,
7622    gpr: usize,
7623    xs: [&[i8]; 4],
7624    gsums: [&[i32]; 4],
7625) -> [f32; 4] {
7626    // SAFETY: callers uphold the 6B-tile and xq/gsum length contracts.
7627    unsafe {
7628        use core::arch::x86_64::*;
7629        let expand = _mm256_setr_epi8(
7630            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,
7631            3, 3, 3,
7632        );
7633        let bitsel = _mm256_setr_epi8(
7634            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
7635            -128, 1, 2, 4, 8, 16, 32, 64, -128,
7636        );
7637        let ones8 = _mm256_set1_epi8(1);
7638        let mut acc = [0f32; 4];
7639        for gi in 0..gpr {
7640            let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
7641            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
7642            let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
7643            let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
7644            let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
7645            for (k, xq) in xs.iter().enumerate() {
7646                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7647                let msum = dpbusd_hsum(ones8, _mm256_and_si256(x, mask));
7648                let d = 2 * msum - gsums[k][gi];
7649                acc[k] += d as f32 * s;
7650            }
7651        }
7652        acc
7653    }
7654}
7655
7656/// The blocked 1×4 flavor: the expanded bit mask serves four activation
7657/// streams per group (mask build once, four select+reduce chains).
7658#[cfg(target_arch = "x86_64")]
7659#[target_feature(enable = "avx2")]
7660unsafe fn dot_q1_row_1x4_avx2(
7661    bytes: &[u8],
7662    r: usize,
7663    gpr: usize,
7664    xs: [&[i8]; 4],
7665    gsums: [&[i32]; 4],
7666) -> [f32; 4] {
7667    // SAFETY: callers uphold the 6B-tile and xq/gsum length contracts.
7668    unsafe {
7669        use core::arch::x86_64::*;
7670        let expand = _mm256_setr_epi8(
7671            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,
7672            3, 3, 3,
7673        );
7674        let bitsel = _mm256_setr_epi8(
7675            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
7676            -128, 1, 2, 4, 8, 16, 32, 64, -128,
7677        );
7678        let ones8 = _mm256_set1_epi8(1);
7679        let ones16 = _mm256_set1_epi16(1);
7680        let mut acc = [0f32; 4];
7681        for gi in 0..gpr {
7682            let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
7683            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
7684            let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
7685            let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
7686            let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
7687            for (k, xq) in xs.iter().enumerate() {
7688                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7689                let sel = _mm256_and_si256(x, mask);
7690                let p16 = _mm256_maddubs_epi16(ones8, sel);
7691                let d32 = _mm256_madd_epi16(p16, ones16);
7692                let hi128 = _mm256_extracti128_si256::<1>(d32);
7693                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
7694                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7695                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7696                let msum = _mm_cvtsi128_si32(s32);
7697                let d = 2 * msum - gsums[k][gi];
7698                acc[k] += d as f32 * s;
7699            }
7700        }
7701        acc
7702    }
7703}
7704
7705#[allow(unreachable_code)]
7706fn dot_q1_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
7707    #[cfg(target_arch = "aarch64")]
7708    unsafe {
7709        return dot_q1_row_sdot(bytes, r, gpr, xq, gsum);
7710    }
7711    #[cfg(target_arch = "x86_64")]
7712    if avx2_enabled() {
7713        unsafe {
7714            if vnni_tiles_enabled() {
7715                return dot_q1_row_vnni(bytes, r, gpr, xq, gsum);
7716            }
7717            return dot_q1_row_avx2(bytes, r, gpr, xq, gsum);
7718        }
7719    }
7720    let _ = gsum;
7721    let mut acc = 0f32;
7722    for gi in 0..gpr {
7723        let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
7724        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
7725        let mut d = 0i32;
7726        for (j, &b) in tile[2..].iter().enumerate() {
7727            for k in 0..8 {
7728                let w = ((b >> k) & 1) as i32 * 2 - 1;
7729                d += w * xq[gi * GROUP_SIZE + j * 8 + k] as i32;
7730            }
7731        }
7732        acc += d as f32 * s;
7733    }
7734    acc
7735}
7736
7737/// SDOT q1 row via the ±1 identity: the vtst mask (0xFF where the bit
7738/// is set, i.e. −1 as i8) feeds `sdot` DIRECTLY — no expansion to ±1
7739/// lanes at all — and `dot = −(2·sdot(mask, x) + Σx_group)`, with the
7740/// per-group activation sums shared across every row of the matvec.
7741/// Four tiles (128 weights) per iteration: integer dots reduce through
7742/// a vpaddq tree into ONE i32x4 that meets its four scales in a single
7743/// fused f32 multiply-add. Integer math throughout — bit-identical to
7744/// the scalar ±1 reference.
7745#[cfg(target_arch = "aarch64")]
7746#[target_feature(enable = "neon,dotprod")]
7747unsafe fn dot_q1_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
7748    // SAFETY: callers uphold slice-length contracts (6B tile per group,
7749    // xq.len() == gpr·GROUP_SIZE, gsum.len() == gpr).
7750    unsafe {
7751        use core::arch::aarch64::*;
7752        use core::arch::asm;
7753        const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
7754        let m = vld1q_u8(MASKS.as_ptr());
7755        // One tile's −Σ_set(x) as an UNREDUCED i32x4 (two mask-sdots).
7756        macro_rules! tile_dot {
7757            ($t:expr, $x:expr) => {{
7758                let v0 = vcombine_u8(vdup_n_u8(*$t.add(2)), vdup_n_u8(*$t.add(3)));
7759                let v1 = vcombine_u8(vdup_n_u8(*$t.add(4)), vdup_n_u8(*$t.add(5)));
7760                let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
7761                let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
7762                let x0 = vld1q_s8($x);
7763                let x1 = vld1q_s8($x.add(16));
7764                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
7765                asm!(
7766                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
7767                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
7768                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
7769                    w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
7770                    options(pure, nomem, nostack),
7771                );
7772                vaddq_s32(a0, a1)
7773            }};
7774        }
7775        // TBL unpack over PAIR loads: one vld1q covers two 6B tiles
7776        // ([s s b b b b][s s b b b b] + 4B slack), TBL replicates each
7777        // bit-byte across 8 lanes for vtst, and the four scales gather
7778        // through tbl2 into one fcvtl — the 16 ld1r broadcast loads and
7779        // 4 branchy software f16 conversions per 128 weights (the
7780        // measured load-port wall of this kernel) become 2 vector
7781        // loads + 9 table lookups. Integer math order is unchanged —
7782        // bit-identical results (FCVTL is exact on every f16).
7783        const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
7784        const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
7785        const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
7786        const IW11: [u8; 16] = [
7787            10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
7788        ];
7789        const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
7790        let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
7791        let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
7792        let isc = vld1_u8(ISC.as_ptr());
7793        // One tile's −Σ_set(x) from a TBL-unpacked pair load.
7794        macro_rules! tile_dot_tbl {
7795            ($ld:expr, $i0:expr, $i1:expr, $x:expr) => {{
7796                let w0 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i0), m));
7797                let w1 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i1), m));
7798                let x0 = vld1q_s8($x);
7799                let x1 = vld1q_s8($x.add(16));
7800                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
7801                asm!(
7802                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
7803                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
7804                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
7805                    w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
7806                    options(pure, nomem, nostack),
7807                );
7808                vaddq_s32(a0, a1)
7809            }};
7810        }
7811        let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
7812        let row_base = r * gpr * Q1_TILE;
7813        let abs_end = bytes.len();
7814        let xp = xq.as_ptr();
7815        let gp = gsum.as_ptr();
7816        let mut accv = vdupq_n_f32(0.0);
7817        let mut gi = 0;
7818        // The second pair load reads 4B past tile gi+3 — stay inside
7819        // the payload slice (only the file's final tiles fall back).
7820        while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
7821            let t0 = base.add(gi * Q1_TILE);
7822            let ld_a = vld1q_u8(t0);
7823            let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
7824            let d0 = tile_dot_tbl!(ld_a, iw00, iw01, xp.add(gi * GROUP_SIZE));
7825            let d1 = tile_dot_tbl!(ld_a, iw10, iw11, xp.add((gi + 1) * GROUP_SIZE));
7826            let d2 = tile_dot_tbl!(ld_b, iw00, iw01, xp.add((gi + 2) * GROUP_SIZE));
7827            let d3 = tile_dot_tbl!(ld_b, iw10, iw11, xp.add((gi + 3) * GROUP_SIZE));
7828            // [−Σ0, −Σ1, −Σ2, −Σ3] → dots = −(2·Σset_neg + gsum)
7829            let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
7830            let g = vld1q_s32(gp.add(gi));
7831            let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
7832            let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
7833            let scf: float32x4_t;
7834            asm!(
7835                "fcvtl {o:v}.4s, {i:v}.4h",
7836                o = out(vreg) scf, i = in(vreg) sc16,
7837                options(pure, nomem, nostack),
7838            );
7839            accv = vfmaq_f32(accv, vcvtq_f32_s32(dots), scf);
7840            gi += 4;
7841        }
7842        let mut acc = vaddvq_f32(accv);
7843        while gi < gpr {
7844            let t = base.add(gi * Q1_TILE);
7845            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
7846            let d = vaddvq_s32(tile_dot!(t, xp.add(gi * GROUP_SIZE)));
7847            acc += (-(2 * d + *gp.add(gi))) as f32 * s;
7848            gi += 1;
7849        }
7850        acc
7851    }
7852}
7853
7854/// Blocked q1 1×4: one TBL unpack of the tile pair serves FOUR
7855/// activation streams (prefill amortization — the same idea as the
7856/// AVX2 twin; per stream the group order, fma order and tail match the
7857/// single-row kernel exactly, so batch == matvec bit-for-bit).
7858#[cfg(target_arch = "aarch64")]
7859#[target_feature(enable = "neon,dotprod")]
7860unsafe fn dot_q1_row_1x4_sdot(
7861    bytes: &[u8],
7862    r: usize,
7863    gpr: usize,
7864    xs: [&[i8]; 4],
7865    gs: [&[i32]; 4],
7866) -> [f32; 4] {
7867    // SAFETY: same slice-length contracts as `dot_q1_row_sdot`, ×4.
7868    unsafe {
7869        use core::arch::aarch64::*;
7870        use core::arch::asm;
7871        const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
7872        const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
7873        const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
7874        const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
7875        const IW11: [u8; 16] = [
7876            10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
7877        ];
7878        const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
7879        let m = vld1q_u8(MASKS.as_ptr());
7880        let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
7881        let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
7882        let isc = vld1_u8(ISC.as_ptr());
7883        macro_rules! sdot2 {
7884            ($w0:expr, $w1:expr, $x:expr) => {{
7885                let x0 = vld1q_s8($x);
7886                let x1 = vld1q_s8($x.add(16));
7887                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
7888                asm!(
7889                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
7890                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
7891                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
7892                    w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
7893                    options(pure, nomem, nostack),
7894                );
7895                vaddq_s32(a0, a1)
7896            }};
7897        }
7898        let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
7899        let row_base = r * gpr * Q1_TILE;
7900        let abs_end = bytes.len();
7901        let mut accv = [vdupq_n_f32(0.0); 4];
7902        let mut gi = 0;
7903        while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
7904            let t0 = base.add(gi * Q1_TILE);
7905            let ld_a = vld1q_u8(t0);
7906            let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
7907            // Unpack ONCE — eight ±mask vectors serve all four streams.
7908            let w00 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw00), m));
7909            let w01 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw01), m));
7910            let w10 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw10), m));
7911            let w11 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw11), m));
7912            let w20 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw00), m));
7913            let w21 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw01), m));
7914            let w30 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw10), m));
7915            let w31 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw11), m));
7916            let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
7917            let scf: float32x4_t;
7918            asm!(
7919                "fcvtl {o:v}.4s, {i:v}.4h",
7920                o = out(vreg) scf, i = in(vreg) sc16,
7921                options(pure, nomem, nostack),
7922            );
7923            for k in 0..4 {
7924                let xp = xs[k].as_ptr();
7925                let d0 = sdot2!(w00, w01, xp.add(gi * GROUP_SIZE));
7926                let d1 = sdot2!(w10, w11, xp.add((gi + 1) * GROUP_SIZE));
7927                let d2 = sdot2!(w20, w21, xp.add((gi + 2) * GROUP_SIZE));
7928                let d3 = sdot2!(w30, w31, xp.add((gi + 3) * GROUP_SIZE));
7929                let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
7930                let g = vld1q_s32(gs[k].as_ptr().add(gi));
7931                let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
7932                accv[k] = vfmaq_f32(accv[k], vcvtq_f32_s32(dots), scf);
7933            }
7934            gi += 4;
7935        }
7936        let mut acc = [
7937            vaddvq_f32(accv[0]),
7938            vaddvq_f32(accv[1]),
7939            vaddvq_f32(accv[2]),
7940            vaddvq_f32(accv[3]),
7941        ];
7942        while gi < gpr {
7943            let t = base.add(gi * Q1_TILE);
7944            let sc = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
7945            let v0 = vcombine_u8(vdup_n_u8(*t.add(2)), vdup_n_u8(*t.add(3)));
7946            let v1 = vcombine_u8(vdup_n_u8(*t.add(4)), vdup_n_u8(*t.add(5)));
7947            let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
7948            let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
7949            for k in 0..4 {
7950                let d = vaddvq_s32(sdot2!(w0, w1, xs[k].as_ptr().add(gi * GROUP_SIZE)));
7951                acc[k] += (-(2 * d + *gs[k].as_ptr().add(gi))) as f32 * sc;
7952            }
7953            gi += 1;
7954        }
7955        acc
7956    }
7957}
7958
7959/// (weight ±1, scale) of one q1 element — the exact outlier term.
7960#[inline]
7961fn q1_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
7962    let gi = j / GROUP_SIZE;
7963    let k = j % GROUP_SIZE;
7964    let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
7965    let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
7966    let bit = (tile[2 + k / 8] >> (k % 8)) & 1;
7967    ((bit as i32 * 2 - 1) as f32, s)
7968}
7969
7970/// Exact scalar q1 row (CMF_SDOT=0 contract).
7971#[inline]
7972fn q1_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
7973    let mut acc = 0f32;
7974    for gi in 0..gpr {
7975        let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
7976        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
7977        let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
7978        let mut ga = 0f32;
7979        for (j, &b) in tile[2..].iter().enumerate() {
7980            for k in 0..8 {
7981                ga += (((b >> k) & 1) as f32 * 2.0 - 1.0) * xg[j * 8 + k];
7982            }
7983        }
7984        acc += ga * s;
7985    }
7986    acc
7987}
7988
7989/// One binary-weight row range via A8W8 (the body of `bin_matvec`'s hot
7990/// loop, extracted so multi-matrix jobs drive the same kernel).
7991#[allow(clippy::too_many_arguments)]
7992fn bin_range_a8w8(
7993    w: BinW,
7994    gpr: usize,
7995    act: &SplitAct,
7996    gsum: &[i32],
7997    out: SendMut,
7998    start: usize,
7999    end: usize,
8000) {
8001    for r in start..end {
8002        let mut acc = w.dot_i8(r, gpr, &act.xq, gsum) * act.sx;
8003        for &(j, xv) in &act.outliers {
8004            let (wv, s) = w.outlier(r, gpr, j);
8005            acc += wv * s * xv;
8006        }
8007        // SAFETY: disjoint row ranges per worker.
8008        unsafe { *out.at(r) = acc };
8009    }
8010}
8011
8012/// Exact-scalar binary-weight row range (CMF_SDOT=0 contract).
8013fn bin_range_f32(w: BinW, gpr: usize, x: &[f32], out: SendMut, start: usize, end: usize) {
8014    for r in start..end {
8015        // SAFETY: disjoint row ranges per worker.
8016        unsafe { *out.at(r) = w.row_exact(r, gpr, x) };
8017    }
8018}
8019
8020/// A binary (±s) weight in either of its layouts: `q1` 6-byte tiles, or
8021/// `q1g` planes with one scale per 128-group. The row drivers — matvec,
8022/// matvec2, matmat, the multi-projection job and the fused gate/up pair —
8023/// are written once against this; only the per-row kernels differ, and
8024/// the match costs one predictable branch per row. `gpr` is always the
8025/// row's count of 32-tiles (cols / 32): the activation group sums and the
8026/// tile walk are shared by both layouts.
8027#[derive(Clone, Copy)]
8028pub(crate) enum BinW<'a> {
8029    Q1(&'a [u8]),
8030    Q1g {
8031        bytes: &'a [u8],
8032        /// Byte offset of the f16 scale plane (= rows · cols/8).
8033        sc_off: usize,
8034        /// Row count: the q1 kernels treat the payload's final row
8035        /// specially (see `dot_q1g_row_sdot`), and q1g mirrors that.
8036        rows: usize,
8037    },
8038}
8039
8040impl<'a> BinW<'a> {
8041    fn q1g(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
8042        Self::Q1g {
8043            bytes,
8044            sc_off: cortiq_core::quant::q1g_scale_offset(rows, cols),
8045            rows,
8046        }
8047    }
8048
8049    #[inline]
8050    fn dot_i8(self, r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
8051        match self {
8052            Self::Q1(b) => dot_q1_row_i8(b, r, gpr, xq, gsum),
8053            Self::Q1g {
8054                bytes,
8055                sc_off,
8056                rows,
8057            } => dot_q1g_row_i8(bytes, sc_off, rows, r, gpr, xq, gsum),
8058        }
8059    }
8060
8061    #[inline]
8062    fn row_exact(self, r: usize, gpr: usize, x: &[f32]) -> f32 {
8063        match self {
8064            Self::Q1(b) => q1_row_exact(b, r, gpr, x),
8065            Self::Q1g { bytes, sc_off, .. } => q1g_row_exact(bytes, sc_off, r, gpr, x),
8066        }
8067    }
8068
8069    #[inline]
8070    fn outlier(self, r: usize, gpr: usize, j: usize) -> (f32, f32) {
8071        match self {
8072            Self::Q1(b) => q1_outlier(b, r, gpr, j),
8073            Self::Q1g { bytes, sc_off, .. } => q1g_outlier(bytes, sc_off, r, gpr, j),
8074        }
8075    }
8076
8077    /// Blocked 1×4 row (one unpack, four activation streams).
8078    ///
8079    /// # Safety
8080    /// The SDOT/dotprod target feature must be present (`sdot_enabled`).
8081    #[cfg(target_arch = "aarch64")]
8082    #[inline]
8083    unsafe fn dot_1x4_sdot(
8084        self,
8085        r: usize,
8086        gpr: usize,
8087        xs: [&[i8]; 4],
8088        gs: [&[i32]; 4],
8089    ) -> [f32; 4] {
8090        // SAFETY: forwarded caller contract.
8091        unsafe {
8092            match self {
8093                Self::Q1(b) => dot_q1_row_1x4_sdot(b, r, gpr, xs, gs),
8094                Self::Q1g {
8095                    bytes,
8096                    sc_off,
8097                    rows,
8098                } => dot_q1g_row_1x4_sdot(bytes, sc_off, rows, r, gpr, xs, gs),
8099            }
8100        }
8101    }
8102
8103    /// Blocked 1×4 row on x86 (VNNI when the tile switch is on).
8104    ///
8105    /// # Safety
8106    /// AVX2 must be present (`avx2_enabled`).
8107    #[cfg(target_arch = "x86_64")]
8108    #[inline]
8109    unsafe fn dot_1x4_x86(self, r: usize, gpr: usize, xs: [&[i8]; 4], gs: [&[i32]; 4]) -> [f32; 4] {
8110        // SAFETY: forwarded caller contract.
8111        unsafe {
8112            match (self, vnni_tiles_enabled()) {
8113                (Self::Q1(b), true) => dot_q1_row_1x4_vnni(b, r, gpr, xs, gs),
8114                (Self::Q1(b), false) => dot_q1_row_1x4_avx2(b, r, gpr, xs, gs),
8115                (Self::Q1g { bytes, sc_off, .. }, true) => {
8116                    dot_q1g_row_1x4_vnni(bytes, sc_off, r, gpr, xs, gs)
8117                }
8118                (Self::Q1g { bytes, sc_off, .. }, false) => {
8119                    dot_q1g_row_1x4_avx2(bytes, sc_off, r, gpr, xs, gs)
8120                }
8121            }
8122        }
8123    }
8124}
8125
8126// ── q1g (dtype 20): binary weights, ONE f16 scale per 128-group, in sign
8127// and scale PLANES (`[bits: rows·cols/8][f16: rows·cols/128]`). Same
8128// kernel family as q1 and — deliberately — the same float accumulation
8129// order (one `acc += d·s` per 32-tile, the group's scale reused four
8130// times), so a q1g repack of a q1 tensor whose tile scales agree gives
8131// the same f32 outputs bit for bit on every path. The win is the stream:
8132// 18 bytes per 128 weights instead of 24, the bits 16-byte aligned. ──
8133
8134/// Sign bits of q1g tile `t` (flat 32-tile index `r·gpr + gi`).
8135#[inline]
8136fn q1g_bits(bytes: &[u8], t: usize) -> u32 {
8137    u32::from_le_bytes([
8138        bytes[t * 4],
8139        bytes[t * 4 + 1],
8140        bytes[t * 4 + 2],
8141        bytes[t * 4 + 3],
8142    ])
8143}
8144
8145/// Raw f16 bits of the scale of the 128-group holding q1g tile `t`.
8146#[inline]
8147fn q1g_scale_bits(bytes: &[u8], sc_off: usize, t: usize) -> u16 {
8148    let o = sc_off + (t / 4) * 2;
8149    u16::from_le_bytes([bytes[o], bytes[o + 1]])
8150}
8151
8152/// One q1g row via the A8W8 int8 path (dispatch twin of `dot_q1_row_i8`).
8153#[allow(unreachable_code)]
8154#[allow(clippy::too_many_arguments)]
8155fn dot_q1g_row_i8(
8156    bytes: &[u8],
8157    sc_off: usize,
8158    rows: usize,
8159    r: usize,
8160    gpr: usize,
8161    xq: &[i8],
8162    gsum: &[i32],
8163) -> f32 {
8164    #[cfg(target_arch = "aarch64")]
8165    unsafe {
8166        return dot_q1g_row_sdot(bytes, sc_off, rows, r, gpr, xq, gsum);
8167    }
8168    #[cfg(target_arch = "x86_64")]
8169    if avx2_enabled() {
8170        unsafe {
8171            if vnni_tiles_enabled() {
8172                return dot_q1g_row_vnni(bytes, sc_off, r, gpr, xq, gsum);
8173            }
8174            return dot_q1g_row_avx2(bytes, sc_off, r, gpr, xq, gsum);
8175        }
8176    }
8177    let _ = (gsum, rows);
8178    let mut acc = 0f32;
8179    for gi in 0..gpr {
8180        let t = r * gpr + gi;
8181        let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, t));
8182        let bits = q1g_bits(bytes, t);
8183        let mut d = 0i32;
8184        for k in 0..GROUP_SIZE {
8185            let w = ((bits >> k) & 1) as i32 * 2 - 1;
8186            d += w * xq[gi * GROUP_SIZE + k] as i32;
8187        }
8188        acc += d as f32 * s;
8189    }
8190    acc
8191}
8192
8193/// Four tiles' unreduced `i32×8` partial sums → `[Σ0, Σ1, Σ2, Σ3]` in one
8194/// hadd tree: 3 hadds + 1 add for four tiles where the per-tile reduce of
8195/// the q1 kernels costs 7 ops each. Integer math, so any order gives the
8196/// q1 kernel's per-tile sums exactly.
8197#[cfg(target_arch = "x86_64")]
8198#[target_feature(enable = "avx2")]
8199#[inline]
8200unsafe fn q1g_hsum4(
8201    d0: core::arch::x86_64::__m256i,
8202    d1: core::arch::x86_64::__m256i,
8203    d2: core::arch::x86_64::__m256i,
8204    d3: core::arch::x86_64::__m256i,
8205) -> core::arch::x86_64::__m128i {
8206    // SAFETY: pure register math.
8207    unsafe {
8208        use core::arch::x86_64::*;
8209        let h = _mm256_hadd_epi32(_mm256_hadd_epi32(d0, d1), _mm256_hadd_epi32(d2, d3));
8210        _mm_add_epi32(_mm256_castsi256_si128(h), _mm256_extracti128_si256::<1>(h))
8211    }
8212}
8213
8214/// One 128-group's epilogue, in q1's float order: `d = 2·Σset − Σx` per
8215/// tile (exact i32 → exact f32), then `acc += d·s` tile by tile — four
8216/// separately rounded multiply/add pairs, the q1 kernels' sequence.
8217#[cfg(target_arch = "x86_64")]
8218#[target_feature(enable = "avx2")]
8219#[inline]
8220unsafe fn q1g_group_acc(
8221    acc: f32,
8222    msum4: core::arch::x86_64::__m128i,
8223    gsum4: &[i32],
8224    s: f32,
8225) -> f32 {
8226    // SAFETY: gsum4 holds the group's four tile sums (callers slice it).
8227    unsafe {
8228        use core::arch::x86_64::*;
8229        let g = _mm_loadu_si128(gsum4.as_ptr() as *const __m128i);
8230        let d = _mm_cvtepi32_ps(_mm_sub_epi32(_mm_slli_epi32::<1>(msum4), g));
8231        let mut f = [0f32; 4];
8232        _mm_storeu_ps(f.as_mut_ptr(), d);
8233        let mut acc = acc;
8234        for v in f {
8235            acc += v * s;
8236        }
8237        acc
8238    }
8239}
8240
8241/// AVX2 q1g row: `dot_q1_row_avx2`'s mask/select/maddubs per tile, but the
8242/// four tiles of a 128-group reduce together (`q1g_hsum4`) and the scale
8243/// converts once per group — bit-identical results (integer sums, then
8244/// q1's per-tile float sequence), fewer horizontal ops and f16 decodes.
8245#[cfg(target_arch = "x86_64")]
8246#[target_feature(enable = "avx2")]
8247unsafe fn dot_q1g_row_avx2(
8248    bytes: &[u8],
8249    sc_off: usize,
8250    r: usize,
8251    gpr: usize,
8252    xq: &[i8],
8253    gsum: &[i32],
8254) -> f32 {
8255    // SAFETY: callers uphold the plane-length and xq/gsum contracts
8256    // (gpr % 4 == 0, xq.len() == gpr·32, gsum.len() == gpr).
8257    unsafe {
8258        use core::arch::x86_64::*;
8259        let expand = _mm256_setr_epi8(
8260            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,
8261            3, 3, 3,
8262        );
8263        let bitsel = _mm256_setr_epi8(
8264            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
8265            -128, 1, 2, 4, 8, 16, 32, 64, -128,
8266        );
8267        let ones8 = _mm256_set1_epi8(1);
8268        let ones16 = _mm256_set1_epi16(1);
8269        let bp = bytes.as_ptr().add(r * gpr * 4);
8270        let mut acc = 0f32;
8271        let mut gi = 0;
8272        while gi < gpr {
8273            let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
8274            let mut d = [_mm256_setzero_si256(); 4];
8275            for (k, dk) in d.iter_mut().enumerate() {
8276                let t = gi + k;
8277                let bits = (bp.add(t * 4) as *const u32).read_unaligned();
8278                let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
8279                let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
8280                let x = _mm256_loadu_si256(xq.as_ptr().add(t * GROUP_SIZE) as *const __m256i);
8281                let sel = _mm256_and_si256(x, mask);
8282                *dk = _mm256_madd_epi16(_mm256_maddubs_epi16(ones8, sel), ones16);
8283            }
8284            let m4 = q1g_hsum4(d[0], d[1], d[2], d[3]);
8285            acc = q1g_group_acc(acc, m4, &gsum[gi..gi + 4], s);
8286            gi += 4;
8287        }
8288        acc
8289    }
8290}
8291
8292/// VNNI twin of `dot_q1g_row_avx2`: `vpdpbusd(1u8, sel)` per tile (the
8293/// same i32 lanes as maddubs+madd — bounded products), same group reduce.
8294#[cfg(target_arch = "x86_64")]
8295#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
8296unsafe fn dot_q1g_row_vnni(
8297    bytes: &[u8],
8298    sc_off: usize,
8299    r: usize,
8300    gpr: usize,
8301    xq: &[i8],
8302    gsum: &[i32],
8303) -> f32 {
8304    // SAFETY: callers uphold the plane-length and xq/gsum contracts.
8305    unsafe {
8306        use core::arch::x86_64::*;
8307        let expand = _mm256_setr_epi8(
8308            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,
8309            3, 3, 3,
8310        );
8311        let bitsel = _mm256_setr_epi8(
8312            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
8313            -128, 1, 2, 4, 8, 16, 32, 64, -128,
8314        );
8315        let ones8 = _mm256_set1_epi8(1);
8316        let bp = bytes.as_ptr().add(r * gpr * 4);
8317        let mut acc = 0f32;
8318        let mut gi = 0;
8319        while gi < gpr {
8320            let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
8321            let mut d = [_mm256_setzero_si256(); 4];
8322            for (k, dk) in d.iter_mut().enumerate() {
8323                let t = gi + k;
8324                let bits = (bp.add(t * 4) as *const u32).read_unaligned();
8325                let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
8326                let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
8327                let x = _mm256_loadu_si256(xq.as_ptr().add(t * GROUP_SIZE) as *const __m256i);
8328                *dk = _mm256_dpbusd_epi32(_mm256_setzero_si256(), ones8, _mm256_and_si256(x, mask));
8329            }
8330            let m4 = q1g_hsum4(d[0], d[1], d[2], d[3]);
8331            acc = q1g_group_acc(acc, m4, &gsum[gi..gi + 4], s);
8332            gi += 4;
8333        }
8334        acc
8335    }
8336}
8337
8338/// Blocked 1×4 AVX2 q1g row: each tile's mask is built once and serves
8339/// four activation streams; per stream, the group reduce and float order
8340/// of `dot_q1g_row_avx2` (so batch == matvec bit for bit).
8341#[cfg(target_arch = "x86_64")]
8342#[target_feature(enable = "avx2")]
8343unsafe fn dot_q1g_row_1x4_avx2(
8344    bytes: &[u8],
8345    sc_off: usize,
8346    r: usize,
8347    gpr: usize,
8348    xs: [&[i8]; 4],
8349    gsums: [&[i32]; 4],
8350) -> [f32; 4] {
8351    // SAFETY: callers uphold the plane-length and xq/gsum contracts.
8352    unsafe {
8353        use core::arch::x86_64::*;
8354        let expand = _mm256_setr_epi8(
8355            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,
8356            3, 3, 3,
8357        );
8358        let bitsel = _mm256_setr_epi8(
8359            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
8360            -128, 1, 2, 4, 8, 16, 32, 64, -128,
8361        );
8362        let ones8 = _mm256_set1_epi8(1);
8363        let ones16 = _mm256_set1_epi16(1);
8364        let bp = bytes.as_ptr().add(r * gpr * 4);
8365        let mut acc = [0f32; 4];
8366        let mut gi = 0;
8367        while gi < gpr {
8368            let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
8369            // d[stream][tile]
8370            let mut d = [[_mm256_setzero_si256(); 4]; 4];
8371            for t in 0..4 {
8372                let bits = (bp.add((gi + t) * 4) as *const u32).read_unaligned();
8373                let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
8374                let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
8375                for (k, xq) in xs.iter().enumerate() {
8376                    let x = _mm256_loadu_si256(
8377                        xq.as_ptr().add((gi + t) * GROUP_SIZE) as *const __m256i
8378                    );
8379                    let sel = _mm256_and_si256(x, mask);
8380                    d[k][t] = _mm256_madd_epi16(_mm256_maddubs_epi16(ones8, sel), ones16);
8381                }
8382            }
8383            for k in 0..4 {
8384                let m4 = q1g_hsum4(d[k][0], d[k][1], d[k][2], d[k][3]);
8385                acc[k] = q1g_group_acc(acc[k], m4, &gsums[k][gi..gi + 4], s);
8386            }
8387            gi += 4;
8388        }
8389        acc
8390    }
8391}
8392
8393/// VNNI twin of `dot_q1g_row_1x4_avx2`.
8394#[cfg(target_arch = "x86_64")]
8395#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
8396unsafe fn dot_q1g_row_1x4_vnni(
8397    bytes: &[u8],
8398    sc_off: usize,
8399    r: usize,
8400    gpr: usize,
8401    xs: [&[i8]; 4],
8402    gsums: [&[i32]; 4],
8403) -> [f32; 4] {
8404    // SAFETY: callers uphold the plane-length and xq/gsum contracts.
8405    unsafe {
8406        use core::arch::x86_64::*;
8407        let expand = _mm256_setr_epi8(
8408            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,
8409            3, 3, 3,
8410        );
8411        let bitsel = _mm256_setr_epi8(
8412            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
8413            -128, 1, 2, 4, 8, 16, 32, 64, -128,
8414        );
8415        let ones8 = _mm256_set1_epi8(1);
8416        let bp = bytes.as_ptr().add(r * gpr * 4);
8417        let mut acc = [0f32; 4];
8418        let mut gi = 0;
8419        while gi < gpr {
8420            let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
8421            let mut d = [[_mm256_setzero_si256(); 4]; 4];
8422            for t in 0..4 {
8423                let bits = (bp.add((gi + t) * 4) as *const u32).read_unaligned();
8424                let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
8425                let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
8426                for (k, xq) in xs.iter().enumerate() {
8427                    let x = _mm256_loadu_si256(
8428                        xq.as_ptr().add((gi + t) * GROUP_SIZE) as *const __m256i
8429                    );
8430                    d[k][t] = _mm256_dpbusd_epi32(
8431                        _mm256_setzero_si256(),
8432                        ones8,
8433                        _mm256_and_si256(x, mask),
8434                    );
8435                }
8436            }
8437            for k in 0..4 {
8438                let m4 = q1g_hsum4(d[k][0], d[k][1], d[k][2], d[k][3]);
8439                acc[k] = q1g_group_acc(acc[k], m4, &gsums[k][gi..gi + 4], s);
8440            }
8441            gi += 4;
8442        }
8443        acc
8444    }
8445}
8446
8447/// SDOT q1g row: the `dot_q1_row_sdot` math on the plane layout. One
8448/// aligned 16-byte load carries a whole 128-group's signs (q1 needs two
8449/// pair loads and a TBL gather of four scales for the same weights), TBL
8450/// replicates each sign byte across 8 lanes for vtst, and the group's ONE
8451/// scale converts once (FCVTL, exact) and broadcasts. Lane order, the
8452/// vpaddq tree, the fma into the four lane accumulators and the tail are
8453/// the q1 kernel's — including its quirk: q1's pair load reads 4 bytes
8454/// past the 4th tile, so the payload's final four tiles (last row) take
8455/// the scalar tail there. q1g has no overread, but takes the same tail on
8456/// the last row to stay bit-identical to q1.
8457#[cfg(target_arch = "aarch64")]
8458#[target_feature(enable = "neon,dotprod")]
8459#[allow(clippy::too_many_arguments)]
8460unsafe fn dot_q1g_row_sdot(
8461    bytes: &[u8],
8462    sc_off: usize,
8463    rows: usize,
8464    r: usize,
8465    gpr: usize,
8466    xq: &[i8],
8467    gsum: &[i32],
8468) -> f32 {
8469    // SAFETY: callers uphold the plane-length and xq/gsum contracts
8470    // (gpr % 4 == 0, xq.len() == gpr·32, gsum.len() == gpr).
8471    unsafe {
8472        use core::arch::aarch64::*;
8473        use core::arch::asm;
8474        const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
8475        // Tile t of the group = sign bytes 4t..4t+4 of the 16-byte load.
8476        const IW: [[u8; 16]; 8] = [
8477            [0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1],
8478            [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3],
8479            [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5],
8480            [6, 6, 6, 6, 6, 6, 6, 6, 7, 7, 7, 7, 7, 7, 7, 7],
8481            [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9],
8482            [
8483                10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
8484            ],
8485            [
8486                12, 12, 12, 12, 12, 12, 12, 12, 13, 13, 13, 13, 13, 13, 13, 13,
8487            ],
8488            [
8489                14, 14, 14, 14, 14, 14, 14, 14, 15, 15, 15, 15, 15, 15, 15, 15,
8490            ],
8491        ];
8492        let m = vld1q_u8(MASKS.as_ptr());
8493        let iw: [uint8x16_t; 8] = core::array::from_fn(|i| vld1q_u8(IW[i].as_ptr()));
8494        macro_rules! tile_dot {
8495            ($w0:expr, $w1:expr, $x:expr) => {{
8496                let x0 = vld1q_s8($x);
8497                let x1 = vld1q_s8($x.add(16));
8498                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
8499                asm!(
8500                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
8501                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
8502                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
8503                    w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
8504                    options(pure, nomem, nostack),
8505                );
8506                vaddq_s32(a0, a1)
8507            }};
8508        }
8509        macro_rules! tbl_mask {
8510            ($ld:expr, $i:expr) => {
8511                vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, iw[$i]), m))
8512            };
8513        }
8514        let bp = bytes.as_ptr().add(r * gpr * 4);
8515        let xp = xq.as_ptr();
8516        let gp = gsum.as_ptr();
8517        let vec_end = if r + 1 == rows { gpr - 4 } else { gpr };
8518        let mut accv = vdupq_n_f32(0.0);
8519        let mut gi = 0;
8520        while gi + 4 <= vec_end {
8521            let ld = vld1q_u8(bp.add(gi * 4));
8522            let d0 = tile_dot!(tbl_mask!(ld, 0), tbl_mask!(ld, 1), xp.add(gi * GROUP_SIZE));
8523            let d1 = tile_dot!(
8524                tbl_mask!(ld, 2),
8525                tbl_mask!(ld, 3),
8526                xp.add((gi + 1) * GROUP_SIZE)
8527            );
8528            let d2 = tile_dot!(
8529                tbl_mask!(ld, 4),
8530                tbl_mask!(ld, 5),
8531                xp.add((gi + 2) * GROUP_SIZE)
8532            );
8533            let d3 = tile_dot!(
8534                tbl_mask!(ld, 6),
8535                tbl_mask!(ld, 7),
8536                xp.add((gi + 3) * GROUP_SIZE)
8537            );
8538            let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
8539            let g = vld1q_s32(gp.add(gi));
8540            let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
8541            let h = vdup_n_u16(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
8542            let scf: float32x4_t;
8543            asm!(
8544                "fcvtl {o:v}.4s, {i:v}.4h",
8545                o = out(vreg) scf, i = in(vreg) h,
8546                options(pure, nomem, nostack),
8547            );
8548            accv = vfmaq_f32(accv, vcvtq_f32_s32(dots), scf);
8549            gi += 4;
8550        }
8551        let mut acc = vaddvq_f32(accv);
8552        while gi < gpr {
8553            let t = bp.add(gi * 4);
8554            let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
8555            let v0 = vcombine_u8(vdup_n_u8(*t), vdup_n_u8(*t.add(1)));
8556            let v1 = vcombine_u8(vdup_n_u8(*t.add(2)), vdup_n_u8(*t.add(3)));
8557            let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
8558            let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
8559            let d = vaddvq_s32(tile_dot!(w0, w1, xp.add(gi * GROUP_SIZE)));
8560            acc += (-(2 * d + *gp.add(gi))) as f32 * s;
8561            gi += 1;
8562        }
8563        acc
8564    }
8565}
8566
8567/// Blocked q1g 1×4: one unpack of the group serves four activation
8568/// streams; per stream bit-identical to `dot_q1g_row_sdot` (and so to the
8569/// q1 1×4 kernel, same last-row tail).
8570#[cfg(target_arch = "aarch64")]
8571#[target_feature(enable = "neon,dotprod")]
8572#[allow(clippy::too_many_arguments)]
8573unsafe fn dot_q1g_row_1x4_sdot(
8574    bytes: &[u8],
8575    sc_off: usize,
8576    rows: usize,
8577    r: usize,
8578    gpr: usize,
8579    xs: [&[i8]; 4],
8580    gs: [&[i32]; 4],
8581) -> [f32; 4] {
8582    // SAFETY: same contracts as `dot_q1g_row_sdot`, ×4.
8583    unsafe {
8584        use core::arch::aarch64::*;
8585        use core::arch::asm;
8586        const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
8587        const IW: [[u8; 16]; 8] = [
8588            [0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1],
8589            [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3],
8590            [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5],
8591            [6, 6, 6, 6, 6, 6, 6, 6, 7, 7, 7, 7, 7, 7, 7, 7],
8592            [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9],
8593            [
8594                10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
8595            ],
8596            [
8597                12, 12, 12, 12, 12, 12, 12, 12, 13, 13, 13, 13, 13, 13, 13, 13,
8598            ],
8599            [
8600                14, 14, 14, 14, 14, 14, 14, 14, 15, 15, 15, 15, 15, 15, 15, 15,
8601            ],
8602        ];
8603        let m = vld1q_u8(MASKS.as_ptr());
8604        let iw: [uint8x16_t; 8] = core::array::from_fn(|i| vld1q_u8(IW[i].as_ptr()));
8605        macro_rules! sdot2 {
8606            ($w0:expr, $w1:expr, $x:expr) => {{
8607                let x0 = vld1q_s8($x);
8608                let x1 = vld1q_s8($x.add(16));
8609                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
8610                asm!(
8611                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
8612                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
8613                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
8614                    w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
8615                    options(pure, nomem, nostack),
8616                );
8617                vaddq_s32(a0, a1)
8618            }};
8619        }
8620        let bp = bytes.as_ptr().add(r * gpr * 4);
8621        let vec_end = if r + 1 == rows { gpr - 4 } else { gpr };
8622        let mut accv = [vdupq_n_f32(0.0); 4];
8623        let mut gi = 0;
8624        while gi + 4 <= vec_end {
8625            let ld = vld1q_u8(bp.add(gi * 4));
8626            // Unpack ONCE — eight ±mask vectors serve all four streams.
8627            let w: [int8x16_t; 8] =
8628                core::array::from_fn(|i| vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld, iw[i]), m)));
8629            let h = vdup_n_u16(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
8630            let scf: float32x4_t;
8631            asm!(
8632                "fcvtl {o:v}.4s, {i:v}.4h",
8633                o = out(vreg) scf, i = in(vreg) h,
8634                options(pure, nomem, nostack),
8635            );
8636            for k in 0..4 {
8637                let xp = xs[k].as_ptr();
8638                let d0 = sdot2!(w[0], w[1], xp.add(gi * GROUP_SIZE));
8639                let d1 = sdot2!(w[2], w[3], xp.add((gi + 1) * GROUP_SIZE));
8640                let d2 = sdot2!(w[4], w[5], xp.add((gi + 2) * GROUP_SIZE));
8641                let d3 = sdot2!(w[6], w[7], xp.add((gi + 3) * GROUP_SIZE));
8642                let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
8643                let g = vld1q_s32(gs[k].as_ptr().add(gi));
8644                let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
8645                accv[k] = vfmaq_f32(accv[k], vcvtq_f32_s32(dots), scf);
8646            }
8647            gi += 4;
8648        }
8649        let mut acc = [
8650            vaddvq_f32(accv[0]),
8651            vaddvq_f32(accv[1]),
8652            vaddvq_f32(accv[2]),
8653            vaddvq_f32(accv[3]),
8654        ];
8655        while gi < gpr {
8656            let t = bp.add(gi * 4);
8657            let sc = f16_to_f32(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
8658            let v0 = vcombine_u8(vdup_n_u8(*t), vdup_n_u8(*t.add(1)));
8659            let v1 = vcombine_u8(vdup_n_u8(*t.add(2)), vdup_n_u8(*t.add(3)));
8660            let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
8661            let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
8662            for k in 0..4 {
8663                let d = vaddvq_s32(sdot2!(w0, w1, xs[k].as_ptr().add(gi * GROUP_SIZE)));
8664                acc[k] += (-(2 * d + *gs[k].as_ptr().add(gi))) as f32 * sc;
8665            }
8666            gi += 1;
8667        }
8668        acc
8669    }
8670}
8671
8672/// (weight ±1, scale) of one q1g element — the exact outlier term.
8673#[inline]
8674fn q1g_outlier(bytes: &[u8], sc_off: usize, r: usize, gpr: usize, j: usize) -> (f32, f32) {
8675    let t = r * gpr + j / GROUP_SIZE;
8676    let k = j % GROUP_SIZE;
8677    let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, t));
8678    let bit = (bytes[t * 4 + k / 8] >> (k % 8)) & 1;
8679    ((bit as i32 * 2 - 1) as f32, s)
8680}
8681
8682/// Exact scalar q1g row (CMF_SDOT=0 contract) — `q1_row_exact`'s order.
8683#[inline]
8684fn q1g_row_exact(bytes: &[u8], sc_off: usize, r: usize, gpr: usize, x: &[f32]) -> f32 {
8685    let mut acc = 0f32;
8686    for gi in 0..gpr {
8687        let t = r * gpr + gi;
8688        let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, t));
8689        let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
8690        let mut ga = 0f32;
8691        for (j, &b) in bytes[t * 4..t * 4 + 4].iter().enumerate() {
8692            for k in 0..8 {
8693                ga += (((b >> k) & 1) as f32 * 2.0 - 1.0) * xg[j * 8 + k];
8694            }
8695        }
8696        acc += ga * s;
8697    }
8698    acc
8699}
8700
8701/// q1t per-row overlay locator. After the base (`base_len`) come
8702/// `[u32 row_ptr[rows+1]]` then `[(u16 col, f16 val)]` grouped by row (row
8703/// `r`'s entries are `[row_ptr[r], row_ptr[r+1])`). Returns
8704/// `(row_ptr offset, entries offset, present)`.
8705fn q1t_overlay(bytes: &[u8], base_len: usize, rows: usize) -> (usize, usize, bool) {
8706    let entries = base_len + (rows + 1) * 4;
8707    (base_len, entries, entries <= bytes.len())
8708}
8709
8710/// Read `row_ptr[r]` from the overlay's prefix-sum table.
8711#[inline]
8712fn q1t_rowptr(bytes: &[u8], rp_off: usize, r: usize) -> usize {
8713    let o = rp_off + r * 4;
8714    u32::from_le_bytes([bytes[o], bytes[o + 1], bytes[o + 2], bytes[o + 3]]) as usize
8715}
8716
8717/// Byte → the 5 ternary signs it packs `{−1,0,+1}` as f32, precomputed so
8718/// decoding a q1t code is a table load, not the base-3 divide/modulo per
8719/// weight (division is ~20–40× the cost of a load). Built at compile time.
8720const SIGN5: [[f32; 5]; 256] = {
8721    let mut lut = [[0.0f32; 5]; 256];
8722    let pow3 = [1u16, 3, 9, 27, 81];
8723    let mut byte = 0usize;
8724    while byte < 256 {
8725        let mut i = 0usize;
8726        while i < 5 {
8727            let code = (byte as u16 / pow3[i]) % 3;
8728            lut[byte][i] = if code == 1 {
8729                1.0
8730            } else if code == 2 {
8731                -1.0
8732            } else {
8733                0.0
8734            };
8735            i += 1;
8736        }
8737        byte += 1;
8738    }
8739    lut
8740};
8741
8742/// Same table, as i8 signs — the operand for the int8 SDOT base kernel.
8743const SIGN5_I8: [[i8; 5]; 256] = {
8744    let mut lut = [[0i8; 5]; 256];
8745    let pow3 = [1u16, 3, 9, 27, 81];
8746    let mut byte = 0usize;
8747    while byte < 256 {
8748        let mut i = 0usize;
8749        while i < 5 {
8750            let code = (byte as u16 / pow3[i]) % 3;
8751            lut[byte][i] = if code == 1 {
8752                1
8753            } else if code == 2 {
8754                -1
8755            } else {
8756                0
8757            };
8758            i += 1;
8759        }
8760        byte += 1;
8761    }
8762    lut
8763};
8764
8765/// The same 5 i8 signs packed into a u64 (`[s0 s1 s2 s3 s4 0 0 0]`, LE) so the
8766/// group unpack is 7 unaligned u64 stores at offsets 0,5,10,…,30 instead of
8767/// six 5-byte copies + LUT indexing — each store's trailing zeros are fixed by
8768/// the next store, and the last one runs 6 B past the 32nd weight (the unpack
8769/// buffer is padded to 40). This is the decode/prefill hot inner op.
8770const SIGN5_U64: [u64; 256] = {
8771    let mut lut = [0u64; 256];
8772    let pow3 = [1u16, 3, 9, 27, 81];
8773    let mut byte = 0usize;
8774    while byte < 256 {
8775        let mut v = 0u64;
8776        let mut i = 0usize;
8777        while i < 5 {
8778            let code = (byte as u16 / pow3[i]) % 3;
8779            let s: u8 = if code == 1 {
8780                1
8781            } else if code == 2 {
8782                0xFF
8783            } else {
8784                0
8785            };
8786            v |= (s as u64) << (i * 8);
8787            i += 1;
8788        }
8789        lut[byte] = v;
8790        byte += 1;
8791    }
8792    lut
8793};
8794
8795/// Ternary base weight at `(row r, col j)` = `sign(code)·s_group`. Used to add
8796/// back activation-outlier columns, whose `x` was zeroed for the int8 bulk dot
8797/// (`split_act`). At a weight-outlier position the code is 0, so this is 0 and
8798/// the overlay correction owns that column — no double counting.
8799#[inline]
8800fn q1t_base_weight(bytes: &[u8], r: usize, gpr: usize, j: usize) -> f32 {
8801    const TILE: usize = cortiq_core::quant::Q1T_TILE;
8802    let off = (r * gpr + j / GROUP_SIZE) * TILE;
8803    let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
8804    let within = j % GROUP_SIZE;
8805    SIGN5[bytes[off + 2 + within / 5] as usize][within % 5] * s
8806}
8807
8808/// One 32-group int8 dot via two SDOTs. Bit-exact vs the scalar i8 sum
8809/// (integer accumulation is order-independent).
8810#[cfg(target_arch = "aarch64")]
8811#[target_feature(enable = "neon,dotprod")]
8812#[inline]
8813unsafe fn sdot32_i8(w: *const i8, x: *const i8) -> i32 {
8814    // SAFETY: caller guarantees 32 readable i8 at each pointer.
8815    unsafe {
8816        use core::arch::aarch64::*;
8817        use core::arch::asm;
8818        let w0 = vld1q_s8(w);
8819        let w1 = vld1q_s8(w.add(16));
8820        let x0 = vld1q_s8(x);
8821        let x1 = vld1q_s8(x.add(16));
8822        let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
8823        asm!(
8824            "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
8825            "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
8826            a0 = inout(vreg) a0, a1 = inout(vreg) a1,
8827            w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
8828            options(pure, nomem, nostack),
8829        );
8830        vaddvq_s32(vaddq_s32(a0, a1))
8831    }
8832}
8833
8834/// One 32-group int8 dot via AVX2: signed·signed as `maddubs(|w|, sign(x,w))`
8835/// then `madd` and a horizontal reduce (the same idiom as `dot_q4t_row_avx2`).
8836#[cfg(target_arch = "x86_64")]
8837#[target_feature(enable = "avx2")]
8838#[inline]
8839unsafe fn i8dot32_avx2(w: *const i8, x: *const i8) -> i32 {
8840    // SAFETY: caller guarantees 32 readable i8 at each pointer.
8841    unsafe {
8842        use core::arch::x86_64::*;
8843        let wv = _mm256_loadu_si256(w as *const __m256i);
8844        let xv = _mm256_loadu_si256(x as *const __m256i);
8845        let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
8846        let d = _mm256_madd_epi16(p16, _mm256_set1_epi16(1));
8847        let hi128 = _mm256_extracti128_si256::<1>(d);
8848        let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
8849        let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8850        let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8851        _mm_cvtsi128_si32(s32)
8852    }
8853}
8854
8855/// Unpack one q1t group's base-3 codes into 32 i8 signs via 7 unaligned u64
8856/// stores (see `SIGN5_U64`). `dst` MUST have ≥ 40 bytes: the 7th store writes
8857/// `dst[30..38]`. Stores go in order so each one's trailing zeros are
8858/// overwritten by the next; the final 6 padding bytes are unused by the dot.
8859#[inline]
8860fn q1t_unpack_group_i8(codes: *const u8, dst: &mut [i8]) {
8861    debug_assert!(dst.len() >= 40);
8862    // SAFETY: codes points at 7 readable bytes; dst has ≥ 40 bytes so every
8863    // 8-byte store at offset bi*5 (bi ≤ 6 → ≤ 30) stays in bounds.
8864    unsafe {
8865        let p = dst.as_mut_ptr();
8866        for bi in 0..7 {
8867            core::ptr::write_unaligned(
8868                p.add(bi * 5) as *mut u64,
8869                SIGN5_U64[*codes.add(bi) as usize],
8870            );
8871        }
8872    }
8873}
8874
8875/// One 32-group int8 dot, arch-dispatched (the matmat inner loop, where the
8876/// row's signs are unpacked once and dotted against every batch input).
8877/// Callers are gated by `a8w8_enabled()`, so the target-feature arms are
8878/// reachable; the scalar arm is a non-SIMD-arch fallback.
8879#[inline]
8880fn q1t_i8dot32(w: *const i8, x: *const i8) -> i32 {
8881    #[cfg(target_arch = "aarch64")]
8882    unsafe {
8883        return sdot32_i8(w, x);
8884    }
8885    #[cfg(target_arch = "x86_64")]
8886    unsafe {
8887        return i8dot32_avx2(w, x);
8888    }
8889    #[allow(unreachable_code)]
8890    unsafe {
8891        let mut s = 0i32;
8892        for k in 0..GROUP_SIZE {
8893            s += *w.add(k) as i32 * *x.add(k) as i32;
8894        }
8895        s
8896    }
8897}
8898
8899#[inline]
8900unsafe fn q1t_unpack_reg_u64s(codes: *const u8) -> (u64, u64, u64, u64) {
8901    let (s0, s1, s2, s3, s4, s5, s6) = unsafe {
8902        (
8903            SIGN5_U64[*codes as usize],
8904            SIGN5_U64[*codes.add(1) as usize],
8905            SIGN5_U64[*codes.add(2) as usize],
8906            SIGN5_U64[*codes.add(3) as usize],
8907            SIGN5_U64[*codes.add(4) as usize],
8908            SIGN5_U64[*codes.add(5) as usize],
8909            SIGN5_U64[*codes.add(6) as usize],
8910        )
8911    };
8912
8913    let u0 = s0 | (s1 << 40);
8914    let u1 = (s1 >> 24) | (s2 << 16) | (s3 << 56);
8915    let u2 = (s3 >> 8) | (s4 << 32);
8916    let u3 = (s4 >> 32) | (s5 << 8) | (s6 << 48);
8917
8918    (u0, u1, u2, u3)
8919}
8920
8921/// One q1t row's int8 base dot: `Σ_group s·dot(signs, xq)` (before the shared
8922/// `sx`). Direct register unpacking (zero stack stores/loads, no STLF stalls).
8923/// ARM SDOT.
8924#[cfg(target_arch = "aarch64")]
8925#[target_feature(enable = "neon,dotprod")]
8926unsafe fn q1t_dot_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
8927    use core::arch::aarch64::*;
8928    use core::arch::asm;
8929    unsafe {
8930        const TILE: usize = cortiq_core::quant::Q1T_TILE;
8931        let mut acc = 0f32;
8932        let bytes_ptr = bytes.as_ptr();
8933        let xq_ptr = xq.as_ptr();
8934        let row_off = r * gpr * TILE;
8935
8936        let gpr2 = gpr & !1;
8937        let mut gi = 0;
8938        while gi < gpr2 {
8939            let off0 = row_off + gi * TILE;
8940            let off1 = off0 + TILE;
8941            let s0 = f16_to_f32(u16::from_le_bytes([
8942                *bytes_ptr.add(off0),
8943                *bytes_ptr.add(off0 + 1),
8944            ]));
8945            let s1 = f16_to_f32(u16::from_le_bytes([
8946                *bytes_ptr.add(off1),
8947                *bytes_ptr.add(off1 + 1),
8948            ]));
8949
8950            let (u0_0, u1_0, u2_0, u3_0) = q1t_unpack_reg_u64s(bytes_ptr.add(off0 + 2));
8951            let (u0_1, u1_1, u2_1, u3_1) = q1t_unpack_reg_u64s(bytes_ptr.add(off1 + 2));
8952
8953            let w0_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_0), vcreate_u64(u1_0)));
8954            let w1_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_0), vcreate_u64(u3_0)));
8955            let w0_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_1), vcreate_u64(u1_1)));
8956            let w1_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_1), vcreate_u64(u3_1)));
8957
8958            let x0_0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE));
8959            let x1_0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE + 16));
8960            let x0_1 = vld1q_s8(xq_ptr.add((gi + 1) * GROUP_SIZE));
8961            let x1_1 = vld1q_s8(xq_ptr.add((gi + 1) * GROUP_SIZE + 16));
8962
8963            let (mut a0_0, mut a1_0) = (vdupq_n_s32(0), vdupq_n_s32(0));
8964            let (mut a0_1, mut a1_1) = (vdupq_n_s32(0), vdupq_n_s32(0));
8965            asm!(
8966                "sdot {a0_0:v}.4s, {w0_0:v}.16b, {x0_0:v}.16b",
8967                "sdot {a1_0:v}.4s, {w1_0:v}.16b, {x1_0:v}.16b",
8968                "sdot {a0_1:v}.4s, {w0_1:v}.16b, {x0_1:v}.16b",
8969                "sdot {a1_1:v}.4s, {w1_1:v}.16b, {x1_1:v}.16b",
8970                a0_0 = inout(vreg) a0_0, a1_0 = inout(vreg) a1_0,
8971                a0_1 = inout(vreg) a0_1, a1_1 = inout(vreg) a1_1,
8972                w0_0 = in(vreg) w0_0, x0_0 = in(vreg) x0_0, w1_0 = in(vreg) w1_0, x1_0 = in(vreg) x1_0,
8973                w0_1 = in(vreg) w0_1, x0_1 = in(vreg) x0_1, w1_1 = in(vreg) w1_1, x1_1 = in(vreg) x1_1,
8974                options(pure, nomem, nostack),
8975            );
8976            let d0 = vaddvq_s32(vaddq_s32(a0_0, a1_0));
8977            let d1 = vaddvq_s32(vaddq_s32(a0_1, a1_1));
8978            acc += d0 as f32 * s0 + d1 as f32 * s1;
8979            gi += 2;
8980        }
8981
8982        if gi < gpr {
8983            let off = row_off + gi * TILE;
8984            let s = f16_to_f32(u16::from_le_bytes([
8985                *bytes_ptr.add(off),
8986                *bytes_ptr.add(off + 1),
8987            ]));
8988            let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
8989            let w0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0), vcreate_u64(u1)));
8990            let w1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2), vcreate_u64(u3)));
8991            let x0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE));
8992            let x1 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE + 16));
8993            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
8994            asm!(
8995                "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
8996                "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
8997                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
8998                w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
8999                options(pure, nomem, nostack),
9000            );
9001            let d = vaddvq_s32(vaddq_s32(a0, a1));
9002            acc += d as f32 * s;
9003        }
9004        acc
9005    }
9006}
9007
9008/// x86 AVX2 mirror of `q1t_dot_row_sdot` (maddubs int8 dot per group).
9009#[cfg(target_arch = "x86_64")]
9010#[target_feature(enable = "avx2")]
9011unsafe fn q1t_dot_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
9012    use core::arch::x86_64::*;
9013    unsafe {
9014        const TILE: usize = cortiq_core::quant::Q1T_TILE;
9015        let mut acc = 0f32;
9016        let bytes_ptr = bytes.as_ptr();
9017        let xq_ptr = xq.as_ptr();
9018        let row_off = r * gpr * TILE;
9019
9020        let ones = _mm256_set1_epi16(1);
9021        for gi in 0..gpr {
9022            let off = row_off + gi * TILE;
9023            let s = f16_to_f32(u16::from_le_bytes([
9024                *bytes_ptr.add(off),
9025                *bytes_ptr.add(off + 1),
9026            ]));
9027            let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
9028            let wv = _mm256_set_epi64x(u3 as i64, u2 as i64, u1 as i64, u0 as i64);
9029            let xv = _mm256_loadu_si256(xq_ptr.add(gi * GROUP_SIZE) as *const __m256i);
9030            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
9031            let d256 = _mm256_madd_epi16(p16, ones);
9032            let d128 = _mm_add_epi32(
9033                _mm256_castsi256_si128(d256),
9034                _mm256_extracti128_si256(d256, 1),
9035            );
9036            let d64 = _mm_add_epi32(d128, _mm_shuffle_epi32(d128, 0xee));
9037            let d32 = _mm_cvtsi128_si32(_mm_add_epi32(d64, _mm_shuffle_epi32(d64, 0x55)));
9038            acc += d32 as f32 * s;
9039        }
9040        acc
9041    }
9042}
9043
9044/// VNNI twin of `q1t_dot_row_avx2` (see `dpbusd_hsum`).
9045#[cfg(target_arch = "x86_64")]
9046#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
9047unsafe fn q1t_dot_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
9048    use core::arch::x86_64::*;
9049    // SAFETY: same tile/xq contracts as `q1t_dot_row_avx2`.
9050    unsafe {
9051        const TILE: usize = cortiq_core::quant::Q1T_TILE;
9052        let mut acc = 0f32;
9053        let bytes_ptr = bytes.as_ptr();
9054        let xq_ptr = xq.as_ptr();
9055        let row_off = r * gpr * TILE;
9056        for gi in 0..gpr {
9057            let off = row_off + gi * TILE;
9058            let s = f16_to_f32(u16::from_le_bytes([
9059                *bytes_ptr.add(off),
9060                *bytes_ptr.add(off + 1),
9061            ]));
9062            let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
9063            let wv = _mm256_set_epi64x(u3 as i64, u2 as i64, u1 as i64, u0 as i64);
9064            let xv = _mm256_loadu_si256(xq_ptr.add(gi * GROUP_SIZE) as *const __m256i);
9065            let d = dpbusd_hsum(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
9066            acc += d as f32 * s;
9067        }
9068        acc
9069    }
9070}
9071
9072/// Per-row int8 base dot, dispatched once per row (matvec decode hot path).
9073/// Callers are gated by `a8w8_enabled()`, so the target-feature kernels are
9074/// reachable.
9075#[inline]
9076fn q1t_dot_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
9077    #[cfg(target_arch = "aarch64")]
9078    unsafe {
9079        return q1t_dot_row_sdot(bytes, r, gpr, xq);
9080    }
9081    #[cfg(target_arch = "x86_64")]
9082    unsafe {
9083        if vnni_tiles_enabled() {
9084            return q1t_dot_row_vnni(bytes, r, gpr, xq);
9085        }
9086        return q1t_dot_row_avx2(bytes, r, gpr, xq);
9087    }
9088    #[allow(unreachable_code)]
9089    {
9090        const TILE: usize = cortiq_core::quant::Q1T_TILE;
9091        let mut acc = 0f32;
9092        let mut sg = [0i8; GROUP_SIZE + 8]; // +8 slack for the u64-store unpack
9093        for gi in 0..gpr {
9094            let off = (r * gpr + gi) * TILE;
9095            let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
9096            q1t_unpack_group_i8(bytes.as_ptr().wrapping_add(off + 2), &mut sg);
9097            let mut d = 0i32;
9098            for k in 0..GROUP_SIZE {
9099                d += sg[k] as i32 * xq[gi * GROUP_SIZE + k] as i32;
9100            }
9101            acc += d as f32 * s;
9102        }
9103        acc
9104    }
9105}
9106
9107/// Σ over a row's outliers of `value·x[col]` — the correction that adds the
9108/// overlay's exact weights on top of the base dot. INVARIANT: the encoder
9109/// writes ternary code 0 at every outlier position (`quantize_q1t`), so the
9110/// base contributes nothing there and this is a plain `value·x`, not
9111/// `(value − base)·x` — no scattered per-outlier scale read. Row `r`'s entries
9112/// are the contiguous slice `[row_ptr[r], row_ptr[r+1])`, so no binary search.
9113fn q1t_row_outlier_correction(
9114    bytes: &[u8],
9115    r: usize,
9116    rp_off: usize,
9117    entries_off: usize,
9118    has_ov: bool,
9119    x: &[f32],
9120) -> f32 {
9121    if !has_ov {
9122        return 0.0;
9123    }
9124    let (c0, c1) = (
9125        q1t_rowptr(bytes, rp_off, r),
9126        q1t_rowptr(bytes, rp_off, r + 1),
9127    );
9128    let mut corr = 0f32;
9129    for p in c0..c1 {
9130        let e = entries_off + p * 4;
9131        let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
9132        let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
9133        corr += val * x[col];
9134    }
9135    corr
9136}
9137
9138/// Dequantize one q1t row into `buf[..cols]` via the sign LUT (no division),
9139/// then apply the row's outliers (its `[row_ptr[r], row_ptr[r+1])` slice).
9140/// Used by the batched (prefill) path where the decode amortizes over the batch.
9141fn q1t_dequant_row(
9142    bytes: &[u8],
9143    r: usize,
9144    gpr: usize,
9145    rp_off: usize,
9146    entries_off: usize,
9147    has_ov: bool,
9148    buf: &mut [f32],
9149) {
9150    const TILE: usize = cortiq_core::quant::Q1T_TILE;
9151    for g in 0..gpr {
9152        let off = (r * gpr + g) * TILE;
9153        let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
9154        let codes = &bytes[off + 2..off + TILE];
9155        let bc = g * GROUP_SIZE;
9156        // 6 full bytes (30 codes) + a 7th byte holding the last 2.
9157        for bi in 0..6 {
9158            let lut = &SIGN5[codes[bi] as usize];
9159            let d = &mut buf[bc + bi * 5..bc + bi * 5 + 5];
9160            for i in 0..5 {
9161                d[i] = lut[i] * s;
9162            }
9163        }
9164        let lut = &SIGN5[codes[6] as usize];
9165        buf[bc + 30] = lut[0] * s;
9166        buf[bc + 31] = lut[1] * s;
9167    }
9168    if !has_ov {
9169        return;
9170    }
9171    let (c0, c1) = (
9172        q1t_rowptr(bytes, rp_off, r),
9173        q1t_rowptr(bytes, rp_off, r + 1),
9174    );
9175    for p in c0..c1 {
9176        let e = entries_off + p * 4;
9177        let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
9178        buf[col] = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
9179    }
9180}
9181
9182/// Add the sparse outlier overlay onto a base dot already in `out` (the GPU
9183/// computes the ternary base; the overlay stays on the CPU — its entries are
9184/// few and its per-row gather doesn't vectorize on the GPU). Row-parallel.
9185fn q1t_add_overlay(
9186    bytes: &[u8],
9187    x: &[f32],
9188    rows: usize,
9189    cols: usize,
9190    out: &mut [f32],
9191    pool: Option<&Pool>,
9192) {
9193    const TILE: usize = cortiq_core::quant::Q1T_TILE;
9194    let gpr = cols / GROUP_SIZE;
9195    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
9196    if !has_ov {
9197        return;
9198    }
9199    let out_addr = SendMut(out.as_mut_ptr());
9200    let run = move |start: usize, end: usize| {
9201        for r in start..end {
9202            let corr = q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
9203            // SAFETY: disjoint rows; add onto the base the GPU already wrote.
9204            unsafe { *out_addr.at(r) += corr };
9205        }
9206    };
9207    dispatch_rows(pool, rows, &run);
9208}
9209
9210/// Q1T row range via the A8W8 int8 path — shared activation split,
9211/// per-row: base SDOT dot + outlier correction + overlay.
9212#[allow(clippy::too_many_arguments)]
9213fn q1t_range_a8w8(
9214    bytes: &[u8],
9215    gpr: usize,
9216    rp_off: usize,
9217    ent_off: usize,
9218    has_ov: bool,
9219    act: &SplitAct,
9220    x: &[f32],
9221    out: SendMut,
9222    start: usize,
9223    end: usize,
9224) {
9225    for r in start..end {
9226        let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
9227        for &(j, xv) in &act.outliers {
9228            acc += q1t_base_weight(bytes, r, gpr, j) * xv;
9229        }
9230        acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
9231        // SAFETY: disjoint row ranges per worker.
9232        unsafe { *out.at(r) = acc };
9233    }
9234}
9235
9236/// Q1T row range via the f32 path (no SDOT) — for matvec_many batched
9237/// dispatch when a8w8 is unavailable.
9238#[allow(clippy::too_many_arguments)]
9239fn q1t_range_f32_batch(
9240    bytes: &[u8],
9241    gpr: usize,
9242    rp_off: usize,
9243    ent_off: usize,
9244    has_ov: bool,
9245    x: &[f32],
9246    out: SendMut,
9247    start: usize,
9248    end: usize,
9249) {
9250    const TILE: usize = cortiq_core::quant::Q1T_TILE;
9251    let mut sg = [0f32; GROUP_SIZE];
9252    for r in start..end {
9253        let mut acc = 0f32;
9254        for g in 0..gpr {
9255            let off = (r * gpr + g) * TILE;
9256            let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
9257            let codes = &bytes[off + 2..off + TILE];
9258            let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
9259            for bi in 0..6 {
9260                sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
9261            }
9262            let lut = &SIGN5[codes[6] as usize];
9263            sg[30] = lut[0];
9264            sg[31] = lut[1];
9265            let mut gsum = 0f32;
9266            for k in 0..GROUP_SIZE {
9267                gsum += sg[k] * xg[k];
9268            }
9269            acc += s * gsum;
9270        }
9271        acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
9272        // SAFETY: disjoint row ranges per worker.
9273        unsafe { *out.at(r) = acc };
9274    }
9275}
9276
9277/// Ternary (q1t) matvec — decode+dot straight from mmap, one group at a time:
9278/// no per-ROW buffer, no division (the sign LUT), and a tiny per-group sign
9279/// buffer so the 32-wide dot vectorizes. This is the decode hot path.
9280fn q1t_matvec(
9281    bytes: &[u8],
9282    x: &[f32],
9283    rows: usize,
9284    cols: usize,
9285    out: &mut [f32],
9286    pool: Option<&Pool>,
9287) {
9288    debug_assert_eq!(out.len(), rows);
9289    const TILE: usize = cortiq_core::quant::Q1T_TILE;
9290    let gpr = cols / GROUP_SIZE;
9291    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
9292    let out_addr = SendMut(out.as_mut_ptr());
9293    // int8 SDOT base dot (ARM dotprod): ~4× the f32 arithmetic. x → i8 once
9294    // (`split_act`), activation outliers added back exactly in f32, weight
9295    // overlay on top. ARM SDOT / x86 AVX2; CMF_SDOT=0 keeps the exact f32 path.
9296    if a8w8_enabled() {
9297        let act = split_act(x);
9298        let act = &act;
9299        let run = move |start: usize, end: usize| {
9300            for r in start..end {
9301                let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
9302                for &(j, xv) in &act.outliers {
9303                    acc += q1t_base_weight(bytes, r, gpr, j) * xv;
9304                }
9305                acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
9306                // SAFETY: disjoint row ranges per worker.
9307                unsafe { *out_addr.at(r) = acc };
9308            }
9309        };
9310        dispatch_rows(pool, rows, &run);
9311        return;
9312    }
9313    let run = move |start: usize, end: usize| {
9314        // Per-group signs, unpacked contiguously so the dot below is a clean
9315        // 32-wide reduction the autovectorizer turns into f32x4 FMAs — the
9316        // 5-values-per-byte base-3 layout won't SIMD in place.
9317        let mut sg = [0f32; GROUP_SIZE];
9318        for r in start..end {
9319            let mut acc = 0f32;
9320            for g in 0..gpr {
9321                let off = (r * gpr + g) * TILE;
9322                let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
9323                let codes = &bytes[off + 2..off + TILE];
9324                let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
9325                for bi in 0..6 {
9326                    sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
9327                }
9328                let lut = &SIGN5[codes[6] as usize];
9329                sg[30] = lut[0];
9330                sg[31] = lut[1];
9331                let mut gsum = 0f32;
9332                for k in 0..GROUP_SIZE {
9333                    gsum += sg[k] * xg[k];
9334                }
9335                acc += s * gsum;
9336            }
9337            acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
9338            unsafe { *out_addr.at(r) = acc };
9339        }
9340    };
9341    dispatch_rows(pool, rows, &run);
9342}
9343
9344/// Fused-pair twin of `q1t_dot_row_sdot`: ONE register unpack of the
9345/// ternary codes serves BOTH activation streams (the unpack chain is
9346/// the dominant per-row cost — MTP verify pairs paid it twice). Per
9347/// stream the group order and f32 accumulation match the single-row
9348/// kernel exactly, so pair == 2×matvec bit-for-bit.
9349#[cfg(target_arch = "aarch64")]
9350#[target_feature(enable = "neon,dotprod")]
9351unsafe fn q1t_dot_row_sdot2(bytes: &[u8], r: usize, gpr: usize, xa: &[i8], xb: &[i8]) -> [f32; 2] {
9352    use core::arch::aarch64::*;
9353    use core::arch::asm;
9354    // SAFETY: same slice-length contracts as `q1t_dot_row_sdot`, ×2.
9355    unsafe {
9356        const TILE: usize = cortiq_core::quant::Q1T_TILE;
9357        let bytes_ptr = bytes.as_ptr();
9358        let row_off = r * gpr * TILE;
9359        let xp = [xa.as_ptr(), xb.as_ptr()];
9360        let mut acc = [0f32; 2];
9361        macro_rules! sdot2 {
9362            ($w0:expr, $w1:expr, $x:expr) => {{
9363                let x0 = vld1q_s8($x);
9364                let x1 = vld1q_s8($x.add(16));
9365                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
9366                asm!(
9367                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
9368                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
9369                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
9370                    w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
9371                    options(pure, nomem, nostack),
9372                );
9373                vaddvq_s32(vaddq_s32(a0, a1))
9374            }};
9375        }
9376        let gpr2 = gpr & !1;
9377        let mut gi = 0;
9378        while gi < gpr2 {
9379            let off0 = row_off + gi * TILE;
9380            let off1 = off0 + TILE;
9381            let s0 = f16_to_f32(u16::from_le_bytes([
9382                *bytes_ptr.add(off0),
9383                *bytes_ptr.add(off0 + 1),
9384            ]));
9385            let s1 = f16_to_f32(u16::from_le_bytes([
9386                *bytes_ptr.add(off1),
9387                *bytes_ptr.add(off1 + 1),
9388            ]));
9389            let (u0_0, u1_0, u2_0, u3_0) = q1t_unpack_reg_u64s(bytes_ptr.add(off0 + 2));
9390            let (u0_1, u1_1, u2_1, u3_1) = q1t_unpack_reg_u64s(bytes_ptr.add(off1 + 2));
9391            let w0_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_0), vcreate_u64(u1_0)));
9392            let w1_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_0), vcreate_u64(u3_0)));
9393            let w0_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_1), vcreate_u64(u1_1)));
9394            let w1_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_1), vcreate_u64(u3_1)));
9395            for k in 0..2 {
9396                let d0 = sdot2!(w0_0, w1_0, xp[k].add(gi * GROUP_SIZE));
9397                let d1 = sdot2!(w0_1, w1_1, xp[k].add((gi + 1) * GROUP_SIZE));
9398                acc[k] += d0 as f32 * s0 + d1 as f32 * s1;
9399            }
9400            gi += 2;
9401        }
9402        if gi < gpr {
9403            let off = row_off + gi * TILE;
9404            let s = f16_to_f32(u16::from_le_bytes([
9405                *bytes_ptr.add(off),
9406                *bytes_ptr.add(off + 1),
9407            ]));
9408            let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
9409            let w0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0), vcreate_u64(u1)));
9410            let w1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2), vcreate_u64(u3)));
9411            for k in 0..2 {
9412                let d = sdot2!(w0, w1, xp[k].add(gi * GROUP_SIZE));
9413                acc[k] += d as f32 * s;
9414            }
9415        }
9416        acc
9417    }
9418}
9419
9420/// Fused Q1T pair matvec: ONE pass over the rows serves both
9421/// activation streams — on ARM the ternary register unpack happens
9422/// once per tile pair (`q1t_dot_row_sdot2`); elsewhere the second dot
9423/// rides the row's L1-warm tile bytes. Per stream the math matches
9424/// `q1t_matvec` exactly.
9425fn q1t_matvec2(
9426    bytes: &[u8],
9427    x1: &[f32],
9428    x2: &[f32],
9429    rows: usize,
9430    cols: usize,
9431    o1: &mut [f32],
9432    o2: &mut [f32],
9433    pool: Option<&Pool>,
9434) {
9435    debug_assert_eq!(o1.len(), rows);
9436    debug_assert_eq!(o2.len(), rows);
9437    const TILE: usize = cortiq_core::quant::Q1T_TILE;
9438    let gpr = cols / GROUP_SIZE;
9439    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
9440    let out1 = SendMut(o1.as_mut_ptr());
9441    let out2 = SendMut(o2.as_mut_ptr());
9442    if a8w8_enabled() {
9443        let a1 = split_act(x1);
9444        let a2 = split_act(x2);
9445        let (a1, a2) = (&a1, &a2);
9446        let run = move |start: usize, end: usize| {
9447            for r in start..end {
9448                #[cfg(target_arch = "aarch64")]
9449                // a8w8 on aarch64 ⇔ sdot_enabled(), so the kernel's
9450                // target features are present.
9451                let ds = unsafe { q1t_dot_row_sdot2(bytes, r, gpr, &a1.xq, &a2.xq) };
9452                #[cfg(not(target_arch = "aarch64"))]
9453                let ds = [
9454                    q1t_dot_row_i8(bytes, r, gpr, &a1.xq),
9455                    q1t_dot_row_i8(bytes, r, gpr, &a2.xq),
9456                ];
9457                let mut acc1 = ds[0] * a1.sx;
9458                for &(j, xv) in &a1.outliers {
9459                    acc1 += q1t_base_weight(bytes, r, gpr, j) * xv;
9460                }
9461                acc1 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x1);
9462                let mut acc2 = ds[1] * a2.sx;
9463                for &(j, xv) in &a2.outliers {
9464                    acc2 += q1t_base_weight(bytes, r, gpr, j) * xv;
9465                }
9466                acc2 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x2);
9467                // SAFETY: disjoint row ranges per worker.
9468                unsafe {
9469                    *out1.at(r) = acc1;
9470                    *out2.at(r) = acc2;
9471                }
9472            }
9473        };
9474        dispatch_rows(pool, rows, &run);
9475        return;
9476    }
9477    let run = move |start: usize, end: usize| {
9478        // Exact path (CMF_SDOT=0): unpack the sign LUT once per group,
9479        // dot both streams — same op order per stream as `q1t_matvec`.
9480        let mut sg = [0f32; GROUP_SIZE];
9481        for r in start..end {
9482            let mut acc1 = 0f32;
9483            let mut acc2 = 0f32;
9484            for g in 0..gpr {
9485                let off = (r * gpr + g) * TILE;
9486                let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
9487                let codes = &bytes[off + 2..off + TILE];
9488                for bi in 0..6 {
9489                    sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
9490                }
9491                let lut = &SIGN5[codes[6] as usize];
9492                sg[30] = lut[0];
9493                sg[31] = lut[1];
9494                let xg1 = &x1[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
9495                let xg2 = &x2[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
9496                let mut gsum1 = 0f32;
9497                for k in 0..GROUP_SIZE {
9498                    gsum1 += sg[k] * xg1[k];
9499                }
9500                acc1 += s * gsum1;
9501                let mut gsum2 = 0f32;
9502                for k in 0..GROUP_SIZE {
9503                    gsum2 += sg[k] * xg2[k];
9504                }
9505                acc2 += s * gsum2;
9506            }
9507            acc1 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x1);
9508            acc2 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x2);
9509            // SAFETY: disjoint row ranges per worker.
9510            unsafe {
9511                *out1.at(r) = acc1;
9512                *out2.at(r) = acc2;
9513            }
9514        }
9515    };
9516    dispatch_rows(pool, rows, &run);
9517}
9518
9519/// Ternary (q1t) matmat (prefill) — dequant each row once, dot the whole
9520/// batch against it (amortizes the per-row decode).
9521fn q1t_matmat(
9522    bytes: &[u8],
9523    xs: &[f32],
9524    b: usize,
9525    rows: usize,
9526    cols: usize,
9527    out: &mut [f32],
9528    pool: Option<&Pool>,
9529) {
9530    debug_assert_eq!(out.len(), b * rows);
9531    const TILE: usize = cortiq_core::quant::Q1T_TILE;
9532    let gpr = cols / GROUP_SIZE;
9533    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
9534    let out_addr = SendMut(out.as_mut_ptr());
9535    // int8 prefill (ARM SDOT / x86 AVX2): quantize the B inputs once, unpack
9536    // each weight row's signs to i8 ONCE, then int8-dot against every input —
9537    // the row sign-decode amortizes over the whole batch. CMF_SDOT=0 → f32.
9538    if a8w8_enabled() {
9539        let acts: Vec<SplitAct> = (0..b)
9540            .map(|bi| split_act(&xs[bi * cols..(bi + 1) * cols]))
9541            .collect();
9542        let acts = &acts;
9543        let run = move |start: usize, end: usize| {
9544            let mut sg = vec![0i8; cols + 8]; // row signs, i8 (+8 unpack slack)
9545            let mut sc = vec![0f32; gpr]; // per-group scales
9546            let mut accs = vec![0f32; b]; // per-batch accumulators, reused per row
9547            for r in start..end {
9548                for g in 0..gpr {
9549                    let off = (r * gpr + g) * TILE;
9550                    sc[g] = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
9551                    q1t_unpack_group_i8(
9552                        bytes.as_ptr().wrapping_add(off + 2),
9553                        &mut sg[g * GROUP_SIZE..],
9554                    );
9555                }
9556                for bi in 0..b {
9557                    let act = &acts[bi];
9558                    let mut isum = 0f32;
9559                    for g in 0..gpr {
9560                        let d = q1t_i8dot32(
9561                            sg.as_ptr().wrapping_add(g * GROUP_SIZE),
9562                            act.xq.as_ptr().wrapping_add(g * GROUP_SIZE),
9563                        );
9564                        isum += d as f32 * sc[g];
9565                    }
9566                    let mut acc = isum * act.sx;
9567                    for &(j, xv) in &act.outliers {
9568                        acc += q1t_base_weight(bytes, r, gpr, j) * xv;
9569                    }
9570                    accs[bi] = acc;
9571                }
9572                // Overlay ONCE per row for the whole batch: read each (col, val)
9573                // from mmap a single time (was b× — the re-read dominated prefill)
9574                // and fan it out over the batch via the cached inputs.
9575                if has_ov {
9576                    let (c0, c1) = (
9577                        q1t_rowptr(bytes, rp_off, r),
9578                        q1t_rowptr(bytes, rp_off, r + 1),
9579                    );
9580                    for p in c0..c1 {
9581                        let e = ent_off + p * 4;
9582                        let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
9583                        let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
9584                        for bi in 0..b {
9585                            accs[bi] += val * xs[bi * cols + col];
9586                        }
9587                    }
9588                }
9589                for bi in 0..b {
9590                    unsafe { *out_addr.at(bi * rows + r) = accs[bi] };
9591                }
9592            }
9593        };
9594        dispatch_rows(pool, rows, &run);
9595        return;
9596    }
9597    let run = move |start: usize, end: usize| {
9598        let mut buf = vec![0f32; cols];
9599        for r in start..end {
9600            q1t_dequant_row(bytes, r, gpr, rp_off, ent_off, has_ov, &mut buf);
9601            for bi in 0..b {
9602                let xr = &xs[bi * cols..(bi + 1) * cols];
9603                let mut acc = 0f32;
9604                for j in 0..cols {
9605                    acc += buf[j] * xr[j];
9606                }
9607                unsafe { *out_addr.at(bi * rows + r) = acc };
9608            }
9609        }
9610    };
9611    dispatch_rows(pool, rows, &run);
9612}
9613
9614/// q1-tile entry points of the binary drivers (the tests' handles; the
9615/// runtime goes through `QTensor::bin_w`).
9616#[cfg(test)]
9617fn q1_matvec(
9618    bytes: &[u8],
9619    x: &[f32],
9620    rows: usize,
9621    cols: usize,
9622    out: &mut [f32],
9623    pool: Option<&Pool>,
9624) {
9625    bin_matvec(BinW::Q1(bytes), x, rows, cols, out, pool)
9626}
9627
9628/// Binary-weight matvec (q1 or q1g): A8W8 rows with the shared group
9629/// sums, or the exact scalar walk under CMF_SDOT=0.
9630fn bin_matvec(w: BinW, x: &[f32], rows: usize, cols: usize, out: &mut [f32], pool: Option<&Pool>) {
9631    debug_assert_eq!(out.len(), rows);
9632    let gpr = cols / GROUP_SIZE;
9633    let out_addr = SendMut(out.as_mut_ptr());
9634    if a8w8_enabled() {
9635        let act = split_act(x);
9636        let gsum = q1_group_sums(&act.xq, gpr);
9637        let (act, gsum) = (&act, &gsum);
9638        let run =
9639            move |start: usize, end: usize| bin_range_a8w8(w, gpr, act, gsum, out_addr, start, end);
9640        dispatch_rows(pool, rows, &run);
9641        return;
9642    }
9643    let run = move |start: usize, end: usize| bin_range_f32(w, gpr, x, out_addr, start, end);
9644    dispatch_rows(pool, rows, &run);
9645}
9646
9647/// Fused two-input q1 matvec (weights read once per pair).
9648#[cfg(test)]
9649#[allow(clippy::too_many_arguments)]
9650fn q1_matvec2(
9651    bytes: &[u8],
9652    x1: &[f32],
9653    x2: &[f32],
9654    rows: usize,
9655    cols: usize,
9656    o1: &mut [f32],
9657    o2: &mut [f32],
9658    pool: Option<&Pool>,
9659) {
9660    bin_matvec2(BinW::Q1(bytes), x1, x2, rows, cols, o1, o2, pool)
9661}
9662
9663/// Fused two-input binary-weight matvec (q1 or q1g).
9664#[allow(clippy::too_many_arguments)]
9665fn bin_matvec2(
9666    w: BinW,
9667    x1: &[f32],
9668    x2: &[f32],
9669    rows: usize,
9670    cols: usize,
9671    o1: &mut [f32],
9672    o2: &mut [f32],
9673    pool: Option<&Pool>,
9674) {
9675    let gpr = cols / GROUP_SIZE;
9676    let p1 = SendMut(o1.as_mut_ptr());
9677    let p2 = SendMut(o2.as_mut_ptr());
9678    if a8w8_enabled() {
9679        let a1 = split_act(x1);
9680        let a2 = split_act(x2);
9681        let g1 = q1_group_sums(&a1.xq, gpr);
9682        let g2 = q1_group_sums(&a2.xq, gpr);
9683        let (a1, a2, g1, g2) = (&a1, &a2, &g1, &g2);
9684        let run = move |start: usize, end: usize| {
9685            for r in start..end {
9686                let mut v1 = w.dot_i8(r, gpr, &a1.xq, g1) * a1.sx;
9687                let mut v2 = w.dot_i8(r, gpr, &a2.xq, g2) * a2.sx;
9688                for &(j, xv) in &a1.outliers {
9689                    let (wv, s) = w.outlier(r, gpr, j);
9690                    v1 += wv * s * xv;
9691                }
9692                for &(j, xv) in &a2.outliers {
9693                    let (wv, s) = w.outlier(r, gpr, j);
9694                    v2 += wv * s * xv;
9695                }
9696                // SAFETY: disjoint row ranges per worker.
9697                unsafe {
9698                    *p1.at(r) = v1;
9699                    *p2.at(r) = v2;
9700                }
9701            }
9702        };
9703        dispatch_rows(pool, rows, &run);
9704        return;
9705    }
9706    let run = move |start: usize, end: usize| {
9707        for r in start..end {
9708            // SAFETY: disjoint row ranges per worker.
9709            unsafe {
9710                *p1.at(r) = w.row_exact(r, gpr, x1);
9711                *p2.at(r) = w.row_exact(r, gpr, x2);
9712            }
9713        }
9714    };
9715    dispatch_rows(pool, rows, &run);
9716}
9717
9718/// Batched q1 matmat: each row's tiles stream once per microbatch.
9719#[cfg(test)]
9720#[allow(clippy::too_many_arguments)]
9721fn q1_matmat(
9722    bytes: &[u8],
9723    xs_all: &[f32],
9724    b: usize,
9725    rows: usize,
9726    cols: usize,
9727    out: &mut [f32],
9728    pool: Option<&Pool>,
9729) {
9730    bin_matmat(BinW::Q1(bytes), xs_all, b, rows, cols, out, pool)
9731}
9732
9733/// Batched binary-weight matmat (q1 or q1g).
9734#[allow(clippy::too_many_arguments)]
9735fn bin_matmat(
9736    w: BinW,
9737    xs_all: &[f32],
9738    b: usize,
9739    rows: usize,
9740    cols: usize,
9741    out: &mut [f32],
9742    pool: Option<&Pool>,
9743) {
9744    debug_assert_eq!(out.len(), b * rows);
9745    let gpr = cols / GROUP_SIZE;
9746    let out_addr = SendMut(out.as_mut_ptr());
9747    if a8w8_enabled() {
9748        let acts: Vec<(SplitAct, Vec<i32>)> = (0..b)
9749            .map(|bi| {
9750                let act = split_act(&xs_all[bi * cols..(bi + 1) * cols]);
9751                let gsum = q1_group_sums(&act.xq, gpr);
9752                (act, gsum)
9753            })
9754            .collect();
9755        let acts = &acts;
9756        #[cfg(target_arch = "x86_64")]
9757        let blocked_ok = avx2_enabled() && blocked_enabled();
9758        #[cfg(target_arch = "aarch64")]
9759        let blocked_ok = sdot_enabled() && blocked_enabled();
9760        let run = move |start: usize, end: usize| {
9761            for r in start..end {
9762                let mut bi = 0usize;
9763                // Blocked 1×4: the unpacked bit mask serves four
9764                // activation streams per group.
9765                #[cfg(target_arch = "aarch64")]
9766                if blocked_ok {
9767                    while bi + 4 <= acts.len() {
9768                        let xs = [
9769                            acts[bi].0.xq.as_slice(),
9770                            acts[bi + 1].0.xq.as_slice(),
9771                            acts[bi + 2].0.xq.as_slice(),
9772                            acts[bi + 3].0.xq.as_slice(),
9773                        ];
9774                        let gs = [
9775                            acts[bi].1.as_slice(),
9776                            acts[bi + 1].1.as_slice(),
9777                            acts[bi + 2].1.as_slice(),
9778                            acts[bi + 3].1.as_slice(),
9779                        ];
9780                        // SAFETY: blocked_ok checked SDOT.
9781                        let d = unsafe { w.dot_1x4_sdot(r, gpr, xs, gs) };
9782                        for k in 0..4 {
9783                            let (act, _) = &acts[bi + k];
9784                            let mut acc = d[k] * act.sx;
9785                            for &(j, xv) in &act.outliers {
9786                                let (wv, sc) = w.outlier(r, gpr, j);
9787                                acc += wv * sc * xv;
9788                            }
9789                            // SAFETY: disjoint (bi, r) cells per worker.
9790                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
9791                        }
9792                        bi += 4;
9793                    }
9794                }
9795                #[cfg(target_arch = "x86_64")]
9796                if blocked_ok {
9797                    while bi + 4 <= acts.len() {
9798                        let xs = [
9799                            acts[bi].0.xq.as_slice(),
9800                            acts[bi + 1].0.xq.as_slice(),
9801                            acts[bi + 2].0.xq.as_slice(),
9802                            acts[bi + 3].0.xq.as_slice(),
9803                        ];
9804                        let gs = [
9805                            acts[bi].1.as_slice(),
9806                            acts[bi + 1].1.as_slice(),
9807                            acts[bi + 2].1.as_slice(),
9808                            acts[bi + 3].1.as_slice(),
9809                        ];
9810                        // SAFETY: blocked_ok checked AVX2.
9811                        let d = unsafe { w.dot_1x4_x86(r, gpr, xs, gs) };
9812                        for k in 0..4 {
9813                            let (act, _) = &acts[bi + k];
9814                            let mut acc = d[k] * act.sx;
9815                            for &(j, xv) in &act.outliers {
9816                                let (wv, sc) = w.outlier(r, gpr, j);
9817                                acc += wv * sc * xv;
9818                            }
9819                            // SAFETY: disjoint (bi, r) cells per worker.
9820                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
9821                        }
9822                        bi += 4;
9823                    }
9824                }
9825                while bi < acts.len() {
9826                    let (act, gsum) = &acts[bi];
9827                    let mut acc = w.dot_i8(r, gpr, &act.xq, gsum) * act.sx;
9828                    for &(j, xv) in &act.outliers {
9829                        let (wv, s) = w.outlier(r, gpr, j);
9830                        acc += wv * s * xv;
9831                    }
9832                    // SAFETY: disjoint (bi, r) cells per worker range.
9833                    unsafe { *out_addr.at(bi * rows + r) = acc };
9834                    bi += 1;
9835                }
9836            }
9837        };
9838        dispatch_rows(pool, rows, &run);
9839        return;
9840    }
9841    let run = move |start: usize, end: usize| {
9842        for r in start..end {
9843            for bi in 0..b {
9844                let x = &xs_all[bi * cols..(bi + 1) * cols];
9845                // SAFETY: disjoint (bi, r) cells per worker range.
9846                unsafe { *out_addr.at(bi * rows + r) = w.row_exact(r, gpr, x) };
9847            }
9848        }
9849    };
9850    dispatch_rows(pool, rows, &run);
9851}
9852
9853/// Fused q4_block matvec straight from the mapped bytes. SDOT path when
9854/// dotprod is available (port of vmfcore `dot_q4_block_sdot`, measured
9855/// +23% on q4 decode): nibbles → centered i8, int8×int8 `sdot` per
9856/// 32-group, exact outlier correction — the same A8W8 contract as q8.
9857/// `CMF_SDOT=0` keeps the exact scalar path.
9858fn q4matvec(
9859    bytes: &[u8],
9860    x: &[f32],
9861    rows: usize,
9862    cols: usize,
9863    out: &mut [f32],
9864    pool: Option<&Pool>,
9865) {
9866    debug_assert_eq!(out.len(), rows);
9867    let (packed, scales) = q4_split(bytes, rows, cols);
9868    let gpr = cols / GROUP_SIZE;
9869    let out_addr = SendMut(out.as_mut_ptr());
9870
9871    if a8w8_enabled() {
9872        let act = split_act(x);
9873        let run = move |start: usize, end: usize| {
9874            q4_range_a8w8(packed, scales, gpr, cols, &act, out_addr, start, end)
9875        };
9876        dispatch_rows(pool, rows, &run);
9877        return;
9878    }
9879
9880    let run =
9881        move |start: usize, end: usize| q4_range_f32(packed, scales, gpr, x, out_addr, start, end);
9882    dispatch_rows(pool, rows, &run);
9883}
9884
9885/// One q4 row via the A8W8 int8 path — SDOT on ARM, AVX2 maddubs on
9886/// x86 (scalar fallback is unreachable: callers gate on a8w8_enabled).
9887#[inline]
9888#[allow(unreachable_code)]
9889/// One UNPACKED q4 row (centered i8 in `buf`) against four activation
9890/// streams: the 32-byte weight chunk and its abs() load once per group,
9891/// the per-group f16 scale decodes once — four maddubs+reduce chains
9892/// instead of four full (load, abs, dot) rounds.
9893#[cfg(target_arch = "x86_64")]
9894#[target_feature(enable = "avx2")]
9895unsafe fn dot_q4b_row_1x4_avx2(
9896    buf: &[u8],
9897    scales: &[u8],
9898    g0: usize,
9899    gpr: usize,
9900    xs: [&[i8]; 4],
9901) -> [f32; 4] {
9902    // SAFETY: callers uphold buffer contracts (buf.len() == gpr·32).
9903    unsafe {
9904        use core::arch::x86_64::*;
9905        let ones = _mm256_set1_epi16(1);
9906        let mut acc = [0f32; 4];
9907        for gi in 0..gpr {
9908            let s = f16_to_f32(u16::from_le_bytes([
9909                scales[(g0 + gi) * 2],
9910                scales[(g0 + gi) * 2 + 1],
9911            ]));
9912            let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9913            let aw = _mm256_abs_epi8(w);
9914            for (k, xq) in xs.iter().enumerate() {
9915                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9916                let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
9917                let d = _mm256_madd_epi16(p16, ones);
9918                let hi128 = _mm256_extracti128_si256::<1>(d);
9919                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
9920                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
9921                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
9922                acc[k] += _mm_cvtsi128_si32(s32) as f32 * s;
9923            }
9924        }
9925        acc
9926    }
9927}
9928
9929/// VNNI twin of `dot_q4b_row_1x4_avx2` (see `dpbusd_hsum`).
9930#[cfg(target_arch = "x86_64")]
9931#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
9932unsafe fn dot_q4b_row_1x4_vnni(
9933    buf: &[u8],
9934    scales: &[u8],
9935    g0: usize,
9936    gpr: usize,
9937    xs: [&[i8]; 4],
9938) -> [f32; 4] {
9939    // SAFETY: callers uphold buffer contracts (buf.len() == gpr·32).
9940    unsafe {
9941        use core::arch::x86_64::*;
9942        let mut acc = [0f32; 4];
9943        for gi in 0..gpr {
9944            let s = f16_to_f32(u16::from_le_bytes([
9945                scales[(g0 + gi) * 2],
9946                scales[(g0 + gi) * 2 + 1],
9947            ]));
9948            let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9949            let aw = _mm256_abs_epi8(w);
9950            for (k, xq) in xs.iter().enumerate() {
9951                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9952                let d = dpbusd_hsum(aw, _mm256_sign_epi8(x, w));
9953                acc[k] += d as f32 * s;
9954            }
9955        }
9956        acc
9957    }
9958}
9959
9960/// The vbit flavor of the blocked 1×4: the per-activation A8W8 scale
9961/// folds in PER GROUP as `(d·sx)·s` — bit-matching the single-matvec
9962/// accumulation order (the q4_block flavor applies sx once at the end,
9963/// matching ITS single path; the two conventions are historical and
9964/// each blocked leg must mirror its own).
9965#[cfg(target_arch = "x86_64")]
9966#[target_feature(enable = "avx2")]
9967unsafe fn dot_q4b_row_1x4_sx_avx2(
9968    buf: &[u8],
9969    scales: &[u8],
9970    g0: usize,
9971    gpr: usize,
9972    xs: [&[i8]; 4],
9973    sxs: [f32; 4],
9974) -> [f32; 4] {
9975    // SAFETY: callers uphold buffer contracts (buf.len() == gpr·32).
9976    unsafe {
9977        use core::arch::x86_64::*;
9978        let ones = _mm256_set1_epi16(1);
9979        let mut acc = [0f32; 4];
9980        for gi in 0..gpr {
9981            let s = f16_to_f32(u16::from_le_bytes([
9982                scales[(g0 + gi) * 2],
9983                scales[(g0 + gi) * 2 + 1],
9984            ]));
9985            let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9986            let aw = _mm256_abs_epi8(w);
9987            for (k, xq) in xs.iter().enumerate() {
9988                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9989                let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
9990                let d = _mm256_madd_epi16(p16, ones);
9991                let hi128 = _mm256_extracti128_si256::<1>(d);
9992                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
9993                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
9994                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
9995                acc[k] += (_mm_cvtsi128_si32(s32) as f32 * sxs[k]) * s;
9996            }
9997        }
9998        acc
9999    }
10000}
10001
10002/// VNNI twin of `dot_q4b_row_1x4_sx_avx2` (see `dpbusd_hsum`; the
10003/// per-group `(d·sx)·s` fold mirrors the vbit single path).
10004#[cfg(target_arch = "x86_64")]
10005#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
10006unsafe fn dot_q4b_row_1x4_sx_vnni(
10007    buf: &[u8],
10008    scales: &[u8],
10009    g0: usize,
10010    gpr: usize,
10011    xs: [&[i8]; 4],
10012    sxs: [f32; 4],
10013) -> [f32; 4] {
10014    // SAFETY: callers uphold buffer contracts (buf.len() == gpr·32).
10015    unsafe {
10016        use core::arch::x86_64::*;
10017        let mut acc = [0f32; 4];
10018        for gi in 0..gpr {
10019            let s = f16_to_f32(u16::from_le_bytes([
10020                scales[(g0 + gi) * 2],
10021                scales[(g0 + gi) * 2 + 1],
10022            ]));
10023            let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
10024            let aw = _mm256_abs_epi8(w);
10025            for (k, xq) in xs.iter().enumerate() {
10026                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
10027                let d = dpbusd_hsum(aw, _mm256_sign_epi8(x, w));
10028                acc[k] += (d as f32 * sxs[k]) * s;
10029            }
10030        }
10031        acc
10032    }
10033}
10034
10035#[allow(unreachable_code)]
10036fn dot_q4_row_i8(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
10037    #[cfg(target_arch = "aarch64")]
10038    unsafe {
10039        return dot_q4_row_sdot(packed, scales, g0, gpr, xq);
10040    }
10041    #[cfg(target_arch = "x86_64")]
10042    unsafe {
10043        return dot_q4_row_avx2(packed, scales, g0, gpr, xq);
10044    }
10045    let mut acc = 0f32;
10046    for gi in 0..gpr {
10047        let g = g0 + gi;
10048        let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
10049        let mut d = 0i32;
10050        for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
10051            d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
10052                + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
10053        }
10054        acc += d as f32 * s;
10055    }
10056    acc
10057}
10058
10059/// Two-activation q4 row via the A8W8 int8 path (see `dot_q4_row_i8`).
10060#[inline]
10061#[allow(unreachable_code)]
10062fn dot_q4_row_i8_2(
10063    packed: &[u8],
10064    scales: &[u8],
10065    g0: usize,
10066    gpr: usize,
10067    xq1: &[i8],
10068    xq2: &[i8],
10069) -> (f32, f32) {
10070    #[cfg(target_arch = "aarch64")]
10071    unsafe {
10072        return dot_q4_row_sdot2(packed, scales, g0, gpr, xq1, xq2);
10073    }
10074    #[cfg(target_arch = "x86_64")]
10075    unsafe {
10076        return dot_q4_row_avx2_2(packed, scales, g0, gpr, xq1, xq2);
10077    }
10078    (
10079        dot_q4_row_i8(packed, scales, g0, gpr, xq1),
10080        dot_q4_row_i8(packed, scales, g0, gpr, xq2),
10081    )
10082}
10083
10084/// One q4 row range via SDOT (kernel body of `q4matvec`, extracted so
10085/// multi-matrix jobs can drive it for several tensors in one dispatch).
10086#[allow(clippy::too_many_arguments)]
10087fn q4_range_a8w8(
10088    packed: &[u8],
10089    scales: &[u8],
10090    gpr: usize,
10091    cols: usize,
10092    act: &SplitAct,
10093    out: SendMut,
10094    start: usize,
10095    end: usize,
10096) {
10097    for r in start..end {
10098        let mut acc = dot_q4_row_i8(packed, scales, r * gpr, gpr, &act.xq) * act.sx;
10099        // xq is zeroed at outlier slots — add the exact terms.
10100        for &(j, xv) in &act.outliers {
10101            let flat = r * cols + j;
10102            let byte = packed[flat / 2];
10103            let nib = if flat & 1 == 0 {
10104                byte & 0x0F
10105            } else {
10106                byte >> 4
10107            };
10108            let s = f16_to_f32(u16::from_le_bytes([
10109                scales[(flat / GROUP_SIZE) * 2],
10110                scales[(flat / GROUP_SIZE) * 2 + 1],
10111            ]));
10112            acc += ((nib as i32 - 8) as f32) * s * xv;
10113        }
10114        // SAFETY: disjoint row ranges per worker.
10115        unsafe { *out.at(r) = acc };
10116    }
10117}
10118
10119/// Two-input q4 row range via the A8W8 int8 path — kernel body of
10120/// `q4matvec2`, extracted for pair multi-matrix jobs.
10121#[allow(clippy::too_many_arguments)]
10122fn q4_range2_a8w8(
10123    packed: &[u8],
10124    scales: &[u8],
10125    gpr: usize,
10126    cols: usize,
10127    a1: &SplitAct,
10128    a2: &SplitAct,
10129    p1: SendMut,
10130    p2: SendMut,
10131    start: usize,
10132    end: usize,
10133) {
10134    for r in start..end {
10135        let (s1, s2) = dot_q4_row_i8_2(packed, scales, r * gpr, gpr, &a1.xq, &a2.xq);
10136        let mut acc1 = s1 * a1.sx;
10137        let mut acc2 = s2 * a2.sx;
10138        // xq is zeroed at outlier slots — add the exact terms.
10139        let fix = |outliers: &[(usize, f32)], acc: &mut f32| {
10140            for &(j, xv) in outliers {
10141                let flat = r * cols + j;
10142                let byte = packed[flat / 2];
10143                let nib = if flat & 1 == 0 {
10144                    byte & 0x0F
10145                } else {
10146                    byte >> 4
10147                };
10148                let s = f16_to_f32(u16::from_le_bytes([
10149                    scales[(flat / GROUP_SIZE) * 2],
10150                    scales[(flat / GROUP_SIZE) * 2 + 1],
10151                ]));
10152                *acc += ((nib as i32 - 8) as f32) * s * xv;
10153            }
10154        };
10155        fix(&a1.outliers, &mut acc1);
10156        fix(&a2.outliers, &mut acc2);
10157        // SAFETY: disjoint row ranges per worker.
10158        unsafe {
10159            *p1.at(r) = acc1;
10160            *p2.at(r) = acc2;
10161        }
10162    }
10163}
10164
10165/// Exact scalar q4 row range (same extraction, non-SDOT path).
10166fn q4_range_f32(
10167    packed: &[u8],
10168    scales: &[u8],
10169    gpr: usize,
10170    x: &[f32],
10171    out: SendMut,
10172    start: usize,
10173    end: usize,
10174) {
10175    for r in start..end {
10176        let mut acc = 0f32;
10177        for gi in 0..gpr {
10178            let g = r * gpr + gi;
10179            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
10180            let pk = &packed[g * 16..(g + 1) * 16];
10181            let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
10182            let mut ga = 0f32;
10183            for (k, &b) in pk.iter().enumerate() {
10184                ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
10185                    + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
10186            }
10187            acc += ga * s;
10188        }
10189        // SAFETY: disjoint row ranges per worker.
10190        unsafe { *out.at(r) = acc };
10191    }
10192}
10193
10194/// Fused two-input q4 matvec: nibbles are unpacked ONCE per group and
10195/// dotted against both activations (was: two full matvecs — double
10196/// weight traffic). Per-lane math matches `q4matvec` exactly.
10197#[allow(clippy::too_many_arguments)]
10198fn q4matvec2(
10199    bytes: &[u8],
10200    x1: &[f32],
10201    x2: &[f32],
10202    rows: usize,
10203    cols: usize,
10204    o1: &mut [f32],
10205    o2: &mut [f32],
10206    pool: Option<&Pool>,
10207) {
10208    debug_assert_eq!(o1.len(), rows);
10209    debug_assert_eq!(o2.len(), rows);
10210    let (packed, scales) = q4_split(bytes, rows, cols);
10211    let gpr = cols / GROUP_SIZE;
10212
10213    if a8w8_enabled() {
10214        let a1 = split_act(x1);
10215        let a2 = split_act(x2);
10216        let p1 = SendMut(o1.as_mut_ptr());
10217        let p2 = SendMut(o2.as_mut_ptr());
10218        let run = move |start: usize, end: usize| {
10219            q4_range2_a8w8(packed, scales, gpr, cols, &a1, &a2, p1, p2, start, end)
10220        };
10221        dispatch_rows(pool, rows, &run);
10222        return;
10223    }
10224
10225    let p1 = SendMut(o1.as_mut_ptr());
10226    let p2 = SendMut(o2.as_mut_ptr());
10227    let run = move |start: usize, end: usize| {
10228        q4_range2_f32(packed, scales, gpr, x1, x2, p1, p2, start, end)
10229    };
10230    dispatch_rows(pool, rows, &run);
10231}
10232
10233/// Two-input exact scalar q4 row range (same extraction).
10234#[allow(clippy::too_many_arguments)]
10235fn q4_range2_f32(
10236    packed: &[u8],
10237    scales: &[u8],
10238    gpr: usize,
10239    x1: &[f32],
10240    x2: &[f32],
10241    p1: SendMut,
10242    p2: SendMut,
10243    start: usize,
10244    end: usize,
10245) {
10246    for r in start..end {
10247        let (mut acc1, mut acc2) = (0f32, 0f32);
10248        for gi in 0..gpr {
10249            let g = r * gpr + gi;
10250            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
10251            let pk = &packed[g * 16..(g + 1) * 16];
10252            let x1g = &x1[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
10253            let x2g = &x2[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
10254            let (mut g1, mut g2) = (0f32, 0f32);
10255            for (k, &b) in pk.iter().enumerate() {
10256                let wl = (b & 0x0F) as f32 - 8.0;
10257                let wh = ((b >> 4) & 0x0F) as f32 - 8.0;
10258                g1 += wl * x1g[k * 2] + wh * x1g[k * 2 + 1];
10259                g2 += wl * x2g[k * 2] + wh * x2g[k * 2 + 1];
10260            }
10261            acc1 += g1 * s;
10262            acc2 += g2 * s;
10263        }
10264        // SAFETY: disjoint row ranges per worker.
10265        unsafe {
10266            *p1.at(r) = acc1;
10267            *p2.at(r) = acc2;
10268        }
10269    }
10270}
10271
10272thread_local! {
10273    /// Per-worker decoded-row scratch for the batched q4/vbit kernels
10274    /// (centered i8 for SDOT, f32 for the exact/scalar paths).
10275    static ROW_I8: std::cell::RefCell<Vec<u8>> = const { std::cell::RefCell::new(Vec::new()) };
10276    static ROW_F32: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
10277}
10278
10279/// Batched q4 matmat: each weight row is unpacked from the mmap ONCE
10280/// and dotted against ALL b activations (prefill used to fall back to b
10281/// full matvecs — b× weight traffic and b× nibble decode). Per-position
10282/// math matches `q4matvec` exactly: same group order, same accumulation.
10283/// `out` is row-major [b, rows] like `qmatmat`.
10284#[allow(clippy::too_many_arguments)]
10285fn q4matmat(
10286    bytes: &[u8],
10287    xs_all: &[f32],
10288    b: usize,
10289    rows: usize,
10290    cols: usize,
10291    out: &mut [f32],
10292    pool: Option<&Pool>,
10293) {
10294    debug_assert_eq!(xs_all.len(), b * cols);
10295    debug_assert_eq!(out.len(), b * rows);
10296    let (packed, scales) = q4_split(bytes, rows, cols);
10297    let gpr = cols / GROUP_SIZE;
10298    let gscale = |g: usize| f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
10299
10300    if a8w8_enabled() {
10301        let acts: Vec<SplitAct> = (0..b)
10302            .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
10303            .collect();
10304        let acts = &acts;
10305        let out_addr = SendMut(out.as_mut_ptr());
10306        let run = move |start: usize, end: usize| {
10307            ROW_I8.with(|rb| {
10308                let mut buf = rb.borrow_mut();
10309                buf.resize(cols, 0);
10310                for r in start..end {
10311                    // Unpack the row's nibbles to centered i8 once
10312                    // (element 2k = low nibble, 2k+1 = high — flat order,
10313                    // same as dot_q4_row_sdot's zip).
10314                    for gi in 0..gpr {
10315                        let g = r * gpr + gi;
10316                        for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
10317                            buf[gi * GROUP_SIZE + k * 2] = ((bt & 0x0F) as i32 - 8) as i8 as u8;
10318                            buf[gi * GROUP_SIZE + k * 2 + 1] =
10319                                (((bt >> 4) & 0x0F) as i32 - 8) as i8 as u8;
10320                        }
10321                    }
10322                    let mut bi = 0usize;
10323                    #[cfg(target_arch = "x86_64")]
10324                    if avx2_enabled() && blocked_enabled() {
10325                        while bi + 4 <= acts.len() {
10326                            let xs = [
10327                                acts[bi].xq.as_slice(),
10328                                acts[bi + 1].xq.as_slice(),
10329                                acts[bi + 2].xq.as_slice(),
10330                                acts[bi + 3].xq.as_slice(),
10331                            ];
10332                            let d = unsafe {
10333                                if vnni_tiles_enabled() {
10334                                    dot_q4b_row_1x4_vnni(&buf, scales, r * gpr, gpr, xs)
10335                                } else {
10336                                    dot_q4b_row_1x4_avx2(&buf, scales, r * gpr, gpr, xs)
10337                                }
10338                            };
10339                            for k in 0..4 {
10340                                let act = &acts[bi + k];
10341                                let mut acc = d[k] * act.sx;
10342                                for &(j, xv) in &act.outliers {
10343                                    acc += (buf[j] as i8) as f32
10344                                        * gscale((r * cols + j) / GROUP_SIZE)
10345                                        * xv;
10346                                }
10347                                // SAFETY: disjoint (bi, r) cells per worker.
10348                                unsafe { *out_addr.at((bi + k) * rows + r) = acc };
10349                            }
10350                            bi += 4;
10351                        }
10352                    }
10353                    while bi < acts.len() {
10354                        let act = &acts[bi];
10355                        let mut acc = 0f32;
10356                        for gi in 0..gpr {
10357                            let d = dot_i8_i8(
10358                                &buf[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
10359                                &act.xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
10360                            );
10361                            acc += d as f32 * gscale(r * gpr + gi);
10362                        }
10363                        acc *= act.sx;
10364                        // xq is zeroed at outlier slots — exact terms.
10365                        for &(j, xv) in &act.outliers {
10366                            acc += (buf[j] as i8) as f32 * gscale((r * cols + j) / GROUP_SIZE) * xv;
10367                        }
10368                        // SAFETY: disjoint (bi, r) cells per worker row range.
10369                        unsafe { *out_addr.at(bi * rows + r) = acc };
10370                        bi += 1;
10371                    }
10372                }
10373            })
10374        };
10375        dispatch_rows(pool, rows, &run);
10376        return;
10377    }
10378
10379    let out_addr = SendMut(out.as_mut_ptr());
10380    let run = move |start: usize, end: usize| {
10381        ROW_F32.with(|rb| {
10382            let mut buf = rb.borrow_mut();
10383            buf.resize(cols, 0.0);
10384            for r in start..end {
10385                // Decode raw (nib − 8) values once; scales stay per-group
10386                // so the accumulation order matches q4matvec bit-for-bit.
10387                for gi in 0..gpr {
10388                    let g = r * gpr + gi;
10389                    for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
10390                        buf[gi * GROUP_SIZE + k * 2] = (bt & 0x0F) as f32 - 8.0;
10391                        buf[gi * GROUP_SIZE + k * 2 + 1] = ((bt >> 4) & 0x0F) as f32 - 8.0;
10392                    }
10393                }
10394                for bi in 0..b {
10395                    let x = &xs_all[bi * cols..(bi + 1) * cols];
10396                    let mut acc = 0f32;
10397                    for gi in 0..gpr {
10398                        let mut ga = 0f32;
10399                        // Pairwise (lo + hi) addition, matching
10400                        // q4matvec's `ga += lo·x + hi·x` shape exactly —
10401                        // a flat one-per-element loop rounds differently
10402                        // and broke bit-parity on the scalar (x86) path.
10403                        for k in 0..GROUP_SIZE / 2 {
10404                            let e = gi * GROUP_SIZE + k * 2;
10405                            ga += buf[e] * x[e] + buf[e + 1] * x[e + 1];
10406                        }
10407                        acc += ga * gscale(r * gpr + gi);
10408                    }
10409                    // SAFETY: disjoint (bi, r) cells per worker row range.
10410                    unsafe { *out_addr.at(bi * rows + r) = acc };
10411                }
10412            }
10413        })
10414    };
10415    dispatch_rows(pool, rows, &run);
10416}
10417
10418/// Batched vbit matmat: each variable-bit row is decoded from the mmap
10419/// ONCE for the whole microbatch. Same per-position math as
10420/// `vbitmatvec` (SDOT A8W8 with exact outliers / exact f32 for b=8 rows
10421/// and the scalar path).
10422#[allow(clippy::too_many_arguments)]
10423fn vbitmatmat(
10424    bytes: &[u8],
10425    offsets: &[usize],
10426    xs_all: &[f32],
10427    b: usize,
10428    rows: usize,
10429    cols: usize,
10430    out: &mut [f32],
10431    pool: Option<&Pool>,
10432) {
10433    debug_assert_eq!(xs_all.len(), b * cols);
10434    debug_assert_eq!(out.len(), b * rows);
10435    debug_assert_eq!(offsets.len(), rows + 1);
10436    let ng = cols / GROUP_SIZE;
10437    let bits = &bytes[..rows];
10438    let sc_off = rows;
10439    let gscale = |r: usize, g: usize| {
10440        let so = (r * ng + g) * 2;
10441        f16_to_f32(u16::from_le_bytes([
10442            bytes[sc_off + so],
10443            bytes[sc_off + so + 1],
10444        ]))
10445    };
10446
10447    // Decode row r's raw (u − L) values into `dst` (f32, unscaled).
10448    let decode_f32 = |r: usize, dst: &mut [f32]| {
10449        let bw = bits[r] as usize;
10450        let l = ((1i32 << (bw - 1)) - 1) as f32;
10451        let data = &bytes[offsets[r]..offsets[r + 1]];
10452        let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
10453        for d in dst.iter_mut() {
10454            while nbits < bw {
10455                acc = (acc << 8) | data[idx] as u64;
10456                idx += 1;
10457                nbits += 8;
10458            }
10459            let u = ((acc >> (nbits - bw)) & ((1u64 << bw) - 1)) as f32;
10460            nbits -= bw;
10461            *d = u - l;
10462        }
10463    };
10464
10465    if a8w8_enabled() {
10466        let acts: Vec<SplitAct> = (0..b)
10467            .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
10468            .collect();
10469        let acts = &acts;
10470        let out_addr = SendMut(out.as_mut_ptr());
10471        let run = move |start: usize, end: usize| {
10472            for r in start..end {
10473                let bw = bits[r] as usize;
10474                if bw == 8 {
10475                    // u−L reaches 128 → no i8 path; decode once, exact
10476                    // f32 dots for every position (same as vbitmatvec).
10477                    ROW_F32.with(|rb| {
10478                        let mut buf = rb.borrow_mut();
10479                        buf.resize(cols, 0.0);
10480                        decode_f32(r, &mut buf);
10481                        for bi in 0..b {
10482                            let x = &xs_all[bi * cols..(bi + 1) * cols];
10483                            let mut dot = 0f32;
10484                            for g in 0..ng {
10485                                let mut gd = 0f32;
10486                                for k in 0..GROUP_SIZE {
10487                                    gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
10488                                }
10489                                dot += gd * gscale(r, g);
10490                            }
10491                            // SAFETY: disjoint (bi, r) cells per worker range.
10492                            unsafe { *out_addr.at(bi * rows + r) = dot };
10493                        }
10494                    });
10495                    continue;
10496                }
10497                let l = (1i32 << (bw - 1)) - 1;
10498                let data = &bytes[offsets[r]..offsets[r + 1]];
10499                ROW_I8.with(|rb| {
10500                    let mut buf = rb.borrow_mut();
10501                    buf.resize(cols, 0);
10502                    #[inline(always)]
10503                    fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
10504                        for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
10505                            let u = unpack8::<B>(&data[blk * B..]);
10506                            for k in 0..8 {
10507                                chunk[k] = (u[k] - l) as i8 as u8;
10508                            }
10509                        }
10510                    }
10511                    match bw {
10512                        3 => fill::<3>(data, l, &mut buf),
10513                        4 => vbit_fill4(data, &mut buf),
10514                        5 => fill::<5>(data, l, &mut buf),
10515                        6 => fill::<6>(data, l, &mut buf),
10516                        _ => unreachable!("vbit bit-width {bw} (validated at load)"),
10517                    }
10518                    let mut bi = 0usize;
10519                    // The vbit scale table shares q4_block's layout
10520                    // (contiguous f16 per (row·ng + g)), so the same
10521                    // blocked 1×4 kernel serves the decoded row.
10522                    #[cfg(target_arch = "x86_64")]
10523                    if avx2_enabled() && blocked_enabled() {
10524                        while bi + 4 <= acts.len() {
10525                            let xs = [
10526                                acts[bi].xq.as_slice(),
10527                                acts[bi + 1].xq.as_slice(),
10528                                acts[bi + 2].xq.as_slice(),
10529                                acts[bi + 3].xq.as_slice(),
10530                            ];
10531                            let sxs = [
10532                                acts[bi].sx,
10533                                acts[bi + 1].sx,
10534                                acts[bi + 2].sx,
10535                                acts[bi + 3].sx,
10536                            ];
10537                            let d = unsafe {
10538                                if vnni_tiles_enabled() {
10539                                    dot_q4b_row_1x4_sx_vnni(
10540                                        &buf,
10541                                        &bytes[sc_off..],
10542                                        r * ng,
10543                                        ng,
10544                                        xs,
10545                                        sxs,
10546                                    )
10547                                } else {
10548                                    dot_q4b_row_1x4_sx_avx2(
10549                                        &buf,
10550                                        &bytes[sc_off..],
10551                                        r * ng,
10552                                        ng,
10553                                        xs,
10554                                        sxs,
10555                                    )
10556                                }
10557                            };
10558                            for k in 0..4 {
10559                                let act = &acts[bi + k];
10560                                let mut dot = d[k];
10561                                for &(j, xv) in &act.outliers {
10562                                    dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
10563                                }
10564                                // SAFETY: disjoint (bi, r) cells per worker.
10565                                unsafe { *out_addr.at((bi + k) * rows + r) = dot };
10566                            }
10567                            bi += 4;
10568                        }
10569                    }
10570                    while bi < acts.len() {
10571                        let act = &acts[bi];
10572                        let mut dot = 0f32;
10573                        for g in 0..ng {
10574                            let d = dot_i8_i8(
10575                                &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
10576                                &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
10577                            ) as f32
10578                                * act.sx;
10579                            dot += d * gscale(r, g);
10580                        }
10581                        for &(j, xv) in &act.outliers {
10582                            dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
10583                        }
10584                        // SAFETY: disjoint (bi, r) cells per worker range.
10585                        unsafe { *out_addr.at(bi * rows + r) = dot };
10586                        bi += 1;
10587                    }
10588                });
10589            }
10590        };
10591        dispatch_rows(pool, rows, &run);
10592        return;
10593    }
10594
10595    let out_addr = SendMut(out.as_mut_ptr());
10596    let run = move |start: usize, end: usize| {
10597        ROW_F32.with(|rb| {
10598            let mut buf = rb.borrow_mut();
10599            buf.resize(cols, 0.0);
10600            for r in start..end {
10601                decode_f32(r, &mut buf);
10602                for bi in 0..b {
10603                    let x = &xs_all[bi * cols..(bi + 1) * cols];
10604                    let mut dot = 0f32;
10605                    for g in 0..ng {
10606                        let mut gd = 0f32;
10607                        for k in 0..GROUP_SIZE {
10608                            gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
10609                        }
10610                        dot += gd * gscale(r, g);
10611                    }
10612                    // SAFETY: disjoint (bi, r) cells per worker range.
10613                    unsafe { *out_addr.at(bi * rows + r) = dot };
10614                }
10615            }
10616        })
10617    };
10618    dispatch_rows(pool, rows, &run);
10619}
10620
10621/// Build a GPU batch job for a q8-family mapped tensor (primary
10622/// shard): prescaled input + directory coordinates. None → not
10623/// GPU-eligible, caller stays on the CPU.
10624pub(crate) fn gpu_batch_job<'a>(
10625    t: &'a QTensor,
10626    x: &[f32],
10627) -> Option<(std::sync::Arc<CmfModel>, crate::gpu::BatchJob<'a>)> {
10628    match t {
10629        QTensor::Mapped {
10630            model,
10631            idx,
10632            dtype: dt @ (TensorDtype::Q8Row | TensorDtype::Q8_2f),
10633            rows,
10634            cols,
10635            row_scale,
10636            col_field,
10637            ..
10638        } => Some((
10639            model.clone(),
10640            crate::gpu::BatchJob {
10641                idx: *idx,
10642                rows: *rows,
10643                cols: *cols,
10644                row_scale,
10645                xs: prescale(x, col_field, *dt).into_owned(),
10646                layout: crate::gpu::BatchLayout::Q8,
10647            },
10648        )),
10649        // q1: raw f32 activations, tile-embedded scales.
10650        QTensor::Mapped {
10651            model,
10652            idx,
10653            dtype: TensorDtype::Q1,
10654            rows,
10655            cols,
10656            ..
10657        } => Some((
10658            model.clone(),
10659            crate::gpu::BatchJob {
10660                idx: *idx,
10661                rows: *rows,
10662                cols: *cols,
10663                row_scale: &[],
10664                xs: x.to_vec(),
10665                layout: crate::gpu::BatchLayout::Q1,
10666            },
10667        )),
10668        // q1g: same raw-f32 contract, sign + scale planes.
10669        QTensor::Mapped {
10670            model,
10671            idx,
10672            dtype: TensorDtype::Q1G,
10673            rows,
10674            cols,
10675            ..
10676        } => Some((
10677            model.clone(),
10678            crate::gpu::BatchJob {
10679                idx: *idx,
10680                rows: *rows,
10681                cols: *cols,
10682                row_scale: &[],
10683                xs: x.to_vec(),
10684                layout: crate::gpu::BatchLayout::Q1g,
10685            },
10686        )),
10687        // q4_tiled / q4tp: raw f32 activations; the scales live in the
10688        // payload (inline tiles / row ladder), so row_scale stays empty.
10689        // The GDN projection batch already runs these layouts on Metal —
10690        // this arm lets the attention QKV batch reach the same kernels.
10691        QTensor::Mapped {
10692            model,
10693            idx,
10694            dtype: dt @ (TensorDtype::Q4Tiled | TensorDtype::Q4TiledP),
10695            rows,
10696            cols,
10697            ..
10698        } => Some((
10699            model.clone(),
10700            crate::gpu::BatchJob {
10701                idx: *idx,
10702                rows: *rows,
10703                cols: *cols,
10704                row_scale: &[],
10705                xs: x.to_vec(),
10706                layout: if *dt == TensorDtype::Q4Tiled {
10707                    crate::gpu::BatchLayout::Q4t
10708                } else {
10709                    crate::gpu::BatchLayout::Q4tp
10710                },
10711            },
10712        )),
10713        _ => None,
10714    }
10715}
10716
10717thread_local! {
10718    static PRESCALE_BUF1: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
10719    static PRESCALE_BUF2: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
10720}
10721
10722pub(crate) fn prescale<'a>(
10723    x: &'a [f32],
10724    col_field: &[f32],
10725    dtype: TensorDtype,
10726) -> std::borrow::Cow<'a, [f32]> {
10727    if dtype == TensorDtype::Q8_2f {
10728        x.iter().zip(col_field).map(|(a, c)| a * c).collect()
10729    } else {
10730        std::borrow::Cow::Borrowed(x)
10731    }
10732}
10733
10734/// θ col-field fold for q8_2f activations. Borrowed pass-through for
10735/// every other dtype, using thread-local buffers to eliminate per-matvec allocations.
10736pub(crate) fn prescale_with<R, F: FnOnce(&[f32]) -> R>(
10737    x: &[f32],
10738    col_field: &[f32],
10739    dtype: TensorDtype,
10740    buf_id: u8,
10741    f: F,
10742) -> R {
10743    if dtype == TensorDtype::Q8_2f {
10744        if buf_id == 1 {
10745            PRESCALE_BUF1.with(|b| {
10746                let mut buf = b.borrow_mut();
10747                buf.clear();
10748                buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
10749                f(&buf)
10750            })
10751        } else {
10752            PRESCALE_BUF2.with(|b| {
10753                let mut buf = b.borrow_mut();
10754                buf.clear();
10755                buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
10756                f(&buf)
10757            })
10758        }
10759    } else {
10760        f(x)
10761    }
10762}
10763
10764// ───────────────────── x86-64 AVX2 kernels (roadmap этап 2) ─────────────────────
10765
10766/// AVX2+FMA available? Default ON when the CPU supports both;
10767/// `CMF_AVX2=0` disables (falls back to the autovectorized loops).
10768#[cfg(target_arch = "x86_64")]
10769pub(crate) fn avx2_enabled() -> bool {
10770    use std::sync::OnceLock;
10771    static ON: OnceLock<bool> = OnceLock::new();
10772    *ON.get_or_init(|| {
10773        std::env::var("CMF_AVX2").map(|v| v != "0").unwrap_or(true)
10774            && std::arch::is_x86_feature_detected!("avx2")
10775            && std::arch::is_x86_feature_detected!("fma")
10776    })
10777}
10778
10779/// AVX2 A8W8 allowed? The quantized-activation contract is switched by
10780/// the SAME env as the ARM SDOT path: `CMF_SDOT=0` keeps exact kernels
10781/// (the golden-parity exact gate relies on it) — AVX2 f32 kernels stay
10782/// active either way, they are exact (regrouped sums only).
10783#[cfg(target_arch = "x86_64")]
10784fn avx2_a8w8_enabled() -> bool {
10785    if FLOAT_ACTIVATIONS.get() {
10786        return false;
10787    }
10788    use std::sync::OnceLock;
10789    static ON: OnceLock<bool> = OnceLock::new();
10790    *ON.get_or_init(|| {
10791        avx2_enabled() && std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true)
10792    })
10793}
10794
10795thread_local! {
10796    static FULL_GPU_Q8: std::cell::Cell<bool> = const { std::cell::Cell::new(false) };
10797}
10798
10799/// Match the graph's full-device q8 projection precision on MiMo's host
10800/// tail. Does not enable the GPU or bypass a CPU-only/device-refusal gate.
10801pub(crate) fn enter_full_gpu_q8_scope() -> impl Drop {
10802    struct Restore(bool, std::marker::PhantomData<std::rc::Rc<()>>);
10803    impl Drop for Restore {
10804        fn drop(&mut self) {
10805            FULL_GPU_Q8.set(self.0);
10806        }
10807    }
10808    Restore(FULL_GPU_Q8.replace(true), std::marker::PhantomData)
10809}
10810
10811// Dynamic MiMo experts must not change activation precision when a cache
10812// fill moves them from CPU to GPU. Thread-local: only the cold-expert
10813// dispatch selects float kernels; concurrent pipelines keep their policy.
10814thread_local! {
10815    static FLOAT_ACTIVATIONS: std::cell::Cell<bool> = const { std::cell::Cell::new(false) };
10816}
10817
10818/// Flint's released Q2/Q4 profiles are unusually sensitive to activation
10819/// quantization in the host-side attention path. With the normal A8W8
10820/// default, a short greedy decode can diverge after the first token even
10821/// though the selected-expert Metal path itself is unchanged. Keep the
10822/// reference float-activation contract for Flint unless the operator has
10823/// explicitly selected an `CMF_SDOT` policy.
10824fn flint_exact_cpu_default_enabled() -> bool {
10825    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
10826    *ON.get_or_init(|| std::env::var_os("CMF_SDOT").is_none())
10827}
10828
10829impl QTensor {
10830    #[inline]
10831    fn flint_exact_cpu_default(&self) -> bool {
10832        !FLOAT_ACTIVATIONS.get()
10833            && flint_exact_cpu_default_enabled()
10834            && matches!(
10835                self,
10836                Self::Mapped { model, .. } if model.header.arch.arch_name == "flint"
10837            )
10838    }
10839
10840    #[inline]
10841    fn any_flint_exact_cpu_default(ts: &[&QTensor]) -> bool {
10842        ts.iter().any(|t| t.flint_exact_cpu_default())
10843    }
10844
10845    #[inline]
10846    fn any_flint_exact_cpu_default_pairs(pairs: &[(&QTensor, &QTensor)]) -> bool {
10847        pairs
10848            .iter()
10849            .any(|(gate, up)| gate.flint_exact_cpu_default() || up.flint_exact_cpu_default())
10850    }
10851}
10852
10853/// Run `f` with the exact float-activation contract on this thread: the
10854/// 4-bit and 8-bit kernels skip the int8 activation split (`CMF_SDOT=0`
10855/// semantics, scoped). For parity tests and diagnostics — e.g. an `mxfp4`
10856/// model then computes bit for bit what its dequantized f32 twin computes.
10857pub fn with_float_activations<R>(f: impl FnOnce() -> R) -> R {
10858    float_activations_scope(f)
10859}
10860
10861/// Active on the current thread. `Pool` snapshots this with its job
10862/// descriptor and installs it on persistent workers, so a scoped quantized
10863/// matvec cannot accidentally choose a different kernel on a worker.
10864#[inline]
10865pub(crate) fn float_activations_active() -> bool {
10866    FLOAT_ACTIVATIONS.get()
10867}
10868
10869/// Restores the current thread's float-activation state on drop. The `Rc`
10870/// marker deliberately keeps the guard on the thread that created it: this
10871/// policy is thread-local, not a portable capability.
10872pub(crate) struct FloatActivationsScope {
10873    previous: bool,
10874    _not_send: std::marker::PhantomData<std::rc::Rc<()>>,
10875}
10876
10877impl Drop for FloatActivationsScope {
10878    fn drop(&mut self) {
10879        FLOAT_ACTIVATIONS.set(self.previous);
10880    }
10881}
10882
10883#[inline]
10884pub(crate) fn enter_float_activations_scope() -> FloatActivationsScope {
10885    FloatActivationsScope {
10886        previous: FLOAT_ACTIVATIONS.replace(true),
10887        _not_send: std::marker::PhantomData,
10888    }
10889}
10890
10891/// Run `f` with exact float activations (no A8W8 int8 activation
10892/// quantization in the quantized matvecs). The persistent row pool carries
10893/// this scope into a dispatched job and restores each worker afterwards.
10894/// Public for parity tests that compare the default kernels with the strict
10895/// ones in one process; `CMF_SDOT=0` is the process-wide equivalent.
10896#[doc(hidden)]
10897pub fn float_activations_scope<R>(f: impl FnOnce() -> R) -> R {
10898    let _scope = enter_float_activations_scope();
10899    f()
10900}
10901
10902/// Row-exact batching: while set, the x86 batched kernels (`qmatmat`,
10903/// `q4tp_matmat`) compute every (weight row, token) cell with the
10904/// single-token kernel instead of the blocked 2×4 / 1×4 / 1×8 tiles, so a
10905/// token's result does not depend on the batch it rides in and equals its
10906/// matvec. On ARM `q4tp_matmat` keeps its 1×4 tile but in the matvec's
10907/// reduction order, and no q4tp batch takes the AMX or device GEMM. The
10908/// MiMo speculative verify holds it (`row_exact_scope`) — its accepted
10909/// rows must be the rows plain decode would have produced.
10910// Shared with pool workers, so overlapping requests must keep the mode
10911// enabled until the LAST scope leaves. Saving/restoring a global bool is
10912// incorrect when two threads enter and leave in a non-LIFO order.
10913static ROW_EXACT: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
10914
10915pub(crate) fn row_exact() -> bool {
10916    ROW_EXACT.load(std::sync::atomic::Ordering::Acquire) != 0
10917}
10918
10919fn counted_row_exact_scope<R>(active: &std::sync::atomic::AtomicUsize, f: impl FnOnce() -> R) -> R {
10920    struct Restore<'a>(&'a std::sync::atomic::AtomicUsize);
10921    impl Drop for Restore<'_> {
10922        fn drop(&mut self) {
10923            self.0.fetch_sub(1, std::sync::atomic::Ordering::AcqRel);
10924        }
10925    }
10926    active.fetch_add(1, std::sync::atomic::Ordering::AcqRel);
10927    let _restore = Restore(active);
10928    f()
10929}
10930
10931/// Run `f` with row-exact batching on (also released on unwind).
10932pub(crate) fn row_exact_scope<R>(f: impl FnOnce() -> R) -> R {
10933    counted_row_exact_scope(&ROW_EXACT, f)
10934}
10935
10936/// A8W8 quantized-activation path available on THIS machine? One
10937/// switch across architectures: ARM dotprod (CMF_SDOT) or x86 AVX2
10938/// (CMF_AVX2 + the same CMF_SDOT exact-contract override).
10939#[inline]
10940pub(crate) fn a8w8_enabled() -> bool {
10941    #[cfg(target_arch = "aarch64")]
10942    {
10943        sdot_enabled()
10944    }
10945    #[cfg(target_arch = "x86_64")]
10946    {
10947        avx2_a8w8_enabled()
10948    }
10949    #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
10950    {
10951        false
10952    }
10953}
10954
10955/// int8·int8 dot dispatch: SDOT on ARM; AVX-512 VNNI (vpdpbusd) or AVX2
10956/// maddubs on x86. Callers are gated by `a8w8_enabled()`.
10957#[inline]
10958#[allow(unreachable_code)]
10959fn dot_i8_i8(w: &[u8], xq: &[i8]) -> i32 {
10960    #[cfg(target_arch = "aarch64")]
10961    unsafe {
10962        return dot_i8_sdot(w, xq);
10963    }
10964    #[cfg(target_arch = "x86_64")]
10965    unsafe {
10966        if avx512vnni_enabled() {
10967            return dot_i8_i8_vnni(w, xq);
10968        }
10969        return dot_i8_i8_avx2(w, xq);
10970    }
10971    w.iter()
10972        .zip(xq)
10973        .map(|(&a, &b)| (a as i8) as i32 * b as i32)
10974        .sum()
10975}
10976
10977/// AVX-512 VNNI available? (F+BW+VL+VNNI; `CMF_AVX512=0` falls back to
10978/// AVX2.) VL matters: short 32-byte groups (q4/vbit) ride the 256-bit
10979/// `vpdpbusd` encoding.
10980#[cfg(target_arch = "x86_64")]
10981fn avx512vnni_enabled() -> bool {
10982    use std::sync::OnceLock;
10983    static ON: OnceLock<bool> = OnceLock::new();
10984    *ON.get_or_init(|| {
10985        std::env::var("CMF_AVX512")
10986            .map(|v| v != "0")
10987            .unwrap_or(true)
10988            && std::arch::is_x86_feature_detected!("avx512f")
10989            && std::arch::is_x86_feature_detected!("avx512bw")
10990            && std::arch::is_x86_feature_detected!("avx512vl")
10991            && std::arch::is_x86_feature_detected!("avx512vnni")
10992    })
10993}
10994
10995/// Grouped-codec VNNI arms (the q4t/q4b/q1/q1t tile kernels): default
10996/// ON where AVX-512 VNNI exists (`CMF_VNNI_TILES=0` opt-out). Measured
10997/// on Ryzen 7950X (Zen4, 3 alternating process pairs, blocked GEMM
10998/// 4864×896 b=256): q4t 63→68 GF/s (+8%), q1 53→56 (+6%), q4b 72→75
10999/// (+4%) — consistent, no leg regressed. The tile kernels keep a
11000/// horizontal reduce per 32-weight group, so the `vpdpbusd` saving is
11001/// smaller than the long-dot q8 win (+13%), but it is real and free.
11002#[cfg(target_arch = "x86_64")]
11003fn vnni_tiles_enabled() -> bool {
11004    use std::sync::OnceLock;
11005    static ON: OnceLock<bool> = OnceLock::new();
11006    *ON.get_or_init(|| {
11007        std::env::var("CMF_VNNI_TILES")
11008            .map(|v| v != "0")
11009            .unwrap_or(true)
11010            && avx512vnni_enabled()
11011    })
11012}
11013
11014/// One 256-bit u8×i8 dot → i32 via `vpdpbusd` into a fresh accumulator
11015/// plus the same horizontal reduce the AVX2 kernels use. Products are
11016/// bounded (|w| ≤ 8 or ≤ 1), so maddubs never saturated — the i32 sum
11017/// is bit-identical to the maddubs+madd pair it replaces.
11018#[cfg(target_arch = "x86_64")]
11019#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
11020#[inline]
11021unsafe fn dpbusd_hsum(aw: core::arch::x86_64::__m256i, xs: core::arch::x86_64::__m256i) -> i32 {
11022    // SAFETY: pure register math.
11023    unsafe {
11024        use core::arch::x86_64::*;
11025        let d = _mm256_dpbusd_epi32(_mm256_setzero_si256(), aw, xs);
11026        let hi128 = _mm256_extracti128_si256::<1>(d);
11027        let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
11028        let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
11029        let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
11030        _mm_cvtsi128_si32(s32)
11031    }
11032}
11033
11034/// int8·int8 via AVX-512 VNNI: `vpdpbusd` fuses the maddubs+madd+add
11035/// triple into one u8×i8 dot-accumulate. AVX-512 has no vpsignb, so the
11036/// |w|·sign(x,w) trick becomes |w| × (x negated where w<0) via a mask
11037/// subtract — w==0 lanes contribute 0 through |w|=0 either way.
11038#[cfg(target_arch = "x86_64")]
11039#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
11040unsafe fn dot_i8_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
11041    // SAFETY: callers uphold slice-length contracts (see call sites).
11042    unsafe {
11043        use core::arch::x86_64::*;
11044        let n = w.len();
11045        let mut j = 0usize;
11046        let mut total: i32;
11047        // 4 independent accumulators: vpdpbusd is its own loop-carried
11048        // dependency (~5-cycle latency) — a single-acc loop runs
11049        // latency-bound and LOSES to the AVX2 maddubs kernel, measured
11050        // on Granite Rapids.
11051        {
11052            #[inline(always)]
11053            unsafe fn step(
11054                w: *const u8,
11055                x: *const i8,
11056                acc: core::arch::x86_64::__m512i,
11057            ) -> core::arch::x86_64::__m512i {
11058                unsafe {
11059                    use core::arch::x86_64::*;
11060                    let wv = _mm512_loadu_si512(w as *const _);
11061                    let xv = _mm512_loadu_si512(x as *const _);
11062                    let aw = _mm512_abs_epi8(wv);
11063                    let neg = _mm512_movepi8_mask(wv);
11064                    let sx = _mm512_mask_sub_epi8(xv, neg, _mm512_setzero_si512(), xv);
11065                    _mm512_dpbusd_epi32(acc, aw, sx)
11066                }
11067            }
11068            let (mut a0, mut a1, mut a2, mut a3) = (
11069                _mm512_setzero_si512(),
11070                _mm512_setzero_si512(),
11071                _mm512_setzero_si512(),
11072                _mm512_setzero_si512(),
11073            );
11074            while j + 256 <= n {
11075                a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
11076                a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), a1);
11077                a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), a2);
11078                a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), a3);
11079                j += 256;
11080            }
11081            while j + 64 <= n {
11082                a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
11083                j += 64;
11084            }
11085            let s01 = _mm512_add_epi32(a0, a1);
11086            let s23 = _mm512_add_epi32(a2, a3);
11087            total = _mm512_reduce_add_epi32(_mm512_add_epi32(s01, s23));
11088        }
11089        // 32-wide (q4/vbit groups are exactly 32 bytes).
11090        if j + 32 <= n {
11091            let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
11092            let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
11093            let d = _mm256_dpbusd_epi32(
11094                _mm256_setzero_si256(),
11095                _mm256_abs_epi8(wv),
11096                _mm256_sign_epi8(xv, wv),
11097            );
11098            let hi128 = _mm256_extracti128_si256::<1>(d);
11099            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
11100            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
11101            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
11102            total += _mm_cvtsi128_si32(s32);
11103            j += 32;
11104        }
11105        while j < n {
11106            total += (w[j] as i8) as i32 * xq[j] as i32;
11107            j += 1;
11108        }
11109        total
11110    }
11111}
11112
11113/// i8 row · f32 x via AVX2/FMA (x86 mirror of `dot_i8_f32_neon`).
11114#[cfg(target_arch = "x86_64")]
11115#[target_feature(enable = "avx2,fma")]
11116unsafe fn dot_i8_f32_avx2(w: &[u8], x: &[f32]) -> f32 {
11117    // SAFETY: callers uphold slice-length contracts (see call sites).
11118    unsafe {
11119        use core::arch::x86_64::*;
11120        let n = x.len();
11121        let wp = w.as_ptr();
11122        let xp = x.as_ptr();
11123        let (mut a0, mut a1) = (_mm256_setzero_ps(), _mm256_setzero_ps());
11124        let mut j = 0usize;
11125        while j + 16 <= n {
11126            let wb = _mm_loadu_si128(wp.add(j) as *const __m128i);
11127            let lo = _mm256_cvtepi8_epi32(wb);
11128            let hi = _mm256_cvtepi8_epi32(_mm_srli_si128::<8>(wb));
11129            a0 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(lo), _mm256_loadu_ps(xp.add(j)), a0);
11130            a1 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(hi), _mm256_loadu_ps(xp.add(j + 8)), a1);
11131            j += 16;
11132        }
11133        let acc = _mm256_add_ps(a0, a1);
11134        let hi128 = _mm256_extractf128_ps::<1>(acc);
11135        let s128 = _mm_add_ps(_mm256_castps256_ps128(acc), hi128);
11136        let s64 = _mm_add_ps(s128, _mm_movehl_ps(s128, s128));
11137        let s32 = _mm_add_ss(s64, _mm_shuffle_ps::<1>(s64, s64));
11138        let mut sum = _mm_cvtss_f32(s32);
11139        while j < n {
11140            sum += (*wp.add(j) as i8) as f32 * *xp.add(j);
11141            j += 1;
11142        }
11143        sum
11144    }
11145}
11146
11147/// int8(weight)·int8(activation) → i32 via AVX2 maddubs — the x86
11148/// analogue of the SDOT A8W8 path. `maddubs` takes u8×i8, so the
11149/// standard sign trick applies: |w| × sign(x, w) ≡ w × x per lane.
11150/// Pair saturation is safe: |w|≤128, |x|≤127 → 2·128·127 < 32767.
11151#[cfg(target_arch = "x86_64")]
11152#[target_feature(enable = "avx2")]
11153unsafe fn dot_i8_i8_avx2(w: &[u8], xq: &[i8]) -> i32 {
11154    // SAFETY: callers uphold slice-length contracts (see call sites).
11155    unsafe {
11156        use core::arch::x86_64::*;
11157        let n = w.len();
11158        let ones = _mm256_set1_epi16(1);
11159        let mut acc = _mm256_setzero_si256();
11160        let mut j = 0usize;
11161        while j + 32 <= n {
11162            let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
11163            let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
11164            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
11165            acc = _mm256_add_epi32(acc, _mm256_madd_epi16(p16, ones));
11166            j += 32;
11167        }
11168        let hi128 = _mm256_extracti128_si256::<1>(acc);
11169        let s128 = _mm_add_epi32(_mm256_castsi256_si128(acc), hi128);
11170        let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
11171        let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
11172        let mut s = _mm_cvtsi128_si32(s32);
11173        while j < n {
11174            s += (w[j] as i8) as i32 * xq[j] as i32;
11175            j += 1;
11176        }
11177        s
11178    }
11179}
11180
11181/// smmla 2×4: one instruction covers a 2-row × 2-activation × 8-deep
11182/// tile (32 MACs vs sdot's 16) — the weight pair loads once per 8-k
11183/// slice as a combined 2×8 register and meets two activation pairs.
11184#[cfg(target_arch = "aarch64")]
11185#[target_feature(enable = "neon,i8mm")]
11186unsafe fn dot_i8_smmla_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
11187    // SAFETY: callers uphold slice-length contracts.
11188    unsafe {
11189        use core::arch::aarch64::*;
11190        use core::arch::asm;
11191        let n = w0.len();
11192        let w0p = w0.as_ptr() as *const i8;
11193        let w1p = w1.as_ptr() as *const i8;
11194        // acc01 holds [c(r0,x0) c(r0,x1) c(r1,x0) c(r1,x1)]; acc23 the
11195        // same for x2/x3.
11196        let mut acc01 = vdupq_n_s32(0);
11197        let mut acc23 = vdupq_n_s32(0);
11198        let mut i = 0usize;
11199        while i + 8 <= n {
11200            let wa = vcombine_s8(vld1_s8(w0p.add(i)), vld1_s8(w1p.add(i)));
11201            let xb01 = vcombine_s8(
11202                vld1_s8(xs[0].as_ptr().add(i)),
11203                vld1_s8(xs[1].as_ptr().add(i)),
11204            );
11205            let xb23 = vcombine_s8(
11206                vld1_s8(xs[2].as_ptr().add(i)),
11207                vld1_s8(xs[3].as_ptr().add(i)),
11208            );
11209            asm!(
11210                "smmla {a01:v}.4s, {w:v}.16b, {x01:v}.16b",
11211                "smmla {a23:v}.4s, {w:v}.16b, {x23:v}.16b",
11212                a01 = inout(vreg) acc01, a23 = inout(vreg) acc23,
11213                w = in(vreg) wa, x01 = in(vreg) xb01, x23 = in(vreg) xb23,
11214                options(pure, nomem, nostack),
11215            );
11216            i += 8;
11217        }
11218        let mut out = [[0i32; 4]; 2];
11219        let a01: [i32; 4] = core::mem::transmute(acc01);
11220        let a23: [i32; 4] = core::mem::transmute(acc23);
11221        out[0][0] = a01[0];
11222        out[0][1] = a01[1];
11223        out[1][0] = a01[2];
11224        out[1][1] = a01[3];
11225        out[0][2] = a23[0];
11226        out[0][3] = a23[1];
11227        out[1][2] = a23[2];
11228        out[1][3] = a23[3];
11229        if i < n {
11230            for (k, x) in xs.iter().enumerate() {
11231                for j in i..n {
11232                    out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
11233                    out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
11234                }
11235            }
11236        }
11237        out
11238    }
11239}
11240
11241/// ARM twin of the x86 blocked prefill GEMM: two weight rows stay in
11242/// registers across four activation streams, eight sdot accumulators.
11243/// (The per-row form re-read each W row once per activation.)
11244#[cfg(target_arch = "aarch64")]
11245#[target_feature(enable = "neon,dotprod")]
11246unsafe fn dot_i8_sdot_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
11247    // SAFETY: callers uphold slice-length contracts.
11248    unsafe {
11249        use core::arch::aarch64::*;
11250        use core::arch::asm;
11251        let n = w0.len();
11252        let w0p = w0.as_ptr() as *const i8;
11253        let w1p = w1.as_ptr() as *const i8;
11254        let mut acc = [[vdupq_n_s32(0); 4]; 2];
11255        let mut i = 0usize;
11256        while i + 16 <= n {
11257            let wv0 = vld1q_s8(w0p.add(i));
11258            let wv1 = vld1q_s8(w1p.add(i));
11259            for (k, x) in xs.iter().enumerate() {
11260                let xv = vld1q_s8(x.as_ptr().add(i));
11261                let (mut a0, mut a1) = (acc[0][k], acc[1][k]);
11262                asm!(
11263                    "sdot {a0:v}.4s, {w0:v}.16b, {x:v}.16b",
11264                    "sdot {a1:v}.4s, {w1:v}.16b, {x:v}.16b",
11265                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
11266                    w0 = in(vreg) wv0, w1 = in(vreg) wv1, x = in(vreg) xv,
11267                    options(pure, nomem, nostack),
11268                );
11269                acc[0][k] = a0;
11270                acc[1][k] = a1;
11271            }
11272            i += 16;
11273        }
11274        let mut out = [[0i32; 4]; 2];
11275        for r in 0..2 {
11276            for k in 0..4 {
11277                out[r][k] = vaddvq_s32(acc[r][k]);
11278            }
11279        }
11280        if i < n {
11281            for (k, x) in xs.iter().enumerate() {
11282                for j in i..n {
11283                    out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
11284                    out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
11285                }
11286            }
11287        }
11288        out
11289    }
11290}
11291
11292/// Blocked 2 weight rows × 4 activations for the prefill GEMM
11293/// (roadmap P0: packed panels + multi-row accumulators). The two rows'
11294/// abs() live in registers across all four activation streams; the
11295/// sign-fixup is recomputed per pair (the price of the maddubs trick).
11296/// Returns raw i8·i8 dots; the caller applies scales and outliers.
11297#[cfg(target_arch = "x86_64")]
11298#[target_feature(enable = "avx2")]
11299unsafe fn dot_i8_i8_avx2_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
11300    // SAFETY: callers uphold slice-length contracts.
11301    unsafe {
11302        use core::arch::x86_64::*;
11303        let n = w0.len();
11304        let ones = _mm256_set1_epi16(1);
11305        let mut acc = [[_mm256_setzero_si256(); 4]; 2];
11306        let mut j = 0usize;
11307        while j + 32 <= n {
11308            let wv0 = _mm256_loadu_si256(w0.as_ptr().add(j) as *const __m256i);
11309            let wv1 = _mm256_loadu_si256(w1.as_ptr().add(j) as *const __m256i);
11310            let aw0 = _mm256_abs_epi8(wv0);
11311            let aw1 = _mm256_abs_epi8(wv1);
11312            for (k, x) in xs.iter().enumerate() {
11313                let xv = _mm256_loadu_si256(x.as_ptr().add(j) as *const __m256i);
11314                let p0 = _mm256_maddubs_epi16(aw0, _mm256_sign_epi8(xv, wv0));
11315                acc[0][k] = _mm256_add_epi32(acc[0][k], _mm256_madd_epi16(p0, ones));
11316                let p1 = _mm256_maddubs_epi16(aw1, _mm256_sign_epi8(xv, wv1));
11317                acc[1][k] = _mm256_add_epi32(acc[1][k], _mm256_madd_epi16(p1, ones));
11318            }
11319            j += 32;
11320        }
11321        let mut out = [[0i32; 4]; 2];
11322        for r in 0..2 {
11323            for k in 0..4 {
11324                let a = acc[r][k];
11325                let hi128 = _mm256_extracti128_si256::<1>(a);
11326                let s128 = _mm_add_epi32(_mm256_castsi256_si128(a), hi128);
11327                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
11328                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
11329                out[r][k] = _mm_cvtsi128_si32(s32);
11330            }
11331        }
11332        if j < n {
11333            for (k, x) in xs.iter().enumerate() {
11334                for i in j..n {
11335                    out[0][k] += (w0[i] as i8) as i32 * x[i] as i32;
11336                    out[1][k] += (w1[i] as i8) as i32 * x[i] as i32;
11337                }
11338            }
11339        }
11340        out
11341    }
11342}
11343
11344/// AVX2/VNNI q8 row dot with exact outlier correction (x86 mirror of
11345/// `row_dot_sdot` — same A8W8 contract). With AVX-512 VNNI the row goes
11346/// through the bias trick: Σ(w+128)·x via pure `vpdpbusd` (no per-lane
11347/// sign fixups), corrected by −128·Σx with Σx precomputed per split.
11348#[cfg(target_arch = "x86_64")]
11349#[inline]
11350fn row_dot_avx2(row: &[u8], act: &SplitAct) -> f32 {
11351    let dot = if avx512vnni_enabled() && row.len() >= 64 {
11352        (unsafe { dot_u8p128_i8_vnni(row, &act.xq) }) - 128 * act.xsum
11353    } else {
11354        unsafe { dot_i8_i8_avx2(row, &act.xq) }
11355    };
11356    let mut acc = dot as f32 * act.sx;
11357    for &(j, xv) in &act.outliers {
11358        acc += (row[j] as i8) as f32 * xv;
11359    }
11360    acc
11361}
11362
11363/// Σ (w[i]+128)·x[i] via pure `vpdpbusd` — the caller subtracts
11364/// 128·Σx. Four independent accumulators (dpbusd is ~5-cycle latency;
11365/// a single-acc loop runs latency-bound, measured on Granite Rapids).
11366#[cfg(target_arch = "x86_64")]
11367#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
11368unsafe fn dot_u8p128_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
11369    // SAFETY: callers uphold slice-length contracts (see call sites).
11370    unsafe {
11371        use core::arch::x86_64::*;
11372        let n = w.len();
11373        let flip = _mm512_set1_epi8(-128); // XOR 0x80: i8 w → u8 (w+128)
11374        #[inline(always)]
11375        unsafe fn step(
11376            w: *const u8,
11377            x: *const i8,
11378            flip: core::arch::x86_64::__m512i,
11379            acc: core::arch::x86_64::__m512i,
11380        ) -> core::arch::x86_64::__m512i {
11381            unsafe {
11382                use core::arch::x86_64::*;
11383                let wv = _mm512_xor_si512(_mm512_loadu_si512(w as *const _), flip);
11384                _mm512_dpbusd_epi32(acc, wv, _mm512_loadu_si512(x as *const _))
11385            }
11386        }
11387        let (mut a0, mut a1, mut a2, mut a3) = (
11388            _mm512_setzero_si512(),
11389            _mm512_setzero_si512(),
11390            _mm512_setzero_si512(),
11391            _mm512_setzero_si512(),
11392        );
11393        let mut j = 0usize;
11394        while j + 256 <= n {
11395            a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
11396            a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), flip, a1);
11397            a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), flip, a2);
11398            a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), flip, a3);
11399            j += 256;
11400        }
11401        while j + 64 <= n {
11402            a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
11403            j += 64;
11404        }
11405        let mut total = _mm512_reduce_add_epi32(_mm512_add_epi32(
11406            _mm512_add_epi32(a0, a1),
11407            _mm512_add_epi32(a2, a3),
11408        ));
11409        // Scalar tail: (w as i8) + 128 ≡ (w as u8) ^ 0x80.
11410        while j < n {
11411            total += ((w[j] ^ 0x80) as i32) * xq[j] as i32;
11412            j += 1;
11413        }
11414        total
11415    }
11416}
11417
11418/// One q4 row via AVX2: nibbles → centered i8 (unpacklo/hi restores the
11419/// writer's flat order, same as the NEON vzip pair), maddubs against
11420/// the pre-quantized activation group, × the group's f16 scale. Pair
11421/// saturation safe: |w|≤8, |x|≤127 → 2·8·127 ≪ 32767. Mirror of
11422/// `dot_q4_row_sdot`.
11423#[cfg(target_arch = "x86_64")]
11424#[target_feature(enable = "avx2")]
11425unsafe fn dot_q4_row_avx2(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
11426    // SAFETY: callers uphold slice-length contracts (16 packed bytes and
11427    // 2 scale bytes per group; xq.len() == gpr·GROUP_SIZE).
11428    unsafe {
11429        use core::arch::x86_64::*;
11430        let lomask = _mm_set1_epi8(0x0F);
11431        let eight = _mm256_set1_epi8(8);
11432        let ones = _mm256_set1_epi16(1);
11433        let mut acc = 0f32;
11434        for gi in 0..gpr {
11435            let g = g0 + gi;
11436            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
11437            let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
11438            let lo = _mm_and_si128(b, lomask);
11439            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
11440            let w = _mm256_sub_epi8(
11441                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
11442                eight,
11443            );
11444            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
11445            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
11446            let d = _mm256_madd_epi16(p16, ones);
11447            let hi128 = _mm256_extracti128_si256::<1>(d);
11448            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
11449            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
11450            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
11451            acc += _mm_cvtsi128_si32(s32) as f32 * s;
11452        }
11453        acc
11454    }
11455}
11456
11457/// Two-activation q4 row via AVX2: nibbles unpacked ONCE per group,
11458/// both activations dotted against the same centered i8 register.
11459#[cfg(target_arch = "x86_64")]
11460#[target_feature(enable = "avx2")]
11461unsafe fn dot_q4_row_avx2_2(
11462    packed: &[u8],
11463    scales: &[u8],
11464    g0: usize,
11465    gpr: usize,
11466    xq1: &[i8],
11467    xq2: &[i8],
11468) -> (f32, f32) {
11469    // SAFETY: callers uphold slice-length contracts (see dot_q4_row_avx2).
11470    unsafe {
11471        use core::arch::x86_64::*;
11472        let lomask = _mm_set1_epi8(0x0F);
11473        let eight = _mm256_set1_epi8(8);
11474        let ones = _mm256_set1_epi16(1);
11475        let (mut acc1, mut acc2) = (0f32, 0f32);
11476        #[inline(always)]
11477        unsafe fn hsum(d: core::arch::x86_64::__m256i) -> i32 {
11478            unsafe {
11479                use core::arch::x86_64::*;
11480                let hi128 = _mm256_extracti128_si256::<1>(d);
11481                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
11482                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
11483                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
11484                _mm_cvtsi128_si32(s32)
11485            }
11486        }
11487        for gi in 0..gpr {
11488            let g = g0 + gi;
11489            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
11490            let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
11491            let lo = _mm_and_si128(b, lomask);
11492            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
11493            let w = _mm256_sub_epi8(
11494                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
11495                eight,
11496            );
11497            let aw = _mm256_abs_epi8(w);
11498            let x1 = _mm256_loadu_si256(xq1.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
11499            let x2 = _mm256_loadu_si256(xq2.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
11500            let d1 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x1, w)), ones);
11501            let d2 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x2, w)), ones);
11502            acc1 += hsum(d1) as f32 * s;
11503            acc2 += hsum(d2) as f32 * s;
11504        }
11505        (acc1, acc2)
11506    }
11507}
11508
11509/// One q8 row range via AVX2 (x86 mirror of `q8_range_sdot`).
11510#[cfg(target_arch = "x86_64")]
11511fn q8_range_avx2(
11512    q: &[u8],
11513    row_scale: &[f32],
11514    act: &SplitAct,
11515    cols: usize,
11516    out_addr: SendMut,
11517    start: usize,
11518    end: usize,
11519) {
11520    for o in start..end {
11521        let v = row_dot_avx2(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
11522        // SAFETY: disjoint row ranges per worker.
11523        unsafe { *out_addr.at(o) = v };
11524    }
11525}
11526
11527/// Two-input q8 row range via AVX2 (x86 mirror of `q8_range2_sdot`).
11528#[cfg(target_arch = "x86_64")]
11529#[allow(clippy::too_many_arguments)]
11530fn q8_range2_avx2(
11531    q: &[u8],
11532    row_scale: &[f32],
11533    a1: &SplitAct,
11534    a2: &SplitAct,
11535    cols: usize,
11536    p1: SendMut,
11537    p2: SendMut,
11538    start: usize,
11539    end: usize,
11540) {
11541    for o in start..end {
11542        let row = &q[o * cols..(o + 1) * cols];
11543        // SAFETY: disjoint row ranges per worker.
11544        unsafe {
11545            *p1.at(o) = row_dot_avx2(row, a1) * row_scale[o];
11546            *p2.at(o) = row_dot_avx2(row, a2) * row_scale[o];
11547        }
11548    }
11549}
11550
11551// ───────────────────── A8W8 SDOT path (port of vmfcore, ×1.78 decode) ─────────────────────
11552
11553/// ARMv8.6 i8mm (smmla): 32 int8 MACs per instruction vs sdot's 16 —
11554/// yet MEASURED 2.4× SLOWER than the blocked sdot on Apple silicon
11555/// (108 vs 264 GF/s): the on-the-fly vcombine packing and the two-
11556/// accumulator dependency chain swamp the MAC advantage, and Apple's
11557/// four SIMD pipes already keep sdot fed. OPT-IN (CMF_I8MM=1) for
11558/// field trials on Cortex-A710/X-class parts with two pipes, where the
11559/// balance may differ; a pre-interleaved weight layout (repack infra)
11560/// is the known path if it ever earns its keep.
11561#[cfg(target_arch = "aarch64")]
11562fn i8mm_enabled() -> bool {
11563    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
11564    *ON.get_or_init(|| {
11565        std::env::var("CMF_I8MM").map(|v| v == "1").unwrap_or(false)
11566            && std::arch::is_aarch64_feature_detected!("i8mm")
11567    })
11568}
11569
11570/// SDOT enabled? Default ON when the CPU has ARMv8.2 dotprod;
11571/// `CMF_SDOT=0` disables (falls back to i8×f32 NEON).
11572/// (On non-ARM release builds only the test tolerance switch calls it.)
11573#[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
11574fn sdot_enabled() -> bool {
11575    if FLOAT_ACTIVATIONS.get() {
11576        return false;
11577    }
11578    use std::sync::OnceLock;
11579    static ON: OnceLock<bool> = OnceLock::new();
11580    *ON.get_or_init(|| {
11581        let want = std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true);
11582        if !want {
11583            return false;
11584        }
11585
11586        #[cfg(target_arch = "aarch64")]
11587        {
11588            if std::arch::is_aarch64_feature_detected!("dotprod") {
11589                return true;
11590            }
11591            #[cfg(target_os = "android")]
11592            {
11593                if let Ok(cpuinfo) = std::fs::read_to_string("/proc/cpuinfo") {
11594                    if cpuinfo.lines().any(|l| {
11595                        (l.starts_with("Features") || l.starts_with("features"))
11596                            && l.contains("asimddp")
11597                    }) {
11598                        return true;
11599                    }
11600                }
11601            }
11602            false
11603        }
11604        #[cfg(not(target_arch = "aarch64"))]
11605        {
11606            false
11607        }
11608    })
11609}
11610
11611/// Two-field activation split (≡ vmfcore `q8_split_prep`): outlier
11612/// channels (>8·rms) are computed exactly in f32; the bulk (outliers
11613/// zeroed → clean absmax) goes through int8 SDOT. Computed ONCE per
11614/// matvec, shared by all rows/workers.
11615struct SplitAct {
11616    xq: Vec<i8>,
11617    sx: f32,
11618    outliers: Vec<(usize, f32)>,
11619    /// Σ xq — the VNNI bias-trick correction (`(w+128)·x` sums need
11620    /// `−128·Σx`); one i32 per split, computed once per matvec.
11621    #[cfg_attr(not(target_arch = "x86_64"), allow(dead_code))]
11622    xsum: i32,
11623}
11624
11625thread_local! {
11626    /// Recycled xq buffers: split_act runs for every matvec (~200/token)
11627    /// and its hidden-size allocation was steady-state heap churn.
11628    static XQ_FREE: std::cell::RefCell<Vec<Vec<i8>>> =
11629        const { std::cell::RefCell::new(Vec::new()) };
11630}
11631
11632impl Drop for SplitAct {
11633    fn drop(&mut self) {
11634        let buf = std::mem::take(&mut self.xq);
11635        if buf.capacity() > 0 {
11636            XQ_FREE.with(|f| {
11637                let mut f = f.borrow_mut();
11638                if f.len() < 16 {
11639                    f.push(buf);
11640                }
11641            });
11642        }
11643    }
11644}
11645
11646thread_local! {
11647    /// One scratch row per WORKER, kept for the life of the thread.
11648    ///
11649    /// The kernels take a row of group scales per dispatch, and a fresh
11650    /// `vec![0f32; gpr]` inside the closure is one allocation per worker per
11651    /// dispatch — on the release checkpoint about six thousand a token, a
11652    /// quarter of everything the benchmark counts.
11653    static KROW: UnsafeCell<Vec<f32>> = const { UnsafeCell::new(Vec::new()) };
11654    /// Two scratch rows for a gate/up pair.  Q2TP and Q4TP fused SwiGLU
11655    /// decode each projection's ladder independently, but both ladders have
11656    /// the same lifetime; retaining them per worker removes two heap trips
11657    /// from every dispatched expert pair.
11658    static KROWS: UnsafeCell<[Vec<f32>; 2]> =
11659        const { UnsafeCell::new([Vec::new(), Vec::new()]) };
11660}
11661
11662/// Borrow `n` floats of the calling worker's scratch. Nothing inside a
11663/// kernel body borrows it again, which is what keeps the RefCell honest.
11664#[inline]
11665fn with_krow<R>(n: usize, f: impl FnOnce(&mut [f32]) -> R) -> R {
11666    KROW.with(|s| {
11667        // SAFETY: `KROW` is thread-local and this function does not recurse;
11668        // each worker has exclusive access to its own scratch row.
11669        let b = unsafe { &mut *s.get() };
11670        if b.len() < n {
11671            b.resize(n, 0.0);
11672        }
11673        f(&mut b[..n])
11674    })
11675}
11676
11677/// `t.round().clamp(-127.0, 127.0) as i8`, bit for bit, without the libm
11678/// call. On baseline x86-64 (no SSE4.1 `roundps`) `f32::round` is a
11679/// function call per element, and split_act runs it over every hidden
11680/// state before every matvec: measured 27 us a call on a 2048-wide
11681/// activation on an EPYC 7763 — 5.4 ms of a 55 ms decode token, all of
11682/// it on the caller's thread while thirty workers wait. Clamping first is
11683/// equivalent (round is monotonic and ±127 are integers), and after the
11684/// clamp `t - trunc(t)` is exact, so the half-away-from-zero decision is
11685/// the one `round` makes. NaN clamps to NaN and converts to 0, as before.
11686/// The loop vectorizes (cvttps2dq + compare/select).
11687#[inline(always)]
11688fn q8_round(t: f32) -> i8 {
11689    let t = t.clamp(-127.0, 127.0);
11690    let i = t as i32;
11691    let f = t - i as f32;
11692    let r = if f >= 0.5 {
11693        i + 1
11694    } else if f <= -0.5 {
11695        i - 1
11696    } else {
11697        i
11698    };
11699    r as i8
11700}
11701
11702#[inline]
11703fn with_krows<R>(n: usize, f: impl FnOnce(&mut [f32], &mut [f32]) -> R) -> R {
11704    KROWS.with(|s| {
11705        // SAFETY: KROWS is thread-local and no kernel body recursively calls
11706        // this helper; each participant owns its two vectors exclusively.
11707        let b = unsafe { &mut *s.get() };
11708        for row in b.iter_mut() {
11709            if row.len() < n {
11710                row.resize(n, 0.0);
11711            }
11712        }
11713        let (a, b) = b.split_at_mut(1);
11714        f(&mut a[0][..n], &mut b[0][..n])
11715    })
11716}
11717
11718#[inline]
11719fn silu_mul_limited(mut gate: f32, mut up: f32, limit: f32) -> f32 {
11720    if limit > 0.0 {
11721        up = up.clamp(-limit, limit);
11722        gate = gate.min(limit);
11723    }
11724    gate / (1.0 + (-gate).exp()) * up
11725}
11726
11727fn split_act(x: &[f32]) -> SplitAct {
11728    let _prof = crate::cpuprof::time(crate::cpuprof::Slot::SplitAct);
11729    let n = x.len();
11730    let rms = (x.iter().map(|&v| (v * v) as f64).sum::<f64>() / n.max(1) as f64).sqrt() as f32;
11731    let thr = 8.0 * rms;
11732    // One pass: collect outliers and the bulk absmax (outliers excluded —
11733    // identical to the old zero-then-fold over a copied buffer, minus the
11734    // full-vector copy).
11735    let mut outliers: Vec<(usize, f32)> = Vec::new();
11736    let mut amax = 0f32;
11737    for (j, &v) in x.iter().enumerate() {
11738        let a = v.abs();
11739        if a > thr {
11740            outliers.push((j, v));
11741        } else if a > amax {
11742            amax = a;
11743        }
11744    }
11745    let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
11746    let inv = 1.0 / sx;
11747    let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
11748    xq.clear();
11749    xq.reserve(n);
11750    if outliers.is_empty() {
11751        xq.extend(
11752            x.iter()
11753                .map(|&v| q8_round(v * inv)),
11754        );
11755    } else {
11756        // Outlier slots quantize to 0 (their exact term is added later).
11757        xq.extend(x.iter().map(|&v| {
11758            if v.abs() > thr {
11759                0
11760            } else {
11761                q8_round(v * inv)
11762            }
11763        }));
11764    }
11765    let xsum = xq.iter().map(|&v| v as i32).sum();
11766    SplitAct {
11767        xq,
11768        sx,
11769        outliers,
11770        xsum,
11771    }
11772}
11773
11774fn split_act_q8_2f(x: &[f32], col: &[f32]) -> SplitAct {
11775    let _prof = crate::cpuprof::time(crate::cpuprof::Slot::SplitAct);
11776    let n = x.len();
11777    let rms = (x
11778        .iter()
11779        .zip(col)
11780        .map(|(&a, &c)| {
11781            let v = a * c;
11782            (v * v) as f64
11783        })
11784        .sum::<f64>()
11785        / n.max(1) as f64)
11786        .sqrt() as f32;
11787    let thr = 8.0 * rms;
11788
11789    let mut outliers = Vec::new();
11790    let mut amax = 0f32;
11791    for (j, (&a, &c)) in x.iter().zip(col).enumerate() {
11792        let v = a * c;
11793        let s = v.abs();
11794        if s > thr {
11795            outliers.push((j, v));
11796        } else if s > amax {
11797            amax = s;
11798        }
11799    }
11800
11801    let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
11802    let inv = 1.0 / sx;
11803    let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
11804    xq.clear();
11805    xq.reserve(n);
11806    if outliers.is_empty() {
11807        xq.extend(
11808            x.iter()
11809                .zip(col)
11810                .map(|(&a, &c)| q8_round((a * c) * inv)),
11811        );
11812    } else {
11813        xq.extend(x.iter().zip(col).map(|(&a, &c)| {
11814            let v = a * c;
11815            if v.abs() > thr {
11816                0
11817            } else {
11818                q8_round(v * inv)
11819            }
11820        }));
11821    }
11822    let xsum = xq.iter().map(|&v| v as i32).sum();
11823    SplitAct {
11824        xq,
11825        sx,
11826        outliers,
11827        xsum,
11828    }
11829}
11830
11831/// int8(weight)·int8(activation) → i32 via `sdot` (inline asm — the
11832/// vdotq intrinsic is unstable; port of vmfcore `dot_i8_sdot`).
11833#[cfg(target_arch = "aarch64")]
11834#[target_feature(enable = "neon,dotprod")]
11835unsafe fn dot_i8_sdot(w: &[u8], xq: &[i8]) -> i32 {
11836    // SAFETY: callers uphold slice-length contracts (see call sites).
11837    unsafe {
11838        use core::arch::aarch64::*;
11839        use core::arch::asm;
11840        let wp = w.as_ptr() as *const i8;
11841        let n = w.len();
11842        let (mut a0, mut a1, mut a2, mut a3) = (
11843            vdupq_n_s32(0),
11844            vdupq_n_s32(0),
11845            vdupq_n_s32(0),
11846            vdupq_n_s32(0),
11847        );
11848        let mut i = 0;
11849        while i + 64 <= n {
11850            let (w0, x0) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
11851            let (w1, x1) = (vld1q_s8(wp.add(i + 16)), vld1q_s8(xq.as_ptr().add(i + 16)));
11852            let (w2, x2) = (vld1q_s8(wp.add(i + 32)), vld1q_s8(xq.as_ptr().add(i + 32)));
11853            let (w3, x3) = (vld1q_s8(wp.add(i + 48)), vld1q_s8(xq.as_ptr().add(i + 48)));
11854            asm!(
11855                "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
11856                "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
11857                "sdot {a2:v}.4s, {w2:v}.16b, {x2:v}.16b",
11858                "sdot {a3:v}.4s, {w3:v}.16b, {x3:v}.16b",
11859                a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
11860                w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
11861                w2 = in(vreg) w2, x2 = in(vreg) x2, w3 = in(vreg) w3, x3 = in(vreg) x3,
11862                options(pure, nomem, nostack),
11863            );
11864            i += 64;
11865        }
11866        while i + 16 <= n {
11867            let (wv, xv) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
11868            asm!("sdot {a:v}.4s, {w:v}.16b, {x:v}.16b",
11869                 a = inout(vreg) a0, w = in(vreg) wv, x = in(vreg) xv, options(pure, nomem, nostack));
11870            i += 16;
11871        }
11872        let mut s = vaddvq_s32(vaddq_s32(vaddq_s32(a0, a1), vaddq_s32(a2, a3)));
11873        while i < n {
11874            s += (*wp.add(i)) as i32 * xq[i] as i32;
11875            i += 1;
11876        }
11877        s
11878    }
11879}
11880
11881/// Row-blocked SDOT: 4 output rows per pass — the activation chunk is
11882/// loaded once and reused, 4 independent accumulators hide sdot latency
11883/// (port of vmfcore `dot_i8_sdot_4rows`).
11884#[cfg(target_arch = "aarch64")]
11885#[target_feature(enable = "neon,dotprod")]
11886unsafe fn dot_i8_sdot_4rows(w0: &[u8], w1: &[u8], w2: &[u8], w3: &[u8], xq: &[i8]) -> [i32; 4] {
11887    // SAFETY: callers uphold slice-length contracts (see call sites).
11888    unsafe {
11889        use core::arch::aarch64::*;
11890        use core::arch::asm;
11891        let n = xq.len();
11892        let px = xq.as_ptr();
11893        let (p0, p1, p2, p3) = (
11894            w0.as_ptr() as *const i8,
11895            w1.as_ptr() as *const i8,
11896            w2.as_ptr() as *const i8,
11897            w3.as_ptr() as *const i8,
11898        );
11899        let (mut a0, mut a1, mut a2, mut a3) = (
11900            vdupq_n_s32(0),
11901            vdupq_n_s32(0),
11902            vdupq_n_s32(0),
11903            vdupq_n_s32(0),
11904        );
11905        let mut i = 0;
11906        while i + 16 <= n {
11907            let x = vld1q_s8(px.add(i));
11908            let v0 = vld1q_s8(p0.add(i));
11909            let v1 = vld1q_s8(p1.add(i));
11910            let v2 = vld1q_s8(p2.add(i));
11911            let v3 = vld1q_s8(p3.add(i));
11912            asm!(
11913                "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
11914                "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
11915                "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
11916                "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
11917                a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
11918                v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
11919                options(pure, nomem, nostack),
11920            );
11921            i += 16;
11922        }
11923        let mut r = [
11924            vaddvq_s32(a0),
11925            vaddvq_s32(a1),
11926            vaddvq_s32(a2),
11927            vaddvq_s32(a3),
11928        ];
11929        while i < n {
11930            let xi = *px.add(i) as i32;
11931            r[0] += (*p0.add(i)) as i32 * xi;
11932            r[1] += (*p1.add(i)) as i32 * xi;
11933            r[2] += (*p2.add(i)) as i32 * xi;
11934            r[3] += (*p3.add(i)) as i32 * xi;
11935            i += 1;
11936        }
11937        r
11938    }
11939}
11940
11941/// 4 interleaved rows in one pass: the repacked group is [r0[c], r1[c],
11942/// r2[c], r3[c]] per 16-byte chunk, so each iteration reads ONE 64-byte
11943/// line plus the shared activation chunk — a single sequential weight
11944/// stream per worker. Per-row accumulation is the same one-accumulator
11945/// scheme as `dot_i8_sdot_4rows`; integer sums are exact, so outputs
11946/// are bit-identical to the mmap-layout kernel.
11947#[cfg(target_arch = "aarch64")]
11948#[target_feature(enable = "neon,dotprod")]
11949unsafe fn dot_i8_sdot_4rows_il(g: &[u8], xq: &[i8]) -> [i32; 4] {
11950    // SAFETY: callers uphold slice-length contracts (g.len() == 4·n,
11951    // n % 16 == 0 — guaranteed by the repack gate).
11952    unsafe {
11953        use core::arch::aarch64::*;
11954        use core::arch::asm;
11955        let n = xq.len();
11956        let px = xq.as_ptr();
11957        let pg = g.as_ptr() as *const i8;
11958        let (mut a0, mut a1, mut a2, mut a3) = (
11959            vdupq_n_s32(0),
11960            vdupq_n_s32(0),
11961            vdupq_n_s32(0),
11962            vdupq_n_s32(0),
11963        );
11964        let mut i = 0;
11965        while i + 16 <= n {
11966            let x = vld1q_s8(px.add(i));
11967            let base = pg.add(4 * i);
11968            let v0 = vld1q_s8(base);
11969            let v1 = vld1q_s8(base.add(16));
11970            let v2 = vld1q_s8(base.add(32));
11971            let v3 = vld1q_s8(base.add(48));
11972            asm!(
11973                "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
11974                "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
11975                "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
11976                "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
11977                a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
11978                v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
11979                options(pure, nomem, nostack),
11980            );
11981            i += 16;
11982        }
11983        [
11984            vaddvq_s32(a0),
11985            vaddvq_s32(a1),
11986            vaddvq_s32(a2),
11987            vaddvq_s32(a3),
11988        ]
11989    }
11990}
11991
11992/// One q8 row range via SDOT (4-row blocks + tail) — the body of
11993/// `qmatvec`'s hot loop, extracted so multi-matrix jobs can drive the
11994/// SAME kernel for several tensors under one pool dispatch. `rep` — the
11995/// load-time interleaved repack (empty = mmap layout only); rows outside
11996/// full 4-row groups always come from the mmap layout.
11997#[cfg(target_arch = "aarch64")]
11998fn q8_range_sdot(
11999    q: &[u8],
12000    rep: &[u8],
12001    row_scale: &[f32],
12002    act: &SplitAct,
12003    cols: usize,
12004    out_addr: SendMut,
12005    start: usize,
12006    end: usize,
12007) {
12008    let mut o = start;
12009    // Leading rows to the group boundary (repack path only): the pool
12010    // splits row ranges arbitrarily, groups are absolute.
12011    if !rep.is_empty() {
12012        while o < end && o % 4 != 0 {
12013            let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
12014            unsafe { *out_addr.at(o) = v };
12015            o += 1;
12016        }
12017    }
12018    while o + 4 <= end {
12019        let r = if rep.is_empty() {
12020            unsafe {
12021                dot_i8_sdot_4rows(
12022                    &q[o * cols..(o + 1) * cols],
12023                    &q[(o + 1) * cols..(o + 2) * cols],
12024                    &q[(o + 2) * cols..(o + 3) * cols],
12025                    &q[(o + 3) * cols..(o + 4) * cols],
12026                    &act.xq,
12027                )
12028            }
12029        } else {
12030            unsafe { dot_i8_sdot_4rows_il(&rep[o * cols..(o + 4) * cols], &act.xq) }
12031        };
12032        for k in 0..4 {
12033            let mut acc = r[k] as f32 * act.sx;
12034            for &(j, xv) in &act.outliers {
12035                acc += (q[(o + k) * cols + j] as i8) as f32 * xv;
12036            }
12037            // SAFETY: disjoint row ranges per worker.
12038            unsafe { *out_addr.at(o + k) = acc * row_scale[o + k] };
12039        }
12040        o += 4;
12041    }
12042    while o < end {
12043        let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
12044        unsafe { *out_addr.at(o) = v };
12045        o += 1;
12046    }
12047}
12048
12049/// Two-input q8 row range via SDOT — `qmatvec2`'s hot loop, extracted
12050/// for the fused pair multi-matrix job (`matvec2_many`).
12051#[cfg(target_arch = "aarch64")]
12052#[allow(clippy::too_many_arguments)]
12053fn q8_range2_sdot(
12054    q: &[u8],
12055    row_scale: &[f32],
12056    a1: &SplitAct,
12057    a2: &SplitAct,
12058    cols: usize,
12059    p1: SendMut,
12060    p2: SendMut,
12061    start: usize,
12062    end: usize,
12063) {
12064    for o in start..end {
12065        let row = &q[o * cols..(o + 1) * cols];
12066        // SAFETY: disjoint row ranges per worker.
12067        unsafe {
12068            *p1.at(o) = row_dot_sdot(row, a1) * row_scale[o];
12069            *p2.at(o) = row_dot_sdot(row, a2) * row_scale[o];
12070        }
12071    }
12072}
12073
12074/// Two-input q8 row range, f32 kernel (non-SDOT) — same extraction.
12075#[allow(clippy::too_many_arguments)]
12076fn q8_range2_f32(
12077    q: &[u8],
12078    row_scale: &[f32],
12079    x1: &[f32],
12080    x2: &[f32],
12081    cols: usize,
12082    p1: SendMut,
12083    p2: SendMut,
12084    start: usize,
12085    end: usize,
12086) {
12087    for o in start..end {
12088        let row = &q[o * cols..(o + 1) * cols];
12089        // SAFETY: disjoint row ranges per worker.
12090        unsafe {
12091            *p1.at(o) = dot_i8_f32(row, x1) * row_scale[o];
12092            *p2.at(o) = dot_i8_f32(row, x2) * row_scale[o];
12093        }
12094    }
12095}
12096
12097/// Scalar/NEON-f32 q8 row range (non-SDOT platforms) — same extraction.
12098fn q8_range_f32(
12099    q: &[u8],
12100    row_scale: &[f32],
12101    xs: &[f32],
12102    cols: usize,
12103    out_addr: SendMut,
12104    start: usize,
12105    end: usize,
12106) {
12107    for o in start..end {
12108        let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
12109        // SAFETY: disjoint row ranges per worker.
12110        unsafe { *out_addr.at(o) = v };
12111    }
12112}
12113
12114/// One q8 row against a split activation, portable: the per-arch fast
12115/// dots where they exist, the exact scalar loop elsewhere. The scalar
12116/// arm is also the test oracle for both fast arms.
12117#[inline]
12118fn q8_row_dot(row: &[u8], act: &SplitAct) -> f32 {
12119    #[cfg(target_arch = "aarch64")]
12120    return row_dot_sdot(row, act);
12121    #[cfg(target_arch = "x86_64")]
12122    return row_dot_avx2(row, act);
12123    #[allow(unreachable_code)]
12124    q8_row_dot_scalar(row, act)
12125}
12126
12127#[allow(dead_code)]
12128fn q8_row_dot_scalar(row: &[u8], act: &SplitAct) -> f32 {
12129    let mut acc = 0i32;
12130    for (k, &b) in row.iter().enumerate() {
12131        acc += (b as i8) as i32 * act.xq[k] as i32;
12132    }
12133    let mut acc = acc as f32 * act.sx;
12134    for &(j, xv) in &act.outliers {
12135        acc += (row[j] as i8) as f32 * xv;
12136    }
12137    acc
12138}
12139
12140/// SDOT row dot with exact outlier correction:
12141/// `dot = sdot(w, xq)·sx + Σ_outl w[j]·x[j]` (then × row_scale by caller).
12142#[cfg(target_arch = "aarch64")]
12143#[inline]
12144fn row_dot_sdot(row: &[u8], act: &SplitAct) -> f32 {
12145    let mut acc = unsafe { dot_i8_sdot(row, &act.xq) } as f32 * act.sx;
12146    for &(j, xv) in &act.outliers {
12147        acc += (row[j] as i8) as f32 * xv;
12148    }
12149    acc
12150}
12151
12152/// One q4 row via SDOT: each 32-group's nibbles unpack to centered i8
12153/// (nib−8 ∈ [−8,7]), int8×int8 `sdot` against the pre-quantized
12154/// activation group, × the group's f16 scale. Returns Σ_g dot_g·s_g;
12155/// the caller multiplies by the activation scale and adds the exact
12156/// outlier terms (port of vmfcore `dot_q4_block_sdot`, +23% measured).
12157/// Nibble order matches the writer: element 2k = low nibble, 2k+1 = high
12158/// → zip(lo,hi) restores flat order.
12159#[cfg(target_arch = "aarch64")]
12160#[target_feature(enable = "neon,dotprod")]
12161unsafe fn dot_q4_row_sdot(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
12162    // SAFETY: callers uphold slice-length contracts (16 packed bytes and
12163    // 2 scale bytes per group; xq.len() == gpr·GROUP_SIZE).
12164    unsafe {
12165        use core::arch::aarch64::*;
12166        use core::arch::asm;
12167        let lomask = vdupq_n_u8(0x0F);
12168        let eight = vdupq_n_s8(8);
12169        let mut acc = 0f32;
12170        for gi in 0..gpr {
12171            let g = g0 + gi;
12172            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
12173            let b = vld1q_u8(packed.as_ptr().add(g * 16));
12174            let lo = vandq_u8(b, lomask);
12175            let hi = vshrq_n_u8::<4>(b);
12176            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
12177            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
12178            let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
12179            let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
12180            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
12181            asm!(
12182                "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
12183                "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
12184                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
12185                e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
12186                options(pure, nomem, nostack),
12187            );
12188            acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
12189        }
12190        acc
12191    }
12192}
12193
12194/// Two-activation q4 row via SDOT: the nibble unpack (the expensive
12195/// part) happens ONCE per group; both pre-quantized activations are
12196/// dotted against the same centered i8 registers. Per-lane math matches
12197/// `dot_q4_row_sdot` exactly.
12198#[cfg(target_arch = "aarch64")]
12199#[target_feature(enable = "neon,dotprod")]
12200unsafe fn dot_q4_row_sdot2(
12201    packed: &[u8],
12202    scales: &[u8],
12203    g0: usize,
12204    gpr: usize,
12205    xq1: &[i8],
12206    xq2: &[i8],
12207) -> (f32, f32) {
12208    // SAFETY: callers uphold slice-length contracts (16 packed bytes and
12209    // 2 scale bytes per group; xq*.len() == gpr·GROUP_SIZE).
12210    unsafe {
12211        use core::arch::aarch64::*;
12212        use core::arch::asm;
12213        let lomask = vdupq_n_u8(0x0F);
12214        let eight = vdupq_n_s8(8);
12215        let (mut acc1, mut acc2) = (0f32, 0f32);
12216        for gi in 0..gpr {
12217            let g = g0 + gi;
12218            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
12219            let b = vld1q_u8(packed.as_ptr().add(g * 16));
12220            let lo = vandq_u8(b, lomask);
12221            let hi = vshrq_n_u8::<4>(b);
12222            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
12223            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
12224            let x10 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE));
12225            let x11 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE + 16));
12226            let x20 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE));
12227            let x21 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE + 16));
12228            let (mut a0, mut a1, mut b0, mut b1) = (
12229                vdupq_n_s32(0),
12230                vdupq_n_s32(0),
12231                vdupq_n_s32(0),
12232                vdupq_n_s32(0),
12233            );
12234            asm!(
12235                "sdot {a0:v}.4s, {e0:v}.16b, {x10:v}.16b",
12236                "sdot {a1:v}.4s, {e1:v}.16b, {x11:v}.16b",
12237                "sdot {b0:v}.4s, {e0:v}.16b, {x20:v}.16b",
12238                "sdot {b1:v}.4s, {e1:v}.16b, {x21:v}.16b",
12239                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
12240                b0 = inout(vreg) b0, b1 = inout(vreg) b1,
12241                e0 = in(vreg) e0, e1 = in(vreg) e1,
12242                x10 = in(vreg) x10, x11 = in(vreg) x11,
12243                x20 = in(vreg) x20, x21 = in(vreg) x21,
12244                options(pure, nomem, nostack),
12245            );
12246            acc1 += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
12247            acc2 += vaddvq_s32(vaddq_s32(b0, b1)) as f32 * s;
12248        }
12249        (acc1, acc2)
12250    }
12251}
12252
12253// ───────────────────── fused int8 kernels ─────────────────────
12254
12255/// `acc += w · row` where the row is centered i8 — NEON widen+fma on
12256/// aarch64, scalar elsewhere. The KV-cache q8 value path rides on this.
12257#[inline]
12258pub(crate) fn axpy_i8_f32(acc: &mut [f32], row: &[i8], w: f32) {
12259    #[cfg(target_arch = "aarch64")]
12260    unsafe {
12261        return axpy_i8_f32_neon(acc, row, w);
12262    }
12263    #[cfg(target_arch = "x86_64")]
12264    if avx2_enabled() {
12265        return unsafe { axpy_i8_f32_avx2(acc, row, w) };
12266    }
12267    #[allow(unreachable_code)]
12268    {
12269        for (a, &b) in acc.iter_mut().zip(row) {
12270            *a += w * b as f32;
12271        }
12272    }
12273}
12274
12275/// i8→f32 axpy via AVX2/FMA (x86 mirror of `axpy_i8_f32_neon`).
12276#[cfg(target_arch = "x86_64")]
12277#[target_feature(enable = "avx2,fma")]
12278unsafe fn axpy_i8_f32_avx2(acc: &mut [f32], row: &[i8], w: f32) {
12279    // SAFETY: callers uphold slice-length contracts (see call sites).
12280    unsafe {
12281        use core::arch::x86_64::*;
12282        let n = acc.len().min(row.len());
12283        let ap = acc.as_mut_ptr();
12284        let rp = row.as_ptr();
12285        let wv = _mm256_set1_ps(w);
12286        let mut j = 0usize;
12287        while j + 16 <= n {
12288            let rb = _mm_loadu_si128(rp.add(j) as *const __m128i);
12289            let lo = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(rb));
12290            let hi = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(rb)));
12291            let v0 = _mm256_fmadd_ps(wv, lo, _mm256_loadu_ps(ap.add(j)));
12292            let v1 = _mm256_fmadd_ps(wv, hi, _mm256_loadu_ps(ap.add(j + 8)));
12293            _mm256_storeu_ps(ap.add(j), v0);
12294            _mm256_storeu_ps(ap.add(j + 8), v1);
12295            j += 16;
12296        }
12297        while j < n {
12298            *ap.add(j) += w * (*rp.add(j)) as f32;
12299            j += 1;
12300        }
12301    }
12302}
12303
12304#[cfg(target_arch = "aarch64")]
12305#[target_feature(enable = "neon")]
12306unsafe fn axpy_i8_f32_neon(acc: &mut [f32], row: &[i8], w: f32) {
12307    // SAFETY: callers uphold slice-length contracts (see call sites).
12308    unsafe {
12309        use core::arch::aarch64::*;
12310        let n = acc.len().min(row.len());
12311        let ap = acc.as_mut_ptr();
12312        let rp = row.as_ptr();
12313        let wv = vdupq_n_f32(w);
12314        let mut j = 0usize;
12315        while j + 16 <= n {
12316            let rb = vld1q_s8(rp.add(j));
12317            let lo = vmovl_s8(vget_low_s8(rb));
12318            let hi = vmovl_s8(vget_high_s8(rb));
12319            for (off, half) in [(0, lo), (8, hi)] {
12320                let f0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(half)));
12321                let f1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(half)));
12322                let o = j + off;
12323                vst1q_f32(ap.add(o), vfmaq_f32(vld1q_f32(ap.add(o)), wv, f0));
12324                vst1q_f32(ap.add(o + 4), vfmaq_f32(vld1q_f32(ap.add(o + 4)), wv, f1));
12325            }
12326            j += 16;
12327        }
12328        while j < n {
12329            *ap.add(j) += w * (*rp.add(j)) as f32;
12330            j += 1;
12331        }
12332    }
12333}
12334
12335/// i8 row · f32 x. NEON on aarch64 (ported from vmfcore `dot_i8_f32_neon`,
12336/// ≈9× scalar), scalar elsewhere.
12337#[inline]
12338pub(crate) fn dot_i8_f32(w: &[u8], x: &[f32]) -> f32 {
12339    #[cfg(target_arch = "aarch64")]
12340    unsafe {
12341        return dot_i8_f32_neon(w, x);
12342    }
12343    #[cfg(target_arch = "x86_64")]
12344    if avx2_enabled() {
12345        return unsafe { dot_i8_f32_avx2(w, x) };
12346    }
12347    #[allow(unreachable_code)]
12348    {
12349        let mut sum = 0.0f32;
12350        for (j, &b) in w.iter().enumerate() {
12351            sum += (b as i8) as f32 * x[j];
12352        }
12353        sum
12354    }
12355}
12356
12357/// i8 row · (x ⊙ col_field) — the q8_2f row dot with the θ col-field
12358/// folded into the product (no prescaled copy of x). NEON on aarch64,
12359/// scalar elsewhere. Used by the active-neuron path `row_dot`.
12360#[inline]
12361fn dot_i8_col_f32(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
12362    #[cfg(target_arch = "aarch64")]
12363    unsafe {
12364        return dot_i8_col_f32_neon(w, x, col);
12365    }
12366    #[allow(unreachable_code)]
12367    {
12368        let mut sum = 0.0f32;
12369        for (j, &b) in w.iter().enumerate() {
12370            sum += (b as i8) as f32 * x[j] * col[j];
12371        }
12372        sum
12373    }
12374}
12375
12376#[cfg(target_arch = "aarch64")]
12377#[target_feature(enable = "neon")]
12378unsafe fn dot_i8_col_f32_neon(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
12379    // SAFETY: callers uphold slice-length contracts (see call sites).
12380    unsafe {
12381        use core::arch::aarch64::*;
12382        let n = x.len();
12383        let wp = w.as_ptr() as *const i8;
12384        let xp = x.as_ptr();
12385        let cp = col.as_ptr();
12386        let (mut a0, mut a1, mut a2, mut a3) = (
12387            vdupq_n_f32(0.0),
12388            vdupq_n_f32(0.0),
12389            vdupq_n_f32(0.0),
12390            vdupq_n_f32(0.0),
12391        );
12392        let mut j = 0usize;
12393        while j + 16 <= n {
12394            let wb = vld1q_s8(wp.add(j));
12395            let lo = vmovl_s8(vget_low_s8(wb));
12396            let hi = vmovl_s8(vget_high_s8(wb));
12397            let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
12398            let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
12399            let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
12400            let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
12401            a0 = vfmaq_f32(
12402                a0,
12403                w0,
12404                vmulq_f32(vld1q_f32(xp.add(j)), vld1q_f32(cp.add(j))),
12405            );
12406            a1 = vfmaq_f32(
12407                a1,
12408                w1,
12409                vmulq_f32(vld1q_f32(xp.add(j + 4)), vld1q_f32(cp.add(j + 4))),
12410            );
12411            a2 = vfmaq_f32(
12412                a2,
12413                w2,
12414                vmulq_f32(vld1q_f32(xp.add(j + 8)), vld1q_f32(cp.add(j + 8))),
12415            );
12416            a3 = vfmaq_f32(
12417                a3,
12418                w3,
12419                vmulq_f32(vld1q_f32(xp.add(j + 12)), vld1q_f32(cp.add(j + 12))),
12420            );
12421            j += 16;
12422        }
12423        let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
12424        while j < n {
12425            sum += (*wp.add(j)) as f32 * *xp.add(j) * *cp.add(j);
12426            j += 1;
12427        }
12428        sum
12429    }
12430}
12431
12432#[cfg(target_arch = "aarch64")]
12433#[target_feature(enable = "neon")]
12434unsafe fn dot_i8_f32_neon(w: &[u8], x: &[f32]) -> f32 {
12435    // SAFETY: callers uphold slice-length contracts (see call sites).
12436    unsafe {
12437        use core::arch::aarch64::*;
12438        let n = x.len();
12439        let wp = w.as_ptr() as *const i8;
12440        let xp = x.as_ptr();
12441        let (mut a0, mut a1, mut a2, mut a3) = (
12442            vdupq_n_f32(0.0),
12443            vdupq_n_f32(0.0),
12444            vdupq_n_f32(0.0),
12445            vdupq_n_f32(0.0),
12446        );
12447        let mut j = 0usize;
12448        while j + 16 <= n {
12449            let wb = vld1q_s8(wp.add(j));
12450            let lo = vmovl_s8(vget_low_s8(wb));
12451            let hi = vmovl_s8(vget_high_s8(wb));
12452            let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
12453            let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
12454            let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
12455            let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
12456            a0 = vfmaq_f32(a0, w0, vld1q_f32(xp.add(j)));
12457            a1 = vfmaq_f32(a1, w1, vld1q_f32(xp.add(j + 4)));
12458            a2 = vfmaq_f32(a2, w2, vld1q_f32(xp.add(j + 8)));
12459            a3 = vfmaq_f32(a3, w3, vld1q_f32(xp.add(j + 12)));
12460            j += 16;
12461        }
12462        let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
12463        while j < n {
12464            sum += (*wp.add(j)) as f32 * *xp.add(j);
12465            j += 1;
12466        }
12467        sum
12468    }
12469}
12470
12471#[allow(clippy::too_many_arguments)]
12472fn qmatvec(
12473    q: &[u8],
12474    rep: &[u8],
12475    row_scale: &[f32],
12476    x: &[f32],
12477    col_field: &[f32],
12478    dtype: TensorDtype,
12479    rows: usize,
12480    cols: usize,
12481    out: &mut [f32],
12482    pool: Option<&Pool>,
12483) {
12484    debug_assert_eq!(out.len(), rows);
12485    #[cfg(not(target_arch = "aarch64"))]
12486    let _ = rep;
12487
12488    #[cfg(target_arch = "aarch64")]
12489    if sdot_enabled() {
12490        let act = if dtype == TensorDtype::Q8_2f {
12491            split_act_q8_2f(x, col_field)
12492        } else {
12493            split_act(x)
12494        };
12495        let out_addr = SendMut(out.as_mut_ptr());
12496        let run_range = |start: usize, end: usize| {
12497            q8_range_sdot(q, rep, row_scale, &act, cols, out_addr, start, end)
12498        };
12499        match pool {
12500            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
12501            _ => run_range(0, rows),
12502        }
12503        return;
12504    }
12505    // x86 A8W8 via AVX2 maddubs — same quantized-activation contract as
12506    // the SDOT path (CMF_AVX2=0 keeps the exact i8×f32 loop).
12507    #[cfg(target_arch = "x86_64")]
12508    if avx2_a8w8_enabled() {
12509        let act = if dtype == TensorDtype::Q8_2f {
12510            split_act_q8_2f(x, col_field)
12511        } else {
12512            split_act(x)
12513        };
12514        let out_addr = SendMut(out.as_mut_ptr());
12515        let run_range = |start: usize, end: usize| {
12516            q8_range_avx2(q, row_scale, &act, cols, out_addr, start, end)
12517        };
12518        match pool {
12519            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
12520            _ => run_range(0, rows),
12521        }
12522        return;
12523    }
12524
12525    prescale_with(x, col_field, dtype, 1, |xs| {
12526        let out_addr = SendMut(out.as_mut_ptr());
12527        let run_range = move |start: usize, end: usize| {
12528            for o in start..end {
12529                let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
12530                // SAFETY: disjoint row ranges per worker.
12531                unsafe { *out_addr.at(o) = v };
12532            }
12533        };
12534        match pool {
12535            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
12536            _ => run_range(0, rows),
12537        }
12538    });
12539}
12540
12541#[allow(clippy::too_many_arguments)]
12542fn qmatvec2(
12543    q: &[u8],
12544    row_scale: &[f32],
12545    x1: &[f32],
12546    x2: &[f32],
12547    col_field: &[f32],
12548    dtype: TensorDtype,
12549    rows: usize,
12550    cols: usize,
12551    o1: &mut [f32],
12552    o2: &mut [f32],
12553    pool: Option<&Pool>,
12554) {
12555    #[cfg(target_arch = "aarch64")]
12556    if sdot_enabled() {
12557        let a1s = if dtype == TensorDtype::Q8_2f {
12558            split_act_q8_2f(x1, col_field)
12559        } else {
12560            split_act(x1)
12561        };
12562        let a2s = if dtype == TensorDtype::Q8_2f {
12563            split_act_q8_2f(x2, col_field)
12564        } else {
12565            split_act(x2)
12566        };
12567        let p1 = SendMut(o1.as_mut_ptr());
12568        let p2 = SendMut(o2.as_mut_ptr());
12569        let run_range = |start: usize, end: usize| {
12570            q8_range2_sdot(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
12571        };
12572        match pool {
12573            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
12574            _ => run_range(0, rows),
12575        }
12576        return;
12577    }
12578    #[cfg(target_arch = "x86_64")]
12579    if avx2_a8w8_enabled() {
12580        let a1s = if dtype == TensorDtype::Q8_2f {
12581            split_act_q8_2f(x1, col_field)
12582        } else {
12583            split_act(x1)
12584        };
12585        let a2s = if dtype == TensorDtype::Q8_2f {
12586            split_act_q8_2f(x2, col_field)
12587        } else {
12588            split_act(x2)
12589        };
12590        let p1 = SendMut(o1.as_mut_ptr());
12591        let p2 = SendMut(o2.as_mut_ptr());
12592        let run_range = |start: usize, end: usize| {
12593            q8_range2_avx2(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
12594        };
12595        match pool {
12596            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
12597            _ => run_range(0, rows),
12598        }
12599        return;
12600    }
12601
12602    prescale_with(x1, col_field, dtype, 1, |x1s| {
12603        prescale_with(x2, col_field, dtype, 2, |x2s| {
12604            let p1 = SendMut(o1.as_mut_ptr());
12605            let p2 = SendMut(o2.as_mut_ptr());
12606            let run_range = move |start: usize, end: usize| {
12607                for o in start..end {
12608                    let row = &q[o * cols..(o + 1) * cols];
12609                    let s1 = dot_i8_f32(row, x1s) * row_scale[o];
12610                    let s2 = dot_i8_f32(row, x2s) * row_scale[o];
12611                    // SAFETY: disjoint row ranges per worker.
12612                    unsafe {
12613                        *p1.at(o) = s1;
12614                        *p2.at(o) = s2;
12615                    }
12616                }
12617            };
12618            match pool {
12619                Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
12620                _ => run_range(0, rows),
12621            }
12622        });
12623    });
12624}
12625
12626#[derive(Clone, Copy)]
12627struct SendMut(*mut f32);
12628unsafe impl Send for SendMut {}
12629unsafe impl Sync for SendMut {}
12630
12631impl SendMut {
12632    #[inline]
12633    fn at(self, i: usize) -> *mut f32 {
12634        unsafe { self.0.add(i) }
12635    }
12636}
12637
12638#[cfg(test)]
12639mod tests {
12640    /// `q8_round` must be `round().clamp(±127) as i8` bit for bit: every
12641    /// half-integer, their neighbours one ulp either side, the clamp
12642    /// boundary, huge values, infinities and NaN, plus a dense sweep.
12643    #[test]
12644    fn q8_round_is_round_clamp() {
12645        let reference = |t: f32| t.round().clamp(-127.0, 127.0) as i8;
12646        let mut probe = vec![
12647            0.0f32,
12648            -0.0,
12649            f32::NAN,
12650            f32::INFINITY,
12651            f32::NEG_INFINITY,
12652            f32::MAX,
12653            f32::MIN,
12654            1e30,
12655            -1e30,
12656            f32::MIN_POSITIVE,
12657            -f32::MIN_POSITIVE,
12658        ];
12659        for k in -300i32..=300 {
12660            let h = k as f32 * 0.5;
12661            let up = f32::from_bits(h.to_bits() + 1);
12662            let down = f32::from_bits(h.to_bits().wrapping_sub(1));
12663            for t in [h, up, down] {
12664                probe.push(t);
12665                probe.push(-t);
12666            }
12667        }
12668        let mut t = -140.0f32;
12669        while t < 140.0 {
12670            probe.push(t);
12671            t += 0.000_731;
12672        }
12673        for t in probe {
12674            assert_eq!(q8_round(t), reference(t), "t = {t:e} ({:#x})", t.to_bits());
12675        }
12676    }
12677
12678    use super::*;
12679
12680    /// The f32 `matmat` walks the batch four positions a row at a time; every
12681    /// output must still be the scalar-order sum `matvec` gives, bit for bit
12682    /// (a 7-row batch covers the four-wide block and the remainder).
12683    #[test]
12684    fn f32_matmat_equals_matvec_bitwise() {
12685        let (rows, cols, b) = (5usize, 37usize, 7usize);
12686        let w: Vec<f32> = (0..rows * cols).map(|i| ((i * 7919 % 113) as f32 - 56.0) / 37.0).collect();
12687        let xs: Vec<f32> = (0..b * cols).map(|i| ((i * 104_729 % 97) as f32 - 48.0) / 29.0).collect();
12688        let t = QTensor::from_f32(w, rows, cols);
12689        let mut out = vec![0f32; b * rows];
12690        t.matmat(&xs, b, &mut out, None);
12691        for bi in 0..b {
12692            let mut one = vec![0f32; rows];
12693            t.matvec(&xs[bi * cols..(bi + 1) * cols], &mut one, None);
12694            for o in 0..rows {
12695                assert_eq!(out[bi * rows + o].to_bits(), one[o].to_bits(), "bi {bi} row {o}");
12696            }
12697        }
12698    }
12699
12700    #[test]
12701    fn q2tp_i8_dot_matches_exact_on_grid() {
12702        // On-grid activations (±1 → sx=1/127, xq=±127 dequantizes
12703        // exactly, no outliers) must make the integer path agree with
12704        // the exact scalar walk to f32 rounding.
12705        let (rows, cols) = (5, 64);
12706        let gpr = cols / GROUP_SIZE;
12707        // Synthetic codes plane + a flat ladder: scales_into is not under
12708        // test here, so drive dot_q2tp_row_i8 / q2tp_row_exact directly
12709        // with hand-made scales.
12710        let chunks: Vec<u8> = (0..rows * gpr * Q2TP_CHUNK)
12711            .map(|i| (i as u32).wrapping_mul(2654435761) as u8)
12712            .collect();
12713        let scales: Vec<f32> = (0..gpr).map(|g| 0.5 + g as f32 * 0.25).collect();
12714        let x: Vec<f32> = (0..cols)
12715            .map(|i| if i % 3 == 0 { -1.0 } else { 1.0 })
12716            .collect();
12717        let act = split_act(&x);
12718        assert!(
12719            act.outliers.is_empty(),
12720            "on-grid input must have no outliers"
12721        );
12722        let gsum = q1_group_sums(&act.xq, gpr);
12723        for r in 0..rows {
12724            let exact = q2tp_row_exact(&chunks, r, gpr, &x, &scales);
12725            let fast = dot_q2tp_row_i8(&chunks, r, gpr, &act.xq, &gsum, &scales) * act.sx;
12726            assert!(
12727                (exact - fast).abs() <= exact.abs() * 1e-5 + 1e-5,
12728                "row {r}: exact {exact} vs i8 {fast}"
12729            );
12730        }
12731    }
12732
12733    #[cfg(target_arch = "x86_64")]
12734    #[test]
12735    fn q2tp_avx2_dot_matches_scalar_for_random_patterns() {
12736        // Compare the release AVX2 integer dot against the scalar oracle over
12737        // arbitrary packed bytes/activation signs.  This guards the exact
12738        // table-load path used after rejecting a faster-looking decoder whose
12739        // full-checkpoint greedy output drifted.
12740        if !std::arch::is_x86_feature_detected!("avx2") {
12741            return;
12742        }
12743        let mut seed = 0x9e3779b9u32;
12744        let mut next = || {
12745            seed = seed.wrapping_mul(1664525).wrapping_add(1013904223);
12746            seed
12747        };
12748        for _ in 0..20_000 {
12749            let mut ch = [0u8; Q2TP_CHUNK];
12750            let mut x = [0i8; GROUP_SIZE];
12751            for b in &mut ch {
12752                *b = next() as u8;
12753            }
12754            for v in &mut x {
12755                *v = (next() >> 24) as i8;
12756            }
12757            let mut reference = 0i32;
12758            for (k, &b) in ch.iter().enumerate() {
12759                reference += (b & 3) as i32 * x[k * 4] as i32;
12760                reference += ((b >> 2) & 3) as i32 * x[k * 4 + 1] as i32;
12761                reference += ((b >> 4) & 3) as i32 * x[k * 4 + 2] as i32;
12762                reference += ((b >> 6) & 3) as i32 * x[k * 4 + 3] as i32;
12763            }
12764            // SAFETY: guarded by the runtime AVX2 feature check and fixed
12765            // 8-byte/32-byte slice lengths above.
12766            let got = unsafe { q2tp_code_dot_avx2(&ch, &x) };
12767            assert_eq!(got, reference, "packed q2 lane mismatch");
12768        }
12769    }
12770
12771    #[test]
12772    fn q2tp_affine_fuses_half_scale_correction_without_changing_raw_decode() {
12773        let (rows, cols) = (1usize, GROUP_SIZE);
12774        let mut bytes = vec![0u8; Q2TP_CHUNK + 4 + 1];
12775        // Repeating symbols 0,1,2,0 at unit scale.  q2tp's raw B is
12776        // (c-1.5), while the affine Prism operator is (c-1.0).
12777        bytes[..Q2TP_CHUNK].fill(0x24); // codes 0,1,2,0 in LSB-first order
12778        bytes[Q2TP_CHUNK..Q2TP_CHUNK + 2].copy_from_slice(&0u16.to_le_bytes());
12779        bytes[Q2TP_CHUNK + 2..Q2TP_CHUNK + 4].copy_from_slice(&0u16.to_le_bytes());
12780        bytes[Q2TP_CHUNK + 4] = 1; // dtype16 rung 1 = 1.0
12781        let x = vec![1.0f32; cols];
12782        let mut raw = vec![0.0f32; rows];
12783        let mut affine = vec![0.0f32; rows];
12784        q2tp_matvec_for_test(&bytes, &x, rows, cols, &mut raw);
12785        q2tp_affine_matvec_for_test(&bytes, &x, rows, cols, &mut affine);
12786        assert_eq!(raw, vec![-24.0]);
12787        assert_eq!(affine, vec![-8.0]);
12788        assert!((affine[0] - (raw[0] + 0.5 * cols as f32)).abs() < 1e-6);
12789    }
12790
12791    #[test]
12792    fn q8_row_dot_fast_matches_scalar() {
12793        // The per-arch fast dot must agree with the exact scalar oracle
12794        // (same contract the fused q8 FFN arm rides on).
12795        let cols = 96;
12796        let row: Vec<u8> = (0..cols)
12797            .map(|i| ((i * 37 % 251) - 125) as i8 as u8)
12798            .collect();
12799        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.13).sin()).collect();
12800        let act = split_act(&x);
12801        let fast = q8_row_dot(&row, &act);
12802        let scalar = q8_row_dot_scalar(&row, &act);
12803        assert!(
12804            (fast - scalar).abs() <= scalar.abs() * 1e-5 + 1e-5,
12805            "fast {fast} vs scalar {scalar}"
12806        );
12807    }
12808
12809    #[test]
12810    fn f32_matvec_matches_matvec_rows_bitexact() {
12811        let (rows, cols) = (300, 40);
12812        let w: Vec<f32> = (0..rows * cols).map(|i| (i as f32 * 0.017).sin()).collect();
12813        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.05).cos()).collect();
12814        let qt = QTensor::from_f32(w.clone(), rows, cols);
12815
12816        let mut a = vec![0.0f32; rows];
12817        matvec_rows(None, &w, &x, &mut a);
12818        let mut b = vec![0.0f32; rows];
12819        qt.matvec(&x, &mut b, None);
12820        assert_eq!(a, b);
12821    }
12822
12823    #[test]
12824    fn sdot_kernel_exact_on_grid() {
12825        // Activations already on the i8 grid (±1 with amax=1 → sx=1/127,
12826        // xq=±127 dequantizes EXACTLY) → the SDOT path must match the
12827        // exact f32 dot to float rounding. This isolates kernel
12828        // correctness from quantization noise.
12829        eprintln!("sdot_enabled = {}", sdot_enabled());
12830        let (rows, cols) = (9, 80); // odd rows → exercises 4-row + tail
12831        let w: Vec<u8> = (0..rows * cols)
12832            .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
12833            .collect();
12834        let scales: Vec<f32> = (0..rows).map(|o| 0.005 + o as f32 * 0.001).collect();
12835        let x: Vec<f32> = (0..cols)
12836            .map(|i| match i % 3 {
12837                0 => 1.0,
12838                1 => -1.0,
12839                _ => 0.0,
12840            })
12841            .collect();
12842        let mut a = vec![0.0f32; rows];
12843        qmatvec(
12844            &w,
12845            &[],
12846            &scales,
12847            &x,
12848            &[],
12849            TensorDtype::Q8Row,
12850            rows,
12851            cols,
12852            &mut a,
12853            None,
12854        );
12855        for o in 0..rows {
12856            let mut acc = 0.0f32;
12857            for j in 0..cols {
12858                acc += (w[o * cols + j] as i8) as f32 * x[j];
12859            }
12860            let expect = acc * scales[o];
12861            assert!(
12862                (a[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
12863                "row {o}: {} vs {expect}",
12864                a[o]
12865            );
12866        }
12867    }
12868
12869    #[test]
12870    fn q1_tbl_fast_path_matches_reference() {
12871        // gpr = 8 exercises the TBL pair-load fast loop, and the LAST
12872        // row's final 4-tile window trips the 4B-overread guard (the
12873        // payload ends exactly at the last tile) — both paths must
12874        // agree with the dequant reference.
12875        let (rows, cols) = (5, 256);
12876        let gpr = cols / GROUP_SIZE;
12877        let mut bytes = Vec::new();
12878        for t in 0..rows * gpr {
12879            let s = 0.007 + (t % 11) as f32 * 0.004;
12880            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
12881            for j in 0..4 {
12882                bytes.push(((t * 53 + j * 89 + 7) % 249) as u8);
12883            }
12884        }
12885        let x: Vec<f32> = (0..cols)
12886            .map(|i| if (i * 5) % 7 < 3 { 1.0 } else { -1.0 })
12887            .collect();
12888        let mut w = vec![0.0f32; rows * cols];
12889        cortiq_core::quant::dequant_q1(&bytes, &mut w);
12890        let mut got = vec![0.0f32; rows];
12891        q1_matvec(&bytes, &x, rows, cols, &mut got, None);
12892        for o in 0..rows {
12893            let expect: f32 = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
12894            assert!(
12895                (got[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
12896                "row {o}: {} vs {expect}",
12897                got[o]
12898            );
12899        }
12900        // Blocked 1×4 batch (b=5: one quad + remainder) must equal the
12901        // single-matvec path bit-for-bit.
12902        let b = 5usize;
12903        let mut xs_all = Vec::new();
12904        for bi in 0..b {
12905            xs_all.extend(x.iter().map(|v| if bi % 2 == 0 { *v } else { -*v }));
12906        }
12907        let mut mm = vec![0.0f32; b * rows];
12908        q1_matmat(&bytes, &xs_all, b, rows, cols, &mut mm, None);
12909        for bi in 0..b {
12910            let mut single = vec![0.0f32; rows];
12911            q1_matvec(
12912                &bytes,
12913                &xs_all[bi * cols..(bi + 1) * cols],
12914                rows,
12915                cols,
12916                &mut single,
12917                None,
12918            );
12919            assert_eq!(&mm[bi * rows..(bi + 1) * rows], &single[..], "stream {bi}");
12920        }
12921    }
12922
12923    #[test]
12924    fn q1_kernels_match_exact_reference() {
12925        // Synthetic q1 payload: 6-byte tiles [f16 scale][4B bits].
12926        let (rows, cols) = (7, 96);
12927        let gpr = cols / GROUP_SIZE;
12928        let mut bytes = Vec::new();
12929        for t in 0..rows * gpr {
12930            let s = 0.01 + (t % 13) as f32 * 0.003;
12931            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
12932            for j in 0..4 {
12933                bytes.push(((t * 31 + j * 97) % 251) as u8);
12934            }
12935        }
12936        // On-grid activations (±1, amax 1) → the SDOT path is exact.
12937        let x: Vec<f32> = (0..cols)
12938            .map(|i| if i % 3 == 0 { 1.0 } else { -1.0 })
12939            .collect();
12940        // Reference through the core dequant.
12941        let mut w = vec![0.0f32; rows * cols];
12942        cortiq_core::quant::dequant_q1(&bytes, &mut w);
12943        let mut expect = vec![0.0f32; rows];
12944        for o in 0..rows {
12945            expect[o] = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
12946        }
12947        let mut got = vec![0.0f32; rows];
12948        q1_matvec(&bytes, &x, rows, cols, &mut got, None);
12949        for o in 0..rows {
12950            assert!(
12951                (got[o] - expect[o]).abs() < 1e-3 * expect[o].abs().max(1e-3),
12952                "row {o}: {} vs {}",
12953                got[o],
12954                expect[o]
12955            );
12956        }
12957        // Pair and batch paths agree with the single path.
12958        let x2: Vec<f32> = x.iter().map(|v| -v).collect();
12959        let (mut a1, mut a2) = (vec![0.0f32; rows], vec![0.0f32; rows]);
12960        q1_matvec2(&bytes, &x, &x2, rows, cols, &mut a1, &mut a2, None);
12961        assert_eq!(a1, got);
12962        let mut xs = x.clone();
12963        xs.extend_from_slice(&x2);
12964        let mut mm = vec![0.0f32; 2 * rows];
12965        q1_matmat(&bytes, &xs, 2, rows, cols, &mut mm, None);
12966        assert_eq!(&mm[..rows], got.as_slice());
12967        assert_eq!(&mm[rows..], a2.as_slice());
12968    }
12969
12970    /// A Bonsai-shaped q1 payload (one f16 scale per 128-group, stored in
12971    /// all four tiles; pseudo-random signs) and its q1g repack.
12972    fn q1_q1g_pair(rows: usize, cols: usize, salt: usize) -> (Vec<u8>, Vec<u8>) {
12973        let tiles = rows * cols / GROUP_SIZE;
12974        let mut q1 = Vec::with_capacity(tiles * Q1_TILE);
12975        for t in 0..tiles {
12976            let s = 0.004 + ((t / 4 * 37 + salt) % 29) as f32 * 0.0011;
12977            q1.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
12978            let bits = ((t * 2_654_435_761 + salt * 40_503) as u32).rotate_left((t % 31) as u32);
12979            q1.extend_from_slice(&bits.to_le_bytes());
12980        }
12981        let rp = cortiq_core::quant::repack_q1_to_q1g(&q1, rows, cols, 0.0);
12982        assert_eq!(rp.differing, 0, "uniform scales repack exactly");
12983        (q1, rp.data)
12984    }
12985
12986    /// Activations off the int8 grid with two outlier channels, so the
12987    /// A8W8 path quantizes AND takes the exact outlier correction.
12988    fn q1g_test_x(cols: usize, salt: usize) -> Vec<f32> {
12989        let mut x: Vec<f32> = (0..cols)
12990            .map(|i| ((i * 13 + salt * 7) as f32 * 0.071).sin() * 0.9)
12991            .collect();
12992        x[3 % cols] = 41.0;
12993        x[(cols / 2 + 5) % cols] = -37.5;
12994        x
12995    }
12996
12997    /// Every CPU row kernel of q1g returns the q1 kernel's f32 BIT FOR BIT
12998    /// on a uniform-scale repack: A8W8 rows (SIMD of this machine, and the
12999    /// x86 AVX2/VNNI twins directly), the exact scalar rows (CMF_SDOT=0
13000    /// contract), the outlier term, the drivers (matvec, matvec2, matmat
13001    /// with the blocked 1×4 quad + remainder) — including gpr = 4 and the
13002    /// last row, where q1's pair load falls back to its scalar tail.
13003    #[test]
13004    fn q1g_cpu_kernels_equal_q1_bit_for_bit() {
13005        for &(rows, cols) in &[(3usize, 128usize), (5, 256), (9, 2048), (2, 6144)] {
13006            let gpr = cols / GROUP_SIZE;
13007            let (q1, q1g) = q1_q1g_pair(rows, cols, rows + cols);
13008            let (wa, wb) = (BinW::Q1(&q1), BinW::q1g(&q1g, rows, cols));
13009            let x = q1g_test_x(cols, rows);
13010            let act = split_act(&x);
13011            assert!(!act.outliers.is_empty(), "test wants the outlier path");
13012            let gsum = q1_group_sums(&act.xq, gpr);
13013            for r in 0..rows {
13014                let (a, b) = (wa.row_exact(r, gpr, &x), wb.row_exact(r, gpr, &x));
13015                assert_eq!(a.to_bits(), b.to_bits(), "{rows}x{cols} exact row {r}");
13016                let (a, b) = (
13017                    wa.dot_i8(r, gpr, &act.xq, &gsum),
13018                    wb.dot_i8(r, gpr, &act.xq, &gsum),
13019                );
13020                assert_eq!(a.to_bits(), b.to_bits(), "{rows}x{cols} i8 row {r}");
13021                for j in [0, 31, 32, 127, cols - 1] {
13022                    let (a, b) = (wa.outlier(r, gpr, j), wb.outlier(r, gpr, j));
13023                    assert_eq!(
13024                        (a.0.to_bits(), a.1.to_bits()),
13025                        (b.0.to_bits(), b.1.to_bits())
13026                    );
13027                }
13028                #[cfg(target_arch = "x86_64")]
13029                if avx2_enabled() {
13030                    let sc = cortiq_core::quant::q1g_scale_offset(rows, cols);
13031                    // SAFETY: AVX2 checked; payloads sized for (rows, cols).
13032                    unsafe {
13033                        let a = dot_q1_row_avx2(&q1, r, gpr, &act.xq, &gsum);
13034                        let b = dot_q1g_row_avx2(&q1g, sc, r, gpr, &act.xq, &gsum);
13035                        assert_eq!(a.to_bits(), b.to_bits(), "avx2 row {r}");
13036                    }
13037                    if std::is_x86_feature_detected!("avx512vnni")
13038                        && std::is_x86_feature_detected!("avx512bw")
13039                        && std::is_x86_feature_detected!("avx512vl")
13040                    {
13041                        // SAFETY: VNNI family detected just above.
13042                        unsafe {
13043                            let a = dot_q1_row_vnni(&q1, r, gpr, &act.xq, &gsum);
13044                            let b = dot_q1g_row_vnni(&q1g, sc, r, gpr, &act.xq, &gsum);
13045                            assert_eq!(a.to_bits(), b.to_bits(), "vnni row {r}");
13046                        }
13047                    }
13048                }
13049            }
13050            // Drivers: single, pair, batch (b = 5: one blocked quad + one).
13051            let (mut ga, mut gb) = (vec![0f32; rows], vec![0f32; rows]);
13052            bin_matvec(wa, &x, rows, cols, &mut ga, None);
13053            bin_matvec(wb, &x, rows, cols, &mut gb, None);
13054            assert_eq!(ga, gb, "{rows}x{cols} matvec");
13055            let x2 = q1g_test_x(cols, rows + 1);
13056            let (mut a1, mut a2, mut b1, mut b2) = (
13057                vec![0f32; rows],
13058                vec![0f32; rows],
13059                vec![0f32; rows],
13060                vec![0f32; rows],
13061            );
13062            bin_matvec2(wa, &x, &x2, rows, cols, &mut a1, &mut a2, None);
13063            bin_matvec2(wb, &x, &x2, rows, cols, &mut b1, &mut b2, None);
13064            assert_eq!((&a1, &a2), (&b1, &b2), "{rows}x{cols} matvec2");
13065            let b = 5usize;
13066            let xs: Vec<f32> = (0..b).flat_map(|k| q1g_test_x(cols, k * 3 + 2)).collect();
13067            let (mut ma, mut mb) = (vec![0f32; b * rows], vec![0f32; b * rows]);
13068            bin_matmat(wa, &xs, b, rows, cols, &mut ma, None);
13069            bin_matmat(wb, &xs, b, rows, cols, &mut mb, None);
13070            assert_eq!(ma, mb, "{rows}x{cols} matmat");
13071            // The blocked 1×4 kernels directly (the driver only takes them
13072            // under CMF_X86_BLOCKED / SDOT, a process-wide switch).
13073            let acts: Vec<(SplitAct, Vec<i32>)> = (0..4)
13074                .map(|k| {
13075                    let a = split_act(&xs[k * cols..(k + 1) * cols]);
13076                    let g = q1_group_sums(&a.xq, gpr);
13077                    (a, g)
13078                })
13079                .collect();
13080            let xq4: [&[i8]; 4] = std::array::from_fn(|k| acts[k].0.xq.as_slice());
13081            let gs4: [&[i32]; 4] = std::array::from_fn(|k| acts[k].1.as_slice());
13082            for r in 0..rows {
13083                #[cfg(target_arch = "x86_64")]
13084                if avx2_enabled() {
13085                    // SAFETY: AVX2 checked; payloads sized for (rows, cols).
13086                    let (a, b) = unsafe {
13087                        (
13088                            wa.dot_1x4_x86(r, gpr, xq4, gs4),
13089                            wb.dot_1x4_x86(r, gpr, xq4, gs4),
13090                        )
13091                    };
13092                    assert_eq!(a.map(f32::to_bits), b.map(f32::to_bits), "1x4 row {r}");
13093                }
13094                #[cfg(target_arch = "aarch64")]
13095                if sdot_enabled() {
13096                    // SAFETY: SDOT checked; payloads sized for (rows, cols).
13097                    let (a, b) = unsafe {
13098                        (
13099                            wa.dot_1x4_sdot(r, gpr, xq4, gs4),
13100                            wb.dot_1x4_sdot(r, gpr, xq4, gs4),
13101                        )
13102                    };
13103                    assert_eq!(a.map(f32::to_bits), b.map(f32::to_bits), "1x4 row {r}");
13104                }
13105                let _ = (r, &xq4, &gs4);
13106            }
13107        }
13108    }
13109
13110    /// CPU kernel A/B, q1 vs q1g matvec (A8W8 path of this machine, the
13111    /// exact path too), Bonsai-1.7B decode shapes, arms alternated every
13112    /// round, medians printed. `CMF_THREADS` sets the pool (default 8).
13113    /// `cargo test -p cortiq-engine --release --lib q1g_cpu_matvec_speed --
13114    /// --ignored --nocapture`
13115    #[test]
13116    #[ignore]
13117    fn q1g_cpu_matvec_speed() {
13118        use crate::pool::Pool;
13119        let threads = std::env::var("CMF_THREADS")
13120            .ok()
13121            .and_then(|v| v.parse().ok())
13122            .unwrap_or(8usize);
13123        let pool = Pool::new(threads);
13124        for &(name, rows, cols, reps) in &[
13125            ("gate/up_proj", 6144usize, 2048usize, 41usize),
13126            ("down_proj", 2048, 6144, 41),
13127            ("lm_head", 151_669, 2048, 9),
13128        ] {
13129            let (q1, q1g) = q1_q1g_pair(rows, cols, 3);
13130            let (wa, wb) = (BinW::Q1(&q1), BinW::q1g(&q1g, rows, cols));
13131            let x = q1g_test_x(cols, 1);
13132            let mut out = vec![0f32; rows];
13133            let mut time = |w: BinW| {
13134                let t0 = std::time::Instant::now();
13135                bin_matvec(w, &x, rows, cols, &mut out, Some(&pool));
13136                t0.elapsed().as_secs_f64()
13137            };
13138            let _ = (time(wa), time(wb));
13139            let (mut ta, mut tb) = (Vec::new(), Vec::new());
13140            for r in 0..reps {
13141                if r % 2 == 0 {
13142                    ta.push(time(wa));
13143                    tb.push(time(wb));
13144                } else {
13145                    tb.push(time(wb));
13146                    ta.push(time(wa));
13147                }
13148            }
13149            ta.sort_by(|a, b| a.partial_cmp(b).unwrap());
13150            tb.sort_by(|a, b| a.partial_cmp(b).unwrap());
13151            let (ma, mb) = (ta[reps / 2], tb[reps / 2]);
13152            eprintln!(
13153                "cpu x{threads} a8w8={} {name:>12} {rows}x{cols}: q1 {:8.1} us | q1g {:8.1} us | q1g/q1 time {:.3}",
13154                a8w8_enabled(),
13155                ma * 1e6,
13156                mb * 1e6,
13157                mb / ma
13158            );
13159        }
13160    }
13161
13162    /// The mapped QTensor paths a real Bonsai q1g file takes — matvec, the
13163    /// multi-projection job (all-q1g and q1/q1g mixed), the fused gate/up
13164    /// with SiLU, matmat, matvec2 and the embedding row — equal the q1
13165    /// tensor they were repacked from, bit for bit.
13166    #[test]
13167    fn q1g_mapped_tensor_paths_equal_q1() {
13168        use crate::pool::Pool;
13169        use cortiq_core::{CmfHeader, CmfModel, TensorSpec};
13170        let (rows, cols) = (300usize, 256usize);
13171        let (gq1, gq1g) = q1_q1g_pair(rows, cols, 1);
13172        let (uq1, uq1g) = q1_q1g_pair(rows, cols, 2);
13173        let header: CmfHeader = serde_json::from_value(serde_json::json!({
13174            "format": "cmf", "version": cortiq_core::CMF_VERSION,
13175            "arch": {"arch_name": "tiny-q1g", "hidden_size": cols,
13176                "intermediate_size": rows, "num_layers": 1, "num_attention_heads": 2,
13177                "num_kv_heads": 1, "head_dim": 128, "vocab_size": rows,
13178                "layer_types": ["FullAttention"], "rms_norm_eps": 1e-6,
13179                "max_position_embeddings": 8},
13180            "quant_type": "VBIT"
13181        }))
13182        .unwrap();
13183        let spec = |name: &str, dtype, data: &Vec<u8>| TensorSpec {
13184            name: name.into(),
13185            dtype,
13186            shape: vec![rows, cols],
13187            data: data.clone(),
13188        };
13189        let specs = [
13190            spec("g1", TensorDtype::Q1, &gq1),
13191            spec("u1", TensorDtype::Q1, &uq1),
13192            spec("gg", TensorDtype::Q1G, &gq1g),
13193            spec("ug", TensorDtype::Q1G, &uq1g),
13194        ];
13195        let dir = std::env::temp_dir().join(format!("cmf-q1g-mapped-{}", std::process::id()));
13196        std::fs::create_dir_all(&dir).unwrap();
13197        let path = dir.join("m.cmf");
13198        CmfModel::write(&path, &header, &specs, None, None).unwrap();
13199        let model = Arc::new(CmfModel::open(&path).unwrap());
13200        assert!(model.verify().is_empty(), "{:?}", model.verify());
13201        let t = |n: &str| QTensor::from_model(&model, n).unwrap();
13202        let (g1, u1, gg, ug) = (t("g1"), t("u1"), t("gg"), t("ug"));
13203        assert_eq!(gg.model_dtype(), Some(TensorDtype::Q1G));
13204        assert!(gg.is_q1() && gg.graph_weight().map(|w| w.2) == Some(10));
13205        let pool = Pool::new(3);
13206        let x = q1g_test_x(cols, 4);
13207        let run = |a: &QTensor| {
13208            let mut o = vec![0f32; rows];
13209            a.matvec(&x, &mut o, Some(&pool));
13210            o
13211        };
13212        assert_eq!(run(&g1), run(&gg), "matvec");
13213        let many = |a: &QTensor, b: &QTensor| {
13214            let (mut oa, mut ob) = (vec![0f32; rows], vec![0f32; rows]);
13215            QTensor::matvec_many([a, b], &x, [&mut oa, &mut ob], Some(&pool));
13216            (oa, ob)
13217        };
13218        let want = many(&g1, &u1);
13219        assert_eq!(want, many(&gg, &ug), "matvec_many q1g");
13220        assert_eq!(want, many(&g1, &ug), "matvec_many mixed");
13221        assert_eq!(want, (run(&g1), run(&u1)), "matvec_many == singles");
13222        let fused = |g: &QTensor, u: &QTensor| {
13223            let mut o = vec![0f32; rows];
13224            let ok = QTensor::matvec_silu_mul(g, u, &x, &mut o, Some(&pool));
13225            (ok, o)
13226        };
13227        let (ok_a, fa) = fused(&g1, &u1);
13228        let (ok_b, fb) = fused(&gg, &ug);
13229        assert_eq!(ok_a, ok_b, "q1g takes the fused arm whenever q1 does");
13230        assert_eq!(fa, fb, "fused gate/up");
13231        let b = 6usize;
13232        let xs: Vec<f32> = (0..b).flat_map(|k| q1g_test_x(cols, k + 9)).collect();
13233        let mm = |a: &QTensor| {
13234            let mut o = vec![0f32; b * rows];
13235            a.matmat(&xs, b, &mut o, Some(&pool));
13236            o
13237        };
13238        assert_eq!(mm(&g1), mm(&gg), "matmat");
13239        let x2 = q1g_test_x(cols, 77);
13240        let pair = |a: &QTensor| {
13241            let (mut o1, mut o2) = (vec![0f32; rows], vec![0f32; rows]);
13242            a.matvec2(&x, &x2, &mut o1, &mut o2, Some(&pool));
13243            (o1, o2)
13244        };
13245        assert_eq!(pair(&g1), pair(&gg), "matvec2");
13246        for r in [0, 1, rows / 2, rows - 1] {
13247            let (mut ra, mut rb) = (vec![0f32; cols], vec![0f32; cols]);
13248            g1.row_f32(r, &mut ra);
13249            gg.row_f32(r, &mut rb);
13250            assert_eq!(ra, rb, "embedding row {r}");
13251        }
13252        drop((g1, u1, gg, ug, model));
13253        let _ = std::fs::remove_dir_all(&dir);
13254    }
13255
13256    /// A `blr` tensor read from a real CMF file: it maps to the factored
13257    /// arm (not a dense copy), declines every device-graph accessor, and
13258    /// every QTensor entry point a layer may take — matvec, matvec2,
13259    /// matmat (table loop and GEMM sizes), the multi-tensor job mixed with
13260    /// a q8_row sibling, the fused gate/up refusal, rows and columns — equals
13261    /// the dense reference built by the core dequant.
13262    #[test]
13263    fn blr_mapped_tensor_paths_equal_dense() {
13264        use crate::pool::Pool;
13265        use cortiq_core::{CmfHeader, CmfModel, TensorSpec};
13266        let (rows, cols, r) = (96usize, 128usize, 16usize);
13267        let mut rng = crate::blr::tests::Rng(41);
13268        let gate = crate::blr::tests::make(&mut rng, rows, cols, r, 2);
13269        let up = crate::blr::tests::make(&mut rng, rows, cols, r, 1);
13270        // q8_row sibling: [int8 rows·cols][f16 row scale · rows].
13271        let mut q8: Vec<u8> = (0..rows * cols)
13272            .map(|i| (((i * 89) % 253) as i32 - 126) as i8 as u8)
13273            .collect();
13274        for o in 0..rows {
13275            q8.extend_from_slice(&cortiq_core::quant::f32_to_f16(0.002 + o as f32 * 1e-4).to_le_bytes());
13276        }
13277        let header: CmfHeader = serde_json::from_value(serde_json::json!({
13278            "format": "cmf", "version": cortiq_core::CMF_VERSION,
13279            "arch": {"arch_name": "tiny-blr", "hidden_size": cols,
13280                "intermediate_size": rows, "num_layers": 1, "num_attention_heads": 2,
13281                "num_kv_heads": 1, "head_dim": 64, "vocab_size": rows,
13282                "layer_types": ["FullAttention"], "rms_norm_eps": 1e-6,
13283                "max_position_embeddings": 8},
13284            "quant_type": "Q8_2F"
13285        }))
13286        .unwrap();
13287        let specs = [
13288            TensorSpec {
13289                name: "g".into(),
13290                dtype: TensorDtype::BinaryLowRank,
13291                shape: vec![rows, cols],
13292                data: gate.clone(),
13293            },
13294            TensorSpec {
13295                name: "u".into(),
13296                dtype: TensorDtype::BinaryLowRank,
13297                shape: vec![rows, cols],
13298                data: up.clone(),
13299            },
13300            TensorSpec {
13301                name: "q".into(),
13302                dtype: TensorDtype::Q8Row,
13303                shape: vec![rows, cols],
13304                data: q8,
13305            },
13306        ];
13307        let dir = std::env::temp_dir().join(format!("cmf-blr-mapped-{}", std::process::id()));
13308        std::fs::create_dir_all(&dir).unwrap();
13309        let path = dir.join("m.cmf");
13310        CmfModel::write(&path, &header, &specs, None, None).unwrap();
13311        let model = Arc::new(CmfModel::open(&path).unwrap());
13312        assert!(model.verify().is_empty(), "{:?}", model.verify());
13313        let t = |n: &str| QTensor::from_model(&model, n).unwrap();
13314        let (g, u, q) = (t("g"), t("u"), t("q"));
13315        assert!(g.is_blr() && u.is_blr() && !q.is_blr());
13316        assert!(g.graph_weight().is_none() && g.graph_weight_descriptor().is_none());
13317        assert!(g.as_f32().is_none() && g.model_idx().is_none());
13318        assert!(g.mapped_device_gemm().is_none() && g.q1_parts().is_none());
13319        assert!(!g.sparse_col_ok());
13320        assert_eq!((g.rows(), g.cols()), (rows, cols));
13321        let deq = |bytes: &[u8]| {
13322            let mut w = vec![0f32; rows * cols];
13323            cortiq_core::blr::dequant_blr(bytes, rows, cols, &mut w).unwrap();
13324            w
13325        };
13326        let (wg, wu) = (deq(&gate), deq(&up));
13327        let mv = |w: &[f32], x: &[f32]| -> Vec<f32> {
13328            (0..rows)
13329                .map(|i| (0..cols).map(|j| w[i * cols + j] * x[j]).sum())
13330                .collect()
13331        };
13332        let close = |a: &[f32], b: &[f32], what: &str| {
13333            let s = b.iter().map(|v| v * v).sum::<f32>().sqrt().max(1e-30);
13334            let e = a.iter().zip(b).map(|(x, y)| (x - y) * (x - y)).sum::<f32>().sqrt();
13335            assert!(e / s < 1e-5, "{what}: rel {}", e / s);
13336        };
13337        let pool = Pool::new(3);
13338        let x: Vec<f32> = (0..cols).map(|_| rng.f()).collect();
13339        let x2: Vec<f32> = (0..cols).map(|_| rng.f()).collect();
13340        let mut o = vec![0f32; rows];
13341        g.matvec(&x, &mut o, Some(&pool));
13342        close(&o, &mv(&wg, &x), "matvec");
13343        let (mut o1, mut o2) = (vec![0f32; rows], vec![0f32; rows]);
13344        g.matvec2(&x, &x2, &mut o1, &mut o2, Some(&pool));
13345        close(&o1, &mv(&wg, &x), "matvec2 a");
13346        close(&o2, &mv(&wg, &x2), "matvec2 b");
13347        for b in [2usize, 5] {
13348            let xs: Vec<f32> = (0..b * cols).map(|_| rng.f()).collect();
13349            let mut om = vec![0f32; b * rows];
13350            g.matmat(&xs, b, &mut om, Some(&pool));
13351            for t in 0..b {
13352                close(
13353                    &om[t * rows..(t + 1) * rows],
13354                    &mv(&wg, &xs[t * cols..(t + 1) * cols]),
13355                    "matmat",
13356                );
13357            }
13358        }
13359        // Multi-tensor job mixed with a q8_row sibling: per-tensor fallback,
13360        // the q8_row output identical to its own matvec.
13361        let (mut a, mut bq, mut c) = (vec![0f32; rows], vec![0f32; rows], vec![0f32; rows]);
13362        QTensor::matvec_many([&g, &q, &u], &x, [&mut a, &mut bq, &mut c], Some(&pool));
13363        close(&a, &mv(&wg, &x), "matvec_many blr gate");
13364        close(&c, &mv(&wu, &x), "matvec_many blr up");
13365        let mut qo = vec![0f32; rows];
13366        q.matvec(&x, &mut qo, Some(&pool));
13367        assert_eq!(bq, qo, "matvec_many q8_row sibling");
13368        let mut f = vec![0f32; rows];
13369        assert!(!QTensor::matvec_silu_mul(&g, &u, &x, &mut f, Some(&pool)));
13370        let mut row = vec![0f32; cols];
13371        g.row_f32(rows / 3, &mut row);
13372        close(&row, &wg[(rows / 3) * cols..(rows / 3 + 1) * cols], "row_f32");
13373        let mut scratch = vec![0f32; cols];
13374        let d = g.row_dot(7, &x, &mut scratch);
13375        let want: f32 = (0..cols).map(|j| wg[7 * cols + j] * x[j]).sum();
13376        assert!((d - want).abs() < 1e-4 * (1.0 + want.abs()), "row_dot");
13377        let mut col = vec![0f32; rows];
13378        g.add_col_scaled(cols / 2, 1.5, &mut col);
13379        let want: Vec<f32> = (0..rows).map(|i| 1.5 * wg[i * cols + cols / 2]).collect();
13380        close(&col, &want, "add_col_scaled");
13381        drop((g, u, q, model));
13382        let _ = std::fs::remove_dir_all(&dir);
13383    }
13384
13385    #[test]
13386    fn repack_is_bit_identical() {
13387        // The interleaved-repack kernel must produce EXACTLY the same
13388        // bits as the mmap-layout kernel: integer accumulation is order-
13389        // exact, the f32 epilogue is identical. Odd rows exercise the
13390        // tail; direct range calls exercise unaligned pool splits.
13391        let (rows, cols) = (267, 96); // 66 groups + 3 tail rows, cols % 16 == 0
13392        let w: Vec<u8> = (0..rows * cols)
13393            .map(|i| (((i * 89) % 253) as i32 - 126) as i8 as u8)
13394            .collect();
13395        let scales: Vec<f32> = (0..rows).map(|o| 0.003 + o as f32 * 0.0007).collect();
13396        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.37).sin() * 2.0).collect();
13397        let rep = q8_repack_layout(&w, rows, cols);
13398        // Group interleave round-trips.
13399        for g in 0..rows / 4 {
13400            for c in 0..cols / 16 {
13401                for lane in 0..4 {
13402                    assert_eq!(
13403                        &rep[g * 4 * cols + c * 64 + lane * 16
13404                            ..g * 4 * cols + c * 64 + lane * 16 + 16],
13405                        &w[(g * 4 + lane) * cols + c * 16..(g * 4 + lane) * cols + c * 16 + 16],
13406                    );
13407                }
13408            }
13409        }
13410        let mut a = vec![0.0f32; rows];
13411        qmatvec(
13412            &w,
13413            &[],
13414            &scales,
13415            &x,
13416            &[],
13417            TensorDtype::Q8Row,
13418            rows,
13419            cols,
13420            &mut a,
13421            None,
13422        );
13423        let mut b = vec![0.0f32; rows];
13424        qmatvec(
13425            &w,
13426            &rep,
13427            &scales,
13428            &x,
13429            &[],
13430            TensorDtype::Q8Row,
13431            rows,
13432            cols,
13433            &mut b,
13434            None,
13435        );
13436        assert_eq!(a, b, "full-range repack output diverged");
13437
13438        #[cfg(target_arch = "aarch64")]
13439        if sdot_enabled() {
13440            // Unaligned range split (pool workers get arbitrary bounds).
13441            let act = split_act(&x);
13442            let mut c1 = vec![0.0f32; rows];
13443            let mut c2 = vec![0.0f32; rows];
13444            q8_range_sdot(
13445                &w,
13446                &[],
13447                &scales,
13448                &act,
13449                cols,
13450                SendMut(c1.as_mut_ptr()),
13451                3,
13452                rows - 2,
13453            );
13454            q8_range_sdot(
13455                &w,
13456                &rep,
13457                &scales,
13458                &act,
13459                cols,
13460                SendMut(c2.as_mut_ptr()),
13461                3,
13462                rows - 2,
13463            );
13464            assert_eq!(c1, c2, "unaligned-range repack output diverged");
13465        }
13466    }
13467
13468    #[test]
13469    fn sdot_a8w8_noise_is_bounded() {
13470        // Off-grid activations: A8 quantization noise must stay small in
13471        // relative L2 over the whole output (realistic accuracy contract;
13472        // vmfcore measured argmax-identical decode on real models).
13473        let (rows, cols) = (16, 512);
13474        let w: Vec<u8> = (0..rows * cols)
13475            .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
13476            .collect();
13477        let scales = vec![0.01f32; rows];
13478        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.21).sin()).collect();
13479        let mut a = vec![0.0f32; rows];
13480        qmatvec(
13481            &w,
13482            &[],
13483            &scales,
13484            &x,
13485            &[],
13486            TensorDtype::Q8Row,
13487            rows,
13488            cols,
13489            &mut a,
13490            None,
13491        );
13492        let (mut num, mut den) = (0f64, 0f64);
13493        for o in 0..rows {
13494            let mut acc = 0.0f32;
13495            for j in 0..cols {
13496                acc += (w[o * cols + j] as i8) as f32 * x[j];
13497            }
13498            let expect = acc * scales[o];
13499            num += ((a[o] - expect) as f64).powi(2);
13500            den += (expect as f64).powi(2);
13501        }
13502        let rel = (num / den.max(1e-12)).sqrt();
13503        assert!(rel < 0.05, "A8W8 relative L2 error too high: {rel}");
13504    }
13505
13506    #[test]
13507    fn i8_dot_neon_matches_scalar() {
13508        let n = 100;
13509        let w: Vec<u8> = (0..n).map(|i| ((i * 37 + 11) % 251) as u8).collect();
13510        let x: Vec<f32> = (0..n).map(|i| (i as f32 * 0.13).sin()).collect();
13511        let mut scalar = 0.0f32;
13512        for j in 0..n {
13513            scalar += (w[j] as i8) as f32 * x[j];
13514        }
13515        let fast = dot_i8_f32(&w, &x);
13516        assert!((scalar - fast).abs() < 1e-3 * scalar.abs().max(1.0));
13517    }
13518
13519    /// Fused vbit matvec must match full dequant_vbit + dense matvec.
13520    #[test]
13521    fn vbitmatvec_matches_full_dequant() {
13522        let (rows, cols) = (6, 64);
13523        let ng = cols / GROUP_SIZE;
13524        // Hand-craft: bits per row, f16 scales, packed rows.
13525        let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
13526        let mut bytes = bits.clone();
13527        for g in 0..rows * ng {
13528            let s = 0.02 + 0.001 * g as f32;
13529            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
13530        }
13531        for r in 0..rows {
13532            let b = bits[r] as usize;
13533            let (mut acc, mut nb) = (0u64, 0usize);
13534            let mut rowbytes = Vec::new();
13535            for i in 0..cols {
13536                let v = ((i * 7 + r * 13) % (1 << b)) as u64;
13537                acc = (acc << b) | v;
13538                nb += b;
13539                while nb >= 8 {
13540                    nb -= 8;
13541                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
13542                }
13543            }
13544            if nb > 0 {
13545                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
13546            }
13547            bytes.extend_from_slice(&rowbytes);
13548        }
13549        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
13550
13551        let mut reference = vec![0f32; rows * cols];
13552        cortiq_core::quant::dequant_vbit(&bytes, rows, cols, &mut reference).unwrap();
13553        let mut expect = vec![0f32; rows];
13554        for r in 0..rows {
13555            expect[r] = reference[r * cols..(r + 1) * cols]
13556                .iter()
13557                .zip(&x)
13558                .map(|(w, xv)| w * xv)
13559                .sum();
13560        }
13561        let mut got = vec![0f32; rows];
13562        let offsets = vbit_row_offsets(&bytes, rows, cols);
13563        vbitmatvec(&bytes, &offsets, &x, rows, cols, &mut got, None);
13564        // SDOT path quantizes activations to i8 (A8W8): bounded noise,
13565        // same contract as q8 (exact path is pinned by CMF_SDOT=0 in
13566        // the golden-parity gate).
13567        let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
13568        let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
13569        for r in 0..rows {
13570            assert!(
13571                (got[r] - expect[r]).abs() < tol * scale,
13572                "row {r}: {} vs {}",
13573                got[r],
13574                expect[r]
13575            );
13576        }
13577    }
13578
13579    /// Fused q4 matvec must match the reference full-dequant + dense
13580    /// matvec bit-for-bit in structure (same f32 math, group order).
13581    /// vbit matmat: the blocked 1×4 leg must match the per-row path
13582    /// (paired env toggle; larger shape so both code paths engage).
13583    #[test]
13584    #[cfg(target_arch = "x86_64")]
13585    fn vbit_matmat_blocked_matches_per_row() {
13586        let (rows, cols, b) = (64usize, 128usize, 9usize);
13587        let ng = cols / GROUP_SIZE;
13588        let bits: Vec<u8> = (0..rows).map(|r| [3u8, 4, 5, 6][r % 4]).collect();
13589        let mut bytes = bits.clone();
13590        for g in 0..rows * ng {
13591            let sc = 0.02 + 0.0005 * g as f32;
13592            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
13593        }
13594        for r in 0..rows {
13595            let bw = bits[r] as usize;
13596            let (mut acc, mut nb) = (0u64, 0usize);
13597            let mut rowbytes = Vec::new();
13598            for i in 0..cols {
13599                let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
13600                acc = (acc << bw) | v;
13601                nb += bw;
13602                while nb >= 8 {
13603                    nb -= 8;
13604                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
13605                }
13606            }
13607            if nb > 0 {
13608                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
13609            }
13610            bytes.extend_from_slice(&rowbytes);
13611        }
13612        let x: Vec<f32> = (0..b * cols)
13613            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
13614            .collect();
13615        let offsets = vbit_row_offsets(&bytes, rows, cols);
13616        let mut y_a = vec![0f32; b * rows];
13617        let mut y_b = vec![0f32; b * rows];
13618        unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
13619        vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_a, None);
13620        unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
13621        vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_b, None);
13622        unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
13623        let max_d = y_a
13624            .iter()
13625            .zip(&y_b)
13626            .map(|(p, q)| (p - q).abs())
13627            .fold(0.0f32, f32::max);
13628        assert!(max_d < 1e-4, "vbit blocked ≠ per-row: max|Δ| = {max_d}");
13629    }
13630
13631    /// q4t blocked 1×4 (SDOT on ARM, AVX2 on x86) must equal the
13632    /// per-row path exactly: same nibble unpack, same group order,
13633    /// same f32 accumulation — batch == matvec bit-for-bit. b=9 covers
13634    /// two full 1×4 blocks plus a remainder through the single-row
13635    /// kernel. (Both paths produce identical output, so the shared
13636    /// CMF_X86_BLOCKED env var racing with other tests cannot flip
13637    /// the verdict — worst case both sides take the same path.)
13638    #[test]
13639    fn q4t_matmat_blocked_matches_per_row() {
13640        let (rows, cols, b) = (16usize, 64usize, 9usize);
13641        let gpr = cols / GROUP_SIZE;
13642        let mut bytes = vec![0u8; rows * gpr * Q4_TILE];
13643        for r in 0..rows {
13644            for g in 0..gpr {
13645                let t = (r * gpr + g) * Q4_TILE;
13646                let sc = 0.02 + 0.001 * (r * gpr + g) as f32;
13647                bytes[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
13648                for k in 0..16 {
13649                    bytes[t + 2 + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
13650                }
13651            }
13652        }
13653        let x: Vec<f32> = (0..b * cols)
13654            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
13655            .collect();
13656        let mut y_blk = vec![0f32; b * rows];
13657        let mut y_row = vec![0f32; b * rows];
13658        unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
13659        q4t_matmat(&bytes, &x, b, rows, cols, &mut y_blk, None);
13660        unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
13661        q4t_matmat(&bytes, &x, b, rows, cols, &mut y_row, None);
13662        unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
13663        assert_eq!(y_blk, y_row, "q4t blocked 1x4 ≠ per-row");
13664    }
13665
13666    /// The wide-batch Accelerate arm of q4t_matmat vs a brute-force
13667    /// f32 dequant matmul: both are f32 GEMMs, so only reduction
13668    /// order differs — tight tolerance.
13669    /// A synthetic q4tp payload: random nibbles plus a per-row ladder whose
13670    /// span varies row to row, so the codes actually exercise the full 0..31
13671    /// range rather than clustering on one rung.
13672    fn synth_q4tp(rows: usize, cols: usize) -> Vec<u8> {
13673        use cortiq_core::quant::{f32_to_f16, q4tp_code_stride, q4tp_put_code};
13674        let gpr = cols / GROUP_SIZE;
13675        let stride = q4tp_code_stride(gpr);
13676        let (params_off, codes_off, _) = q4tp_sections(rows, cols);
13677        let mut b = vec![0u8; codes_off + rows * stride];
13678        for r in 0..rows {
13679            for g in 0..gpr {
13680                let t = (r * gpr + g) * Q4TP_NIB;
13681                for k in 0..16 {
13682                    b[t + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
13683                }
13684            }
13685            let lo = -6.0 - 0.03 * (r % 17) as f32;
13686            let step = 0.01 + 0.004 * (r % 11) as f32;
13687            let p = params_off + r * 4;
13688            b[p..p + 2].copy_from_slice(&f32_to_f16(lo).to_le_bytes());
13689            b[p + 2..p + 4].copy_from_slice(&f32_to_f16(step).to_le_bytes());
13690            let crow = &mut b[codes_off + r * stride..codes_off + (r + 1) * stride];
13691            for g in 0..gpr {
13692                q4tp_put_code(crow, g, (r * 5 + g * 3) % 32);
13693            }
13694        }
13695        b
13696    }
13697
13698    /// The same weights re-expressed as q4_tiled, so the proven kernel can
13699    /// be the reference: each tile stores the ladder scale its code selects.
13700    /// Only the f16 rounding of that scale separates the two payloads.
13701    fn q4tp_as_q4t(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
13702        let gpr = cols / GROUP_SIZE;
13703        let v = Q4tpView::new(bytes, rows, cols);
13704        let mut out = vec![0u8; rows * gpr * Q4_TILE];
13705        let mut sc = vec![0f32; gpr];
13706        for r in 0..rows {
13707            v.scales_into(r, gpr, &mut sc);
13708            for g in 0..gpr {
13709                let t = (r * gpr + g) * Q4_TILE;
13710                let s = sc[g];
13711                out[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
13712                let src = (r * gpr + g) * Q4TP_NIB;
13713                out[t + 2..t + Q4_TILE].copy_from_slice(&v.nib[src..src + Q4TP_NIB]);
13714            }
13715        }
13716        out
13717    }
13718
13719    /// The exact (`CMF_SDOT=0`) path must reproduce `dequant_q4tp` to f32
13720    /// rounding — that scalar routine is the format's definition, and the
13721    /// kernels re-derive the scale from the ladder independently. Call the
13722    /// row kernel directly: `matmat` picks the int8 arm when a8w8 is on,
13723    /// so routing through it would test the other path by accident.
13724    #[test]
13725    fn q4tp_exact_path_matches_dequant_reference() {
13726        let (rows, cols) = (256usize, 512usize);
13727        let gpr = cols / GROUP_SIZE;
13728        let bytes = synth_q4tp(rows, cols);
13729        let mut w = vec![0f32; rows * cols];
13730        cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
13731
13732        let x: Vec<f32> = (0..cols)
13733            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
13734            .collect();
13735        let v = Q4tpView::new(&bytes, rows, cols);
13736        let mut sc = vec![0f32; gpr];
13737        for r in 0..rows {
13738            v.scales_into(r, gpr, &mut sc);
13739            let got = q4tp_row_exact(v.nib, r, gpr, &x, &sc);
13740            let want: f32 = (0..cols).map(|c| w[r * cols + c] * x[c]).sum();
13741            // These dot products cancel down to ~1e-3 from terms of ~5e-2, so
13742            // the meaningful yardstick is the summed magnitude, not the result:
13743            // against the result any reordering of a 512-term f32 sum "fails".
13744            let mag: f32 = (0..cols).map(|c| (w[r * cols + c] * x[c]).abs()).sum();
13745            assert!(
13746                (got - want).abs() <= 1e-5 * mag,
13747                "row {r}: kernel {got} vs dequant {want}"
13748            );
13749        }
13750    }
13751
13752    /// The int8 (a8w8) path can't be checked against an f32 reference — the
13753    /// activation quantization dominates. Check it against the q4t kernel it
13754    /// was ported from instead, on payloads holding the same weights: that
13755    /// isolates exactly what the port could break (16 B stride, ladder
13756    /// lookup, nibble unpack) from what it deliberately shares.
13757    #[test]
13758    fn q4tp_matvec_matches_the_q4t_kernel_it_was_ported_from() {
13759        let (rows, cols) = (256usize, 512usize);
13760        let bytes = synth_q4tp(rows, cols);
13761        let twin = q4tp_as_q4t(&bytes, rows, cols);
13762        let x: Vec<f32> = (0..cols)
13763            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
13764            .collect();
13765
13766        let mut got = vec![0f32; rows];
13767        q4tp_matvec(&bytes, &x, rows, cols, &mut got, None);
13768        let mut want = vec![0f32; rows];
13769        q4t_matvec(&twin, &x, rows, cols, &mut want, None);
13770
13771        // Scale is f16 in the twin and f32 here, so allow that rounding on
13772        // top of the summed magnitude (same cancellation argument as above).
13773        let mut w = vec![0f32; rows * cols];
13774        cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
13775        for r in 0..rows {
13776            let mag: f32 = (0..cols).map(|c| (w[r * cols + c] * x[c]).abs()).sum();
13777            assert!(
13778                (got[r] - want[r]).abs() <= 1e-3 * mag,
13779                "row {r}: q4tp {} vs q4t {}",
13780                got[r],
13781                want[r]
13782            );
13783        }
13784    }
13785
13786    /// `matmat` carries three arms (Accelerate, blocked int8 1x4, scalar).
13787    /// Batch 5 crosses the blocked kernel's stride, so this exercises the
13788    /// 1x4 path AND its scalar tail in one run — the blocked kernel is new
13789    /// code and its four accumulators are exactly what tends to go wrong.
13790    #[test]
13791    fn q4tp_matmat_matches_the_q4t_kernel_it_was_ported_from() {
13792        let (rows, cols, b) = (256usize, 512usize, 5usize);
13793        let bytes = synth_q4tp(rows, cols);
13794        let twin = q4tp_as_q4t(&bytes, rows, cols);
13795        let xs: Vec<f32> = (0..b * cols)
13796            .map(|i| ((i * 29 + 11) % 89) as f32 / 89.0 - 0.5)
13797            .collect();
13798
13799        let mut got = vec![0f32; b * rows];
13800        q4tp_matmat(&bytes, &xs, b, rows, cols, &mut got, None);
13801        let mut want = vec![0f32; b * rows];
13802        q4t_matmat(&twin, &xs, b, rows, cols, &mut want, None);
13803
13804        let mut w = vec![0f32; rows * cols];
13805        cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
13806        for t in 0..b {
13807            for r in 0..rows {
13808                let mag: f32 = (0..cols)
13809                    .map(|c| (w[r * cols + c] * xs[t * cols + c]).abs())
13810                    .sum();
13811                let (g, wa) = (got[t * rows + r], want[t * rows + r]);
13812                assert!(
13813                    (g - wa).abs() <= 1e-3 * mag,
13814                    "batch {t} row {r}: q4tp {g} vs q4t {wa}"
13815                );
13816            }
13817        }
13818    }
13819
13820    #[test]
13821    fn q4tp_matvec2_matches_the_single_stream_kernel() {
13822        let (rows, cols) = (128usize, 256usize);
13823        let gpr = cols / GROUP_SIZE;
13824        let bytes = synth_q4tp(rows, cols);
13825        let xs: Vec<f32> = (0..2 * cols)
13826            .map(|i| ((i * 29 + 11) % 89) as f32 / 89.0 - 0.5)
13827            .collect();
13828
13829        let (mut o1, mut o2) = (vec![0f32; rows], vec![0f32; rows]);
13830        q4tp_matvec2(
13831            &bytes,
13832            &xs[..cols],
13833            &xs[cols..],
13834            rows,
13835            cols,
13836            &mut o1,
13837            &mut o2,
13838            None,
13839        );
13840
13841        // matvec2 takes the exact path for both streams, so the single-row
13842        // kernel is an exact reference — no tolerance for path differences.
13843        let v = Q4tpView::new(&bytes, rows, cols);
13844        let mut sc = vec![0f32; gpr];
13845        for r in 0..rows {
13846            v.scales_into(r, gpr, &mut sc);
13847            assert_eq!(o1[r], q4tp_row_exact(v.nib, r, gpr, &xs[..cols], &sc));
13848            assert_eq!(o2[r], q4tp_row_exact(v.nib, r, gpr, &xs[cols..], &sc));
13849        }
13850    }
13851
13852    /// q4tp must not COST speed — it exists to save bytes, and a format that
13853    /// trades 7% of a file for a slower model is a bad trade. This guard is
13854    /// here because correctness tests happily passed while `q4tp_matmat` was
13855    /// missing its int8 and Accelerate arms and the model ran 5x slower.
13856    /// Measured on M-series: 0.97-1.04x, i.e. parity (16 B tiles are better
13857    /// aligned than q4t's 18 B, which pays for the scale indirection).
13858    #[test]
13859    fn q4tp_matvec_keeps_pace_with_q4t() {
13860        let (rows, cols) = (4096usize, 3072usize);
13861        let bytes = synth_q4tp(rows, cols);
13862        let twin = q4tp_as_q4t(&bytes, rows, cols);
13863        let x: Vec<f32> = (0..cols).map(|i| (i % 97) as f32 / 97.0 - 0.5).collect();
13864        let mut o = vec![0f32; rows];
13865        let n = 12;
13866        let mut best = (f64::MAX, f64::MAX);
13867        // Interleaved A/B, minimum statistic: this machine throttles, and a
13868        // mean over a thermal ramp reliably indicts whichever ran second.
13869        for _ in 0..3 {
13870            let t0 = std::time::Instant::now();
13871            for _ in 0..n {
13872                q4t_matvec(&twin, &x, rows, cols, &mut o, None);
13873            }
13874            best.0 = best.0.min(t0.elapsed().as_secs_f64());
13875            let t0 = std::time::Instant::now();
13876            for _ in 0..n {
13877                q4tp_matvec(&bytes, &x, rows, cols, &mut o, None);
13878            }
13879            best.1 = best.1.min(t0.elapsed().as_secs_f64());
13880        }
13881        let ratio = best.1 / best.0;
13882        println!(
13883            "q4t {:.3} ms | q4tp {:.3} ms | {ratio:.2}x",
13884            best.0 * 1e3 / n as f64,
13885            best.1 * 1e3 / n as f64
13886        );
13887        assert!(ratio < 2.0, "q4tp matvec {ratio:.2}x slower than q4t");
13888    }
13889
13890    #[cfg(target_os = "macos")]
13891    #[test]
13892    fn q4t_matmat_accel_matches_dequant_reference() {
13893        if !accel_gemm_enabled() {
13894            return; // CMF_ACCEL=0
13895        }
13896        let (rows, cols, b) = (512usize, 1024usize, 8usize); // ≥500K → accel arm
13897        let gpr = cols / GROUP_SIZE;
13898        let mut bytes = vec![0u8; rows * gpr * Q4_TILE];
13899        for r in 0..rows {
13900            for g in 0..gpr {
13901                let t = (r * gpr + g) * Q4_TILE;
13902                let sc = 0.02 + 0.0005 * ((r * gpr + g) % 64) as f32;
13903                bytes[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
13904                for k in 0..16 {
13905                    bytes[t + 2 + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
13906                }
13907            }
13908        }
13909        let x: Vec<f32> = (0..b * cols)
13910            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
13911            .collect();
13912        let mut got = vec![0f32; b * rows];
13913        q4t_matmat(&bytes, &x, b, rows, cols, &mut got, None);
13914        // Brute-force reference off the same tiles.
13915        let mut w = vec![0f32; rows * cols];
13916        for r in 0..rows {
13917            for g in 0..gpr {
13918                let t = (r * gpr + g) * Q4_TILE;
13919                let s = f16_to_f32(u16::from_le_bytes([bytes[t], bytes[t + 1]]));
13920                for (k, &bb) in bytes[t + 2..t + Q4_TILE].iter().enumerate() {
13921                    w[r * cols + g * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s;
13922                    w[r * cols + g * GROUP_SIZE + k * 2 + 1] =
13923                        (((bb >> 4) & 0x0F) as f32 - 8.0) * s;
13924                }
13925            }
13926        }
13927        for bi in 0..b {
13928            for r in 0..rows {
13929                let want: f32 = (0..cols).map(|j| x[bi * cols + j] * w[r * cols + j]).sum();
13930                let d = (got[bi * rows + r] - want).abs();
13931                assert!(
13932                    d <= want.abs().max(1.0) * 1e-4,
13933                    "accel q4t GEMM diverged at ({bi},{r}): {} vs {want}",
13934                    got[bi * rows + r]
13935                );
13936            }
13937        }
13938    }
13939
13940    #[test]
13941    fn q4matvec_matches_full_dequant() {
13942        let (rows, cols) = (8, 64);
13943        let groups = rows * cols / GROUP_SIZE;
13944        // Hand-craft a q4_block blob: nibbles then f16 scales.
13945        let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
13946        for i in 0..groups * 16 {
13947            bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
13948        }
13949        for g in 0..groups {
13950            let s = 0.01 + 0.003 * g as f32;
13951            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
13952        }
13953        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
13954
13955        let mut reference = vec![0.0f32; rows * cols];
13956        cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
13957        let mut expect = vec![0.0f32; rows];
13958        for r in 0..rows {
13959            expect[r] = reference[r * cols..(r + 1) * cols]
13960                .iter()
13961                .zip(&x)
13962                .map(|(w, xv)| w * xv)
13963                .sum();
13964        }
13965
13966        let mut got = vec![0.0f32; rows];
13967        q4matvec(&bytes, &x, rows, cols, &mut got, None);
13968        // SDOT path quantizes activations to i8 (A8W8): bounded noise,
13969        // same contract as q8/vbit (exact path is pinned by CMF_SDOT=0
13970        // in the golden-parity gate).
13971        let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
13972        let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
13973        for r in 0..rows {
13974            assert!(
13975                (got[r] - expect[r]).abs() < tol * scale,
13976                "row {r}: {} vs {}",
13977                got[r],
13978                expect[r]
13979            );
13980        }
13981    }
13982
13983    /// Fused two-input vbit matvec must equal two single matvecs exactly
13984    /// (same per-lane accumulation order on both scalar and SDOT paths).
13985    #[test]
13986    fn vbitmatvec2_equals_two_singles() {
13987        let (rows, cols) = (6, 64);
13988        let ng = cols / GROUP_SIZE;
13989        let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
13990        let mut bytes = bits.clone();
13991        for g in 0..rows * ng {
13992            let s = 0.02 + 0.001 * g as f32;
13993            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
13994        }
13995        for r in 0..rows {
13996            let b = bits[r] as usize;
13997            let (mut acc, mut nb) = (0u64, 0usize);
13998            let mut rowbytes = Vec::new();
13999            for i in 0..cols {
14000                let v = ((i * 7 + r * 13) % (1 << b)) as u64;
14001                acc = (acc << b) | v;
14002                nb += b;
14003                while nb >= 8 {
14004                    nb -= 8;
14005                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
14006                }
14007            }
14008            if nb > 0 {
14009                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
14010            }
14011            bytes.extend_from_slice(&rowbytes);
14012        }
14013        let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
14014        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).cos()).collect();
14015        let offsets = vbit_row_offsets(&bytes, rows, cols);
14016
14017        let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
14018        vbitmatvec(&bytes, &offsets, &x1, rows, cols, &mut a1, None);
14019        vbitmatvec(&bytes, &offsets, &x2, rows, cols, &mut a2, None);
14020        let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
14021        vbitmatvec2(
14022            &bytes, &offsets, &x1, &x2, rows, cols, &mut b1, &mut b2, None,
14023        );
14024        assert_eq!(a1, b1, "fused vbit lane 1 must be bit-identical");
14025        assert_eq!(a2, b2, "fused vbit lane 2 must be bit-identical");
14026    }
14027
14028    /// Fused two-input q4 matvec must equal two single matvecs exactly.
14029    #[test]
14030    fn q4matvec2_equals_two_singles() {
14031        let (rows, cols) = (8, 128);
14032        let groups = rows * cols / GROUP_SIZE;
14033        let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
14034        for i in 0..groups * 16 {
14035            bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
14036        }
14037        for g in 0..groups {
14038            let s = 0.01 + 0.003 * g as f32;
14039            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
14040        }
14041        // Include an outlier channel so the SDOT correction path is
14042        // exercised in the pair kernel too.
14043        let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
14044        x1[9] = 250.0;
14045        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
14046
14047        let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
14048        q4matvec(&bytes, &x1, rows, cols, &mut a1, None);
14049        q4matvec(&bytes, &x2, rows, cols, &mut a2, None);
14050        let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
14051        q4matvec2(&bytes, &x1, &x2, rows, cols, &mut b1, &mut b2, None);
14052        assert_eq!(a1, b1, "fused q4 lane 1 must be bit-identical");
14053        assert_eq!(a2, b2, "fused q4 lane 2 must be bit-identical");
14054    }
14055
14056    /// Multi-matrix job must equal separate matvecs exactly — same
14057    /// kernels, only the dispatch is fused.
14058    #[test]
14059    fn matvec_many_equals_separate_matvecs() {
14060        use crate::pool::Pool;
14061        let (r1, r2, cols) = (300, 200, 64);
14062        let mk = |salt: usize, rows: usize| {
14063            QTensor::from_f32(
14064                (0..rows * cols)
14065                    .map(|i| ((i * 7 + salt) % 97) as f32 / 97.0 - 0.5)
14066                    .collect(),
14067                rows,
14068                cols,
14069            )
14070        };
14071        let (a, b) = (mk(1, r1), mk(5, r2));
14072        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).sin()).collect();
14073        let pool = Pool::new(3);
14074
14075        let (mut ea, mut eb) = (vec![0f32; r1], vec![0f32; r2]);
14076        a.matvec(&x, &mut ea, Some(&pool));
14077        b.matvec(&x, &mut eb, Some(&pool));
14078        let (mut ga, mut gb) = (vec![0f32; r1], vec![0f32; r2]);
14079        QTensor::matvec_many([&a, &b], &x, [&mut ga, &mut gb], Some(&pool));
14080        assert_eq!(ea, ga, "fused multi-matrix lane 1 must be bit-identical");
14081        assert_eq!(eb, gb, "fused multi-matrix lane 2 must be bit-identical");
14082    }
14083
14084    /// The public Q4TP operator must take the real mapped matvec_many arm,
14085    /// rather than the F32 fallback above.  Build a tiny valid CMF so both
14086    /// handles retain their mmap payloads, then compare the fused dispatch
14087    /// with two ordinary mapped matvec calls bit-for-bit.
14088    #[test]
14089    fn q4tp_matvec_many_equals_separate_matvecs() {
14090        use crate::pool::Pool;
14091        use cortiq_core::{CMF_VERSION, CmfHeader, CmfModel, QuantType, TensorSpec};
14092
14093        let (r1, r2, cols) = (300usize, 200usize, 64usize);
14094        let arch: cortiq_core::ModelArch = serde_json::from_value(serde_json::json!({
14095            "arch_name": "tiny-q4tp",
14096            "hidden_size": cols,
14097            "intermediate_size": cols * 2,
14098            "num_layers": 1,
14099            "num_attention_heads": 2,
14100            "num_kv_heads": 1,
14101            "head_dim": 32,
14102            "vocab_size": r1,
14103            "layer_types": ["FullAttention"],
14104            "rms_norm_eps": 1e-6,
14105            "max_position_embeddings": 8,
14106            "linear_conv_kernel_dim": 0,
14107            "linear_num_key_heads": 0,
14108            "linear_num_value_heads": 0
14109        }))
14110        .unwrap();
14111        let header = CmfHeader {
14112            format: "cmf".into(),
14113            version: CMF_VERSION,
14114            arch,
14115            quant_type: QuantType::Q4Block,
14116            provenance: None,
14117            tokenizer_config: None,
14118            section_hashes: None,
14119            skills: Vec::new(),
14120            shard: None,
14121            calibration: None,
14122            routing: None,
14123            genome: None,
14124            lineage: Vec::new(),
14125            router: None,
14126            segments: Vec::new(),
14127        };
14128        let specs = [
14129            TensorSpec {
14130                name: "q".into(),
14131                dtype: TensorDtype::Q4TiledP,
14132                shape: vec![r1, cols],
14133                data: synth_q4tp(r1, cols),
14134            },
14135            TensorSpec {
14136                name: "kv".into(),
14137                dtype: TensorDtype::Q4TiledP,
14138                shape: vec![r2, cols],
14139                data: synth_q4tp(r2, cols),
14140            },
14141        ];
14142        let dir = std::env::temp_dir().join(format!("cmf-q4tp-many-{}", std::process::id()));
14143        std::fs::create_dir_all(&dir).unwrap();
14144        let path = dir.join("m.cmf");
14145        CmfModel::write(&path, &header, &specs, None, None).unwrap();
14146        let model = Arc::new(CmfModel::open(&path).unwrap());
14147        let (a, b) = (
14148            QTensor::from_model(&model, "q").unwrap(),
14149            QTensor::from_model(&model, "kv").unwrap(),
14150        );
14151        assert_eq!(a.model_dtype(), Some(TensorDtype::Q4TiledP));
14152        assert_eq!(b.model_dtype(), Some(TensorDtype::Q4TiledP));
14153        let x: Vec<f32> = (0..cols)
14154            .map(|i| ((i * 17 + 3) % 97) as f32 / 97.0 - 0.5)
14155            .collect();
14156        let pool = Pool::new(3);
14157        let (mut ea, mut eb) = (vec![0.0f32; r1], vec![0.0f32; r2]);
14158        a.matvec(&x, &mut ea, Some(&pool));
14159        b.matvec(&x, &mut eb, Some(&pool));
14160        let (mut ga, mut gb) = (vec![0.0f32; r1], vec![0.0f32; r2]);
14161        QTensor::matvec_many([&a, &b], &x, [&mut ga, &mut gb], Some(&pool));
14162        assert_eq!(ea, ga, "Q4TP fused lane 1 must be bit-identical");
14163        assert_eq!(eb, gb, "Q4TP fused lane 2 must be bit-identical");
14164        let _ = std::fs::remove_dir_all(&dir);
14165    }
14166
14167    /// The MiMo speculative verify's kernels: several tokens' MoE through
14168    /// `moe_gate_up_rows` / `moe_down_rows` (+ the caller's route-order sum)
14169    /// is bit-identical to each token alone through `moe_gate_up_many` /
14170    /// `moe_down_many` (decode), and a row-exact `q4tp_matmat` of five
14171    /// tokens (wide enough for the blocked tiles) equals five matvecs.
14172    #[test]
14173    fn multi_token_moe_rows_equal_single_token_decode() {
14174        use crate::pool::Pool;
14175        use cortiq_core::{CMF_VERSION, CmfHeader, CmfModel, QuantType, TensorSpec};
14176
14177        let (h, inter, ne) = (64usize, 128usize, 3usize);
14178        let arch: cortiq_core::ModelArch = serde_json::from_value(serde_json::json!({
14179            "arch_name": "tiny-q4tp-moe",
14180            "hidden_size": h,
14181            "intermediate_size": inter,
14182            "num_layers": 1,
14183            "num_attention_heads": 2,
14184            "num_kv_heads": 1,
14185            "head_dim": 32,
14186            "vocab_size": 8,
14187            "layer_types": ["FullAttention"],
14188            "rms_norm_eps": 1e-6,
14189            "max_position_embeddings": 8,
14190            "linear_conv_kernel_dim": 0,
14191            "linear_num_key_heads": 0,
14192            "linear_num_value_heads": 0
14193        }))
14194        .unwrap();
14195        let header = CmfHeader {
14196            format: "cmf".into(),
14197            version: CMF_VERSION,
14198            arch,
14199            quant_type: QuantType::Q4Block,
14200            provenance: None,
14201            tokenizer_config: None,
14202            section_hashes: None,
14203            skills: Vec::new(),
14204            shard: None,
14205            calibration: None,
14206            routing: None,
14207            genome: None,
14208            lineage: Vec::new(),
14209            router: None,
14210            segments: Vec::new(),
14211        };
14212        let mut specs = Vec::new();
14213        for e in 0..ne {
14214            for (k, (n, r, c)) in [("g", inter, h), ("u", inter, h), ("d", h, inter)]
14215                .into_iter()
14216                .enumerate()
14217            {
14218                // Distinct experts: perturb only the nibble plane (any byte
14219                // is a valid pair of codes; the ladder stays intact).
14220                let mut data = synth_q4tp(r, c);
14221                for (i, byte) in data[..r * (c / GROUP_SIZE) * Q4TP_NIB]
14222                    .iter_mut()
14223                    .enumerate()
14224                {
14225                    *byte ^= ((i * (e * 3 + k + 1)) % 251) as u8;
14226                }
14227                specs.push(TensorSpec {
14228                    name: format!("{n}{e}"),
14229                    dtype: TensorDtype::Q4TiledP,
14230                    shape: vec![r, c],
14231                    data,
14232                });
14233            }
14234        }
14235        let dir = std::env::temp_dir().join(format!(
14236            "cmf-moe-rows-{}-{}",
14237            std::process::id(),
14238            FLOAT_ACTIVATIONS.get()
14239        ));
14240        std::fs::create_dir_all(&dir).unwrap();
14241        let path = dir.join("m.cmf");
14242        CmfModel::write(&path, &header, &specs, None, None).unwrap();
14243        let model = Arc::new(CmfModel::open(&path).unwrap());
14244        let t = |n: String| QTensor::from_model(&model, &n).unwrap();
14245        let g: Vec<QTensor> = (0..ne).map(|e| t(format!("g{e}"))).collect();
14246        let u: Vec<QTensor> = (0..ne).map(|e| t(format!("u{e}"))).collect();
14247        let d: Vec<QTensor> = (0..ne).map(|e| t(format!("d{e}"))).collect();
14248        let b = 4usize;
14249        let mut xs: Vec<f32> = (0..b * h)
14250            .map(|i| ((i * 31 + 7) % 89) as f32 / 89.0 - 0.5)
14251            .collect();
14252        xs[5] = 9.0; // an activation outlier on token 0
14253        // Token -> (experts in route order, weights).
14254        let routes: Vec<(Vec<usize>, Vec<f32>)> = vec![
14255            (vec![2, 0], vec![0.6, 0.4]),
14256            (vec![0, 1, 2], vec![0.2, 0.5, 0.3]),
14257            (vec![1], vec![1.0]),
14258            (vec![2, 1, 0], vec![0.25, 0.25, 0.5]),
14259        ];
14260        let pool = Pool::new(3);
14261        // Decode reference, token by token.
14262        let mut want = vec![0f32; b * h];
14263        for (tk, (idx, w)) in routes.iter().enumerate() {
14264            let x = &xs[tk * h..(tk + 1) * h];
14265            let pairs: Vec<(&QTensor, &QTensor)> = idx.iter().map(|&e| (&g[e], &u[e])).collect();
14266            let mut gs: Vec<Vec<f32>> = idx.iter().map(|_| vec![0f32; inter]).collect();
14267            assert!(QTensor::moe_gate_up_many(&pairs, x, &mut gs, Some(&pool)));
14268            if FLOAT_ACTIVATIONS.get() {
14269                for (slot, &e) in idx.iter().enumerate() {
14270                    let (mut gate, mut up) = (vec![0.0; inter], vec![0.0; inter]);
14271                    g[e].matvec(x, &mut gate, Some(&pool));
14272                    u[e].matvec(x, &mut up, Some(&pool));
14273                    for (v, u) in gate.iter_mut().zip(up) {
14274                        *v = (*v / (1.0 + (-*v).exp())) * u;
14275                    }
14276                    assert_eq!(gs[slot], gate, "float gate/up must equal ordinary matvecs");
14277                }
14278            }
14279            let downs: Vec<&QTensor> = idx.iter().map(|&e| &d[e]).collect();
14280            assert!(QTensor::moe_down_many(
14281                &downs,
14282                &gs,
14283                w,
14284                &mut want[tk * h..(tk + 1) * h],
14285                Some(&pool)
14286            ));
14287        }
14288        if FLOAT_ACTIVATIONS.get() {
14289            for (tk, (idx, w)) in routes.iter().enumerate() {
14290                let mut scalar = vec![0.0; h];
14291                for (&e, &weight) in idx.iter().zip(w) {
14292                    let (mut gate, mut up, mut down) =
14293                        (vec![0.0; inter], vec![0.0; inter], vec![0.0; h]);
14294                    g[e].matvec(&xs[tk * h..(tk + 1) * h], &mut gate, Some(&pool));
14295                    u[e].matvec(&xs[tk * h..(tk + 1) * h], &mut up, Some(&pool));
14296                    for (v, u) in gate.iter_mut().zip(up) {
14297                        *v = (*v / (1.0 + (-*v).exp())) * u;
14298                    }
14299                    d[e].matvec(&gate, &mut down, Some(&pool));
14300                    for (v, d) in scalar.iter_mut().zip(down) {
14301                        *v += weight * d;
14302                    }
14303                }
14304                assert_eq!(
14305                    &want[tk * h..(tk + 1) * h],
14306                    scalar,
14307                    "float many equals scalar experts"
14308                );
14309            }
14310        }
14311        // All four tokens at once, grouped by expert.
14312        let mut experts: Vec<usize> = Vec::new();
14313        let mut groups: Vec<Vec<usize>> = Vec::new();
14314        for (tk, (idx, _)) in routes.iter().enumerate() {
14315            for &e in idx {
14316                match experts.iter().position(|&x| x == e) {
14317                    Some(k) => groups[k].push(tk),
14318                    None => {
14319                        experts.push(e);
14320                        groups.push(vec![tk]);
14321                    }
14322                }
14323            }
14324        }
14325        let n_pairs: usize = groups.iter().map(|g| g.len()).sum();
14326        let pairs: Vec<(&QTensor, &QTensor)> = experts.iter().map(|&e| (&g[e], &u[e])).collect();
14327        let mut gs: Vec<Vec<f32>> = (0..n_pairs).map(|_| vec![0f32; inter]).collect();
14328        assert!(QTensor::moe_gate_up_rows(
14329            &pairs,
14330            &groups,
14331            &xs,
14332            &mut gs,
14333            Some(&pool)
14334        ));
14335        let downs: Vec<&QTensor> = experts.iter().map(|&e| &d[e]).collect();
14336        let lens: Vec<usize> = groups.iter().map(|g| g.len()).collect();
14337        let mut ds: Vec<Vec<f32>> = (0..n_pairs).map(|_| vec![0f32; h]).collect();
14338        assert!(QTensor::moe_down_rows(
14339            &downs,
14340            &lens,
14341            &gs,
14342            &mut ds,
14343            Some(&pool)
14344        ));
14345        let slot = |tk: usize, e: usize| {
14346            let k = experts.iter().position(|&x| x == e).unwrap();
14347            groups[..k].iter().map(|g| g.len()).sum::<usize>()
14348                + groups[k].iter().position(|&x| x == tk).unwrap()
14349        };
14350        let mut got = vec![0f32; b * h];
14351        for (tk, (idx, w)) in routes.iter().enumerate() {
14352            for i in 0..h {
14353                let mut acc = 0f32;
14354                for (&e, &we) in idx.iter().zip(w) {
14355                    acc += we * ds[slot(tk, e)][i];
14356                }
14357                got[tk * h + i] = acc;
14358            }
14359        }
14360        assert!(want.iter().any(|v| *v != 0.0));
14361        assert_eq!(
14362            want.iter().map(|v| v.to_bits()).collect::<Vec<_>>(),
14363            got.iter().map(|v| v.to_bits()).collect::<Vec<_>>(),
14364            "multi-token MoE must equal decode bit for bit"
14365        );
14366
14367        // Row-exact q4tp_matmat: five tokens (a blocked 1x4 tile + tail
14368        // otherwise) equal five matvecs.
14369        let b5 = 5usize;
14370        let x5: Vec<f32> = (0..b5 * h)
14371            .map(|i| ((i * 13 + 5) % 71) as f32 / 71.0 - 0.5)
14372            .collect();
14373        let mut mm = vec![0f32; b5 * inter];
14374        row_exact_scope(|| g[1].matmat(&x5, b5, &mut mm, Some(&pool)));
14375        for tk in 0..b5 {
14376            let mut mv = vec![0f32; inter];
14377            g[1].matvec(&x5[tk * h..(tk + 1) * h], &mut mv, Some(&pool));
14378            assert_eq!(
14379                mv.iter().map(|v| v.to_bits()).collect::<Vec<_>>(),
14380                mm[tk * inter..(tk + 1) * inter]
14381                    .iter()
14382                    .map(|v| v.to_bits())
14383                    .collect::<Vec<_>>(),
14384                "row-exact matmat token {tk}"
14385            );
14386        }
14387        // Other concurrent tests/requests may still hold the shared mode.
14388        // Nested, overlapping and unwind restoration is checked separately.
14389        let _ = std::fs::remove_dir_all(&dir);
14390    }
14391
14392    /// The row-exact fix must not touch the fast path. Outside the scope
14393    /// `q4tp_matmat` has to produce exactly what it did before, and on ARM
14394    /// "before" is spelled out below: the tuned 1x4 SDOT tile for every
14395    /// four columns and the single-row kernel for the tail. Inside the
14396    /// scope every column equals its token's matvec. `q4tp_matmat_with`
14397    /// takes the mode as an argument, so a concurrent test holding the
14398    /// shared scope cannot flip it under this one.
14399    #[test]
14400    fn q4tp_matmat_fast_path_unchanged_outside_row_exact() {
14401        use crate::pool::Pool;
14402        use std::sync::atomic::Ordering::Relaxed;
14403        let _alt = Q4TP_ALT_TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner());
14404        // The tuned ARM shape, which is also what an unset switch picks
14405        // unless CMF_Q4TP_V1 is exported.
14406        Q4TP_ALT.store(2, Relaxed);
14407        let pool = Pool::new(3);
14408        // Under 500k cells, so macOS keeps the matmat off the AMX; the
14409        // second shape runs across pool workers (rows >= 256) with 32
14410        // groups of accumulation and a tail after two 1x4 tiles.
14411        for &(rows, cols, b) in &[(64usize, 256usize, 7usize), (320, 1024, 9)] {
14412            let bytes = synth_q4tp(rows, cols);
14413            let mut xs: Vec<f32> = (0..b * cols)
14414                .map(|i| ((i * 29 + 11) % 83) as f32 / 83.0 - 0.5)
14415                .collect();
14416            xs[3] = 7.5; // an activation outlier on token 0
14417            let run = |exact: bool| {
14418                let mut out = vec![0f32; b * rows];
14419                q4tp_matmat_with(&bytes, &xs, b, rows, cols, &mut out, Some(&pool), exact);
14420                out
14421            };
14422            let (fast, exact) = (run(false), run(true));
14423            #[cfg(not(target_arch = "aarch64"))]
14424            let _ = fast;
14425            let bits = |v: &[f32]| v.iter().map(|x| x.to_bits()).collect::<Vec<_>>();
14426            let mut matvecs = vec![0f32; b * rows];
14427            for (bi, o) in matvecs.chunks_mut(rows).enumerate() {
14428                q4tp_matvec(
14429                    &bytes,
14430                    &xs[bi * cols..(bi + 1) * cols],
14431                    rows,
14432                    cols,
14433                    o,
14434                    Some(&pool),
14435                );
14436            }
14437            assert!(matvecs.iter().any(|v| *v != 0.0));
14438            assert_eq!(
14439                bits(&exact),
14440                bits(&matvecs),
14441                "{rows}x{cols} b={b}: row-exact matmat must equal per-token matvecs"
14442            );
14443            #[cfg(target_arch = "aarch64")]
14444            {
14445                // The pre-fix ARM loop, cell for cell.
14446                let gpr = cols / GROUP_SIZE;
14447                let v = Q4tpView::new(&bytes, rows, cols);
14448                let mut old = vec![0f32; b * rows];
14449                let mut sc = vec![0f32; gpr];
14450                let a8w8 = a8w8_enabled();
14451                let blocked = sdot_enabled() && blocked_enabled();
14452                let acts: Vec<SplitAct> = (0..b)
14453                    .map(|bi| split_act(&xs[bi * cols..(bi + 1) * cols]))
14454                    .collect();
14455                for r in 0..rows {
14456                    v.scales_into(r, gpr, &mut sc);
14457                    if !a8w8 {
14458                        for bi in 0..b {
14459                            let x = &xs[bi * cols..(bi + 1) * cols];
14460                            old[bi * rows + r] = q4tp_row_exact(v.nib, r, gpr, x, &sc);
14461                        }
14462                        continue;
14463                    }
14464                    let finish = |d: f32, act: &SplitAct| {
14465                        let mut acc = d * act.sx;
14466                        for &(j, xv) in &act.outliers {
14467                            let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
14468                            acc += w * s * xv;
14469                        }
14470                        acc
14471                    };
14472                    let mut bi = 0usize;
14473                    while blocked && bi + 4 <= b {
14474                        let xs4 = [
14475                            acts[bi].xq.as_slice(),
14476                            acts[bi + 1].xq.as_slice(),
14477                            acts[bi + 2].xq.as_slice(),
14478                            acts[bi + 3].xq.as_slice(),
14479                        ];
14480                        let d = unsafe { dot_q4tp_row_1x4_sdot(v.nib, r, gpr, xs4, &sc) };
14481                        for k in 0..4 {
14482                            old[(bi + k) * rows + r] = finish(d[k], &acts[bi + k]);
14483                        }
14484                        bi += 4;
14485                    }
14486                    for (bi, act) in acts.iter().enumerate().skip(bi) {
14487                        let d = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, &sc);
14488                        old[bi * rows + r] = finish(d, act);
14489                    }
14490                }
14491                assert_eq!(
14492                    bits(&fast),
14493                    bits(&old),
14494                    "{rows}x{cols} b={b}: the fast path changed outside row_exact"
14495                );
14496                // And it is still the fast tile that runs: its lane-parallel
14497                // fma sum rounds differently from the matvec somewhere.
14498                if blocked {
14499                    assert_ne!(
14500                        bits(&fast),
14501                        bits(&matvecs),
14502                        "{rows}x{cols} b={b}: the tuned tile no longer runs outside row_exact"
14503                    );
14504                }
14505            }
14506        }
14507        Q4TP_ALT.store(0, Relaxed);
14508    }
14509
14510    #[test]
14511    fn row_exact_scopes_survive_overlap_nesting_and_unwind() {
14512        use std::sync::{Barrier, atomic::{AtomicUsize, Ordering}};
14513        // A private counter makes this restoration test independent of
14514        // numerical tests concurrently using the production counter.
14515        let active = AtomicUsize::new(0);
14516        counted_row_exact_scope(&active, || {
14517            assert_eq!(active.load(Ordering::Acquire), 1);
14518            counted_row_exact_scope(&active, || {
14519                assert_eq!(active.load(Ordering::Acquire), 2);
14520            });
14521            assert_eq!(active.load(Ordering::Acquire), 1);
14522        });
14523        assert_eq!(active.load(Ordering::Acquire), 0);
14524
14525        let both_entered = Barrier::new(2);
14526        let release_last = Barrier::new(2);
14527        std::thread::scope(|s| {
14528            let first = s.spawn(|| counted_row_exact_scope(&active, || {
14529                both_entered.wait();
14530            }));
14531            let last = s.spawn(|| counted_row_exact_scope(&active, || {
14532                both_entered.wait();
14533                release_last.wait();
14534            }));
14535            first.join().unwrap();
14536            let after_first = active.load(Ordering::Acquire);
14537            release_last.wait();
14538            last.join().unwrap();
14539            assert_eq!(after_first, 1, "second request must remain exact");
14540        });
14541        assert_eq!(active.load(Ordering::Acquire), 0);
14542        let panic = std::panic::catch_unwind(|| {
14543            counted_row_exact_scope(&active, || panic!("scope unwind"));
14544        });
14545        assert!(panic.is_err());
14546        assert_eq!(active.load(Ordering::Acquire), 0);
14547    }
14548
14549    #[test]
14550    #[cfg(target_arch = "x86_64")]
14551    fn q4tp_float_avx2_is_bitwise_scalar() {
14552        if !avx2_enabled() {
14553            return;
14554        }
14555        for cols in [32, 64, 96, 2048, 4096] {
14556            let rows = 9;
14557            let bytes = synth_q4tp(rows, cols);
14558            let v = Q4tpView::new(&bytes, rows, cols);
14559            let gpr = cols / GROUP_SIZE;
14560            let mut sc = vec![0.0; gpr];
14561            for seed in 1..=5 {
14562                let xs: Vec<f32> = (0..cols)
14563                    .map(|i| (((i * 104729 + seed * 8191) % 100003) as f32 - 50001.0) / 7919.0)
14564                    .collect();
14565                for r in 0..rows {
14566                    v.scales_into(r, gpr, &mut sc);
14567                    let scalar = q4tp_row_float_scalar(v.nib, r, gpr, &xs, &sc);
14568                    let vector = unsafe { q4tp_row_float_avx2(v.nib, r, gpr, &xs, &sc) };
14569                    assert_eq!(
14570                        scalar.to_bits(),
14571                        vector.to_bits(),
14572                        "cols={cols} row={r} seed={seed}"
14573                    );
14574                }
14575            }
14576        }
14577    }
14578
14579    #[test]
14580    fn multi_token_moe_rows_float_equal_single_token_decode() {
14581        float_activations_scope(multi_token_moe_rows_equal_single_token_decode);
14582    }
14583
14584    #[test]
14585    fn full_gpu_q8_scope_is_nested_and_thread_local() {
14586        assert!(!FULL_GPU_Q8.get());
14587        let before = gpu_split_frac();
14588        {
14589            let _guard = enter_full_gpu_q8_scope();
14590            assert_eq!(gpu_split_frac(), 1.0);
14591            {
14592                let _nested = enter_full_gpu_q8_scope();
14593            }
14594            assert_eq!(gpu_split_frac(), 1.0);
14595            std::thread::spawn(|| assert!(!FULL_GPU_Q8.get())).join().unwrap();
14596        }
14597        assert!(!FULL_GPU_Q8.get());
14598        assert_eq!(gpu_split_frac(), before);
14599    }
14600
14601    #[test]
14602    fn float_activation_scope_is_nested_thread_local_and_unwind_safe() {
14603        assert!(!FLOAT_ACTIVATIONS.get());
14604        let before = a8w8_enabled();
14605        float_activations_scope(|| {
14606            assert!(!a8w8_enabled());
14607            float_activations_scope(|| assert!(!a8w8_enabled()));
14608            assert!(FLOAT_ACTIVATIONS.get());
14609            std::thread::spawn(|| assert!(!FLOAT_ACTIVATIONS.get()))
14610                .join()
14611                .unwrap();
14612        });
14613        assert!(!FLOAT_ACTIVATIONS.get());
14614        assert_eq!(a8w8_enabled(), before);
14615        let _ = std::panic::catch_unwind(|| float_activations_scope(|| panic!("test unwind")));
14616        assert!(!FLOAT_ACTIVATIONS.get());
14617    }
14618
14619    /// Batched q4/vbit matmat must equal per-position matvec calls
14620    /// exactly (the fallback it replaced) — same kernels, same order.
14621    #[test]
14622    fn batched_matmat_equals_per_position_matvec() {
14623        let (rows, cols, b) = (8, 64, 5);
14624        // q4 blob.
14625        let groups = rows * cols / GROUP_SIZE;
14626        let mut q4 = Vec::new();
14627        for i in 0..groups * 16 {
14628            q4.push((((i * 7 + 3) % 256) & 0xFF) as u8);
14629        }
14630        for g in 0..groups {
14631            q4.extend_from_slice(
14632                &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
14633            );
14634        }
14635        // vbit blob (mixed widths incl. 8).
14636        let ng = cols / GROUP_SIZE;
14637        let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4, 5, 3];
14638        let mut vb = bits.clone();
14639        for g in 0..rows * ng {
14640            vb.extend_from_slice(
14641                &cortiq_core::quant::f32_to_f16(0.02 + 0.001 * g as f32).to_le_bytes(),
14642            );
14643        }
14644        for r in 0..rows {
14645            let bw = bits[r] as usize;
14646            let (mut acc, mut nb) = (0u64, 0usize);
14647            let mut rowbytes = Vec::new();
14648            for i in 0..cols {
14649                let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
14650                acc = (acc << bw) | v;
14651                nb += bw;
14652                while nb >= 8 {
14653                    nb -= 8;
14654                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
14655                }
14656            }
14657            if nb > 0 {
14658                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
14659            }
14660            vb.extend_from_slice(&rowbytes);
14661        }
14662        let offsets = vbit_row_offsets(&vb, rows, cols);
14663
14664        let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
14665
14666        // q4: batch vs singles.
14667        let mut got = vec![0f32; b * rows];
14668        q4matmat(&q4, &xs, b, rows, cols, &mut got, None);
14669        for bi in 0..b {
14670            let mut expect = vec![0f32; rows];
14671            q4matvec(
14672                &q4,
14673                &xs[bi * cols..(bi + 1) * cols],
14674                rows,
14675                cols,
14676                &mut expect,
14677                None,
14678            );
14679            assert_eq!(
14680                &got[bi * rows..(bi + 1) * rows],
14681                &expect[..],
14682                "q4 batch pos {bi}"
14683            );
14684        }
14685
14686        // vbit: batch vs singles.
14687        let mut got = vec![0f32; b * rows];
14688        vbitmatmat(&vb, &offsets, &xs, b, rows, cols, &mut got, None);
14689        for bi in 0..b {
14690            let mut expect = vec![0f32; rows];
14691            vbitmatvec(
14692                &vb,
14693                &offsets,
14694                &xs[bi * cols..(bi + 1) * cols],
14695                rows,
14696                cols,
14697                &mut expect,
14698                None,
14699            );
14700            assert_eq!(
14701                &got[bi * rows..(bi + 1) * rows],
14702                &expect[..],
14703                "vbit batch pos {bi}"
14704            );
14705        }
14706    }
14707
14708    /// q4_tiled kernels must produce BIT-identical outputs to the q4
14709    /// split kernels on the same values (same ints, same order — only
14710    /// the byte placement differs).
14711    #[test]
14712    fn q4_tiled_matches_q4_block_bitexact() {
14713        let (rows, cols, b) = (8usize, 128usize, 3usize);
14714        let groups = rows * cols / GROUP_SIZE;
14715        let mut split = Vec::with_capacity(groups * 18);
14716        for i in 0..groups * 16 {
14717            split.push((((i * 7 + 3) % 256) & 0xFF) as u8);
14718        }
14719        for g in 0..groups {
14720            split.extend_from_slice(
14721                &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
14722            );
14723        }
14724        // Re-tile: [scale][nibbles] per group.
14725        let (packed, scales) = split.split_at(groups * 16);
14726        let mut tiled = Vec::with_capacity(groups * Q4_TILE);
14727        for g in 0..groups {
14728            tiled.extend_from_slice(&scales[g * 2..g * 2 + 2]);
14729            tiled.extend_from_slice(&packed[g * 16..(g + 1) * 16]);
14730        }
14731
14732        let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
14733        x1[9] = 250.0; // exercise the outlier path
14734        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
14735
14736        let (mut a, mut t) = (vec![0f32; rows], vec![0f32; rows]);
14737        q4matvec(&split, &x1, rows, cols, &mut a, None);
14738        q4t_matvec(&tiled, &x1, rows, cols, &mut t, None);
14739        assert_eq!(a, t, "q4t matvec must match q4 bit-for-bit");
14740
14741        let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
14742        let (mut t1, mut t2) = (vec![0f32; rows], vec![0f32; rows]);
14743        q4matvec2(&split, &x1, &x2, rows, cols, &mut a1, &mut a2, None);
14744        q4t_matvec2(&tiled, &x1, &x2, rows, cols, &mut t1, &mut t2, None);
14745        assert_eq!(a1, t1);
14746        assert_eq!(a2, t2);
14747
14748        let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
14749        let (mut am, mut tm) = (vec![0f32; b * rows], vec![0f32; b * rows]);
14750        q4matmat(&split, &xs, b, rows, cols, &mut am, None);
14751        q4t_matmat(&tiled, &xs, b, rows, cols, &mut tm, None);
14752        assert_eq!(am, tm, "q4t matmat must match q4 bit-for-bit");
14753    }
14754
14755    /// q4 SDOT outlier correction: a single huge activation channel
14756    /// (>8·rms → outlier, zeroed in xq) must still contribute its EXACT
14757    /// term. On-grid bulk (±1/0 → xq dequantizes exactly) isolates the
14758    /// correction from A8W8 noise. cols must exceed 64: at n=64 the
14759    /// 8·rms threshold equals sqrt(v²+rest) ≥ v, so a single outlier
14760    /// can never qualify (8² = n).
14761    #[test]
14762    fn q4matvec_sdot_outlier_exact() {
14763        let (rows, cols) = (4, 128);
14764        let groups = rows * cols / GROUP_SIZE;
14765        let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
14766        for i in 0..groups * 16 {
14767            bytes.push(((i * 11 + 5) % 256) as u8);
14768        }
14769        for g in 0..groups {
14770            let s = 0.02 + 0.002 * g as f32;
14771            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
14772        }
14773        let mut x: Vec<f32> = (0..cols)
14774            .map(|i| match i % 3 {
14775                0 => 1.0,
14776                1 => -1.0,
14777                _ => 0.0,
14778            })
14779            .collect();
14780        x[17] = 300.0; // ≫ 8·rms → outlier channel
14781
14782        let mut reference = vec![0.0f32; rows * cols];
14783        cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
14784        let mut expect = vec![0.0f32; rows];
14785        for r in 0..rows {
14786            expect[r] = reference[r * cols..(r + 1) * cols]
14787                .iter()
14788                .zip(&x)
14789                .map(|(w, xv)| w * xv)
14790                .sum();
14791        }
14792        let mut got = vec![0.0f32; rows];
14793        q4matvec(&bytes, &x, rows, cols, &mut got, None);
14794        let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
14795        for r in 0..rows {
14796            assert!(
14797                (got[r] - expect[r]).abs() < 2e-3 * scale,
14798                "row {r}: {} vs {} (outlier term must be exact)",
14799                got[r],
14800                expect[r]
14801            );
14802        }
14803    }
14804
14805    /// The fused q1t matvec must equal the reference (dequant_q1t → dot),
14806    /// including the ternary zero level and the binary-searched outlier
14807    /// overlay. Guards the mmap kernel that makes a 12B q1t runnable.
14808    #[test]
14809    fn q1t_matvec_matches_reference() {
14810        use cortiq_core::quant::{dequant_q1t, f32_to_f16};
14811        let (rows, cols) = (3usize, 64usize); // gpr = 2
14812        let gpr = cols / GROUP_SIZE;
14813        let scales = [0.5f32, 0.3, 0.7, 0.2, 0.6, 0.15];
14814        // Overlay (must be sorted by flat index): a few spikes across rows.
14815        let outliers: [(u32, f32); 3] = [(5, 9.0), (70, -4.5), (150, 3.25)];
14816        let is_out = |flat: usize| outliers.iter().any(|&(i, _)| i as usize == flat);
14817        let mut bytes = Vec::new();
14818        for r in 0..rows {
14819            for g in 0..gpr {
14820                bytes.extend_from_slice(&f32_to_f16(scales[r * gpr + g]).to_le_bytes());
14821                let mut c = [0u8; 7];
14822                for k in 0..GROUP_SIZE {
14823                    // Encoder invariant: code 0 at outlier positions.
14824                    let code = if is_out(r * cols + g * GROUP_SIZE + k) {
14825                        0
14826                    } else {
14827                        ((k + r * 3 + g) % 3) as u8 // 0,1,2
14828                    };
14829                    cortiq_core::quant::q1t_pack(&mut c, k, code);
14830                }
14831                bytes.extend_from_slice(&c);
14832            }
14833        }
14834        // Per-row overlay: [u32 row_ptr[rows+1]] then [(u16 col, f16 val)] by
14835        // row (outliers are sorted by flat index → already grouped by row).
14836        let mut row_ptr = vec![0u32; rows + 1];
14837        for &(idx, _) in &outliers {
14838            row_ptr[idx as usize / cols + 1] += 1;
14839        }
14840        for r in 0..rows {
14841            row_ptr[r + 1] += row_ptr[r];
14842        }
14843        for &p in &row_ptr {
14844            bytes.extend_from_slice(&p.to_le_bytes());
14845        }
14846        for &(idx, v) in &outliers {
14847            bytes.extend_from_slice(&((idx as usize % cols) as u16).to_le_bytes());
14848            bytes.extend_from_slice(&f32_to_f16(v).to_le_bytes());
14849        }
14850
14851        let mut refw = vec![0f32; rows * cols];
14852        dequant_q1t(&bytes, rows, cols, &mut refw);
14853        // On-grid activations (±1, amax 1) so the int8 SDOT path reconstructs
14854        // x exactly and matches the f32 reference (same trick as the q1 test).
14855        let x: Vec<f32> = (0..cols)
14856            .map(|j| if j % 3 == 0 { 1.0 } else { -1.0 })
14857            .collect();
14858        let mut expect = vec![0f32; rows];
14859        for r in 0..rows {
14860            let mut a = 0.0f32;
14861            for j in 0..cols {
14862                a += refw[r * cols + j] * x[j];
14863            }
14864            expect[r] = a;
14865        }
14866        let tol = |e: f32| 1e-3 * e.abs().max(1e-3);
14867        let mut got = vec![0f32; rows];
14868        q1t_matvec(&bytes, &x, rows, cols, &mut got, None);
14869        for r in 0..rows {
14870            assert!(
14871                (got[r] - expect[r]).abs() < tol(expect[r]),
14872                "row {r}: {} vs {}",
14873                got[r],
14874                expect[r]
14875            );
14876        }
14877        // matmat (b=2, f32 decode path) must agree too.
14878        let x2: Vec<f32> = x.iter().chain(x.iter()).copied().collect();
14879        let mut gm = vec![0f32; 2 * rows];
14880        q1t_matmat(&bytes, &x2, 2, rows, cols, &mut gm, None);
14881        for r in 0..rows {
14882            assert!((gm[r] - expect[r]).abs() < tol(expect[r]));
14883            assert!((gm[rows + r] - expect[r]).abs() < tol(expect[r]));
14884        }
14885        // Fused pair (q1t_matvec2) must equal two single matvecs
14886        // bit-for-bit: same unpack, same group order, same f32
14887        // accumulation per stream. Distinct x2 exercises both lanes.
14888        let xb: Vec<f32> = (0..cols)
14889            .map(|j| if j % 5 == 0 { -1.0 } else { 1.0 })
14890            .collect();
14891        let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
14892        q1t_matvec(&bytes, &x, rows, cols, &mut s1, None);
14893        q1t_matvec(&bytes, &xb, rows, cols, &mut s2, None);
14894        let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
14895        q1t_matvec2(&bytes, &x, &xb, rows, cols, &mut p1, &mut p2, None);
14896        assert_eq!(p1, s1, "q1t pair lane 1 ≠ single matvec");
14897        assert_eq!(p2, s2, "q1t pair lane 2 ≠ single matvec");
14898    }
14899
14900    /// Pair == 2×matvec with an ODD group count (the kernel's tail
14901    /// group) and no overlay section.
14902    #[test]
14903    fn q1t_matvec2_odd_gpr_matches_singles() {
14904        use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_pack};
14905        let (rows, cols) = (5usize, 96usize); // gpr = 3 → paired + tail
14906        let gpr = cols / GROUP_SIZE;
14907        let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE);
14908        for r in 0..rows {
14909            for g in 0..gpr {
14910                bytes.extend_from_slice(&f32_to_f16(0.1 + 0.05 * (r + g) as f32).to_le_bytes());
14911                let mut c = [0u8; 7];
14912                for k in 0..GROUP_SIZE {
14913                    q1t_pack(&mut c, k, ((k * 7 + r * 5 + g * 3) % 3) as u8);
14914                }
14915                bytes.extend_from_slice(&c);
14916            }
14917        }
14918        let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.31).sin()).collect();
14919        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).cos()).collect();
14920        let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
14921        q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
14922        q1t_matvec(&bytes, &x2, rows, cols, &mut s2, None);
14923        let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
14924        q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
14925        assert_eq!(p1, s1, "odd-gpr pair lane 1 ≠ single");
14926        assert_eq!(p2, s2, "odd-gpr pair lane 2 ≠ single");
14927    }
14928
14929    // Speed A/B: fused pair (one unpack, two streams) vs two single
14930    // matvecs. Single-threaded, FFN-sized, min-of paired in-process.
14931    //   cargo test -p cortiq-engine --release q1t_matvec2_speed -- --ignored --nocapture
14932    #[test]
14933    #[ignore]
14934    fn q1t_matvec2_speed() {
14935        use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_pack};
14936        use std::time::Instant;
14937        let (rows, cols) = (8192usize, 4096usize);
14938        let gpr = cols / GROUP_SIZE;
14939        let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE);
14940        for r in 0..rows {
14941            for g in 0..gpr {
14942                let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
14943                bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
14944                let mut c = [0u8; 7];
14945                for k in 0..GROUP_SIZE {
14946                    q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
14947                }
14948                bytes.extend_from_slice(&c);
14949            }
14950        }
14951        let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.31).sin()).collect();
14952        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).cos()).collect();
14953        let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
14954        let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
14955        // Warm both paths once.
14956        q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
14957        q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
14958        let (mut t_pair, mut t_two) = (f64::MAX, f64::MAX);
14959        for _ in 0..8 {
14960            let t0 = Instant::now();
14961            q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
14962            t_pair = t_pair.min(t0.elapsed().as_secs_f64() * 1000.0);
14963            let t1 = Instant::now();
14964            q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
14965            q1t_matvec(&bytes, &x2, rows, cols, &mut s2, None);
14966            t_two = t_two.min(t1.elapsed().as_secs_f64() * 1000.0);
14967        }
14968        assert_eq!(p1, s1);
14969        assert_eq!(p2, s2);
14970        println!("q1t pair {rows}x{cols}: fused {t_pair:.2} ms | two singles {t_two:.2} ms");
14971    }
14972
14973    // Speed A/B: the base-3-division decode (what the packing commit left in
14974    // place) vs the fused sign-LUT matvec. Both single-threaded, same bytes.
14975    //   cargo test -p cortiq-engine q1t_matvec_speed -- --ignored --nocapture
14976    #[test]
14977    #[ignore]
14978    fn q1t_matvec_speed() {
14979        use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_code, q1t_pack};
14980        use std::time::Instant;
14981        let (rows, cols) = (8192usize, 4096usize); // FFN-sized
14982        let gpr = cols / GROUP_SIZE;
14983        let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE + 16);
14984        for r in 0..rows {
14985            for g in 0..gpr {
14986                let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
14987                bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
14988                let mut c = [0u8; 7];
14989                for k in 0..GROUP_SIZE {
14990                    q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
14991                }
14992                bytes.extend_from_slice(&c);
14993            }
14994        }
14995        let (n, stride) = (rows * cols, 40usize); // ~2.5% outliers, per-row overlay
14996        let mut row_ptr = vec![0u32; rows + 1];
14997        let mut idx = 0usize;
14998        while idx < n {
14999            row_ptr[idx / cols + 1] += 1;
15000            idx += stride;
15001        }
15002        for r in 0..rows {
15003            row_ptr[r + 1] += row_ptr[r];
15004        }
15005        for &p in &row_ptr {
15006            bytes.extend_from_slice(&p.to_le_bytes());
15007        }
15008        let mut idx = 0usize;
15009        while idx < n {
15010            bytes.extend_from_slice(&((idx % cols) as u16).to_le_bytes());
15011            bytes.extend_from_slice(&f32_to_f16((idx % 13) as f32 * 0.1 - 0.6).to_le_bytes());
15012            idx += stride;
15013        }
15014        // On-grid ±1 so the fast path's int8 SDOT is exact vs the f32 "slow"
15015        // reference (the A/B is a timing check; values must still agree).
15016        let x: Vec<f32> = (0..cols)
15017            .map(|j| if j % 3 == 0 { 1.0 } else { -1.0 })
15018            .collect();
15019        let (rp_off, ent_off, has_ov) = q1t_overlay(&bytes, rows * gpr * Q1T_TILE, rows);
15020
15021        // "before": base-3 division decode into a buffer, then dot.
15022        let slow = |out: &mut [f32]| {
15023            let mut buf = vec![0f32; cols];
15024            for r in 0..rows {
15025                for g in 0..gpr {
15026                    let off = (r * gpr + g) * Q1T_TILE;
15027                    let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
15028                    let codes = &bytes[off + 2..off + Q1T_TILE];
15029                    for k in 0..GROUP_SIZE {
15030                        buf[g * GROUP_SIZE + k] = match q1t_code(codes, k) {
15031                            1 => s,
15032                            2 => -s,
15033                            _ => 0.0,
15034                        };
15035                    }
15036                }
15037                out[r] = q1t_row_outlier_correction(&bytes, r, rp_off, ent_off, has_ov, &x)
15038                    + (0..cols).map(|j| buf[j] * x[j]).sum::<f32>();
15039            }
15040        };
15041        let iters = 5;
15042        let mut a = vec![0f32; rows];
15043        slow(&mut a); // warm
15044        let t = Instant::now();
15045        for _ in 0..iters {
15046            slow(&mut a);
15047        }
15048        let slow_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
15049
15050        let mut b = vec![0f32; rows];
15051        q1t_matvec(&bytes, &x, rows, cols, &mut b, None); // warm
15052        let t = Instant::now();
15053        for _ in 0..iters {
15054            q1t_matvec(&bytes, &x, rows, cols, &mut b, None);
15055        }
15056        let fast_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
15057
15058        for r in 0..rows {
15059            assert!((a[r] - b[r]).abs() < 1e-2, "mismatch row {r}");
15060        }
15061        println!(
15062            "q1t matvec {rows}x{cols} (1 thread): div-decode {slow_ms:.2} ms  fused-LUT {fast_ms:.2} ms  => {:.2}x",
15063            slow_ms / fast_ms
15064        );
15065    }
15066}
15067
15068#[cfg(test)]
15069mod gemm_bench {
15070    /// `cargo test -p cortiq-engine --release q4tp_matmat_throughput -- --ignored --nocapture`
15071    /// Times the batched q4tp GEMM at the shapes the image DiT runs
15072    /// (b=296 tokens, 2304 -> 9216), on synthetic bytes: no model, no
15073    /// mmap, no thermal drift over minutes — a kernel change shows up
15074    /// here in seconds where a full render hides it in noise.
15075    ///
15076    /// On macOS add `CMF_ACCEL=0`: this shape is over the 500k-cell mark
15077    /// where the matmat hands off to Accelerate's dequant sgemm, and
15078    /// without the opt-out both rows below measure the AMX, not the
15079    /// kernel under test.
15080    #[test]
15081    #[ignore]
15082    fn q4tp_matmat_throughput() {
15083        let _alt = super::Q4TP_ALT_TEST_LOCK
15084            .lock()
15085            .unwrap_or_else(|e| e.into_inner());
15086        // 296 is a prompt-encode batch; the image DiT runs 2085 at
15087        // 512x512, where the activation panel stops fitting L2 and the
15088        // loop's shape starts to matter more than its instructions.
15089        let b: usize = std::env::var("CMF_BENCH_B")
15090            .ok()
15091            .and_then(|v| v.parse().ok())
15092            .unwrap_or(296);
15093        let (rows, cols) = (9216usize, 2304usize);
15094        let (_, _, _) = (rows, cols, b);
15095        let total =
15096            cortiq_core::quant::expected_nbytes(cortiq_core::TensorDtype::Q4TiledP, &[rows, cols])
15097                .unwrap();
15098        // Random nibbles are fine, but the row params are f16 (lo, step)
15099        // of a geometric ladder: garbage there gives exp2 of a huge
15100        // exponent, the scales come back inf, and the whole bench times
15101        // NaN arithmetic instead of the kernel.
15102        let (params_off, codes_off, _) = cortiq_core::quant::q4tp_sections(rows, cols);
15103        let mut bytes: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
15104        let lo = cortiq_core::quant::f32_to_f16(-4.0);
15105        let step = cortiq_core::quant::f32_to_f16(0.1);
15106        for r in 0..rows {
15107            let o = params_off + r * 4;
15108            bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
15109            bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
15110        }
15111        let _ = codes_off;
15112        let xs: Vec<f32> = (0..b * cols)
15113            .map(|i| ((i % 97) as f32 - 48.0) / 48.0)
15114            .collect();
15115        let mut out = vec![0f32; b * rows];
15116        let pool = crate::pool::Pool::from_env();
15117        // A shared 48-core stand drifts ±25% run to run, which is wider
15118        // than any kernel change worth making. So: alternate the two
15119        // kernels inside one process and keep the BEST time for
15120        // each. Interleaving makes both see the same interference, and a
15121        // minimum is the one statistic another tenant cannot inflate.
15122        super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
15123        let reps: usize = std::env::var("CMF_BENCH_REPS")
15124            .ok()
15125            .and_then(|v| v.parse().ok())
15126            .unwrap_or(10);
15127        let mut best = [f64::MAX; 2];
15128        let mut sums = [0f32; 2];
15129        for _ in 0..reps {
15130            for (k, w) in [(0usize, 1u8), (1usize, 2u8)] {
15131                super::Q4TP_ALT.store(w, std::sync::atomic::Ordering::Relaxed);
15132                let t = std::time::Instant::now();
15133                super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
15134                best[k] = best[k].min(t.elapsed().as_secs_f64());
15135                sums[k] = out.iter().take(64).sum::<f32>();
15136            }
15137        }
15138        let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
15139        for (k, name) in ["previous", "tuned   "].iter().enumerate() {
15140            println!(
15141                "q4tp matmat {rows}x{cols} b={b} {name}: {:.1} ms  {:.1} GFLOP/s  (checksum {:.3})",
15142                best[k] * 1e3,
15143                flops / best[k] / 1e9,
15144                sums[k]
15145            );
15146        }
15147        assert!(
15148            (sums[0] - sums[1]).abs() < 1e-2,
15149            "the tuned kernel changed the result: {} vs {}",
15150            sums[0],
15151            sums[1]
15152        );
15153    }
15154
15155    /// The blocked kernel must agree with the per-column path exactly —
15156    /// same weights, same activation split, only a different instruction
15157    /// mix. Shapes are chosen to hit the awkward cases: a column count
15158    /// that leaves an odd group (the 512-bit kernel does two at a time),
15159    /// and a batch that does not divide by four.
15160    #[test]
15161    fn q4tp_matmat_blocked_matches_scalar() {
15162        use std::sync::atomic::Ordering::Relaxed;
15163        let _alt = super::Q4TP_ALT_TEST_LOCK
15164            .lock()
15165            .unwrap_or_else(|e| e.into_inner());
15166        // The last shape carries the image DiT's column count — 2304, so
15167        // 72 groups of accumulation, which is where a reordered sum can
15168        // actually drift — and runs through the thread pool, since the
15169        // blocked path splits rows across workers. Its row count stays
15170        // under 500k cells on purpose: above that, macOS diverts the whole
15171        // matmat to the Accelerate/AMX dequant sgemm and neither kernel
15172        // here would run.
15173        for &(rows, cols, b) in &[
15174            (64usize, 128usize, 7usize),
15175            (33, 96, 4),
15176            (16, 256, 9),
15177            (192, 2304, 37),
15178        ] {
15179            let total = cortiq_core::quant::expected_nbytes(
15180                cortiq_core::TensorDtype::Q4TiledP,
15181                &[rows, cols],
15182            )
15183            .unwrap();
15184            let (params_off, _, _) = cortiq_core::quant::q4tp_sections(rows, cols);
15185            let mut bytes: Vec<u8> = (0..total).map(|i| (i * 61 % 251) as u8).collect();
15186            let lo = cortiq_core::quant::f32_to_f16(-4.0);
15187            let step = cortiq_core::quant::f32_to_f16(0.1);
15188            for r in 0..rows {
15189                let o = params_off + r * 4;
15190                bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
15191                bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
15192            }
15193            let xs: Vec<f32> = (0..b * cols)
15194                .map(|i| ((i % 89) as f32 - 44.0) / 44.0)
15195                .collect();
15196            let mut got = vec![0f32; b * rows];
15197            let mut want = vec![0f32; b * rows];
15198            let gpr = cols / 32;
15199            let view = super::Q4tpView::new(&bytes, rows, cols);
15200            let pool = crate::pool::Pool::from_env();
15201            super::Q4TP_ALT.store(2, Relaxed);
15202            super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut got, pool.as_deref());
15203            super::Q4TP_ALT.store(1, Relaxed);
15204            super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut want, pool.as_deref());
15205            super::Q4TP_ALT.store(0, Relaxed);
15206            // Measured against the output's scale, not cell by cell: a
15207            // dot product of 2304 terms lands near zero wherever the row
15208            // and the activation nearly cancel, and there a per-cell
15209            // ratio reports 1e-3 for an absolute error of 5e-6 — f32's
15210            // own rounding, reordered. What must stay small is the error
15211            // relative to what the layer actually outputs.
15212            let scale = want.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
15213            let (mut worst, mut at) = (0f32, 0usize);
15214            for (i, (g, w)) in got.iter().zip(&want).enumerate() {
15215                if (g - w).abs() > worst {
15216                    worst = (g - w).abs();
15217                    at = i;
15218                }
15219            }
15220            assert!(
15221                worst <= 1e-4 * scale,
15222                "{rows}x{cols} b={b}: blocked and scalar disagree by {worst:.3e} \
15223                 (scale {scale:.3e}) at cell {at}: {} vs {}",
15224                got[at],
15225                want[at]
15226            );
15227
15228            // "Same speed, no quality loss" is a claim about which answer
15229            // is RIGHT, not about which two agree. Both paths sum the same
15230            // 2304 products in different orders, so f64 decides: the
15231            // blocked kernel keeps sixteen partial sums and folds them at
15232            // the end, which is a shallower addition tree than the
15233            // per-column path's running scalar, and it must not be worse.
15234            let (mut e_blocked, mut e_scalar) = (0f64, 0f64);
15235            for bi in 0..b {
15236                let act = super::split_act(&xs[bi * cols..(bi + 1) * cols]);
15237                for r in 0..rows {
15238                    let mut sc = vec![0f32; gpr];
15239                    view.scales_into(r, gpr, &mut sc);
15240                    let mut exact = 0f64;
15241                    for j in 0..cols {
15242                        let (w, sq) = super::q4tp_outlier(view.nib, r, gpr, j, &sc);
15243                        exact += w as f64 * sq as f64 * act.xq[j] as f64;
15244                    }
15245                    exact *= act.sx as f64;
15246                    for &(j, xv) in &act.outliers {
15247                        let (w, sq) = super::q4tp_outlier(view.nib, r, gpr, j, &sc);
15248                        exact += w as f64 * sq as f64 * xv as f64;
15249                    }
15250                    let i = bi * rows + r;
15251                    e_blocked = e_blocked.max((got[i] as f64 - exact).abs());
15252                    e_scalar = e_scalar.max((want[i] as f64 - exact).abs());
15253                }
15254            }
15255            println!(
15256                "{rows}x{cols} b={b}: worst error vs f64 — blocked {e_blocked:.3e}, \
15257                 per-column {e_scalar:.3e}"
15258            );
15259            // An absolute bar, not a race between the two: at these
15260            // magnitudes both sit in f32's last bits, and on a small shape
15261            // whichever one happens to round the unluckiest cell "wins" by
15262            // a factor the next seed reverses.
15263            assert!(
15264                e_blocked <= 1e-5 * scale as f64 && e_scalar <= 1e-5 * scale as f64,
15265                "{rows}x{cols} b={b}: error against f64 too large — blocked \
15266                 {e_blocked:.3e}, per-column {e_scalar:.3e}, scale {scale:.3e}"
15267            );
15268        }
15269    }
15270
15271    /// What the row-exact contract costs a speculative-verify panel on the
15272    /// host: the same batch through the fast arms, through the row-exact
15273    /// arms, and as one matvec per token (the other way to be exact).
15274    /// Arms alternate inside one process and keep their best time.
15275    /// `CMF_GPU=0 cargo test -p cortiq-engine --release q4tp_matmat_row_exact_cost -- --ignored --nocapture`
15276    #[test]
15277    #[ignore]
15278    fn q4tp_matmat_row_exact_cost() {
15279        let _alt = super::Q4TP_ALT_TEST_LOCK
15280            .lock()
15281            .unwrap_or_else(|e| e.into_inner());
15282        let pool = crate::pool::Pool::from_env();
15283        let reps: usize = std::env::var("CMF_BENCH_REPS")
15284            .ok()
15285            .and_then(|v| v.parse().ok())
15286            .unwrap_or(30);
15287        for &(rows, cols) in &[(2048usize, 4096usize), (4096, 2048), (4096, 4096)] {
15288            let total = cortiq_core::quant::expected_nbytes(
15289                cortiq_core::TensorDtype::Q4TiledP,
15290                &[rows, cols],
15291            )
15292            .unwrap();
15293            let (params_off, _, _) = cortiq_core::quant::q4tp_sections(rows, cols);
15294            let mut bytes: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
15295            let lo = cortiq_core::quant::f32_to_f16(-4.0);
15296            let step = cortiq_core::quant::f32_to_f16(0.1);
15297            for r in 0..rows {
15298                let o = params_off + r * 4;
15299                bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
15300                bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
15301            }
15302            for &b in &[2usize, 4, 5, 8] {
15303                let xs: Vec<f32> = (0..b * cols)
15304                    .map(|i| ((i % 97) as f32 - 48.0) / 48.0)
15305                    .collect();
15306                let mut out = vec![0f32; b * rows];
15307                let mut best = [f64::MAX; 3];
15308                for _ in 0..reps {
15309                    for (k, best_k) in best.iter_mut().enumerate() {
15310                        let t = std::time::Instant::now();
15311                        match k {
15312                            0 | 1 => super::q4tp_matmat_with(
15313                                &bytes,
15314                                &xs,
15315                                b,
15316                                rows,
15317                                cols,
15318                                &mut out,
15319                                pool.as_deref(),
15320                                k == 1,
15321                            ),
15322                            _ => {
15323                                for (bi, o) in out.chunks_mut(rows).enumerate() {
15324                                    super::q4tp_matvec(
15325                                        &bytes,
15326                                        &xs[bi * cols..(bi + 1) * cols],
15327                                        rows,
15328                                        cols,
15329                                        o,
15330                                        pool.as_deref(),
15331                                    );
15332                                }
15333                            }
15334                        }
15335                        *best_k = best_k.min(t.elapsed().as_secs_f64());
15336                    }
15337                }
15338                println!(
15339                    "q4tp {rows}x{cols} b={b}: fast {:.3} ms, row-exact {:.3} ms, \
15340                     {b} matvecs {:.3} ms",
15341                    best[0] * 1e3,
15342                    best[1] * 1e3,
15343                    best[2] * 1e3
15344                );
15345            }
15346        }
15347    }
15348
15349    /// The q4t twin of the throughput bench, same shape and rules, so the
15350    /// two quantisations' batch kernels can be read against each other.
15351    /// `cargo test -p cortiq-engine --release q4t_matmat_throughput -- --ignored --nocapture`
15352    #[test]
15353    #[ignore]
15354    fn q4t_matmat_throughput() {
15355        let (rows, cols, b) = (9216usize, 2304usize, 296usize);
15356        let total =
15357            cortiq_core::quant::expected_nbytes(cortiq_core::TensorDtype::Q4Tiled, &[rows, cols])
15358                .unwrap();
15359        // q4t carries a per-group f16 scale in the tile's first two bytes;
15360        // random bytes there decode to inf and the bench would time NaNs.
15361        let mut bytes: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
15362        let sc = cortiq_core::quant::f32_to_f16(0.02);
15363        for t in bytes.chunks_mut(super::Q4_TILE) {
15364            t[..2].copy_from_slice(&sc.to_le_bytes());
15365        }
15366        let xs: Vec<f32> = (0..b * cols)
15367            .map(|i| ((i % 97) as f32 - 48.0) / 48.0)
15368            .collect();
15369        let mut out = vec![0f32; b * rows];
15370        let pool = crate::pool::Pool::from_env();
15371        super::q4t_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
15372        let reps: usize = std::env::var("CMF_BENCH_REPS")
15373            .ok()
15374            .and_then(|v| v.parse().ok())
15375            .unwrap_or(10);
15376        let mut best = f64::MAX;
15377        for _ in 0..reps {
15378            let t = std::time::Instant::now();
15379            super::q4t_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
15380            best = best.min(t.elapsed().as_secs_f64());
15381        }
15382        let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
15383        println!(
15384            "q4t matmat {rows}x{cols} b={b}: {:.1} ms  {:.1} GFLOP/s  (checksum {:.3})",
15385            best * 1e3,
15386            flops / best / 1e9,
15387            out.iter().take(64).sum::<f32>()
15388        );
15389    }
15390}