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;
22use std::sync::atomic::{AtomicU64, Ordering};
23
24/// `mxfp4` (OCP MX FP4) CPU kernels.
25mod mxfp4;
26
27/// D5 GPU-matmat gate telemetry: large-GEMM calls that reached the gate
28/// vs those the q8_2f device kernel actually consumed (SAM3 profiling).
29pub static MATMAT_GATE: AtomicU64 = AtomicU64::new(0);
30pub static MATMAT_GPU_HIT: AtomicU64 = AtomicU64::new(0);
31
32pub enum QTensor {
33    F32 {
34        data: Vec<f32>,
35        rows: usize,
36        cols: usize,
37    },
38    Mapped {
39        model: Arc<CmfModel>,
40        /// Index into the model's tensor directory.
41        idx: usize,
42        dtype: TensorDtype,
43        rows: usize,
44        cols: usize,
45        /// Per-row scales, dequantized to f32 up front (tiny).
46        row_scale: Vec<f32>,
47        /// q8_2f column field (θ), dequantized up front; empty for q8_row.
48        col_field: Vec<f32>,
49        /// Vbit only: byte offset of each row's packed data within the
50        /// tensor blob (`[rows + 1]`, computed once at load — the per-
51        /// matvec prefix scan over row bit-widths was O(rows) each call).
52        vbit_offsets: Vec<usize>,
53        /// q8-family decode repack (load-time, optional): rows in groups
54        /// of 4, interleaved in 16-byte units — one 64-byte line per
55        /// iteration feeds all 4 sdot lanes, ONE sequential weight
56        /// stream per worker instead of four (this is where llama.cpp's
57        /// repacked Q8 kernels get their bandwidth). Empty = off
58        /// (CMF_REPACK=0, non-SDOT arch, or an ineligible shape). Trades
59        /// an anonymous copy of the quants for mmap pages that go cold.
60        repack: Vec<u8>,
61    },
62    /// Binary low-rank (`blr`): ±1 paths, signs from the mmap, scales as
63    /// f32 (docs/BLR_NVG.md §3). Every device-graph accessor answers `None`
64    /// for it (whole-token graphs decline to per-op/host paths); the CPU
65    /// kernels live in `crate::blr`, the per-op wgpu kernels in
66    /// `gpu_wgpu::blr`.
67    Blr(Box<crate::blr::BlrTensor>),
68}
69
70/// Load-time q8 repack gate (see `Mapped::repack`). OPT-IN
71/// (`CMF_REPACK=1`): the single-stream hypothesis LOST on Apple Silicon
72/// (M4, interleaved A/B: decode 101 vs 94 tok/s — four adjacent row
73/// streams per worker feed the prefetcher MORE memory-level parallelism
74/// than one); kept as an experiment flag for x86, where the tradeoff
75/// may land differently.
76fn repack_enabled() -> bool {
77    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
78    *ON.get_or_init(|| {
79        std::env::var("CMF_REPACK")
80            .map(|v| v == "1")
81            .unwrap_or(cfg!(target_os = "android"))
82    })
83}
84
85/// Interleave q8 rows for the decode kernel: group g holds rows
86/// 4g..4g+4 as [r0[c], r1[c], r2[c], r3[c]] per 16-byte chunk c. Only
87/// full groups are packed — tail rows keep reading the mmap layout.
88fn q8_repack(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
89    #[cfg(target_arch = "aarch64")]
90    let arch_ok = sdot_enabled();
91    #[cfg(not(target_arch = "aarch64"))]
92    let arch_ok = false;
93    if !arch_ok || !repack_enabled() || rows < 256 || cols % 16 != 0 {
94        return Vec::new();
95    }
96    q8_repack_layout(bytes, rows, cols)
97}
98
99/// The pure layout transform behind `q8_repack` (tested directly —
100/// the gate depends on arch and env).
101fn q8_repack_layout(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
102    let groups = rows / 4;
103    let mut rep = vec![0u8; groups * 4 * cols];
104    for g in 0..groups {
105        let dst = &mut rep[g * 4 * cols..(g + 1) * 4 * cols];
106        for c in 0..cols / 16 {
107            for lane in 0..4 {
108                let src = (g * 4 + lane) * cols + c * 16;
109                dst[c * 64 + lane * 16..c * 64 + lane * 16 + 16]
110                    .copy_from_slice(&bytes[src..src + 16]);
111            }
112        }
113    }
114    rep
115}
116
117/// Prefix-sum of vbit row payload offsets (absolute within the tensor
118/// bytes). `offsets[r]..offsets[r+1]` is row r's packed data.
119fn vbit_row_offsets(bytes: &[u8], rows: usize, cols: usize) -> Vec<usize> {
120    let ng = cols / GROUP_SIZE;
121    let bits = &bytes[..rows];
122    let mut offsets = Vec::with_capacity(rows + 1);
123    let mut off = rows + rows * ng * 2;
124    for r in 0..rows {
125        offsets.push(off);
126        off += (cols * bits[r] as usize).div_ceil(8);
127    }
128    offsets.push(off);
129    offsets
130}
131
132/// `CMF_X86_BLOCKED` / `CMF_GPU_LMHEAD` / `CMF_GPU_SPLIT`, read once. They
133/// used to be read from the environment on every large matvec and on every
134/// matmat in six places — microseconds each, but also a knob that could
135/// change under a running process, which is not a thing a kernel choice
136/// should be able to do mid-sequence.
137fn blocked_enabled() -> bool {
138    use std::sync::atomic::Ordering::Relaxed;
139    match BLOCKED_OVERRIDE.load(Relaxed) {
140        1 => false,
141        2 => true,
142        _ => {
143            static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
144            *ON.get_or_init(|| {
145                std::env::var("CMF_X86_BLOCKED")
146                    .map(|v| v != "0")
147                    .unwrap_or(true)
148            })
149        }
150    }
151}
152
153static BLOCKED_OVERRIDE: std::sync::atomic::AtomicU8 = std::sync::atomic::AtomicU8::new(0);
154
155/// Force the blocked GEMM on or off, ignoring the environment; `None`
156/// restores it. For tests that need to run BOTH paths and compare them:
157/// `blocked_enabled` caches its answer for the life of the process, which
158/// is right when the environment is the only input, but leaves a test that
159/// flips `CMF_X86_BLOCKED` between two calls comparing a path against
160/// itself — or against whatever a test running in parallel latched first.
161pub fn set_blocked_override(on: Option<bool>) {
162    let v = match on {
163        None => 0,
164        Some(false) => 1,
165        Some(true) => 2,
166    };
167    BLOCKED_OVERRIDE.store(v, std::sync::atomic::Ordering::Relaxed);
168}
169
170fn gpu_lmhead_enabled() -> bool {
171    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
172    *ON.get_or_init(|| {
173        std::env::var("CMF_GPU_LMHEAD")
174            .map(|v| v != "0")
175            .unwrap_or(true)
176    })
177}
178
179fn gpu_split_frac() -> f32 {
180    // MiMo's banked verification projects every row on the device. Plain
181    // decode must not quantize half of its O/head activation on the CPU.
182    if FULL_GPU_Q8.get() {
183        return 1.0;
184    }
185    static F: std::sync::OnceLock<f32> = std::sync::OnceLock::new();
186    *F.get_or_init(|| {
187        std::env::var("CMF_GPU_SPLIT")
188            .ok()
189            .and_then(|v| v.parse::<f32>().ok())
190            .unwrap_or(0.5)
191            .clamp(0.0, 1.0)
192    })
193}
194
195impl QTensor {
196    pub fn from_f32(data: Vec<f32>, rows: usize, cols: usize) -> Self {
197        debug_assert_eq!(data.len(), rows * cols);
198        Self::F32 { data, rows, cols }
199    }
200
201    /// Wrap a directory tensor without dequantizing the payload.
202    /// Falls back to dequantized f32 for dtypes without a fused kernel.
203    pub fn from_model(model: &Arc<CmfModel>, name: &str) -> Result<Self, String> {
204        // Indexed lookup: the linear directory scan made pipeline build
205        // O(N²) on MoE/skills files with thousands of tensors.
206        let idx = model
207            .tensor_index(name)
208            .ok_or_else(|| format!("tensor '{name}' not found in CMF directory"))?;
209        let entry = &model.tensors[idx];
210        if entry.shape.len() != 2 {
211            return Err(format!("QTensor::from_model needs 2-D, got '{name}'"));
212        }
213        let (rows, cols) = (entry.shape[0], entry.shape[1]);
214        let bytes = model.entry_bytes(entry);
215
216        match entry.dtype {
217            TensorDtype::Q8Row | TensorDtype::Q8_2f => {
218                let n = rows * cols;
219                let scales_off = n;
220                let row_scale: Vec<f32> = (0..rows)
221                    .map(|o| {
222                        f16_to_f32(u16::from_le_bytes([
223                            bytes[scales_off + o * 2],
224                            bytes[scales_off + o * 2 + 1],
225                        ]))
226                    })
227                    .collect();
228                let col_field: Vec<f32> = if entry.dtype == TensorDtype::Q8_2f {
229                    let col_off = n + rows * 2;
230                    (0..cols)
231                        .map(|i| {
232                            f16_to_f32(u16::from_le_bytes([
233                                bytes[col_off + i * 2],
234                                bytes[col_off + i * 2 + 1],
235                            ]))
236                        })
237                        .collect()
238                } else {
239                    Vec::new()
240                };
241                Ok(Self::Mapped {
242                    model: model.clone(),
243                    idx,
244                    dtype: entry.dtype,
245                    rows,
246                    cols,
247                    row_scale,
248                    col_field,
249                    vbit_offsets: Vec::new(),
250                    repack: q8_repack(bytes, rows, cols),
251                })
252            }
253            // vbit: fused kernel unpacks variable-bit rows from mmap.
254            TensorDtype::Vbit if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
255                model: model.clone(),
256                idx,
257                dtype: entry.dtype,
258                rows,
259                cols,
260                row_scale: Vec::new(),
261                col_field: Vec::new(),
262                vbit_offsets: vbit_row_offsets(bytes, rows, cols),
263                repack: Vec::new(),
264            }),
265            // vbit_ro (§4.2): the offset table comes straight from the
266            // file — no load-time prefix scan; kernels are shared with
267            // legacy vbit (they consume absolute offsets either way).
268            TensorDtype::VbitRo if cols % GROUP_SIZE == 0 => {
269                let (_, off_off, packed_off) = cortiq_core::quant::vbit_ro_sections(rows, cols);
270                let offsets: Vec<usize> = (0..=rows)
271                    .map(|r| packed_off + cortiq_core::quant::vbit_ro_offset(bytes, off_off, r))
272                    .collect();
273                Ok(Self::Mapped {
274                    model: model.clone(),
275                    idx,
276                    dtype: entry.dtype,
277                    rows,
278                    cols,
279                    row_scale: Vec::new(),
280                    col_field: Vec::new(),
281                    vbit_offsets: offsets,
282                    repack: Vec::new(),
283                })
284            }
285            // q4_block: fused kernel reads nibbles straight from mmap —
286            // a 14B q4 file no longer explodes into ×8 f32 RAM.
287            // q4_tiled (§4.3): interleaved [scale][nibbles] tiles — one
288            // sequential memory stream (measured ×1.66 ARM / ×1.13 AVX2
289            // at kernel level over the split layout).
290            TensorDtype::Q4Tiled if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
291                model: model.clone(),
292                idx,
293                dtype: entry.dtype,
294                rows,
295                cols,
296                row_scale: Vec::new(),
297                col_field: Vec::new(),
298                vbit_offsets: Vec::new(),
299                repack: Vec::new(),
300            }),
301            // q4tp (§4.10): nibbles from mmap, scale from the row ladder —
302            // 7.3% less file than q4t at the same 4-bit grid.
303            TensorDtype::Q4TiledP if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
304                model: model.clone(),
305                idx,
306                dtype: entry.dtype,
307                rows,
308                cols,
309                row_scale: Vec::new(),
310                col_field: Vec::new(),
311                vbit_offsets: Vec::new(),
312                repack: Vec::new(),
313            }),
314            // q2tp: 2-bit chunks from mmap, scale from the same row ladder.
315            TensorDtype::Q2TiledP if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
316                model: model.clone(),
317                idx,
318                dtype: entry.dtype,
319                rows,
320                cols,
321                row_scale: Vec::new(),
322                col_field: Vec::new(),
323                vbit_offsets: Vec::new(),
324                repack: Vec::new(),
325            }),
326            TensorDtype::Q4Block if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
327                model: model.clone(),
328                idx,
329                dtype: entry.dtype,
330                rows,
331                cols,
332                row_scale: Vec::new(),
333                col_field: Vec::new(),
334                vbit_offsets: Vec::new(),
335                repack: Vec::new(),
336            }),
337            // q1: binary sign-bit tiles from mmap (1-bit-trained models).
338            // q1g: the same signs as planes with one scale per 128-group.
339            TensorDtype::Q1 if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
340                model: model.clone(),
341                idx,
342                dtype: entry.dtype,
343                rows,
344                cols,
345                row_scale: Vec::new(),
346                col_field: Vec::new(),
347                vbit_offsets: Vec::new(),
348                repack: Vec::new(),
349            }),
350            TensorDtype::Q1G if cols % cortiq_core::quant::Q1G_GROUP == 0 => Ok(Self::Mapped {
351                model: model.clone(),
352                idx,
353                dtype: entry.dtype,
354                rows,
355                cols,
356                row_scale: Vec::new(),
357                col_field: Vec::new(),
358                vbit_offsets: Vec::new(),
359                repack: Vec::new(),
360            }),
361            // mxfp4: E2M1 codes + E8M0 scales straight from the mmap (the
362            // native FP4 expert codec; `qtensor::mxfp4`).
363            TensorDtype::Mxfp4 if cols % cortiq_core::mxfp4::GROUP == 0 => Ok(Self::Mapped {
364                model: model.clone(),
365                idx,
366                dtype: entry.dtype,
367                rows,
368                cols,
369                row_scale: Vec::new(),
370                col_field: Vec::new(),
371                vbit_offsets: Vec::new(),
372                repack: Vec::new(),
373            }),
374            // q1t (ternary + outlier overlay): fused per-row dequant kernel
375            // reads straight from mmap — a 12B q1t stays ~its file size in
376            // RAM instead of dequantizing to ~48 GB of f32.
377            TensorDtype::Q1T if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
378                model: model.clone(),
379                idx,
380                dtype: entry.dtype,
381                rows,
382                cols,
383                row_scale: Vec::new(),
384                col_field: Vec::new(),
385                vbit_offsets: Vec::new(),
386                repack: Vec::new(),
387            }),
388            // blr: factored kernels over the packed signs. CMF_BLR_DENSE=1
389            // dequantizes to f32 instead (A/B parity runs).
390            TensorDtype::BinaryLowRank if std::env::var("CMF_BLR_DENSE").as_deref() != Ok("1") => {
391                Ok(Self::Blr(Box::new(crate::blr::BlrTensor::from_model(
392                    model, idx,
393                )?)))
394            }
395            // No fused kernel yet → dequantize once (correct, more RAM).
396            _ => {
397                let mut data = vec![0.0f32; rows * cols];
398                cortiq_core::quant::dequant_tensor(entry, bytes, &mut data)?;
399                Ok(Self::from_f32(data, rows, cols))
400            }
401        }
402    }
403
404    /// Binary low-rank tensor? Fused multi-tensor kernels refuse it and fall
405    /// back to per-tensor calls.
406    pub fn is_blr(&self) -> bool {
407        matches!(self, Self::Blr(_))
408    }
409
410    /// Memory-mapped `mxfp4` tensor? Like `blr`, it has no row-range kernel
411    /// in the fused multi-tensor jobs: those run it per tensor.
412    pub fn is_mxfp4(&self) -> bool {
413        matches!(
414            self,
415            Self::Mapped {
416                dtype: TensorDtype::Mxfp4,
417                ..
418            }
419        )
420    }
421
422    /// Codecs the fused multi-tensor jobs (`matvec_many`, the SwiGLU and
423    /// MoE row kernels) do not carry: they fall back to per-tensor calls.
424    fn per_tensor_only(&self) -> bool {
425        self.is_blr() || self.is_mxfp4()
426    }
427
428    /// (model, idx) of a `blr` tensor, for the wgpu graphs only (their
429    /// `blr` stage kernels, kind [`crate::gpu::GRAPH_KIND_BLR`]). Kept apart
430    /// from [`graph_weight`](Self::graph_weight), which stays None for
431    /// `blr`: every other graph consumer keys its kernels on that answer.
432    /// A file with a Prism contract keeps the per-op path.
433    pub fn graph_blr(&self) -> Option<(&Arc<CmfModel>, usize)> {
434        match self {
435            Self::Blr(t) if !crate::prism::has_contract(&t.model) => Some((&t.model, t.idx)),
436            _ => None,
437        }
438    }
439
440    /// [`graph_weight`](Self::graph_weight), then the Prism descriptor
441    /// form, then a `blr` tensor as kind [`crate::gpu::GRAPH_KIND_BLR`] —
442    /// the wgpu token and batch graph builders' lookup.
443    pub(crate) fn graph_weight_or_blr(
444        &self,
445    ) -> Option<(&std::sync::Arc<CmfModel>, usize, u8, &[f32])> {
446        self.graph_weight()
447            .or_else(|| self.graph_weight_descriptor())
448            .or_else(|| {
449                self.graph_blr()
450                    .map(|(m, i)| (m, i, crate::gpu::GRAPH_KIND_BLR, &[][..]))
451            })
452    }
453
454    /// Binary-weight-mapped tensor, q1 or q1g? (GPU gates: the binary CPU
455    /// kernels are compute-bound, so offload pays at much smaller shapes
456    /// than q8.)
457    pub(crate) fn is_q1(&self) -> bool {
458        matches!(
459            self,
460            Self::Mapped {
461                dtype: TensorDtype::Q1 | TensorDtype::Q1G,
462                ..
463            }
464        )
465    }
466
467    /// The binary-weight view (q1 tiles or q1g planes) the shared CPU row
468    /// drivers take; None for every other storage.
469    fn bin_w(&self) -> Option<BinW<'_>> {
470        match self {
471            Self::Mapped {
472                dtype: TensorDtype::Q1,
473                ..
474            } => Some(BinW::Q1(self.quant_bytes())),
475            Self::Mapped {
476                dtype: TensorDtype::Q1G,
477                rows,
478                cols,
479                ..
480            } => Some(BinW::q1g(self.quant_bytes(), *rows, *cols)),
481            _ => None,
482        }
483    }
484
485    /// Owned-f32 view (data, rows, cols) — the GDN a/b gate projections
486    /// arrive dequantized (force-f16 in the converter → F32 in RAM).
487    pub(crate) fn f32_parts(&self) -> Option<(&[f32], usize, usize)> {
488        match self {
489            Self::F32 { data, rows, cols } => Some((data, *rows, *cols)),
490            _ => None,
491        }
492    }
493
494    /// (directory idx, rows, cols) of a q1-mapped tensor — the
495    /// whole-block GPU path resolves offsets itself.
496    /// (idx, rows, cols) of a mapped tensor the whole-token GPU graph can drive
497    /// — Q1, Q1T or Q4-block (it resolves the offset and picks the kernel by
498    /// dtype). Q4-block lets a precise down_proj/lm_head stay on-device.
499    /// Named `q1_parts` for historical reasons.
500    pub(crate) fn q1_parts(&self) -> Option<(usize, usize, usize)> {
501        if self.has_prism_contract() {
502            return None;
503        }
504        match self {
505            #[cfg(target_os = "macos")]
506            Self::Mapped {
507                dtype: TensorDtype::Q1T,
508                ..
509            } if !crate::gpu::metal_q1t_enabled() => None,
510            Self::Mapped {
511                idx,
512                dtype:
513                    TensorDtype::Q1
514                    | TensorDtype::Q1G
515                    | TensorDtype::Q1T
516                    | TensorDtype::Q4Block
517                    | TensorDtype::Q4Tiled
518                    // Q2TiledP deliberately absent: the Metal graph has no
519                    // q2tp kernel, and advertising it here made the block
520                    // plan truncate mid-run at the first q2tp layer.
521                    | TensorDtype::Q4TiledP
522                    | TensorDtype::Q8Row
523                    | TensorDtype::Q8_2f,
524                rows,
525                cols,
526                ..
527            } => Some((*idx, *rows, *cols)),
528            _ => None,
529        }
530    }
531
532    /// `(directory idx, rows, cols)` for the native Metal token graph.  The
533    /// historical q1 graph gate intentionally refuses every Prism tensor so
534    /// an untransformed q2 payload cannot slip into the resident path.  The
535    /// Metal2 graph is descriptor-aware and admits only the production
536    /// q2tp-affine forward targets; ordinary q1/q4 callers retain the old
537    /// `q1_parts` behaviour.
538    #[cfg(target_os = "macos")]
539    pub(crate) fn metal_graph_parts(&self) -> Option<(usize, usize, usize)> {
540        if let Some((model, idx, kind, _)) = self.graph_weight_descriptor() {
541            let name = &model.tensors[idx].name;
542            let forward = kind == 9 && crate::prism::is_forward_weight(model, name);
543            let affine = kind == 9 && crate::prism::is_affine_target(model, name);
544            if forward && affine {
545                let e = model.tensors.get(idx)?;
546                return Some((idx, *e.shape.first()?, *e.shape.get(1)?));
547            }
548        }
549        self.q1_parts()
550    }
551
552    /// (directory idx, rows, cols) of a q4_tiled mapped tensor. The
553    /// chunk-prefill graph takes it in the same 4-tuple slot as
554    /// `q8_row_parts` with an EMPTY row_scale — q4t carries its scales
555    /// inside the 18-byte tiles, and the empty slice is what tells the
556    /// encoder to reach for the q4t kernels.
557    pub(crate) fn q4t_parts(&self) -> Option<(usize, usize, usize)> {
558        if self.has_prism_contract() {
559            return None;
560        }
561        match self {
562            Self::Mapped {
563                idx,
564                dtype: TensorDtype::Q4Tiled,
565                rows,
566                cols,
567                ..
568            } => Some((*idx, *rows, *cols)),
569            _ => None,
570        }
571    }
572
573    /// (directory idx, rows, cols) of a q4tp mapped tensor. Same empty-scale
574    /// slot as `q4t_parts` in the chunk graph — the encoder tells the two
575    /// apart by the tensor's dtype, not by the slot.
576    pub(crate) fn q4tp_parts(&self) -> Option<(usize, usize, usize)> {
577        if self.has_prism_contract() {
578            return None;
579        }
580        match self {
581            Self::Mapped {
582                idx,
583                dtype: TensorDtype::Q4TiledP,
584                rows,
585                cols,
586                ..
587            } => Some((*idx, *rows, *cols)),
588            _ => None,
589        }
590    }
591
592    /// (directory idx, rows, cols, row_scale) of a plain q8_row mapped
593    /// tensor — the chunk-prefill GPU graph resolves offsets itself.
594    /// q8_2f is excluded on purpose: its column field would need a
595    /// prescale stage on the device.
596    pub(crate) fn q8_row_parts(&self) -> Option<(usize, usize, usize, &[f32])> {
597        if self.has_prism_contract() {
598            return None;
599        }
600        match self {
601            Self::Mapped {
602                idx,
603                dtype: TensorDtype::Q8Row,
604                rows,
605                cols,
606                row_scale,
607                col_field,
608                ..
609            } if col_field.is_empty() => Some((*idx, *rows, *cols, row_scale)),
610            _ => None,
611        }
612    }
613
614    /// (directory idx, rows, cols, row_scale, col_field) of a q8_2f mapped
615    /// tensor — the chunk-prefill graph stages x·col on the device first.
616    pub(crate) fn q8_2f_parts(&self) -> Option<(usize, usize, usize, &[f32], &[f32])> {
617        if self.has_prism_contract() {
618            return None;
619        }
620        match self {
621            Self::Mapped {
622                idx,
623                dtype: TensorDtype::Q8_2f,
624                rows,
625                cols,
626                row_scale,
627                col_field,
628                ..
629            } if col_field.len() == *cols => Some((*idx, *rows, *cols, row_scale, col_field)),
630            _ => None,
631        }
632    }
633
634    /// The layout this tensor is stored in, when it is mapped from a model.
635    /// The frames branch on it — a q2tp gate against a q4tp down is a real
636    /// combination in the 2-bit profile and needs a different kernel.
637    pub fn model_dtype(&self) -> Option<cortiq_core::TensorDtype> {
638        match self {
639            Self::Mapped { dtype, .. } => Some(*dtype),
640            _ => None,
641        }
642    }
643
644    /// The tensor's index in the model directory, when it is mapped from one.
645    /// The GPU frames bind by index rather than by name — a name lookup per
646    /// layer per token is not free, and the index is what the device cache is
647    /// keyed on anyway.
648    pub fn model_idx(&self) -> Option<usize> {
649        match self {
650            Self::Mapped { idx, .. } => Some(*idx),
651            _ => None,
652        }
653    }
654
655    /// The model this tensor is mapped from, when it is mapped at all. The
656    /// GPU frames need the container to reach the bytes; a QTensor already
657    /// holds it, and threading a second handle down every call site to say
658    /// the same thing invites the two to disagree.
659    pub fn model_arc(&self) -> Option<std::sync::Arc<cortiq_core::CmfModel>> {
660        match self {
661            Self::Mapped { model, .. } => Some(model.clone()),
662            _ => None,
663        }
664    }
665
666    /// Whether this mapped tensor belongs to the Prism/Bonsai transform
667    /// contract.  Device graphs do not carry the descriptor, so callers use
668    /// this conservative predicate to stay on the descriptor-aware CPU path
669    /// instead of silently executing an unrotated matrix.
670    pub(crate) fn has_prism_contract(&self) -> bool {
671        matches!(self, Self::Mapped { model, .. } if crate::prism::has_contract(model))
672    }
673
674    pub fn rows(&self) -> usize {
675        match self {
676            Self::F32 { rows, .. } | Self::Mapped { rows, .. } => *rows,
677            Self::Blr(t) => t.rows,
678        }
679    }
680
681    /// Mapped q4t handle (model + directory index) — the fused GPU FFN
682    /// needs the raw file coordinates of its three projections.
683    pub(crate) fn mapped_q4t(&self) -> Option<(&Arc<CmfModel>, usize)> {
684        if self.has_prism_contract() {
685            return None;
686        }
687        match self {
688            Self::Mapped {
689                model,
690                idx,
691                dtype: TensorDtype::Q4Tiled,
692                ..
693            } => Some((model, *idx)),
694            _ => None,
695        }
696    }
697
698    /// Same slot as `mapped_q4t` for a q4tp tensor — the fused DiT FFN picks
699    /// its kernels by which of the two answers.
700    pub fn mapped_q4tp(&self) -> Option<(&Arc<CmfModel>, usize)> {
701        if self.has_prism_contract() {
702            return None;
703        }
704        match self {
705            Self::Mapped {
706                model,
707                idx,
708                dtype: TensorDtype::Q4TiledP,
709                ..
710            } => Some((model, *idx)),
711            _ => None,
712        }
713    }
714
715    /// (model, tensor idx) for a mapped weight in ANY codec the fused device
716    /// paths can run — four-bit tiled or either int8 layout.
717    ///
718    /// The fused DiT chains asked for `mapped_q4tp` by name, so an eight-bit
719    /// container never reached them and rendered through per-op GEMMs even
720    /// after those kernels learned its codec. The gate is what the codec has
721    /// a device GEMM for, not which codec it is.
722    pub fn mapped_device_gemm(&self) -> Option<(&Arc<CmfModel>, usize)> {
723        if self.has_prism_contract() {
724            return None;
725        }
726        match self {
727            Self::Mapped {
728                model,
729                idx,
730                dtype: TensorDtype::Q4TiledP | TensorDtype::Q8Row | TensorDtype::Q8_2f,
731                ..
732            } => Some((model, *idx)),
733            _ => None,
734        }
735    }
736
737    /// (model, tensor idx) for a q8_2f mapped weight — the eight-bit slot of
738    /// the MoE graph ladder (int8 body, f16 row scales, f16 column field).
739    pub fn mapped_q8_2f(&self) -> Option<(&Arc<CmfModel>, usize)> {
740        if self.has_prism_contract() {
741            return None;
742        }
743        match self {
744            Self::Mapped {
745                model,
746                idx,
747                dtype: TensorDtype::Q8_2f,
748                ..
749            } => Some((model, *idx)),
750            _ => None,
751        }
752    }
753
754    /// (model, tensor idx) for a q2tp mapped weight — the 2-bit twin of
755    /// `mapped_q4tp`, used by the mixed MoE profile.
756    pub fn mapped_q2tp(&self) -> Option<(&Arc<CmfModel>, usize)> {
757        if self.has_prism_contract() {
758            return None;
759        }
760        match self {
761            Self::Mapped {
762                model,
763                idx,
764                dtype: TensorDtype::Q2TiledP,
765                ..
766            } => Some((model, *idx)),
767            _ => None,
768        }
769    }
770
771    pub fn cols(&self) -> usize {
772        match self {
773            Self::F32 { cols, .. } | Self::Mapped { cols, .. } => *cols,
774            Self::Blr(t) => t.cols,
775        }
776    }
777
778    /// (model, tensor idx) for a q1 mapped weight — the wgpu token graph
779    /// keys its resident VRAM cache by idx. None for any other dtype/kind.
780    pub fn mapped_q1(&self) -> Option<(&std::sync::Arc<CmfModel>, usize)> {
781        if self.has_prism_contract() {
782            return None;
783        }
784        match self {
785            Self::Mapped {
786                model,
787                idx,
788                dtype: TensorDtype::Q1,
789                ..
790            } => Some((model, *idx)),
791            _ => None,
792        }
793    }
794
795    /// (model, idx, kind, row_scale) for a graph-capable mapped weight.
796    /// kind: 0=q8_row (per-row scales), 1=q1, 2=q4_block, 3=q1t
797    /// (tile-embedded, no rs), 5=q4_tiled, 6=q4tp, 7=q8_2f (both scale
798    /// planes live inside the tensor), 9=q2tp, 10=q1g (sign + scale
799    /// planes). None only for `vbit`.
800    ///
801    /// The old comment here claimed q4_block was unhandled while the arm
802    /// right below mapped it, and it named q8_2f as unhandled after that
803    /// stopped being true — a stale comment on this function is how a
804    /// model silently loses the graph, so it is worth keeping honest.
805    pub fn graph_weight(&self) -> Option<(&std::sync::Arc<CmfModel>, usize, u8, &[f32])> {
806        if self.has_prism_contract() {
807            return None;
808        }
809        self.graph_weight_descriptor()
810    }
811
812    /// Descriptor-aware graph handle used only by the Prism token graph.
813    /// Ordinary graph callers continue to use [`graph_weight`] and therefore
814    /// remain fail-closed until they provide the same explicit transform
815    /// contract.
816    pub(crate) fn graph_weight_descriptor(
817        &self,
818    ) -> Option<(&std::sync::Arc<CmfModel>, usize, u8, &[f32])> {
819        match self {
820            Self::Mapped {
821                model,
822                idx,
823                dtype: TensorDtype::Q8Row,
824                row_scale,
825                ..
826            } => Some((model, *idx, 0, row_scale.as_slice())),
827            Self::Mapped {
828                model,
829                idx,
830                dtype: TensorDtype::Q1,
831                ..
832            } => Some((model, *idx, 1, &[])),
833            // Kind 10, never 1: q1g's sign and scale planes are a different
834            // walk from q1's 6-byte tiles (the kind-2/5 mixup taught that).
835            Self::Mapped {
836                model,
837                idx,
838                dtype: TensorDtype::Q1G,
839                ..
840            } => Some((model, *idx, 10, &[])),
841            // Q4Tiled is kind 5, NOT 2: both carried 2 historically, and
842            // the wgpu token graph fed 18B interleaved tiles to the
843            // split-layout q4b kernel — garbage output on q4t models
844            // (caught by an end-to-end answer check on real Vulkan).
845            Self::Mapped {
846                model,
847                idx,
848                dtype: TensorDtype::Q4Tiled,
849                ..
850            } => Some((model, *idx, 5, &[])),
851            // Kind 6, not 5: q4tp's nibble stride and scale planes differ,
852            // and feeding them to the q4t kernel is exactly the mistake that
853            // produced garbage when Q4Tiled shared kind 2 with Q4Block.
854            Self::Mapped {
855                model,
856                idx,
857                dtype: TensorDtype::Q4TiledP,
858                ..
859            } => Some((model, *idx, 6, &[])),
860            Self::Mapped {
861                model,
862                idx,
863                dtype: TensorDtype::Q4Block,
864                ..
865            } => Some((model, *idx, 2, &[])),
866            // q8_2f carries BOTH scale planes after the int8 body (rows
867            // f16, then cols f16), so the graph takes the whole tensor
868            // and the kernel reads them where they lie — no host-side
869            // prescale, which is what the per-op path does instead.
870            Self::Mapped {
871                model,
872                idx,
873                dtype: TensorDtype::Q8_2f,
874                ..
875            } => Some((model, *idx, 7, &[])),
876            Self::Mapped {
877                model,
878                idx,
879                dtype: TensorDtype::Q1T,
880                ..
881            } => Some((model, *idx, 3, &[])),
882            // Kind 9: the 2-bit plane on the q4tp ladder (dense FFN gate/up
883            // of the q2tp profile). Its own kernel — 8 bytes a group where
884            // q4tp has 16, and rung 0 is the exact zero.
885            Self::Mapped {
886                model,
887                idx,
888                dtype: TensorDtype::Q2TiledP,
889                ..
890            } => Some((model, *idx, 9, &[])),
891            _ => None,
892        }
893    }
894
895    /// Dense f32 view — only for owned tensors. Masked/sparse execution
896    /// paths require it; quantized weights don't support masks yet.
897    pub fn as_f32(&self) -> Option<&[f32]> {
898        match self {
899            Self::F32 { data, .. } => Some(data),
900            Self::Mapped { .. } | Self::Blr(_) => None,
901        }
902    }
903
904    fn quant_bytes(&self) -> &[u8] {
905        match self {
906            Self::Mapped { model, idx, .. } => model.entry_bytes(&model.tensors[*idx]),
907            Self::Blr(t) => t.bytes(),
908            Self::F32 { .. } => unreachable!("quant_bytes on F32"),
909        }
910    }
911
912    /// Dequantize one row into `dst` (embedding lookup).
913    pub fn row_f32(&self, r: usize, dst: &mut [f32]) {
914        let cols = self.cols();
915        debug_assert_eq!(dst.len(), cols);
916        match self {
917            Self::F32 { data, .. } => dst.copy_from_slice(&data[r * cols..(r + 1) * cols]),
918            Self::Blr(t) => t.row_f32(r, dst),
919            Self::Mapped {
920                model,
921                idx,
922                dtype,
923                row_scale,
924                col_field,
925                vbit_offsets,
926                ..
927            } => {
928                if *dtype == TensorDtype::Q4Tiled {
929                    let bytes = self.quant_bytes();
930                    let gpr = cols / GROUP_SIZE;
931                    for gi in 0..gpr {
932                        let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
933                        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
934                        for (k, &b) in tile[2..].iter().enumerate() {
935                            dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
936                            dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
937                        }
938                    }
939                    if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
940                        crate::prism::inverse_embedding(model, dst);
941                    }
942                    return;
943                }
944                if *dtype == TensorDtype::Q4TiledP {
945                    let bytes = self.quant_bytes();
946                    let gpr = cols / GROUP_SIZE;
947                    let v = Q4tpView::new(bytes, self.rows(), cols);
948                    let mut sc = vec![0f32; gpr];
949                    v.scales_into(r, gpr, &mut sc);
950                    for gi in 0..gpr {
951                        let tile = &v.nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
952                        let s = sc[gi];
953                        for (k, &b) in tile.iter().enumerate() {
954                            dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
955                            dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
956                        }
957                    }
958                    if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
959                        crate::prism::inverse_embedding(model, dst);
960                    }
961                    return;
962                }
963                if *dtype == TensorDtype::Q2TiledP {
964                    let bytes = self.quant_bytes();
965                    let gpr = cols / GROUP_SIZE;
966                    let v = Q4tpView::new_q2(bytes, self.rows(), cols);
967                    let mut sc = vec![0f32; gpr];
968                    v.scales_into(r, gpr, &mut sc);
969                    for gi in 0..gpr {
970                        let ch =
971                            &v.nib[(r * gpr + gi) * Q2TP_CHUNK..(r * gpr + gi + 1) * Q2TP_CHUNK];
972                        let s = sc[gi];
973                        for (k, &b) in ch.iter().enumerate() {
974                            for j in 0..4 {
975                                let center = if crate::prism::is_affine_target(
976                                    model,
977                                    &model.tensors[*idx].name,
978                                ) {
979                                    1.0
980                                } else {
981                                    1.5
982                                };
983                                dst[gi * GROUP_SIZE + k * 4 + j] =
984                                    (((b >> (2 * j)) & 3) as f32 - center) * s;
985                            }
986                        }
987                    }
988                    if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
989                        crate::prism::inverse_embedding(model, dst);
990                    }
991                    return;
992                }
993                if *dtype == TensorDtype::Q4Block {
994                    let (packed, scales) = q4_split(self.quant_bytes(), self.rows(), cols);
995                    let gpr = cols / GROUP_SIZE;
996                    for gi in 0..gpr {
997                        let g = r * gpr + gi;
998                        let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
999                        for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
1000                            dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
1001                            dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
1002                        }
1003                    }
1004                    if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
1005                        crate::prism::inverse_embedding(model, dst);
1006                    }
1007                    return;
1008                }
1009                if *dtype == TensorDtype::Mxfp4 {
1010                    mxfp4::Mx::new(self.quant_bytes(), self.rows(), cols).row_f32(r, dst);
1011                    if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
1012                        crate::prism::inverse_embedding(model, dst);
1013                    }
1014                    return;
1015                }
1016                if *dtype == TensorDtype::Q1 {
1017                    let bytes = self.quant_bytes();
1018                    let gpr = cols / GROUP_SIZE;
1019                    for gi in 0..gpr {
1020                        let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
1021                        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
1022                        for (j, &b) in tile[2..].iter().enumerate() {
1023                            for k in 0..8 {
1024                                dst[gi * GROUP_SIZE + j * 8 + k] =
1025                                    (((b >> k) & 1) as f32 * 2.0 - 1.0) * s;
1026                            }
1027                        }
1028                    }
1029                    if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
1030                        crate::prism::inverse_embedding(model, dst);
1031                    }
1032                    return;
1033                }
1034                if *dtype == TensorDtype::Q1G {
1035                    // The tied Bonsai embedding is q1g too: one row = cols/8
1036                    // sign bytes + cols/128 scales, same ±s values as q1.
1037                    let bytes = self.quant_bytes();
1038                    let sc_off = cortiq_core::quant::q1g_scale_offset(self.rows(), cols);
1039                    let gpr = cols / GROUP_SIZE;
1040                    for gi in 0..gpr {
1041                        let t = r * gpr + gi;
1042                        let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, t));
1043                        for (j, &b) in bytes[t * 4..t * 4 + 4].iter().enumerate() {
1044                            for k in 0..8 {
1045                                dst[gi * GROUP_SIZE + j * 8 + k] =
1046                                    (((b >> k) & 1) as f32 * 2.0 - 1.0) * s;
1047                            }
1048                        }
1049                    }
1050                    if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
1051                        crate::prism::inverse_embedding(model, dst);
1052                    }
1053                    return;
1054                }
1055                if *dtype == TensorDtype::Q1T {
1056                    let bytes = self.quant_bytes();
1057                    let gpr = cols / GROUP_SIZE;
1058                    let base_len = self.rows() * gpr * cortiq_core::quant::Q1T_TILE;
1059                    for gi in 0..gpr {
1060                        let off = (r * gpr + gi) * cortiq_core::quant::Q1T_TILE;
1061                        let s = cortiq_core::quant::f16_to_f32(u16::from_le_bytes([
1062                            bytes[off],
1063                            bytes[off + 1],
1064                        ]));
1065                        let codes = &bytes[off + 2..off + cortiq_core::quant::Q1T_TILE];
1066                        for k in 0..GROUP_SIZE {
1067                            dst[gi * GROUP_SIZE + k] = match cortiq_core::quant::q1t_code(codes, k)
1068                            {
1069                                1 => s,
1070                                2 => -s,
1071                                _ => 0.0,
1072                            };
1073                        }
1074                    }
1075                    // Overlay
1076                    let rows = self.rows();
1077                    let entries = base_len + (rows + 1) * 4;
1078                    if entries <= bytes.len() {
1079                        let ptrs = &bytes[base_len..base_len + (rows + 1) * 4];
1080                        let r0 = u32::from_le_bytes([
1081                            ptrs[r * 4],
1082                            ptrs[r * 4 + 1],
1083                            ptrs[r * 4 + 2],
1084                            ptrs[r * 4 + 3],
1085                        ]) as usize;
1086                        let r1 = u32::from_le_bytes([
1087                            ptrs[(r + 1) * 4],
1088                            ptrs[(r + 1) * 4 + 1],
1089                            ptrs[(r + 1) * 4 + 2],
1090                            ptrs[(r + 1) * 4 + 3],
1091                        ]) as usize;
1092                        let off = entries + r0 * 4;
1093                        for i in 0..r1 - r0 {
1094                            let item = &bytes[off + i * 4..off + i * 4 + 4];
1095                            let c = u16::from_le_bytes([item[0], item[1]]) as usize;
1096                            let v = cortiq_core::quant::f16_to_f32(u16::from_le_bytes([
1097                                item[2], item[3],
1098                            ]));
1099                            if c < cols {
1100                                dst[c] = v;
1101                            }
1102                        }
1103                    }
1104                    if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
1105                        crate::prism::inverse_embedding(model, dst);
1106                    }
1107                    return;
1108                }
1109                if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1110                    let bytes = self.quant_bytes();
1111                    let rows = self.rows();
1112                    let ng = cols / GROUP_SIZE;
1113                    let bits = &bytes[..rows];
1114                    let sc_off = rows;
1115                    // Precomputed at load — embedding lookup used to scan
1116                    // the bit-widths of every preceding row (O(token_id)).
1117                    let off = vbit_offsets[r];
1118                    let b = bits[r] as usize;
1119                    let l = ((1usize << (b - 1)) - 1) as f32;
1120                    let data = &bytes[off..];
1121                    let (mut acc, mut nbits, mut byte_idx) = (0u64, 0usize, 0usize);
1122                    for (i, d) in dst.iter_mut().enumerate() {
1123                        while nbits < b {
1124                            acc = (acc << 8) | data[byte_idx] as u64;
1125                            byte_idx += 1;
1126                            nbits += 8;
1127                        }
1128                        let u = ((acc >> (nbits - b)) & ((1u64 << b) - 1)) as f32;
1129                        nbits -= b;
1130                        let so = (r * ng + i / GROUP_SIZE) * 2;
1131                        let sv = f16_to_f32(u16::from_le_bytes([
1132                            bytes[sc_off + so],
1133                            bytes[sc_off + so + 1],
1134                        ]));
1135                        *d = (u - l) * sv;
1136                    }
1137                    if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
1138                        crate::prism::inverse_embedding(model, dst);
1139                    }
1140                    return;
1141                }
1142                let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
1143                let s = row_scale[r];
1144                match dtype {
1145                    TensorDtype::Q8Row => {
1146                        for (d, &b) in dst.iter_mut().zip(q) {
1147                            *d = (b as i8) as f32 * s;
1148                        }
1149                    }
1150                    TensorDtype::Q8_2f => {
1151                        for (i, (d, &b)) in dst.iter_mut().zip(q).enumerate() {
1152                            *d = (b as i8) as f32 * s * col_field[i];
1153                        }
1154                    }
1155                    _ => unreachable!(),
1156                }
1157                if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
1158                    crate::prism::inverse_embedding(model, dst);
1159                }
1160            }
1161        }
1162    }
1163
1164    /// Can this tensor's columns be read cheaply (for sparse down_proj)?
1165    /// True for F32/Q8Row/Q8_2f (per-row scale, direct strided access);
1166    /// false for group-packed q4/vbit (column access would unpack whole
1167    /// groups — sparse execution falls back to f32 for those).
1168    pub fn sparse_col_ok(&self) -> bool {
1169        match self {
1170            Self::F32 { .. } => true,
1171            Self::Blr(_) => false,
1172            Self::Mapped { dtype, .. } => {
1173                matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
1174            }
1175        }
1176    }
1177
1178    /// down_proj [hidden, inter]: accumulate `w · col(c)` into `out`
1179    /// [hidden] — reads ONLY column `c` (one neuron) from the mmap,
1180    /// no full-matrix dequant. `out[k] += w · down[k, c]`.
1181    pub fn add_col_scaled(&self, c: usize, w: f32, out: &mut [f32]) {
1182        let inter = self.cols();
1183        let hidden = self.rows();
1184        debug_assert_eq!(out.len(), hidden);
1185        match self {
1186            Self::F32 { data, .. } => {
1187                for (k, o) in out.iter_mut().enumerate() {
1188                    *o += w * data[k * inter + c];
1189                }
1190            }
1191            Self::Blr(t) => t.add_col_scaled(c, w, out),
1192            Self::Mapped {
1193                dtype,
1194                row_scale,
1195                col_field,
1196                ..
1197            } => {
1198                let q = self.quant_bytes();
1199                let colf = if *dtype == TensorDtype::Q8_2f {
1200                    col_field[c]
1201                } else {
1202                    1.0
1203                };
1204                let wc = w * colf;
1205                for (k, o) in out.iter_mut().enumerate() {
1206                    let b = q[k * inter + c] as i8 as f32;
1207                    *o += wc * b * row_scale[k];
1208                }
1209            }
1210        }
1211    }
1212
1213    /// Touch the head of row `r` so the DRAM latency of the next
1214    /// neuron's weights overlaps the current one's arithmetic.
1215    ///
1216    /// Scattered rows are what per-token sparsity reads, and a 2 KB
1217    /// stride is past what the hardware prefetcher follows: without this
1218    /// every row starts with a cold miss that nothing hides. One touch
1219    /// per 512 bytes is enough — the rest of the row is a sequential run
1220    /// the prefetcher does pick up.
1221    #[inline]
1222    pub fn prefetch_row(&self, r: usize) {
1223        let Self::Mapped { dtype, .. } = self else {
1224            return;
1225        };
1226        if !matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f) {
1227            return;
1228        }
1229        let cols = self.cols();
1230        let q = self.quant_bytes();
1231        let (a, b) = (r * cols, (r + 1) * cols);
1232        if b > q.len() {
1233            return;
1234        }
1235        let mut j = a;
1236        while j < b {
1237            unsafe { std::ptr::read_volatile(q.as_ptr().add(j)) };
1238            j += 512;
1239        }
1240    }
1241
1242    /// `out += w · row(r)` — the transposed twin of `add_col_scaled`.
1243    ///
1244    /// A neuron's `down` weights are a COLUMN of `[hidden, inter]`, and a
1245    /// column is strided: reading one costs a cache line per element, so
1246    /// per-neuron dynamic sparsity saves arithmetic and no bytes. Stored
1247    /// transposed (`down_proj.t.weight`, `[inter, hidden]`) the same
1248    /// weights are a contiguous ROW, and this accumulate reads exactly
1249    /// the neurons the token asked for.
1250    pub fn add_row_scaled(&self, r: usize, w: f32, out: &mut [f32], scratch: &mut [f32]) {
1251        let cols = self.cols();
1252        debug_assert_eq!(out.len(), cols);
1253        match self {
1254            Self::F32 { data, .. } => {
1255                let row = &data[r * cols..(r + 1) * cols];
1256                for (o, v) in out.iter_mut().zip(row) {
1257                    *o += w * v;
1258                }
1259            }
1260            Self::Blr(t) => {
1261                let tmp = &mut scratch[..cols];
1262                t.row_f32(r, tmp);
1263                for (o, v) in out.iter_mut().zip(tmp.iter()) {
1264                    *o += w * v;
1265                }
1266            }
1267            Self::Mapped {
1268                dtype,
1269                row_scale,
1270                col_field,
1271                ..
1272            } => match dtype {
1273                TensorDtype::Q8Row => {
1274                    let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
1275                    let ws = w * row_scale[r];
1276                    let row: &[i8] =
1277                        unsafe { std::slice::from_raw_parts(q.as_ptr() as *const i8, q.len()) };
1278                    axpy_i8_f32(out, row, ws);
1279                }
1280                TensorDtype::Q8_2f => {
1281                    let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
1282                    let ws = w * row_scale[r];
1283                    for ((o, b), c) in out.iter_mut().zip(q).zip(col_field) {
1284                        *o += ws * c * (*b as i8 as f32);
1285                    }
1286                }
1287                _ => {
1288                    self.row_f32(r, scratch);
1289                    for (o, v) in out.iter_mut().zip(scratch.iter()) {
1290                        *o += w * v;
1291                    }
1292                }
1293            },
1294        }
1295    }
1296
1297    /// Dot of row `r` with `x` (gate/up active-neuron path). Reads only
1298    /// row `r` from the mmap — no full dequant. q4/vbit dequant the row
1299    /// into `scratch` first (rare for active-FFN weights).
1300    pub fn row_dot(&self, r: usize, x: &[f32], scratch: &mut [f32]) -> f32 {
1301        let cols = self.cols();
1302        match self {
1303            Self::F32 { data, .. } => {
1304                let row = &data[r * cols..(r + 1) * cols];
1305                row.iter().zip(x).map(|(w, v)| w * v).sum()
1306            }
1307            Self::Blr(t) => {
1308                let tmp = &mut scratch[..cols];
1309                t.row_f32(r, tmp);
1310                tmp.iter().zip(x).map(|(w, v)| w * v).sum()
1311            }
1312            Self::Mapped {
1313                model,
1314                idx,
1315                dtype,
1316                row_scale,
1317                col_field,
1318                ..
1319            } => {
1320                let prism_forward =
1321                    crate::prism::is_forward_weight(model, &model.tensors[*idx].name);
1322                if prism_forward {
1323                    let transformed = crate::prism::forward(model, &x[..cols]);
1324                    let gpr = cols / GROUP_SIZE;
1325                    match dtype {
1326                        TensorDtype::Q2TiledP => {
1327                            let v = Q4tpView::new_q2(self.quant_bytes(), self.rows(), cols);
1328                            let mut sc = vec![0f32; gpr];
1329                            v.scales_into(r, gpr, &mut sc);
1330                            if crate::prism::is_affine_target(model, &model.tensors[*idx].name) {
1331                                return q2tp_affine_row_exact(v.nib, r, gpr, &transformed, &sc);
1332                            }
1333                            return q2tp_row_exact(v.nib, r, gpr, &transformed, &sc);
1334                        }
1335                        _ => {
1336                            self.row_f32(r, scratch);
1337                            return scratch.iter().zip(&transformed).map(|(w, v)| w * v).sum();
1338                        }
1339                    }
1340                }
1341                match dtype {
1342                    TensorDtype::Q8Row => {
1343                        let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
1344                        dot_i8_f32(q, x) * row_scale[r]
1345                    }
1346                    TensorDtype::Q8_2f => {
1347                        let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
1348                        dot_i8_col_f32(q, x, col_field) * row_scale[r]
1349                    }
1350                    _ => {
1351                        self.row_f32(r, scratch);
1352                        scratch.iter().zip(x).map(|(w, v)| w * v).sum()
1353                    }
1354                }
1355            }
1356        }
1357    }
1358
1359    /// `out = W · x` (row-major). F32 delegates to the historical
1360    /// bit-exact path; Mapped runs the fused int8 kernel.
1361    pub fn matvec(&self, x: &[f32], out: &mut [f32], pool: Option<&Pool>) {
1362        if self.flint_exact_cpu_default() {
1363            return float_activations_scope(|| self.matvec(x, out, pool));
1364        }
1365        match self {
1366            // NOTE: `out.len()` DRIVES this arm — it computes that many rows,
1367            // and `x.len()` is the stride. A short `out` is legitimate here,
1368            // which is why the check below lives in the Mapped arm only.
1369            Self::F32 { data, .. } => {
1370                if !crate::f32_backend::matvec(data, x, out) {
1371                    matvec_rows(pool, data, x, out);
1372                }
1373            }
1374            Self::Blr(t) => {
1375                if crate::gpu::enabled_here()
1376                    && crate::blr::gpu_wanted()
1377                    && crate::gpu::blr_matmat(t, &x[..t.cols], 1, &mut out[..t.rows])
1378                {
1379                    return;
1380                }
1381                t.matvec(x, out, pool)
1382            }
1383            Self::Mapped {
1384                model,
1385                idx,
1386                dtype,
1387                rows,
1388                cols,
1389                row_scale,
1390                col_field,
1391                vbit_offsets,
1392                repack,
1393            } => {
1394                let _ = (model, idx);
1395                // Every kernel below writes `rows` entries through a raw
1396                // pointer, so a short `out` is an out-of-bounds WRITE, not a
1397                // wrong answer: it scribbles on the allocator's metadata and
1398                // the process aborts much later, somewhere innocent
1399                // (`double free or corruption`, `corrupted double-linked
1400                // list`). The debug_assert two of the kernels carried is
1401                // compiled out of the release — exactly the build where it
1402                // matters. Fail here instead, while the caller is still on
1403                // the stack to be named.
1404                assert!(
1405                    out.len() >= *rows && x.len() >= *cols,
1406                    "matvec {rows}x{cols}: out {} (need {rows}), x {} (need {cols})",
1407                    out.len(),
1408                    x.len(),
1409                );
1410                let prism_forward =
1411                    crate::prism::is_forward_weight(model, &model.tensors[*idx].name);
1412                if *dtype == TensorDtype::Q2TiledP
1413                    && std::env::var("CMF_Q2TP_TRACE").as_deref() == Ok("1")
1414                {
1415                    use std::sync::atomic::{AtomicUsize, Ordering};
1416                    static N: AtomicUsize = AtomicUsize::new(0);
1417                    let n = N.fetch_add(1, Ordering::Relaxed);
1418                    if n < 128 {
1419                        eprintln!(
1420                            "q2tp-dispatch #{n} name={} prism={} rows={} cols={} gpu={} optin={} layer={}",
1421                            model.tensors[*idx].name,
1422                            prism_forward,
1423                            rows,
1424                            cols,
1425                            crate::gpu::enabled_here(),
1426                            crate::gpu::q2tp_gpu_opt_in(),
1427                            crate::gpu::cur_layer(),
1428                        );
1429                    }
1430                }
1431                // Prism stores every manifest-listed forward matrix in the
1432                // signed-Hadamard basis.  The q2tp WGSL path receives that
1433                // transformed vector and an explicit affine bit; codecs
1434                // without a descriptor-aware kernel remain on CPU below.
1435                if prism_forward {
1436                    let transformed = crate::prism::forward(model, &x[..*cols]);
1437                    match dtype {
1438                        TensorDtype::Q4Block => {
1439                            q4matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1440                        }
1441                        TensorDtype::Q4Tiled => {
1442                            q4t_matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1443                        }
1444                        TensorDtype::Q4TiledP => {
1445                            q4tp_matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1446                        }
1447                        TensorDtype::Q2TiledP => {
1448                            let affine =
1449                                crate::prism::is_affine_target(model, &model.tensors[*idx].name);
1450                            if *rows * *cols >= 8_388_608
1451                                && crate::gpu::enabled_here()
1452                                && crate::gpu::q2tp_gpu_opt_in()
1453                            {
1454                                let gpu_ok = if affine {
1455                                    crate::gpu::q2tp_affine_matvec(
1456                                        model,
1457                                        *idx,
1458                                        &transformed,
1459                                        *rows,
1460                                        *cols,
1461                                        out,
1462                                    )
1463                                } else {
1464                                    crate::gpu::q2tp_matvec(
1465                                        model,
1466                                        *idx,
1467                                        &transformed,
1468                                        *rows,
1469                                        *cols,
1470                                        out,
1471                                    )
1472                                };
1473                                if gpu_ok {
1474                                    return;
1475                                }
1476                            }
1477                            if affine {
1478                                q2tp_affine_matvec(
1479                                    self.quant_bytes(),
1480                                    &transformed,
1481                                    *rows,
1482                                    *cols,
1483                                    out,
1484                                    pool,
1485                                )
1486                            } else {
1487                                q2tp_matvec(
1488                                    self.quant_bytes(),
1489                                    &transformed,
1490                                    *rows,
1491                                    *cols,
1492                                    out,
1493                                    pool,
1494                                )
1495                            }
1496                        }
1497                        TensorDtype::Q1 | TensorDtype::Q1G => bin_matvec(
1498                            self.bin_w().expect("binary dtype"),
1499                            &transformed,
1500                            *rows,
1501                            *cols,
1502                            out,
1503                            pool,
1504                        ),
1505                        TensorDtype::Q1T => {
1506                            q1t_matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1507                        }
1508                        TensorDtype::Mxfp4 => {
1509                            mxfp4::matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1510                        }
1511                        TensorDtype::Vbit | TensorDtype::VbitRo => vbitmatvec(
1512                            self.quant_bytes(),
1513                            vbit_offsets,
1514                            &transformed,
1515                            *rows,
1516                            *cols,
1517                            out,
1518                            pool,
1519                        ),
1520                        TensorDtype::Q8Row | TensorDtype::Q8_2f => qmatvec(
1521                            self.quant_bytes(),
1522                            repack,
1523                            row_scale,
1524                            &transformed,
1525                            col_field,
1526                            *dtype,
1527                            *rows,
1528                            *cols,
1529                            out,
1530                            pool,
1531                        ),
1532                        _ => unreachable!("unsupported mapped Prism dtype {dtype:?}"),
1533                    }
1534                    return;
1535                }
1536                if *dtype == TensorDtype::Q4Block {
1537                    // GPU route (wgpu q4b kernel) for large q4_block matvecs —
1538                    // gives NVIDIA/AMD/Intel q4 models a GPU path. Probe keeps
1539                    // the winner; Metal returns false → the CPU kernel below.
1540                    if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1541                        let t0 = std::time::Instant::now();
1542                        match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
1543                            crate::gpu::ProbeArm::Gpu => {
1544                                if crate::gpu::q4b_matvec(model, *idx, x, *rows, *cols, out) {
1545                                    crate::gpu::probe_record(
1546                                        crate::gpu::OpClass::Matvec,
1547                                        true,
1548                                        t0.elapsed(),
1549                                    );
1550                                    return;
1551                                }
1552                            }
1553                            crate::gpu::ProbeArm::CpuTimed => {
1554                                q4matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1555                                crate::gpu::probe_record(
1556                                    crate::gpu::OpClass::Matvec,
1557                                    false,
1558                                    t0.elapsed(),
1559                                );
1560                                return;
1561                            }
1562                            crate::gpu::ProbeArm::Cpu => {}
1563                        }
1564                    }
1565                    q4matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1566                    return;
1567                }
1568                if *dtype == TensorDtype::Q4Tiled {
1569                    // GPU route for large q4t matvecs — the lm_head class,
1570                    // same shape as the q4tp arm below. The probe keeps the
1571                    // winner; a backend without the kernel refuses and the
1572                    // CPU path stays.
1573                    if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1574                        let t0 = std::time::Instant::now();
1575                        let cls = crate::gpu::matvec_class(*rows, *cols);
1576                        match crate::gpu::probe_arm(cls) {
1577                            crate::gpu::ProbeArm::Gpu => {
1578                                if crate::gpu::q4t_matvec(model, *idx, x, *rows, *cols, out) {
1579                                    crate::gpu::probe_record(cls, true, t0.elapsed());
1580                                    return;
1581                                }
1582                            }
1583                            crate::gpu::ProbeArm::CpuTimed => {
1584                                q4t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1585                                crate::gpu::probe_record(cls, false, t0.elapsed());
1586                                return;
1587                            }
1588                            crate::gpu::ProbeArm::Cpu => {}
1589                        }
1590                    }
1591                    q4t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1592                    return;
1593                }
1594                if *dtype == TensorDtype::Q4TiledP {
1595                    // GPU route for large q4tp matvecs — the lm_head class.
1596                    // On a q4tp checkpoint the head is the biggest single
1597                    // host matvec left in the decode step, and the batched
1598                    // kernel at b=1 already exists on both backends. Probe
1599                    // keeps the winner, same as q4_block above.
1600                    if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1601                        let t0 = std::time::Instant::now();
1602                        let cls = crate::gpu::matvec_class(*rows, *cols);
1603                        match crate::gpu::probe_arm(cls) {
1604                            crate::gpu::ProbeArm::Gpu => {
1605                                if crate::gpu::q4tp_matvec(model, *idx, x, *rows, *cols, out) {
1606                                    crate::gpu::probe_record(cls, true, t0.elapsed());
1607                                    return;
1608                                }
1609                            }
1610                            crate::gpu::ProbeArm::CpuTimed => {
1611                                q4tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1612                                crate::gpu::probe_record(cls, false, t0.elapsed());
1613                                return;
1614                            }
1615                            crate::gpu::ProbeArm::Cpu => {}
1616                        }
1617                    }
1618                    q4tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1619                    return;
1620                }
1621                if *dtype == TensorDtype::Q2TiledP {
1622                    q2tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1623                    return;
1624                }
1625                if *dtype == TensorDtype::Mxfp4 {
1626                    self.mxfp4_matvec(x, out, pool);
1627                    return;
1628                }
1629                if *dtype == TensorDtype::Q1 || *dtype == TensorDtype::Q1G {
1630                    // GPU route for large q1 matvecs (out_proj / lm_head
1631                    // class): the CPU q1 kernel is load-port-bound at
1632                    // ~4 GB/s/core, the GPU one is bandwidth-bound — the
1633                    // probe measures both arms and keeps the winner.
1634                    let w = self.bin_w().expect("binary dtype");
1635                    if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1636                        let t0 = std::time::Instant::now();
1637                        let arm = if crate::gpu::q1_force() {
1638                            crate::gpu::ProbeArm::Gpu
1639                        } else {
1640                            crate::gpu::probe_arm(crate::gpu::OpClass::Matvec)
1641                        };
1642                        match arm {
1643                            crate::gpu::ProbeArm::Gpu => {
1644                                let ok = if *dtype == TensorDtype::Q1G {
1645                                    crate::gpu::q1g_matvec(model, *idx, x, *rows, *cols, out)
1646                                } else {
1647                                    crate::gpu::q1_matvec(model, *idx, x, *rows, *cols, out)
1648                                };
1649                                if ok {
1650                                    crate::gpu::probe_record(
1651                                        crate::gpu::OpClass::Matvec,
1652                                        true,
1653                                        t0.elapsed(),
1654                                    );
1655                                    return;
1656                                }
1657                            }
1658                            crate::gpu::ProbeArm::CpuTimed => {
1659                                bin_matvec(w, x, *rows, *cols, out, pool);
1660                                crate::gpu::probe_record(
1661                                    crate::gpu::OpClass::Matvec,
1662                                    false,
1663                                    t0.elapsed(),
1664                                );
1665                                return;
1666                            }
1667                            crate::gpu::ProbeArm::Cpu => {}
1668                        }
1669                    }
1670                    bin_matvec(w, x, *rows, *cols, out, pool);
1671                    return;
1672                }
1673                if *dtype == TensorDtype::Q1T {
1674                    // GPU route for large q1t matvecs: the ternary BASE dot runs
1675                    // on the GPU (load-port-bound on CPU, like q1), then the
1676                    // sparse overlay is added on the CPU. Probe keeps the winner.
1677                    if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1678                        let t0 = std::time::Instant::now();
1679                        match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
1680                            crate::gpu::ProbeArm::Gpu => {
1681                                if crate::gpu::q1t_matvec(model, *idx, x, *rows, *cols, out) {
1682                                    q1t_add_overlay(self.quant_bytes(), x, *rows, *cols, out, pool);
1683                                    crate::gpu::probe_record(
1684                                        crate::gpu::OpClass::Matvec,
1685                                        true,
1686                                        t0.elapsed(),
1687                                    );
1688                                    return;
1689                                }
1690                            }
1691                            crate::gpu::ProbeArm::CpuTimed => {
1692                                q1t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1693                                crate::gpu::probe_record(
1694                                    crate::gpu::OpClass::Matvec,
1695                                    false,
1696                                    t0.elapsed(),
1697                                );
1698                                return;
1699                            }
1700                            crate::gpu::ProbeArm::Cpu => {}
1701                        }
1702                    }
1703                    q1t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1704                    return;
1705                }
1706                if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1707                    vbitmatvec(self.quant_bytes(), vbit_offsets, x, *rows, *cols, out, pool);
1708                    return;
1709                }
1710                let xs = prescale(x, col_field, *dtype);
1711                // D5: large q8 matrices (lm_head-class) — hybrid
1712                // CPU∥GPU: split the rows, both sides compute
1713                // SIMULTANEOUSLY (same math, shared prescale).
1714                // GPU share: CMF_GPU_SPLIT (0..1, default 0.5).
1715                if *rows >= crate::gpu::min_rows()
1716                    && matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
1717                    && gpu_lmhead_enabled()
1718                    && crate::gpu::enabled_here()
1719                {
1720                    // Runtime probe: alternate the hybrid against the
1721                    // pure-CPU matvec, keep whichever is faster HERE.
1722                    let t0 = std::time::Instant::now();
1723                    match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
1724                        crate::gpu::ProbeArm::Gpu => {}
1725                        crate::gpu::ProbeArm::CpuTimed => {
1726                            qmatvec(
1727                                self.quant_bytes(),
1728                                repack,
1729                                row_scale,
1730                                x,
1731                                col_field,
1732                                *dtype,
1733                                *rows,
1734                                *cols,
1735                                out,
1736                                pool,
1737                            );
1738                            crate::gpu::probe_record(
1739                                crate::gpu::OpClass::Matvec,
1740                                false,
1741                                t0.elapsed(),
1742                            );
1743                            return;
1744                        }
1745                        crate::gpu::ProbeArm::Cpu => {
1746                            qmatvec(
1747                                self.quant_bytes(),
1748                                repack,
1749                                row_scale,
1750                                x,
1751                                col_field,
1752                                *dtype,
1753                                *rows,
1754                                *cols,
1755                                out,
1756                                pool,
1757                            );
1758                            return;
1759                        }
1760                    }
1761                    let frac = gpu_split_frac();
1762                    let cpu_rows = ((*rows as f32) * (1.0 - frac)) as usize;
1763                    let (out_cpu, out_gpu) = out.split_at_mut(cpu_rows);
1764                    let bytes = self.quant_bytes();
1765                    let ok = std::thread::scope(|sc| {
1766                        let g = sc.spawn(|| {
1767                            crate::gpu::q8_matvec_range(
1768                                model,
1769                                *idx,
1770                                cpu_rows,
1771                                &row_scale[cpu_rows..],
1772                                &xs,
1773                                *rows - cpu_rows,
1774                                *cols,
1775                                out_gpu,
1776                            )
1777                        });
1778                        if cpu_rows > 0 {
1779                            // Repack prefix covers the full groups of the
1780                            // CPU half (the split starts at row 0).
1781                            let rep_cpu = if repack.is_empty() {
1782                                &[][..]
1783                            } else {
1784                                &repack[..(cpu_rows / 4) * 4 * *cols]
1785                            };
1786                            qmatvec(
1787                                &bytes[..cpu_rows * *cols],
1788                                rep_cpu,
1789                                &row_scale[..cpu_rows],
1790                                x,
1791                                col_field,
1792                                *dtype,
1793                                cpu_rows,
1794                                *cols,
1795                                out_cpu,
1796                                pool,
1797                            );
1798                        }
1799                        g.join().unwrap_or(false)
1800                    });
1801                    if ok {
1802                        crate::gpu::probe_record(crate::gpu::OpClass::Matvec, true, t0.elapsed());
1803                        return;
1804                    }
1805                    // GPU failed — CPU finishes its half (rows rebased —
1806                    // group offsets don't line up, mmap layout only).
1807                    qmatvec(
1808                        &bytes[cpu_rows * *cols..(*rows) * *cols],
1809                        &[],
1810                        &row_scale[cpu_rows..],
1811                        x,
1812                        col_field,
1813                        *dtype,
1814                        *rows - cpu_rows,
1815                        *cols,
1816                        out_gpu,
1817                        pool,
1818                    );
1819                    return;
1820                }
1821                qmatvec(
1822                    self.quant_bytes(),
1823                    repack,
1824                    row_scale,
1825                    x,
1826                    col_field,
1827                    *dtype,
1828                    *rows,
1829                    *cols,
1830                    out,
1831                    pool,
1832                );
1833            }
1834        }
1835    }
1836
1837    /// Fused two-input matvec (MTP verify pair): weights streamed once.
1838    pub fn matvec2(
1839        &self,
1840        x1: &[f32],
1841        x2: &[f32],
1842        o1: &mut [f32],
1843        o2: &mut [f32],
1844        pool: Option<&Pool>,
1845    ) {
1846        if self.flint_exact_cpu_default() {
1847            return float_activations_scope(|| self.matvec2(x1, x2, o1, o2, pool));
1848        }
1849        match self {
1850            Self::F32 { data, .. } => matvec_rows2(pool, data, x1, x2, o1, o2),
1851            Self::Blr(_) => {
1852                self.matvec(x1, o1, pool);
1853                self.matvec(x2, o2, pool);
1854            }
1855            Self::Mapped {
1856                model,
1857                idx,
1858                dtype,
1859                rows,
1860                cols,
1861                row_scale,
1862                col_field,
1863                vbit_offsets,
1864                ..
1865            } => {
1866                if crate::prism::is_forward_weight(model, &model.tensors[*idx].name) {
1867                    let tx1 = crate::prism::forward(model, &x1[..*cols]);
1868                    let tx2 = crate::prism::forward(model, &x2[..*cols]);
1869                    match dtype {
1870                        TensorDtype::Q4Block => {
1871                            q4matvec2(self.quant_bytes(), &tx1, &tx2, *rows, *cols, o1, o2, pool)
1872                        }
1873                        TensorDtype::Q4Tiled => {
1874                            q4t_matvec2(self.quant_bytes(), &tx1, &tx2, *rows, *cols, o1, o2, pool)
1875                        }
1876                        TensorDtype::Q4TiledP => {
1877                            q4tp_matvec2(self.quant_bytes(), &tx1, &tx2, *rows, *cols, o1, o2, pool)
1878                        }
1879                        TensorDtype::Q2TiledP => {
1880                            if crate::prism::is_affine_target(model, &model.tensors[*idx].name) {
1881                                q2tp_affine_matvec2(
1882                                    self.quant_bytes(),
1883                                    &tx1,
1884                                    &tx2,
1885                                    *rows,
1886                                    *cols,
1887                                    o1,
1888                                    o2,
1889                                    pool,
1890                                )
1891                            } else {
1892                                q2tp_matvec2(
1893                                    self.quant_bytes(),
1894                                    &tx1,
1895                                    &tx2,
1896                                    *rows,
1897                                    *cols,
1898                                    o1,
1899                                    o2,
1900                                    pool,
1901                                )
1902                            }
1903                        }
1904                        TensorDtype::Q1 | TensorDtype::Q1G => bin_matvec2(
1905                            self.bin_w().expect("binary dtype"),
1906                            &tx1,
1907                            &tx2,
1908                            *rows,
1909                            *cols,
1910                            o1,
1911                            o2,
1912                            pool,
1913                        ),
1914                        TensorDtype::Q1T => {
1915                            q1t_matvec2(self.quant_bytes(), &tx1, &tx2, *rows, *cols, o1, o2, pool)
1916                        }
1917                        TensorDtype::Mxfp4 => mxfp4::matvec2(
1918                            self.quant_bytes(),
1919                            &tx1,
1920                            &tx2,
1921                            *rows,
1922                            *cols,
1923                            o1,
1924                            o2,
1925                            pool,
1926                        ),
1927                        TensorDtype::Vbit | TensorDtype::VbitRo => vbitmatvec2(
1928                            self.quant_bytes(),
1929                            vbit_offsets,
1930                            &tx1,
1931                            &tx2,
1932                            *rows,
1933                            *cols,
1934                            o1,
1935                            o2,
1936                            pool,
1937                        ),
1938                        TensorDtype::Q8Row | TensorDtype::Q8_2f => qmatvec2(
1939                            self.quant_bytes(),
1940                            row_scale,
1941                            &tx1,
1942                            &tx2,
1943                            col_field,
1944                            *dtype,
1945                            *rows,
1946                            *cols,
1947                            o1,
1948                            o2,
1949                            pool,
1950                        ),
1951                        _ => unreachable!("unsupported mapped Prism dtype {dtype:?}"),
1952                    }
1953                    return;
1954                }
1955                if *dtype == TensorDtype::Q4Block {
1956                    q4matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1957                    return;
1958                }
1959                if *dtype == TensorDtype::Q4Tiled {
1960                    q4t_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1961                    return;
1962                }
1963                if *dtype == TensorDtype::Q4TiledP {
1964                    q4tp_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1965                    return;
1966                }
1967                if *dtype == TensorDtype::Q2TiledP {
1968                    q2tp_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1969                    return;
1970                }
1971                if *dtype == TensorDtype::Mxfp4 {
1972                    mxfp4::matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1973                    return;
1974                }
1975                if let Some(w) = self.bin_w() {
1976                    bin_matvec2(w, x1, x2, *rows, *cols, o1, o2, pool);
1977                    return;
1978                }
1979                if *dtype == TensorDtype::Q1T {
1980                    // Fused ternary pair: one row pass, the register
1981                    // unpack shared across both streams on ARM. (Q1T
1982                    // lacks a row_scale array — scales live inline in
1983                    // the tiles — so it must not fall through to the
1984                    // q8 qmatvec2 below.)
1985                    q1t_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1986                    return;
1987                }
1988                if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1989                    vbitmatvec2(
1990                        self.quant_bytes(),
1991                        vbit_offsets,
1992                        x1,
1993                        x2,
1994                        *rows,
1995                        *cols,
1996                        o1,
1997                        o2,
1998                        pool,
1999                    );
2000                    return;
2001                }
2002                qmatvec2(
2003                    self.quant_bytes(),
2004                    row_scale,
2005                    x1,
2006                    x2,
2007                    col_field,
2008                    *dtype,
2009                    *rows,
2010                    *cols,
2011                    o1,
2012                    o2,
2013                    pool,
2014                );
2015            }
2016        }
2017    }
2018}
2019
2020impl QTensor {
2021    /// Batched matvec (prefill-GEMM): xs — row-major [b, cols],
2022    /// out — row-major [b, rows]. Element-wise semantics are IDENTICAL
2023    /// to b matvec calls (same dot kernels in the same order); the win —
2024    /// the weight row streams from DRAM once per batch, not b times.
2025    /// `(model, index)` when this is a memory-mapped q4tp tensor — the
2026    /// identity a device-resident chain needs to hand `tp_matmat` the
2027    /// weight without going through this struct's own dispatch.
2028    pub fn q4tp_mapped(&self) -> Option<(&std::sync::Arc<CmfModel>, usize)> {
2029        if self.has_prism_contract() {
2030            return None;
2031        }
2032        match self {
2033            Self::Mapped {
2034                model, idx, dtype, ..
2035            } if *dtype == TensorDtype::Q4TiledP => Some((model, *idx)),
2036            _ => None,
2037        }
2038    }
2039
2040    pub fn matmat(&self, xs_all: &[f32], b: usize, out: &mut [f32], pool: Option<&Pool>) {
2041        if self.flint_exact_cpu_default() {
2042            return float_activations_scope(|| self.matmat(xs_all, b, out, pool));
2043        }
2044        let cols = self.cols();
2045        let rows = self.rows();
2046        debug_assert_eq!(xs_all.len(), b * cols);
2047        debug_assert_eq!(out.len(), b * rows);
2048        let _prof = crate::cpuprof::time(crate::cpuprof::Slot::Matmat);
2049        // GPTQ calibration: fold this layer's inputs into its Hessian. Only
2050        // Mapped tensors carry a directory name; the check is a relaxed
2051        // atomic load, free when not calibrating.
2052        if crate::gptq_capture::capturing() {
2053            if let Self::Mapped { model, idx, .. } = self {
2054                crate::gptq_capture::accumulate(&model.tensors[*idx].name, xs_all, b, cols);
2055            }
2056            if let Self::Blr(t) = self {
2057                crate::gptq_capture::accumulate(t.name(), xs_all, b, cols);
2058            }
2059        }
2060        match self {
2061            Self::Blr(t) => {
2062                if crate::gpu::enabled_here()
2063                    && crate::blr::gpu_wanted()
2064                    && crate::gpu::blr_matmat(t, &xs_all[..b * cols], b, &mut out[..b * rows])
2065                {
2066                    return;
2067                }
2068                t.matmat(xs_all, b, out, pool)
2069            }
2070            Self::F32 { data, .. } => {
2071                if crate::f32_backend::matmat(data, xs_all, b, rows, cols, out) {
2072                    return;
2073                }
2074                let out_addr = SendMut(out.as_mut_ptr());
2075                // Over the batch, four positions a row at a time: four
2076                // independent chains, each still summing j = 0..cols in the
2077                // scalar order (bit-identical to the one-at-a-time loop and
2078                // to `matvec_rows`). The batch is the wide axis of a prefill;
2079                // the rows may be few (Spark-X2.5's 16-row g_proj spent
2080                // 17 ms a layer here split over its rows).
2081                let run = |start: usize, end: usize| {
2082                    let mut bi = start;
2083                    while bi < end {
2084                        let n = (end - bi).min(4);
2085                        for o in 0..rows {
2086                            let row = &data[o * cols..(o + 1) * cols];
2087                            if n == 4 {
2088                                let x = |k: usize| &xs_all[(bi + k) * cols..(bi + k + 1) * cols];
2089                                let (x0, x1, x2, x3) = (x(0), x(1), x(2), x(3));
2090                                let (mut a0, mut a1, mut a2, mut a3) = (0f32, 0f32, 0f32, 0f32);
2091                                for j in 0..cols {
2092                                    let w = row[j];
2093                                    a0 += w * x0[j];
2094                                    a1 += w * x1[j];
2095                                    a2 += w * x2[j];
2096                                    a3 += w * x3[j];
2097                                }
2098                                for (k, acc) in [a0, a1, a2, a3].into_iter().enumerate() {
2099                                    unsafe { *out_addr.at((bi + k) * rows + o) = acc };
2100                                }
2101                            } else {
2102                                for k in 0..n {
2103                                    let x = &xs_all[(bi + k) * cols..(bi + k + 1) * cols];
2104                                    let mut acc = 0f32;
2105                                    for j in 0..cols {
2106                                        acc += row[j] * x[j];
2107                                    }
2108                                    unsafe { *out_addr.at((bi + k) * rows + o) = acc };
2109                                }
2110                            }
2111                        }
2112                        bi += n;
2113                    }
2114                };
2115                dispatch_rows(pool, b, &run);
2116            }
2117            Self::Mapped {
2118                model,
2119                idx,
2120                dtype,
2121                row_scale,
2122                col_field,
2123                vbit_offsets,
2124                ..
2125            } => {
2126                if crate::prism::is_forward_weight(model, &model.tensors[*idx].name) {
2127                    let mut transformed = Vec::with_capacity(xs_all.len());
2128                    for bi in 0..b {
2129                        transformed.extend_from_slice(&crate::prism::forward(
2130                            model,
2131                            &xs_all[bi * cols..(bi + 1) * cols],
2132                        ));
2133                    }
2134                    match dtype {
2135                        TensorDtype::Q4Block => {
2136                            q4matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
2137                        }
2138                        TensorDtype::Q4Tiled => {
2139                            q4t_matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
2140                        }
2141                        TensorDtype::Q4TiledP => {
2142                            q4tp_matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
2143                        }
2144                        TensorDtype::Q2TiledP => {
2145                            let affine =
2146                                crate::prism::is_affine_target(model, &model.tensors[*idx].name);
2147                            // Affine Prism Q2TP has a descriptor-aware GPU
2148                            // kernel for short/tail batches too.  Unlike the
2149                            // ordinary Q2TP path, don't force b<32 back to a
2150                            // scalar CPU matmat: prefill chunks and the final
2151                            // tail both need to stay on the tested GPU arm.
2152                            let gpu_batch_ok = if affine {
2153                                b >= 2
2154                            } else {
2155                                b >= 32 && b * rows * cols >= 128_000_000
2156                            };
2157                            if gpu_batch_ok
2158                                && cols % 32 == 0
2159                                && crate::gpu::enabled_here()
2160                                && crate::gpu::q2tp_gpu_opt_in()
2161                            {
2162                                let gpu_ok = if affine {
2163                                    crate::gpu::q2tp_affine_matmat(
2164                                        model,
2165                                        *idx,
2166                                        &transformed,
2167                                        b,
2168                                        rows,
2169                                        cols,
2170                                        out,
2171                                    )
2172                                } else {
2173                                    crate::gpu::q2tp_matmat(
2174                                        model,
2175                                        *idx,
2176                                        &transformed,
2177                                        b,
2178                                        rows,
2179                                        cols,
2180                                        out,
2181                                    )
2182                                };
2183                                if gpu_ok {
2184                                    return;
2185                                }
2186                            }
2187                            // A one-token Prism decode is the other short
2188                            // case.  Use the descriptor-aware matvec kernel
2189                            // before falling back to the exact CPU path.
2190                            if affine
2191                                && b == 1
2192                                && cols % 32 == 0
2193                                && crate::gpu::enabled_here()
2194                                && crate::gpu::q2tp_gpu_opt_in()
2195                                && crate::gpu::q2tp_affine_matvec(
2196                                    model,
2197                                    *idx,
2198                                    &transformed[..cols],
2199                                    rows,
2200                                    cols,
2201                                    &mut out[..rows],
2202                                )
2203                            {
2204                                return;
2205                            }
2206                            if affine {
2207                                q2tp_affine_matmat(
2208                                    self.quant_bytes(),
2209                                    &transformed,
2210                                    b,
2211                                    rows,
2212                                    cols,
2213                                    out,
2214                                    pool,
2215                                )
2216                            } else {
2217                                q2tp_matmat(
2218                                    self.quant_bytes(),
2219                                    &transformed,
2220                                    b,
2221                                    rows,
2222                                    cols,
2223                                    out,
2224                                    pool,
2225                                )
2226                            }
2227                        }
2228                        TensorDtype::Q1 | TensorDtype::Q1G => bin_matmat(
2229                            self.bin_w().expect("binary dtype"),
2230                            &transformed,
2231                            b,
2232                            rows,
2233                            cols,
2234                            out,
2235                            pool,
2236                        ),
2237                        TensorDtype::Q1T => {
2238                            q1t_matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
2239                        }
2240                        TensorDtype::Mxfp4 => {
2241                            mxfp4::matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
2242                        }
2243                        TensorDtype::Vbit | TensorDtype::VbitRo => vbitmatmat(
2244                            self.quant_bytes(),
2245                            vbit_offsets,
2246                            &transformed,
2247                            b,
2248                            rows,
2249                            cols,
2250                            out,
2251                            pool,
2252                        ),
2253                        TensorDtype::Q8Row | TensorDtype::Q8_2f => {
2254                            let pre: Vec<std::borrow::Cow<'_, [f32]>> = (0..b)
2255                                .map(|bi| {
2256                                    prescale(
2257                                        &transformed[bi * cols..(bi + 1) * cols],
2258                                        col_field,
2259                                        *dtype,
2260                                    )
2261                                })
2262                                .collect();
2263                            qmatmat(self.quant_bytes(), row_scale, &pre, rows, cols, out, pool)
2264                        }
2265                        _ => unreachable!("unsupported mapped Prism dtype {dtype:?}"),
2266                    }
2267                    return;
2268                }
2269                if *dtype == TensorDtype::Q4Block {
2270                    q4matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2271                    return;
2272                }
2273                if *dtype == TensorDtype::Mxfp4 {
2274                    self.mxfp4_matmat(xs_all, b, out, pool);
2275                    return;
2276                }
2277                if *dtype == TensorDtype::Q4TiledP {
2278                    // GPU batched q4tp GEMM (dequant + f32nt mul_mm on the
2279                    // device); the probe keeps whichever beats the CPU arm.
2280                    // Narrow (prompt-encode) and wide (DiT) batches probe
2281                    // as separate classes — the regimes have opposite
2282                    // winners and one shared verdict locked the wrong arm.
2283                    // Kill switch (gpu::mm_kill): one grossly slow GPU op
2284                    // (a fair-condition op is ≤~100 ms even at 1024px)
2285                    // means the device is contended by another process
2286                    // (e.g. a simulator) — verdicts are per-process, so
2287                    // without the bail the whole render crawls behind
2288                    // someone else's queue.
2289                    // Row-exact batches stay on the host: the device GEMM
2290                    // is an f32 dequant-sgemm, not the host matvec's sum.
2291                    if b >= 32
2292                        && b * rows * cols >= 128_000_000
2293                        && cols % 32 == 0
2294                        && !row_exact()
2295                        && !crate::gpu::mm_killed()
2296                        && crate::gpu::enabled_here()
2297                    {
2298                        let class = if b >= 128 {
2299                            crate::gpu::OpClass::MatmatWide
2300                        } else {
2301                            crate::gpu::OpClass::Matmat
2302                        };
2303                        if let Self::Mapped { model, idx, .. } = self {
2304                            // In-process A/B (`CMF_MM_AB=1`). Three
2305                            // wall-clock A/Bs on a shared stand disagreed
2306                            // with each other by 25% on the same change,
2307                            // because the machine drifts between processes
2308                            // and interleaving whole renders does not fix
2309                            // that. Here both arms run back to back on the
2310                            // SAME data inside one call, so whatever the
2311                            // machine is doing, it does to both — and the
2312                            // disagreement between their outputs falls out
2313                            // for free. Doubles the work; a diagnostic,
2314                            // not a mode.
2315                            if crate::mm_ab::on() {
2316                                let mut g = vec![0f32; b * rows];
2317                                let t = std::time::Instant::now();
2318                                let took = crate::gpu::q4tp_matmat(
2319                                    model, *idx, xs_all, b, rows, cols, &mut g,
2320                                );
2321                                let dg = t.elapsed();
2322                                let t = std::time::Instant::now();
2323                                q4tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2324                                let dc = t.elapsed();
2325                                crate::mm_ab::record(b, rows, cols, took, dg, dc, &g, out);
2326                                return;
2327                            }
2328                            let t0 = std::time::Instant::now();
2329                            // A cold call takes the device arm: its sample
2330                            // is discarded either way, and the upload is
2331                            // what the next step needs.
2332                            let resident = crate::gpu::weight_is_resident(model, *idx);
2333                            match crate::gpu::probe_arm_cold_prefers_gpu(class, resident) {
2334                                crate::gpu::ProbeArm::Gpu => {
2335                                    if crate::gpu::q4tp_matmat(
2336                                        model, *idx, xs_all, b, rows, cols, out,
2337                                    ) {
2338                                        let el = t0.elapsed();
2339                                        // Work-proportional budget: ~8× the
2340                                        // fair-device estimate (+20 ms slack).
2341                                        // An absolute cap missed the worst
2342                                        // case — contended ops sit at
2343                                        // 100–240 ms each and still bury a
2344                                        // render whose fair op is 3–9 ms.
2345                                        // Cold ops (first PSO build, buffer
2346                                        // alloc) are exempt: a one-off
2347                                        // ~50 ms compile is not contention.
2348                                        let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
2349                                        let budget = std::time::Duration::from_secs_f64(
2350                                            flops / 1.5e12 * 8.0 + 0.020,
2351                                        );
2352                                        crate::gpu::mm_budget_check(
2353                                            "q4tp matmat",
2354                                            el,
2355                                            budget,
2356                                            crate::gpu::probe_was_cold() || !resident,
2357                                        );
2358                                        crate::gpu::probe_record(class, true, el);
2359                                        return;
2360                                    }
2361                                }
2362                                crate::gpu::ProbeArm::CpuTimed => {
2363                                    q4tp_matmat(
2364                                        self.quant_bytes(),
2365                                        xs_all,
2366                                        b,
2367                                        rows,
2368                                        cols,
2369                                        out,
2370                                        pool,
2371                                    );
2372                                    crate::gpu::probe_record(class, false, t0.elapsed());
2373                                    return;
2374                                }
2375                                crate::gpu::ProbeArm::Cpu => {}
2376                            }
2377                        }
2378                    }
2379                    q4tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2380                    return;
2381                }
2382                if *dtype == TensorDtype::Q2TiledP {
2383                    // Same device arm as q4tp, behind the same probe:
2384                    // the planes differ, the dispatch does not. Without
2385                    // this a q2tp file ran its widest projections on the
2386                    // host while the 4-bit one had the card, which is a
2387                    // codec paying for its size twice.
2388                    if b >= 32
2389                        && b * rows * cols >= 128_000_000
2390                        && cols % 32 == 0
2391                        && !crate::gpu::mm_killed()
2392                        && crate::gpu::enabled_here()
2393                    {
2394                        let class = if b >= 128 {
2395                            crate::gpu::OpClass::MatmatWide
2396                        } else {
2397                            crate::gpu::OpClass::Matmat
2398                        };
2399                        if let Self::Mapped { model, idx, .. } = self {
2400                            let t0 = std::time::Instant::now();
2401                            match crate::gpu::probe_arm(class) {
2402                                crate::gpu::ProbeArm::Gpu => {
2403                                    if crate::gpu::q2tp_matmat(
2404                                        model, *idx, xs_all, b, rows, cols, out,
2405                                    ) {
2406                                        crate::gpu::probe_record(class, true, t0.elapsed());
2407                                        return;
2408                                    }
2409                                }
2410                                crate::gpu::ProbeArm::CpuTimed => {
2411                                    q2tp_matmat(
2412                                        self.quant_bytes(),
2413                                        xs_all,
2414                                        b,
2415                                        rows,
2416                                        cols,
2417                                        out,
2418                                        pool,
2419                                    );
2420                                    crate::gpu::probe_record(class, false, t0.elapsed());
2421                                    return;
2422                                }
2423                                crate::gpu::ProbeArm::Cpu => {}
2424                            }
2425                        }
2426                    }
2427                    // Without a host arm a q2tp tensor falls through to
2428                    // the q8 fallback, which reads it at one BYTE per
2429                    // weight — a 2x overrun that killed pool workers
2430                    // mid-prefill while the dispatcher waited forever.
2431                    q2tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2432                    return;
2433                }
2434                if *dtype == TensorDtype::Q4Tiled {
2435                    // GPU batched q4t GEMM (dequant + f32nt mul_mm on the
2436                    // device); the probe keeps whichever beats the CPU arm.
2437                    // Narrow (prompt-encode) and wide (DiT) batches probe
2438                    // as separate classes — the regimes have opposite
2439                    // winners and one shared verdict locked the wrong arm.
2440                    // Kill switch (gpu::mm_kill): one grossly slow GPU op
2441                    // (a fair-condition op is ≤~100 ms even at 1024px)
2442                    // means the device is contended by another process
2443                    // (e.g. a simulator) — verdicts are per-process, so
2444                    // without the bail the whole render crawls behind
2445                    // someone else's queue.
2446                    if b >= 32
2447                        && b * rows * cols >= 128_000_000
2448                        && cols % 32 == 0
2449                        && !crate::gpu::mm_killed()
2450                        && crate::gpu::enabled_here()
2451                    {
2452                        let class = if b >= 128 {
2453                            crate::gpu::OpClass::MatmatWide
2454                        } else {
2455                            crate::gpu::OpClass::Matmat
2456                        };
2457                        if let Self::Mapped { model, idx, .. } = self {
2458                            let t0 = std::time::Instant::now();
2459                            match crate::gpu::probe_arm(class) {
2460                                crate::gpu::ProbeArm::Gpu => {
2461                                    if crate::gpu::q4t_matmat(
2462                                        model, *idx, xs_all, b, rows, cols, out,
2463                                    ) {
2464                                        let el = t0.elapsed();
2465                                        // Work-proportional budget: ~8× the
2466                                        // fair-device estimate (+20 ms slack).
2467                                        // An absolute cap missed the worst
2468                                        // case — contended ops sit at
2469                                        // 100–240 ms each and still bury a
2470                                        // render whose fair op is 3–9 ms.
2471                                        // Cold ops (first PSO build, buffer
2472                                        // alloc) are exempt: a one-off
2473                                        // ~50 ms compile is not contention.
2474                                        let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
2475                                        let budget = std::time::Duration::from_secs_f64(
2476                                            flops / 1.5e12 * 8.0 + 0.020,
2477                                        );
2478                                        crate::gpu::mm_budget_check(
2479                                            "q4t matmat",
2480                                            el,
2481                                            budget,
2482                                            crate::gpu::probe_was_cold(),
2483                                        );
2484                                        crate::gpu::probe_record(class, true, el);
2485                                        return;
2486                                    }
2487                                }
2488                                crate::gpu::ProbeArm::CpuTimed => {
2489                                    q4t_matmat(
2490                                        self.quant_bytes(),
2491                                        xs_all,
2492                                        b,
2493                                        rows,
2494                                        cols,
2495                                        out,
2496                                        pool,
2497                                    );
2498                                    crate::gpu::probe_record(class, false, t0.elapsed());
2499                                    return;
2500                                }
2501                                crate::gpu::ProbeArm::Cpu => {}
2502                            }
2503                        }
2504                    }
2505                    q4t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2506                    return;
2507                }
2508                if *dtype == TensorDtype::Q1 || *dtype == TensorDtype::Q1G {
2509                    // GPU batched q1 GEMM for wide prefill (q1_mul_mm on the
2510                    // device); the probe keeps whichever beats the CPU matmat.
2511                    // q1g has the twin `q1g_mul_mm` (cols % 128 by format).
2512                    let w = self.bin_w().expect("binary dtype");
2513                    if b >= 32
2514                        && b * rows * cols >= 128_000_000
2515                        && cols % 64 == 0
2516                        && crate::gpu::enabled_here()
2517                    {
2518                        if let Self::Mapped { model, idx, .. } = self {
2519                            let t0 = std::time::Instant::now();
2520                            match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
2521                                crate::gpu::ProbeArm::Gpu => {
2522                                    let ok = if *dtype == TensorDtype::Q1G {
2523                                        crate::gpu::q1g_matmat(
2524                                            model, *idx, xs_all, b, rows, cols, out,
2525                                        )
2526                                    } else {
2527                                        crate::gpu::q1_matmat(
2528                                            model, *idx, xs_all, b, rows, cols, out,
2529                                        )
2530                                    };
2531                                    if ok {
2532                                        crate::gpu::probe_record(
2533                                            crate::gpu::OpClass::Matmat,
2534                                            true,
2535                                            t0.elapsed(),
2536                                        );
2537                                        return;
2538                                    }
2539                                }
2540                                crate::gpu::ProbeArm::CpuTimed => {
2541                                    bin_matmat(w, xs_all, b, rows, cols, out, pool);
2542                                    crate::gpu::probe_record(
2543                                        crate::gpu::OpClass::Matmat,
2544                                        false,
2545                                        t0.elapsed(),
2546                                    );
2547                                    return;
2548                                }
2549                                crate::gpu::ProbeArm::Cpu => {}
2550                            }
2551                        }
2552                    }
2553                    bin_matmat(w, xs_all, b, rows, cols, out, pool);
2554                    return;
2555                }
2556                if *dtype == TensorDtype::Q1T {
2557                    // GPU batched GEMM for wide prefill (base + overlay on the
2558                    // device); probe keeps the winner vs the CPU matmat.
2559                    if b >= 32 && b * rows * cols >= 128_000_000 && crate::gpu::enabled_here() {
2560                        if let Self::Mapped { model, idx, .. } = self {
2561                            let t0 = std::time::Instant::now();
2562                            match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
2563                                crate::gpu::ProbeArm::Gpu => {
2564                                    if crate::gpu::q1t_matmat(
2565                                        model, *idx, xs_all, b, rows, cols, out,
2566                                    ) {
2567                                        crate::gpu::probe_record(
2568                                            crate::gpu::OpClass::Matmat,
2569                                            true,
2570                                            t0.elapsed(),
2571                                        );
2572                                        return;
2573                                    }
2574                                }
2575                                crate::gpu::ProbeArm::CpuTimed => {
2576                                    q1t_matmat(
2577                                        self.quant_bytes(),
2578                                        xs_all,
2579                                        b,
2580                                        rows,
2581                                        cols,
2582                                        out,
2583                                        pool,
2584                                    );
2585                                    crate::gpu::probe_record(
2586                                        crate::gpu::OpClass::Matmat,
2587                                        false,
2588                                        t0.elapsed(),
2589                                    );
2590                                    return;
2591                                }
2592                                crate::gpu::ProbeArm::Cpu => {}
2593                            }
2594                        }
2595                    }
2596                    q1t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2597                    return;
2598                }
2599                if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
2600                    vbitmatmat(
2601                        self.quant_bytes(),
2602                        vbit_offsets,
2603                        xs_all,
2604                        b,
2605                        rows,
2606                        cols,
2607                        out,
2608                        pool,
2609                    );
2610                    return;
2611                }
2612                let _pmp = crate::cpuprof::spans::span("mm.prescale");
2613                let pre: Vec<std::borrow::Cow<'_, [f32]>> = (0..b)
2614                    .map(|bi| prescale(&xs_all[bi * cols..(bi + 1) * cols], col_field, *dtype))
2615                    .collect();
2616                drop(_pmp);
2617                // MiMo verification is a 2–4 row decode panel, not a wide
2618                // prompt GEMM. Keep q8 projections on the same device as
2619                // decode; the generic b>=8 gate otherwise silently moves
2620                // every projection back to CPU. The short wgpu matmat uses
2621                // the same 64-lane reduction as its single-token matvec.
2622                if row_exact()
2623                    && (1..=4).contains(&b)
2624                    && matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
2625                    && crate::gpu::enabled_here()
2626                    && crate::gpu::wgpu_active()
2627                {
2628                    let flat: Vec<f32> = pre.iter().flat_map(|v| v.iter().copied()).collect();
2629                    if crate::gpu::q8_matmat(model, *idx, row_scale, &flat, b, rows, cols, out) {
2630                        return;
2631                    }
2632                }
2633                // D5: large prefill-batch GEMMs — on the GPU (threshold by
2634                // work volume: submission carries b×rows×cols MACs).
2635                // Runtime probe: the naive GEMM shader + sync readback
2636                // lose to the CPU GEMM on slow driver stacks — alternate
2637                // both arms and keep the winner.
2638                // CMF_MATMAT_GPU=0: diagnostic off-switch for this D5
2639                // device arm only (default unchanged). The GPU int8 kernel
2640                // and the CPU int8 kernel are not bit-comparable and differ
2641                // ~1e-2/layer (measured on SAM3), which compounds through
2642                // deep stacks — validation harnesses that need a single
2643                // changed block per run turn the others off with this.
2644                if b >= 8
2645                    && b * rows * cols >= 128_000_000
2646                    && crate::gpu::enabled_here()
2647                    && std::env::var("CMF_MATMAT_GPU").as_deref() != Ok("0")
2648                {
2649                    if let Self::Mapped { model, idx, .. } = self {
2650                        MATMAT_GATE.fetch_add(1, Ordering::Relaxed);
2651                        let t0 = std::time::Instant::now();
2652                        match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
2653                            crate::gpu::ProbeArm::Gpu
2654                                if crate::gpu::probe_deciding(crate::gpu::OpClass::Matmat)
2655                                    && !crate::gpu::q8_resident_or_upload(model, *idx) =>
2656                            {
2657                                // Cold weights during probing: the upload
2658                                // has started, the count runs on the CPU —
2659                                // the GPU arm samples on the next touch.
2660                                let q = self.quant_bytes();
2661                                qmatmat(q, row_scale, &pre, rows, cols, out, pool);
2662                                return;
2663                            }
2664                            crate::gpu::ProbeArm::Gpu => {
2665                                // The two-field codec's own entry first, as
2666                                // `device_matmat` takes it: the column field
2667                                // folds into the weight plane and the GEMM
2668                                // runs on the matrix units. The plain int8
2669                                // entry below is the scalar f32 GEMM — on a
2670                                // Spark-X2.5 4B prefill it ran q|k|v|o at
2671                                // host speed.
2672                                if *dtype == TensorDtype::Q8_2f
2673                                    && std::env::var("CMF_Q8_2F_DEV").as_deref() != Ok("0")
2674                                    && crate::gpu::q8_matmat_2f(
2675                                        model, *idx, row_scale, col_field, xs_all, b, rows,
2676                                        cols, out,
2677                                    )
2678                                {
2679                                    MATMAT_GPU_HIT.fetch_add(1, Ordering::Relaxed);
2680                                    crate::gpu::probe_record(
2681                                        crate::gpu::OpClass::Matmat,
2682                                        true,
2683                                        t0.elapsed(),
2684                                    );
2685                                    return;
2686                                }
2687                                let flat: Vec<f32> =
2688                                    pre.iter().flat_map(|v| v.iter().copied()).collect();
2689                                if crate::gpu::q8_matmat(
2690                                    model, *idx, row_scale, &flat, b, rows, cols, out,
2691                                ) {
2692                                    crate::gpu::probe_record(
2693                                        crate::gpu::OpClass::Matmat,
2694                                        true,
2695                                        t0.elapsed(),
2696                                    );
2697                                    return;
2698                                }
2699                            }
2700                            crate::gpu::ProbeArm::CpuTimed => {
2701                                let q = self.quant_bytes();
2702                                qmatmat(q, row_scale, &pre, rows, cols, out, pool);
2703                                crate::gpu::probe_record(
2704                                    crate::gpu::OpClass::Matmat,
2705                                    false,
2706                                    t0.elapsed(),
2707                                );
2708                                return;
2709                            }
2710                            crate::gpu::ProbeArm::Cpu => {}
2711                        }
2712                    }
2713                }
2714                let q = self.quant_bytes();
2715                qmatmat(q, row_scale, &pre, rows, cols, out, pool);
2716            }
2717        }
2718    }
2719}
2720
2721impl QTensor {
2722    /// `mxfp4` matvec: the wgpu kernel for lm_head/expert-class shapes when
2723    /// the probe prefers the device, else the CPU kernels of
2724    /// `qtensor::mxfp4` (always on native Metal, which has no mxfp4 kernel).
2725    fn mxfp4_matvec(&self, x: &[f32], out: &mut [f32], pool: Option<&Pool>) {
2726        let (rows, cols) = (self.rows(), self.cols());
2727        if let Self::Mapped { model, idx, .. } = self {
2728            if rows * cols >= 8_388_608 && crate::gpu::enabled_here() {
2729                let t0 = std::time::Instant::now();
2730                let cls = crate::gpu::matvec_class(rows, cols);
2731                match crate::gpu::probe_arm(cls) {
2732                    crate::gpu::ProbeArm::Gpu => {
2733                        if crate::gpu::mxfp4_matvec(model, *idx, x, rows, cols, out) {
2734                            crate::gpu::probe_record(cls, true, t0.elapsed());
2735                            return;
2736                        }
2737                    }
2738                    crate::gpu::ProbeArm::CpuTimed => {
2739                        mxfp4::matvec(self.quant_bytes(), x, rows, cols, out, pool);
2740                        crate::gpu::probe_record(cls, false, t0.elapsed());
2741                        return;
2742                    }
2743                    crate::gpu::ProbeArm::Cpu => {}
2744                }
2745            }
2746        }
2747        mxfp4::matvec(self.quant_bytes(), x, rows, cols, out, pool);
2748    }
2749
2750    /// `mxfp4` batched matvec (`xs` `[b, cols]` → `out` `[b, rows]`): the
2751    /// device's batched kernel behind the same probe as the other codecs'
2752    /// GEMMs (never for row-exact batches — the device sums in another
2753    /// order than the host matvec), else the CPU kernels.
2754    fn mxfp4_matmat(&self, xs: &[f32], b: usize, out: &mut [f32], pool: Option<&Pool>) {
2755        let (rows, cols) = (self.rows(), self.cols());
2756        if let Self::Mapped { model, idx, .. } = self {
2757            if b >= 2
2758                && b * rows * cols >= 33_554_432
2759                && !row_exact()
2760                && !crate::gpu::mm_killed()
2761                && crate::gpu::enabled_here()
2762            {
2763                let class = if b >= 128 {
2764                    crate::gpu::OpClass::MatmatWide
2765                } else {
2766                    crate::gpu::OpClass::Matmat
2767                };
2768                let t0 = std::time::Instant::now();
2769                match crate::gpu::probe_arm(class) {
2770                    crate::gpu::ProbeArm::Gpu => {
2771                        if crate::gpu::mxfp4_matmat(model, *idx, xs, b, rows, cols, out) {
2772                            crate::gpu::probe_record(class, true, t0.elapsed());
2773                            return;
2774                        }
2775                    }
2776                    crate::gpu::ProbeArm::CpuTimed => {
2777                        mxfp4::matmat(self.quant_bytes(), xs, b, rows, cols, out, pool);
2778                        crate::gpu::probe_record(class, false, t0.elapsed());
2779                        return;
2780                    }
2781                    crate::gpu::ProbeArm::Cpu => {}
2782                }
2783            }
2784        }
2785        mxfp4::matmat(self.quant_bytes(), xs, b, rows, cols, out, pool);
2786    }
2787}
2788
2789impl QTensor {
2790    /// The device GEMM this tensor would take, run once on the caller's
2791    /// data — the startup parity probe's arm, and the one place that knows
2792    /// which entry point each codec has.
2793    ///
2794    /// It exists because the probe used to look for a `q4tp` weight by
2795    /// name AND dtype, and a container packed any other way was declared
2796    /// "host path" for the whole render even though its codec had a device
2797    /// GEMM of its own. A gate that only recognizes one codec is a gate
2798    /// that silently downgrades every other one.
2799    pub fn device_matmat(&self, xs: &[f32], b: usize, out: &mut [f32]) -> bool {
2800        let (rows, cols) = (self.rows(), self.cols());
2801        let Self::Mapped {
2802            model,
2803            idx,
2804            dtype,
2805            row_scale,
2806            col_field,
2807            ..
2808        } = self
2809        else {
2810            return false;
2811        };
2812        if crate::prism::has_contract(model) {
2813            return false;
2814        }
2815        match *dtype {
2816            TensorDtype::Q4TiledP => crate::gpu::q4tp_matmat(model, *idx, xs, b, rows, cols, out),
2817            // The two-field codec folds its column field into the
2818            // activation, which leaves a plain per-row int8 GEMM — the
2819            // same kernel `q8_row` uses, on both backends.
2820            TensorDtype::Q8Row | TensorDtype::Q8_2f => {
2821                // The field belongs to the weight; only a backend that cannot
2822                // apply it there makes a scaled copy of the activation.
2823                if *dtype == TensorDtype::Q8_2f
2824                    && std::env::var("CMF_Q8_2F_DEV").as_deref() != Ok("0")
2825                    && crate::gpu::q8_matmat_2f(
2826                        model, *idx, row_scale, col_field, xs, b, rows, cols, out,
2827                    )
2828                {
2829                    return true;
2830                }
2831                let flat: Vec<f32> = (0..b)
2832                    .flat_map(|bi| {
2833                        prescale(&xs[bi * cols..(bi + 1) * cols], col_field, *dtype).into_owned()
2834                    })
2835                    .collect();
2836                crate::gpu::q8_matmat(model, *idx, row_scale, &flat, b, rows, cols, out)
2837            }
2838            _ => false,
2839        }
2840    }
2841
2842    /// Multi-matrix job (roadmap §3 P0): N tensors sharing one input
2843    /// run under a SINGLE pool dispatch — QKV or gate+up cost one
2844    /// barrier instead of N. Per-row math is the exact same kernel as
2845    /// `matvec` (bit-identical outputs); only the dispatch is fused.
2846    /// Falls back to N sequential matvecs when the set is not a uniform
2847    /// q8-family/F32 group or there is no pool.
2848    pub fn matvec_many<const N: usize>(
2849        ts: [&QTensor; N],
2850        x: &[f32],
2851        mut outs: [&mut [f32]; N],
2852        pool: Option<&Pool>,
2853    ) {
2854        if Self::any_flint_exact_cpu_default(&ts) {
2855            return float_activations_scope(|| Self::matvec_many(ts, x, outs, pool));
2856        }
2857        let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
2858        if ts.iter().any(|t| t.has_prism_contract() || t.per_tensor_only()) {
2859            // The fused range kernels have no transform descriptor.  Let
2860            // each tensor's ordinary matvec dispatch perform the explicit
2861            // signed FWHT (and retain CPU fallback for mixed q2tp/q4tp).
2862            // blr has no row-range kernel either: per-tensor calls.
2863            for (t, o) in ts.iter().zip(outs.iter_mut()) {
2864                t.matvec(x, o, pool);
2865            }
2866            return;
2867        }
2868        let uniform_q8 = ts.iter().all(|t| {
2869            matches!(
2870                t,
2871                Self::Mapped {
2872                    dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f,
2873                    ..
2874                }
2875            )
2876        });
2877        let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
2878        if uniform_f32 && crate::f32_backend::active() {
2879            for (t, o) in ts.iter().zip(outs.iter_mut()) {
2880                t.matvec(x, o, pool);
2881            }
2882            return;
2883        }
2884        let uniform_q4 = ts.iter().all(|t| {
2885            matches!(
2886                t,
2887                Self::Mapped {
2888                    dtype: TensorDtype::Q4Block,
2889                    ..
2890                }
2891            )
2892        });
2893        let uniform_vbit = ts.iter().all(|t| {
2894            matches!(
2895                t,
2896                Self::Mapped {
2897                    dtype: TensorDtype::Vbit | TensorDtype::VbitRo,
2898                    ..
2899                }
2900            )
2901        });
2902        // q1 and q1g may mix freely: the shared activation split and group
2903        // sums are layout-independent, only the row kernel differs.
2904        let uniform_q1 =
2905            ts.iter().all(|t| t.bin_w().is_some()) && ts.iter().all(|t| t.cols() == ts[0].cols());
2906        let uniform_q1t = ts.iter().all(|t| {
2907            matches!(
2908                t,
2909                Self::Mapped {
2910                    dtype: TensorDtype::Q1T,
2911                    ..
2912                }
2913            )
2914        });
2915        // q4tp is the skeleton dtype of the big MoE files, and without an arm
2916        // here every projection that shares an input paid its own pool
2917        // barrier: DeepSeek-V4's attention step alone hands this function
2918        // wq_a, wkv and both compressors' pairs off the same hidden state.
2919        let uniform_q4tp = ts.iter().all(|t| {
2920            matches!(
2921                t,
2922                Self::Mapped {
2923                    dtype: TensorDtype::Q4TiledP,
2924                    ..
2925                }
2926            )
2927        }) && ts
2928            .iter()
2929            .all(|t| t.cols() == ts[0].cols() && t.cols() % GROUP_SIZE == 0);
2930        let Some(pool) = pool else {
2931            for (t, o) in ts.iter().zip(outs.iter_mut()) {
2932                t.matvec(x, o, None);
2933            }
2934            return;
2935        };
2936        if total_rows < 256
2937            || !(uniform_q8
2938                || uniform_f32
2939                || uniform_q4
2940                || uniform_vbit
2941                || uniform_q1
2942                || uniform_q1t
2943                || uniform_q4tp)
2944        {
2945            for (t, o) in ts.iter().zip(outs.iter_mut()) {
2946                t.matvec(x, o, Some(pool));
2947            }
2948            return;
2949        }
2950
2951        if uniform_q4tp {
2952            // Every tensor's rows laid end to end in one virtual row space,
2953            // so the whole set is ONE dispatch. The per-row body is the
2954            // `q4tp_matvec` arm verbatim — same activation split, same
2955            // accumulation order — so the outputs are bit-identical to the
2956            // sequential calls this replaces.
2957            let cols = ts[0].cols();
2958            let gpr = cols / GROUP_SIZE;
2959            let views: [Q4tpView; N] =
2960                std::array::from_fn(|i| Q4tpView::new(ts[i].quant_bytes(), ts[i].rows(), cols));
2961            let rows_of: [usize; N] = std::array::from_fn(|i| ts[i].rows());
2962            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2963            // flat index -> (which tensor, which of its rows)
2964            let locate = |flat: usize| -> (usize, usize) {
2965                let mut acc = 0;
2966                for (i, &r) in rows_of.iter().enumerate() {
2967                    if flat < acc + r {
2968                        return (i, flat - acc);
2969                    }
2970                    acc += r;
2971                }
2972                (rows_of.len() - 1, 0)
2973            };
2974            let (views, outs_addr) = (&views, &outs_addr);
2975            if a8w8_enabled() {
2976                let act = split_act(x);
2977                let act = &act;
2978                #[cfg(target_arch = "x86_64")]
2979                let lanes = q4tp_lanes();
2980                let run = |start: usize, end: usize| {
2981                    #[cfg(target_arch = "x86_64")]
2982                    if lanes > 1 && act.xg8.len() == gpr {
2983                        // L independent accumulator chains per worker —
2984                        // the plain body's math, lane-interleaved.
2985                        with_krow(gpr * lanes, |sc| {
2986                            // SAFETY: a8w8 implies AVX2; each (tensor, row)
2987                            // pair still belongs to exactly one worker.
2988                            unsafe {
2989                                match lanes {
2990                                    2 => q4tp_many_lanes::<2>(
2991                                        &views[..],
2992                                        &locate,
2993                                        start,
2994                                        end,
2995                                        gpr,
2996                                        &act.xq,
2997                                        &act.xg8,
2998                                        act.sx,
2999                                        &act.outliers,
3000                                        sc,
3001                                        &outs_addr[..],
3002                                    ),
3003                                    8 => q4tp_many_lanes::<8>(
3004                                        &views[..],
3005                                        &locate,
3006                                        start,
3007                                        end,
3008                                        gpr,
3009                                        &act.xq,
3010                                        &act.xg8,
3011                                        act.sx,
3012                                        &act.outliers,
3013                                        sc,
3014                                        &outs_addr[..],
3015                                    ),
3016                                    _ => q4tp_many_lanes::<4>(
3017                                        &views[..],
3018                                        &locate,
3019                                        start,
3020                                        end,
3021                                        gpr,
3022                                        &act.xq,
3023                                        &act.xg8,
3024                                        act.sx,
3025                                        &act.outliers,
3026                                        sc,
3027                                        &outs_addr[..],
3028                                    ),
3029                                }
3030                            }
3031                        });
3032                        return;
3033                    }
3034                    crate::cpuprof::count_plain_rows((end - start) as u64);
3035                    with_krow(gpr, |sc| {
3036                        for flat in start..end {
3037                            let (t, r) = locate(flat);
3038                            let v = &views[t];
3039                            v.scales_into(r, gpr, sc);
3040                            let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, sc) * act.sx;
3041                            for &(j, xv) in &act.outliers {
3042                                let (w, s) = q4tp_outlier(v.nib, r, gpr, j, sc);
3043                                acc += w * s * xv;
3044                            }
3045                            // SAFETY: one worker owns each (tensor, row) pair.
3046                            unsafe { *outs_addr[t].at(r) = acc };
3047                        }
3048                    });
3049                };
3050                pool.run_rows(total_rows, &run);
3051            } else {
3052                let run = |start: usize, end: usize| {
3053                    with_krow(gpr, |sc| {
3054                        for flat in start..end {
3055                            let (t, r) = locate(flat);
3056                            let v = &views[t];
3057                            v.scales_into(r, gpr, sc);
3058                            // SAFETY: one worker owns each (tensor, row) pair.
3059                            unsafe { *outs_addr[t].at(r) = q4tp_row_exact(v.nib, r, gpr, x, sc) };
3060                        }
3061                    });
3062                };
3063                pool.run_rows(total_rows, &run);
3064            }
3065            return;
3066        }
3067
3068        if uniform_q1 {
3069            // One shared activation split + group sums (q1 has no col
3070            // field; the same input feeds every tensor).
3071            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
3072            if a8w8_enabled() {
3073                let act = split_act(x);
3074                let gsum = q1_group_sums(&act.xq, ts[0].cols() / GROUP_SIZE);
3075                let (act, gsum) = (&act, &gsum);
3076                let closures: [_; N] = std::array::from_fn(|i| {
3077                    let (w, gpr, out) = (
3078                        ts[i].bin_w().expect("uniform_q1 checked"),
3079                        ts[i].cols() / GROUP_SIZE,
3080                        outs_addr[i],
3081                    );
3082                    move |s: usize, e: usize| bin_range_a8w8(w, gpr, act, gsum, out, s, e)
3083                });
3084                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3085                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3086                pool.run_many(&parts);
3087            } else {
3088                let closures: [_; N] = std::array::from_fn(|i| {
3089                    let (w, gpr, out) = (
3090                        ts[i].bin_w().expect("uniform_q1 checked"),
3091                        ts[i].cols() / GROUP_SIZE,
3092                        outs_addr[i],
3093                    );
3094                    move |s: usize, e: usize| bin_range_f32(w, gpr, x, out, s, e)
3095                });
3096                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3097                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3098                pool.run_many(&parts);
3099            }
3100            return;
3101        }
3102
3103        if uniform_q1t {
3104            // Q1T batched: one shared activation split + overlay decode,
3105            // all tensors' rows in ONE pool dispatch (saves N−1 dispatches
3106            // and N−1 redundant split_act calls per layer).
3107            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
3108            const TILE: usize = cortiq_core::quant::Q1T_TILE;
3109            if a8w8_enabled() {
3110                let act = split_act(x);
3111                let act = &act;
3112                let x_ref = x;
3113                let closures: [_; N] = std::array::from_fn(|i| {
3114                    let bytes = ts[i].quant_bytes();
3115                    let (rows, cols) = (ts[i].rows(), ts[i].cols());
3116                    let gpr = cols / GROUP_SIZE;
3117                    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
3118                    let out = outs_addr[i];
3119                    move |s: usize, e: usize| {
3120                        q1t_range_a8w8(bytes, gpr, rp_off, ent_off, has_ov, act, x_ref, out, s, e)
3121                    }
3122                });
3123                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3124                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3125                pool.run_many(&parts);
3126            } else {
3127                let x_ref = x;
3128                let closures: [_; N] = std::array::from_fn(|i| {
3129                    let bytes = ts[i].quant_bytes();
3130                    let (rows, cols) = (ts[i].rows(), ts[i].cols());
3131                    let gpr = cols / GROUP_SIZE;
3132                    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
3133                    let out = outs_addr[i];
3134                    move |s: usize, e: usize| {
3135                        q1t_range_f32_batch(bytes, gpr, rp_off, ent_off, has_ov, x_ref, out, s, e)
3136                    }
3137                });
3138                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3139                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3140                pool.run_many(&parts);
3141            }
3142            return;
3143        }
3144
3145        if uniform_q4 || uniform_vbit {
3146            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
3147            // q4/vbit share one activation split — no per-tensor col field.
3148            if a8w8_enabled() {
3149                let act = split_act(x);
3150                let act = &act;
3151                if uniform_q4 {
3152                    let closures: [_; N] = std::array::from_fn(|i| {
3153                        let (packed, scales) =
3154                            q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
3155                        let (gpr, cols, out) =
3156                            (ts[i].cols() / GROUP_SIZE, ts[i].cols(), outs_addr[i]);
3157                        move |s: usize, e: usize| {
3158                            q4_range_a8w8(packed, scales, gpr, cols, act, out, s, e)
3159                        }
3160                    });
3161                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3162                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3163                    pool.run_many(&parts);
3164                } else {
3165                    let closures: [_; N] = std::array::from_fn(|i| {
3166                        let Self::Mapped { vbit_offsets, .. } = ts[i] else {
3167                            unreachable!()
3168                        };
3169                        let (bytes, rows, cols, out) = (
3170                            ts[i].quant_bytes(),
3171                            ts[i].rows(),
3172                            ts[i].cols(),
3173                            outs_addr[i],
3174                        );
3175                        move |s: usize, e: usize| {
3176                            vbit_range_a8w8(bytes, vbit_offsets, x, act, rows, cols, out, s, e)
3177                        }
3178                    });
3179                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3180                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3181                    pool.run_many(&parts);
3182                }
3183                return;
3184            }
3185            if uniform_q4 {
3186                let closures: [_; N] = std::array::from_fn(|i| {
3187                    let (packed, scales) =
3188                        q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
3189                    let (gpr, out) = (ts[i].cols() / GROUP_SIZE, outs_addr[i]);
3190                    move |s: usize, e: usize| q4_range_f32(packed, scales, gpr, x, out, s, e)
3191                });
3192                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3193                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3194                pool.run_many(&parts);
3195            } else {
3196                let closures: [_; N] = std::array::from_fn(|i| {
3197                    let Self::Mapped { vbit_offsets, .. } = ts[i] else {
3198                        unreachable!()
3199                    };
3200                    let (bytes, rows, cols, out) = (
3201                        ts[i].quant_bytes(),
3202                        ts[i].rows(),
3203                        ts[i].cols(),
3204                        outs_addr[i],
3205                    );
3206                    move |s: usize, e: usize| {
3207                        vbit_range_f32(bytes, vbit_offsets, x, rows, cols, out, s, e)
3208                    }
3209                });
3210                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3211                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3212                pool.run_many(&parts);
3213            }
3214            return;
3215        }
3216
3217        if uniform_f32 {
3218            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
3219            let closures: [_; N] = std::array::from_fn(|i| {
3220                let Self::F32 { data, cols, .. } = ts[i] else {
3221                    unreachable!()
3222                };
3223                let out = outs_addr[i];
3224                move |start: usize, end: usize| {
3225                    for o in start..end {
3226                        let row = &data[o * cols..(o + 1) * cols];
3227                        let mut sum = 0.0f32;
3228                        for j in 0..*cols {
3229                            sum += row[j] * x[j];
3230                        }
3231                        // SAFETY: disjoint (tensor, row) cells per worker.
3232                        unsafe { *out.at(o) = sum };
3233                    }
3234                }
3235            });
3236            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3237                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3238            pool.run_many(&parts);
3239            return;
3240        }
3241
3242        // Uniform q8-family: per-tensor prescale (q8_2f col fields
3243        // differ per tensor) + the shared range kernels.
3244        struct Ctx<'a> {
3245            bytes: &'a [u8],
3246            #[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
3247            rep: &'a [u8],
3248            row_scale: &'a [f32],
3249            cols: usize,
3250            xs: std::borrow::Cow<'a, [f32]>,
3251        }
3252        let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
3253            let Self::Mapped {
3254                dtype,
3255                cols,
3256                row_scale,
3257                col_field,
3258                repack,
3259                ..
3260            } = ts[i]
3261            else {
3262                unreachable!()
3263            };
3264            Ctx {
3265                bytes: ts[i].quant_bytes(),
3266                rep: repack,
3267                row_scale,
3268                cols: *cols,
3269                xs: prescale(x, col_field, *dtype),
3270            }
3271        });
3272        let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
3273        #[cfg(target_arch = "aarch64")]
3274        if sdot_enabled() {
3275            let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
3276            let closures: [_; N] = std::array::from_fn(|i| {
3277                let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
3278                move |start: usize, end: usize| {
3279                    q8_range_sdot(c.bytes, c.rep, c.row_scale, act, c.cols, out, start, end)
3280                }
3281            });
3282            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3283                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3284            pool.run_many(&parts);
3285            return;
3286        }
3287        #[cfg(target_arch = "x86_64")]
3288        if avx2_a8w8_enabled() {
3289            let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
3290            let closures: [_; N] = std::array::from_fn(|i| {
3291                let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
3292                move |start: usize, end: usize| {
3293                    q8_range_avx2(c.bytes, c.row_scale, act, c.cols, out, start, end)
3294                }
3295            });
3296            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3297                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3298            pool.run_many(&parts);
3299            return;
3300        }
3301        let closures: [_; N] = std::array::from_fn(|i| {
3302            let (c, out) = (&ctxs[i], outs_addr[i]);
3303            move |start: usize, end: usize| {
3304                q8_range_f32(c.bytes, c.row_scale, &c.xs, c.cols, out, start, end)
3305            }
3306        });
3307        let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3308            std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3309        pool.run_many(&parts);
3310    }
3311}
3312
3313impl QTensor {
3314    /// Pair-input multi-matrix job: N tensors × 2 shared inputs under a
3315    /// single pool dispatch — the MTP/pair decode path publishes one job
3316    /// for Q/K/V (and one for gate+up) instead of one per tensor.
3317    /// Per-row math is exactly `matvec2`'s kernels; bit-identical.
3318    #[allow(clippy::needless_range_loop)]
3319    pub fn matvec2_many<const N: usize>(
3320        ts: [&QTensor; N],
3321        x1: &[f32],
3322        x2: &[f32],
3323        mut o1s: [&mut [f32]; N],
3324        mut o2s: [&mut [f32]; N],
3325        pool: Option<&Pool>,
3326    ) {
3327        if Self::any_flint_exact_cpu_default(&ts) {
3328            return float_activations_scope(|| {
3329                Self::matvec2_many(ts, x1, x2, o1s, o2s, pool)
3330            });
3331        }
3332        let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
3333        if ts.iter().any(|t| t.has_prism_contract() || t.per_tensor_only()) {
3334            for i in 0..N {
3335                ts[i].matvec2(x1, x2, o1s[i], o2s[i], pool);
3336            }
3337            return;
3338        }
3339        let uniform_q8 = ts.iter().all(|t| {
3340            matches!(
3341                t,
3342                Self::Mapped {
3343                    dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f,
3344                    ..
3345                }
3346            )
3347        });
3348        let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
3349        let uniform_q4 = ts.iter().all(|t| {
3350            matches!(
3351                t,
3352                Self::Mapped {
3353                    dtype: TensorDtype::Q4Block,
3354                    ..
3355                }
3356            )
3357        });
3358        let uniform_vbit = ts.iter().all(|t| {
3359            matches!(
3360                t,
3361                Self::Mapped {
3362                    dtype: TensorDtype::Vbit | TensorDtype::VbitRo,
3363                    ..
3364                }
3365            )
3366        });
3367        let fusable = pool.is_some()
3368            && total_rows >= 256
3369            && (uniform_q8 || uniform_f32 || uniform_q4 || uniform_vbit);
3370        if !fusable {
3371            for i in 0..N {
3372                ts[i].matvec2(x1, x2, o1s[i], o2s[i], pool);
3373            }
3374            return;
3375        }
3376        let pool = pool.unwrap();
3377
3378        if uniform_q4 || uniform_vbit {
3379            let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
3380            let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
3381            // q4/vbit share activation splits — no per-tensor col field.
3382            if a8w8_enabled() {
3383                let a1 = split_act(x1);
3384                let a2 = split_act(x2);
3385                let (a1, a2) = (&a1, &a2);
3386                if uniform_q4 {
3387                    let closures: [_; N] = std::array::from_fn(|i| {
3388                        let (packed, scales) =
3389                            q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
3390                        let (gpr, cols, o1, o2) =
3391                            (ts[i].cols() / GROUP_SIZE, ts[i].cols(), p1[i], p2[i]);
3392                        move |s: usize, e: usize| {
3393                            q4_range2_a8w8(packed, scales, gpr, cols, a1, a2, o1, o2, s, e)
3394                        }
3395                    });
3396                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3397                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3398                    pool.run_many(&parts);
3399                } else {
3400                    let closures: [_; N] = std::array::from_fn(|i| {
3401                        let Self::Mapped { vbit_offsets, .. } = ts[i] else {
3402                            unreachable!()
3403                        };
3404                        let (bytes, rows, cols, o1, o2) = (
3405                            ts[i].quant_bytes(),
3406                            ts[i].rows(),
3407                            ts[i].cols(),
3408                            p1[i],
3409                            p2[i],
3410                        );
3411                        move |s: usize, e: usize| {
3412                            vbit_range2_a8w8(
3413                                bytes,
3414                                vbit_offsets,
3415                                x1,
3416                                x2,
3417                                a1,
3418                                a2,
3419                                rows,
3420                                cols,
3421                                o1,
3422                                o2,
3423                                s,
3424                                e,
3425                            )
3426                        }
3427                    });
3428                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3429                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3430                    pool.run_many(&parts);
3431                }
3432                return;
3433            }
3434            if uniform_q4 {
3435                let closures: [_; N] = std::array::from_fn(|i| {
3436                    let (packed, scales) =
3437                        q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
3438                    let (gpr, o1, o2) = (ts[i].cols() / GROUP_SIZE, p1[i], p2[i]);
3439                    move |s: usize, e: usize| {
3440                        q4_range2_f32(packed, scales, gpr, x1, x2, o1, o2, s, e)
3441                    }
3442                });
3443                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3444                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3445                pool.run_many(&parts);
3446            } else {
3447                let closures: [_; N] = std::array::from_fn(|i| {
3448                    let Self::Mapped { vbit_offsets, .. } = ts[i] else {
3449                        unreachable!()
3450                    };
3451                    let (bytes, rows, cols, o1, o2) = (
3452                        ts[i].quant_bytes(),
3453                        ts[i].rows(),
3454                        ts[i].cols(),
3455                        p1[i],
3456                        p2[i],
3457                    );
3458                    move |s: usize, e: usize| {
3459                        vbit_range2_f32(bytes, vbit_offsets, x1, x2, rows, cols, o1, o2, s, e)
3460                    }
3461                });
3462                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3463                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3464                pool.run_many(&parts);
3465            }
3466            return;
3467        }
3468
3469        if uniform_f32 {
3470            let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
3471            let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
3472            let closures: [_; N] = std::array::from_fn(|i| {
3473                let Self::F32 { data, cols, .. } = ts[i] else {
3474                    unreachable!()
3475                };
3476                let (o1, o2) = (p1[i], p2[i]);
3477                move |start: usize, end: usize| {
3478                    for o in start..end {
3479                        let row = &data[o * cols..(o + 1) * cols];
3480                        let (mut s1, mut s2) = (0.0f32, 0.0f32);
3481                        for j in 0..*cols {
3482                            s1 += row[j] * x1[j];
3483                            s2 += row[j] * x2[j];
3484                        }
3485                        // SAFETY: disjoint (tensor, row) cells per worker.
3486                        unsafe {
3487                            *o1.at(o) = s1;
3488                            *o2.at(o) = s2;
3489                        }
3490                    }
3491                }
3492            });
3493            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3494                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3495            pool.run_many(&parts);
3496            return;
3497        }
3498
3499        struct Ctx<'a> {
3500            bytes: &'a [u8],
3501            row_scale: &'a [f32],
3502            cols: usize,
3503            xs1: std::borrow::Cow<'a, [f32]>,
3504            xs2: std::borrow::Cow<'a, [f32]>,
3505        }
3506        let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
3507            let Self::Mapped {
3508                dtype,
3509                cols,
3510                row_scale,
3511                col_field,
3512                ..
3513            } = ts[i]
3514            else {
3515                unreachable!()
3516            };
3517            Ctx {
3518                bytes: ts[i].quant_bytes(),
3519                row_scale,
3520                cols: *cols,
3521                xs1: prescale(x1, col_field, *dtype),
3522                xs2: prescale(x2, col_field, *dtype),
3523            }
3524        });
3525        let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
3526        let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
3527        #[cfg(target_arch = "aarch64")]
3528        if sdot_enabled() {
3529            let acts: [(SplitAct, SplitAct); N] =
3530                std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
3531            let closures: [_; N] = std::array::from_fn(|i| {
3532                let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
3533                move |start: usize, end: usize| {
3534                    q8_range2_sdot(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
3535                }
3536            });
3537            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3538                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3539            pool.run_many(&parts);
3540            return;
3541        }
3542        #[cfg(target_arch = "x86_64")]
3543        if avx2_a8w8_enabled() {
3544            let acts: [(SplitAct, SplitAct); N] =
3545                std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
3546            let closures: [_; N] = std::array::from_fn(|i| {
3547                let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
3548                move |start: usize, end: usize| {
3549                    q8_range2_avx2(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
3550                }
3551            });
3552            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3553                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3554            pool.run_many(&parts);
3555            return;
3556        }
3557        let closures: [_; N] = std::array::from_fn(|i| {
3558            let (c, o1, o2) = (&ctxs[i], p1[i], p2[i]);
3559            move |start: usize, end: usize| {
3560                q8_range2_f32(
3561                    c.bytes,
3562                    c.row_scale,
3563                    &c.xs1,
3564                    &c.xs2,
3565                    c.cols,
3566                    o1,
3567                    o2,
3568                    start,
3569                    end,
3570                )
3571            }
3572        });
3573        let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3574            std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3575        pool.run_many(&parts);
3576    }
3577
3578    /// Fused gate+up matvec with SiLU·mul: for each row r, computes
3579    /// `silu(gate·x) * (up·x)` and writes to `out[r]`. ONE pool dispatch,
3580    /// no intermediate g/u buffers, no separate silu pass. Falls back
3581    /// (returns false) for unsupported dtype combos.
3582    pub fn matvec_silu_mul(
3583        gate: &QTensor,
3584        up: &QTensor,
3585        x: &[f32],
3586        out: &mut [f32],
3587        pool: Option<&Pool>,
3588    ) -> bool {
3589        Self::matvec_silu_mul_limited(gate, up, x, out, 0.0, pool)
3590    }
3591
3592    /// Fused gate+up+SiLU with the GLM asymmetrical clamp.  `limit == 0`
3593    /// preserves the historical unclamped helper; a positive limit clamps
3594    /// `up` to both sides and `gate` only from above, matching the GLM
3595    /// SwiGLU reference.  Keeping the limit in the row kernel avoids the two
3596    /// intermediate vectors and the extra combine pass on the Q2TP experts.
3597    pub fn matvec_silu_mul_limited(
3598        gate: &QTensor,
3599        up: &QTensor,
3600        x: &[f32],
3601        out: &mut [f32],
3602        limit: f32,
3603        pool: Option<&Pool>,
3604    ) -> bool {
3605        if gate.flint_exact_cpu_default() || up.flint_exact_cpu_default() {
3606            return float_activations_scope(|| {
3607                Self::matvec_silu_mul_limited(gate, up, x, out, limit, pool)
3608            });
3609        }
3610        if gate.has_prism_contract()
3611            || up.has_prism_contract()
3612            || gate.per_tensor_only()
3613            || up.per_tensor_only()
3614        {
3615            // The fused gate/up kernels consume x directly.  Prism requires
3616            // a per-matrix signed FWHT, so the caller must use two ordinary
3617            // descriptor-aware matvecs instead of an unrotated fast path.
3618            return false;
3619        }
3620        let inter = gate.rows();
3621        debug_assert_eq!(up.rows(), inter);
3622        debug_assert_eq!(out.len(), inter);
3623        debug_assert_eq!(gate.cols(), up.cols());
3624        if !a8w8_enabled() {
3625            return false;
3626        }
3627        let act = split_act(x);
3628        let act = &act;
3629        let x_ref = x;
3630        let out_addr = SendMut(out.as_mut_ptr());
3631
3632        match (gate, up) {
3633            // Q4Block gate + Q4Block up (most common mobile q4 models)
3634            (
3635                Self::Mapped {
3636                    dtype: TensorDtype::Q4Block,
3637                    ..
3638                },
3639                Self::Mapped {
3640                    dtype: TensorDtype::Q4Block,
3641                    ..
3642                },
3643            ) => {
3644                let (gp, gs) = q4_split(gate.quant_bytes(), gate.rows(), gate.cols());
3645                let (up_p, up_s) = q4_split(up.quant_bytes(), up.rows(), up.cols());
3646                let gpr = gate.cols() / GROUP_SIZE;
3647                let cols = gate.cols();
3648                let run = move |start: usize, end: usize| {
3649                    for r in start..end {
3650                        let mut gv = dot_q4_row_i8(gp, gs, r * gpr, gpr, &act.xq) * act.sx;
3651                        let mut uv = dot_q4_row_i8(up_p, up_s, r * gpr, gpr, &act.xq) * act.sx;
3652                        for &(j, xv) in &act.outliers {
3653                            let flat = r * cols + j;
3654                            let gb = gp[flat / 2];
3655                            let gn = if flat & 1 == 0 { gb & 0x0F } else { gb >> 4 };
3656                            let gsc = f16_to_f32(u16::from_le_bytes([
3657                                gs[(flat / GROUP_SIZE) * 2],
3658                                gs[(flat / GROUP_SIZE) * 2 + 1],
3659                            ]));
3660                            gv += ((gn as i32 - 8) as f32) * gsc * xv;
3661                            let ub = up_p[flat / 2];
3662                            let un = if flat & 1 == 0 { ub & 0x0F } else { ub >> 4 };
3663                            let usc = f16_to_f32(u16::from_le_bytes([
3664                                up_s[(flat / GROUP_SIZE) * 2],
3665                                up_s[(flat / GROUP_SIZE) * 2 + 1],
3666                            ]));
3667                            uv += ((un as i32 - 8) as f32) * usc * xv;
3668                        }
3669                        // SAFETY: disjoint row ranges per worker.
3670                        unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3671                    }
3672                };
3673                dispatch_rows(pool, inter, &run);
3674                true
3675            }
3676            // Q4Tiled gate + Q4Tiled up — one row pass, both tile
3677            // streams sequential, silu·mul fused (same per-row math as
3678            // `q4t_matvec`).
3679            (
3680                Self::Mapped {
3681                    dtype: TensorDtype::Q4Tiled,
3682                    ..
3683                },
3684                Self::Mapped {
3685                    dtype: TensorDtype::Q4Tiled,
3686                    ..
3687                },
3688            ) => {
3689                let g_bytes = gate.quant_bytes();
3690                let u_bytes = up.quant_bytes();
3691                let gpr = gate.cols() / GROUP_SIZE;
3692                let run = move |start: usize, end: usize| {
3693                    for r in start..end {
3694                        let mut gv = dot_q4t_row_i8(g_bytes, r, gpr, &act.xq) * act.sx;
3695                        let mut uv = dot_q4t_row_i8(u_bytes, r, gpr, &act.xq) * act.sx;
3696                        for &(j, xv) in &act.outliers {
3697                            let (w, s) = q4t_outlier(g_bytes, r, gpr, j);
3698                            gv += w * s * xv;
3699                            let (w, s) = q4t_outlier(u_bytes, r, gpr, j);
3700                            uv += w * s * xv;
3701                        }
3702                        // SAFETY: disjoint row ranges per worker.
3703                        unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3704                    }
3705                };
3706                dispatch_rows(pool, inter, &run);
3707                true
3708            }
3709            // Q4TiledP gate + Q4TiledP up — the same fused row pass, with
3710            // each row's two ladders built once and spent on both streams.
3711            (
3712                Self::Mapped {
3713                    dtype: TensorDtype::Q4TiledP,
3714                    ..
3715                },
3716                Self::Mapped {
3717                    dtype: TensorDtype::Q4TiledP,
3718                    ..
3719                },
3720            ) => {
3721                let cols = gate.cols();
3722                let gpr = cols / GROUP_SIZE;
3723                let gv_view = Q4tpView::new(gate.quant_bytes(), inter, cols);
3724                let uv_view = Q4tpView::new(up.quant_bytes(), inter, cols);
3725                #[cfg(target_arch = "x86_64")]
3726                let lanes = q4tp_lanes();
3727                let run = |start: usize, end: usize| {
3728                    #[cfg(target_arch = "x86_64")]
3729                    if lanes > 1 && act.xg8.len() == gpr {
3730                        // L independent accumulator chains per worker —
3731                        // the plain arm's math, lane-interleaved.
3732                        with_krows(gpr * lanes, |gsc, usc| {
3733                            // SAFETY: a8w8 implies AVX2; disjoint row ranges.
3734                            unsafe {
3735                                match lanes {
3736                                    2 => q4tp_gateup_lanes::<2>(
3737                                        &gv_view, &uv_view, start, end, gpr, &act.xq,
3738                                        &act.xg8, act.sx, limit, &act.outliers, gsc, usc,
3739                                        out_addr,
3740                                    ),
3741                                    8 => q4tp_gateup_lanes::<8>(
3742                                        &gv_view, &uv_view, start, end, gpr, &act.xq,
3743                                        &act.xg8, act.sx, limit, &act.outliers, gsc, usc,
3744                                        out_addr,
3745                                    ),
3746                                    _ => q4tp_gateup_lanes::<4>(
3747                                        &gv_view, &uv_view, start, end, gpr, &act.xq,
3748                                        &act.xg8, act.sx, limit, &act.outliers, gsc, usc,
3749                                        out_addr,
3750                                    ),
3751                                }
3752                            }
3753                        });
3754                        return;
3755                    }
3756                    crate::cpuprof::count_plain_rows((end - start) as u64);
3757                    with_krows(gpr, |gsc, usc| {
3758                        for r in start..end {
3759                            gv_view.scales_into(r, gpr, gsc);
3760                            uv_view.scales_into(r, gpr, usc);
3761                            let mut gv =
3762                                dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsc) * act.sx;
3763                            let mut uv =
3764                                dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, usc) * act.sx;
3765                            for &(j, xv) in &act.outliers {
3766                                let (w, s) = q4tp_outlier(gv_view.nib, r, gpr, j, gsc);
3767                                gv += w * s * xv;
3768                                let (w, s) = q4tp_outlier(uv_view.nib, r, gpr, j, usc);
3769                                uv += w * s * xv;
3770                            }
3771                            // SAFETY: disjoint row ranges per worker.
3772                            unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3773                        }
3774                    });
3775                };
3776                dispatch_rows(pool, inter, &run);
3777                true
3778            }
3779            // Q1 gate + Q1 up — one row pass over both sign streams,
3780            // silu·mul fused (the per-row math of `q1_range_a8w8`); the
3781            // activation group sums are shared by both streams. Without
3782            // this arm a q1 dense FFN paid two dispatches + a combine
3783            // loop — the exact barrier this function exists to remove.
3784            // q1g (and a q1/q1g mix) rides the same arm through `BinW`.
3785            (
3786                Self::Mapped {
3787                    dtype: TensorDtype::Q1 | TensorDtype::Q1G,
3788                    ..
3789                },
3790                Self::Mapped {
3791                    dtype: TensorDtype::Q1 | TensorDtype::Q1G,
3792                    ..
3793                },
3794            ) => {
3795                let gw = gate.bin_w().expect("binary dtype matched");
3796                let uw = up.bin_w().expect("binary dtype matched");
3797                let gpr = gate.cols() / GROUP_SIZE;
3798                let gsum = q1_group_sums(&act.xq, gpr);
3799                let gsum = &gsum;
3800                let run = move |start: usize, end: usize| {
3801                    for r in start..end {
3802                        let mut gv = gw.dot_i8(r, gpr, &act.xq, gsum) * act.sx;
3803                        let mut uv = uw.dot_i8(r, gpr, &act.xq, gsum) * act.sx;
3804                        for &(j, xv) in &act.outliers {
3805                            let (w, s) = gw.outlier(r, gpr, j);
3806                            gv += w * s * xv;
3807                            let (w, s) = uw.outlier(r, gpr, j);
3808                            uv += w * s * xv;
3809                        }
3810                        // SAFETY: disjoint row ranges per worker.
3811                        unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3812                    }
3813                };
3814                dispatch_rows(pool, inter, &run);
3815                true
3816            }
3817            // Q2TiledP gate + Q2TiledP up — the 2-bit expert pair (MoE
3818            // FFNs of the W2 class): one row pass, both ladders built
3819            // once, integer code dots with shared group sums.
3820            (
3821                Self::Mapped {
3822                    dtype: TensorDtype::Q2TiledP,
3823                    ..
3824                },
3825                Self::Mapped {
3826                    dtype: TensorDtype::Q2TiledP,
3827                    ..
3828                },
3829            ) => {
3830                let cols = gate.cols();
3831                let gpr = cols / GROUP_SIZE;
3832                let gv_view = Q4tpView::new_q2(gate.quant_bytes(), inter, cols);
3833                let uv_view = Q4tpView::new_q2(up.quant_bytes(), inter, cols);
3834                let gsum = q1_group_sums(&act.xq, gpr);
3835                let gsum = &gsum;
3836                let run = move |start: usize, end: usize| {
3837                    with_krows(gpr, |gsc, usc| {
3838                        for r in start..end {
3839                            gv_view.scales_into(r, gpr, gsc);
3840                            uv_view.scales_into(r, gpr, usc);
3841                            let mut gv =
3842                                dot_q2tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsum, gsc) * act.sx;
3843                            let mut uv =
3844                                dot_q2tp_row_i8(uv_view.nib, r, gpr, &act.xq, gsum, usc) * act.sx;
3845                            for &(j, xv) in &act.outliers {
3846                                let (w, s) = q2tp_outlier(gv_view.nib, r, gpr, j, gsc);
3847                                gv += w * s * xv;
3848                                let (w, s) = q2tp_outlier(uv_view.nib, r, gpr, j, usc);
3849                                uv += w * s * xv;
3850                            }
3851                            // SAFETY: disjoint row ranges per worker.
3852                            unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3853                        }
3854                    });
3855                };
3856                dispatch_rows(pool, inter, &run);
3857                true
3858            }
3859            // Q8Row gate + Q8Row up — one row pass over both i8 streams.
3860            // Q8_2f stays out on purpose: its column field prescales the
3861            // activations PER TENSOR, which breaks this fn's shared
3862            // split_act contract — it keeps the two-dispatch path.
3863            (
3864                Self::Mapped {
3865                    dtype: TensorDtype::Q8Row,
3866                    row_scale: g_rs,
3867                    ..
3868                },
3869                Self::Mapped {
3870                    dtype: TensorDtype::Q8Row,
3871                    row_scale: u_rs,
3872                    ..
3873                },
3874            ) => {
3875                let g_bytes = gate.quant_bytes();
3876                let u_bytes = up.quant_bytes();
3877                let cols = gate.cols();
3878                let run = move |start: usize, end: usize| {
3879                    for r in start..end {
3880                        let gv = q8_row_dot(&g_bytes[r * cols..(r + 1) * cols], act) * g_rs[r];
3881                        let uv = q8_row_dot(&u_bytes[r * cols..(r + 1) * cols], act) * u_rs[r];
3882                        // SAFETY: disjoint row ranges per worker.
3883                        unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3884                    }
3885                };
3886                dispatch_rows(pool, inter, &run);
3887                true
3888            }
3889            // Q1T gate + Q1T up
3890            (
3891                Self::Mapped {
3892                    dtype: TensorDtype::Q1T,
3893                    ..
3894                },
3895                Self::Mapped {
3896                    dtype: TensorDtype::Q1T,
3897                    ..
3898                },
3899            ) => {
3900                const TILE: usize = cortiq_core::quant::Q1T_TILE;
3901                let g_bytes = gate.quant_bytes();
3902                let u_bytes = up.quant_bytes();
3903                let gpr = gate.cols() / GROUP_SIZE;
3904                let (g_rp, g_ent, g_ov) = q1t_overlay(g_bytes, inter * gpr * TILE, inter);
3905                let (u_rp, u_ent, u_ov) = q1t_overlay(u_bytes, inter * gpr * TILE, inter);
3906                let run = move |start: usize, end: usize| {
3907                    for r in start..end {
3908                        let mut gv = q1t_dot_row_i8(g_bytes, r, gpr, &act.xq) * act.sx;
3909                        let mut uv = q1t_dot_row_i8(u_bytes, r, gpr, &act.xq) * act.sx;
3910                        for &(j, xv) in &act.outliers {
3911                            gv += q1t_base_weight(g_bytes, r, gpr, j) * xv;
3912                            uv += q1t_base_weight(u_bytes, r, gpr, j) * xv;
3913                        }
3914                        gv += q1t_row_outlier_correction(g_bytes, r, g_rp, g_ent, g_ov, x_ref);
3915                        uv += q1t_row_outlier_correction(u_bytes, r, u_rp, u_ent, u_ov, x_ref);
3916                        // SAFETY: disjoint row ranges per worker.
3917                        unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3918                    }
3919                };
3920                dispatch_rows(pool, inter, &run);
3921                true
3922            }
3923            _ => false,
3924        }
3925    }
3926
3927    /// Every routed expert's fused gate/up/SiLU under ONE pool dispatch.
3928    ///
3929    /// The per-expert path pays a pool barrier per expert per stage: at 9
3930    /// experts over 40 layers that is ~720 barriers a token, and a decode
3931    /// profile of Qwen3.6-35B-A3B showed the pool parked in
3932    /// `psynch_cvwait` about twice as long as it spent computing. Laying
3933    /// every expert's rows end-to-end in one virtual row space collapses
3934    /// the stage to a single dispatch. The per-row body is the
3935    /// single-expert q4tp arm verbatim, so outputs are bit-identical.
3936    ///
3937    /// `false` = something is outside the fused q4tp kernel (dtype, shape,
3938    /// or a transformed tensor); the caller walks the ordinary per-expert
3939    /// path. Float activations use the same exact scalar rows, still fused
3940    /// under one pool dispatch.
3941    pub fn moe_gate_up_many(
3942        pairs: &[(&QTensor, &QTensor)],
3943        x: &[f32],
3944        outs: &mut [Vec<f32>],
3945        pool: Option<&Pool>,
3946    ) -> bool {
3947        Self::moe_gate_up_many_limited(pairs, x, outs, 0.0, pool)
3948    }
3949
3950    /// Batched gate/up/SiLU with the optional GLM clamp.  The public legacy
3951    /// helper above keeps its historical unclamped semantics; callers that
3952    /// implement a reference with a positive SwiGLU limit use this variant.
3953    pub fn moe_gate_up_many_limited(
3954        pairs: &[(&QTensor, &QTensor)],
3955        x: &[f32],
3956        outs: &mut [Vec<f32>],
3957        limit: f32,
3958        pool: Option<&Pool>,
3959    ) -> bool {
3960        if Self::any_flint_exact_cpu_default_pairs(pairs) {
3961            return float_activations_scope(|| {
3962                Self::moe_gate_up_many_limited(pairs, x, outs, limit, pool)
3963            });
3964        }
3965        if pairs.is_empty() || pairs.len() != outs.len() {
3966            return false;
3967        }
3968        if pairs.iter().any(|(g, u)| g.per_tensor_only() || u.per_tensor_only()) {
3969            return false;
3970        }
3971        if !a8w8_enabled() {
3972            if limit > 0.0 {
3973                // The exact-row fallback applies no SwiGLU clamp; the
3974                // caller's per-expert path carries it instead.
3975                return false;
3976            }
3977            let groups = vec![vec![0]; pairs.len()];
3978            return Self::moe_gate_up_rows(pairs, &groups, x, outs, pool);
3979        }
3980        let inter = pairs[0].0.rows();
3981        let cols = pairs[0].0.cols();
3982        if cols % GROUP_SIZE != 0 {
3983            return false;
3984        }
3985        let gpr = cols / GROUP_SIZE;
3986        // Uniform layout across every routed pair: q4tp, or the 2-bit
3987        // profile's q2tp gate/up (the W2 class). Mixed sets refuse.
3988        let q2 = matches!(
3989            pairs[0].0,
3990            Self::Mapped {
3991                dtype: TensorDtype::Q2TiledP,
3992                ..
3993            }
3994        );
3995        let want = if q2 {
3996            TensorDtype::Q2TiledP
3997        } else {
3998            TensorDtype::Q4TiledP
3999        };
4000        let mut views = Vec::with_capacity(pairs.len() * 2);
4001        for ((g, u), o) in pairs.iter().zip(outs.iter()) {
4002            let both = matches!(g, Self::Mapped { dtype, .. } if *dtype == want)
4003                && matches!(u, Self::Mapped { dtype, .. } if *dtype == want);
4004            if !both
4005                || g.rows() != inter
4006                || u.rows() != inter
4007                || g.cols() != cols
4008                || u.cols() != cols
4009                || o.len() != inter
4010            {
4011                return false;
4012            }
4013            let mk = if q2 { Q4tpView::new_q2 } else { Q4tpView::new };
4014            views.push(mk(g.quant_bytes(), inter, cols));
4015            views.push(mk(u.quant_bytes(), inter, cols));
4016        }
4017        let act = split_act(x);
4018        let gsum = if q2 {
4019            q1_group_sums(&act.xq, gpr)
4020        } else {
4021            Vec::new()
4022        };
4023        let (act, gsum) = (&act, &gsum);
4024        let ptrs: Vec<SendMut> = outs.iter_mut().map(|o| SendMut(o.as_mut_ptr())).collect();
4025        let (views, ptrs) = (&views, &ptrs);
4026        let run = |start: usize, end: usize| {
4027            with_krows(gpr, |gsc, usc| {
4028                for flat in start..end {
4029                    let (e, r) = (flat / inter, flat % inter);
4030                    let gv_view = &views[e * 2];
4031                    let uv_view = &views[e * 2 + 1];
4032                    gv_view.scales_into(r, gpr, gsc);
4033                    uv_view.scales_into(r, gpr, usc);
4034                    let (mut gv, mut uv) = if q2 {
4035                        (
4036                            dot_q2tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsum, gsc) * act.sx,
4037                            dot_q2tp_row_i8(uv_view.nib, r, gpr, &act.xq, gsum, usc) * act.sx,
4038                        )
4039                    } else {
4040                        (
4041                            dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsc) * act.sx,
4042                            dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, usc) * act.sx,
4043                        )
4044                    };
4045                    for &(j, xv) in &act.outliers {
4046                        let (og, ou) = if q2 {
4047                            (
4048                                q2tp_outlier(gv_view.nib, r, gpr, j, gsc),
4049                                q2tp_outlier(uv_view.nib, r, gpr, j, usc),
4050                            )
4051                        } else {
4052                            (
4053                                q4tp_outlier(gv_view.nib, r, gpr, j, gsc),
4054                                q4tp_outlier(uv_view.nib, r, gpr, j, usc),
4055                            )
4056                        };
4057                        gv += og.0 * og.1 * xv;
4058                        uv += ou.0 * ou.1 * xv;
4059                    }
4060                    // SAFETY: one worker owns each (expert, row) pair.
4061                    unsafe { *ptrs[e].at(r) = silu_mul_limited(gv, uv, limit) };
4062                }
4063            });
4064        };
4065        dispatch_rows(pool, pairs.len() * inter, &run);
4066        true
4067    }
4068
4069    /// Every routed expert's down projection, weighted and summed into
4070    /// `out`, under ONE pool dispatch.
4071    ///
4072    /// Partitioned by OUTPUT row rather than by expert: each row is owned
4073    /// by a single worker, so the experts are summed in the caller's order
4074    /// — the same sequence of f32 adds the serial `out[i] += w·eo[i]` loop
4075    /// performs, hence bit-identical. Partitioning by expert instead would
4076    /// race on the shared accumulator.
4077    pub fn moe_down_many(
4078        downs: &[&QTensor],
4079        gs: &[Vec<f32>],
4080        weights: &[f32],
4081        out: &mut [f32],
4082        pool: Option<&Pool>,
4083    ) -> bool {
4084        if Self::any_flint_exact_cpu_default(downs) {
4085            return float_activations_scope(|| {
4086                Self::moe_down_many(downs, gs, weights, out, pool)
4087            });
4088        }
4089        if downs.is_empty() || downs.len() != gs.len() || downs.len() != weights.len() {
4090            return false;
4091        }
4092        if downs.iter().any(|d| d.per_tensor_only()) {
4093            return false;
4094        }
4095        if !a8w8_enabled() {
4096            let mut terms = vec![vec![0.0; out.len()]; downs.len()];
4097            if !Self::moe_down_rows(downs, &vec![1; downs.len()], gs, &mut terms, pool) {
4098                return false;
4099            }
4100            out.fill(0.0);
4101            for (row, &w) in terms.iter().zip(weights) {
4102                for (o, &v) in out.iter_mut().zip(row) {
4103                    *o += w * v;
4104                }
4105            }
4106            return true;
4107        }
4108        let rows = out.len();
4109        let cols = downs[0].cols();
4110        if cols % GROUP_SIZE != 0 {
4111            return false;
4112        }
4113        let gpr = cols / GROUP_SIZE;
4114        let mut views = Vec::with_capacity(downs.len());
4115        for (d, g) in downs.iter().zip(gs.iter()) {
4116            if !matches!(
4117                d,
4118                Self::Mapped {
4119                    dtype: TensorDtype::Q4TiledP,
4120                    ..
4121                }
4122            ) || d.rows() != rows
4123                || d.cols() != cols
4124                || g.len() != cols
4125            {
4126                return false;
4127            }
4128            views.push(Q4tpView::new(d.quant_bytes(), rows, cols));
4129        }
4130        // One int8 split per expert — the activation vectors differ.
4131        let acts: Vec<SplitAct> = gs.iter().map(|g| split_act(g)).collect();
4132        // Partitioned by OUTPUT row, with the experts folded inside: each
4133        // row is owned by one worker, so they are summed in the caller's
4134        // order — the same f32 sequence the serial `out[i] += w·eo[i]`
4135        // loop produces. Partitioning by expert instead would either race
4136        // on the accumulator or need a scratch plane and a second pass;
4137        // measured, that variant was a wash, so this keeps the simpler
4138        // shape.
4139        let out_addr = SendMut(out.as_mut_ptr());
4140        let (views, acts, weights) = (&views, &acts, &weights);
4141        let run = |start: usize, end: usize| {
4142            with_krow(gpr, |sc| {
4143                for r in start..end {
4144                    let mut acc = 0f32;
4145                    for (e, v) in views.iter().enumerate() {
4146                        v.scales_into(r, gpr, sc);
4147                        let a = &acts[e];
4148                        let mut d = dot_q4tp_row_i8(v.nib, r, gpr, &a.xq, sc) * a.sx;
4149                        for &(j, xv) in &a.outliers {
4150                            let (w, s) = q4tp_outlier(v.nib, r, gpr, j, sc);
4151                            d += w * s * xv;
4152                        }
4153                        acc += weights[e] * d;
4154                    }
4155                    // SAFETY: disjoint row ranges per worker.
4156                    unsafe { *out_addr.at(r) = acc };
4157                }
4158            });
4159        };
4160        dispatch_rows(pool, rows, &run);
4161        true
4162    }
4163
4164    /// `moe_gate_up_many` for SEVERAL tokens at once, decode-exact: expert
4165    /// `e` (`pairs[e]`, q4tp) serves the tokens `groups[e]` (row indices
4166    /// into `xs`, each `cols` wide). Every (expert, token) output is
4167    /// bit-identical to `moe_gate_up_many` run on that token alone — the
4168    /// same int8 activation split, VNNI dots, outlier terms and inline
4169    /// SiLU — while each weight row is read once for all the tokens routed
4170    /// to its expert (the speculative verify's expert sharing). `outs` is
4171    /// flat in (expert, token-of-group) order. False = not covered (not
4172    /// q4tp): the caller takes the per-token path. With float activations,
4173    /// the exact scalar row kernel replaces the int8 dot without changing
4174    /// the shared dispatch or route-order reduction.
4175    pub fn moe_gate_up_rows(
4176        pairs: &[(&QTensor, &QTensor)],
4177        groups: &[Vec<usize>],
4178        xs: &[f32],
4179        outs: &mut [Vec<f32>],
4180        pool: Option<&Pool>,
4181    ) -> bool {
4182        if Self::any_flint_exact_cpu_default_pairs(pairs) {
4183            return float_activations_scope(|| {
4184                Self::moe_gate_up_rows(pairs, groups, xs, outs, pool)
4185            });
4186        }
4187        if pairs.is_empty() || pairs.len() != groups.len() {
4188            return false;
4189        }
4190        if pairs.iter().any(|(g, u)| g.per_tensor_only() || u.per_tensor_only()) {
4191            return false;
4192        }
4193        let inter = pairs[0].0.rows();
4194        let cols = pairs[0].0.cols();
4195        let n_pairs: usize = groups.iter().map(|g| g.len()).sum();
4196        if cols == 0 || cols % GROUP_SIZE != 0 || outs.len() != n_pairs || xs.len() % cols != 0 {
4197            return false;
4198        }
4199        let b = xs.len() / cols;
4200        let gpr = cols / GROUP_SIZE;
4201        let mut views = Vec::with_capacity(pairs.len() * 2);
4202        for (g, u) in pairs {
4203            let q4tp = |t: &QTensor| {
4204                matches!(
4205                    t,
4206                    Self::Mapped {
4207                        dtype: TensorDtype::Q4TiledP,
4208                        ..
4209                    }
4210                )
4211            };
4212            if g.has_prism_contract()
4213                || u.has_prism_contract()
4214                || !q4tp(g)
4215                || !q4tp(u)
4216                || g.rows() != inter
4217                || u.rows() != inter
4218                || g.cols() != cols
4219                || u.cols() != cols
4220            {
4221                return false;
4222            }
4223            views.push(Q4tpView::new(g.quant_bytes(), inter, cols));
4224            views.push(Q4tpView::new(u.quant_bytes(), inter, cols));
4225        }
4226        if outs.iter().any(|o| o.len() != inter) || groups.iter().flatten().any(|&t| t >= b) {
4227            return false;
4228        }
4229        let quantized = a8w8_enabled();
4230        let acts: Vec<SplitAct> = if quantized {
4231            (0..b)
4232                .map(|t| split_act(&xs[t * cols..(t + 1) * cols]))
4233                .collect()
4234        } else {
4235            Vec::new()
4236        };
4237        let mut offs = Vec::with_capacity(groups.len());
4238        let mut o = 0usize;
4239        for g in groups {
4240            offs.push(o);
4241            o += g.len();
4242        }
4243        let ptrs: Vec<SendMut> = outs.iter_mut().map(|o| SendMut(o.as_mut_ptr())).collect();
4244        let (views, ptrs, acts, offs) = (&views, &ptrs, &acts, &offs);
4245        let run = |start: usize, end: usize| {
4246            let (mut gsc, mut usc) = (vec![0f32; gpr], vec![0f32; gpr]);
4247            for flat in start..end {
4248                let (e, r) = (flat / inter, flat % inter);
4249                let (gv_view, uv_view) = (&views[e * 2], &views[e * 2 + 1]);
4250                gv_view.scales_into(r, gpr, &mut gsc);
4251                uv_view.scales_into(r, gpr, &mut usc);
4252                for (k, &t) in groups[e].iter().enumerate() {
4253                    if !quantized {
4254                        let x = &xs[t * cols..(t + 1) * cols];
4255                        let gv = q4tp_row_exact(gv_view.nib, r, gpr, x, &gsc);
4256                        let uv = q4tp_row_exact(uv_view.nib, r, gpr, x, &usc);
4257                        unsafe { *ptrs[offs[e] + k].at(r) = (gv / (1.0 + (-gv).exp())) * uv };
4258                        continue;
4259                    }
4260                    let act = &acts[t];
4261                    let mut gv = dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, &gsc) * act.sx;
4262                    let mut uv = dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, &usc) * act.sx;
4263                    for &(j, xv) in &act.outliers {
4264                        let og = q4tp_outlier(gv_view.nib, r, gpr, j, &gsc);
4265                        let ou = q4tp_outlier(uv_view.nib, r, gpr, j, &usc);
4266                        gv += og.0 * og.1 * xv;
4267                        uv += ou.0 * ou.1 * xv;
4268                    }
4269                    let silu_g = gv / (1.0 + (-gv).exp());
4270                    // SAFETY: one worker owns each (expert, row) cell of
4271                    // every output of the expert's group.
4272                    unsafe { *ptrs[offs[e] + k].at(r) = silu_g * uv };
4273                }
4274            }
4275        };
4276        dispatch_rows(pool, pairs.len() * inter, &run);
4277        true
4278    }
4279
4280    /// The per-(expert, token) down terms `moe_down_many` weights and sums,
4281    /// for SEVERAL tokens: `outs[p][o] = down_e[o] · gs[p]` (int8 split of
4282    /// `gs[p]`, VNNI dot, outlier terms — bit-identical to that kernel's
4283    /// `d`), each down row read once for its expert's whole group. The
4284    /// caller sums `w·d` per token in its route order, which reproduces
4285    /// `moe_down_many`'s f32 sequence exactly. Layout as `moe_gate_up_rows`.
4286    pub fn moe_down_rows(
4287        downs: &[&QTensor],
4288        group_lens: &[usize],
4289        gs: &[Vec<f32>],
4290        outs: &mut [Vec<f32>],
4291        pool: Option<&Pool>,
4292    ) -> bool {
4293        if Self::any_flint_exact_cpu_default(downs) {
4294            return float_activations_scope(|| {
4295                Self::moe_down_rows(downs, group_lens, gs, outs, pool)
4296            });
4297        }
4298        if downs.is_empty() || downs.len() != group_lens.len() {
4299            return false;
4300        }
4301        if downs.iter().any(|d| d.per_tensor_only()) {
4302            return false;
4303        }
4304        let rows = downs[0].rows();
4305        let cols = downs[0].cols();
4306        let n_pairs: usize = group_lens.iter().sum();
4307        if cols == 0 || cols % GROUP_SIZE != 0 || gs.len() != n_pairs || outs.len() != n_pairs {
4308            return false;
4309        }
4310        let gpr = cols / GROUP_SIZE;
4311        let mut views = Vec::with_capacity(downs.len());
4312        for d in downs {
4313            if d.has_prism_contract()
4314                || !matches!(
4315                    d,
4316                    Self::Mapped {
4317                        dtype: TensorDtype::Q4TiledP,
4318                        ..
4319                    }
4320                )
4321                || d.rows() != rows
4322                || d.cols() != cols
4323            {
4324                return false;
4325            }
4326            views.push(Q4tpView::new(d.quant_bytes(), rows, cols));
4327        }
4328        if gs.iter().any(|g| g.len() != cols) || outs.iter().any(|o| o.len() != rows) {
4329            return false;
4330        }
4331        let quantized = a8w8_enabled();
4332        let acts: Vec<SplitAct> = if quantized {
4333            gs.iter().map(|g| split_act(g)).collect()
4334        } else {
4335            Vec::new()
4336        };
4337        let mut offs = Vec::with_capacity(group_lens.len());
4338        let mut o = 0usize;
4339        for &l in group_lens {
4340            offs.push(o);
4341            o += l;
4342        }
4343        let ptrs: Vec<SendMut> = outs.iter_mut().map(|o| SendMut(o.as_mut_ptr())).collect();
4344        let (views, ptrs, acts, offs) = (&views, &ptrs, &acts, &offs);
4345        let run = |start: usize, end: usize| {
4346            let mut sc = vec![0f32; gpr];
4347            for flat in start..end {
4348                let (e, r) = (flat / rows, flat % rows);
4349                let v = &views[e];
4350                v.scales_into(r, gpr, &mut sc);
4351                for k in 0..group_lens[e] {
4352                    if !quantized {
4353                        let d = q4tp_row_exact(v.nib, r, gpr, &gs[offs[e] + k], &sc);
4354                        unsafe { *ptrs[offs[e] + k].at(r) = d };
4355                        continue;
4356                    }
4357                    let a = &acts[offs[e] + k];
4358                    let mut d = dot_q4tp_row_i8(v.nib, r, gpr, &a.xq, &sc) * a.sx;
4359                    for &(j, xv) in &a.outliers {
4360                        let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
4361                        d += w * s * xv;
4362                    }
4363                    // SAFETY: one worker owns each (expert, row) cell.
4364                    unsafe { *ptrs[offs[e] + k].at(r) = d };
4365                }
4366            }
4367        };
4368        dispatch_rows(pool, downs.len() * rows, &run);
4369        true
4370    }
4371}
4372
4373/// Batched q8 kernel: same math as qmatvec, the row makes a single
4374/// pass from memory for the whole batch.
4375/// Accelerate CBLAS — the Apple AMX matrix units, the same engine
4376/// llama.cpp's `-ngl 0` prefill rides via ggml-blas.
4377#[cfg(target_os = "macos")]
4378mod accel_blas {
4379    #[link(name = "Accelerate", kind = "framework")]
4380    unsafe extern "C" {
4381        pub fn cblas_sgemm(
4382            order: i32,
4383            trans_a: i32,
4384            trans_b: i32,
4385            m: i32,
4386            n: i32,
4387            k: i32,
4388            alpha: f32,
4389            a: *const f32,
4390            lda: i32,
4391            b: *const f32,
4392            ldb: i32,
4393            beta: f32,
4394            c: *mut f32,
4395            ldc: i32,
4396        );
4397    }
4398}
4399
4400#[cfg(target_os = "macos")]
4401pub(crate) fn accel_gemm_enabled() -> bool {
4402    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4403    *ON.get_or_init(|| std::env::var("CMF_ACCEL").map(|v| v != "0").unwrap_or(true))
4404}
4405
4406/// Off macOS the "accel" GEMM is the portable NEON micro-kernel below —
4407/// same entry point, so the batched-attention path opens on mobile.
4408#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
4409pub(crate) fn accel_gemm_enabled() -> bool {
4410    true
4411}
4412
4413/// Portable NEON f32 GEMM (row-major, optional Bᵀ): a 4×8 fmla
4414/// micro-kernel with A broadcast against B panels — the mobile stand-in
4415/// for Accelerate in the batched causal attention (QKᵀ and P·V). Not a
4416/// BLAS: shapes here are the attention panels (m ≤ heads·chunk,
4417/// k = head_dim or context), and the goal is removing the per-position
4418/// quadratic wall, not peak GEMM.
4419#[cfg(target_arch = "aarch64")]
4420#[allow(clippy::too_many_arguments)]
4421pub(crate) fn neon_gemm_rm(
4422    m: usize,
4423    n: usize,
4424    k: usize,
4425    alpha: f32,
4426    a: &[f32],
4427    lda: usize,
4428    b_mat: &[f32],
4429    ldb: usize,
4430    b_rows_are_n: bool,
4431    c: &mut [f32],
4432    ldc: usize,
4433) {
4434    debug_assert!(a.len() >= (m - 1) * lda + k);
4435    debug_assert!(c.len() >= (m - 1) * ldc + n);
4436    // SAFETY: bounds asserted above; NEON is baseline on aarch64.
4437    unsafe {
4438        use core::arch::aarch64::*;
4439        let mut i = 0usize;
4440        while i < m {
4441            let mi = (m - i).min(4);
4442            let mut j = 0usize;
4443            while j < n {
4444                let nj = (n - j).min(8);
4445                if mi == 4 && nj == 8 {
4446                    let (mut c0a, mut c0b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
4447                    let (mut c1a, mut c1b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
4448                    let (mut c2a, mut c2b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
4449                    let (mut c3a, mut c3b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
4450                    for p in 0..k {
4451                        let (b0, b1) = if b_rows_are_n {
4452                            // B is [n, k]: column p of Bᵀ = element p of
4453                            // eight consecutive B rows — gathered.
4454                            let base = b_mat.as_ptr().add(j * ldb + p);
4455                            let g = |o: usize| *base.add(o * ldb);
4456                            ([g(0), g(1), g(2), g(3)], [g(4), g(5), g(6), g(7)])
4457                        } else {
4458                            let base = b_mat.as_ptr().add(p * ldb + j);
4459                            (
4460                                [*base, *base.add(1), *base.add(2), *base.add(3)],
4461                                [*base.add(4), *base.add(5), *base.add(6), *base.add(7)],
4462                            )
4463                        };
4464                        let bv0 = vld1q_f32(b0.as_ptr());
4465                        let bv1 = vld1q_f32(b1.as_ptr());
4466                        let a0 = vdupq_n_f32(*a.as_ptr().add(i * lda + p));
4467                        let a1 = vdupq_n_f32(*a.as_ptr().add((i + 1) * lda + p));
4468                        let a2 = vdupq_n_f32(*a.as_ptr().add((i + 2) * lda + p));
4469                        let a3 = vdupq_n_f32(*a.as_ptr().add((i + 3) * lda + p));
4470                        c0a = vfmaq_f32(c0a, a0, bv0);
4471                        c0b = vfmaq_f32(c0b, a0, bv1);
4472                        c1a = vfmaq_f32(c1a, a1, bv0);
4473                        c1b = vfmaq_f32(c1b, a1, bv1);
4474                        c2a = vfmaq_f32(c2a, a2, bv0);
4475                        c2b = vfmaq_f32(c2b, a2, bv1);
4476                        c3a = vfmaq_f32(c3a, a3, bv0);
4477                        c3b = vfmaq_f32(c3b, a3, bv1);
4478                    }
4479                    let al = vdupq_n_f32(alpha);
4480                    for (r, (ca, cb)) in [(c0a, c0b), (c1a, c1b), (c2a, c2b), (c3a, c3b)]
4481                        .iter()
4482                        .enumerate()
4483                    {
4484                        let dst = c.as_mut_ptr().add((i + r) * ldc + j);
4485                        vst1q_f32(dst, vmulq_f32(*ca, al));
4486                        vst1q_f32(dst.add(4), vmulq_f32(*cb, al));
4487                    }
4488                } else {
4489                    for r in 0..mi {
4490                        for q in 0..nj {
4491                            let mut acc = 0f32;
4492                            for p in 0..k {
4493                                let bv = if b_rows_are_n {
4494                                    b_mat[(j + q) * ldb + p]
4495                                } else {
4496                                    b_mat[p * ldb + j + q]
4497                                };
4498                                acc += a[(i + r) * lda + p] * bv;
4499                            }
4500                            c[(i + r) * ldc + j + q] = acc * alpha;
4501                        }
4502                    }
4503                }
4504                j += nj;
4505            }
4506            i += mi;
4507        }
4508    }
4509}
4510
4511/// Off-macOS aarch64: the batched attention rides the NEON micro-GEMM.
4512#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
4513#[allow(clippy::too_many_arguments)]
4514pub(crate) fn sgemm_rm(
4515    m: usize,
4516    n: usize,
4517    k: usize,
4518    alpha: f32,
4519    a: &[f32],
4520    lda: usize,
4521    b_mat: &[f32],
4522    ldb: usize,
4523    b_rows_are_n: bool,
4524    c: &mut [f32],
4525    ldc: usize,
4526) {
4527    neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
4528}
4529
4530/// Row-major f32 GEMM, exposed for offline tools (the AWNP pass builds a
4531/// per-layer projection and applies it to every expert; a naive triple loop
4532/// would turn a two-minute job into half an hour).
4533#[allow(clippy::too_many_arguments)]
4534pub fn sgemm_public(
4535    m: usize,
4536    n: usize,
4537    k: usize,
4538    alpha: f32,
4539    a: &[f32],
4540    lda: usize,
4541    b_mat: &[f32],
4542    ldb: usize,
4543    b_rows_are_n: bool,
4544    c: &mut [f32],
4545    ldc: usize,
4546) {
4547    #[cfg(any(target_os = "macos", target_arch = "aarch64"))]
4548    {
4549        sgemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
4550    }
4551    // x86 without Accelerate has no sgemm_rm (the specialized paths there are
4552    // quantized kernels): the packed AVX/FMA/AVX-512 kernels of
4553    // matrixmultiply, single-threaded — offline callers split rows across
4554    // their own threads.
4555    #[cfg(not(any(target_os = "macos", target_arch = "aarch64")))]
4556    {
4557        if m == 0 || n == 0 {
4558            return;
4559        }
4560        assert!(a.len() >= (m - 1) * lda + k && c.len() >= (m - 1) * ldc + n);
4561        let (rsb, csb) = if b_rows_are_n { (1, ldb as isize) } else { (ldb as isize, 1) };
4562        assert!(k == 0 || b_mat.len() > if b_rows_are_n { (n - 1) * ldb + k - 1 } else { (k - 1) * ldb + n - 1 });
4563        // SAFETY: the asserts above bound every element the strides reach.
4564        unsafe {
4565            matrixmultiply::sgemm(
4566                m,
4567                k,
4568                n,
4569                alpha,
4570                a.as_ptr(),
4571                lda as isize,
4572                1,
4573                b_mat.as_ptr(),
4574                rsb,
4575                csb,
4576                0.0,
4577                c.as_mut_ptr(),
4578                ldc as isize,
4579                1,
4580            );
4581        }
4582    }
4583}
4584
4585/// Row-major f32 GEMM on Accelerate: C[m,n] = alpha·A[m,k] × B(ᵀ).
4586/// `b_rows_are_n` = true multiplies by Bᵀ where B is stored [n, k].
4587#[cfg(target_os = "macos")]
4588#[allow(clippy::too_many_arguments)]
4589pub(crate) fn sgemm_rm(
4590    m: usize,
4591    n: usize,
4592    k: usize,
4593    alpha: f32,
4594    a: &[f32],
4595    lda: usize,
4596    b_mat: &[f32],
4597    ldb: usize,
4598    b_rows_are_n: bool,
4599    c: &mut [f32],
4600    ldc: usize,
4601) {
4602    debug_assert!(a.len() >= (m - 1) * lda + k);
4603    debug_assert!(c.len() >= (m - 1) * ldc + n);
4604    // Test hook: route the attention GEMMs through the portable NEON
4605    // micro-kernel ON APPLE SILICON — how the mobile batched attend is
4606    // measured without a phone in the loop. (Intel macOS has no NEON —
4607    // the hook is a no-op there, Accelerate continues below.)
4608    #[cfg(target_arch = "aarch64")]
4609    if std::env::var("CMF_FORCE_NEON_GEMM")
4610        .map(|v| v == "1")
4611        .unwrap_or(false)
4612    {
4613        return neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
4614    }
4615    unsafe {
4616        accel_blas::cblas_sgemm(
4617            101, // RowMajor
4618            111, // NoTrans A
4619            if b_rows_are_n { 112 } else { 111 },
4620            m as i32,
4621            n as i32,
4622            k as i32,
4623            alpha,
4624            a.as_ptr(),
4625            lda as i32,
4626            b_mat.as_ptr(),
4627            ldb as i32,
4628            0.0,
4629            c.as_mut_ptr(),
4630            ldc as i32,
4631        );
4632    }
4633}
4634
4635/// Prefill GEMM through Accelerate (macOS): dequantize q8 rows into
4636/// f32 tiles (scale folded in, pool-parallel) and multiply each tile
4637/// on the AMX with one row-major sgemm. Tiles live in cache, weights
4638/// stream once. Numerics are f32-GEMM (not the int8 dot): prefill
4639/// logits shift within f32 rounding — tolerance-class, like every
4640/// reduction-order change; decode (M=1) never takes this path.
4641#[cfg(target_os = "macos")]
4642fn qmatmat_accel(
4643    q: &[u8],
4644    row_scale: &[f32],
4645    pre: &[std::borrow::Cow<'_, [f32]>],
4646    rows: usize,
4647    cols: usize,
4648    out: &mut [f32],
4649    pool: Option<&Pool>,
4650) {
4651    // NOTE: double-buffering the dequant against the sgemm (a scoped
4652    // thread driving the pool on tile k+1 while the caller multiplies
4653    // tile k) was tried and LOST ~6%: Accelerate's sgemm is itself
4654    // multithreaded, and the dequant workers just steal its cores.
4655    const TR: usize = 2048;
4656    let b = pre.len();
4657    thread_local! {
4658        static XPANEL: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
4659        static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
4660    }
4661    XPANEL.with(|xp| {
4662        WTILE.with(|wt| {
4663            let mut xpanel = xp.borrow_mut();
4664            xpanel.clear();
4665            for x in pre {
4666                xpanel.extend_from_slice(x);
4667            }
4668            let mut wtile = wt.borrow_mut();
4669            wtile.resize(TR * cols, 0.0);
4670            let mut r0 = 0usize;
4671            while r0 < rows {
4672                let tr = TR.min(rows - r0);
4673                // Dequant the tile (scale folded) — pool-parallel.
4674                let wt_addr = SendMut(wtile.as_mut_ptr());
4675                let run = |start: usize, end: usize| {
4676                    for r in start..end {
4677                        let row = &q[(r0 + r) * cols..(r0 + r + 1) * cols];
4678                        let s = row_scale[r0 + r];
4679                        // SAFETY: workers cover disjoint r ranges.
4680                        let dst =
4681                            unsafe { std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols) };
4682                        for (d, &v) in dst.iter_mut().zip(row) {
4683                            *d = (v as i8) as f32 * s;
4684                        }
4685                    }
4686                };
4687                dispatch_rows(pool, tr, &run);
4688                // C[b, tr] (at column r0 of out[b, rows]) = X · Wtileᵀ
4689                unsafe {
4690                    accel_blas::cblas_sgemm(
4691                        101, // RowMajor
4692                        111, // NoTrans A
4693                        112, // Trans B
4694                        b as i32,
4695                        tr as i32,
4696                        cols as i32,
4697                        1.0,
4698                        xpanel.as_ptr(),
4699                        cols as i32,
4700                        wtile.as_ptr(),
4701                        cols as i32,
4702                        0.0,
4703                        out.as_mut_ptr().add(r0),
4704                        rows as i32,
4705                    );
4706                }
4707                r0 += tr;
4708            }
4709        })
4710    });
4711}
4712
4713fn qmatmat(
4714    q: &[u8],
4715    row_scale: &[f32],
4716    pre: &[std::borrow::Cow<'_, [f32]>],
4717    rows: usize,
4718    cols: usize,
4719    out: &mut [f32],
4720    pool: Option<&Pool>,
4721) {
4722    let b = pre.len();
4723    debug_assert_eq!(out.len(), b * rows);
4724    let _p = crate::cpuprof::spans::span("mm.qmatmat");
4725    // Big prefill batches ride the AMX (roadmap PR3): the row×batch
4726    // SDOT loop below peaks near the CPU's dot throughput, an order
4727    // below the matrix units. Small tensors and tiny test models stay
4728    // on the exact integer path.
4729    #[cfg(target_os = "macos")]
4730    if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
4731        qmatmat_accel(q, row_scale, pre, rows, cols, out, pool);
4732        return;
4733    }
4734    #[cfg(target_arch = "aarch64")]
4735    if sdot_enabled() {
4736        let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
4737        let out_addr = SendMut(out.as_mut_ptr());
4738        // Blocked 2×4 (mobile prefill: no AMX to fall back on — this
4739        // path IS the ARM prefill GEMM off Apple silicon).
4740        let blocked_ok = blocked_enabled();
4741        let use_i8mm = i8mm_enabled();
4742        if blocked_ok {
4743            let run = |start: usize, end: usize| {
4744                let mut o = start;
4745                while o < end {
4746                    if o + 2 <= end {
4747                        let r0 = &q[o * cols..(o + 1) * cols];
4748                        let r1 = &q[(o + 1) * cols..(o + 2) * cols];
4749                        let mut bi = 0usize;
4750                        while bi + 4 <= acts.len() {
4751                            let xs = [
4752                                acts[bi].xq.as_slice(),
4753                                acts[bi + 1].xq.as_slice(),
4754                                acts[bi + 2].xq.as_slice(),
4755                                acts[bi + 3].xq.as_slice(),
4756                            ];
4757                            let d = if use_i8mm {
4758                                unsafe { dot_i8_smmla_2x4(r0, r1, xs) }
4759                            } else {
4760                                unsafe { dot_i8_sdot_2x4(r0, r1, xs) }
4761                            };
4762                            for (r, row) in [r0, r1].into_iter().enumerate() {
4763                                for k in 0..4 {
4764                                    let act = &acts[bi + k];
4765                                    let mut v = d[r][k] as f32 * act.sx;
4766                                    for &(j, xv) in &act.outliers {
4767                                        v += (row[j] as i8) as f32 * xv;
4768                                    }
4769                                    unsafe {
4770                                        *out_addr.at((bi + k) * rows + o + r) = v * row_scale[o + r]
4771                                    };
4772                                }
4773                            }
4774                            bi += 4;
4775                        }
4776                        while bi < acts.len() {
4777                            for (r, row) in [r0, r1].into_iter().enumerate() {
4778                                let v = row_dot_sdot(row, &acts[bi]) * row_scale[o + r];
4779                                unsafe { *out_addr.at(bi * rows + o + r) = v };
4780                            }
4781                            bi += 1;
4782                        }
4783                        o += 2;
4784                    } else {
4785                        let row = &q[o * cols..(o + 1) * cols];
4786                        for (bi, act) in acts.iter().enumerate() {
4787                            let v = row_dot_sdot(row, act) * row_scale[o];
4788                            unsafe { *out_addr.at(bi * rows + o) = v };
4789                        }
4790                        o += 1;
4791                    }
4792                }
4793            };
4794            dispatch_rows(pool, rows, &run);
4795            return;
4796        }
4797        let run = |start: usize, end: usize| {
4798            for o in start..end {
4799                let row = &q[o * cols..(o + 1) * cols];
4800                for (bi, act) in acts.iter().enumerate() {
4801                    let v = row_dot_sdot(row, act) * row_scale[o];
4802                    unsafe { *out_addr.at(bi * rows + o) = v };
4803                }
4804            }
4805        };
4806        dispatch_rows(pool, rows, &run);
4807        return;
4808    }
4809    // x86 A8W8 batch. Non-VNNI parts take the BLOCKED 2×4 kernel
4810    // (roadmap P0: two weight rows' abs() stay in registers across four
4811    // activation streams); VNNI machines keep the per-row bias-trick
4812    // dot, which is already throughput-bound there.
4813    #[cfg(target_arch = "x86_64")]
4814    if avx2_a8w8_enabled() {
4815        let acts = {
4816            let _psa = crate::cpuprof::spans::span("mm.split_act");
4817            let a: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
4818            drop(_psa);
4819            a
4820        };
4821        let out_addr = SendMut(out.as_mut_ptr());
4822        // CMF_X86_BLOCKED=0 forces the per-row path (paired in-process
4823        // A/B on noisy shared-vCPU hosts).
4824        let blocked_ok = blocked_enabled();
4825        if !avx512vnni_enabled() && blocked_ok && !row_exact() {
4826            let run = |start: usize, end: usize| {
4827                let mut o = start;
4828                while o < end {
4829                    if o + 2 <= end {
4830                        let r0 = &q[o * cols..(o + 1) * cols];
4831                        let r1 = &q[(o + 1) * cols..(o + 2) * cols];
4832                        let mut bi = 0usize;
4833                        while bi + 4 <= acts.len() {
4834                            let xs = [
4835                                acts[bi].xq.as_slice(),
4836                                acts[bi + 1].xq.as_slice(),
4837                                acts[bi + 2].xq.as_slice(),
4838                                acts[bi + 3].xq.as_slice(),
4839                            ];
4840                            let d = unsafe { dot_i8_i8_avx2_2x4(r0, r1, xs) };
4841                            for (r, row) in [r0, r1].into_iter().enumerate() {
4842                                for k in 0..4 {
4843                                    let act = &acts[bi + k];
4844                                    let mut v = d[r][k] as f32 * act.sx;
4845                                    for &(j, xv) in &act.outliers {
4846                                        v += (row[j] as i8) as f32 * xv;
4847                                    }
4848                                    unsafe {
4849                                        *out_addr.at((bi + k) * rows + o + r) = v * row_scale[o + r]
4850                                    };
4851                                }
4852                            }
4853                            bi += 4;
4854                        }
4855                        while bi < acts.len() {
4856                            for (r, row) in [r0, r1].into_iter().enumerate() {
4857                                let v = row_dot_avx2(row, &acts[bi]) * row_scale[o + r];
4858                                unsafe { *out_addr.at(bi * rows + o + r) = v };
4859                            }
4860                            bi += 1;
4861                        }
4862                        o += 2;
4863                    } else {
4864                        let row = &q[o * cols..(o + 1) * cols];
4865                        for (bi, act) in acts.iter().enumerate() {
4866                            let v = row_dot_avx2(row, act) * row_scale[o];
4867                            unsafe { *out_addr.at(bi * rows + o) = v };
4868                        }
4869                        o += 1;
4870                    }
4871                }
4872            };
4873            dispatch_rows(pool, rows, &run);
4874            return;
4875        }
4876        let run = |start: usize, end: usize| {
4877            for o in start..end {
4878                let row = &q[o * cols..(o + 1) * cols];
4879                for (bi, act) in acts.iter().enumerate() {
4880                    let v = row_dot_avx2(row, act) * row_scale[o];
4881                    unsafe { *out_addr.at(bi * rows + o) = v };
4882                }
4883            }
4884        };
4885        dispatch_rows(pool, rows, &run);
4886        return;
4887    }
4888    let out_addr = SendMut(out.as_mut_ptr());
4889    let run = |start: usize, end: usize| {
4890        for o in start..end {
4891            let row = &q[o * cols..(o + 1) * cols];
4892            for (bi, x) in pre.iter().enumerate() {
4893                let mut acc = 0f32;
4894                for j in 0..cols {
4895                    acc += (row[j] as i8) as f32 * x[j];
4896                }
4897                unsafe { *out_addr.at(bi * rows + o) = acc * row_scale[o] };
4898            }
4899        }
4900    };
4901    dispatch_rows(pool, rows, &run);
4902}
4903
4904/// Split rows across pool workers (shared qmatvec pattern). Self-balancing
4905/// — see `Pool::run_rows` for why a static 1/n split is wrong here.
4906fn dispatch_rows(pool: Option<&Pool>, rows: usize, run: &(dyn Fn(usize, usize) + Sync)) {
4907    match pool {
4908        Some(pool) if rows >= 256 => pool.run_rows(rows, run),
4909        _ => run(0, rows),
4910    }
4911}
4912
4913/// Split a q4_block blob into (packed nibbles, f16 group scales).
4914fn q4_split(bytes: &[u8], rows: usize, cols: usize) -> (&[u8], &[u8]) {
4915    let groups = rows * cols / GROUP_SIZE;
4916    bytes.split_at(groups * 16)
4917}
4918
4919/// SIMD unpack for the dominant vbit width B=4 (94% of rows on the
4920/// log2-shape calibration): 16 packed bytes -> 32 centered i8 values.
4921/// vbit packs MSB-first, so the HIGH nibble is the even element
4922/// (opposite of q4_block's lo-first interleave). Centering is u-7.
4923#[inline]
4924fn vbit_fill4(data: &[u8], buf: &mut [u8]) {
4925    #[cfg(target_arch = "aarch64")]
4926    unsafe {
4927        return vbit_fill4_neon(data, buf);
4928    }
4929    #[cfg(target_arch = "x86_64")]
4930    if avx2_enabled() {
4931        return unsafe { vbit_fill4_avx2(data, buf) };
4932    }
4933    #[allow(unreachable_code)]
4934    for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
4935        let u = unpack8::<4>(&data[blk * 4..]);
4936        for k in 0..8 {
4937            chunk[k] = (u[k] - 7) as i8 as u8;
4938        }
4939    }
4940}
4941
4942#[cfg(target_arch = "aarch64")]
4943#[target_feature(enable = "neon")]
4944unsafe fn vbit_fill4_neon(data: &[u8], buf: &mut [u8]) {
4945    // SAFETY: buf.len() is a multiple of GROUP_SIZE=32; data holds
4946    // buf.len()/2 packed bytes (validated at load).
4947    unsafe {
4948        use core::arch::aarch64::*;
4949        let n = buf.len();
4950        let mask = vdupq_n_u8(0x0F);
4951        let seven = vdupq_n_s8(7);
4952        let mut g = 0usize;
4953        while g * 32 + 32 <= n {
4954            let b = vld1q_u8(data.as_ptr().add(g * 16));
4955            let hi = vshrq_n_u8::<4>(b);
4956            let lo = vandq_u8(b, mask);
4957            let z0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(hi, lo)), seven);
4958            let z1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(hi, lo)), seven);
4959            vst1q_u8(buf.as_mut_ptr().add(g * 32), vreinterpretq_u8_s8(z0));
4960            vst1q_u8(buf.as_mut_ptr().add(g * 32 + 16), vreinterpretq_u8_s8(z1));
4961            g += 1;
4962        }
4963    }
4964}
4965
4966#[cfg(target_arch = "x86_64")]
4967#[target_feature(enable = "avx2")]
4968unsafe fn vbit_fill4_avx2(data: &[u8], buf: &mut [u8]) {
4969    // SAFETY: see vbit_fill4_neon.
4970    unsafe {
4971        use core::arch::x86_64::*;
4972        let n = buf.len();
4973        let mask = _mm_set1_epi8(0x0F);
4974        let seven = _mm256_set1_epi8(7);
4975        let mut g = 0usize;
4976        while g * 32 + 32 <= n {
4977            let b = _mm_loadu_si128(data.as_ptr().add(g * 16) as *const __m128i);
4978            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), mask);
4979            let lo = _mm_and_si128(b, mask);
4980            let z = _mm256_sub_epi8(
4981                _mm256_set_m128i(_mm_unpackhi_epi8(hi, lo), _mm_unpacklo_epi8(hi, lo)),
4982                seven,
4983            );
4984            _mm256_storeu_si256(buf.as_mut_ptr().add(g * 32) as *mut __m256i, z);
4985            g += 1;
4986        }
4987    }
4988}
4989
4990/// Unpack 8 MSB-first B-bit values from exactly B bytes (fixed shifts —
4991/// no serial bit-buffer, auto-vectorizable). Every 32-value group starts
4992/// byte-aligned (32·B/8 is integral for B∈3..8), so groups decompose
4993/// into 4 such blocks.
4994#[inline(always)]
4995fn unpack8<const B: usize>(data: &[u8]) -> [i32; 8] {
4996    let mut acc = 0u64;
4997    for i in 0..B {
4998        acc = (acc << 8) | data[i] as u64;
4999    }
5000    let mask = (1u64 << B) - 1;
5001    let mut out = [0i32; 8];
5002    for (k, o) in out.iter_mut().enumerate() {
5003        *o = ((acc >> ((7 - k) * B)) & mask) as i32;
5004    }
5005    out
5006}
5007
5008/// Fused vbit matvec straight from the mapped bytes (spec §3, P13
5009/// FIG.3): [u8 bits: rows][f16 scales: rows·cols/32][bit-packed rows,
5010/// MSB-first, byte-padded]. Row data offsets are precomputed at load
5011/// (`vbit_row_offsets`) — the per-call prefix scan was O(rows) pure
5012/// overhead on every matvec.
5013#[allow(clippy::too_many_arguments)]
5014fn vbitmatvec(
5015    bytes: &[u8],
5016    offsets: &[usize],
5017    x: &[f32],
5018    rows: usize,
5019    cols: usize,
5020    out: &mut [f32],
5021    pool: Option<&Pool>,
5022) {
5023    debug_assert_eq!(out.len(), rows);
5024    debug_assert_eq!(offsets.len(), rows + 1);
5025
5026    // SDOT path: unpack the row to centered i8 once, then per-group
5027    // int8 dot against the quantized activations — same A8W8 contract
5028    // as q8 (bounded noise; CMF_SDOT=0 keeps the exact scalar path).
5029    if a8w8_enabled() {
5030        let act = split_act(x);
5031        let out_addr = SendMut(out.as_mut_ptr());
5032        let run = move |start: usize, end: usize| {
5033            vbit_range_a8w8(bytes, offsets, x, &act, rows, cols, out_addr, start, end)
5034        };
5035        dispatch_rows(pool, rows, &run);
5036        return;
5037    }
5038
5039    let out_addr = SendMut(out.as_mut_ptr());
5040    let run = move |start: usize, end: usize| {
5041        vbit_range_f32(bytes, offsets, x, rows, cols, out_addr, start, end)
5042    };
5043    dispatch_rows(pool, rows, &run);
5044}
5045
5046/// One vbit row range via the A8W8 int8 path — kernel body of
5047/// `vbitmatvec`, extracted so multi-matrix jobs can drive it for
5048/// several tensors in one dispatch (b=8 rows go exact f32).
5049#[allow(clippy::too_many_arguments)]
5050fn vbit_range_a8w8(
5051    bytes: &[u8],
5052    offsets: &[usize],
5053    x: &[f32],
5054    act: &SplitAct,
5055    rows: usize,
5056    cols: usize,
5057    out: SendMut,
5058    start: usize,
5059    end: usize,
5060) {
5061    let ng = cols / GROUP_SIZE;
5062    let bits = &bytes[..rows];
5063    let sc_off = rows;
5064    let row_dot = |r: usize| -> f32 {
5065        let b = bits[r] as usize;
5066        let l = (1i32 << (b - 1)) - 1;
5067        let mask = (1u64 << b) - 1;
5068        let data = &bytes[offsets[r]..offsets[r + 1]];
5069        if b == 8 {
5070            // u−L reaches 128 → does not fit i8; exact f32 path.
5071            let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
5072            let mut dot = 0f32;
5073            for g in 0..ng {
5074                let so = (r * ng + g) * 2;
5075                let sgf = f16_to_f32(u16::from_le_bytes([
5076                    bytes[sc_off + so],
5077                    bytes[sc_off + so + 1],
5078                ]));
5079                let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
5080                let mut gd = 0f32;
5081                for &xv in xg.iter() {
5082                    if nbits < 8 {
5083                        acc = (acc << 8) | data[idx] as u64;
5084                        idx += 1;
5085                        nbits += 8;
5086                    }
5087                    let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
5088                    nbits -= 8;
5089                    gd += (u - l) as f32 * xv;
5090                }
5091                dot += gd * sgf;
5092            }
5093            return dot;
5094        }
5095        // Per-worker scratch: this closure runs for every row of the
5096        // tensor (lm_head ≈ 150k rows/token) — a heap allocation per
5097        // row was measurable pure overhead.
5098        thread_local! {
5099            static VBIT_SCRATCH: std::cell::RefCell<Vec<u8>> =
5100                const { std::cell::RefCell::new(Vec::new()) };
5101        }
5102        #[inline(always)]
5103        fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
5104            for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
5105                let u = unpack8::<B>(&data[blk * B..]);
5106                for k in 0..8 {
5107                    chunk[k] = (u[k] - l) as i8 as u8;
5108                }
5109            }
5110        }
5111        let _ = mask;
5112        VBIT_SCRATCH.with(|scratch| {
5113            let mut buf = scratch.borrow_mut();
5114            buf.resize(cols, 0);
5115            match b {
5116                3 => fill::<3>(data, l, &mut buf),
5117                4 => vbit_fill4(data, &mut buf),
5118                5 => fill::<5>(data, l, &mut buf),
5119                6 => fill::<6>(data, l, &mut buf),
5120                _ => unreachable!(),
5121            }
5122            let mut dot = 0f32;
5123            for g in 0..ng {
5124                let so = (r * ng + g) * 2;
5125                let s = f16_to_f32(u16::from_le_bytes([
5126                    bytes[sc_off + so],
5127                    bytes[sc_off + so + 1],
5128                ]));
5129                let d = dot_i8_i8(
5130                    &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
5131                    &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
5132                ) as f32
5133                    * act.sx;
5134                dot += d * s;
5135            }
5136            for &(j, xv) in &act.outliers {
5137                let so = (r * ng + j / GROUP_SIZE) * 2;
5138                let s = f16_to_f32(u16::from_le_bytes([
5139                    bytes[sc_off + so],
5140                    bytes[sc_off + so + 1],
5141                ]));
5142                // xq is zeroed at outlier slots — add the exact term.
5143                dot += (buf[j] as i8) as f32 * s * xv;
5144            }
5145            dot
5146        })
5147    };
5148    for r in start..end {
5149        // SAFETY: disjoint row ranges per worker.
5150        unsafe { *out.at(r) = row_dot(r) };
5151    }
5152}
5153
5154/// Exact scalar vbit row range (same extraction, non-SDOT path).
5155#[allow(clippy::too_many_arguments)]
5156fn vbit_range_f32(
5157    bytes: &[u8],
5158    offsets: &[usize],
5159    x: &[f32],
5160    rows: usize,
5161    cols: usize,
5162    out: SendMut,
5163    start: usize,
5164    end: usize,
5165) {
5166    let ng = cols / GROUP_SIZE;
5167    let bits = &bytes[..rows];
5168    let sc_off = rows;
5169    // Per-bit-width specialized inner loops: the compiler unrolls the
5170    // constant shifts (the generic bit-buffer loop was branch-bound —
5171    // 5.6 vs 13.2 tok/s q4 on the 0.8B).
5172    #[inline(always)]
5173    fn dot_row<const B: usize>(
5174        data: &[u8],
5175        bytes: &[u8],
5176        sc_off: usize,
5177        r: usize,
5178        ng: usize,
5179        x: &[f32],
5180    ) -> f32 {
5181        let l = ((1i32 << (B - 1)) - 1) as f32;
5182        let gbytes = GROUP_SIZE * B / 8;
5183        let mut dot = 0f32;
5184        for g in 0..ng {
5185            let so = (r * ng + g) * 2;
5186            let s = f16_to_f32(u16::from_le_bytes([
5187                bytes[sc_off + so],
5188                bytes[sc_off + so + 1],
5189            ]));
5190            let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
5191            let gd0 = &data[g * gbytes..(g + 1) * gbytes];
5192            let mut gd = 0f32;
5193            for blk in 0..GROUP_SIZE / 8 {
5194                let u = unpack8::<B>(&gd0[blk * B..]);
5195                let xb = &xg[blk * 8..blk * 8 + 8];
5196                for k in 0..8 {
5197                    gd += (u[k] as f32 - l) * xb[k];
5198                }
5199            }
5200            dot += gd * s;
5201        }
5202        dot
5203    }
5204    for r in start..end {
5205        let data = &bytes[offsets[r]..offsets[r + 1]];
5206        let v = match bits[r] {
5207            3 => dot_row::<3>(data, bytes, sc_off, r, ng, x),
5208            4 => dot_row::<4>(data, bytes, sc_off, r, ng, x),
5209            5 => dot_row::<5>(data, bytes, sc_off, r, ng, x),
5210            6 => dot_row::<6>(data, bytes, sc_off, r, ng, x),
5211            8 => dot_row::<8>(data, bytes, sc_off, r, ng, x),
5212            b => unreachable!("vbit bit-width {b} (validated at load)"),
5213        };
5214        // SAFETY: disjoint row ranges per worker.
5215        unsafe { *out.at(r) = v };
5216    }
5217}
5218
5219/// Fused two-input vbit matvec: each row is unpacked from the mmap ONCE
5220/// and dotted against BOTH activations (MTP verify / pair prefill used
5221/// to run two full matvecs — double weight traffic and double unpack).
5222/// Per-input math is identical to `vbitmatvec` → same accuracy contract.
5223#[allow(clippy::too_many_arguments)]
5224fn vbitmatvec2(
5225    bytes: &[u8],
5226    offsets: &[usize],
5227    x1: &[f32],
5228    x2: &[f32],
5229    rows: usize,
5230    cols: usize,
5231    o1: &mut [f32],
5232    o2: &mut [f32],
5233    pool: Option<&Pool>,
5234) {
5235    debug_assert_eq!(o1.len(), rows);
5236    debug_assert_eq!(o2.len(), rows);
5237
5238    if a8w8_enabled() {
5239        let a1 = split_act(x1);
5240        let a2 = split_act(x2);
5241        let p1 = SendMut(o1.as_mut_ptr());
5242        let p2 = SendMut(o2.as_mut_ptr());
5243        let run = move |start: usize, end: usize| {
5244            vbit_range2_a8w8(
5245                bytes, offsets, x1, x2, &a1, &a2, rows, cols, p1, p2, start, end,
5246            )
5247        };
5248        dispatch_rows(pool, rows, &run);
5249        return;
5250    }
5251
5252    let p1 = SendMut(o1.as_mut_ptr());
5253    let p2 = SendMut(o2.as_mut_ptr());
5254    let run = move |start: usize, end: usize| {
5255        vbit_range2_f32(bytes, offsets, x1, x2, rows, cols, p1, p2, start, end)
5256    };
5257    dispatch_rows(pool, rows, &run);
5258}
5259
5260/// Two-input vbit row range via the A8W8 int8 path — kernel body of
5261/// `vbitmatvec2`, extracted for pair multi-matrix jobs (b=8 rows go
5262/// exact f32 for both lanes, bits streamed once).
5263#[allow(clippy::too_many_arguments)]
5264fn vbit_range2_a8w8(
5265    bytes: &[u8],
5266    offsets: &[usize],
5267    x1: &[f32],
5268    x2: &[f32],
5269    a1: &SplitAct,
5270    a2: &SplitAct,
5271    rows: usize,
5272    cols: usize,
5273    p1: SendMut,
5274    p2: SendMut,
5275    start: usize,
5276    end: usize,
5277) {
5278    let ng = cols / GROUP_SIZE;
5279    let bits = &bytes[..rows];
5280    let sc_off = rows;
5281    let row_dots = |r: usize| -> (f32, f32) {
5282        let b = bits[r] as usize;
5283        let l = (1i32 << (b - 1)) - 1;
5284        let data = &bytes[offsets[r]..offsets[r + 1]];
5285        if b == 8 {
5286            // u−L reaches 128 → does not fit i8; exact f32 path,
5287            // bits still streamed once for both lanes.
5288            let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
5289            let (mut d1, mut d2) = (0f32, 0f32);
5290            for g in 0..ng {
5291                let so = (r * ng + g) * 2;
5292                let sgf = f16_to_f32(u16::from_le_bytes([
5293                    bytes[sc_off + so],
5294                    bytes[sc_off + so + 1],
5295                ]));
5296                let (mut g1, mut g2) = (0f32, 0f32);
5297                for k in 0..GROUP_SIZE {
5298                    if nbits < 8 {
5299                        acc = (acc << 8) | data[idx] as u64;
5300                        idx += 1;
5301                        nbits += 8;
5302                    }
5303                    let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
5304                    nbits -= 8;
5305                    let w = (u - l) as f32;
5306                    g1 += w * x1[g * GROUP_SIZE + k];
5307                    g2 += w * x2[g * GROUP_SIZE + k];
5308                }
5309                d1 += g1 * sgf;
5310                d2 += g2 * sgf;
5311            }
5312            return (d1, d2);
5313        }
5314        thread_local! {
5315            static VBIT_SCRATCH2: std::cell::RefCell<Vec<u8>> =
5316                const { std::cell::RefCell::new(Vec::new()) };
5317        }
5318        #[inline(always)]
5319        fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
5320            for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
5321                let u = unpack8::<B>(&data[blk * B..]);
5322                for k in 0..8 {
5323                    chunk[k] = (u[k] - l) as i8 as u8;
5324                }
5325            }
5326        }
5327        VBIT_SCRATCH2.with(|scratch| {
5328            let mut buf = scratch.borrow_mut();
5329            buf.resize(cols, 0);
5330            match b {
5331                3 => fill::<3>(data, l, &mut buf),
5332                4 => vbit_fill4(data, &mut buf),
5333                5 => fill::<5>(data, l, &mut buf),
5334                6 => fill::<6>(data, l, &mut buf),
5335                _ => unreachable!(),
5336            }
5337            let (mut d1, mut d2) = (0f32, 0f32);
5338            for g in 0..ng {
5339                let so = (r * ng + g) * 2;
5340                let s = f16_to_f32(u16::from_le_bytes([
5341                    bytes[sc_off + so],
5342                    bytes[sc_off + so + 1],
5343                ]));
5344                let wg = &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
5345                let v1 = dot_i8_i8(wg, &a1.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a1.sx;
5346                let v2 = dot_i8_i8(wg, &a2.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a2.sx;
5347                d1 += v1 * s;
5348                d2 += v2 * s;
5349            }
5350            for &(j, xv) in &a1.outliers {
5351                let so = (r * ng + j / GROUP_SIZE) * 2;
5352                let s = f16_to_f32(u16::from_le_bytes([
5353                    bytes[sc_off + so],
5354                    bytes[sc_off + so + 1],
5355                ]));
5356                d1 += (buf[j] as i8) as f32 * s * xv;
5357            }
5358            for &(j, xv) in &a2.outliers {
5359                let so = (r * ng + j / GROUP_SIZE) * 2;
5360                let s = f16_to_f32(u16::from_le_bytes([
5361                    bytes[sc_off + so],
5362                    bytes[sc_off + so + 1],
5363                ]));
5364                d2 += (buf[j] as i8) as f32 * s * xv;
5365            }
5366            (d1, d2)
5367        })
5368    };
5369    for r in start..end {
5370        let (v1, v2) = row_dots(r);
5371        // SAFETY: disjoint row ranges per worker.
5372        unsafe {
5373            *p1.at(r) = v1;
5374            *p2.at(r) = v2;
5375        }
5376    }
5377}
5378
5379/// Two-input exact scalar vbit row range (same extraction) —
5380/// per-bit-width specialized, two accumulators per row; per-lane
5381/// accumulation order matches `vbitmatvec` exactly.
5382#[allow(clippy::too_many_arguments)]
5383fn vbit_range2_f32(
5384    bytes: &[u8],
5385    offsets: &[usize],
5386    x1: &[f32],
5387    x2: &[f32],
5388    rows: usize,
5389    cols: usize,
5390    p1: SendMut,
5391    p2: SendMut,
5392    start: usize,
5393    end: usize,
5394) {
5395    let ng = cols / GROUP_SIZE;
5396    let bits = &bytes[..rows];
5397    let sc_off = rows;
5398    #[inline(always)]
5399    #[allow(clippy::too_many_arguments)]
5400    fn dot_row2<const B: usize>(
5401        data: &[u8],
5402        bytes: &[u8],
5403        sc_off: usize,
5404        r: usize,
5405        ng: usize,
5406        x1: &[f32],
5407        x2: &[f32],
5408    ) -> (f32, f32) {
5409        let l = ((1i32 << (B - 1)) - 1) as f32;
5410        let gbytes = GROUP_SIZE * B / 8;
5411        let (mut d1, mut d2) = (0f32, 0f32);
5412        for g in 0..ng {
5413            let so = (r * ng + g) * 2;
5414            let s = f16_to_f32(u16::from_le_bytes([
5415                bytes[sc_off + so],
5416                bytes[sc_off + so + 1],
5417            ]));
5418            let x1g = &x1[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
5419            let x2g = &x2[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
5420            let gd0 = &data[g * gbytes..(g + 1) * gbytes];
5421            let (mut g1, mut g2) = (0f32, 0f32);
5422            for blk in 0..GROUP_SIZE / 8 {
5423                let u = unpack8::<B>(&gd0[blk * B..]);
5424                for k in 0..8 {
5425                    let w = u[k] as f32 - l;
5426                    g1 += w * x1g[blk * 8 + k];
5427                    g2 += w * x2g[blk * 8 + k];
5428                }
5429            }
5430            d1 += g1 * s;
5431            d2 += g2 * s;
5432        }
5433        (d1, d2)
5434    }
5435    for r in start..end {
5436        let data = &bytes[offsets[r]..offsets[r + 1]];
5437        let (v1, v2) = match bits[r] {
5438            3 => dot_row2::<3>(data, bytes, sc_off, r, ng, x1, x2),
5439            4 => dot_row2::<4>(data, bytes, sc_off, r, ng, x1, x2),
5440            5 => dot_row2::<5>(data, bytes, sc_off, r, ng, x1, x2),
5441            6 => dot_row2::<6>(data, bytes, sc_off, r, ng, x1, x2),
5442            8 => dot_row2::<8>(data, bytes, sc_off, r, ng, x1, x2),
5443            b => unreachable!("vbit bit-width {b} (validated at load)"),
5444        };
5445        // SAFETY: disjoint row ranges per worker.
5446        unsafe {
5447            *p1.at(r) = v1;
5448            *p2.at(r) = v2;
5449        }
5450    }
5451}
5452
5453// ───────────────────── q4_tiled kernels (§4.3) ─────────────────────
5454
5455/// One q4_tiled row dot on the A8W8 int8 path: per 32-group the tile
5456/// is ONE sequential read — [f16 scale][16B nibbles] — versus the two
5457/// distant streams of the split layout. Values/order identical to the
5458/// split kernels.
5459#[inline]
5460#[allow(unreachable_code)]
5461fn dot_q4t_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
5462    #[cfg(target_arch = "aarch64")]
5463    unsafe {
5464        return dot_q4t_row_sdot(bytes, r, gpr, xq);
5465    }
5466    #[cfg(target_arch = "x86_64")]
5467    unsafe {
5468        if vnni_tiles_enabled() {
5469            return dot_q4t_row_vnni(bytes, r, gpr, xq);
5470        }
5471        return dot_q4t_row_avx2(bytes, r, gpr, xq);
5472    }
5473    let mut acc = 0f32;
5474    for gi in 0..gpr {
5475        let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
5476        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
5477        let mut d = 0i32;
5478        for (k, &b) in tile[2..].iter().enumerate() {
5479            d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
5480                + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
5481        }
5482        acc += d as f32 * s;
5483    }
5484    acc
5485}
5486
5487#[cfg(target_arch = "aarch64")]
5488#[target_feature(enable = "neon,dotprod")]
5489unsafe fn dot_q4t_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
5490    // SAFETY: callers uphold slice-length contracts (18B tile per group,
5491    // xq.len() == gpr·GROUP_SIZE).
5492    unsafe {
5493        use core::arch::aarch64::*;
5494        use core::arch::asm;
5495        let lomask = vdupq_n_u8(0x0F);
5496        let eight = vdupq_n_s8(8);
5497        let mut acc = 0f32;
5498        for gi in 0..gpr {
5499            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
5500            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5501            let b = vld1q_u8(t.add(2));
5502            let lo = vandq_u8(b, lomask);
5503            let hi = vshrq_n_u8::<4>(b);
5504            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
5505            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
5506            let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
5507            let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
5508            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5509            asm!(
5510                "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
5511                "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
5512                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5513                e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
5514                options(pure, nomem, nostack),
5515            );
5516            acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
5517        }
5518        acc
5519    }
5520}
5521
5522#[cfg(target_arch = "x86_64")]
5523#[target_feature(enable = "avx2")]
5524unsafe fn dot_q4t_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
5525    // SAFETY: see dot_q4t_row_sdot.
5526    unsafe {
5527        use core::arch::x86_64::*;
5528        let lomask = _mm_set1_epi8(0x0F);
5529        let eight = _mm256_set1_epi8(8);
5530        let ones = _mm256_set1_epi16(1);
5531        let mut acc = 0f32;
5532        for gi in 0..gpr {
5533            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
5534            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5535            let b = _mm_loadu_si128(t.add(2) as *const __m128i);
5536            let lo = _mm_and_si128(b, lomask);
5537            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
5538            let w = _mm256_sub_epi8(
5539                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5540                eight,
5541            );
5542            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5543            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
5544            let d = _mm256_madd_epi16(p16, ones);
5545            let hi128 = _mm256_extracti128_si256::<1>(d);
5546            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
5547            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
5548            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
5549            acc += _mm_cvtsi128_si32(s32) as f32 * s;
5550        }
5551        acc
5552    }
5553}
5554
5555/// VNNI twin of `dot_q4t_row_avx2`: same unpack, `vpdpbusd` replaces
5556/// the maddubs+madd pair (see `dpbusd_hsum` — sums are bit-identical).
5557/// 256-bit VL encoding, so the VEX `vpsignb` stays usable.
5558#[cfg(target_arch = "x86_64")]
5559#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
5560unsafe fn dot_q4t_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
5561    // SAFETY: see dot_q4t_row_sdot.
5562    unsafe {
5563        use core::arch::x86_64::*;
5564        let lomask = _mm_set1_epi8(0x0F);
5565        let eight = _mm256_set1_epi8(8);
5566        let mut acc = 0f32;
5567        for gi in 0..gpr {
5568            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
5569            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5570            let b = _mm_loadu_si128(t.add(2) as *const __m128i);
5571            let lo = _mm_and_si128(b, lomask);
5572            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
5573            let w = _mm256_sub_epi8(
5574                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5575                eight,
5576            );
5577            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5578            let d = dpbusd_hsum(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
5579            acc += d as f32 * s;
5580        }
5581        acc
5582    }
5583}
5584
5585/// One q4_tiled row against FOUR activation streams: the nibble unpack
5586/// and abs() happen once per group instead of once per (group,
5587/// activation) — the unpack is the dominant per-element cost of the
5588/// tiled format (roadmap P0 portable blocking, q4t leg).
5589#[cfg(target_arch = "x86_64")]
5590// `fma` is NOT implied by `avx2`: without it LLVM lowers _mm256_fmadd_ps
5591// to a libm call per lane — measured 2x slower than the reduction this
5592// kernel replaces. The runtime gate (`avx2_enabled`) already requires
5593// both features, so declaring it here is safe.
5594#[target_feature(enable = "avx2,fma")]
5595unsafe fn dot_q4t_row_1x4_avx2(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
5596    // SAFETY: callers uphold the 18B-tile and xq-length contracts.
5597    unsafe {
5598        use core::arch::x86_64::*;
5599        let lomask = _mm_set1_epi8(0x0F);
5600        let eight = _mm256_set1_epi8(8);
5601        let ones = _mm256_set1_epi16(1);
5602        // One f32 accumulator VECTOR per activation, reduced once at the
5603        // end. Folding each group's i32 lanes to a scalar inside the loop
5604        // costs an extracti128 + three shift/add + a movd — a cross-lane
5605        // dependency chain per (group, activation), 288 of them per row at
5606        // cols=2304. The per-group scale is what forces a float
5607        // accumulator; it does not force a horizontal sum.
5608        //
5609        // The four accumulators are NAMED, not an array: as `[__m256; 4]`
5610        // indexed by a loop variable LLVM keeps them in memory and every
5611        // group pays four 32-byte loads and stores. That alone made this
5612        // kernel 2x SLOWER than the per-group reduction it replaces
5613        // (measured on the EPYC box: 150 s vs 71 s for two 256² steps).
5614        let mut f0 = _mm256_setzero_ps();
5615        let mut f1 = _mm256_setzero_ps();
5616        let mut f2 = _mm256_setzero_ps();
5617        let mut f3 = _mm256_setzero_ps();
5618        for gi in 0..gpr {
5619            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
5620            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5621            let sv = _mm256_set1_ps(s);
5622            let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
5623            let lo = _mm_and_si128(bb, lomask);
5624            let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
5625            let w = _mm256_sub_epi8(
5626                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5627                eight,
5628            );
5629            let aw = _mm256_abs_epi8(w);
5630            let off = gi * GROUP_SIZE;
5631            let dot = |xq: &[i8]| {
5632                let x = _mm256_loadu_si256(xq.as_ptr().add(off) as *const __m256i);
5633                let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
5634                _mm256_cvtepi32_ps(_mm256_madd_epi16(p16, ones))
5635            };
5636            f0 = _mm256_fmadd_ps(dot(xs[0]), sv, f0);
5637            f1 = _mm256_fmadd_ps(dot(xs[1]), sv, f1);
5638            f2 = _mm256_fmadd_ps(dot(xs[2]), sv, f2);
5639            f3 = _mm256_fmadd_ps(dot(xs[3]), sv, f3);
5640        }
5641        [
5642            hsum256_ps(f0),
5643            hsum256_ps(f1),
5644            hsum256_ps(f2),
5645            hsum256_ps(f3),
5646        ]
5647    }
5648}
5649
5650/// Horizontal sum of eight f32 lanes — the one cross-lane reduction the
5651/// blocked kernels pay, once per row instead of once per group.
5652#[cfg(target_arch = "x86_64")]
5653#[target_feature(enable = "avx2")]
5654#[inline]
5655unsafe fn hsum256_ps(v: core::arch::x86_64::__m256) -> f32 {
5656    // SAFETY: pure register arithmetic on the caller's vector.
5657    unsafe {
5658        use core::arch::x86_64::*;
5659        let hi = _mm256_extractf128_ps::<1>(v);
5660        let s = _mm_add_ps(_mm256_castps256_ps128(v), hi);
5661        let s = _mm_add_ps(s, _mm_movehl_ps(s, s));
5662        let s = _mm_add_ss(s, _mm_shuffle_ps::<0x55>(s, s));
5663        _mm_cvtss_f32(s)
5664    }
5665}
5666
5667/// VNNI twin of `dot_q4t_row_1x4_avx2` (see `dpbusd_hsum`).
5668#[cfg(target_arch = "x86_64")]
5669#[target_feature(enable = "avx2,fma,avx512f,avx512bw,avx512vl,avx512vnni")]
5670unsafe fn dot_q4t_row_1x4_vnni(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
5671    // SAFETY: callers uphold the 18B-tile and xq-length contracts.
5672    unsafe {
5673        use core::arch::x86_64::*;
5674        let lomask = _mm_set1_epi8(0x0F);
5675        let eight = _mm256_set1_epi8(8);
5676        // Same shape as the AVX2 twin: accumulate in f32 vectors and pay
5677        // one cross-lane reduction per row, not per (group, activation).
5678        let mut f0 = _mm256_setzero_ps();
5679        let mut f1 = _mm256_setzero_ps();
5680        let mut f2 = _mm256_setzero_ps();
5681        let mut f3 = _mm256_setzero_ps();
5682        for gi in 0..gpr {
5683            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
5684            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5685            let sv = _mm256_set1_ps(s);
5686            let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
5687            let lo = _mm_and_si128(bb, lomask);
5688            let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
5689            let w = _mm256_sub_epi8(
5690                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5691                eight,
5692            );
5693            let aw = _mm256_abs_epi8(w);
5694            let off = gi * GROUP_SIZE;
5695            let dot = |xq: &[i8]| {
5696                let x = _mm256_loadu_si256(xq.as_ptr().add(off) as *const __m256i);
5697                _mm256_cvtepi32_ps(_mm256_dpbusd_epi32(
5698                    _mm256_setzero_si256(),
5699                    aw,
5700                    _mm256_sign_epi8(x, w),
5701                ))
5702            };
5703            f0 = _mm256_fmadd_ps(dot(xs[0]), sv, f0);
5704            f1 = _mm256_fmadd_ps(dot(xs[1]), sv, f1);
5705            f2 = _mm256_fmadd_ps(dot(xs[2]), sv, f2);
5706            f3 = _mm256_fmadd_ps(dot(xs[3]), sv, f3);
5707        }
5708        let acc = [
5709            hsum256_ps(f0),
5710            hsum256_ps(f1),
5711            hsum256_ps(f2),
5712            hsum256_ps(f3),
5713        ];
5714        acc
5715    }
5716}
5717
5718/// ARM twin of `dot_q4t_row_1x4_avx2`: one nibble unpack per group
5719/// serves FOUR activation streams. Per stream the group order and f32
5720/// accumulation match `dot_q4t_row_sdot` exactly — batch == matvec
5721/// bit-for-bit.
5722#[cfg(target_arch = "aarch64")]
5723#[target_feature(enable = "neon,dotprod")]
5724unsafe fn dot_q4t_row_1x4_sdot(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
5725    // SAFETY: callers uphold the 18B-tile and xq-length contracts.
5726    unsafe {
5727        use core::arch::aarch64::*;
5728        use core::arch::asm;
5729        let lomask = vdupq_n_u8(0x0F);
5730        let eight = vdupq_n_s8(8);
5731        let mut acc = [0f32; 4];
5732        for gi in 0..gpr {
5733            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
5734            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5735            let b = vld1q_u8(t.add(2));
5736            let lo = vandq_u8(b, lomask);
5737            let hi = vshrq_n_u8::<4>(b);
5738            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
5739            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
5740            for (k, xq) in xs.iter().enumerate() {
5741                let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
5742                let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
5743                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5744                asm!(
5745                    "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
5746                    "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
5747                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5748                    e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
5749                    options(pure, nomem, nostack),
5750                );
5751                acc[k] += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
5752            }
5753        }
5754        acc
5755    }
5756}
5757
5758/// Exact-term correction for A8W8 outliers on a tiled row.
5759#[inline]
5760fn q4t_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
5761    let gi = j / GROUP_SIZE;
5762    let k = j % GROUP_SIZE;
5763    let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
5764    let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
5765    let byte = tile[2 + k / 2];
5766    let nib = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
5767    ((nib as i32 - 8) as f32, s)
5768}
5769
5770/// Exact scalar q4_tiled row (CMF_SDOT=0 contract) — same pairwise
5771/// accumulation shape as `q4_range_f32`.
5772#[inline]
5773fn q4t_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
5774    let mut acc = 0f32;
5775    for gi in 0..gpr {
5776        let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
5777        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
5778        let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
5779        let mut ga = 0f32;
5780        for (k, &b) in tile[2..].iter().enumerate() {
5781            ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
5782                + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
5783        }
5784        acc += ga * s;
5785    }
5786    acc
5787}
5788
5789/// Split view of a `q4tp` payload. The three planes are resolved once per
5790/// matvec instead of per row — `q4tp_sections` is cheap, but doing it inside
5791/// the row loop would put a division on the hot path for nothing.
5792struct Q4tpView<'a> {
5793    nib: &'a [u8],
5794    params: &'a [u8],
5795    codes: &'a [u8],
5796    stride: usize,
5797    /// q2tp reads the ladder with rung 0 = exact zero.
5798    zero_rung: bool,
5799}
5800
5801impl<'a> Q4tpView<'a> {
5802    fn new(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
5803        let (params_off, codes_off, stride) = q4tp_sections(rows, cols);
5804        Self {
5805            nib: &bytes[..params_off],
5806            params: &bytes[params_off..codes_off],
5807            codes: &bytes[codes_off..],
5808            stride,
5809            zero_rung: false,
5810        }
5811    }
5812
5813    /// The q2tp view: identical params/codes planes, 8 B weight chunks.
5814    fn new_q2(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
5815        let (params_off, codes_off, stride) = q2tp_sections(rows, cols);
5816        Self {
5817            nib: &bytes[..params_off],
5818            params: &bytes[params_off..codes_off],
5819            codes: &bytes[codes_off..],
5820            stride,
5821            zero_rung: true,
5822        }
5823    }
5824
5825    /// Expand row `r`'s per-tile scales into `out` (length `gpr`).
5826    ///
5827    /// Doing this once per row — rather than decoding a 5-bit code inside the
5828    /// tile loop — is what makes the format free at runtime. Random access to
5829    /// a packed 5-bit field costs a division, two bounds checks and a branch;
5830    /// the tile's actual work is two `sdot`s, so per-tile decoding dominated
5831    /// the kernel and cost 5x (measured: 1.4 vs 6.9 tok/s on Nanbeige-3B).
5832    /// Walking the plane sequentially with a bit accumulator is ~3 ops.
5833    /// Eight 5-bit codes are exactly five bytes, so a whole group of
5834    /// eight decodes from one little-endian word at fixed shifts. The
5835    /// bit-accumulator this replaces carried a data-dependent `while
5836    /// have < 5` refill whose branch sat in the innermost loop of every
5837    /// q4tp row; a decode profile put this function above the dot
5838    /// products it feeds. Same bitstream, same codes — just no branch
5839    /// and eight independent extractions.
5840    /// `avx2_enabled()` and `ladder_4` are compiled only for x86-64, so a
5841    /// plain `cfg!(target_arch = ...)` guard is NOT enough: `cfg!` is a
5842    /// runtime `false` on other targets, yet both arms still get compiled
5843    /// and the x86-only symbols do not exist there at all (aarch64 build:
5844    /// E0425). This associated fn IS a compile-time `cfg` — the whole
5845    /// branch, symbol references included, is dropped on non-x86 targets,
5846    /// and the non-x86 arm falls straight through to the scalar ladder.
5847    #[cfg(target_arch = "x86_64")]
5848    #[inline]
5849    fn ladder_via_avx2(&self) -> bool {
5850        avx2_enabled()
5851    }
5852    #[cfg(not(target_arch = "x86_64"))]
5853    #[inline]
5854    fn ladder_via_avx2(&self) -> bool {
5855        false
5856    }
5857
5858    #[inline]
5859    fn scales_into(&self, r: usize, gpr: usize, out: &mut [f32]) {
5860        let tab = if self.zero_rung {
5861            q2tp_ladder(self.params, r)
5862        } else if self.ladder_via_avx2() {
5863            // One lane of the 4-wide vector ladder (the other three are
5864            // duplicates) — every q4tp scale in the process then comes
5865            // from the SAME ladder bits, whichever call site decoded it.
5866            // `ladder_via_avx2()` is cfg-gated, so on non-x86 targets
5867            // this arm is never taken and never even compiled with the
5868            // x86-only symbols in scope.
5869            #[cfg(target_arch = "x86_64")]
5870            let tab = {
5871                let p = &self.params[r * 4..r * 4 + 4];
5872                let lo = f16_to_f32(u16::from_le_bytes([p[0], p[1]]));
5873                let st = f16_to_f32(u16::from_le_bytes([p[2], p[3]]));
5874                let mut tabc = [0f32; 128];
5875                unsafe { ladder_4(&[lo; 4], &[st; 4], &mut tabc) };
5876                let mut t = [0f32; 32];
5877                for (g, o) in t.iter_mut().enumerate() {
5878                    *o = tabc[g * 4];
5879                }
5880                t
5881            };
5882            // On a non-x86 target the arm is unreachable (the `cfg!`
5883            // above is false), so it decodes through the scalar ladder —
5884            // the very path this branch would have fallen back to.
5885            #[cfg(not(target_arch = "x86_64"))]
5886            let tab = q4tp_ladder(self.params, r);
5887            tab
5888        } else {
5889            q4tp_ladder(self.params, r)
5890        };
5891        let codes = &self.codes[r * self.stride..(r + 1) * self.stride];
5892        let out = &mut out[..gpr];
5893        let mut chunks = out.chunks_exact_mut(8);
5894        let mut ci = 0usize;
5895        for c in &mut chunks {
5896            let w = u64::from(codes[ci])
5897                | u64::from(codes[ci + 1]) << 8
5898                | u64::from(codes[ci + 2]) << 16
5899                | u64::from(codes[ci + 3]) << 24
5900                | u64::from(codes[ci + 4]) << 32;
5901            for (k, o) in c.iter_mut().enumerate() {
5902                *o = tab[((w >> (5 * k)) & 31) as usize];
5903            }
5904            ci += 5;
5905        }
5906        // Fewer than eight codes left: the shared total accessor, which
5907        // tolerates a 5-bit field whose spill byte is past the stride.
5908        let tail = &codes[ci..];
5909        for (k, o) in chunks.into_remainder().iter_mut().enumerate() {
5910            *o = tab[q4tp_code(tail, k)];
5911        }
5912    }
5913}
5914
5915#[inline]
5916fn dot_q4tp_row_i8(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
5917    #[cfg(target_arch = "aarch64")]
5918    unsafe {
5919        return dot_q4tp_row_sdot(nib, r, gpr, xq, scales);
5920    }
5921    #[cfg(target_arch = "x86_64")]
5922    unsafe {
5923        if vnni_tiles_enabled() {
5924            return dot_q4tp_row_vnni(nib, r, gpr, xq, scales);
5925        }
5926        return dot_q4tp_row_avx2(nib, r, gpr, xq, scales);
5927    }
5928    #[allow(unreachable_code)]
5929    {
5930        let mut acc = 0f32;
5931        for gi in 0..gpr {
5932            let tile = &nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
5933            let s = scales[gi];
5934            let mut d = 0i32;
5935            for (k, &b) in tile.iter().enumerate() {
5936                d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
5937                    + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
5938            }
5939            acc += d as f32 * s;
5940        }
5941        acc
5942    }
5943}
5944
5945/// q4tp twin of `dot_q4t_row_sdot`: identical nibble math, but the tile
5946/// stride is 16 B (no inline scale) and the scale is a ladder lookup.
5947#[cfg(target_arch = "aarch64")]
5948#[target_feature(enable = "neon,dotprod")]
5949unsafe fn dot_q4tp_row_sdot(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
5950    // SAFETY: callers uphold slice-length contracts (16B tile per group,
5951    // xq.len() == gpr·GROUP_SIZE, codes covering gpr 5-bit fields).
5952    unsafe {
5953        use core::arch::aarch64::*;
5954        use core::arch::asm;
5955        let lomask = vdupq_n_u8(0x0F);
5956        let eight = vdupq_n_s8(8);
5957        let mut acc = 0f32;
5958        for gi in 0..gpr {
5959            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5960            let s = *scales.get_unchecked(gi);
5961            let b = vld1q_u8(t);
5962            let lo = vandq_u8(b, lomask);
5963            let hi = vshrq_n_u8::<4>(b);
5964            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
5965            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
5966            let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
5967            let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
5968            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5969            asm!(
5970                "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
5971                "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
5972                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5973                e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
5974                options(pure, nomem, nostack),
5975            );
5976            acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
5977        }
5978        acc
5979    }
5980}
5981
5982#[cfg(target_arch = "x86_64")]
5983#[target_feature(enable = "avx2")]
5984unsafe fn dot_q4tp_row_avx2(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
5985    // SAFETY: see dot_q4tp_row_sdot.
5986    unsafe {
5987        use core::arch::x86_64::*;
5988        let lomask = _mm_set1_epi8(0x0F);
5989        let eight = _mm256_set1_epi8(8);
5990        let ones = _mm256_set1_epi16(1);
5991        let mut acc = 0f32;
5992        for gi in 0..gpr {
5993            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5994            let s = *scales.get_unchecked(gi);
5995            let b = _mm_loadu_si128(t as *const __m128i);
5996            let lo = _mm_and_si128(b, lomask);
5997            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
5998            let w = _mm256_sub_epi8(
5999                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
6000                eight,
6001            );
6002            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6003            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
6004            let d = _mm256_madd_epi16(p16, ones);
6005            let hi128 = _mm256_extracti128_si256::<1>(d);
6006            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
6007            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
6008            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
6009            acc += _mm_cvtsi128_si32(s32) as f32 * s;
6010        }
6011        acc
6012    }
6013}
6014
6015/// VNNI twin of `dot_q4tp_row_avx2` (see `dot_q4t_row_vnni` for why the
6016/// 256-bit VL encoding is the one to use here).
6017#[cfg(target_arch = "x86_64")]
6018#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
6019unsafe fn dot_q4tp_row_vnni(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
6020    // SAFETY: see dot_q4tp_row_sdot.
6021    unsafe {
6022        use core::arch::x86_64::*;
6023        let lomask = _mm_set1_epi8(0x0F);
6024        let eight = _mm256_set1_epi8(8);
6025        let mut acc = 0f32;
6026        for gi in 0..gpr {
6027            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
6028            let s = *scales.get_unchecked(gi);
6029            let b = _mm_loadu_si128(t as *const __m128i);
6030            let lo = _mm_and_si128(b, lomask);
6031            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
6032            let w = _mm256_sub_epi8(
6033                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
6034                eight,
6035            );
6036            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6037            acc += dpbusd_hsum(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w)) as f32 * s;
6038        }
6039        acc
6040    }
6041}
6042
6043/// `2^x` for a 4-lane vector (range-reduced degree-6 polynomial).
6044/// ≤1 ulp vs libm `exp2f` — f32's own precision rounds the polynomial
6045/// home. Used only by [`q4tp_scales_lanes`]; see the ladder note there.
6046///
6047/// # Safety
6048/// AVX2.
6049#[cfg(target_arch = "x86_64")]
6050#[target_feature(enable = "avx2")]
6051#[inline]
6052unsafe fn exp2_128(x: core::arch::x86_64::__m128) -> core::arch::x86_64::__m128 {
6053    unsafe {
6054        use core::arch::x86_64::*;
6055        // floor via trunc with a below-integer correction.
6056        let tr = _mm_cvttps_epi32(x);
6057        // Compare against cvt(tr), NOT a bit-cast: the bit pattern of a
6058        // negative i32 is a NaN and cmplt against it is always false.
6059        let below = _mm_and_si128(
6060            _mm_castps_si128(_mm_cmplt_ps(x, _mm_cvtepi32_ps(tr))),
6061            _mm_set1_epi32(1),
6062        );
6063        let n = _mm_sub_epi32(tr, below); // floor(x)
6064        let r = _mm_sub_ps(x, _mm_cvtepi32_ps(n)); // x - floor(x) ∈ [0,1)
6065        // 2^r = exp(r·ln2) ≈ Σ aₖrᵏ, aₖ = ln2ᵏ/k!, k = 1..9 (Horner).
6066        // At r ∈ [0,1) the a₁₀ term is 7.1e-9 — below f32's own ulp.
6067        // The coefficients are carried at full decimal precision on
6068        // purpose: truncating them to an f32 spelling (what
6069        // `excessive_precision` suggests) would round the constants before
6070        // they reach the AVX registers. `approx_constant` fires on a₁ =
6071        // LN_2, but the table must stay internally consistent, so both
6072        // lints are allowed for this table only.
6073        #[allow(clippy::excessive_precision, clippy::approx_constant)]
6074        let a = [
6075            6.931_471_805_599_453e-1,
6076            2.402_265_069_591_007e-1,
6077            5.550_410_866_482_158e-2,
6078            9.618_129_107_628_477e-3,
6079            1.333_355_814_642_844e-3,
6080            1.540_353_039_338_161e-4,
6081            1.525_273_380_405_984e-5,
6082            1.321_548_679_014_431e-6,
6083            1.017_808_600_923_969_6e-7,
6084        ];
6085        let mut p = _mm_set1_ps(a[8]);
6086        for &c in a[..8].iter().rev() {
6087            p = _mm_add_ps(_mm_mul_ps(p, r), _mm_set1_ps(c));
6088        }
6089        p = _mm_add_ps(_mm_mul_ps(p, r), _mm_set1_ps(1.0));
6090        // Combine: p ∈ [1,2) carries the mantissa; n the exponent.
6091        // Underflow (n < −126) → 0; our args are f16-derived, so this
6092        // only matters for pathological rows.
6093        // bits(p) already carries exponent 127; adding n<<23 (mod 2^32,
6094        // correct for negative n too) retargets it to 127+n.
6095        let ncl = _mm_min_epi32(
6096            _mm_max_epi32(n, _mm_set1_epi32(-126)),
6097            _mm_set1_epi32(127),
6098        );
6099        let mut bits = _mm_add_epi32(_mm_castps_si128(p), _mm_slli_epi32(ncl, 23));
6100        let over = _mm_cmpgt_epi32(n, _mm_set1_epi32(127));
6101        bits = _mm_or_si128(
6102            _mm_andnot_si128(over, bits),
6103            _mm_and_si128(over, _mm_set1_epi32(0x007F_8000)),
6104        );
6105        let zero = _mm_cmplt_ps(x, _mm_set1_ps(-149.0));
6106        _mm_andnot_ps(zero, _mm_castsi128_ps(bits))
6107    }
6108}
6109
6110/// Horizontal helper: the four lanes of `v` as `[f32; 4]`.
6111#[cfg(target_arch = "x86_64")]
6112#[inline]
6113unsafe fn lanes4(v: core::arch::x86_64::__m128) -> [f32; 4] {
6114    use core::arch::x86_64::*;
6115    unsafe {
6116        [
6117            _mm_cvtss_f32(v),
6118            _mm_cvtss_f32(_mm_shuffle_ps(v, v, 0x55)),
6119            _mm_cvtss_f32(_mm_shuffle_ps(v, v, 0xAA)),
6120            _mm_cvtss_f32(_mm_shuffle_ps(v, v, 0xFF)),
6121        ]
6122    }
6123}
6124
6125/// FOUR rows' 32-run q4tp ladders at once into `[group * 4 + lane]`.
6126/// The chain is `t = exp2(lo); t *= exp2(st)` — per lane this is the
6127/// scalar `q4tp_ladder` recurrence run as `vmulps`, bit-identical; only
6128/// `exp2` is replaced (≤1 ulp, see [`exp2_128`]).
6129///
6130/// # Safety
6131/// AVX2.
6132#[cfg(target_arch = "x86_64")]
6133#[target_feature(enable = "avx2")]
6134#[inline]
6135unsafe fn ladder_4(lo4: &[f32; 4], st4: &[f32; 4], tabc: &mut [f32; 128]) {
6136    unsafe {
6137        use core::arch::x86_64::*;
6138        let ratio = exp2_128(_mm_loadu_ps(st4.as_ptr()));
6139        let mut t = exp2_128(_mm_loadu_ps(lo4.as_ptr()));
6140        tabc[0..4].copy_from_slice(&lanes4(t));
6141        for c in 1..32 {
6142            t = _mm_mul_ps(t, ratio);
6143            tabc[c * 4..c * 4 + 4].copy_from_slice(&lanes4(t));
6144        }
6145    }
6146}
6147
6148/// Decode FOUR rows' q4tp group scales at once. The scalar
6149/// [`Q4tpView::scales_into`] costs ~364 cycles a row — two libm `exp2`
6150/// calls plus the 31-step serial `t *= ratio` ladder chain — and the
6151/// ladder is 17 % of every decode token. Here the chains run as one
6152/// 4-lane vector: `t[c] = t[c-1] * ratio` is `vmulps`, bit-identical to
6153/// the scalar f32 chain (and to `q4tp_ladder`'s definition), so the only
6154/// bit change is `exp2` (≤1 ulp, see [`exp2_128`]); the dequant reference
6155/// shares this function. Output layout is `[row][group]`.
6156///
6157/// Only for `Q4TiledP` views (`!zero_rung`); rows need not be
6158/// contiguous. Falls back per incomplete tail chunk.
6159///
6160/// # Safety
6161/// AVX2; `out` must hold `rows.len() * gpr` floats.
6162#[cfg(target_arch = "x86_64")]
6163#[target_feature(enable = "avx2")]
6164unsafe fn q4tp_scales_lanes(
6165    vs: &[&Q4tpView<'_>],
6166    rows: &[usize],
6167    gpr: usize,
6168    out: &mut [f32],
6169) {
6170    use core::arch::x86_64::*;
6171    debug_assert_eq!(vs.len(), rows.len());
6172    debug_assert!(out.len() >= rows.len() * gpr);
6173    let mut done = 0usize;
6174    for (ci, chunk) in rows.chunks(4).enumerate() {
6175        let vchunk = &vs[ci * 4..ci * 4 + chunk.len()];
6176        if chunk.len() < 4 || vchunk.iter().any(|v| v.zero_rung) {
6177            for (l, &r) in chunk.iter().enumerate() {
6178                vchunk[l].scales_into(r, gpr, &mut out[(done + l) * gpr..(done + l + 1) * gpr]);
6179            }
6180            done += chunk.len();
6181            continue;
6182        }
6183        let (r0, r1, r2, r3) = (chunk[0], chunk[1], chunk[2], chunk[3]);
6184        let v0 = vchunk[0];
6185        let v1 = vchunk[1];
6186        let v2 = vchunk[2];
6187        let v3 = vchunk[3];
6188        // Gather the four (lo, st) f16 pairs — 4 B per row.
6189        let mut lo = [0f32; 4];
6190        let mut st = [0f32; 4];
6191        for (l, r) in [r0, r1, r2, r3].iter().enumerate() {
6192            let o = *r * 4;
6193            let vv = vchunk[l];
6194            lo[l] = cortiq_core::quant::f16_to_f32(u16::from_le_bytes([
6195                vv.params[o],
6196                vv.params[o + 1],
6197            ]));
6198            st[l] = cortiq_core::quant::f16_to_f32(u16::from_le_bytes([
6199                vv.params[o + 2],
6200                vv.params[o + 3],
6201            ]));
6202        }
6203        let lo_v = _mm_loadu_ps(lo.as_ptr());
6204        let st_v = _mm_loadu_ps(st.as_ptr());
6205        let ratio = exp2_128(st_v);
6206        let mut t = exp2_128(lo_v);
6207        // tab layout [group * 4 + lane] — contiguous stores per step.
6208        let mut tabc = [0f32; 32 * 4];
6209        tabc[0..4].copy_from_slice(&[
6210            _mm_cvtss_f32(t),
6211            _mm_cvtss_f32(_mm_shuffle_ps(t, t, 0x55)),
6212            _mm_cvtss_f32(_mm_shuffle_ps(t, t, 0xAA)),
6213            _mm_cvtss_f32(_mm_shuffle_ps(t, t, 0xFF)),
6214        ]);
6215        for c in 1..32 {
6216            t = _mm_mul_ps(t, ratio);
6217            tabc[c * 4..c * 4 + 4].copy_from_slice(&[
6218                _mm_cvtss_f32(t),
6219                _mm_cvtss_f32(_mm_shuffle_ps(t, t, 0x55)),
6220                _mm_cvtss_f32(_mm_shuffle_ps(t, t, 0xAA)),
6221                _mm_cvtss_f32(_mm_shuffle_ps(t, t, 0xFF)),
6222            ]);
6223        }
6224        // Scalar code walks (per row), reading tabc[group * 4 + lane].
6225        let vs4 = [v0, v1, v2, v3];
6226        for (l, &r) in [r0, r1, r2, r3].iter().enumerate() {
6227            let v = vs4[l];
6228            let codes = &v.codes[r * v.stride..(r + 1) * v.stride];
6229            let dst = &mut out[(done + l) * gpr..(done + l + 1) * gpr];
6230            let mut chunks = dst.chunks_exact_mut(8);
6231            let mut ci = 0usize;
6232            for c in &mut chunks {
6233                let w = u64::from(codes[ci])
6234                    | u64::from(codes[ci + 1]) << 8
6235                    | u64::from(codes[ci + 2]) << 16
6236                    | u64::from(codes[ci + 3]) << 24
6237                    | u64::from(codes[ci + 4]) << 32;
6238                for (k, o) in c.iter_mut().enumerate() {
6239                    *o = tabc[(((w >> (5 * k)) & 31) as usize) * 4 + l];
6240                }
6241                ci += 5;
6242            }
6243            let tail = &codes[ci..];
6244            for (k, o) in chunks.into_remainder().iter_mut().enumerate() {
6245                *o = tabc[q4tp_code(tail, k) * 4 + l];
6246            }
6247        }
6248        done += 4;
6249    }
6250}
6251
6252/// Decode-lane count for the q4tp row loops: how many rows one worker
6253/// steps through with independent `acc += d * scale` chains in ONE body.
6254/// The per-group AVX2 dot is a ~17-cycle dependent chain per 32 weights;
6255/// a single row serializes on it, overlapping rows reaches the uop-issue
6256/// bound instead. Each row keeps its own group order, so every output is
6257/// bit-identical to the plain row loop. `CMF_Q4TP_LANES=1` restores it;
6258/// 2/4/8 are the supported lane counts (default 4).
6259fn q4tp_lanes() -> usize {
6260    #[cfg(not(target_arch = "x86_64"))]
6261    {
6262        1
6263    }
6264    #[cfg(target_arch = "x86_64")]
6265    {
6266        static ON: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
6267        *ON.get_or_init(|| {
6268            let n: usize = std::env::var("CMF_Q4TP_LANES")
6269                .ok()
6270                .and_then(|v| v.parse().ok())
6271                .unwrap_or(4);
6272            match n {
6273                1 => 1,
6274                2 => 2,
6275                8 => 8,
6276                _ => 4,
6277            }
6278        })
6279    }
6280}
6281
6282/// One q4tp group tile unpacked to interleaved weights in 0..=16 — the
6283/// raw nibbles, WITHOUT the −8 shift. The correction travels in
6284/// `SplitAct::xg8` instead (each vertical lane subtracts `Σx`, and the
6285/// 8-lane reduction makes that `8·Σx` overall), which lets `vpmaddubsw`
6286/// consume the unsigned codes directly and deletes the per-byte `vpaddb`
6287/// + `vpabsb`/`vpsignb` — two of them port-5 ops in the old row
6288/// kernel's innermost loop.
6289///
6290/// # Safety
6291/// AVX2; `tile` must point at 16 valid weight bytes.
6292#[cfg(target_arch = "x86_64")]
6293#[target_feature(enable = "avx2")]
6294#[inline]
6295unsafe fn q4tp_unpack_w(tile: *const u8) -> core::arch::x86_64::__m256i {
6296    unsafe {
6297        use core::arch::x86_64::*;
6298        let lomask = _mm_set1_epi8(0x0F);
6299        let b = _mm_loadu_si128(tile as *const __m128i);
6300        let lo = _mm_and_si128(b, lomask);
6301        let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
6302        _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi))
6303    }
6304}
6305
6306/// Σ over 32 columns of `w·x` as i32x8 lane partials: `vpmaddubsw`
6307/// (|w| ≤ 15, |x| ≤ 127 — pair sums ≤ 3810, no saturation) then `vpmaddwd`.
6308/// Exact in any association.
6309///
6310/// # Safety
6311/// AVX2; `xp` must point at 32 valid i8 activation bytes.
6312#[cfg(target_arch = "x86_64")]
6313#[target_feature(enable = "avx2")]
6314#[inline]
6315unsafe fn q4tp_madd_wx(
6316    w: core::arch::x86_64::__m256i,
6317    xp: *const i8,
6318) -> core::arch::x86_64::__m256i {
6319    unsafe {
6320        use core::arch::x86_64::*;
6321        let ones = _mm256_set1_epi16(1);
6322        let x = _mm256_loadu_si256(xp as *const __m256i);
6323        _mm256_madd_epi16(_mm256_maddubs_epi16(w, x), ones)
6324    }
6325}
6326
6327/// THE shared end-of-row reduction of the vertical f32 accumulator.
6328/// Every a8w8 q4tp path — the decode lane drivers and the row-exact
6329/// matmat cells — must reduce through this one function, or the
6330/// matvec↔matmat bitwise parity test breaks.
6331#[cfg(target_arch = "x86_64")]
6332#[target_feature(enable = "avx2")]
6333#[inline]
6334unsafe fn q4tp_hsum256(acc: core::arch::x86_64::__m256) -> f32 {
6335    unsafe {
6336        use core::arch::x86_64::*;
6337        let a = _mm_add_ps(_mm256_castps256_ps128(acc), _mm256_extractf128_ps::<1>(acc));
6338        let b = _mm_add_ps(a, _mm_movehl_ps(a, a));
6339        let c = _mm_add_ss(b, _mm_shuffle_ps(b, b, 0x55));
6340        _mm_cvtss_f32(c)
6341    }
6342}
6343
6344/// One (row, activation) cell computed the vertical way: per group,
6345/// `(Σ(w+8)·x − 8·Σx) as f32 · scale` accumulated lane-wise in f32x8,
6346/// one horizontal reduction at the end. Used by the lane drivers' tails
6347/// and by the row-exact matmat column loop so both land on identical bits.
6348///
6349/// # Safety
6350/// AVX2; `gpr` groups must be covered by `xq`/`sc`/`xg8`.
6351#[cfg(target_arch = "x86_64")]
6352#[target_feature(enable = "avx2")]
6353#[inline]
6354unsafe fn q4tp_cell_dot_v(
6355    nib: *const u8,
6356    r: usize,
6357    gpr: usize,
6358    xp: *const i8,
6359    sc: *const f32,
6360    xg8: *const f32,
6361) -> f32 {
6362    unsafe {
6363        use core::arch::x86_64::*;
6364        let mut acc = _mm256_setzero_ps();
6365        for gi in 0..gpr {
6366            let w = q4tp_unpack_w(nib.add((r * gpr + gi) * Q4TP_NIB));
6367            let d = q4tp_madd_wx(w, xp.add(gi * GROUP_SIZE));
6368            let t = _mm256_sub_ps(_mm256_cvtepi32_ps(d), _mm256_set1_ps(*xg8.add(gi)));
6369            acc = _mm256_add_ps(acc, _mm256_mul_ps(t, _mm256_set1_ps(*sc.add(gi))));
6370        }
6371        q4tp_hsum256(acc)
6372    }
6373}
6374
6375/// Row `r` of a row-exact q4tp matmat against one activation panel (up to
6376/// four columns at a time): each weight group is unpacked ONCE and fed to
6377/// four independent column accumulators — the cell math of
6378/// [`q4tp_cell_dot_v`], gi-ascending, so every cell equals the decode
6379/// matvec's bits while the columns overlap their latency chains and the
6380/// unpack is shared. This is the AVX2 twin of the avx512-only blocked
6381/// kernel: without it, Broadwell prefill re-decodes every row once per
6382/// column (7 tok/s measured against q8_2f's 24).
6383///
6384/// # Safety
6385/// AVX2 (caller gates on `a8w8_enabled()` + `xg8.len() == gpr`);
6386/// `out` cells for this row are this worker's alone.
6387#[cfg(target_arch = "x86_64")]
6388#[target_feature(enable = "avx2")]
6389unsafe fn q4tp_matmat_row_cols_v(
6390    v: &Q4tpView<'_>,
6391    r: usize,
6392    gpr: usize,
6393    acts: &[SplitAct],
6394    sc: &[f32],
6395    abase: usize,
6396    rows: usize,
6397    out: SendMut,
6398) {
6399    use core::arch::x86_64::*;
6400    debug_assert_eq!(sc.len(), gpr);
6401    let n = acts.len();
6402    let mut bi = 0usize;
6403    while bi + 4 <= n {
6404        let a: [&SplitAct; 4] = [&acts[bi], &acts[bi + 1], &acts[bi + 2], &acts[bi + 3]];
6405        let mut accs = [_mm256_setzero_ps(); 4];
6406        for gi in 0..gpr {
6407            let w = q4tp_unpack_w(v.nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB));
6408            let sb = _mm256_set1_ps(*sc.get_unchecked(gi));
6409            for (k, act) in a.iter().enumerate() {
6410                let xp = act.xq.as_ptr().add(gi * GROUP_SIZE);
6411                let d = q4tp_madd_wx(w, xp);
6412                let xgb = _mm256_set1_ps(*act.xg8.get_unchecked(gi));
6413                let t = _mm256_sub_ps(_mm256_cvtepi32_ps(d), xgb);
6414                accs[k] = _mm256_add_ps(accs[k], _mm256_mul_ps(t, sb));
6415            }
6416        }
6417        for k in 0..4 {
6418            let act = a[k];
6419            let mut acc = q4tp_hsum256(accs[k]) * act.sx;
6420            for &(j, xv) in &act.outliers {
6421                let (w, s) = q4tp_outlier(v.nib, r, gpr, j, sc);
6422                acc += w * s * xv;
6423            }
6424            // SAFETY: disjoint (bi, r) cells per worker.
6425            unsafe { *out.at((abase + bi + k) * rows + r) = acc };
6426        }
6427        bi += 4;
6428    }
6429    // Remainder columns (1..=3): one vertical cell each.
6430    while bi < n {
6431        let act = &acts[bi];
6432        let mut acc = q4tp_cell_dot_v(
6433            v.nib.as_ptr(),
6434            r,
6435            gpr,
6436            act.xq.as_ptr(),
6437            sc.as_ptr(),
6438            act.xg8.as_ptr(),
6439        ) * act.sx;
6440        for &(j, xv) in &act.outliers {
6441            let (w, s) = q4tp_outlier(v.nib, r, gpr, j, sc);
6442            acc += w * s * xv;
6443        }
6444        // SAFETY: disjoint (bi, r) cells per worker.
6445        unsafe { *out.at((abase + bi) * rows + r) = acc };
6446        bi += 1;
6447    }
6448}
6449
6450/// Rows `start..end` of a single-output q4tp a8w8 matvec, `L` independent
6451/// accumulator chains wide. Per row: scales decoded once, groups stepped
6452/// in ascending `gi` with `acc += d as f32 * scale` (the exact plain-row
6453/// order), then the `* sx` scale and the outlier corrections in list
6454/// order — the plain row loop's arithmetic, line for line.
6455///
6456/// # Safety
6457/// AVX2 (callers gate on `a8w8_enabled()`); `out` disjoint per row.
6458#[cfg(target_arch = "x86_64")]
6459#[target_feature(enable = "avx2")]
6460unsafe fn q4tp_matvec_lanes<const L: usize>(
6461    v: &Q4tpView<'_>,
6462    start: usize,
6463    end: usize,
6464    gpr: usize,
6465    xq: &[i8],
6466    xg8: &[f32],
6467    sx: f32,
6468    outliers: &[(usize, f32)],
6469    sc: &mut [f32],
6470    out: SendMut,
6471) {
6472    use core::arch::x86_64::*;
6473    debug_assert!(L >= 2 && sc.len() >= L * gpr && xg8.len() == gpr);
6474    crate::cpuprof::count_lane_rows((end - start) as u64);
6475    let mut base = start;
6476    while base + L <= end {
6477        let _t0 = crate::cpuprof::time(crate::cpuprof::Slot::GuScales);
6478        let mut rows = [0usize; L];
6479        for l in 0..L {
6480            rows[l] = base + l;
6481        }
6482        let v4: [&Q4tpView; L] = [v; L];
6483        q4tp_scales_lanes(&v4, &rows, gpr, sc);
6484        drop(_t0);
6485        let _t1 = crate::cpuprof::time(crate::cpuprof::Slot::GuDots);
6486        let mut accs = [_mm256_setzero_ps(); L];
6487        for gi in 0..gpr {
6488            let xp = xq.as_ptr().add(gi * GROUP_SIZE);
6489            let xgb = _mm256_set1_ps(*xg8.get_unchecked(gi));
6490            for l in 0..L {
6491                let w = q4tp_unpack_w(v.nib.as_ptr().add((rows[l] * gpr + gi) * Q4TP_NIB));
6492                let d = q4tp_madd_wx(w, xp);
6493                let t = _mm256_sub_ps(_mm256_cvtepi32_ps(d), xgb);
6494                let sb = _mm256_set1_ps(*sc.get_unchecked(l * gpr + gi));
6495                accs[l] = _mm256_add_ps(accs[l], _mm256_mul_ps(t, sb));
6496            }
6497        }
6498        drop(_t1);
6499        let _t2 = crate::cpuprof::time(crate::cpuprof::Slot::GuPost);
6500        for l in 0..L {
6501            let r = rows[l];
6502            let scl = &sc[l * gpr..(l + 1) * gpr];
6503            let mut acc = q4tp_hsum256(accs[l]) * sx;
6504            for &(j, xv) in outliers {
6505                let (w, s) = q4tp_outlier(v.nib, r, gpr, j, scl);
6506                acc += w * s * xv;
6507            }
6508            // SAFETY: one lane owns each row.
6509            unsafe { *out.at(r) = acc };
6510        }
6511        drop(_t2);
6512        base += L;
6513    }
6514    // Partial lane batch (and any tail): one vertical cell per row —
6515    // the same reduction shape as the lanes above and as the row-exact
6516    // matmat, so bits agree across every path.
6517    for r in base..end {
6518        let scl = &mut sc[..gpr];
6519        v.scales_into(r, gpr, scl);
6520        let scl: &[f32] = scl;
6521        let mut acc = q4tp_cell_dot_v(v.nib.as_ptr(), r, gpr, xq.as_ptr(), scl.as_ptr(), xg8.as_ptr())
6522            * sx;
6523        for &(j, xv) in outliers {
6524            let (w, s) = q4tp_outlier(v.nib, r, gpr, j, scl);
6525            acc += w * s * xv;
6526        }
6527        // SAFETY: disjoint row ranges per worker.
6528        unsafe { *out.at(r) = acc };
6529    }
6530}
6531
6532/// Rows `start..end` of the fused SwiGLU gate+up q4tp arm, `L` chains
6533/// wide: both projections' scales decoded per lane, both group loops
6534/// stepped in the plain arm's per-row order, then `silu_mul_limited`
6535/// over the outlier-corrected pair — bit-identical to the plain arm.
6536///
6537/// # Safety
6538/// AVX2 (callers gate on `a8w8_enabled()`); `out` disjoint per row.
6539#[cfg(target_arch = "x86_64")]
6540#[target_feature(enable = "avx2")]
6541unsafe fn q4tp_gateup_lanes<const L: usize>(
6542    gv: &Q4tpView<'_>,
6543    uv: &Q4tpView<'_>,
6544    start: usize,
6545    end: usize,
6546    gpr: usize,
6547    xq: &[i8],
6548    xg8: &[f32],
6549    sx: f32,
6550    limit: f32,
6551    outliers: &[(usize, f32)],
6552    gsc: &mut [f32],
6553    usc: &mut [f32],
6554    out: SendMut,
6555) {
6556    use core::arch::x86_64::*;
6557    debug_assert!(
6558        L >= 2 && gsc.len() >= L * gpr && usc.len() >= L * gpr && xg8.len() == gpr
6559    );
6560    crate::cpuprof::count_lane_rows((end - start) as u64);
6561    let mut base = start;
6562    while base + L <= end {
6563        let _t0 = crate::cpuprof::time(crate::cpuprof::Slot::GuScales);
6564        let mut rows = [0usize; L];
6565        for l in 0..L {
6566            rows[l] = base + l;
6567        }
6568        let gv4: [&Q4tpView; L] = [gv; L];
6569        let uv4: [&Q4tpView; L] = [uv; L];
6570        q4tp_scales_lanes(&gv4, &rows, gpr, gsc);
6571        q4tp_scales_lanes(&uv4, &rows, gpr, usc);
6572        drop(_t0);
6573        let _t1 = crate::cpuprof::time(crate::cpuprof::Slot::GuDots);
6574        let mut ag = [_mm256_setzero_ps(); L];
6575        let mut au = [_mm256_setzero_ps(); L];
6576        for gi in 0..gpr {
6577            let xp = xq.as_ptr().add(gi * GROUP_SIZE);
6578            let xgb = _mm256_set1_ps(*xg8.get_unchecked(gi));
6579            for l in 0..L {
6580                let w = q4tp_unpack_w(gv.nib.as_ptr().add((rows[l] * gpr + gi) * Q4TP_NIB));
6581                let d = q4tp_madd_wx(w, xp);
6582                let t = _mm256_sub_ps(_mm256_cvtepi32_ps(d), xgb);
6583                let sb = _mm256_set1_ps(*gsc.get_unchecked(l * gpr + gi));
6584                ag[l] = _mm256_add_ps(ag[l], _mm256_mul_ps(t, sb));
6585            }
6586            for l in 0..L {
6587                let w = q4tp_unpack_w(uv.nib.as_ptr().add((rows[l] * gpr + gi) * Q4TP_NIB));
6588                let d = q4tp_madd_wx(w, xp);
6589                let t = _mm256_sub_ps(_mm256_cvtepi32_ps(d), xgb);
6590                let sb = _mm256_set1_ps(*usc.get_unchecked(l * gpr + gi));
6591                au[l] = _mm256_add_ps(au[l], _mm256_mul_ps(t, sb));
6592            }
6593        }
6594        drop(_t1);
6595        let _t2 = crate::cpuprof::time(crate::cpuprof::Slot::GuPost);
6596        for l in 0..L {
6597            let r = rows[l];
6598            let gs = &gsc[l * gpr..(l + 1) * gpr];
6599            let us = &usc[l * gpr..(l + 1) * gpr];
6600            let mut g = q4tp_hsum256(ag[l]) * sx;
6601            let mut u = q4tp_hsum256(au[l]) * sx;
6602            for &(j, xv) in outliers {
6603                let (w, s) = q4tp_outlier(gv.nib, r, gpr, j, gs);
6604                g += w * s * xv;
6605                let (w, s) = q4tp_outlier(uv.nib, r, gpr, j, us);
6606                u += w * s * xv;
6607            }
6608            // SAFETY: one lane owns each row.
6609            unsafe { *out.at(r) = silu_mul_limited(g, u, limit) };
6610        }
6611        drop(_t2);
6612        base += L;
6613    }
6614    // Partial lane batch (and any tail): one vertical cell per stream.
6615    for r in base..end {
6616        let g = &mut gsc[..gpr];
6617        let u = &mut usc[..gpr];
6618        gv.scales_into(r, gpr, g);
6619        uv.scales_into(r, gpr, u);
6620        let (g, u): (&[f32], &[f32]) = (g, u);
6621        let mut gv_ =
6622            q4tp_cell_dot_v(gv.nib.as_ptr(), r, gpr, xq.as_ptr(), g.as_ptr(), xg8.as_ptr()) * sx;
6623        let mut uv_ =
6624            q4tp_cell_dot_v(uv.nib.as_ptr(), r, gpr, xq.as_ptr(), u.as_ptr(), xg8.as_ptr()) * sx;
6625        for &(j, xv) in outliers {
6626            let (w, s) = q4tp_outlier(gv.nib, r, gpr, j, g);
6627            gv_ += w * s * xv;
6628            let (w, s) = q4tp_outlier(uv.nib, r, gpr, j, u);
6629            uv_ += w * s * xv;
6630        }
6631        // SAFETY: disjoint row ranges per worker.
6632        unsafe { *out.at(r) = silu_mul_limited(gv_, uv_, limit) };
6633    }
6634}
6635
6636/// Flat-space sibling of [`q4tp_matvec_lanes`] for `matvec_many`: the
6637/// virtual row space spans several tensors, so each lane resolves its
6638/// (tensor, row) pair through the caller's `locate`, then runs the plain
6639/// per-row math of the fused q4tp arm.
6640///
6641/// # Safety
6642/// AVX2 (callers gate on `a8w8_enabled()`); `outs` disjoint per row.
6643#[cfg(target_arch = "x86_64")]
6644#[target_feature(enable = "avx2")]
6645unsafe fn q4tp_many_lanes<const L: usize>(
6646    views: &[Q4tpView<'_>],
6647    locate: &(dyn Fn(usize) -> (usize, usize)),
6648    start: usize,
6649    end: usize,
6650    gpr: usize,
6651    xq: &[i8],
6652    xg8: &[f32],
6653    sx: f32,
6654    outliers: &[(usize, f32)],
6655    sc: &mut [f32],
6656    outs: &[SendMut],
6657) {
6658    use core::arch::x86_64::*;
6659    debug_assert!(L >= 2 && sc.len() >= L * gpr && xg8.len() == gpr);
6660    crate::cpuprof::count_lane_rows((end - start) as u64);
6661    let mut base = start;
6662    while base + L <= end {
6663        let _t0 = crate::cpuprof::time(crate::cpuprof::Slot::GuScales);
6664        let mut tk = [0usize; L];
6665        let mut rows = [0usize; L];
6666        for l in 0..L {
6667            let (t, r) = locate(base + l);
6668            tk[l] = t;
6669            rows[l] = r;
6670        }
6671        let v4: [&Q4tpView; L] = std::array::from_fn(|l| &views[tk[l]]);
6672        q4tp_scales_lanes(&v4, &rows, gpr, sc);
6673        drop(_t0);
6674        let _t1 = crate::cpuprof::time(crate::cpuprof::Slot::GuDots);
6675        let mut accs = [_mm256_setzero_ps(); L];
6676        for gi in 0..gpr {
6677            let xp = xq.as_ptr().add(gi * GROUP_SIZE);
6678            let xgb = _mm256_set1_ps(*xg8.get_unchecked(gi));
6679            for l in 0..L {
6680                let w = q4tp_unpack_w(
6681                    views[tk[l]].nib.as_ptr().add((rows[l] * gpr + gi) * Q4TP_NIB),
6682                );
6683                let d = q4tp_madd_wx(w, xp);
6684                let t = _mm256_sub_ps(_mm256_cvtepi32_ps(d), xgb);
6685                let sb = _mm256_set1_ps(*sc.get_unchecked(l * gpr + gi));
6686                accs[l] = _mm256_add_ps(accs[l], _mm256_mul_ps(t, sb));
6687            }
6688        }
6689        drop(_t1);
6690        let _t2 = crate::cpuprof::time(crate::cpuprof::Slot::GuPost);
6691        for l in 0..L {
6692            let (t, r) = (tk[l], rows[l]);
6693            let scl = &sc[l * gpr..(l + 1) * gpr];
6694            let mut acc = q4tp_hsum256(accs[l]) * sx;
6695            for &(j, xv) in outliers {
6696                let (w, s) = q4tp_outlier(views[t].nib, r, gpr, j, scl);
6697                acc += w * s * xv;
6698            }
6699            // SAFETY: one lane owns each (tensor, row) pair.
6700            unsafe { *outs[t].at(r) = acc };
6701        }
6702        drop(_t2);
6703        base += L;
6704    }
6705    // Partial lane batch (and any tail): one vertical cell per row.
6706    for flat in base..end {
6707        let (t, r) = locate(flat);
6708        let v = &views[t];
6709        let scl = &mut sc[..gpr];
6710        v.scales_into(r, gpr, scl);
6711        let scl: &[f32] = scl;
6712        let mut acc =
6713            q4tp_cell_dot_v(v.nib.as_ptr(), r, gpr, xq.as_ptr(), scl.as_ptr(), xg8.as_ptr()) * sx;
6714        for &(j, xv) in outliers {
6715            let (w, s) = q4tp_outlier(v.nib, r, gpr, j, scl);
6716            acc += w * s * xv;
6717        }
6718        // SAFETY: one worker owns each (tensor, row) pair.
6719        unsafe { *outs[t].at(r) = acc };
6720    }
6721}
6722
6723/// Exact scalar q4tp row — the `CMF_SDOT=0` contract, same pairwise
6724/// accumulation shape as `q4t_row_exact`.
6725#[inline]
6726fn q4tp_row_exact(nib: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
6727    #[cfg(target_arch = "x86_64")]
6728    if avx2_enabled() {
6729        // Keep the scalar pair/group reduction order, not a vector sum.
6730        return unsafe { q4tp_row_float_avx2(nib, r, gpr, x, scales) };
6731    }
6732    q4tp_row_float_scalar(nib, r, gpr, x, scales)
6733}
6734
6735#[inline]
6736fn q4tp_row_float_scalar(nib: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
6737    let mut acc = 0f32;
6738    for gi in 0..gpr {
6739        let tile = &nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
6740        let s = scales[gi];
6741        let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
6742        let mut ga = 0f32;
6743        for (k, &b) in tile.iter().enumerate() {
6744            ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
6745                + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
6746        }
6747        acc += ga * s;
6748    }
6749    acc
6750}
6751
6752/// Vectorize unpack, conversion and multiplication, but preserve every
6753/// pair addition and the scalar accumulation order. No activation rounding
6754/// or FMA: bit-identical to the float scalar row, including its group scale.
6755#[cfg(target_arch = "x86_64")]
6756#[target_feature(enable = "avx2")]
6757unsafe fn q4tp_row_float_avx2(nib: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
6758    // SAFETY: caller checks AVX2 and provides the same complete 32-element
6759    // groups as the scalar row. Loads/stores are explicitly unaligned.
6760    unsafe {
6761        use core::arch::x86_64::*;
6762        let mask = _mm_set1_epi8(15);
6763        let eight = _mm_set1_epi8(8);
6764        let order = _mm256_setr_epi32(0, 1, 4, 5, 2, 3, 6, 7);
6765        let mut acc = 0.0f32;
6766        for gi in 0..gpr {
6767            let packed = _mm_loadu_si128(nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB).cast());
6768            let lo = _mm_and_si128(packed, mask);
6769            let hi = _mm_and_si128(_mm_srli_epi16::<4>(packed), mask);
6770            let w0 = _mm_sub_epi8(_mm_unpacklo_epi8(lo, hi), eight);
6771            let w1 = _mm_sub_epi8(_mm_unpackhi_epi8(lo, hi), eight);
6772            let xp = x.as_ptr().add(gi * GROUP_SIZE);
6773            let a = _mm256_mul_ps(
6774                _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(w0)),
6775                _mm256_loadu_ps(xp),
6776            );
6777            let b = _mm256_mul_ps(
6778                _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(w0))),
6779                _mm256_loadu_ps(xp.add(8)),
6780            );
6781            let c = _mm256_mul_ps(
6782                _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(w1)),
6783                _mm256_loadu_ps(xp.add(16)),
6784            );
6785            let d = _mm256_mul_ps(
6786                _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(w1))),
6787                _mm256_loadu_ps(xp.add(24)),
6788            );
6789            let mut pairs = [0.0f32; 16];
6790            _mm256_storeu_ps(
6791                pairs.as_mut_ptr(),
6792                _mm256_permutevar8x32_ps(_mm256_hadd_ps(a, b), order),
6793            );
6794            _mm256_storeu_ps(
6795                pairs.as_mut_ptr().add(8),
6796                _mm256_permutevar8x32_ps(_mm256_hadd_ps(c, d), order),
6797            );
6798            let mut ga = 0.0f32;
6799            for v in pairs {
6800                ga += v;
6801            }
6802            acc += ga * scales[gi];
6803        }
6804        acc
6805    }
6806}
6807
6808/// Single weight of a q4tp tensor — the a8w8 outlier path, which restores
6809/// activation outliers at full precision after the int8 pass.
6810#[inline]
6811fn q4tp_outlier(nib: &[u8], r: usize, gpr: usize, j: usize, scales: &[f32]) -> (f32, f32) {
6812    let (gi, k) = (j / GROUP_SIZE, j % GROUP_SIZE);
6813    let byte = nib[(r * gpr + gi) * Q4TP_NIB + k / 2];
6814    let n = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
6815    ((n as i32 - 8) as f32, scales[gi])
6816}
6817
6818/// Fused q4tp matvec (dispatch mirrors `q4t_matvec`).
6819fn q4tp_matvec(
6820    bytes: &[u8],
6821    x: &[f32],
6822    rows: usize,
6823    cols: usize,
6824    out: &mut [f32],
6825    pool: Option<&Pool>,
6826) {
6827    debug_assert_eq!(out.len(), rows);
6828    let gpr = cols / GROUP_SIZE;
6829    let v = Q4tpView::new(bytes, rows, cols);
6830    let out_addr = SendMut(out.as_mut_ptr());
6831    if a8w8_enabled() {
6832        let act = split_act(x);
6833        #[cfg(target_arch = "x86_64")]
6834        let lanes = q4tp_lanes();
6835        let run = |start: usize, end: usize| {
6836            #[cfg(target_arch = "x86_64")]
6837            if lanes > 1 && act.xg8.len() == gpr {
6838                // L independent accumulator chains per worker — the plain
6839                // row body's math, lane-interleaved.
6840                with_krow(gpr * lanes, |sc| {
6841                    // SAFETY: a8w8 implies AVX2; disjoint row ranges.
6842                    unsafe {
6843                        match lanes {
6844                            2 => q4tp_matvec_lanes::<2>(
6845                                &v, start, end, gpr, &act.xq, &act.xg8, act.sx,
6846                                &act.outliers, sc, out_addr,
6847                            ),
6848                            8 => q4tp_matvec_lanes::<8>(
6849                                &v, start, end, gpr, &act.xq, &act.xg8, act.sx,
6850                                &act.outliers, sc, out_addr,
6851                            ),
6852                            _ => q4tp_matvec_lanes::<4>(
6853                                &v, start, end, gpr, &act.xq, &act.xg8, act.sx,
6854                                &act.outliers, sc, out_addr,
6855                            ),
6856                        }
6857                    }
6858                });
6859                return;
6860            }
6861            crate::cpuprof::count_plain_rows((end - start) as u64);
6862            // One scratch row of scales per worker — borrowed, not minted.
6863            with_krow(gpr, |sc| {
6864                for r in start..end {
6865                    v.scales_into(r, gpr, sc);
6866                    let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, sc) * act.sx;
6867                    for &(j, xv) in &act.outliers {
6868                        let (w, s) = q4tp_outlier(v.nib, r, gpr, j, sc);
6869                        acc += w * s * xv;
6870                    }
6871                    // SAFETY: disjoint row ranges per worker.
6872                    unsafe { *out_addr.at(r) = acc }
6873                }
6874            })
6875        };
6876        dispatch_rows(pool, rows, &run);
6877        return;
6878    }
6879    let run = |start: usize, end: usize| {
6880        with_krow(gpr, |sc| {
6881            for r in start..end {
6882                v.scales_into(r, gpr, sc);
6883                // SAFETY: disjoint row ranges per worker.
6884                unsafe { *out_addr.at(r) = q4tp_row_exact(v.nib, r, gpr, x, sc) };
6885            }
6886        })
6887    };
6888    dispatch_rows(pool, rows, &run);
6889}
6890
6891/// Fused two-input q4tp matvec — the SwiGLU gate/up pair. Weights and the
6892/// row ladder are read once and spent on both activation streams.
6893#[allow(clippy::too_many_arguments)]
6894fn q4tp_matvec2(
6895    bytes: &[u8],
6896    x1: &[f32],
6897    x2: &[f32],
6898    rows: usize,
6899    cols: usize,
6900    o1: &mut [f32],
6901    o2: &mut [f32],
6902    pool: Option<&Pool>,
6903) {
6904    let gpr = cols / GROUP_SIZE;
6905    let v = Q4tpView::new(bytes, rows, cols);
6906    let (p1, p2) = (SendMut(o1.as_mut_ptr()), SendMut(o2.as_mut_ptr()));
6907    let run = |start: usize, end: usize| {
6908        let mut sc = vec![0f32; gpr];
6909        for r in start..end {
6910            v.scales_into(r, gpr, &mut sc);
6911            // SAFETY: disjoint row ranges per worker.
6912            unsafe {
6913                *p1.at(r) = q4tp_row_exact(v.nib, r, gpr, x1, &sc);
6914                *p2.at(r) = q4tp_row_exact(v.nib, r, gpr, x2, &sc);
6915            }
6916        }
6917    };
6918    dispatch_rows(pool, rows, &run);
6919}
6920
6921/// One q2tp outlier weight at column `j` of row `r`: the 2-bit code and
6922/// its group scale, mirrored on `q4tp_outlier`.
6923#[inline]
6924fn q2tp_outlier(chunks: &[u8], r: usize, gpr: usize, j: usize, scales: &[f32]) -> (f32, f32) {
6925    let (gi, k) = (j / GROUP_SIZE, j % GROUP_SIZE);
6926    let byte = chunks[(r * gpr + gi) * Q2TP_CHUNK + k / 4];
6927    let c = (byte >> (2 * (k % 4))) & 3;
6928    (c as f32 - 1.5, scales[gi])
6929}
6930
6931#[cfg(target_arch = "x86_64")]
6932const Q2TP_DECODE_U32: [u32; 256] = {
6933    let mut tab = [0u32; 256];
6934    let mut b = 0usize;
6935    while b < 256 {
6936        tab[b] = ((b as u32) & 3)
6937            | ((((b as u32) >> 2) & 3) << 8)
6938            | ((((b as u32) >> 4) & 3) << 16)
6939            | ((((b as u32) >> 6) & 3) << 24);
6940        b += 1;
6941    }
6942    tab
6943};
6944
6945/// Eight packed q2tp bytes against 32 signed activation bytes. `maddubs`
6946/// exactly computes unsigned 2-bit code × signed i8; its pair sums cannot
6947/// saturate (2 × 3 × 127 < i16::MAX), and the second madd widens to i32.
6948#[cfg(target_arch = "x86_64")]
6949#[target_feature(enable = "avx2")]
6950unsafe fn q2tp_code_dot_avx2(ch: &[u8], x: &[i8]) -> i32 {
6951    use core::arch::x86_64::*;
6952    debug_assert!(ch.len() >= Q2TP_CHUNK && x.len() >= GROUP_SIZE);
6953    let codes = _mm256_setr_epi32(
6954        Q2TP_DECODE_U32[ch[0] as usize] as i32,
6955        Q2TP_DECODE_U32[ch[1] as usize] as i32,
6956        Q2TP_DECODE_U32[ch[2] as usize] as i32,
6957        Q2TP_DECODE_U32[ch[3] as usize] as i32,
6958        Q2TP_DECODE_U32[ch[4] as usize] as i32,
6959        Q2TP_DECODE_U32[ch[5] as usize] as i32,
6960        Q2TP_DECODE_U32[ch[6] as usize] as i32,
6961        Q2TP_DECODE_U32[ch[7] as usize] as i32,
6962    );
6963    let xv = unsafe { _mm256_loadu_si256(x.as_ptr().cast()) };
6964    let pair = _mm256_maddubs_epi16(codes, xv);
6965    let quad = _mm256_madd_epi16(pair, _mm256_set1_epi16(1));
6966    let sum128 = _mm_add_epi32(
6967        _mm256_castsi256_si128(quad),
6968        _mm256_extracti128_si256(quad, 1),
6969    );
6970    let sum64 = _mm_hadd_epi32(sum128, sum128);
6971    _mm_cvtsi128_si32(_mm_hadd_epi32(sum64, sum64))
6972}
6973
6974#[cfg(target_arch = "x86_64")]
6975#[inline]
6976fn q2tp_avx2_enabled() -> bool {
6977    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6978    *ON.get_or_init(|| std::arch::is_x86_feature_detected!("avx2"))
6979}
6980
6981/// Integer dot of one q2tp row against pre-quantized activations:
6982/// Σ_g s_g · (Σ c·xq − 1.5·Σ xq). The half-integer grid (c − 1.5)
6983/// becomes exact integer math through the group sums — the same trick
6984/// every a8w8 kernel in this file rides. The codes decode into a
6985/// 32-byte scratch in natural order and the dot itself is the shared
6986/// SDOT primitive; elsewhere a scalar integer loop.
6987#[inline]
6988fn dot_q2tp_row_i8(
6989    chunks: &[u8],
6990    r: usize,
6991    gpr: usize,
6992    xq: &[i8],
6993    gsum: &[i32],
6994    scales: &[f32],
6995) -> f32 {
6996    let mut acc = 0f32;
6997    let base = r * gpr * Q2TP_CHUNK;
6998    #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
6999    let mut codes = [0i8; GROUP_SIZE];
7000    #[cfg(target_arch = "x86_64")]
7001    // Cache the process-wide AVX2 decision; probing it for every 32-weight
7002    // group adds a branch to the hottest Q2TP row loop while retaining the
7003    // established table-load AVX2 arithmetic (and without coupling this path
7004    // to the separate FMA-gated A8W8 switch).
7005    let avx2 = q2tp_avx2_enabled();
7006    for gi in 0..gpr {
7007        let ch = &chunks[base + gi * Q2TP_CHUNK..base + (gi + 1) * Q2TP_CHUNK];
7008        let xg = &xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
7009        #[cfg(target_arch = "aarch64")]
7010        // NEON: the byte's four 2-bit fields land in four lane vectors
7011        // (shift+mask), vld4 de-interleaves xq to match (xj[k] =
7012        // xq[4k+j]), widening MACs accumulate exactly in i32. A scalar
7013        // decode here cost as much as the dot it fed — the profile put
7014        // it at the top of the whole W2 decode.
7015        let dot = unsafe {
7016            use core::arch::aarch64::*;
7017            let b = vld1_u8(ch.as_ptr());
7018            let three = vdup_n_u8(3);
7019            let c0 = vreinterpret_s8_u8(vand_u8(b, three));
7020            let c1 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 2), three));
7021            let c2 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 4), three));
7022            let c3 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 6), three));
7023            let x4 = vld4_s8(xg.as_ptr());
7024            let mut acc4 = vdupq_n_s32(0);
7025            acc4 = vpadalq_s16(acc4, vmull_s8(c0, x4.0));
7026            acc4 = vpadalq_s16(acc4, vmull_s8(c1, x4.1));
7027            acc4 = vpadalq_s16(acc4, vmull_s8(c2, x4.2));
7028            acc4 = vpadalq_s16(acc4, vmull_s8(c3, x4.3));
7029            vaddvq_s32(acc4)
7030        };
7031        #[cfg(target_arch = "x86_64")]
7032        let dot: i32 = if avx2 {
7033            // SAFETY: the runtime feature check gates the target-feature body;
7034            // the group slices above are exactly 8 and 32 bytes long.
7035            unsafe { q2tp_code_dot_avx2(ch, xg) }
7036        } else {
7037            ch.iter()
7038                .enumerate()
7039                .map(|(k, &b)| {
7040                    ((b & 3) as i32) * xg[k * 4] as i32
7041                        + (((b >> 2) & 3) as i32) * xg[k * 4 + 1] as i32
7042                        + (((b >> 4) & 3) as i32) * xg[k * 4 + 2] as i32
7043                        + (((b >> 6) & 3) as i32) * xg[k * 4 + 3] as i32
7044                })
7045                .sum()
7046        };
7047        #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
7048        let dot: i32 = {
7049            for (k, &b) in ch.iter().enumerate() {
7050                codes[k * 4] = (b & 3) as i8;
7051                codes[k * 4 + 1] = ((b >> 2) & 3) as i8;
7052                codes[k * 4 + 2] = ((b >> 4) & 3) as i8;
7053                codes[k * 4 + 3] = ((b >> 6) & 3) as i8;
7054            }
7055            codes
7056                .iter()
7057                .zip(xg)
7058                .map(|(&c, &x)| c as i32 * x as i32)
7059                .sum()
7060        };
7061        acc += scales[gi] * (dot as f32 - 1.5 * gsum[gi] as f32);
7062    }
7063    acc
7064}
7065
7066/// Exact f32 dot of one q2tp row: 2-bit fields LSB-first, (c − 1.5)·s.
7067/// Scalar on purpose — the 2-bit class targets the GPU graph; the CPU
7068/// path exists for parity gates and small-machine fallback.
7069fn q2tp_row_exact(chunks: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
7070    q2tp_row_exact_center(chunks, r, gpr, x, scales, 1.5)
7071}
7072
7073/// Fused Prism affine row: the derived correction is applied inside the
7074/// decoded code, avoiding a second accumulated dot and avoiding cancellation
7075/// between `B=(c-1.5)s` and `+.5s` for long 5120/17408 rows.
7076#[inline]
7077fn q2tp_affine_row_exact(chunks: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
7078    q2tp_row_exact_center(chunks, r, gpr, x, scales, 1.0)
7079}
7080
7081#[inline]
7082fn q2tp_row_exact_center(
7083    chunks: &[u8],
7084    r: usize,
7085    gpr: usize,
7086    x: &[f32],
7087    scales: &[f32],
7088    center: f32,
7089) -> f32 {
7090    let mut acc = 0f32;
7091    for gi in 0..gpr {
7092        let ch = &chunks[(r * gpr + gi) * Q2TP_CHUNK..(r * gpr + gi + 1) * Q2TP_CHUNK];
7093        let s = scales[gi];
7094        let xb = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
7095        let mut g = 0f32;
7096        for (k, &b) in ch.iter().enumerate() {
7097            g += ((b & 3) as f32 - center) * xb[k * 4]
7098                + (((b >> 2) & 3) as f32 - center) * xb[k * 4 + 1]
7099                + (((b >> 4) & 3) as f32 - center) * xb[k * 4 + 2]
7100                + (((b >> 6) & 3) as f32 - center) * xb[k * 4 + 3];
7101        }
7102        acc += s * g;
7103    }
7104    acc
7105}
7106
7107fn q2tp_matvec(
7108    bytes: &[u8],
7109    x: &[f32],
7110    rows: usize,
7111    cols: usize,
7112    out: &mut [f32],
7113    pool: Option<&Pool>,
7114) {
7115    q2tp_matvec_mode(bytes, x, rows, cols, out, pool, false);
7116}
7117
7118fn q2tp_affine_matvec(
7119    bytes: &[u8],
7120    x: &[f32],
7121    rows: usize,
7122    cols: usize,
7123    out: &mut [f32],
7124    pool: Option<&Pool>,
7125) {
7126    q2tp_matvec_mode(bytes, x, rows, cols, out, pool, true);
7127}
7128
7129fn q2tp_matvec_mode(
7130    bytes: &[u8],
7131    x: &[f32],
7132    rows: usize,
7133    cols: usize,
7134    out: &mut [f32],
7135    pool: Option<&Pool>,
7136    affine: bool,
7137) {
7138    debug_assert_eq!(out.len(), rows);
7139    let gpr = cols / GROUP_SIZE;
7140    let v = Q4tpView::new_q2(bytes, rows, cols);
7141    let out_addr = SendMut(out.as_mut_ptr());
7142    // a8w8 fast path (CMF_SDOT=0 keeps the exact scalar walk): integer
7143    // code dots + group sums, exact outlier correction — the same
7144    // contract as every sibling kernel; measured 2-bit rows were the
7145    // only scalar holdout in the family.
7146    if !affine && a8w8_enabled() {
7147        let act = split_act(x);
7148        let gsum = q1_group_sums(&act.xq, gpr);
7149        let (act, gsum) = (&act, &gsum);
7150        let run = move |start: usize, end: usize| {
7151            with_krow(gpr, |sc| {
7152                for r in start..end {
7153                    v.scales_into(r, gpr, sc);
7154                    let mut acc = dot_q2tp_row_i8(v.nib, r, gpr, &act.xq, gsum, sc) * act.sx;
7155                    for &(j, xv) in &act.outliers {
7156                        let (w, s) = q2tp_outlier(v.nib, r, gpr, j, sc);
7157                        acc += w * s * xv;
7158                    }
7159                    // SAFETY: disjoint row ranges per worker.
7160                    unsafe { *out_addr.at(r) = acc };
7161                }
7162            })
7163        };
7164        dispatch_rows(pool, rows, &run);
7165        return;
7166    }
7167    let run = |start: usize, end: usize| {
7168        with_krow(gpr, |sc| {
7169            for r in start..end {
7170                v.scales_into(r, gpr, sc);
7171                // SAFETY: disjoint row ranges per worker.
7172                unsafe {
7173                    *out_addr.at(r) = if affine {
7174                        q2tp_affine_row_exact(v.nib, r, gpr, x, sc)
7175                    } else {
7176                        q2tp_row_exact(v.nib, r, gpr, x, sc)
7177                    }
7178                };
7179            }
7180        })
7181    };
7182    dispatch_rows(pool, rows, &run);
7183}
7184
7185/// Fused two-input q2tp matvec — the SwiGLU gate/up pair.
7186#[allow(clippy::too_many_arguments)]
7187fn q2tp_matvec2(
7188    bytes: &[u8],
7189    x1: &[f32],
7190    x2: &[f32],
7191    rows: usize,
7192    cols: usize,
7193    o1: &mut [f32],
7194    o2: &mut [f32],
7195    pool: Option<&Pool>,
7196) {
7197    q2tp_matvec2_mode(bytes, x1, x2, rows, cols, o1, o2, pool, false);
7198}
7199
7200#[allow(clippy::too_many_arguments)]
7201fn q2tp_affine_matvec2(
7202    bytes: &[u8],
7203    x1: &[f32],
7204    x2: &[f32],
7205    rows: usize,
7206    cols: usize,
7207    o1: &mut [f32],
7208    o2: &mut [f32],
7209    pool: Option<&Pool>,
7210) {
7211    q2tp_matvec2_mode(bytes, x1, x2, rows, cols, o1, o2, pool, true);
7212}
7213
7214#[allow(clippy::too_many_arguments)]
7215fn q2tp_matvec2_mode(
7216    bytes: &[u8],
7217    x1: &[f32],
7218    x2: &[f32],
7219    rows: usize,
7220    cols: usize,
7221    o1: &mut [f32],
7222    o2: &mut [f32],
7223    pool: Option<&Pool>,
7224    affine: bool,
7225) {
7226    let gpr = cols / GROUP_SIZE;
7227    let v = Q4tpView::new_q2(bytes, rows, cols);
7228    let (p1, p2) = (SendMut(o1.as_mut_ptr()), SendMut(o2.as_mut_ptr()));
7229    let run = |start: usize, end: usize| {
7230        let mut sc = vec![0f32; gpr];
7231        for r in start..end {
7232            v.scales_into(r, gpr, &mut sc);
7233            // SAFETY: disjoint row ranges per worker.
7234            unsafe {
7235                *p1.at(r) = if affine {
7236                    q2tp_affine_row_exact(v.nib, r, gpr, x1, &sc)
7237                } else {
7238                    q2tp_row_exact(v.nib, r, gpr, x1, &sc)
7239                };
7240                *p2.at(r) = if affine {
7241                    q2tp_affine_row_exact(v.nib, r, gpr, x2, &sc)
7242                } else {
7243                    q2tp_row_exact(v.nib, r, gpr, x2, &sc)
7244                };
7245            }
7246        }
7247    };
7248    dispatch_rows(pool, rows, &run);
7249}
7250
7251/// Batched q2tp matmat: scalar row kernel over every batch column. CPU
7252/// prefill only — decode rides the graph, so plain and correct beats
7253/// clever here.
7254/// Test doors into the host 2-bit kernels: the stand's heap corruption
7255/// pointed at down-shaped tensors, and the private fns need a way to be
7256/// held to a reference without a model file around them.
7257pub fn q2tp_matvec_for_test(bytes: &[u8], x: &[f32], rows: usize, cols: usize, out: &mut [f32]) {
7258    // The facade IS the reference: encoder oracles hold requant output
7259    // to the exact scalar walk. The production dispatch may take the i8
7260    // fast path, whose error scale is the ACTIVATIONS' — a different
7261    // claim than the encoder correctness these tests pin.
7262    let gpr = cols / GROUP_SIZE;
7263    let v = Q4tpView::new_q2(bytes, rows, cols);
7264    with_krow(gpr, |sc| {
7265        for r in 0..rows {
7266            v.scales_into(r, gpr, sc);
7267            out[r] = q2tp_row_exact(v.nib, r, gpr, x, sc);
7268        }
7269    });
7270}
7271
7272/// Test door for the descriptor-specific fused affine decode.  Production
7273/// callers select this through a validated Prism header, never by dtype alone.
7274pub fn q2tp_affine_matvec_for_test(
7275    bytes: &[u8],
7276    x: &[f32],
7277    rows: usize,
7278    cols: usize,
7279    out: &mut [f32],
7280) {
7281    q2tp_affine_matvec(bytes, x, rows, cols, out, None);
7282}
7283
7284pub fn q2tp_matmat_for_test(
7285    bytes: &[u8],
7286    xs_all: &[f32],
7287    b: usize,
7288    rows: usize,
7289    cols: usize,
7290    out: &mut [f32],
7291) {
7292    q2tp_matmat(bytes, xs_all, b, rows, cols, out, None);
7293}
7294
7295fn q2tp_matmat(
7296    bytes: &[u8],
7297    xs_all: &[f32],
7298    b: usize,
7299    rows: usize,
7300    cols: usize,
7301    out: &mut [f32],
7302    pool: Option<&Pool>,
7303) {
7304    q2tp_matmat_mode(bytes, xs_all, b, rows, cols, out, pool, false);
7305}
7306
7307fn q2tp_affine_matmat(
7308    bytes: &[u8],
7309    xs_all: &[f32],
7310    b: usize,
7311    rows: usize,
7312    cols: usize,
7313    out: &mut [f32],
7314    pool: Option<&Pool>,
7315) {
7316    q2tp_matmat_mode(bytes, xs_all, b, rows, cols, out, pool, true);
7317}
7318
7319fn q2tp_matmat_mode(
7320    bytes: &[u8],
7321    xs_all: &[f32],
7322    b: usize,
7323    rows: usize,
7324    cols: usize,
7325    out: &mut [f32],
7326    pool: Option<&Pool>,
7327    affine: bool,
7328) {
7329    debug_assert_eq!(out.len(), b * rows);
7330    let gpr = cols / GROUP_SIZE;
7331    let v = Q4tpView::new_q2(bytes, rows, cols);
7332    let out_addr = SendMut(out.as_mut_ptr());
7333    let run = |start: usize, end: usize| {
7334        let mut sc = vec![0f32; gpr];
7335        for r in start..end {
7336            v.scales_into(r, gpr, &mut sc);
7337            for bi in 0..b {
7338                let x = &xs_all[bi * cols..(bi + 1) * cols];
7339                // SAFETY: disjoint row ranges per worker.
7340                unsafe {
7341                    *out_addr.at(bi * rows + r) = if affine {
7342                        q2tp_affine_row_exact(v.nib, r, gpr, x, &sc)
7343                    } else {
7344                        q2tp_row_exact(v.nib, r, gpr, x, &sc)
7345                    }
7346                };
7347            }
7348        }
7349    };
7350    dispatch_rows(pool, rows, &run);
7351}
7352
7353/// The pre-vectorised shape, kept for A/B (`CMF_Q4TP_V1=1`): the
7354/// horizontal add lands once per group per column instead of once per
7355/// row. Same weights, same activations — only the reduction differs.
7356/// It is also the row-exact batch kernel: per group and column it forms
7357/// `int dot as f32 * scale` and adds it to a scalar running sum, which is
7358/// `dot_q4tp_row_sdot` step for step (Rust never contracts to an fma),
7359/// so each column is bit-identical to that column's matvec.
7360#[cfg(target_arch = "aarch64")]
7361#[target_feature(enable = "neon,dotprod")]
7362unsafe fn dot_q4tp_row_1x4_sdot_v1(
7363    nib: &[u8],
7364    r: usize,
7365    gpr: usize,
7366    xs: [&[i8]; 4],
7367    scales: &[f32],
7368) -> [f32; 4] {
7369    unsafe {
7370        use core::arch::aarch64::*;
7371        use core::arch::asm;
7372        let lomask = vdupq_n_u8(0x0F);
7373        let eight = vdupq_n_s8(8);
7374        let (mut f0, mut f1, mut f2, mut f3) = (0f32, 0f32, 0f32, 0f32);
7375        for gi in 0..gpr {
7376            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
7377            let s = *scales.get_unchecked(gi);
7378            let bb = vld1q_u8(t);
7379            let lo = vandq_u8(bb, lomask);
7380            let hi = vshrq_n_u8::<4>(bb);
7381            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
7382            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
7383            let mut d = [0f32; 4];
7384            for (k, dk) in d.iter_mut().enumerate() {
7385                let x0 = vld1q_s8(xs[k].as_ptr().add(gi * GROUP_SIZE));
7386                let x1 = vld1q_s8(xs[k].as_ptr().add(gi * GROUP_SIZE + 16));
7387                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
7388                asm!(
7389                    "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
7390                    "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
7391                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
7392                    e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
7393                    options(pure, nomem, nostack),
7394                );
7395                *dk = vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
7396            }
7397            f0 += d[0];
7398            f1 += d[1];
7399            f2 += d[2];
7400            f3 += d[3];
7401        }
7402        [f0, f1, f2, f3]
7403    }
7404}
7405
7406/// Which q4tp batch kernel to run: 1 = the previous one, 2 = the tuned
7407/// one, 0 = decide from the CPU. An atomic rather than a `OnceLock` so a
7408/// benchmark can alternate the two inside one process, where the machine's
7409/// mood — a shared box drifts ±25% between runs — is the same for both.
7410/// What the two mean is per-architecture: on x86 the blocked AVX-512 path
7411/// against the per-column one, on ARM the two reduction shapes.
7412#[allow(dead_code)]
7413static Q4TP_ALT: std::sync::atomic::AtomicU8 = std::sync::atomic::AtomicU8::new(0);
7414
7415/// Tests that store `Q4TP_ALT` hold this, so one test's kernel pick does
7416/// not leak into another's bit-exact comparison running in parallel.
7417#[cfg(test)]
7418static Q4TP_ALT_TEST_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
7419
7420/// Blocking pays on x86 only with 512-bit VNNI. With AVX2 alone, four
7421/// columns sharing an unpack still measured slower than the per-column
7422/// path (23.2 ms against 19.4 on a 48-thread EPYC), because that path
7423/// already dequantizes the row once — so the blocked kernel bought a
7424/// second unpack-free pass at the price of half the vector width.
7425#[cfg(target_arch = "x86_64")]
7426fn q4tp_blocked_x86() -> bool {
7427    match Q4TP_ALT.load(std::sync::atomic::Ordering::Relaxed) {
7428        1 => false,
7429        // A forced ON still asks the CPU. The switch exists so a bench can
7430        // pick a kernel, not so it can promise instructions the machine
7431        // does not have — CI caught that as a SIGILL on a runner without
7432        // AVX-512, where the parity test had turned the path on by hand.
7433        2 => avx512vnni_enabled(),
7434        // Deliberately not cached back into the switch: both gates below
7435        // hold their own `OnceLock`, and latching their answer here would
7436        // make a test's override outlive the test that set it.
7437        _ => blocked_enabled() && avx512vnni_enabled(),
7438    }
7439}
7440
7441/// `CMF_Q4TP_V1=1` picks the old reduction shape (A/B only).
7442#[cfg(target_arch = "aarch64")]
7443#[allow(dead_code)]
7444fn q4tp_v1() -> bool {
7445    match Q4TP_ALT.load(std::sync::atomic::Ordering::Relaxed) {
7446        1 => true,
7447        2 => false,
7448        _ => {
7449            static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
7450            *ON.get_or_init(|| std::env::var("CMF_Q4TP_V1").is_ok_and(|v| v != "0"))
7451        }
7452    }
7453}
7454
7455/// Two weight rows against eight columns. The activation load is the
7456/// same for both rows, so it is paid once for twice the arithmetic, and
7457/// sixteen accumulator chains run where eight did — which is what a kernel
7458/// retiring 0.29 instructions a cycle is short of. Register pressure is
7459/// the limit: sixteen `zmm` accumulators, two weight tiles, one
7460/// activation, of thirty-two.
7461///
7462/// Four rows by four columns spends the same sixteen accumulators the
7463/// other way and measured worse — 1488 GFLOP/s against 1644 — so the
7464/// unpack, which four rows pay twice as often, costs more than the extra
7465/// sharing of one activation load buys.
7466#[cfg(target_arch = "x86_64")]
7467#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
7468unsafe fn dot_q4tp_2x8_avx512(
7469    nib: &[u8],
7470    r0: usize,
7471    gpr: usize,
7472    xs: [&[i8]; 8],
7473    sc0: &[f32],
7474    sc1: &[f32],
7475) -> [[f32; 8]; 2] {
7476    // SAFETY: as dot_q4tp_row_1x8_avx512, two adjacent rows at once; the
7477    // caller guarantees r0 + 1 < rows and the ISA.
7478    unsafe {
7479        use core::arch::x86_64::*;
7480        let lomask = _mm256_set1_epi8(0x0F);
7481        let eight = _mm256_set1_epi8(8);
7482        let zero = _mm512_setzero_si512();
7483        let mut v0 = [_mm512_setzero_ps(); 8];
7484        let mut v1 = [_mm512_setzero_ps(); 8];
7485        let pairs = gpr / 2;
7486        let unpack = |r: usize, gi: usize| -> (__m512i, __mmask64) {
7487            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
7488            let bb = _mm256_loadu_si256(t as *const __m256i);
7489            let lo = _mm256_and_si256(bb, lomask);
7490            let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
7491            let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
7492            let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
7493            let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
7494            let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
7495            (_mm512_abs_epi8(w), _mm512_movepi8_mask(w))
7496        };
7497        for gp in 0..pairs {
7498            let gi = gp * 2;
7499            let (wa0, neg0) = unpack(r0, gi);
7500            let (wa1, neg1) = unpack(r0 + 1, gi);
7501            let off = gi * GROUP_SIZE;
7502            let sv = |sc: &[f32]| {
7503                _mm512_insertf32x8::<1>(
7504                    _mm512_castps256_ps512(_mm256_set1_ps(*sc.get_unchecked(gi))),
7505                    _mm256_set1_ps(*sc.get_unchecked(gi + 1)),
7506                )
7507            };
7508            let s0 = sv(sc0);
7509            let s1 = sv(sc1);
7510            for k in 0..8 {
7511                let xv = _mm512_loadu_si512(xs[k].as_ptr().add(off) as *const __m512i);
7512                let d0 = _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(
7513                    zero,
7514                    wa0,
7515                    _mm512_mask_sub_epi8(xv, neg0, zero, xv),
7516                ));
7517                let d1 = _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(
7518                    zero,
7519                    wa1,
7520                    _mm512_mask_sub_epi8(xv, neg1, zero, xv),
7521                ));
7522                v0[k] = _mm512_fmadd_ps(d0, s0, v0[k]);
7523                v1[k] = _mm512_fmadd_ps(d1, s1, v1[k]);
7524            }
7525        }
7526        let mut acc = [[0f32; 8]; 2];
7527        for k in 0..8 {
7528            acc[0][k] = _mm512_reduce_add_ps(v0[k]);
7529            acc[1][k] = _mm512_reduce_add_ps(v1[k]);
7530        }
7531        if gpr % 2 == 1 {
7532            let off = (gpr - 1) * GROUP_SIZE;
7533            for j in off..off + GROUP_SIZE {
7534                let (w0, sa) = q4tp_outlier(nib, r0, gpr, j, sc0);
7535                let (w1, sb) = q4tp_outlier(nib, r0 + 1, gpr, j, sc1);
7536                for k in 0..8 {
7537                    let x = *xs[k].get_unchecked(j) as f32;
7538                    acc[0][k] += w0 * sa * x;
7539                    acc[1][k] += w1 * sb * x;
7540                }
7541            }
7542        }
7543        acc
7544    }
7545}
7546
7547/// The same, eight columns at a time. One unpack then feeds twice as many
7548/// activation streams, so a wide batch reads the weight tile half as
7549/// often; the price is eight accumulators live at once. Measured 9.0 ->
7550/// 8.3 ms at 9216x2304, b=296 on a 48-thread EPYC 9B45.
7551#[cfg(target_arch = "x86_64")]
7552#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
7553unsafe fn dot_q4tp_row_1x8_avx512(
7554    nib: &[u8],
7555    r: usize,
7556    gpr: usize,
7557    xs: [&[i8]; 8],
7558    scales: &[f32],
7559) -> [f32; 8] {
7560    // SAFETY: as dot_q4tp_row_1x4_avx2; caller guarantees the ISA.
7561    unsafe {
7562        use core::arch::x86_64::*;
7563        let lomask = _mm256_set1_epi8(0x0F);
7564        let eight = _mm256_set1_epi8(8);
7565        let zero = _mm512_setzero_si512();
7566        let (mut v0, mut v1, mut v2, mut v3) = (
7567            _mm512_setzero_ps(),
7568            _mm512_setzero_ps(),
7569            _mm512_setzero_ps(),
7570            _mm512_setzero_ps(),
7571        );
7572        let (mut v4, mut v5, mut v6, mut v7) = (
7573            _mm512_setzero_ps(),
7574            _mm512_setzero_ps(),
7575            _mm512_setzero_ps(),
7576            _mm512_setzero_ps(),
7577        );
7578        let pairs = gpr / 2;
7579        for gp in 0..pairs {
7580            let gi = gp * 2;
7581            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
7582            let bb = _mm256_loadu_si256(t as *const __m256i);
7583            let lo = _mm256_and_si256(bb, lomask);
7584            let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
7585            // `unpack` works per 128-bit lane, so the halves come out as
7586            // [A.lo, B.lo] and [A.hi, B.hi]; the shuffle reorders the four
7587            // 128-bit lanes into the weights' natural order, which is what
7588            // the straight activation load expects.
7589            let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
7590            let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
7591            let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
7592            let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
7593            let wabs = _mm512_abs_epi8(w);
7594            let neg = _mm512_movepi8_mask(w);
7595            let off = gi * GROUP_SIZE;
7596            let sv = _mm512_insertf32x8::<1>(
7597                _mm512_castps256_ps512(_mm256_set1_ps(*scales.get_unchecked(gi))),
7598                _mm256_set1_ps(*scales.get_unchecked(gi + 1)),
7599            );
7600            let dot = |x: &[i8]| -> __m512 {
7601                let xv = _mm512_loadu_si512(x.as_ptr().add(off) as *const __m512i);
7602                let sx = _mm512_mask_sub_epi8(xv, neg, zero, xv);
7603                _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(zero, wabs, sx))
7604            };
7605            v0 = _mm512_fmadd_ps(dot(xs[0]), sv, v0);
7606            v1 = _mm512_fmadd_ps(dot(xs[1]), sv, v1);
7607            v2 = _mm512_fmadd_ps(dot(xs[2]), sv, v2);
7608            v3 = _mm512_fmadd_ps(dot(xs[3]), sv, v3);
7609            v4 = _mm512_fmadd_ps(dot(xs[4]), sv, v4);
7610            v5 = _mm512_fmadd_ps(dot(xs[5]), sv, v5);
7611            v6 = _mm512_fmadd_ps(dot(xs[6]), sv, v6);
7612            v7 = _mm512_fmadd_ps(dot(xs[7]), sv, v7);
7613        }
7614        let mut acc = [
7615            _mm512_reduce_add_ps(v0),
7616            _mm512_reduce_add_ps(v1),
7617            _mm512_reduce_add_ps(v2),
7618            _mm512_reduce_add_ps(v3),
7619            _mm512_reduce_add_ps(v4),
7620            _mm512_reduce_add_ps(v5),
7621            _mm512_reduce_add_ps(v6),
7622            _mm512_reduce_add_ps(v7),
7623        ];
7624        // An odd group count leaves one group over; the narrow kernel
7625        // finishes it rather than the tail being a special case here.
7626        if gpr % 2 == 1 {
7627            let off = (gpr - 1) * GROUP_SIZE;
7628            for j in off..off + GROUP_SIZE {
7629                let (w, s) = q4tp_outlier(nib, r, gpr, j, scales);
7630                let ws = w * s;
7631                for k in 0..8 {
7632                    acc[k] += ws * *xs[k].get_unchecked(j) as f32;
7633                }
7634            }
7635        }
7636        acc
7637    }
7638}
7639
7640/// The same four columns, 512 bits wide. Two groups (64 weights) ride one
7641/// unpack and one `vpdpbusd`, where AVX2 needs two unpacks and four
7642/// `maddubs`/`madd` pairs — about 2.3x fewer instructions for the same
7643/// arithmetic. The two groups carry different scales, so the fma takes a
7644/// vector whose halves hold each group's scale rather than a broadcast.
7645///
7646/// There is no 512-bit `vpsignb`, so the activation's sign is applied by
7647/// negating under a mask taken from the weight's sign bits. That mask is
7648/// per-tile, so it is hoisted out of the column loop and the per-column
7649/// cost stays exactly one instruction, as with `sign_epi8`. Weights of
7650/// zero are not zeroed by the mask trick and do not need to be: their
7651/// magnitude is zero, so the product is.
7652#[cfg(target_arch = "x86_64")]
7653#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
7654unsafe fn dot_q4tp_row_1x4_avx512(
7655    nib: &[u8],
7656    r: usize,
7657    gpr: usize,
7658    xs: [&[i8]; 4],
7659    scales: &[f32],
7660) -> [f32; 4] {
7661    // SAFETY: as dot_q4tp_row_1x4_avx2; caller guarantees the ISA.
7662    unsafe {
7663        use core::arch::x86_64::*;
7664        let lomask = _mm256_set1_epi8(0x0F);
7665        let eight = _mm256_set1_epi8(8);
7666        let zero = _mm512_setzero_si512();
7667        let (mut v0, mut v1, mut v2, mut v3) = (
7668            _mm512_setzero_ps(),
7669            _mm512_setzero_ps(),
7670            _mm512_setzero_ps(),
7671            _mm512_setzero_ps(),
7672        );
7673        let pairs = gpr / 2;
7674        for gp in 0..pairs {
7675            let gi = gp * 2;
7676            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
7677            let bb = _mm256_loadu_si256(t as *const __m256i);
7678            let lo = _mm256_and_si256(bb, lomask);
7679            let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
7680            // `unpack` works per 128-bit lane, so the halves come out as
7681            // [A.lo, B.lo] and [A.hi, B.hi]; the shuffle reorders the four
7682            // 128-bit lanes into the weights' natural order, which is what
7683            // the straight activation load expects.
7684            let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
7685            let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
7686            let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
7687            let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
7688            let wabs = _mm512_abs_epi8(w);
7689            let neg = _mm512_movepi8_mask(w);
7690            let off = gi * GROUP_SIZE;
7691            let sv = _mm512_insertf32x8::<1>(
7692                _mm512_castps256_ps512(_mm256_set1_ps(*scales.get_unchecked(gi))),
7693                _mm256_set1_ps(*scales.get_unchecked(gi + 1)),
7694            );
7695            let dot = |x: &[i8]| -> __m512 {
7696                let xv = _mm512_loadu_si512(x.as_ptr().add(off) as *const __m512i);
7697                let sx = _mm512_mask_sub_epi8(xv, neg, zero, xv);
7698                _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(zero, wabs, sx))
7699            };
7700            v0 = _mm512_fmadd_ps(dot(xs[0]), sv, v0);
7701            v1 = _mm512_fmadd_ps(dot(xs[1]), sv, v1);
7702            v2 = _mm512_fmadd_ps(dot(xs[2]), sv, v2);
7703            v3 = _mm512_fmadd_ps(dot(xs[3]), sv, v3);
7704        }
7705        let mut acc = [
7706            _mm512_reduce_add_ps(v0),
7707            _mm512_reduce_add_ps(v1),
7708            _mm512_reduce_add_ps(v2),
7709            _mm512_reduce_add_ps(v3),
7710        ];
7711        // An odd group count leaves one group over; the narrow kernel
7712        // finishes it rather than the tail being a special case here.
7713        if gpr % 2 == 1 {
7714            let off = (gpr - 1) * GROUP_SIZE;
7715            for j in off..off + GROUP_SIZE {
7716                let (w, s) = q4tp_outlier(nib, r, gpr, j, scales);
7717                let ws = w * s;
7718                for k in 0..4 {
7719                    acc[k] += ws * *xs[k].get_unchecked(j) as f32;
7720                }
7721            }
7722        }
7723        acc
7724    }
7725}
7726
7727/// Four batch columns against one q4tp row: the tile is unpacked ONCE and
7728/// spent on four activation streams, which is where a prefill batch stops
7729/// being weight-bandwidth-bound. Twin of `dot_q4t_row_1x4_sdot`.
7730#[cfg(target_arch = "aarch64")]
7731#[target_feature(enable = "neon,dotprod")]
7732unsafe fn dot_q4tp_row_1x4_sdot(
7733    nib: &[u8],
7734    r: usize,
7735    gpr: usize,
7736    xs: [&[i8]; 4],
7737    scales: &[f32],
7738) -> [f32; 4] {
7739    // SAFETY: see dot_q4tp_row_sdot; every xs[k] is gpr·GROUP_SIZE long.
7740    unsafe {
7741        use core::arch::aarch64::*;
7742        use core::arch::asm;
7743        let lomask = vdupq_n_u8(0x0F);
7744        let eight = vdupq_n_s8(8);
7745        // Named accumulators, NOT an array indexed by a loop variable: the
7746        // latter does not stay in registers (the same defect cost 2x in the
7747        // AVX2 q4t kernel and again in WGSL).
7748        //
7749        // They are VECTORS, and the horizontal add happens once at the end
7750        // instead of once per group per column. `vaddvq` is a cross-lane
7751        // reduction — with 72 groups and four columns the old shape paid
7752        // 288 of them per row, each one a dependency stall the pipeline
7753        // cannot hide, to save four float adds. The group's scale now
7754        // rides an fma into the lane accumulators, so the arithmetic per
7755        // group is one convert and one fma. Summation order changes (the
7756        // lanes carry independent partial sums), which is the same
7757        // round-off class the SDOT path already lives in — the strict
7758        // kernel (`CMF_SDOT=0`, what `cortiq ppl` runs) is unchanged and
7759        // stays the reference.
7760        let (mut v0, mut v1, mut v2, mut v3) = (
7761            vdupq_n_f32(0.0),
7762            vdupq_n_f32(0.0),
7763            vdupq_n_f32(0.0),
7764            vdupq_n_f32(0.0),
7765        );
7766        for gi in 0..gpr {
7767            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
7768            let s = *scales.get_unchecked(gi);
7769            let bb = vld1q_u8(t);
7770            let lo = vandq_u8(bb, lomask);
7771            let hi = vshrq_n_u8::<4>(bb);
7772            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
7773            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
7774            let off = gi * GROUP_SIZE;
7775            let dot4 = |x: &[i8]| -> int32x4_t {
7776                let x0 = vld1q_s8(x.as_ptr().add(off));
7777                let x1 = vld1q_s8(x.as_ptr().add(off + 16));
7778                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
7779                asm!(
7780                    "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
7781                    "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
7782                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
7783                    e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
7784                    options(pure, nomem, nostack),
7785                );
7786                vaddq_s32(a0, a1)
7787            };
7788            v0 = vfmaq_n_f32(v0, vcvtq_f32_s32(dot4(xs[0])), s);
7789            v1 = vfmaq_n_f32(v1, vcvtq_f32_s32(dot4(xs[1])), s);
7790            v2 = vfmaq_n_f32(v2, vcvtq_f32_s32(dot4(xs[2])), s);
7791            v3 = vfmaq_n_f32(v3, vcvtq_f32_s32(dot4(xs[3])), s);
7792        }
7793        [
7794            vaddvq_f32(v0),
7795            vaddvq_f32(v1),
7796            vaddvq_f32(v2),
7797            vaddvq_f32(v3),
7798        ]
7799    }
7800}
7801
7802/// Fused q4tp matmat — the same three arms `q4t_matmat` has. Shipping only
7803/// the scalar one made Nanbeige-3B decode at 1.2 tok/s against q4t's 5.9:
7804/// the format was fine, the missing arms were the whole regression.
7805///
7806/// Under `row_exact()` every cell is summed in `q4tp_matvec`'s order on
7807/// every architecture; the mode is read once, so one call never mixes
7808/// the two contracts when a concurrent scope opens or closes mid-call.
7809fn q4tp_matmat(
7810    bytes: &[u8],
7811    xs_all: &[f32],
7812    b: usize,
7813    rows: usize,
7814    cols: usize,
7815    out: &mut [f32],
7816    pool: Option<&Pool>,
7817) {
7818    q4tp_matmat_with(bytes, xs_all, b, rows, cols, out, pool, row_exact())
7819}
7820
7821/// `q4tp_matmat` with the row-exact mode passed in rather than read from
7822/// the shared counter: `exact` = each cell equals its token's matvec bit
7823/// for bit, otherwise the fast blocked / AMX arms are free to reorder.
7824#[allow(clippy::too_many_arguments)]
7825fn q4tp_matmat_with(
7826    bytes: &[u8],
7827    xs_all: &[f32],
7828    b: usize,
7829    rows: usize,
7830    cols: usize,
7831    out: &mut [f32],
7832    pool: Option<&Pool>,
7833    exact: bool,
7834) {
7835    debug_assert_eq!(out.len(), b * rows);
7836    let gpr = cols / GROUP_SIZE;
7837    let v = Q4tpView::new(bytes, rows, cols);
7838
7839    // Wide batches ride the AMX through a dequant-tile sgemm, as in q4t.
7840    // An f32 GEMM over dequantized weights is not the matvec's int8 sum,
7841    // so a row-exact batch never takes it.
7842    #[cfg(target_os = "macos")]
7843    if !exact && b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
7844        dequant_matmat_accel(
7845            &|r, dst| {
7846                let mut sc = [0f32; 32];
7847                let mut scv;
7848                let s: &[f32] = if gpr <= 32 {
7849                    v.scales_into(r, gpr, &mut sc);
7850                    &sc[..gpr]
7851                } else {
7852                    scv = vec![0f32; gpr];
7853                    v.scales_into(r, gpr, &mut scv);
7854                    &scv
7855                };
7856                for gi in 0..gpr {
7857                    let tile = &v.nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
7858                    for (k, &bb) in tile.iter().enumerate() {
7859                        dst[gi * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s[gi];
7860                        dst[gi * GROUP_SIZE + k * 2 + 1] =
7861                            (((bb >> 4) & 0x0F) as f32 - 8.0) * s[gi];
7862                    }
7863                }
7864            },
7865            xs_all,
7866            b,
7867            rows,
7868            cols,
7869            out,
7870            pool,
7871        );
7872        return;
7873    }
7874
7875    let out_addr = SendMut(out.as_mut_ptr());
7876    if a8w8_enabled() {
7877        let acts: Vec<SplitAct> = (0..b)
7878            .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
7879            .collect();
7880        let acts = &acts;
7881        // ARM stays blocked under `exact` too: the 1x4 kernel then runs in
7882        // its v1 shape, whose per-group `int dot as f32 * scale` and scalar
7883        // running sum are exactly `dot_q4tp_row_sdot`'s, so the tile is
7884        // still unpacked once for four columns and each column equals its
7885        // matvec. The tuned shape (fma into lane partials, one horizontal
7886        // add per row) is 1-16 ulp off the matvec and is kept for `!exact`.
7887        #[cfg(target_arch = "aarch64")]
7888        let blocked_ok = sdot_enabled() && blocked_enabled();
7889        // x86 gets the same blocking: one tile unpack spent on four
7890        // columns. Without it every column re-decoded the row, which is
7891        // why a 48-core EPYC measured a sixth of an M4's per-core rate.
7892        // The gate is `avx2_enabled`, as in q4t — `sdot_enabled` answers
7893        // for ARM's dotprod and is hard-wired false everywhere else, so
7894        // asking it here left the whole blocked path unreachable on x86.
7895        #[cfg(target_arch = "x86_64")]
7896        let blocked_ok = q4tp_blocked_x86() && !exact;
7897        #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
7898        let blocked_ok = {
7899            let _ = exact;
7900            false
7901        };
7902        // Columns are swept in panels that fit L2. Without this a
7903        // row-pair walks every activation in the batch — 4.8 MB at
7904        // 512x512 — and does it again for the next pair, so the whole
7905        // batch streams out of the shared cache once per row. Measured
7906        // 800 GB/s of it, flat across batch sizes, which is the signature
7907        // of a loop bound by traffic rather than by arithmetic. A panel of
7908        // 256 columns is 590 KB beside 221 KB of this worker's weights:
7909        // both stay resident and the batch crosses L3 once instead of
7910        // once per row.
7911        let panel_cols: usize = std::env::var("CMF_Q4TP_PANEL")
7912            .ok()
7913            .and_then(|v| v.parse().ok())
7914            .filter(|v| *v > 0)
7915            .unwrap_or(256);
7916        let run = |start: usize, end: usize| {
7917            for abase in (0..acts.len()).step_by(panel_cols) {
7918                let alen = (acts.len() - abase).min(panel_cols);
7919                let mut sc = vec![0f32; gpr];
7920                #[cfg(target_arch = "x86_64")]
7921                let mut r_lo = start;
7922                #[cfg(target_arch = "x86_64")]
7923                if blocked_ok && alen >= 8 {
7924                    let mut sc1 = vec![0f32; gpr];
7925                    while r_lo + 2 <= end {
7926                        v.scales_into(r_lo, gpr, &mut sc);
7927                        v.scales_into(r_lo + 1, gpr, &mut sc1);
7928                        let mut bi = 0usize;
7929                        while bi + 8 <= alen {
7930                            let xs = [
7931                                acts[abase + bi].xq.as_slice(),
7932                                acts[abase + bi + 1].xq.as_slice(),
7933                                acts[abase + bi + 2].xq.as_slice(),
7934                                acts[abase + bi + 3].xq.as_slice(),
7935                                acts[abase + bi + 4].xq.as_slice(),
7936                                acts[abase + bi + 5].xq.as_slice(),
7937                                acts[abase + bi + 6].xq.as_slice(),
7938                                acts[abase + bi + 7].xq.as_slice(),
7939                            ];
7940                            let d = unsafe { dot_q4tp_2x8_avx512(v.nib, r_lo, gpr, xs, &sc, &sc1) };
7941                            for (row, dr, scr) in [(r_lo, &d[0], &sc), (r_lo + 1, &d[1], &sc1)] {
7942                                for k in 0..8 {
7943                                    let act = &acts[abase + bi + k];
7944                                    let mut acc = dr[k] * act.sx;
7945                                    for &(j, xv) in &act.outliers {
7946                                        let (w, s) = q4tp_outlier(v.nib, row, gpr, j, scr);
7947                                        acc += w * s * xv;
7948                                    }
7949                                    // SAFETY: disjoint (bi, r) cells per worker.
7950                                    unsafe { *out_addr.at((abase + bi + k) * rows + row) = acc };
7951                                }
7952                            }
7953                            bi += 8;
7954                        }
7955                        // Columns past the last group of eight, both rows —
7956                        // the same single-row kernel the tail below uses.
7957                        for row in [r_lo, r_lo + 1] {
7958                            let scr: &[f32] = if row == r_lo { &sc } else { &sc1 };
7959                            for b2 in bi..alen {
7960                                let act = &acts[abase + b2];
7961                                let xs4 = [
7962                                    act.xq.as_slice(),
7963                                    act.xq.as_slice(),
7964                                    act.xq.as_slice(),
7965                                    act.xq.as_slice(),
7966                                ];
7967                                let d =
7968                                    unsafe { dot_q4tp_row_1x4_avx512(v.nib, row, gpr, xs4, scr) };
7969                                let mut acc = d[0] * act.sx;
7970                                for &(j, xv) in &act.outliers {
7971                                    let (w, s) = q4tp_outlier(v.nib, row, gpr, j, scr);
7972                                    acc += w * s * xv;
7973                                }
7974                                // SAFETY: disjoint (bi, r) cells per worker.
7975                                unsafe { *out_addr.at((abase + b2) * rows + row) = acc };
7976                            }
7977                        }
7978                        r_lo += 2;
7979                    }
7980                }
7981                #[cfg(target_arch = "x86_64")]
7982                let row_start = r_lo;
7983                #[cfg(not(target_arch = "x86_64"))]
7984                let row_start = start;
7985                for r in row_start..end {
7986                    v.scales_into(r, gpr, &mut sc);
7987                    let mut bi = 0usize;
7988                    #[cfg(target_arch = "x86_64")]
7989                    if blocked_ok {
7990                        while bi + 8 <= alen {
7991                            let xs = [
7992                                acts[abase + bi].xq.as_slice(),
7993                                acts[abase + bi + 1].xq.as_slice(),
7994                                acts[abase + bi + 2].xq.as_slice(),
7995                                acts[abase + bi + 3].xq.as_slice(),
7996                                acts[abase + bi + 4].xq.as_slice(),
7997                                acts[abase + bi + 5].xq.as_slice(),
7998                                acts[abase + bi + 6].xq.as_slice(),
7999                                acts[abase + bi + 7].xq.as_slice(),
8000                            ];
8001                            let d = unsafe { dot_q4tp_row_1x8_avx512(v.nib, r, gpr, xs, &sc) };
8002                            for k in 0..8 {
8003                                let act = &acts[abase + bi + k];
8004                                let mut acc = d[k] * act.sx;
8005                                for &(j, xv) in &act.outliers {
8006                                    let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
8007                                    acc += w * s * xv;
8008                                }
8009                                // SAFETY: disjoint (bi, r) cells per worker.
8010                                unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
8011                            }
8012                            bi += 8;
8013                        }
8014                        while bi + 4 <= alen {
8015                            let xs = [
8016                                acts[abase + bi].xq.as_slice(),
8017                                acts[abase + bi + 1].xq.as_slice(),
8018                                acts[abase + bi + 2].xq.as_slice(),
8019                                acts[abase + bi + 3].xq.as_slice(),
8020                            ];
8021                            let d = unsafe { dot_q4tp_row_1x4_avx512(v.nib, r, gpr, xs, &sc) };
8022                            for k in 0..4 {
8023                                let act = &acts[abase + bi + k];
8024                                let mut acc = d[k] * act.sx;
8025                                for &(j, xv) in &act.outliers {
8026                                    let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
8027                                    acc += w * s * xv;
8028                                }
8029                                // SAFETY: disjoint (bi, r) cells per worker.
8030                                unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
8031                            }
8032                            bi += 4;
8033                        }
8034                    }
8035                    #[cfg(target_arch = "aarch64")]
8036                    if blocked_ok {
8037                        while bi + 4 <= alen {
8038                            let xs = [
8039                                acts[abase + bi].xq.as_slice(),
8040                                acts[abase + bi + 1].xq.as_slice(),
8041                                acts[abase + bi + 2].xq.as_slice(),
8042                                acts[abase + bi + 3].xq.as_slice(),
8043                            ];
8044                            let d = unsafe {
8045                                if exact || q4tp_v1() {
8046                                    dot_q4tp_row_1x4_sdot_v1(v.nib, r, gpr, xs, &sc)
8047                                } else {
8048                                    dot_q4tp_row_1x4_sdot(v.nib, r, gpr, xs, &sc)
8049                                }
8050                            };
8051                            for k in 0..4 {
8052                                let act = &acts[abase + bi + k];
8053                                let mut acc = d[k] * act.sx;
8054                                for &(j, xv) in &act.outliers {
8055                                    let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
8056                                    acc += w * s * xv;
8057                                }
8058                                // SAFETY: disjoint (bi, r) cells per worker.
8059                                unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
8060                            }
8061                            bi += 4;
8062                        }
8063                    }
8064                    let _ = blocked_ok;
8065                    // AVX2: the vertical column panel shares each group's
8066                    // unpack across four columns and keeps the row-exact
8067                    // cell bits (see `q4tp_matmat_row_cols_v`). Falls
8068                    // through to the single-column loop only when a column
8069                    // lacks its per-group correction.
8070                    #[cfg(target_arch = "x86_64")]
8071                    if acts[abase..abase + alen]
8072                        .iter()
8073                        .all(|a| a.xg8.len() == gpr)
8074                    {
8075                        // SAFETY: a8w8 implies AVX2; cells stay disjoint
8076                        // per (worker, bi, r).
8077                        unsafe {
8078                            q4tp_matmat_row_cols_v(
8079                                &v,
8080                                r,
8081                                gpr,
8082                                &acts[abase + bi..abase + alen],
8083                                &sc,
8084                                abase + bi,
8085                                rows,
8086                                out_addr,
8087                            )
8088                        };
8089                        bi = alen;
8090                    }
8091                    while bi < alen {
8092                        let act = &acts[abase + bi];
8093                        let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, &sc) * act.sx;
8094                        for &(j, xv) in &act.outliers {
8095                            let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
8096                            acc += w * s * xv;
8097                        }
8098                        // SAFETY: disjoint (bi, r) cells per worker range.
8099                        unsafe { *out_addr.at((abase + bi) * rows + r) = acc };
8100                        bi += 1;
8101                    }
8102                }
8103            }
8104        };
8105        dispatch_rows(pool, rows, &run);
8106        return;
8107    }
8108
8109    let run = |start: usize, end: usize| {
8110        let mut sc = vec![0f32; gpr];
8111        for r in start..end {
8112            v.scales_into(r, gpr, &mut sc);
8113            for bi in 0..b {
8114                let x = &xs_all[bi * cols..(bi + 1) * cols];
8115                // SAFETY: disjoint (bi, r) cells per worker range.
8116                unsafe { *out_addr.at(bi * rows + r) = q4tp_row_exact(v.nib, r, gpr, x, &sc) };
8117            }
8118        }
8119    };
8120    dispatch_rows(pool, rows, &run);
8121}
8122
8123/// Fused q4_tiled matvec (dispatch mirrors `q4matvec`).
8124fn q4t_matvec(
8125    bytes: &[u8],
8126    x: &[f32],
8127    rows: usize,
8128    cols: usize,
8129    out: &mut [f32],
8130    pool: Option<&Pool>,
8131) {
8132    debug_assert_eq!(out.len(), rows);
8133    let gpr = cols / GROUP_SIZE;
8134    let out_addr = SendMut(out.as_mut_ptr());
8135    if a8w8_enabled() {
8136        let act = split_act(x);
8137        let run = move |start: usize, end: usize| {
8138            for r in start..end {
8139                let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
8140                for &(j, xv) in &act.outliers {
8141                    let (w, s) = q4t_outlier(bytes, r, gpr, j);
8142                    acc += w * s * xv;
8143                }
8144                // SAFETY: disjoint row ranges per worker.
8145                unsafe { *out_addr.at(r) = acc };
8146            }
8147        };
8148        dispatch_rows(pool, rows, &run);
8149        return;
8150    }
8151    let run = move |start: usize, end: usize| {
8152        for r in start..end {
8153            // SAFETY: disjoint row ranges per worker.
8154            unsafe { *out_addr.at(r) = q4t_row_exact(bytes, r, gpr, x) };
8155        }
8156    };
8157    dispatch_rows(pool, rows, &run);
8158}
8159
8160/// Fused two-input q4_tiled matvec (weights read once per pair).
8161#[allow(clippy::too_many_arguments)]
8162fn q4t_matvec2(
8163    bytes: &[u8],
8164    x1: &[f32],
8165    x2: &[f32],
8166    rows: usize,
8167    cols: usize,
8168    o1: &mut [f32],
8169    o2: &mut [f32],
8170    pool: Option<&Pool>,
8171) {
8172    let gpr = cols / GROUP_SIZE;
8173    let p1 = SendMut(o1.as_mut_ptr());
8174    let p2 = SendMut(o2.as_mut_ptr());
8175    if a8w8_enabled() {
8176        let a1 = split_act(x1);
8177        let a2 = split_act(x2);
8178        let run = move |start: usize, end: usize| {
8179            for r in start..end {
8180                let mut v1 = dot_q4t_row_i8(bytes, r, gpr, &a1.xq) * a1.sx;
8181                let mut v2 = dot_q4t_row_i8(bytes, r, gpr, &a2.xq) * a2.sx;
8182                for &(j, xv) in &a1.outliers {
8183                    let (w, s) = q4t_outlier(bytes, r, gpr, j);
8184                    v1 += w * s * xv;
8185                }
8186                for &(j, xv) in &a2.outliers {
8187                    let (w, s) = q4t_outlier(bytes, r, gpr, j);
8188                    v2 += w * s * xv;
8189                }
8190                // SAFETY: disjoint row ranges per worker.
8191                unsafe {
8192                    *p1.at(r) = v1;
8193                    *p2.at(r) = v2;
8194                }
8195            }
8196        };
8197        dispatch_rows(pool, rows, &run);
8198        return;
8199    }
8200    let run = move |start: usize, end: usize| {
8201        for r in start..end {
8202            // SAFETY: disjoint row ranges per worker.
8203            unsafe {
8204                *p1.at(r) = q4t_row_exact(bytes, r, gpr, x1);
8205                *p2.at(r) = q4t_row_exact(bytes, r, gpr, x2);
8206            }
8207        }
8208    };
8209    dispatch_rows(pool, rows, &run);
8210}
8211
8212/// Batched q4_tiled matmat: each row's tiles stream once per microbatch.
8213#[allow(clippy::too_many_arguments)]
8214/// Prefill GEMM through Accelerate for group-quantized codecs: a
8215/// caller-supplied row dequantizer fills f32 tiles (pool-parallel) and
8216/// each tile rides the AMX with one sgemm — the generic sibling of
8217/// `qmatmat_accel` (q8). Numerics are f32-GEMM (tolerance class);
8218/// decode (b=1) never takes this path.
8219#[cfg(target_os = "macos")]
8220fn dequant_matmat_accel(
8221    dequant_row: &(dyn Fn(usize, &mut [f32]) + Sync),
8222    xs_all: &[f32],
8223    b: usize,
8224    rows: usize,
8225    cols: usize,
8226    out: &mut [f32],
8227    pool: Option<&Pool>,
8228) {
8229    const TR: usize = 2048;
8230    thread_local! {
8231        static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
8232    }
8233    WTILE.with(|wt| {
8234        let mut wtile = wt.borrow_mut();
8235        wtile.resize(TR * cols, 0.0);
8236        let mut r0 = 0usize;
8237        while r0 < rows {
8238            let tr = TR.min(rows - r0);
8239            let wt_addr = SendMut(wtile.as_mut_ptr());
8240            let run = |start: usize, end: usize| {
8241                for r in start..end {
8242                    // SAFETY: workers cover disjoint r ranges.
8243                    let dst = unsafe { std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols) };
8244                    dequant_row(r0 + r, dst);
8245                }
8246            };
8247            dispatch_rows(pool, tr, &run);
8248            unsafe {
8249                accel_blas::cblas_sgemm(
8250                    101, // RowMajor
8251                    111, // NoTrans A
8252                    112, // Trans B
8253                    b as i32,
8254                    tr as i32,
8255                    cols as i32,
8256                    1.0,
8257                    xs_all.as_ptr(),
8258                    cols as i32,
8259                    wtile.as_ptr(),
8260                    cols as i32,
8261                    0.0,
8262                    out.as_mut_ptr().add(r0),
8263                    rows as i32,
8264                );
8265            }
8266            r0 += tr;
8267        }
8268    });
8269}
8270
8271fn q4t_matmat(
8272    bytes: &[u8],
8273    xs_all: &[f32],
8274    b: usize,
8275    rows: usize,
8276    cols: usize,
8277    out: &mut [f32],
8278    pool: Option<&Pool>,
8279) {
8280    debug_assert_eq!(out.len(), b * rows);
8281    let gpr = cols / GROUP_SIZE;
8282    // Wide batches ride the AMX like q8's qmatmat: on Apple silicon
8283    // the dequant-tile sgemm is an order above the SDOT row loop for
8284    // prefill shapes (imagegen DiT forwards are exactly this).
8285    #[cfg(target_os = "macos")]
8286    if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
8287        dequant_matmat_accel(
8288            &|r, dst| {
8289                for gi in 0..gpr {
8290                    let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
8291                    let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
8292                    for (k, &bb) in tile[2..].iter().enumerate() {
8293                        dst[gi * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s;
8294                        dst[gi * GROUP_SIZE + k * 2 + 1] = (((bb >> 4) & 0x0F) as f32 - 8.0) * s;
8295                    }
8296                }
8297            },
8298            xs_all,
8299            b,
8300            rows,
8301            cols,
8302            out,
8303            pool,
8304        );
8305        return;
8306    }
8307    let out_addr = SendMut(out.as_mut_ptr());
8308    if a8w8_enabled() {
8309        let acts: Vec<SplitAct> = (0..b)
8310            .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
8311            .collect();
8312        let acts = &acts;
8313        #[cfg(target_arch = "x86_64")]
8314        let blocked_ok = avx2_enabled() && blocked_enabled();
8315        #[cfg(target_arch = "aarch64")]
8316        let blocked_ok = sdot_enabled() && blocked_enabled();
8317        #[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))]
8318        let blocked_ok = false;
8319        let run = move |start: usize, end: usize| {
8320            for r in start..end {
8321                let mut bi = 0usize;
8322                #[cfg(target_arch = "aarch64")]
8323                if blocked_ok {
8324                    while bi + 4 <= acts.len() {
8325                        let xs = [
8326                            acts[bi].xq.as_slice(),
8327                            acts[bi + 1].xq.as_slice(),
8328                            acts[bi + 2].xq.as_slice(),
8329                            acts[bi + 3].xq.as_slice(),
8330                        ];
8331                        let d = unsafe { dot_q4t_row_1x4_sdot(bytes, r, gpr, xs) };
8332                        for k in 0..4 {
8333                            let act = &acts[bi + k];
8334                            let mut acc = d[k] * act.sx;
8335                            for &(j, xv) in &act.outliers {
8336                                let (w, sc) = q4t_outlier(bytes, r, gpr, j);
8337                                acc += w * sc * xv;
8338                            }
8339                            // SAFETY: disjoint (bi, r) cells per worker.
8340                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
8341                        }
8342                        bi += 4;
8343                    }
8344                }
8345                #[cfg(target_arch = "x86_64")]
8346                if blocked_ok {
8347                    while bi + 4 <= acts.len() {
8348                        let xs = [
8349                            acts[bi].xq.as_slice(),
8350                            acts[bi + 1].xq.as_slice(),
8351                            acts[bi + 2].xq.as_slice(),
8352                            acts[bi + 3].xq.as_slice(),
8353                        ];
8354                        let d = unsafe {
8355                            if vnni_tiles_enabled() {
8356                                dot_q4t_row_1x4_vnni(bytes, r, gpr, xs)
8357                            } else {
8358                                dot_q4t_row_1x4_avx2(bytes, r, gpr, xs)
8359                            }
8360                        };
8361                        for k in 0..4 {
8362                            let act = &acts[bi + k];
8363                            let mut acc = d[k] * act.sx;
8364                            for &(j, xv) in &act.outliers {
8365                                let (w, sc) = q4t_outlier(bytes, r, gpr, j);
8366                                acc += w * sc * xv;
8367                            }
8368                            // SAFETY: disjoint (bi, r) cells per worker.
8369                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
8370                        }
8371                        bi += 4;
8372                    }
8373                }
8374                let _ = blocked_ok;
8375                while bi < acts.len() {
8376                    let act = &acts[bi];
8377                    let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
8378                    for &(j, xv) in &act.outliers {
8379                        let (w, s) = q4t_outlier(bytes, r, gpr, j);
8380                        acc += w * s * xv;
8381                    }
8382                    // SAFETY: disjoint (bi, r) cells per worker range.
8383                    unsafe { *out_addr.at(bi * rows + r) = acc };
8384                    bi += 1;
8385                }
8386            }
8387        };
8388        dispatch_rows(pool, rows, &run);
8389        return;
8390    }
8391    let run = move |start: usize, end: usize| {
8392        for r in start..end {
8393            for bi in 0..b {
8394                let x = &xs_all[bi * cols..(bi + 1) * cols];
8395                // SAFETY: disjoint (bi, r) cells per worker range.
8396                unsafe { *out_addr.at(bi * rows + r) = q4t_row_exact(bytes, r, gpr, x) };
8397            }
8398        }
8399    };
8400    dispatch_rows(pool, rows, &run);
8401}
8402
8403// ── q1 (dtype 12): binary weights, [f16 scale][4B sign bits] per
8404// 32-group tile. The kernel family mirrors q4_tiled: one sequential
8405// stream of 6-byte tiles, per-tile integer dot × scale, exact outlier
8406// correction (A8W8 contract), exact scalar path under CMF_SDOT=0. ──
8407
8408/// Per-32-group sums of the quantized activation — the ±1 identity's
8409/// shared half: `dot = −2·sdot(mask, x) − gsum[g]`, computed ONCE per
8410/// matvec and reused by every row.
8411fn q1_group_sums(xq: &[i8], gpr: usize) -> Vec<i32> {
8412    (0..gpr)
8413        .map(|gi| {
8414            xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE]
8415                .iter()
8416                .map(|&v| v as i32)
8417                .sum()
8418        })
8419        .collect()
8420}
8421
8422/// One q1 row via the A8W8 int8 path — mask-SDOT on ARM (no ±1
8423/// expansion at all), scalar bit loop elsewhere (AVX2 queued with the
8424/// x86 pass).
8425#[inline]
8426#[allow(unreachable_code)]
8427/// AVX2 q1 row via the same ±1 identity as the ARM sdot kernel: the
8428/// sign bits expand to a {0, −1} byte mask through shuffle+cmpeq, the
8429/// masked activation sums through maddubs(1, x&mask), and
8430/// `dot = −(2·masked_sum + Σx_group)` — bit-identical integer math.
8431#[cfg(target_arch = "x86_64")]
8432#[target_feature(enable = "avx2")]
8433unsafe fn dot_q1_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
8434    // SAFETY: callers uphold the 6B-tile and xq/gsum length contracts.
8435    unsafe {
8436        use core::arch::x86_64::*;
8437        // Byte j of the mask must replicate bits-byte j/8.
8438        let expand = _mm256_setr_epi8(
8439            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,
8440            3, 3, 3,
8441        );
8442        let bitsel = _mm256_setr_epi8(
8443            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
8444            -128, 1, 2, 4, 8, 16, 32, 64, -128,
8445        );
8446        let ones8 = _mm256_set1_epi8(1);
8447        let ones16 = _mm256_set1_epi16(1);
8448        let mut acc = 0f32;
8449        for gi in 0..gpr {
8450            let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
8451            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
8452            let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
8453            let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
8454            let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
8455            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8456            let sel = _mm256_and_si256(x, mask);
8457            // Σ of selected i8 lanes: maddubs(1u8, sel_i8) pairs → madd.
8458            let p16 = _mm256_maddubs_epi16(ones8, sel);
8459            let d32 = _mm256_madd_epi16(p16, ones16);
8460            let hi128 = _mm256_extracti128_si256::<1>(d32);
8461            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
8462            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8463            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8464            let msum = _mm_cvtsi128_si32(s32);
8465            // The and-select keeps x UN-negated (unlike ARM's −1-mask
8466            // sdot): d = Σ_set − Σ_unset = 2·Σ_set − Σ_all.
8467            let d = 2 * msum - gsum[gi];
8468            acc += d as f32 * s;
8469        }
8470        acc
8471    }
8472}
8473
8474/// VNNI twin of `dot_q1_row_avx2`: the masked-select sum goes through
8475/// one `vpdpbusd(1u8, sel)` (see `dpbusd_hsum` — bit-identical).
8476#[cfg(target_arch = "x86_64")]
8477#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
8478unsafe fn dot_q1_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
8479    // SAFETY: callers uphold the 6B-tile and xq/gsum length contracts.
8480    unsafe {
8481        use core::arch::x86_64::*;
8482        let expand = _mm256_setr_epi8(
8483            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,
8484            3, 3, 3,
8485        );
8486        let bitsel = _mm256_setr_epi8(
8487            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
8488            -128, 1, 2, 4, 8, 16, 32, 64, -128,
8489        );
8490        let ones8 = _mm256_set1_epi8(1);
8491        let mut acc = 0f32;
8492        for gi in 0..gpr {
8493            let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
8494            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
8495            let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
8496            let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
8497            let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
8498            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8499            let msum = dpbusd_hsum(ones8, _mm256_and_si256(x, mask));
8500            let d = 2 * msum - gsum[gi];
8501            acc += d as f32 * s;
8502        }
8503        acc
8504    }
8505}
8506
8507/// VNNI twin of `dot_q1_row_1x4_avx2` (see `dpbusd_hsum`).
8508#[cfg(target_arch = "x86_64")]
8509#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
8510unsafe fn dot_q1_row_1x4_vnni(
8511    bytes: &[u8],
8512    r: usize,
8513    gpr: usize,
8514    xs: [&[i8]; 4],
8515    gsums: [&[i32]; 4],
8516) -> [f32; 4] {
8517    // SAFETY: callers uphold the 6B-tile and xq/gsum length contracts.
8518    unsafe {
8519        use core::arch::x86_64::*;
8520        let expand = _mm256_setr_epi8(
8521            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,
8522            3, 3, 3,
8523        );
8524        let bitsel = _mm256_setr_epi8(
8525            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
8526            -128, 1, 2, 4, 8, 16, 32, 64, -128,
8527        );
8528        let ones8 = _mm256_set1_epi8(1);
8529        let mut acc = [0f32; 4];
8530        for gi in 0..gpr {
8531            let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
8532            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
8533            let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
8534            let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
8535            let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
8536            for (k, xq) in xs.iter().enumerate() {
8537                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8538                let msum = dpbusd_hsum(ones8, _mm256_and_si256(x, mask));
8539                let d = 2 * msum - gsums[k][gi];
8540                acc[k] += d as f32 * s;
8541            }
8542        }
8543        acc
8544    }
8545}
8546
8547/// The blocked 1×4 flavor: the expanded bit mask serves four activation
8548/// streams per group (mask build once, four select+reduce chains).
8549#[cfg(target_arch = "x86_64")]
8550#[target_feature(enable = "avx2")]
8551unsafe fn dot_q1_row_1x4_avx2(
8552    bytes: &[u8],
8553    r: usize,
8554    gpr: usize,
8555    xs: [&[i8]; 4],
8556    gsums: [&[i32]; 4],
8557) -> [f32; 4] {
8558    // SAFETY: callers uphold the 6B-tile and xq/gsum length contracts.
8559    unsafe {
8560        use core::arch::x86_64::*;
8561        let expand = _mm256_setr_epi8(
8562            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,
8563            3, 3, 3,
8564        );
8565        let bitsel = _mm256_setr_epi8(
8566            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
8567            -128, 1, 2, 4, 8, 16, 32, 64, -128,
8568        );
8569        let ones8 = _mm256_set1_epi8(1);
8570        let ones16 = _mm256_set1_epi16(1);
8571        let mut acc = [0f32; 4];
8572        for gi in 0..gpr {
8573            let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
8574            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
8575            let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
8576            let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
8577            let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
8578            for (k, xq) in xs.iter().enumerate() {
8579                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8580                let sel = _mm256_and_si256(x, mask);
8581                let p16 = _mm256_maddubs_epi16(ones8, sel);
8582                let d32 = _mm256_madd_epi16(p16, ones16);
8583                let hi128 = _mm256_extracti128_si256::<1>(d32);
8584                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
8585                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8586                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8587                let msum = _mm_cvtsi128_si32(s32);
8588                let d = 2 * msum - gsums[k][gi];
8589                acc[k] += d as f32 * s;
8590            }
8591        }
8592        acc
8593    }
8594}
8595
8596#[allow(unreachable_code)]
8597fn dot_q1_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
8598    #[cfg(target_arch = "aarch64")]
8599    unsafe {
8600        return dot_q1_row_sdot(bytes, r, gpr, xq, gsum);
8601    }
8602    #[cfg(target_arch = "x86_64")]
8603    if avx2_enabled() {
8604        unsafe {
8605            if vnni_tiles_enabled() {
8606                return dot_q1_row_vnni(bytes, r, gpr, xq, gsum);
8607            }
8608            return dot_q1_row_avx2(bytes, r, gpr, xq, gsum);
8609        }
8610    }
8611    let _ = gsum;
8612    let mut acc = 0f32;
8613    for gi in 0..gpr {
8614        let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
8615        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
8616        let mut d = 0i32;
8617        for (j, &b) in tile[2..].iter().enumerate() {
8618            for k in 0..8 {
8619                let w = ((b >> k) & 1) as i32 * 2 - 1;
8620                d += w * xq[gi * GROUP_SIZE + j * 8 + k] as i32;
8621            }
8622        }
8623        acc += d as f32 * s;
8624    }
8625    acc
8626}
8627
8628/// SDOT q1 row via the ±1 identity: the vtst mask (0xFF where the bit
8629/// is set, i.e. −1 as i8) feeds `sdot` DIRECTLY — no expansion to ±1
8630/// lanes at all — and `dot = −(2·sdot(mask, x) + Σx_group)`, with the
8631/// per-group activation sums shared across every row of the matvec.
8632/// Four tiles (128 weights) per iteration: integer dots reduce through
8633/// a vpaddq tree into ONE i32x4 that meets its four scales in a single
8634/// fused f32 multiply-add. Integer math throughout — bit-identical to
8635/// the scalar ±1 reference.
8636#[cfg(target_arch = "aarch64")]
8637#[target_feature(enable = "neon,dotprod")]
8638unsafe fn dot_q1_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
8639    // SAFETY: callers uphold slice-length contracts (6B tile per group,
8640    // xq.len() == gpr·GROUP_SIZE, gsum.len() == gpr).
8641    unsafe {
8642        use core::arch::aarch64::*;
8643        use core::arch::asm;
8644        const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
8645        let m = vld1q_u8(MASKS.as_ptr());
8646        // One tile's −Σ_set(x) as an UNREDUCED i32x4 (two mask-sdots).
8647        macro_rules! tile_dot {
8648            ($t:expr, $x:expr) => {{
8649                let v0 = vcombine_u8(vdup_n_u8(*$t.add(2)), vdup_n_u8(*$t.add(3)));
8650                let v1 = vcombine_u8(vdup_n_u8(*$t.add(4)), vdup_n_u8(*$t.add(5)));
8651                let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
8652                let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
8653                let x0 = vld1q_s8($x);
8654                let x1 = vld1q_s8($x.add(16));
8655                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
8656                asm!(
8657                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
8658                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
8659                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
8660                    w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
8661                    options(pure, nomem, nostack),
8662                );
8663                vaddq_s32(a0, a1)
8664            }};
8665        }
8666        // TBL unpack over PAIR loads: one vld1q covers two 6B tiles
8667        // ([s s b b b b][s s b b b b] + 4B slack), TBL replicates each
8668        // bit-byte across 8 lanes for vtst, and the four scales gather
8669        // through tbl2 into one fcvtl — the 16 ld1r broadcast loads and
8670        // 4 branchy software f16 conversions per 128 weights (the
8671        // measured load-port wall of this kernel) become 2 vector
8672        // loads + 9 table lookups. Integer math order is unchanged —
8673        // bit-identical results (FCVTL is exact on every f16).
8674        const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
8675        const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
8676        const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
8677        const IW11: [u8; 16] = [
8678            10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
8679        ];
8680        const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
8681        let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
8682        let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
8683        let isc = vld1_u8(ISC.as_ptr());
8684        // One tile's −Σ_set(x) from a TBL-unpacked pair load.
8685        macro_rules! tile_dot_tbl {
8686            ($ld:expr, $i0:expr, $i1:expr, $x:expr) => {{
8687                let w0 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i0), m));
8688                let w1 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i1), m));
8689                let x0 = vld1q_s8($x);
8690                let x1 = vld1q_s8($x.add(16));
8691                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
8692                asm!(
8693                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
8694                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
8695                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
8696                    w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
8697                    options(pure, nomem, nostack),
8698                );
8699                vaddq_s32(a0, a1)
8700            }};
8701        }
8702        let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
8703        let row_base = r * gpr * Q1_TILE;
8704        let abs_end = bytes.len();
8705        let xp = xq.as_ptr();
8706        let gp = gsum.as_ptr();
8707        let mut accv = vdupq_n_f32(0.0);
8708        let mut gi = 0;
8709        // The second pair load reads 4B past tile gi+3 — stay inside
8710        // the payload slice (only the file's final tiles fall back).
8711        while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
8712            let t0 = base.add(gi * Q1_TILE);
8713            let ld_a = vld1q_u8(t0);
8714            let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
8715            let d0 = tile_dot_tbl!(ld_a, iw00, iw01, xp.add(gi * GROUP_SIZE));
8716            let d1 = tile_dot_tbl!(ld_a, iw10, iw11, xp.add((gi + 1) * GROUP_SIZE));
8717            let d2 = tile_dot_tbl!(ld_b, iw00, iw01, xp.add((gi + 2) * GROUP_SIZE));
8718            let d3 = tile_dot_tbl!(ld_b, iw10, iw11, xp.add((gi + 3) * GROUP_SIZE));
8719            // [−Σ0, −Σ1, −Σ2, −Σ3] → dots = −(2·Σset_neg + gsum)
8720            let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
8721            let g = vld1q_s32(gp.add(gi));
8722            let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
8723            let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
8724            let scf: float32x4_t;
8725            asm!(
8726                "fcvtl {o:v}.4s, {i:v}.4h",
8727                o = out(vreg) scf, i = in(vreg) sc16,
8728                options(pure, nomem, nostack),
8729            );
8730            accv = vfmaq_f32(accv, vcvtq_f32_s32(dots), scf);
8731            gi += 4;
8732        }
8733        let mut acc = vaddvq_f32(accv);
8734        while gi < gpr {
8735            let t = base.add(gi * Q1_TILE);
8736            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
8737            let d = vaddvq_s32(tile_dot!(t, xp.add(gi * GROUP_SIZE)));
8738            acc += (-(2 * d + *gp.add(gi))) as f32 * s;
8739            gi += 1;
8740        }
8741        acc
8742    }
8743}
8744
8745/// Blocked q1 1×4: one TBL unpack of the tile pair serves FOUR
8746/// activation streams (prefill amortization — the same idea as the
8747/// AVX2 twin; per stream the group order, fma order and tail match the
8748/// single-row kernel exactly, so batch == matvec bit-for-bit).
8749#[cfg(target_arch = "aarch64")]
8750#[target_feature(enable = "neon,dotprod")]
8751unsafe fn dot_q1_row_1x4_sdot(
8752    bytes: &[u8],
8753    r: usize,
8754    gpr: usize,
8755    xs: [&[i8]; 4],
8756    gs: [&[i32]; 4],
8757) -> [f32; 4] {
8758    // SAFETY: same slice-length contracts as `dot_q1_row_sdot`, ×4.
8759    unsafe {
8760        use core::arch::aarch64::*;
8761        use core::arch::asm;
8762        const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
8763        const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
8764        const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
8765        const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
8766        const IW11: [u8; 16] = [
8767            10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
8768        ];
8769        const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
8770        let m = vld1q_u8(MASKS.as_ptr());
8771        let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
8772        let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
8773        let isc = vld1_u8(ISC.as_ptr());
8774        macro_rules! sdot2 {
8775            ($w0:expr, $w1:expr, $x:expr) => {{
8776                let x0 = vld1q_s8($x);
8777                let x1 = vld1q_s8($x.add(16));
8778                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
8779                asm!(
8780                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
8781                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
8782                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
8783                    w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
8784                    options(pure, nomem, nostack),
8785                );
8786                vaddq_s32(a0, a1)
8787            }};
8788        }
8789        let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
8790        let row_base = r * gpr * Q1_TILE;
8791        let abs_end = bytes.len();
8792        let mut accv = [vdupq_n_f32(0.0); 4];
8793        let mut gi = 0;
8794        while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
8795            let t0 = base.add(gi * Q1_TILE);
8796            let ld_a = vld1q_u8(t0);
8797            let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
8798            // Unpack ONCE — eight ±mask vectors serve all four streams.
8799            let w00 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw00), m));
8800            let w01 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw01), m));
8801            let w10 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw10), m));
8802            let w11 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw11), m));
8803            let w20 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw00), m));
8804            let w21 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw01), m));
8805            let w30 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw10), m));
8806            let w31 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw11), m));
8807            let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
8808            let scf: float32x4_t;
8809            asm!(
8810                "fcvtl {o:v}.4s, {i:v}.4h",
8811                o = out(vreg) scf, i = in(vreg) sc16,
8812                options(pure, nomem, nostack),
8813            );
8814            for k in 0..4 {
8815                let xp = xs[k].as_ptr();
8816                let d0 = sdot2!(w00, w01, xp.add(gi * GROUP_SIZE));
8817                let d1 = sdot2!(w10, w11, xp.add((gi + 1) * GROUP_SIZE));
8818                let d2 = sdot2!(w20, w21, xp.add((gi + 2) * GROUP_SIZE));
8819                let d3 = sdot2!(w30, w31, xp.add((gi + 3) * GROUP_SIZE));
8820                let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
8821                let g = vld1q_s32(gs[k].as_ptr().add(gi));
8822                let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
8823                accv[k] = vfmaq_f32(accv[k], vcvtq_f32_s32(dots), scf);
8824            }
8825            gi += 4;
8826        }
8827        let mut acc = [
8828            vaddvq_f32(accv[0]),
8829            vaddvq_f32(accv[1]),
8830            vaddvq_f32(accv[2]),
8831            vaddvq_f32(accv[3]),
8832        ];
8833        while gi < gpr {
8834            let t = base.add(gi * Q1_TILE);
8835            let sc = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
8836            let v0 = vcombine_u8(vdup_n_u8(*t.add(2)), vdup_n_u8(*t.add(3)));
8837            let v1 = vcombine_u8(vdup_n_u8(*t.add(4)), vdup_n_u8(*t.add(5)));
8838            let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
8839            let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
8840            for k in 0..4 {
8841                let d = vaddvq_s32(sdot2!(w0, w1, xs[k].as_ptr().add(gi * GROUP_SIZE)));
8842                acc[k] += (-(2 * d + *gs[k].as_ptr().add(gi))) as f32 * sc;
8843            }
8844            gi += 1;
8845        }
8846        acc
8847    }
8848}
8849
8850/// (weight ±1, scale) of one q1 element — the exact outlier term.
8851#[inline]
8852fn q1_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
8853    let gi = j / GROUP_SIZE;
8854    let k = j % GROUP_SIZE;
8855    let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
8856    let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
8857    let bit = (tile[2 + k / 8] >> (k % 8)) & 1;
8858    ((bit as i32 * 2 - 1) as f32, s)
8859}
8860
8861/// Exact scalar q1 row (CMF_SDOT=0 contract).
8862#[inline]
8863fn q1_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
8864    let mut acc = 0f32;
8865    for gi in 0..gpr {
8866        let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
8867        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
8868        let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
8869        let mut ga = 0f32;
8870        for (j, &b) in tile[2..].iter().enumerate() {
8871            for k in 0..8 {
8872                ga += (((b >> k) & 1) as f32 * 2.0 - 1.0) * xg[j * 8 + k];
8873            }
8874        }
8875        acc += ga * s;
8876    }
8877    acc
8878}
8879
8880/// One binary-weight row range via A8W8 (the body of `bin_matvec`'s hot
8881/// loop, extracted so multi-matrix jobs drive the same kernel).
8882#[allow(clippy::too_many_arguments)]
8883fn bin_range_a8w8(
8884    w: BinW,
8885    gpr: usize,
8886    act: &SplitAct,
8887    gsum: &[i32],
8888    out: SendMut,
8889    start: usize,
8890    end: usize,
8891) {
8892    for r in start..end {
8893        let mut acc = w.dot_i8(r, gpr, &act.xq, gsum) * act.sx;
8894        for &(j, xv) in &act.outliers {
8895            let (wv, s) = w.outlier(r, gpr, j);
8896            acc += wv * s * xv;
8897        }
8898        // SAFETY: disjoint row ranges per worker.
8899        unsafe { *out.at(r) = acc };
8900    }
8901}
8902
8903/// Exact-scalar binary-weight row range (CMF_SDOT=0 contract).
8904fn bin_range_f32(w: BinW, gpr: usize, x: &[f32], out: SendMut, start: usize, end: usize) {
8905    for r in start..end {
8906        // SAFETY: disjoint row ranges per worker.
8907        unsafe { *out.at(r) = w.row_exact(r, gpr, x) };
8908    }
8909}
8910
8911/// A binary (±s) weight in either of its layouts: `q1` 6-byte tiles, or
8912/// `q1g` planes with one scale per 128-group. The row drivers — matvec,
8913/// matvec2, matmat, the multi-projection job and the fused gate/up pair —
8914/// are written once against this; only the per-row kernels differ, and
8915/// the match costs one predictable branch per row. `gpr` is always the
8916/// row's count of 32-tiles (cols / 32): the activation group sums and the
8917/// tile walk are shared by both layouts.
8918#[derive(Clone, Copy)]
8919pub(crate) enum BinW<'a> {
8920    Q1(&'a [u8]),
8921    Q1g {
8922        bytes: &'a [u8],
8923        /// Byte offset of the f16 scale plane (= rows · cols/8).
8924        sc_off: usize,
8925        /// Row count: the q1 kernels treat the payload's final row
8926        /// specially (see `dot_q1g_row_sdot`), and q1g mirrors that.
8927        rows: usize,
8928    },
8929}
8930
8931impl<'a> BinW<'a> {
8932    fn q1g(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
8933        Self::Q1g {
8934            bytes,
8935            sc_off: cortiq_core::quant::q1g_scale_offset(rows, cols),
8936            rows,
8937        }
8938    }
8939
8940    #[inline]
8941    fn dot_i8(self, r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
8942        match self {
8943            Self::Q1(b) => dot_q1_row_i8(b, r, gpr, xq, gsum),
8944            Self::Q1g {
8945                bytes,
8946                sc_off,
8947                rows,
8948            } => dot_q1g_row_i8(bytes, sc_off, rows, r, gpr, xq, gsum),
8949        }
8950    }
8951
8952    #[inline]
8953    fn row_exact(self, r: usize, gpr: usize, x: &[f32]) -> f32 {
8954        match self {
8955            Self::Q1(b) => q1_row_exact(b, r, gpr, x),
8956            Self::Q1g { bytes, sc_off, .. } => q1g_row_exact(bytes, sc_off, r, gpr, x),
8957        }
8958    }
8959
8960    #[inline]
8961    fn outlier(self, r: usize, gpr: usize, j: usize) -> (f32, f32) {
8962        match self {
8963            Self::Q1(b) => q1_outlier(b, r, gpr, j),
8964            Self::Q1g { bytes, sc_off, .. } => q1g_outlier(bytes, sc_off, r, gpr, j),
8965        }
8966    }
8967
8968    /// Blocked 1×4 row (one unpack, four activation streams).
8969    ///
8970    /// # Safety
8971    /// The SDOT/dotprod target feature must be present (`sdot_enabled`).
8972    #[cfg(target_arch = "aarch64")]
8973    #[inline]
8974    unsafe fn dot_1x4_sdot(
8975        self,
8976        r: usize,
8977        gpr: usize,
8978        xs: [&[i8]; 4],
8979        gs: [&[i32]; 4],
8980    ) -> [f32; 4] {
8981        // SAFETY: forwarded caller contract.
8982        unsafe {
8983            match self {
8984                Self::Q1(b) => dot_q1_row_1x4_sdot(b, r, gpr, xs, gs),
8985                Self::Q1g {
8986                    bytes,
8987                    sc_off,
8988                    rows,
8989                } => dot_q1g_row_1x4_sdot(bytes, sc_off, rows, r, gpr, xs, gs),
8990            }
8991        }
8992    }
8993
8994    /// Blocked 1×4 row on x86 (VNNI when the tile switch is on).
8995    ///
8996    /// # Safety
8997    /// AVX2 must be present (`avx2_enabled`).
8998    #[cfg(target_arch = "x86_64")]
8999    #[inline]
9000    unsafe fn dot_1x4_x86(self, r: usize, gpr: usize, xs: [&[i8]; 4], gs: [&[i32]; 4]) -> [f32; 4] {
9001        // SAFETY: forwarded caller contract.
9002        unsafe {
9003            match (self, vnni_tiles_enabled()) {
9004                (Self::Q1(b), true) => dot_q1_row_1x4_vnni(b, r, gpr, xs, gs),
9005                (Self::Q1(b), false) => dot_q1_row_1x4_avx2(b, r, gpr, xs, gs),
9006                (Self::Q1g { bytes, sc_off, .. }, true) => {
9007                    dot_q1g_row_1x4_vnni(bytes, sc_off, r, gpr, xs, gs)
9008                }
9009                (Self::Q1g { bytes, sc_off, .. }, false) => {
9010                    dot_q1g_row_1x4_avx2(bytes, sc_off, r, gpr, xs, gs)
9011                }
9012            }
9013        }
9014    }
9015}
9016
9017// ── q1g (dtype 20): binary weights, ONE f16 scale per 128-group, in sign
9018// and scale PLANES (`[bits: rows·cols/8][f16: rows·cols/128]`). Same
9019// kernel family as q1 and — deliberately — the same float accumulation
9020// order (one `acc += d·s` per 32-tile, the group's scale reused four
9021// times), so a q1g repack of a q1 tensor whose tile scales agree gives
9022// the same f32 outputs bit for bit on every path. The win is the stream:
9023// 18 bytes per 128 weights instead of 24, the bits 16-byte aligned. ──
9024
9025/// Sign bits of q1g tile `t` (flat 32-tile index `r·gpr + gi`).
9026#[inline]
9027fn q1g_bits(bytes: &[u8], t: usize) -> u32 {
9028    u32::from_le_bytes([
9029        bytes[t * 4],
9030        bytes[t * 4 + 1],
9031        bytes[t * 4 + 2],
9032        bytes[t * 4 + 3],
9033    ])
9034}
9035
9036/// Raw f16 bits of the scale of the 128-group holding q1g tile `t`.
9037#[inline]
9038fn q1g_scale_bits(bytes: &[u8], sc_off: usize, t: usize) -> u16 {
9039    let o = sc_off + (t / 4) * 2;
9040    u16::from_le_bytes([bytes[o], bytes[o + 1]])
9041}
9042
9043/// One q1g row via the A8W8 int8 path (dispatch twin of `dot_q1_row_i8`).
9044#[allow(unreachable_code)]
9045#[allow(clippy::too_many_arguments)]
9046fn dot_q1g_row_i8(
9047    bytes: &[u8],
9048    sc_off: usize,
9049    rows: usize,
9050    r: usize,
9051    gpr: usize,
9052    xq: &[i8],
9053    gsum: &[i32],
9054) -> f32 {
9055    #[cfg(target_arch = "aarch64")]
9056    unsafe {
9057        return dot_q1g_row_sdot(bytes, sc_off, rows, r, gpr, xq, gsum);
9058    }
9059    #[cfg(target_arch = "x86_64")]
9060    if avx2_enabled() {
9061        unsafe {
9062            if vnni_tiles_enabled() {
9063                return dot_q1g_row_vnni(bytes, sc_off, r, gpr, xq, gsum);
9064            }
9065            return dot_q1g_row_avx2(bytes, sc_off, r, gpr, xq, gsum);
9066        }
9067    }
9068    let _ = (gsum, rows);
9069    let mut acc = 0f32;
9070    for gi in 0..gpr {
9071        let t = r * gpr + gi;
9072        let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, t));
9073        let bits = q1g_bits(bytes, t);
9074        let mut d = 0i32;
9075        for k in 0..GROUP_SIZE {
9076            let w = ((bits >> k) & 1) as i32 * 2 - 1;
9077            d += w * xq[gi * GROUP_SIZE + k] as i32;
9078        }
9079        acc += d as f32 * s;
9080    }
9081    acc
9082}
9083
9084/// Four tiles' unreduced `i32×8` partial sums → `[Σ0, Σ1, Σ2, Σ3]` in one
9085/// hadd tree: 3 hadds + 1 add for four tiles where the per-tile reduce of
9086/// the q1 kernels costs 7 ops each. Integer math, so any order gives the
9087/// q1 kernel's per-tile sums exactly.
9088#[cfg(target_arch = "x86_64")]
9089#[target_feature(enable = "avx2")]
9090#[inline]
9091unsafe fn q1g_hsum4(
9092    d0: core::arch::x86_64::__m256i,
9093    d1: core::arch::x86_64::__m256i,
9094    d2: core::arch::x86_64::__m256i,
9095    d3: core::arch::x86_64::__m256i,
9096) -> core::arch::x86_64::__m128i {
9097    // SAFETY: pure register math.
9098    unsafe {
9099        use core::arch::x86_64::*;
9100        let h = _mm256_hadd_epi32(_mm256_hadd_epi32(d0, d1), _mm256_hadd_epi32(d2, d3));
9101        _mm_add_epi32(_mm256_castsi256_si128(h), _mm256_extracti128_si256::<1>(h))
9102    }
9103}
9104
9105/// One 128-group's epilogue, in q1's float order: `d = 2·Σset − Σx` per
9106/// tile (exact i32 → exact f32), then `acc += d·s` tile by tile — four
9107/// separately rounded multiply/add pairs, the q1 kernels' sequence.
9108#[cfg(target_arch = "x86_64")]
9109#[target_feature(enable = "avx2")]
9110#[inline]
9111unsafe fn q1g_group_acc(
9112    acc: f32,
9113    msum4: core::arch::x86_64::__m128i,
9114    gsum4: &[i32],
9115    s: f32,
9116) -> f32 {
9117    // SAFETY: gsum4 holds the group's four tile sums (callers slice it).
9118    unsafe {
9119        use core::arch::x86_64::*;
9120        let g = _mm_loadu_si128(gsum4.as_ptr() as *const __m128i);
9121        let d = _mm_cvtepi32_ps(_mm_sub_epi32(_mm_slli_epi32::<1>(msum4), g));
9122        let mut f = [0f32; 4];
9123        _mm_storeu_ps(f.as_mut_ptr(), d);
9124        let mut acc = acc;
9125        for v in f {
9126            acc += v * s;
9127        }
9128        acc
9129    }
9130}
9131
9132/// AVX2 q1g row: `dot_q1_row_avx2`'s mask/select/maddubs per tile, but the
9133/// four tiles of a 128-group reduce together (`q1g_hsum4`) and the scale
9134/// converts once per group — bit-identical results (integer sums, then
9135/// q1's per-tile float sequence), fewer horizontal ops and f16 decodes.
9136#[cfg(target_arch = "x86_64")]
9137#[target_feature(enable = "avx2")]
9138unsafe fn dot_q1g_row_avx2(
9139    bytes: &[u8],
9140    sc_off: usize,
9141    r: usize,
9142    gpr: usize,
9143    xq: &[i8],
9144    gsum: &[i32],
9145) -> f32 {
9146    // SAFETY: callers uphold the plane-length and xq/gsum contracts
9147    // (gpr % 4 == 0, xq.len() == gpr·32, gsum.len() == gpr).
9148    unsafe {
9149        use core::arch::x86_64::*;
9150        let expand = _mm256_setr_epi8(
9151            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,
9152            3, 3, 3,
9153        );
9154        let bitsel = _mm256_setr_epi8(
9155            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
9156            -128, 1, 2, 4, 8, 16, 32, 64, -128,
9157        );
9158        let ones8 = _mm256_set1_epi8(1);
9159        let ones16 = _mm256_set1_epi16(1);
9160        let bp = bytes.as_ptr().add(r * gpr * 4);
9161        let mut acc = 0f32;
9162        let mut gi = 0;
9163        while gi < gpr {
9164            let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
9165            let mut d = [_mm256_setzero_si256(); 4];
9166            for (k, dk) in d.iter_mut().enumerate() {
9167                let t = gi + k;
9168                let bits = (bp.add(t * 4) as *const u32).read_unaligned();
9169                let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
9170                let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
9171                let x = _mm256_loadu_si256(xq.as_ptr().add(t * GROUP_SIZE) as *const __m256i);
9172                let sel = _mm256_and_si256(x, mask);
9173                *dk = _mm256_madd_epi16(_mm256_maddubs_epi16(ones8, sel), ones16);
9174            }
9175            let m4 = q1g_hsum4(d[0], d[1], d[2], d[3]);
9176            acc = q1g_group_acc(acc, m4, &gsum[gi..gi + 4], s);
9177            gi += 4;
9178        }
9179        acc
9180    }
9181}
9182
9183/// VNNI twin of `dot_q1g_row_avx2`: `vpdpbusd(1u8, sel)` per tile (the
9184/// same i32 lanes as maddubs+madd — bounded products), same group reduce.
9185#[cfg(target_arch = "x86_64")]
9186#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
9187unsafe fn dot_q1g_row_vnni(
9188    bytes: &[u8],
9189    sc_off: usize,
9190    r: usize,
9191    gpr: usize,
9192    xq: &[i8],
9193    gsum: &[i32],
9194) -> f32 {
9195    // SAFETY: callers uphold the plane-length and xq/gsum contracts.
9196    unsafe {
9197        use core::arch::x86_64::*;
9198        let expand = _mm256_setr_epi8(
9199            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,
9200            3, 3, 3,
9201        );
9202        let bitsel = _mm256_setr_epi8(
9203            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
9204            -128, 1, 2, 4, 8, 16, 32, 64, -128,
9205        );
9206        let ones8 = _mm256_set1_epi8(1);
9207        let bp = bytes.as_ptr().add(r * gpr * 4);
9208        let mut acc = 0f32;
9209        let mut gi = 0;
9210        while gi < gpr {
9211            let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
9212            let mut d = [_mm256_setzero_si256(); 4];
9213            for (k, dk) in d.iter_mut().enumerate() {
9214                let t = gi + k;
9215                let bits = (bp.add(t * 4) as *const u32).read_unaligned();
9216                let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
9217                let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
9218                let x = _mm256_loadu_si256(xq.as_ptr().add(t * GROUP_SIZE) as *const __m256i);
9219                *dk = _mm256_dpbusd_epi32(_mm256_setzero_si256(), ones8, _mm256_and_si256(x, mask));
9220            }
9221            let m4 = q1g_hsum4(d[0], d[1], d[2], d[3]);
9222            acc = q1g_group_acc(acc, m4, &gsum[gi..gi + 4], s);
9223            gi += 4;
9224        }
9225        acc
9226    }
9227}
9228
9229/// Blocked 1×4 AVX2 q1g row: each tile's mask is built once and serves
9230/// four activation streams; per stream, the group reduce and float order
9231/// of `dot_q1g_row_avx2` (so batch == matvec bit for bit).
9232#[cfg(target_arch = "x86_64")]
9233#[target_feature(enable = "avx2")]
9234unsafe fn dot_q1g_row_1x4_avx2(
9235    bytes: &[u8],
9236    sc_off: usize,
9237    r: usize,
9238    gpr: usize,
9239    xs: [&[i8]; 4],
9240    gsums: [&[i32]; 4],
9241) -> [f32; 4] {
9242    // SAFETY: callers uphold the plane-length and xq/gsum contracts.
9243    unsafe {
9244        use core::arch::x86_64::*;
9245        let expand = _mm256_setr_epi8(
9246            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,
9247            3, 3, 3,
9248        );
9249        let bitsel = _mm256_setr_epi8(
9250            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
9251            -128, 1, 2, 4, 8, 16, 32, 64, -128,
9252        );
9253        let ones8 = _mm256_set1_epi8(1);
9254        let ones16 = _mm256_set1_epi16(1);
9255        let bp = bytes.as_ptr().add(r * gpr * 4);
9256        let mut acc = [0f32; 4];
9257        let mut gi = 0;
9258        while gi < gpr {
9259            let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
9260            // d[stream][tile]
9261            let mut d = [[_mm256_setzero_si256(); 4]; 4];
9262            for t in 0..4 {
9263                let bits = (bp.add((gi + t) * 4) as *const u32).read_unaligned();
9264                let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
9265                let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
9266                for (k, xq) in xs.iter().enumerate() {
9267                    let x = _mm256_loadu_si256(
9268                        xq.as_ptr().add((gi + t) * GROUP_SIZE) as *const __m256i
9269                    );
9270                    let sel = _mm256_and_si256(x, mask);
9271                    d[k][t] = _mm256_madd_epi16(_mm256_maddubs_epi16(ones8, sel), ones16);
9272                }
9273            }
9274            for k in 0..4 {
9275                let m4 = q1g_hsum4(d[k][0], d[k][1], d[k][2], d[k][3]);
9276                acc[k] = q1g_group_acc(acc[k], m4, &gsums[k][gi..gi + 4], s);
9277            }
9278            gi += 4;
9279        }
9280        acc
9281    }
9282}
9283
9284/// VNNI twin of `dot_q1g_row_1x4_avx2`.
9285#[cfg(target_arch = "x86_64")]
9286#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
9287unsafe fn dot_q1g_row_1x4_vnni(
9288    bytes: &[u8],
9289    sc_off: usize,
9290    r: usize,
9291    gpr: usize,
9292    xs: [&[i8]; 4],
9293    gsums: [&[i32]; 4],
9294) -> [f32; 4] {
9295    // SAFETY: callers uphold the plane-length and xq/gsum contracts.
9296    unsafe {
9297        use core::arch::x86_64::*;
9298        let expand = _mm256_setr_epi8(
9299            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,
9300            3, 3, 3,
9301        );
9302        let bitsel = _mm256_setr_epi8(
9303            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
9304            -128, 1, 2, 4, 8, 16, 32, 64, -128,
9305        );
9306        let ones8 = _mm256_set1_epi8(1);
9307        let bp = bytes.as_ptr().add(r * gpr * 4);
9308        let mut acc = [0f32; 4];
9309        let mut gi = 0;
9310        while gi < gpr {
9311            let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
9312            let mut d = [[_mm256_setzero_si256(); 4]; 4];
9313            for t in 0..4 {
9314                let bits = (bp.add((gi + t) * 4) as *const u32).read_unaligned();
9315                let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
9316                let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
9317                for (k, xq) in xs.iter().enumerate() {
9318                    let x = _mm256_loadu_si256(
9319                        xq.as_ptr().add((gi + t) * GROUP_SIZE) as *const __m256i
9320                    );
9321                    d[k][t] = _mm256_dpbusd_epi32(
9322                        _mm256_setzero_si256(),
9323                        ones8,
9324                        _mm256_and_si256(x, mask),
9325                    );
9326                }
9327            }
9328            for k in 0..4 {
9329                let m4 = q1g_hsum4(d[k][0], d[k][1], d[k][2], d[k][3]);
9330                acc[k] = q1g_group_acc(acc[k], m4, &gsums[k][gi..gi + 4], s);
9331            }
9332            gi += 4;
9333        }
9334        acc
9335    }
9336}
9337
9338/// SDOT q1g row: the `dot_q1_row_sdot` math on the plane layout. One
9339/// aligned 16-byte load carries a whole 128-group's signs (q1 needs two
9340/// pair loads and a TBL gather of four scales for the same weights), TBL
9341/// replicates each sign byte across 8 lanes for vtst, and the group's ONE
9342/// scale converts once (FCVTL, exact) and broadcasts. Lane order, the
9343/// vpaddq tree, the fma into the four lane accumulators and the tail are
9344/// the q1 kernel's — including its quirk: q1's pair load reads 4 bytes
9345/// past the 4th tile, so the payload's final four tiles (last row) take
9346/// the scalar tail there. q1g has no overread, but takes the same tail on
9347/// the last row to stay bit-identical to q1.
9348#[cfg(target_arch = "aarch64")]
9349#[target_feature(enable = "neon,dotprod")]
9350#[allow(clippy::too_many_arguments)]
9351unsafe fn dot_q1g_row_sdot(
9352    bytes: &[u8],
9353    sc_off: usize,
9354    rows: usize,
9355    r: usize,
9356    gpr: usize,
9357    xq: &[i8],
9358    gsum: &[i32],
9359) -> f32 {
9360    // SAFETY: callers uphold the plane-length and xq/gsum contracts
9361    // (gpr % 4 == 0, xq.len() == gpr·32, gsum.len() == gpr).
9362    unsafe {
9363        use core::arch::aarch64::*;
9364        use core::arch::asm;
9365        const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
9366        // Tile t of the group = sign bytes 4t..4t+4 of the 16-byte load.
9367        const IW: [[u8; 16]; 8] = [
9368            [0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1],
9369            [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3],
9370            [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5],
9371            [6, 6, 6, 6, 6, 6, 6, 6, 7, 7, 7, 7, 7, 7, 7, 7],
9372            [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9],
9373            [
9374                10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
9375            ],
9376            [
9377                12, 12, 12, 12, 12, 12, 12, 12, 13, 13, 13, 13, 13, 13, 13, 13,
9378            ],
9379            [
9380                14, 14, 14, 14, 14, 14, 14, 14, 15, 15, 15, 15, 15, 15, 15, 15,
9381            ],
9382        ];
9383        let m = vld1q_u8(MASKS.as_ptr());
9384        let iw: [uint8x16_t; 8] = core::array::from_fn(|i| vld1q_u8(IW[i].as_ptr()));
9385        macro_rules! tile_dot {
9386            ($w0:expr, $w1:expr, $x:expr) => {{
9387                let x0 = vld1q_s8($x);
9388                let x1 = vld1q_s8($x.add(16));
9389                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
9390                asm!(
9391                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
9392                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
9393                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
9394                    w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
9395                    options(pure, nomem, nostack),
9396                );
9397                vaddq_s32(a0, a1)
9398            }};
9399        }
9400        macro_rules! tbl_mask {
9401            ($ld:expr, $i:expr) => {
9402                vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, iw[$i]), m))
9403            };
9404        }
9405        let bp = bytes.as_ptr().add(r * gpr * 4);
9406        let xp = xq.as_ptr();
9407        let gp = gsum.as_ptr();
9408        let vec_end = if r + 1 == rows { gpr - 4 } else { gpr };
9409        let mut accv = vdupq_n_f32(0.0);
9410        let mut gi = 0;
9411        while gi + 4 <= vec_end {
9412            let ld = vld1q_u8(bp.add(gi * 4));
9413            let d0 = tile_dot!(tbl_mask!(ld, 0), tbl_mask!(ld, 1), xp.add(gi * GROUP_SIZE));
9414            let d1 = tile_dot!(
9415                tbl_mask!(ld, 2),
9416                tbl_mask!(ld, 3),
9417                xp.add((gi + 1) * GROUP_SIZE)
9418            );
9419            let d2 = tile_dot!(
9420                tbl_mask!(ld, 4),
9421                tbl_mask!(ld, 5),
9422                xp.add((gi + 2) * GROUP_SIZE)
9423            );
9424            let d3 = tile_dot!(
9425                tbl_mask!(ld, 6),
9426                tbl_mask!(ld, 7),
9427                xp.add((gi + 3) * GROUP_SIZE)
9428            );
9429            let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
9430            let g = vld1q_s32(gp.add(gi));
9431            let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
9432            let h = vdup_n_u16(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
9433            let scf: float32x4_t;
9434            asm!(
9435                "fcvtl {o:v}.4s, {i:v}.4h",
9436                o = out(vreg) scf, i = in(vreg) h,
9437                options(pure, nomem, nostack),
9438            );
9439            accv = vfmaq_f32(accv, vcvtq_f32_s32(dots), scf);
9440            gi += 4;
9441        }
9442        let mut acc = vaddvq_f32(accv);
9443        while gi < gpr {
9444            let t = bp.add(gi * 4);
9445            let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
9446            let v0 = vcombine_u8(vdup_n_u8(*t), vdup_n_u8(*t.add(1)));
9447            let v1 = vcombine_u8(vdup_n_u8(*t.add(2)), vdup_n_u8(*t.add(3)));
9448            let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
9449            let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
9450            let d = vaddvq_s32(tile_dot!(w0, w1, xp.add(gi * GROUP_SIZE)));
9451            acc += (-(2 * d + *gp.add(gi))) as f32 * s;
9452            gi += 1;
9453        }
9454        acc
9455    }
9456}
9457
9458/// Blocked q1g 1×4: one unpack of the group serves four activation
9459/// streams; per stream bit-identical to `dot_q1g_row_sdot` (and so to the
9460/// q1 1×4 kernel, same last-row tail).
9461#[cfg(target_arch = "aarch64")]
9462#[target_feature(enable = "neon,dotprod")]
9463#[allow(clippy::too_many_arguments)]
9464unsafe fn dot_q1g_row_1x4_sdot(
9465    bytes: &[u8],
9466    sc_off: usize,
9467    rows: usize,
9468    r: usize,
9469    gpr: usize,
9470    xs: [&[i8]; 4],
9471    gs: [&[i32]; 4],
9472) -> [f32; 4] {
9473    // SAFETY: same contracts as `dot_q1g_row_sdot`, ×4.
9474    unsafe {
9475        use core::arch::aarch64::*;
9476        use core::arch::asm;
9477        const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
9478        const IW: [[u8; 16]; 8] = [
9479            [0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1],
9480            [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3],
9481            [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5],
9482            [6, 6, 6, 6, 6, 6, 6, 6, 7, 7, 7, 7, 7, 7, 7, 7],
9483            [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9],
9484            [
9485                10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
9486            ],
9487            [
9488                12, 12, 12, 12, 12, 12, 12, 12, 13, 13, 13, 13, 13, 13, 13, 13,
9489            ],
9490            [
9491                14, 14, 14, 14, 14, 14, 14, 14, 15, 15, 15, 15, 15, 15, 15, 15,
9492            ],
9493        ];
9494        let m = vld1q_u8(MASKS.as_ptr());
9495        let iw: [uint8x16_t; 8] = core::array::from_fn(|i| vld1q_u8(IW[i].as_ptr()));
9496        macro_rules! sdot2 {
9497            ($w0:expr, $w1:expr, $x:expr) => {{
9498                let x0 = vld1q_s8($x);
9499                let x1 = vld1q_s8($x.add(16));
9500                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
9501                asm!(
9502                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
9503                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
9504                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
9505                    w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
9506                    options(pure, nomem, nostack),
9507                );
9508                vaddq_s32(a0, a1)
9509            }};
9510        }
9511        let bp = bytes.as_ptr().add(r * gpr * 4);
9512        let vec_end = if r + 1 == rows { gpr - 4 } else { gpr };
9513        let mut accv = [vdupq_n_f32(0.0); 4];
9514        let mut gi = 0;
9515        while gi + 4 <= vec_end {
9516            let ld = vld1q_u8(bp.add(gi * 4));
9517            // Unpack ONCE — eight ±mask vectors serve all four streams.
9518            let w: [int8x16_t; 8] =
9519                core::array::from_fn(|i| vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld, iw[i]), m)));
9520            let h = vdup_n_u16(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
9521            let scf: float32x4_t;
9522            asm!(
9523                "fcvtl {o:v}.4s, {i:v}.4h",
9524                o = out(vreg) scf, i = in(vreg) h,
9525                options(pure, nomem, nostack),
9526            );
9527            for k in 0..4 {
9528                let xp = xs[k].as_ptr();
9529                let d0 = sdot2!(w[0], w[1], xp.add(gi * GROUP_SIZE));
9530                let d1 = sdot2!(w[2], w[3], xp.add((gi + 1) * GROUP_SIZE));
9531                let d2 = sdot2!(w[4], w[5], xp.add((gi + 2) * GROUP_SIZE));
9532                let d3 = sdot2!(w[6], w[7], xp.add((gi + 3) * GROUP_SIZE));
9533                let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
9534                let g = vld1q_s32(gs[k].as_ptr().add(gi));
9535                let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
9536                accv[k] = vfmaq_f32(accv[k], vcvtq_f32_s32(dots), scf);
9537            }
9538            gi += 4;
9539        }
9540        let mut acc = [
9541            vaddvq_f32(accv[0]),
9542            vaddvq_f32(accv[1]),
9543            vaddvq_f32(accv[2]),
9544            vaddvq_f32(accv[3]),
9545        ];
9546        while gi < gpr {
9547            let t = bp.add(gi * 4);
9548            let sc = f16_to_f32(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
9549            let v0 = vcombine_u8(vdup_n_u8(*t), vdup_n_u8(*t.add(1)));
9550            let v1 = vcombine_u8(vdup_n_u8(*t.add(2)), vdup_n_u8(*t.add(3)));
9551            let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
9552            let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
9553            for k in 0..4 {
9554                let d = vaddvq_s32(sdot2!(w0, w1, xs[k].as_ptr().add(gi * GROUP_SIZE)));
9555                acc[k] += (-(2 * d + *gs[k].as_ptr().add(gi))) as f32 * sc;
9556            }
9557            gi += 1;
9558        }
9559        acc
9560    }
9561}
9562
9563/// (weight ±1, scale) of one q1g element — the exact outlier term.
9564#[inline]
9565fn q1g_outlier(bytes: &[u8], sc_off: usize, r: usize, gpr: usize, j: usize) -> (f32, f32) {
9566    let t = r * gpr + j / GROUP_SIZE;
9567    let k = j % GROUP_SIZE;
9568    let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, t));
9569    let bit = (bytes[t * 4 + k / 8] >> (k % 8)) & 1;
9570    ((bit as i32 * 2 - 1) as f32, s)
9571}
9572
9573/// Exact scalar q1g row (CMF_SDOT=0 contract) — `q1_row_exact`'s order.
9574#[inline]
9575fn q1g_row_exact(bytes: &[u8], sc_off: usize, r: usize, gpr: usize, x: &[f32]) -> f32 {
9576    let mut acc = 0f32;
9577    for gi in 0..gpr {
9578        let t = r * gpr + gi;
9579        let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, t));
9580        let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
9581        let mut ga = 0f32;
9582        for (j, &b) in bytes[t * 4..t * 4 + 4].iter().enumerate() {
9583            for k in 0..8 {
9584                ga += (((b >> k) & 1) as f32 * 2.0 - 1.0) * xg[j * 8 + k];
9585            }
9586        }
9587        acc += ga * s;
9588    }
9589    acc
9590}
9591
9592/// q1t per-row overlay locator. After the base (`base_len`) come
9593/// `[u32 row_ptr[rows+1]]` then `[(u16 col, f16 val)]` grouped by row (row
9594/// `r`'s entries are `[row_ptr[r], row_ptr[r+1])`). Returns
9595/// `(row_ptr offset, entries offset, present)`.
9596fn q1t_overlay(bytes: &[u8], base_len: usize, rows: usize) -> (usize, usize, bool) {
9597    let entries = base_len + (rows + 1) * 4;
9598    (base_len, entries, entries <= bytes.len())
9599}
9600
9601/// Read `row_ptr[r]` from the overlay's prefix-sum table.
9602#[inline]
9603fn q1t_rowptr(bytes: &[u8], rp_off: usize, r: usize) -> usize {
9604    let o = rp_off + r * 4;
9605    u32::from_le_bytes([bytes[o], bytes[o + 1], bytes[o + 2], bytes[o + 3]]) as usize
9606}
9607
9608/// Byte → the 5 ternary signs it packs `{−1,0,+1}` as f32, precomputed so
9609/// decoding a q1t code is a table load, not the base-3 divide/modulo per
9610/// weight (division is ~20–40× the cost of a load). Built at compile time.
9611const SIGN5: [[f32; 5]; 256] = {
9612    let mut lut = [[0.0f32; 5]; 256];
9613    let pow3 = [1u16, 3, 9, 27, 81];
9614    let mut byte = 0usize;
9615    while byte < 256 {
9616        let mut i = 0usize;
9617        while i < 5 {
9618            let code = (byte as u16 / pow3[i]) % 3;
9619            lut[byte][i] = if code == 1 {
9620                1.0
9621            } else if code == 2 {
9622                -1.0
9623            } else {
9624                0.0
9625            };
9626            i += 1;
9627        }
9628        byte += 1;
9629    }
9630    lut
9631};
9632
9633/// Same table, as i8 signs — the operand for the int8 SDOT base kernel.
9634const SIGN5_I8: [[i8; 5]; 256] = {
9635    let mut lut = [[0i8; 5]; 256];
9636    let pow3 = [1u16, 3, 9, 27, 81];
9637    let mut byte = 0usize;
9638    while byte < 256 {
9639        let mut i = 0usize;
9640        while i < 5 {
9641            let code = (byte as u16 / pow3[i]) % 3;
9642            lut[byte][i] = if code == 1 {
9643                1
9644            } else if code == 2 {
9645                -1
9646            } else {
9647                0
9648            };
9649            i += 1;
9650        }
9651        byte += 1;
9652    }
9653    lut
9654};
9655
9656/// The same 5 i8 signs packed into a u64 (`[s0 s1 s2 s3 s4 0 0 0]`, LE) so the
9657/// group unpack is 7 unaligned u64 stores at offsets 0,5,10,…,30 instead of
9658/// six 5-byte copies + LUT indexing — each store's trailing zeros are fixed by
9659/// the next store, and the last one runs 6 B past the 32nd weight (the unpack
9660/// buffer is padded to 40). This is the decode/prefill hot inner op.
9661const SIGN5_U64: [u64; 256] = {
9662    let mut lut = [0u64; 256];
9663    let pow3 = [1u16, 3, 9, 27, 81];
9664    let mut byte = 0usize;
9665    while byte < 256 {
9666        let mut v = 0u64;
9667        let mut i = 0usize;
9668        while i < 5 {
9669            let code = (byte as u16 / pow3[i]) % 3;
9670            let s: u8 = if code == 1 {
9671                1
9672            } else if code == 2 {
9673                0xFF
9674            } else {
9675                0
9676            };
9677            v |= (s as u64) << (i * 8);
9678            i += 1;
9679        }
9680        lut[byte] = v;
9681        byte += 1;
9682    }
9683    lut
9684};
9685
9686/// Ternary base weight at `(row r, col j)` = `sign(code)·s_group`. Used to add
9687/// back activation-outlier columns, whose `x` was zeroed for the int8 bulk dot
9688/// (`split_act`). At a weight-outlier position the code is 0, so this is 0 and
9689/// the overlay correction owns that column — no double counting.
9690#[inline]
9691fn q1t_base_weight(bytes: &[u8], r: usize, gpr: usize, j: usize) -> f32 {
9692    const TILE: usize = cortiq_core::quant::Q1T_TILE;
9693    let off = (r * gpr + j / GROUP_SIZE) * TILE;
9694    let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
9695    let within = j % GROUP_SIZE;
9696    SIGN5[bytes[off + 2 + within / 5] as usize][within % 5] * s
9697}
9698
9699/// One 32-group int8 dot via two SDOTs. Bit-exact vs the scalar i8 sum
9700/// (integer accumulation is order-independent).
9701#[cfg(target_arch = "aarch64")]
9702#[target_feature(enable = "neon,dotprod")]
9703#[inline]
9704unsafe fn sdot32_i8(w: *const i8, x: *const i8) -> i32 {
9705    // SAFETY: caller guarantees 32 readable i8 at each pointer.
9706    unsafe {
9707        use core::arch::aarch64::*;
9708        use core::arch::asm;
9709        let w0 = vld1q_s8(w);
9710        let w1 = vld1q_s8(w.add(16));
9711        let x0 = vld1q_s8(x);
9712        let x1 = vld1q_s8(x.add(16));
9713        let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
9714        asm!(
9715            "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
9716            "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
9717            a0 = inout(vreg) a0, a1 = inout(vreg) a1,
9718            w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
9719            options(pure, nomem, nostack),
9720        );
9721        vaddvq_s32(vaddq_s32(a0, a1))
9722    }
9723}
9724
9725/// One 32-group int8 dot via AVX2: signed·signed as `maddubs(|w|, sign(x,w))`
9726/// then `madd` and a horizontal reduce (the same idiom as `dot_q4t_row_avx2`).
9727#[cfg(target_arch = "x86_64")]
9728#[target_feature(enable = "avx2")]
9729#[inline]
9730unsafe fn i8dot32_avx2(w: *const i8, x: *const i8) -> i32 {
9731    // SAFETY: caller guarantees 32 readable i8 at each pointer.
9732    unsafe {
9733        use core::arch::x86_64::*;
9734        let wv = _mm256_loadu_si256(w as *const __m256i);
9735        let xv = _mm256_loadu_si256(x as *const __m256i);
9736        let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
9737        let d = _mm256_madd_epi16(p16, _mm256_set1_epi16(1));
9738        let hi128 = _mm256_extracti128_si256::<1>(d);
9739        let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
9740        let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
9741        let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
9742        _mm_cvtsi128_si32(s32)
9743    }
9744}
9745
9746/// Unpack one q1t group's base-3 codes into 32 i8 signs via 7 unaligned u64
9747/// stores (see `SIGN5_U64`). `dst` MUST have ≥ 40 bytes: the 7th store writes
9748/// `dst[30..38]`. Stores go in order so each one's trailing zeros are
9749/// overwritten by the next; the final 6 padding bytes are unused by the dot.
9750#[inline]
9751fn q1t_unpack_group_i8(codes: *const u8, dst: &mut [i8]) {
9752    debug_assert!(dst.len() >= 40);
9753    // SAFETY: codes points at 7 readable bytes; dst has ≥ 40 bytes so every
9754    // 8-byte store at offset bi*5 (bi ≤ 6 → ≤ 30) stays in bounds.
9755    unsafe {
9756        let p = dst.as_mut_ptr();
9757        for bi in 0..7 {
9758            core::ptr::write_unaligned(
9759                p.add(bi * 5) as *mut u64,
9760                SIGN5_U64[*codes.add(bi) as usize],
9761            );
9762        }
9763    }
9764}
9765
9766/// One 32-group int8 dot, arch-dispatched (the matmat inner loop, where the
9767/// row's signs are unpacked once and dotted against every batch input).
9768/// Callers are gated by `a8w8_enabled()`, so the target-feature arms are
9769/// reachable; the scalar arm is a non-SIMD-arch fallback.
9770#[inline]
9771fn q1t_i8dot32(w: *const i8, x: *const i8) -> i32 {
9772    #[cfg(target_arch = "aarch64")]
9773    unsafe {
9774        return sdot32_i8(w, x);
9775    }
9776    #[cfg(target_arch = "x86_64")]
9777    unsafe {
9778        return i8dot32_avx2(w, x);
9779    }
9780    #[allow(unreachable_code)]
9781    unsafe {
9782        let mut s = 0i32;
9783        for k in 0..GROUP_SIZE {
9784            s += *w.add(k) as i32 * *x.add(k) as i32;
9785        }
9786        s
9787    }
9788}
9789
9790#[inline]
9791unsafe fn q1t_unpack_reg_u64s(codes: *const u8) -> (u64, u64, u64, u64) {
9792    let (s0, s1, s2, s3, s4, s5, s6) = unsafe {
9793        (
9794            SIGN5_U64[*codes as usize],
9795            SIGN5_U64[*codes.add(1) as usize],
9796            SIGN5_U64[*codes.add(2) as usize],
9797            SIGN5_U64[*codes.add(3) as usize],
9798            SIGN5_U64[*codes.add(4) as usize],
9799            SIGN5_U64[*codes.add(5) as usize],
9800            SIGN5_U64[*codes.add(6) as usize],
9801        )
9802    };
9803
9804    let u0 = s0 | (s1 << 40);
9805    let u1 = (s1 >> 24) | (s2 << 16) | (s3 << 56);
9806    let u2 = (s3 >> 8) | (s4 << 32);
9807    let u3 = (s4 >> 32) | (s5 << 8) | (s6 << 48);
9808
9809    (u0, u1, u2, u3)
9810}
9811
9812/// One q1t row's int8 base dot: `Σ_group s·dot(signs, xq)` (before the shared
9813/// `sx`). Direct register unpacking (zero stack stores/loads, no STLF stalls).
9814/// ARM SDOT.
9815#[cfg(target_arch = "aarch64")]
9816#[target_feature(enable = "neon,dotprod")]
9817unsafe fn q1t_dot_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
9818    use core::arch::aarch64::*;
9819    use core::arch::asm;
9820    unsafe {
9821        const TILE: usize = cortiq_core::quant::Q1T_TILE;
9822        let mut acc = 0f32;
9823        let bytes_ptr = bytes.as_ptr();
9824        let xq_ptr = xq.as_ptr();
9825        let row_off = r * gpr * TILE;
9826
9827        let gpr2 = gpr & !1;
9828        let mut gi = 0;
9829        while gi < gpr2 {
9830            let off0 = row_off + gi * TILE;
9831            let off1 = off0 + TILE;
9832            let s0 = f16_to_f32(u16::from_le_bytes([
9833                *bytes_ptr.add(off0),
9834                *bytes_ptr.add(off0 + 1),
9835            ]));
9836            let s1 = f16_to_f32(u16::from_le_bytes([
9837                *bytes_ptr.add(off1),
9838                *bytes_ptr.add(off1 + 1),
9839            ]));
9840
9841            let (u0_0, u1_0, u2_0, u3_0) = q1t_unpack_reg_u64s(bytes_ptr.add(off0 + 2));
9842            let (u0_1, u1_1, u2_1, u3_1) = q1t_unpack_reg_u64s(bytes_ptr.add(off1 + 2));
9843
9844            let w0_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_0), vcreate_u64(u1_0)));
9845            let w1_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_0), vcreate_u64(u3_0)));
9846            let w0_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_1), vcreate_u64(u1_1)));
9847            let w1_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_1), vcreate_u64(u3_1)));
9848
9849            let x0_0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE));
9850            let x1_0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE + 16));
9851            let x0_1 = vld1q_s8(xq_ptr.add((gi + 1) * GROUP_SIZE));
9852            let x1_1 = vld1q_s8(xq_ptr.add((gi + 1) * GROUP_SIZE + 16));
9853
9854            let (mut a0_0, mut a1_0) = (vdupq_n_s32(0), vdupq_n_s32(0));
9855            let (mut a0_1, mut a1_1) = (vdupq_n_s32(0), vdupq_n_s32(0));
9856            asm!(
9857                "sdot {a0_0:v}.4s, {w0_0:v}.16b, {x0_0:v}.16b",
9858                "sdot {a1_0:v}.4s, {w1_0:v}.16b, {x1_0:v}.16b",
9859                "sdot {a0_1:v}.4s, {w0_1:v}.16b, {x0_1:v}.16b",
9860                "sdot {a1_1:v}.4s, {w1_1:v}.16b, {x1_1:v}.16b",
9861                a0_0 = inout(vreg) a0_0, a1_0 = inout(vreg) a1_0,
9862                a0_1 = inout(vreg) a0_1, a1_1 = inout(vreg) a1_1,
9863                w0_0 = in(vreg) w0_0, x0_0 = in(vreg) x0_0, w1_0 = in(vreg) w1_0, x1_0 = in(vreg) x1_0,
9864                w0_1 = in(vreg) w0_1, x0_1 = in(vreg) x0_1, w1_1 = in(vreg) w1_1, x1_1 = in(vreg) x1_1,
9865                options(pure, nomem, nostack),
9866            );
9867            let d0 = vaddvq_s32(vaddq_s32(a0_0, a1_0));
9868            let d1 = vaddvq_s32(vaddq_s32(a0_1, a1_1));
9869            acc += d0 as f32 * s0 + d1 as f32 * s1;
9870            gi += 2;
9871        }
9872
9873        if gi < gpr {
9874            let off = row_off + gi * TILE;
9875            let s = f16_to_f32(u16::from_le_bytes([
9876                *bytes_ptr.add(off),
9877                *bytes_ptr.add(off + 1),
9878            ]));
9879            let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
9880            let w0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0), vcreate_u64(u1)));
9881            let w1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2), vcreate_u64(u3)));
9882            let x0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE));
9883            let x1 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE + 16));
9884            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
9885            asm!(
9886                "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
9887                "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
9888                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
9889                w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
9890                options(pure, nomem, nostack),
9891            );
9892            let d = vaddvq_s32(vaddq_s32(a0, a1));
9893            acc += d as f32 * s;
9894        }
9895        acc
9896    }
9897}
9898
9899/// x86 AVX2 mirror of `q1t_dot_row_sdot` (maddubs int8 dot per group).
9900#[cfg(target_arch = "x86_64")]
9901#[target_feature(enable = "avx2")]
9902unsafe fn q1t_dot_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
9903    use core::arch::x86_64::*;
9904    unsafe {
9905        const TILE: usize = cortiq_core::quant::Q1T_TILE;
9906        let mut acc = 0f32;
9907        let bytes_ptr = bytes.as_ptr();
9908        let xq_ptr = xq.as_ptr();
9909        let row_off = r * gpr * TILE;
9910
9911        let ones = _mm256_set1_epi16(1);
9912        for gi in 0..gpr {
9913            let off = row_off + gi * TILE;
9914            let s = f16_to_f32(u16::from_le_bytes([
9915                *bytes_ptr.add(off),
9916                *bytes_ptr.add(off + 1),
9917            ]));
9918            let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
9919            let wv = _mm256_set_epi64x(u3 as i64, u2 as i64, u1 as i64, u0 as i64);
9920            let xv = _mm256_loadu_si256(xq_ptr.add(gi * GROUP_SIZE) as *const __m256i);
9921            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
9922            let d256 = _mm256_madd_epi16(p16, ones);
9923            let d128 = _mm_add_epi32(
9924                _mm256_castsi256_si128(d256),
9925                _mm256_extracti128_si256(d256, 1),
9926            );
9927            let d64 = _mm_add_epi32(d128, _mm_shuffle_epi32(d128, 0xee));
9928            let d32 = _mm_cvtsi128_si32(_mm_add_epi32(d64, _mm_shuffle_epi32(d64, 0x55)));
9929            acc += d32 as f32 * s;
9930        }
9931        acc
9932    }
9933}
9934
9935/// VNNI twin of `q1t_dot_row_avx2` (see `dpbusd_hsum`).
9936#[cfg(target_arch = "x86_64")]
9937#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
9938unsafe fn q1t_dot_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
9939    use core::arch::x86_64::*;
9940    // SAFETY: same tile/xq contracts as `q1t_dot_row_avx2`.
9941    unsafe {
9942        const TILE: usize = cortiq_core::quant::Q1T_TILE;
9943        let mut acc = 0f32;
9944        let bytes_ptr = bytes.as_ptr();
9945        let xq_ptr = xq.as_ptr();
9946        let row_off = r * gpr * TILE;
9947        for gi in 0..gpr {
9948            let off = row_off + gi * TILE;
9949            let s = f16_to_f32(u16::from_le_bytes([
9950                *bytes_ptr.add(off),
9951                *bytes_ptr.add(off + 1),
9952            ]));
9953            let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
9954            let wv = _mm256_set_epi64x(u3 as i64, u2 as i64, u1 as i64, u0 as i64);
9955            let xv = _mm256_loadu_si256(xq_ptr.add(gi * GROUP_SIZE) as *const __m256i);
9956            let d = dpbusd_hsum(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
9957            acc += d as f32 * s;
9958        }
9959        acc
9960    }
9961}
9962
9963/// Per-row int8 base dot, dispatched once per row (matvec decode hot path).
9964/// Callers are gated by `a8w8_enabled()`, so the target-feature kernels are
9965/// reachable.
9966#[inline]
9967fn q1t_dot_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
9968    #[cfg(target_arch = "aarch64")]
9969    unsafe {
9970        return q1t_dot_row_sdot(bytes, r, gpr, xq);
9971    }
9972    #[cfg(target_arch = "x86_64")]
9973    unsafe {
9974        if vnni_tiles_enabled() {
9975            return q1t_dot_row_vnni(bytes, r, gpr, xq);
9976        }
9977        return q1t_dot_row_avx2(bytes, r, gpr, xq);
9978    }
9979    #[allow(unreachable_code)]
9980    {
9981        const TILE: usize = cortiq_core::quant::Q1T_TILE;
9982        let mut acc = 0f32;
9983        let mut sg = [0i8; GROUP_SIZE + 8]; // +8 slack for the u64-store unpack
9984        for gi in 0..gpr {
9985            let off = (r * gpr + gi) * TILE;
9986            let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
9987            q1t_unpack_group_i8(bytes.as_ptr().wrapping_add(off + 2), &mut sg);
9988            let mut d = 0i32;
9989            for k in 0..GROUP_SIZE {
9990                d += sg[k] as i32 * xq[gi * GROUP_SIZE + k] as i32;
9991            }
9992            acc += d as f32 * s;
9993        }
9994        acc
9995    }
9996}
9997
9998/// Σ over a row's outliers of `value·x[col]` — the correction that adds the
9999/// overlay's exact weights on top of the base dot. INVARIANT: the encoder
10000/// writes ternary code 0 at every outlier position (`quantize_q1t`), so the
10001/// base contributes nothing there and this is a plain `value·x`, not
10002/// `(value − base)·x` — no scattered per-outlier scale read. Row `r`'s entries
10003/// are the contiguous slice `[row_ptr[r], row_ptr[r+1])`, so no binary search.
10004fn q1t_row_outlier_correction(
10005    bytes: &[u8],
10006    r: usize,
10007    rp_off: usize,
10008    entries_off: usize,
10009    has_ov: bool,
10010    x: &[f32],
10011) -> f32 {
10012    if !has_ov {
10013        return 0.0;
10014    }
10015    let (c0, c1) = (
10016        q1t_rowptr(bytes, rp_off, r),
10017        q1t_rowptr(bytes, rp_off, r + 1),
10018    );
10019    let mut corr = 0f32;
10020    for p in c0..c1 {
10021        let e = entries_off + p * 4;
10022        let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
10023        let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
10024        corr += val * x[col];
10025    }
10026    corr
10027}
10028
10029/// Dequantize one q1t row into `buf[..cols]` via the sign LUT (no division),
10030/// then apply the row's outliers (its `[row_ptr[r], row_ptr[r+1])` slice).
10031/// Used by the batched (prefill) path where the decode amortizes over the batch.
10032fn q1t_dequant_row(
10033    bytes: &[u8],
10034    r: usize,
10035    gpr: usize,
10036    rp_off: usize,
10037    entries_off: usize,
10038    has_ov: bool,
10039    buf: &mut [f32],
10040) {
10041    const TILE: usize = cortiq_core::quant::Q1T_TILE;
10042    for g in 0..gpr {
10043        let off = (r * gpr + g) * TILE;
10044        let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
10045        let codes = &bytes[off + 2..off + TILE];
10046        let bc = g * GROUP_SIZE;
10047        // 6 full bytes (30 codes) + a 7th byte holding the last 2.
10048        for bi in 0..6 {
10049            let lut = &SIGN5[codes[bi] as usize];
10050            let d = &mut buf[bc + bi * 5..bc + bi * 5 + 5];
10051            for i in 0..5 {
10052                d[i] = lut[i] * s;
10053            }
10054        }
10055        let lut = &SIGN5[codes[6] as usize];
10056        buf[bc + 30] = lut[0] * s;
10057        buf[bc + 31] = lut[1] * s;
10058    }
10059    if !has_ov {
10060        return;
10061    }
10062    let (c0, c1) = (
10063        q1t_rowptr(bytes, rp_off, r),
10064        q1t_rowptr(bytes, rp_off, r + 1),
10065    );
10066    for p in c0..c1 {
10067        let e = entries_off + p * 4;
10068        let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
10069        buf[col] = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
10070    }
10071}
10072
10073/// Add the sparse outlier overlay onto a base dot already in `out` (the GPU
10074/// computes the ternary base; the overlay stays on the CPU — its entries are
10075/// few and its per-row gather doesn't vectorize on the GPU). Row-parallel.
10076fn q1t_add_overlay(
10077    bytes: &[u8],
10078    x: &[f32],
10079    rows: usize,
10080    cols: usize,
10081    out: &mut [f32],
10082    pool: Option<&Pool>,
10083) {
10084    const TILE: usize = cortiq_core::quant::Q1T_TILE;
10085    let gpr = cols / GROUP_SIZE;
10086    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
10087    if !has_ov {
10088        return;
10089    }
10090    let out_addr = SendMut(out.as_mut_ptr());
10091    let run = move |start: usize, end: usize| {
10092        for r in start..end {
10093            let corr = q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
10094            // SAFETY: disjoint rows; add onto the base the GPU already wrote.
10095            unsafe { *out_addr.at(r) += corr };
10096        }
10097    };
10098    dispatch_rows(pool, rows, &run);
10099}
10100
10101/// Q1T row range via the A8W8 int8 path — shared activation split,
10102/// per-row: base SDOT dot + outlier correction + overlay.
10103#[allow(clippy::too_many_arguments)]
10104fn q1t_range_a8w8(
10105    bytes: &[u8],
10106    gpr: usize,
10107    rp_off: usize,
10108    ent_off: usize,
10109    has_ov: bool,
10110    act: &SplitAct,
10111    x: &[f32],
10112    out: SendMut,
10113    start: usize,
10114    end: usize,
10115) {
10116    for r in start..end {
10117        let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
10118        for &(j, xv) in &act.outliers {
10119            acc += q1t_base_weight(bytes, r, gpr, j) * xv;
10120        }
10121        acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
10122        // SAFETY: disjoint row ranges per worker.
10123        unsafe { *out.at(r) = acc };
10124    }
10125}
10126
10127/// Q1T row range via the f32 path (no SDOT) — for matvec_many batched
10128/// dispatch when a8w8 is unavailable.
10129#[allow(clippy::too_many_arguments)]
10130fn q1t_range_f32_batch(
10131    bytes: &[u8],
10132    gpr: usize,
10133    rp_off: usize,
10134    ent_off: usize,
10135    has_ov: bool,
10136    x: &[f32],
10137    out: SendMut,
10138    start: usize,
10139    end: usize,
10140) {
10141    const TILE: usize = cortiq_core::quant::Q1T_TILE;
10142    let mut sg = [0f32; GROUP_SIZE];
10143    for r in start..end {
10144        let mut acc = 0f32;
10145        for g in 0..gpr {
10146            let off = (r * gpr + g) * TILE;
10147            let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
10148            let codes = &bytes[off + 2..off + TILE];
10149            let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
10150            for bi in 0..6 {
10151                sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
10152            }
10153            let lut = &SIGN5[codes[6] as usize];
10154            sg[30] = lut[0];
10155            sg[31] = lut[1];
10156            let mut gsum = 0f32;
10157            for k in 0..GROUP_SIZE {
10158                gsum += sg[k] * xg[k];
10159            }
10160            acc += s * gsum;
10161        }
10162        acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
10163        // SAFETY: disjoint row ranges per worker.
10164        unsafe { *out.at(r) = acc };
10165    }
10166}
10167
10168/// Ternary (q1t) matvec — decode+dot straight from mmap, one group at a time:
10169/// no per-ROW buffer, no division (the sign LUT), and a tiny per-group sign
10170/// buffer so the 32-wide dot vectorizes. This is the decode hot path.
10171fn q1t_matvec(
10172    bytes: &[u8],
10173    x: &[f32],
10174    rows: usize,
10175    cols: usize,
10176    out: &mut [f32],
10177    pool: Option<&Pool>,
10178) {
10179    debug_assert_eq!(out.len(), rows);
10180    const TILE: usize = cortiq_core::quant::Q1T_TILE;
10181    let gpr = cols / GROUP_SIZE;
10182    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
10183    let out_addr = SendMut(out.as_mut_ptr());
10184    // int8 SDOT base dot (ARM dotprod): ~4× the f32 arithmetic. x → i8 once
10185    // (`split_act`), activation outliers added back exactly in f32, weight
10186    // overlay on top. ARM SDOT / x86 AVX2; CMF_SDOT=0 keeps the exact f32 path.
10187    if a8w8_enabled() {
10188        let act = split_act(x);
10189        let act = &act;
10190        let run = move |start: usize, end: usize| {
10191            for r in start..end {
10192                let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
10193                for &(j, xv) in &act.outliers {
10194                    acc += q1t_base_weight(bytes, r, gpr, j) * xv;
10195                }
10196                acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
10197                // SAFETY: disjoint row ranges per worker.
10198                unsafe { *out_addr.at(r) = acc };
10199            }
10200        };
10201        dispatch_rows(pool, rows, &run);
10202        return;
10203    }
10204    let run = move |start: usize, end: usize| {
10205        // Per-group signs, unpacked contiguously so the dot below is a clean
10206        // 32-wide reduction the autovectorizer turns into f32x4 FMAs — the
10207        // 5-values-per-byte base-3 layout won't SIMD in place.
10208        let mut sg = [0f32; GROUP_SIZE];
10209        for r in start..end {
10210            let mut acc = 0f32;
10211            for g in 0..gpr {
10212                let off = (r * gpr + g) * TILE;
10213                let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
10214                let codes = &bytes[off + 2..off + TILE];
10215                let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
10216                for bi in 0..6 {
10217                    sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
10218                }
10219                let lut = &SIGN5[codes[6] as usize];
10220                sg[30] = lut[0];
10221                sg[31] = lut[1];
10222                let mut gsum = 0f32;
10223                for k in 0..GROUP_SIZE {
10224                    gsum += sg[k] * xg[k];
10225                }
10226                acc += s * gsum;
10227            }
10228            acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
10229            unsafe { *out_addr.at(r) = acc };
10230        }
10231    };
10232    dispatch_rows(pool, rows, &run);
10233}
10234
10235/// Fused-pair twin of `q1t_dot_row_sdot`: ONE register unpack of the
10236/// ternary codes serves BOTH activation streams (the unpack chain is
10237/// the dominant per-row cost — MTP verify pairs paid it twice). Per
10238/// stream the group order and f32 accumulation match the single-row
10239/// kernel exactly, so pair == 2×matvec bit-for-bit.
10240#[cfg(target_arch = "aarch64")]
10241#[target_feature(enable = "neon,dotprod")]
10242unsafe fn q1t_dot_row_sdot2(bytes: &[u8], r: usize, gpr: usize, xa: &[i8], xb: &[i8]) -> [f32; 2] {
10243    use core::arch::aarch64::*;
10244    use core::arch::asm;
10245    // SAFETY: same slice-length contracts as `q1t_dot_row_sdot`, ×2.
10246    unsafe {
10247        const TILE: usize = cortiq_core::quant::Q1T_TILE;
10248        let bytes_ptr = bytes.as_ptr();
10249        let row_off = r * gpr * TILE;
10250        let xp = [xa.as_ptr(), xb.as_ptr()];
10251        let mut acc = [0f32; 2];
10252        macro_rules! sdot2 {
10253            ($w0:expr, $w1:expr, $x:expr) => {{
10254                let x0 = vld1q_s8($x);
10255                let x1 = vld1q_s8($x.add(16));
10256                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
10257                asm!(
10258                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
10259                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
10260                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
10261                    w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
10262                    options(pure, nomem, nostack),
10263                );
10264                vaddvq_s32(vaddq_s32(a0, a1))
10265            }};
10266        }
10267        let gpr2 = gpr & !1;
10268        let mut gi = 0;
10269        while gi < gpr2 {
10270            let off0 = row_off + gi * TILE;
10271            let off1 = off0 + TILE;
10272            let s0 = f16_to_f32(u16::from_le_bytes([
10273                *bytes_ptr.add(off0),
10274                *bytes_ptr.add(off0 + 1),
10275            ]));
10276            let s1 = f16_to_f32(u16::from_le_bytes([
10277                *bytes_ptr.add(off1),
10278                *bytes_ptr.add(off1 + 1),
10279            ]));
10280            let (u0_0, u1_0, u2_0, u3_0) = q1t_unpack_reg_u64s(bytes_ptr.add(off0 + 2));
10281            let (u0_1, u1_1, u2_1, u3_1) = q1t_unpack_reg_u64s(bytes_ptr.add(off1 + 2));
10282            let w0_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_0), vcreate_u64(u1_0)));
10283            let w1_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_0), vcreate_u64(u3_0)));
10284            let w0_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_1), vcreate_u64(u1_1)));
10285            let w1_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_1), vcreate_u64(u3_1)));
10286            for k in 0..2 {
10287                let d0 = sdot2!(w0_0, w1_0, xp[k].add(gi * GROUP_SIZE));
10288                let d1 = sdot2!(w0_1, w1_1, xp[k].add((gi + 1) * GROUP_SIZE));
10289                acc[k] += d0 as f32 * s0 + d1 as f32 * s1;
10290            }
10291            gi += 2;
10292        }
10293        if gi < gpr {
10294            let off = row_off + gi * TILE;
10295            let s = f16_to_f32(u16::from_le_bytes([
10296                *bytes_ptr.add(off),
10297                *bytes_ptr.add(off + 1),
10298            ]));
10299            let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
10300            let w0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0), vcreate_u64(u1)));
10301            let w1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2), vcreate_u64(u3)));
10302            for k in 0..2 {
10303                let d = sdot2!(w0, w1, xp[k].add(gi * GROUP_SIZE));
10304                acc[k] += d as f32 * s;
10305            }
10306        }
10307        acc
10308    }
10309}
10310
10311/// Fused Q1T pair matvec: ONE pass over the rows serves both
10312/// activation streams — on ARM the ternary register unpack happens
10313/// once per tile pair (`q1t_dot_row_sdot2`); elsewhere the second dot
10314/// rides the row's L1-warm tile bytes. Per stream the math matches
10315/// `q1t_matvec` exactly.
10316fn q1t_matvec2(
10317    bytes: &[u8],
10318    x1: &[f32],
10319    x2: &[f32],
10320    rows: usize,
10321    cols: usize,
10322    o1: &mut [f32],
10323    o2: &mut [f32],
10324    pool: Option<&Pool>,
10325) {
10326    debug_assert_eq!(o1.len(), rows);
10327    debug_assert_eq!(o2.len(), rows);
10328    const TILE: usize = cortiq_core::quant::Q1T_TILE;
10329    let gpr = cols / GROUP_SIZE;
10330    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
10331    let out1 = SendMut(o1.as_mut_ptr());
10332    let out2 = SendMut(o2.as_mut_ptr());
10333    if a8w8_enabled() {
10334        let a1 = split_act(x1);
10335        let a2 = split_act(x2);
10336        let (a1, a2) = (&a1, &a2);
10337        let run = move |start: usize, end: usize| {
10338            for r in start..end {
10339                #[cfg(target_arch = "aarch64")]
10340                // a8w8 on aarch64 ⇔ sdot_enabled(), so the kernel's
10341                // target features are present.
10342                let ds = unsafe { q1t_dot_row_sdot2(bytes, r, gpr, &a1.xq, &a2.xq) };
10343                #[cfg(not(target_arch = "aarch64"))]
10344                let ds = [
10345                    q1t_dot_row_i8(bytes, r, gpr, &a1.xq),
10346                    q1t_dot_row_i8(bytes, r, gpr, &a2.xq),
10347                ];
10348                let mut acc1 = ds[0] * a1.sx;
10349                for &(j, xv) in &a1.outliers {
10350                    acc1 += q1t_base_weight(bytes, r, gpr, j) * xv;
10351                }
10352                acc1 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x1);
10353                let mut acc2 = ds[1] * a2.sx;
10354                for &(j, xv) in &a2.outliers {
10355                    acc2 += q1t_base_weight(bytes, r, gpr, j) * xv;
10356                }
10357                acc2 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x2);
10358                // SAFETY: disjoint row ranges per worker.
10359                unsafe {
10360                    *out1.at(r) = acc1;
10361                    *out2.at(r) = acc2;
10362                }
10363            }
10364        };
10365        dispatch_rows(pool, rows, &run);
10366        return;
10367    }
10368    let run = move |start: usize, end: usize| {
10369        // Exact path (CMF_SDOT=0): unpack the sign LUT once per group,
10370        // dot both streams — same op order per stream as `q1t_matvec`.
10371        let mut sg = [0f32; GROUP_SIZE];
10372        for r in start..end {
10373            let mut acc1 = 0f32;
10374            let mut acc2 = 0f32;
10375            for g in 0..gpr {
10376                let off = (r * gpr + g) * TILE;
10377                let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
10378                let codes = &bytes[off + 2..off + TILE];
10379                for bi in 0..6 {
10380                    sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
10381                }
10382                let lut = &SIGN5[codes[6] as usize];
10383                sg[30] = lut[0];
10384                sg[31] = lut[1];
10385                let xg1 = &x1[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
10386                let xg2 = &x2[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
10387                let mut gsum1 = 0f32;
10388                for k in 0..GROUP_SIZE {
10389                    gsum1 += sg[k] * xg1[k];
10390                }
10391                acc1 += s * gsum1;
10392                let mut gsum2 = 0f32;
10393                for k in 0..GROUP_SIZE {
10394                    gsum2 += sg[k] * xg2[k];
10395                }
10396                acc2 += s * gsum2;
10397            }
10398            acc1 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x1);
10399            acc2 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x2);
10400            // SAFETY: disjoint row ranges per worker.
10401            unsafe {
10402                *out1.at(r) = acc1;
10403                *out2.at(r) = acc2;
10404            }
10405        }
10406    };
10407    dispatch_rows(pool, rows, &run);
10408}
10409
10410/// Ternary (q1t) matmat (prefill) — dequant each row once, dot the whole
10411/// batch against it (amortizes the per-row decode).
10412fn q1t_matmat(
10413    bytes: &[u8],
10414    xs: &[f32],
10415    b: usize,
10416    rows: usize,
10417    cols: usize,
10418    out: &mut [f32],
10419    pool: Option<&Pool>,
10420) {
10421    debug_assert_eq!(out.len(), b * rows);
10422    const TILE: usize = cortiq_core::quant::Q1T_TILE;
10423    let gpr = cols / GROUP_SIZE;
10424    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
10425    let out_addr = SendMut(out.as_mut_ptr());
10426    // int8 prefill (ARM SDOT / x86 AVX2): quantize the B inputs once, unpack
10427    // each weight row's signs to i8 ONCE, then int8-dot against every input —
10428    // the row sign-decode amortizes over the whole batch. CMF_SDOT=0 → f32.
10429    if a8w8_enabled() {
10430        let acts: Vec<SplitAct> = (0..b)
10431            .map(|bi| split_act(&xs[bi * cols..(bi + 1) * cols]))
10432            .collect();
10433        let acts = &acts;
10434        let run = move |start: usize, end: usize| {
10435            let mut sg = vec![0i8; cols + 8]; // row signs, i8 (+8 unpack slack)
10436            let mut sc = vec![0f32; gpr]; // per-group scales
10437            let mut accs = vec![0f32; b]; // per-batch accumulators, reused per row
10438            for r in start..end {
10439                for g in 0..gpr {
10440                    let off = (r * gpr + g) * TILE;
10441                    sc[g] = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
10442                    q1t_unpack_group_i8(
10443                        bytes.as_ptr().wrapping_add(off + 2),
10444                        &mut sg[g * GROUP_SIZE..],
10445                    );
10446                }
10447                for bi in 0..b {
10448                    let act = &acts[bi];
10449                    let mut isum = 0f32;
10450                    for g in 0..gpr {
10451                        let d = q1t_i8dot32(
10452                            sg.as_ptr().wrapping_add(g * GROUP_SIZE),
10453                            act.xq.as_ptr().wrapping_add(g * GROUP_SIZE),
10454                        );
10455                        isum += d as f32 * sc[g];
10456                    }
10457                    let mut acc = isum * act.sx;
10458                    for &(j, xv) in &act.outliers {
10459                        acc += q1t_base_weight(bytes, r, gpr, j) * xv;
10460                    }
10461                    accs[bi] = acc;
10462                }
10463                // Overlay ONCE per row for the whole batch: read each (col, val)
10464                // from mmap a single time (was b× — the re-read dominated prefill)
10465                // and fan it out over the batch via the cached inputs.
10466                if has_ov {
10467                    let (c0, c1) = (
10468                        q1t_rowptr(bytes, rp_off, r),
10469                        q1t_rowptr(bytes, rp_off, r + 1),
10470                    );
10471                    for p in c0..c1 {
10472                        let e = ent_off + p * 4;
10473                        let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
10474                        let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
10475                        for bi in 0..b {
10476                            accs[bi] += val * xs[bi * cols + col];
10477                        }
10478                    }
10479                }
10480                for bi in 0..b {
10481                    unsafe { *out_addr.at(bi * rows + r) = accs[bi] };
10482                }
10483            }
10484        };
10485        dispatch_rows(pool, rows, &run);
10486        return;
10487    }
10488    let run = move |start: usize, end: usize| {
10489        let mut buf = vec![0f32; cols];
10490        for r in start..end {
10491            q1t_dequant_row(bytes, r, gpr, rp_off, ent_off, has_ov, &mut buf);
10492            for bi in 0..b {
10493                let xr = &xs[bi * cols..(bi + 1) * cols];
10494                let mut acc = 0f32;
10495                for j in 0..cols {
10496                    acc += buf[j] * xr[j];
10497                }
10498                unsafe { *out_addr.at(bi * rows + r) = acc };
10499            }
10500        }
10501    };
10502    dispatch_rows(pool, rows, &run);
10503}
10504
10505/// q1-tile entry points of the binary drivers (the tests' handles; the
10506/// runtime goes through `QTensor::bin_w`).
10507#[cfg(test)]
10508fn q1_matvec(
10509    bytes: &[u8],
10510    x: &[f32],
10511    rows: usize,
10512    cols: usize,
10513    out: &mut [f32],
10514    pool: Option<&Pool>,
10515) {
10516    bin_matvec(BinW::Q1(bytes), x, rows, cols, out, pool)
10517}
10518
10519/// Binary-weight matvec (q1 or q1g): A8W8 rows with the shared group
10520/// sums, or the exact scalar walk under CMF_SDOT=0.
10521fn bin_matvec(w: BinW, x: &[f32], rows: usize, cols: usize, out: &mut [f32], pool: Option<&Pool>) {
10522    debug_assert_eq!(out.len(), rows);
10523    let gpr = cols / GROUP_SIZE;
10524    let out_addr = SendMut(out.as_mut_ptr());
10525    if a8w8_enabled() {
10526        let act = split_act(x);
10527        let gsum = q1_group_sums(&act.xq, gpr);
10528        let (act, gsum) = (&act, &gsum);
10529        let run =
10530            move |start: usize, end: usize| bin_range_a8w8(w, gpr, act, gsum, out_addr, start, end);
10531        dispatch_rows(pool, rows, &run);
10532        return;
10533    }
10534    let run = move |start: usize, end: usize| bin_range_f32(w, gpr, x, out_addr, start, end);
10535    dispatch_rows(pool, rows, &run);
10536}
10537
10538/// Fused two-input q1 matvec (weights read once per pair).
10539#[cfg(test)]
10540#[allow(clippy::too_many_arguments)]
10541fn q1_matvec2(
10542    bytes: &[u8],
10543    x1: &[f32],
10544    x2: &[f32],
10545    rows: usize,
10546    cols: usize,
10547    o1: &mut [f32],
10548    o2: &mut [f32],
10549    pool: Option<&Pool>,
10550) {
10551    bin_matvec2(BinW::Q1(bytes), x1, x2, rows, cols, o1, o2, pool)
10552}
10553
10554/// Fused two-input binary-weight matvec (q1 or q1g).
10555#[allow(clippy::too_many_arguments)]
10556fn bin_matvec2(
10557    w: BinW,
10558    x1: &[f32],
10559    x2: &[f32],
10560    rows: usize,
10561    cols: usize,
10562    o1: &mut [f32],
10563    o2: &mut [f32],
10564    pool: Option<&Pool>,
10565) {
10566    let gpr = cols / GROUP_SIZE;
10567    let p1 = SendMut(o1.as_mut_ptr());
10568    let p2 = SendMut(o2.as_mut_ptr());
10569    if a8w8_enabled() {
10570        let a1 = split_act(x1);
10571        let a2 = split_act(x2);
10572        let g1 = q1_group_sums(&a1.xq, gpr);
10573        let g2 = q1_group_sums(&a2.xq, gpr);
10574        let (a1, a2, g1, g2) = (&a1, &a2, &g1, &g2);
10575        let run = move |start: usize, end: usize| {
10576            for r in start..end {
10577                let mut v1 = w.dot_i8(r, gpr, &a1.xq, g1) * a1.sx;
10578                let mut v2 = w.dot_i8(r, gpr, &a2.xq, g2) * a2.sx;
10579                for &(j, xv) in &a1.outliers {
10580                    let (wv, s) = w.outlier(r, gpr, j);
10581                    v1 += wv * s * xv;
10582                }
10583                for &(j, xv) in &a2.outliers {
10584                    let (wv, s) = w.outlier(r, gpr, j);
10585                    v2 += wv * s * xv;
10586                }
10587                // SAFETY: disjoint row ranges per worker.
10588                unsafe {
10589                    *p1.at(r) = v1;
10590                    *p2.at(r) = v2;
10591                }
10592            }
10593        };
10594        dispatch_rows(pool, rows, &run);
10595        return;
10596    }
10597    let run = move |start: usize, end: usize| {
10598        for r in start..end {
10599            // SAFETY: disjoint row ranges per worker.
10600            unsafe {
10601                *p1.at(r) = w.row_exact(r, gpr, x1);
10602                *p2.at(r) = w.row_exact(r, gpr, x2);
10603            }
10604        }
10605    };
10606    dispatch_rows(pool, rows, &run);
10607}
10608
10609/// Batched q1 matmat: each row's tiles stream once per microbatch.
10610#[cfg(test)]
10611#[allow(clippy::too_many_arguments)]
10612fn q1_matmat(
10613    bytes: &[u8],
10614    xs_all: &[f32],
10615    b: usize,
10616    rows: usize,
10617    cols: usize,
10618    out: &mut [f32],
10619    pool: Option<&Pool>,
10620) {
10621    bin_matmat(BinW::Q1(bytes), xs_all, b, rows, cols, out, pool)
10622}
10623
10624/// Batched binary-weight matmat (q1 or q1g).
10625#[allow(clippy::too_many_arguments)]
10626fn bin_matmat(
10627    w: BinW,
10628    xs_all: &[f32],
10629    b: usize,
10630    rows: usize,
10631    cols: usize,
10632    out: &mut [f32],
10633    pool: Option<&Pool>,
10634) {
10635    debug_assert_eq!(out.len(), b * rows);
10636    let gpr = cols / GROUP_SIZE;
10637    let out_addr = SendMut(out.as_mut_ptr());
10638    if a8w8_enabled() {
10639        let acts: Vec<(SplitAct, Vec<i32>)> = (0..b)
10640            .map(|bi| {
10641                let act = split_act(&xs_all[bi * cols..(bi + 1) * cols]);
10642                let gsum = q1_group_sums(&act.xq, gpr);
10643                (act, gsum)
10644            })
10645            .collect();
10646        let acts = &acts;
10647        #[cfg(target_arch = "x86_64")]
10648        let blocked_ok = avx2_enabled() && blocked_enabled();
10649        #[cfg(target_arch = "aarch64")]
10650        let blocked_ok = sdot_enabled() && blocked_enabled();
10651        let run = move |start: usize, end: usize| {
10652            for r in start..end {
10653                let mut bi = 0usize;
10654                // Blocked 1×4: the unpacked bit mask serves four
10655                // activation streams per group.
10656                #[cfg(target_arch = "aarch64")]
10657                if blocked_ok {
10658                    while bi + 4 <= acts.len() {
10659                        let xs = [
10660                            acts[bi].0.xq.as_slice(),
10661                            acts[bi + 1].0.xq.as_slice(),
10662                            acts[bi + 2].0.xq.as_slice(),
10663                            acts[bi + 3].0.xq.as_slice(),
10664                        ];
10665                        let gs = [
10666                            acts[bi].1.as_slice(),
10667                            acts[bi + 1].1.as_slice(),
10668                            acts[bi + 2].1.as_slice(),
10669                            acts[bi + 3].1.as_slice(),
10670                        ];
10671                        // SAFETY: blocked_ok checked SDOT.
10672                        let d = unsafe { w.dot_1x4_sdot(r, gpr, xs, gs) };
10673                        for k in 0..4 {
10674                            let (act, _) = &acts[bi + k];
10675                            let mut acc = d[k] * act.sx;
10676                            for &(j, xv) in &act.outliers {
10677                                let (wv, sc) = w.outlier(r, gpr, j);
10678                                acc += wv * sc * xv;
10679                            }
10680                            // SAFETY: disjoint (bi, r) cells per worker.
10681                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
10682                        }
10683                        bi += 4;
10684                    }
10685                }
10686                #[cfg(target_arch = "x86_64")]
10687                if blocked_ok {
10688                    while bi + 4 <= acts.len() {
10689                        let xs = [
10690                            acts[bi].0.xq.as_slice(),
10691                            acts[bi + 1].0.xq.as_slice(),
10692                            acts[bi + 2].0.xq.as_slice(),
10693                            acts[bi + 3].0.xq.as_slice(),
10694                        ];
10695                        let gs = [
10696                            acts[bi].1.as_slice(),
10697                            acts[bi + 1].1.as_slice(),
10698                            acts[bi + 2].1.as_slice(),
10699                            acts[bi + 3].1.as_slice(),
10700                        ];
10701                        // SAFETY: blocked_ok checked AVX2.
10702                        let d = unsafe { w.dot_1x4_x86(r, gpr, xs, gs) };
10703                        for k in 0..4 {
10704                            let (act, _) = &acts[bi + k];
10705                            let mut acc = d[k] * act.sx;
10706                            for &(j, xv) in &act.outliers {
10707                                let (wv, sc) = w.outlier(r, gpr, j);
10708                                acc += wv * sc * xv;
10709                            }
10710                            // SAFETY: disjoint (bi, r) cells per worker.
10711                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
10712                        }
10713                        bi += 4;
10714                    }
10715                }
10716                while bi < acts.len() {
10717                    let (act, gsum) = &acts[bi];
10718                    let mut acc = w.dot_i8(r, gpr, &act.xq, gsum) * act.sx;
10719                    for &(j, xv) in &act.outliers {
10720                        let (wv, s) = w.outlier(r, gpr, j);
10721                        acc += wv * s * xv;
10722                    }
10723                    // SAFETY: disjoint (bi, r) cells per worker range.
10724                    unsafe { *out_addr.at(bi * rows + r) = acc };
10725                    bi += 1;
10726                }
10727            }
10728        };
10729        dispatch_rows(pool, rows, &run);
10730        return;
10731    }
10732    let run = move |start: usize, end: usize| {
10733        for r in start..end {
10734            for bi in 0..b {
10735                let x = &xs_all[bi * cols..(bi + 1) * cols];
10736                // SAFETY: disjoint (bi, r) cells per worker range.
10737                unsafe { *out_addr.at(bi * rows + r) = w.row_exact(r, gpr, x) };
10738            }
10739        }
10740    };
10741    dispatch_rows(pool, rows, &run);
10742}
10743
10744/// Fused q4_block matvec straight from the mapped bytes. SDOT path when
10745/// dotprod is available (port of vmfcore `dot_q4_block_sdot`, measured
10746/// +23% on q4 decode): nibbles → centered i8, int8×int8 `sdot` per
10747/// 32-group, exact outlier correction — the same A8W8 contract as q8.
10748/// `CMF_SDOT=0` keeps the exact scalar path.
10749fn q4matvec(
10750    bytes: &[u8],
10751    x: &[f32],
10752    rows: usize,
10753    cols: usize,
10754    out: &mut [f32],
10755    pool: Option<&Pool>,
10756) {
10757    debug_assert_eq!(out.len(), rows);
10758    let (packed, scales) = q4_split(bytes, rows, cols);
10759    let gpr = cols / GROUP_SIZE;
10760    let out_addr = SendMut(out.as_mut_ptr());
10761
10762    if a8w8_enabled() {
10763        let act = split_act(x);
10764        let run = move |start: usize, end: usize| {
10765            q4_range_a8w8(packed, scales, gpr, cols, &act, out_addr, start, end)
10766        };
10767        dispatch_rows(pool, rows, &run);
10768        return;
10769    }
10770
10771    let run =
10772        move |start: usize, end: usize| q4_range_f32(packed, scales, gpr, x, out_addr, start, end);
10773    dispatch_rows(pool, rows, &run);
10774}
10775
10776/// One q4 row via the A8W8 int8 path — SDOT on ARM, AVX2 maddubs on
10777/// x86 (scalar fallback is unreachable: callers gate on a8w8_enabled).
10778#[inline]
10779#[allow(unreachable_code)]
10780/// One UNPACKED q4 row (centered i8 in `buf`) against four activation
10781/// streams: the 32-byte weight chunk and its abs() load once per group,
10782/// the per-group f16 scale decodes once — four maddubs+reduce chains
10783/// instead of four full (load, abs, dot) rounds.
10784#[cfg(target_arch = "x86_64")]
10785#[target_feature(enable = "avx2")]
10786unsafe fn dot_q4b_row_1x4_avx2(
10787    buf: &[u8],
10788    scales: &[u8],
10789    g0: usize,
10790    gpr: usize,
10791    xs: [&[i8]; 4],
10792) -> [f32; 4] {
10793    // SAFETY: callers uphold buffer contracts (buf.len() == gpr·32).
10794    unsafe {
10795        use core::arch::x86_64::*;
10796        let ones = _mm256_set1_epi16(1);
10797        let mut acc = [0f32; 4];
10798        for gi in 0..gpr {
10799            let s = f16_to_f32(u16::from_le_bytes([
10800                scales[(g0 + gi) * 2],
10801                scales[(g0 + gi) * 2 + 1],
10802            ]));
10803            let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
10804            let aw = _mm256_abs_epi8(w);
10805            for (k, xq) in xs.iter().enumerate() {
10806                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
10807                let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
10808                let d = _mm256_madd_epi16(p16, ones);
10809                let hi128 = _mm256_extracti128_si256::<1>(d);
10810                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
10811                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
10812                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
10813                acc[k] += _mm_cvtsi128_si32(s32) as f32 * s;
10814            }
10815        }
10816        acc
10817    }
10818}
10819
10820/// VNNI twin of `dot_q4b_row_1x4_avx2` (see `dpbusd_hsum`).
10821#[cfg(target_arch = "x86_64")]
10822#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
10823unsafe fn dot_q4b_row_1x4_vnni(
10824    buf: &[u8],
10825    scales: &[u8],
10826    g0: usize,
10827    gpr: usize,
10828    xs: [&[i8]; 4],
10829) -> [f32; 4] {
10830    // SAFETY: callers uphold buffer contracts (buf.len() == gpr·32).
10831    unsafe {
10832        use core::arch::x86_64::*;
10833        let mut acc = [0f32; 4];
10834        for gi in 0..gpr {
10835            let s = f16_to_f32(u16::from_le_bytes([
10836                scales[(g0 + gi) * 2],
10837                scales[(g0 + gi) * 2 + 1],
10838            ]));
10839            let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
10840            let aw = _mm256_abs_epi8(w);
10841            for (k, xq) in xs.iter().enumerate() {
10842                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
10843                let d = dpbusd_hsum(aw, _mm256_sign_epi8(x, w));
10844                acc[k] += d as f32 * s;
10845            }
10846        }
10847        acc
10848    }
10849}
10850
10851/// The vbit flavor of the blocked 1×4: the per-activation A8W8 scale
10852/// folds in PER GROUP as `(d·sx)·s` — bit-matching the single-matvec
10853/// accumulation order (the q4_block flavor applies sx once at the end,
10854/// matching ITS single path; the two conventions are historical and
10855/// each blocked leg must mirror its own).
10856#[cfg(target_arch = "x86_64")]
10857#[target_feature(enable = "avx2")]
10858unsafe fn dot_q4b_row_1x4_sx_avx2(
10859    buf: &[u8],
10860    scales: &[u8],
10861    g0: usize,
10862    gpr: usize,
10863    xs: [&[i8]; 4],
10864    sxs: [f32; 4],
10865) -> [f32; 4] {
10866    // SAFETY: callers uphold buffer contracts (buf.len() == gpr·32).
10867    unsafe {
10868        use core::arch::x86_64::*;
10869        let ones = _mm256_set1_epi16(1);
10870        let mut acc = [0f32; 4];
10871        for gi in 0..gpr {
10872            let s = f16_to_f32(u16::from_le_bytes([
10873                scales[(g0 + gi) * 2],
10874                scales[(g0 + gi) * 2 + 1],
10875            ]));
10876            let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
10877            let aw = _mm256_abs_epi8(w);
10878            for (k, xq) in xs.iter().enumerate() {
10879                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
10880                let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
10881                let d = _mm256_madd_epi16(p16, ones);
10882                let hi128 = _mm256_extracti128_si256::<1>(d);
10883                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
10884                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
10885                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
10886                acc[k] += (_mm_cvtsi128_si32(s32) as f32 * sxs[k]) * s;
10887            }
10888        }
10889        acc
10890    }
10891}
10892
10893/// VNNI twin of `dot_q4b_row_1x4_sx_avx2` (see `dpbusd_hsum`; the
10894/// per-group `(d·sx)·s` fold mirrors the vbit single path).
10895#[cfg(target_arch = "x86_64")]
10896#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
10897unsafe fn dot_q4b_row_1x4_sx_vnni(
10898    buf: &[u8],
10899    scales: &[u8],
10900    g0: usize,
10901    gpr: usize,
10902    xs: [&[i8]; 4],
10903    sxs: [f32; 4],
10904) -> [f32; 4] {
10905    // SAFETY: callers uphold buffer contracts (buf.len() == gpr·32).
10906    unsafe {
10907        use core::arch::x86_64::*;
10908        let mut acc = [0f32; 4];
10909        for gi in 0..gpr {
10910            let s = f16_to_f32(u16::from_le_bytes([
10911                scales[(g0 + gi) * 2],
10912                scales[(g0 + gi) * 2 + 1],
10913            ]));
10914            let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
10915            let aw = _mm256_abs_epi8(w);
10916            for (k, xq) in xs.iter().enumerate() {
10917                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
10918                let d = dpbusd_hsum(aw, _mm256_sign_epi8(x, w));
10919                acc[k] += (d as f32 * sxs[k]) * s;
10920            }
10921        }
10922        acc
10923    }
10924}
10925
10926#[allow(unreachable_code)]
10927fn dot_q4_row_i8(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
10928    #[cfg(target_arch = "aarch64")]
10929    unsafe {
10930        return dot_q4_row_sdot(packed, scales, g0, gpr, xq);
10931    }
10932    #[cfg(target_arch = "x86_64")]
10933    unsafe {
10934        return dot_q4_row_avx2(packed, scales, g0, gpr, xq);
10935    }
10936    let mut acc = 0f32;
10937    for gi in 0..gpr {
10938        let g = g0 + gi;
10939        let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
10940        let mut d = 0i32;
10941        for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
10942            d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
10943                + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
10944        }
10945        acc += d as f32 * s;
10946    }
10947    acc
10948}
10949
10950/// Two-activation q4 row via the A8W8 int8 path (see `dot_q4_row_i8`).
10951#[inline]
10952#[allow(unreachable_code)]
10953fn dot_q4_row_i8_2(
10954    packed: &[u8],
10955    scales: &[u8],
10956    g0: usize,
10957    gpr: usize,
10958    xq1: &[i8],
10959    xq2: &[i8],
10960) -> (f32, f32) {
10961    #[cfg(target_arch = "aarch64")]
10962    unsafe {
10963        return dot_q4_row_sdot2(packed, scales, g0, gpr, xq1, xq2);
10964    }
10965    #[cfg(target_arch = "x86_64")]
10966    unsafe {
10967        return dot_q4_row_avx2_2(packed, scales, g0, gpr, xq1, xq2);
10968    }
10969    (
10970        dot_q4_row_i8(packed, scales, g0, gpr, xq1),
10971        dot_q4_row_i8(packed, scales, g0, gpr, xq2),
10972    )
10973}
10974
10975/// One q4 row range via SDOT (kernel body of `q4matvec`, extracted so
10976/// multi-matrix jobs can drive it for several tensors in one dispatch).
10977#[allow(clippy::too_many_arguments)]
10978fn q4_range_a8w8(
10979    packed: &[u8],
10980    scales: &[u8],
10981    gpr: usize,
10982    cols: usize,
10983    act: &SplitAct,
10984    out: SendMut,
10985    start: usize,
10986    end: usize,
10987) {
10988    for r in start..end {
10989        let mut acc = dot_q4_row_i8(packed, scales, r * gpr, gpr, &act.xq) * act.sx;
10990        // xq is zeroed at outlier slots — add the exact terms.
10991        for &(j, xv) in &act.outliers {
10992            let flat = r * cols + j;
10993            let byte = packed[flat / 2];
10994            let nib = if flat & 1 == 0 {
10995                byte & 0x0F
10996            } else {
10997                byte >> 4
10998            };
10999            let s = f16_to_f32(u16::from_le_bytes([
11000                scales[(flat / GROUP_SIZE) * 2],
11001                scales[(flat / GROUP_SIZE) * 2 + 1],
11002            ]));
11003            acc += ((nib as i32 - 8) as f32) * s * xv;
11004        }
11005        // SAFETY: disjoint row ranges per worker.
11006        unsafe { *out.at(r) = acc };
11007    }
11008}
11009
11010/// Two-input q4 row range via the A8W8 int8 path — kernel body of
11011/// `q4matvec2`, extracted for pair multi-matrix jobs.
11012#[allow(clippy::too_many_arguments)]
11013fn q4_range2_a8w8(
11014    packed: &[u8],
11015    scales: &[u8],
11016    gpr: usize,
11017    cols: usize,
11018    a1: &SplitAct,
11019    a2: &SplitAct,
11020    p1: SendMut,
11021    p2: SendMut,
11022    start: usize,
11023    end: usize,
11024) {
11025    for r in start..end {
11026        let (s1, s2) = dot_q4_row_i8_2(packed, scales, r * gpr, gpr, &a1.xq, &a2.xq);
11027        let mut acc1 = s1 * a1.sx;
11028        let mut acc2 = s2 * a2.sx;
11029        // xq is zeroed at outlier slots — add the exact terms.
11030        let fix = |outliers: &[(usize, f32)], acc: &mut f32| {
11031            for &(j, xv) in outliers {
11032                let flat = r * cols + j;
11033                let byte = packed[flat / 2];
11034                let nib = if flat & 1 == 0 {
11035                    byte & 0x0F
11036                } else {
11037                    byte >> 4
11038                };
11039                let s = f16_to_f32(u16::from_le_bytes([
11040                    scales[(flat / GROUP_SIZE) * 2],
11041                    scales[(flat / GROUP_SIZE) * 2 + 1],
11042                ]));
11043                *acc += ((nib as i32 - 8) as f32) * s * xv;
11044            }
11045        };
11046        fix(&a1.outliers, &mut acc1);
11047        fix(&a2.outliers, &mut acc2);
11048        // SAFETY: disjoint row ranges per worker.
11049        unsafe {
11050            *p1.at(r) = acc1;
11051            *p2.at(r) = acc2;
11052        }
11053    }
11054}
11055
11056/// Exact scalar q4 row range (same extraction, non-SDOT path).
11057fn q4_range_f32(
11058    packed: &[u8],
11059    scales: &[u8],
11060    gpr: usize,
11061    x: &[f32],
11062    out: SendMut,
11063    start: usize,
11064    end: usize,
11065) {
11066    for r in start..end {
11067        let mut acc = 0f32;
11068        for gi in 0..gpr {
11069            let g = r * gpr + gi;
11070            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
11071            let pk = &packed[g * 16..(g + 1) * 16];
11072            let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
11073            let mut ga = 0f32;
11074            for (k, &b) in pk.iter().enumerate() {
11075                ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
11076                    + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
11077            }
11078            acc += ga * s;
11079        }
11080        // SAFETY: disjoint row ranges per worker.
11081        unsafe { *out.at(r) = acc };
11082    }
11083}
11084
11085/// Fused two-input q4 matvec: nibbles are unpacked ONCE per group and
11086/// dotted against both activations (was: two full matvecs — double
11087/// weight traffic). Per-lane math matches `q4matvec` exactly.
11088#[allow(clippy::too_many_arguments)]
11089fn q4matvec2(
11090    bytes: &[u8],
11091    x1: &[f32],
11092    x2: &[f32],
11093    rows: usize,
11094    cols: usize,
11095    o1: &mut [f32],
11096    o2: &mut [f32],
11097    pool: Option<&Pool>,
11098) {
11099    debug_assert_eq!(o1.len(), rows);
11100    debug_assert_eq!(o2.len(), rows);
11101    let (packed, scales) = q4_split(bytes, rows, cols);
11102    let gpr = cols / GROUP_SIZE;
11103
11104    if a8w8_enabled() {
11105        let a1 = split_act(x1);
11106        let a2 = split_act(x2);
11107        let p1 = SendMut(o1.as_mut_ptr());
11108        let p2 = SendMut(o2.as_mut_ptr());
11109        let run = move |start: usize, end: usize| {
11110            q4_range2_a8w8(packed, scales, gpr, cols, &a1, &a2, p1, p2, start, end)
11111        };
11112        dispatch_rows(pool, rows, &run);
11113        return;
11114    }
11115
11116    let p1 = SendMut(o1.as_mut_ptr());
11117    let p2 = SendMut(o2.as_mut_ptr());
11118    let run = move |start: usize, end: usize| {
11119        q4_range2_f32(packed, scales, gpr, x1, x2, p1, p2, start, end)
11120    };
11121    dispatch_rows(pool, rows, &run);
11122}
11123
11124/// Two-input exact scalar q4 row range (same extraction).
11125#[allow(clippy::too_many_arguments)]
11126fn q4_range2_f32(
11127    packed: &[u8],
11128    scales: &[u8],
11129    gpr: usize,
11130    x1: &[f32],
11131    x2: &[f32],
11132    p1: SendMut,
11133    p2: SendMut,
11134    start: usize,
11135    end: usize,
11136) {
11137    for r in start..end {
11138        let (mut acc1, mut acc2) = (0f32, 0f32);
11139        for gi in 0..gpr {
11140            let g = r * gpr + gi;
11141            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
11142            let pk = &packed[g * 16..(g + 1) * 16];
11143            let x1g = &x1[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
11144            let x2g = &x2[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
11145            let (mut g1, mut g2) = (0f32, 0f32);
11146            for (k, &b) in pk.iter().enumerate() {
11147                let wl = (b & 0x0F) as f32 - 8.0;
11148                let wh = ((b >> 4) & 0x0F) as f32 - 8.0;
11149                g1 += wl * x1g[k * 2] + wh * x1g[k * 2 + 1];
11150                g2 += wl * x2g[k * 2] + wh * x2g[k * 2 + 1];
11151            }
11152            acc1 += g1 * s;
11153            acc2 += g2 * s;
11154        }
11155        // SAFETY: disjoint row ranges per worker.
11156        unsafe {
11157            *p1.at(r) = acc1;
11158            *p2.at(r) = acc2;
11159        }
11160    }
11161}
11162
11163thread_local! {
11164    /// Per-worker decoded-row scratch for the batched q4/vbit kernels
11165    /// (centered i8 for SDOT, f32 for the exact/scalar paths).
11166    static ROW_I8: std::cell::RefCell<Vec<u8>> = const { std::cell::RefCell::new(Vec::new()) };
11167    static ROW_F32: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
11168}
11169
11170/// Batched q4 matmat: each weight row is unpacked from the mmap ONCE
11171/// and dotted against ALL b activations (prefill used to fall back to b
11172/// full matvecs — b× weight traffic and b× nibble decode). Per-position
11173/// math matches `q4matvec` exactly: same group order, same accumulation.
11174/// `out` is row-major [b, rows] like `qmatmat`.
11175#[allow(clippy::too_many_arguments)]
11176fn q4matmat(
11177    bytes: &[u8],
11178    xs_all: &[f32],
11179    b: usize,
11180    rows: usize,
11181    cols: usize,
11182    out: &mut [f32],
11183    pool: Option<&Pool>,
11184) {
11185    debug_assert_eq!(xs_all.len(), b * cols);
11186    debug_assert_eq!(out.len(), b * rows);
11187    let (packed, scales) = q4_split(bytes, rows, cols);
11188    let gpr = cols / GROUP_SIZE;
11189    let gscale = |g: usize| f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
11190
11191    if a8w8_enabled() {
11192        let acts: Vec<SplitAct> = (0..b)
11193            .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
11194            .collect();
11195        let acts = &acts;
11196        let out_addr = SendMut(out.as_mut_ptr());
11197        let run = move |start: usize, end: usize| {
11198            ROW_I8.with(|rb| {
11199                let mut buf = rb.borrow_mut();
11200                buf.resize(cols, 0);
11201                for r in start..end {
11202                    // Unpack the row's nibbles to centered i8 once
11203                    // (element 2k = low nibble, 2k+1 = high — flat order,
11204                    // same as dot_q4_row_sdot's zip).
11205                    for gi in 0..gpr {
11206                        let g = r * gpr + gi;
11207                        for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
11208                            buf[gi * GROUP_SIZE + k * 2] = ((bt & 0x0F) as i32 - 8) as i8 as u8;
11209                            buf[gi * GROUP_SIZE + k * 2 + 1] =
11210                                (((bt >> 4) & 0x0F) as i32 - 8) as i8 as u8;
11211                        }
11212                    }
11213                    let mut bi = 0usize;
11214                    #[cfg(target_arch = "x86_64")]
11215                    if avx2_enabled() && blocked_enabled() {
11216                        while bi + 4 <= acts.len() {
11217                            let xs = [
11218                                acts[bi].xq.as_slice(),
11219                                acts[bi + 1].xq.as_slice(),
11220                                acts[bi + 2].xq.as_slice(),
11221                                acts[bi + 3].xq.as_slice(),
11222                            ];
11223                            let d = unsafe {
11224                                if vnni_tiles_enabled() {
11225                                    dot_q4b_row_1x4_vnni(&buf, scales, r * gpr, gpr, xs)
11226                                } else {
11227                                    dot_q4b_row_1x4_avx2(&buf, scales, r * gpr, gpr, xs)
11228                                }
11229                            };
11230                            for k in 0..4 {
11231                                let act = &acts[bi + k];
11232                                let mut acc = d[k] * act.sx;
11233                                for &(j, xv) in &act.outliers {
11234                                    acc += (buf[j] as i8) as f32
11235                                        * gscale((r * cols + j) / GROUP_SIZE)
11236                                        * xv;
11237                                }
11238                                // SAFETY: disjoint (bi, r) cells per worker.
11239                                unsafe { *out_addr.at((bi + k) * rows + r) = acc };
11240                            }
11241                            bi += 4;
11242                        }
11243                    }
11244                    while bi < acts.len() {
11245                        let act = &acts[bi];
11246                        let mut acc = 0f32;
11247                        for gi in 0..gpr {
11248                            let d = dot_i8_i8(
11249                                &buf[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
11250                                &act.xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
11251                            );
11252                            acc += d as f32 * gscale(r * gpr + gi);
11253                        }
11254                        acc *= act.sx;
11255                        // xq is zeroed at outlier slots — exact terms.
11256                        for &(j, xv) in &act.outliers {
11257                            acc += (buf[j] as i8) as f32 * gscale((r * cols + j) / GROUP_SIZE) * xv;
11258                        }
11259                        // SAFETY: disjoint (bi, r) cells per worker row range.
11260                        unsafe { *out_addr.at(bi * rows + r) = acc };
11261                        bi += 1;
11262                    }
11263                }
11264            })
11265        };
11266        dispatch_rows(pool, rows, &run);
11267        return;
11268    }
11269
11270    let out_addr = SendMut(out.as_mut_ptr());
11271    let run = move |start: usize, end: usize| {
11272        ROW_F32.with(|rb| {
11273            let mut buf = rb.borrow_mut();
11274            buf.resize(cols, 0.0);
11275            for r in start..end {
11276                // Decode raw (nib − 8) values once; scales stay per-group
11277                // so the accumulation order matches q4matvec bit-for-bit.
11278                for gi in 0..gpr {
11279                    let g = r * gpr + gi;
11280                    for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
11281                        buf[gi * GROUP_SIZE + k * 2] = (bt & 0x0F) as f32 - 8.0;
11282                        buf[gi * GROUP_SIZE + k * 2 + 1] = ((bt >> 4) & 0x0F) as f32 - 8.0;
11283                    }
11284                }
11285                for bi in 0..b {
11286                    let x = &xs_all[bi * cols..(bi + 1) * cols];
11287                    let mut acc = 0f32;
11288                    for gi in 0..gpr {
11289                        let mut ga = 0f32;
11290                        // Pairwise (lo + hi) addition, matching
11291                        // q4matvec's `ga += lo·x + hi·x` shape exactly —
11292                        // a flat one-per-element loop rounds differently
11293                        // and broke bit-parity on the scalar (x86) path.
11294                        for k in 0..GROUP_SIZE / 2 {
11295                            let e = gi * GROUP_SIZE + k * 2;
11296                            ga += buf[e] * x[e] + buf[e + 1] * x[e + 1];
11297                        }
11298                        acc += ga * gscale(r * gpr + gi);
11299                    }
11300                    // SAFETY: disjoint (bi, r) cells per worker row range.
11301                    unsafe { *out_addr.at(bi * rows + r) = acc };
11302                }
11303            }
11304        })
11305    };
11306    dispatch_rows(pool, rows, &run);
11307}
11308
11309/// Batched vbit matmat: each variable-bit row is decoded from the mmap
11310/// ONCE for the whole microbatch. Same per-position math as
11311/// `vbitmatvec` (SDOT A8W8 with exact outliers / exact f32 for b=8 rows
11312/// and the scalar path).
11313#[allow(clippy::too_many_arguments)]
11314fn vbitmatmat(
11315    bytes: &[u8],
11316    offsets: &[usize],
11317    xs_all: &[f32],
11318    b: usize,
11319    rows: usize,
11320    cols: usize,
11321    out: &mut [f32],
11322    pool: Option<&Pool>,
11323) {
11324    debug_assert_eq!(xs_all.len(), b * cols);
11325    debug_assert_eq!(out.len(), b * rows);
11326    debug_assert_eq!(offsets.len(), rows + 1);
11327    let ng = cols / GROUP_SIZE;
11328    let bits = &bytes[..rows];
11329    let sc_off = rows;
11330    let gscale = |r: usize, g: usize| {
11331        let so = (r * ng + g) * 2;
11332        f16_to_f32(u16::from_le_bytes([
11333            bytes[sc_off + so],
11334            bytes[sc_off + so + 1],
11335        ]))
11336    };
11337
11338    // Decode row r's raw (u − L) values into `dst` (f32, unscaled).
11339    let decode_f32 = |r: usize, dst: &mut [f32]| {
11340        let bw = bits[r] as usize;
11341        let l = ((1i32 << (bw - 1)) - 1) as f32;
11342        let data = &bytes[offsets[r]..offsets[r + 1]];
11343        let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
11344        for d in dst.iter_mut() {
11345            while nbits < bw {
11346                acc = (acc << 8) | data[idx] as u64;
11347                idx += 1;
11348                nbits += 8;
11349            }
11350            let u = ((acc >> (nbits - bw)) & ((1u64 << bw) - 1)) as f32;
11351            nbits -= bw;
11352            *d = u - l;
11353        }
11354    };
11355
11356    if a8w8_enabled() {
11357        let acts: Vec<SplitAct> = (0..b)
11358            .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
11359            .collect();
11360        let acts = &acts;
11361        let out_addr = SendMut(out.as_mut_ptr());
11362        let run = move |start: usize, end: usize| {
11363            for r in start..end {
11364                let bw = bits[r] as usize;
11365                if bw == 8 {
11366                    // u−L reaches 128 → no i8 path; decode once, exact
11367                    // f32 dots for every position (same as vbitmatvec).
11368                    ROW_F32.with(|rb| {
11369                        let mut buf = rb.borrow_mut();
11370                        buf.resize(cols, 0.0);
11371                        decode_f32(r, &mut buf);
11372                        for bi in 0..b {
11373                            let x = &xs_all[bi * cols..(bi + 1) * cols];
11374                            let mut dot = 0f32;
11375                            for g in 0..ng {
11376                                let mut gd = 0f32;
11377                                for k in 0..GROUP_SIZE {
11378                                    gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
11379                                }
11380                                dot += gd * gscale(r, g);
11381                            }
11382                            // SAFETY: disjoint (bi, r) cells per worker range.
11383                            unsafe { *out_addr.at(bi * rows + r) = dot };
11384                        }
11385                    });
11386                    continue;
11387                }
11388                let l = (1i32 << (bw - 1)) - 1;
11389                let data = &bytes[offsets[r]..offsets[r + 1]];
11390                ROW_I8.with(|rb| {
11391                    let mut buf = rb.borrow_mut();
11392                    buf.resize(cols, 0);
11393                    #[inline(always)]
11394                    fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
11395                        for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
11396                            let u = unpack8::<B>(&data[blk * B..]);
11397                            for k in 0..8 {
11398                                chunk[k] = (u[k] - l) as i8 as u8;
11399                            }
11400                        }
11401                    }
11402                    match bw {
11403                        3 => fill::<3>(data, l, &mut buf),
11404                        4 => vbit_fill4(data, &mut buf),
11405                        5 => fill::<5>(data, l, &mut buf),
11406                        6 => fill::<6>(data, l, &mut buf),
11407                        _ => unreachable!("vbit bit-width {bw} (validated at load)"),
11408                    }
11409                    let mut bi = 0usize;
11410                    // The vbit scale table shares q4_block's layout
11411                    // (contiguous f16 per (row·ng + g)), so the same
11412                    // blocked 1×4 kernel serves the decoded row.
11413                    #[cfg(target_arch = "x86_64")]
11414                    if avx2_enabled() && blocked_enabled() {
11415                        while bi + 4 <= acts.len() {
11416                            let xs = [
11417                                acts[bi].xq.as_slice(),
11418                                acts[bi + 1].xq.as_slice(),
11419                                acts[bi + 2].xq.as_slice(),
11420                                acts[bi + 3].xq.as_slice(),
11421                            ];
11422                            let sxs = [
11423                                acts[bi].sx,
11424                                acts[bi + 1].sx,
11425                                acts[bi + 2].sx,
11426                                acts[bi + 3].sx,
11427                            ];
11428                            let d = unsafe {
11429                                if vnni_tiles_enabled() {
11430                                    dot_q4b_row_1x4_sx_vnni(
11431                                        &buf,
11432                                        &bytes[sc_off..],
11433                                        r * ng,
11434                                        ng,
11435                                        xs,
11436                                        sxs,
11437                                    )
11438                                } else {
11439                                    dot_q4b_row_1x4_sx_avx2(
11440                                        &buf,
11441                                        &bytes[sc_off..],
11442                                        r * ng,
11443                                        ng,
11444                                        xs,
11445                                        sxs,
11446                                    )
11447                                }
11448                            };
11449                            for k in 0..4 {
11450                                let act = &acts[bi + k];
11451                                let mut dot = d[k];
11452                                for &(j, xv) in &act.outliers {
11453                                    dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
11454                                }
11455                                // SAFETY: disjoint (bi, r) cells per worker.
11456                                unsafe { *out_addr.at((bi + k) * rows + r) = dot };
11457                            }
11458                            bi += 4;
11459                        }
11460                    }
11461                    while bi < acts.len() {
11462                        let act = &acts[bi];
11463                        let mut dot = 0f32;
11464                        for g in 0..ng {
11465                            let d = dot_i8_i8(
11466                                &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
11467                                &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
11468                            ) as f32
11469                                * act.sx;
11470                            dot += d * gscale(r, g);
11471                        }
11472                        for &(j, xv) in &act.outliers {
11473                            dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
11474                        }
11475                        // SAFETY: disjoint (bi, r) cells per worker range.
11476                        unsafe { *out_addr.at(bi * rows + r) = dot };
11477                        bi += 1;
11478                    }
11479                });
11480            }
11481        };
11482        dispatch_rows(pool, rows, &run);
11483        return;
11484    }
11485
11486    let out_addr = SendMut(out.as_mut_ptr());
11487    let run = move |start: usize, end: usize| {
11488        ROW_F32.with(|rb| {
11489            let mut buf = rb.borrow_mut();
11490            buf.resize(cols, 0.0);
11491            for r in start..end {
11492                decode_f32(r, &mut buf);
11493                for bi in 0..b {
11494                    let x = &xs_all[bi * cols..(bi + 1) * cols];
11495                    let mut dot = 0f32;
11496                    for g in 0..ng {
11497                        let mut gd = 0f32;
11498                        for k in 0..GROUP_SIZE {
11499                            gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
11500                        }
11501                        dot += gd * gscale(r, g);
11502                    }
11503                    // SAFETY: disjoint (bi, r) cells per worker range.
11504                    unsafe { *out_addr.at(bi * rows + r) = dot };
11505                }
11506            }
11507        })
11508    };
11509    dispatch_rows(pool, rows, &run);
11510}
11511
11512/// Build a GPU batch job for a q8-family mapped tensor (primary
11513/// shard): prescaled input + directory coordinates. None → not
11514/// GPU-eligible, caller stays on the CPU.
11515pub(crate) fn gpu_batch_job<'a>(
11516    t: &'a QTensor,
11517    x: &[f32],
11518) -> Option<(std::sync::Arc<CmfModel>, crate::gpu::BatchJob<'a>)> {
11519    match t {
11520        QTensor::Mapped {
11521            model,
11522            idx,
11523            dtype: dt @ (TensorDtype::Q8Row | TensorDtype::Q8_2f),
11524            rows,
11525            cols,
11526            row_scale,
11527            col_field,
11528            ..
11529        } => Some((
11530            model.clone(),
11531            crate::gpu::BatchJob {
11532                idx: *idx,
11533                rows: *rows,
11534                cols: *cols,
11535                row_scale,
11536                xs: prescale(x, col_field, *dt).into_owned(),
11537                layout: crate::gpu::BatchLayout::Q8,
11538            },
11539        )),
11540        // q1: raw f32 activations, tile-embedded scales.
11541        QTensor::Mapped {
11542            model,
11543            idx,
11544            dtype: TensorDtype::Q1,
11545            rows,
11546            cols,
11547            ..
11548        } => Some((
11549            model.clone(),
11550            crate::gpu::BatchJob {
11551                idx: *idx,
11552                rows: *rows,
11553                cols: *cols,
11554                row_scale: &[],
11555                xs: x.to_vec(),
11556                layout: crate::gpu::BatchLayout::Q1,
11557            },
11558        )),
11559        // q1g: same raw-f32 contract, sign + scale planes.
11560        QTensor::Mapped {
11561            model,
11562            idx,
11563            dtype: TensorDtype::Q1G,
11564            rows,
11565            cols,
11566            ..
11567        } => Some((
11568            model.clone(),
11569            crate::gpu::BatchJob {
11570                idx: *idx,
11571                rows: *rows,
11572                cols: *cols,
11573                row_scale: &[],
11574                xs: x.to_vec(),
11575                layout: crate::gpu::BatchLayout::Q1g,
11576            },
11577        )),
11578        // q4_tiled / q4tp: raw f32 activations; the scales live in the
11579        // payload (inline tiles / row ladder), so row_scale stays empty.
11580        // The GDN projection batch already runs these layouts on Metal —
11581        // this arm lets the attention QKV batch reach the same kernels.
11582        QTensor::Mapped {
11583            model,
11584            idx,
11585            dtype: dt @ (TensorDtype::Q4Tiled | TensorDtype::Q4TiledP),
11586            rows,
11587            cols,
11588            ..
11589        } => Some((
11590            model.clone(),
11591            crate::gpu::BatchJob {
11592                idx: *idx,
11593                rows: *rows,
11594                cols: *cols,
11595                row_scale: &[],
11596                xs: x.to_vec(),
11597                layout: if *dt == TensorDtype::Q4Tiled {
11598                    crate::gpu::BatchLayout::Q4t
11599                } else {
11600                    crate::gpu::BatchLayout::Q4tp
11601                },
11602            },
11603        )),
11604        _ => None,
11605    }
11606}
11607
11608thread_local! {
11609    static PRESCALE_BUF1: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
11610    static PRESCALE_BUF2: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
11611}
11612
11613pub(crate) fn prescale<'a>(
11614    x: &'a [f32],
11615    col_field: &[f32],
11616    dtype: TensorDtype,
11617) -> std::borrow::Cow<'a, [f32]> {
11618    if dtype == TensorDtype::Q8_2f {
11619        x.iter().zip(col_field).map(|(a, c)| a * c).collect()
11620    } else {
11621        std::borrow::Cow::Borrowed(x)
11622    }
11623}
11624
11625/// θ col-field fold for q8_2f activations. Borrowed pass-through for
11626/// every other dtype, using thread-local buffers to eliminate per-matvec allocations.
11627pub(crate) fn prescale_with<R, F: FnOnce(&[f32]) -> R>(
11628    x: &[f32],
11629    col_field: &[f32],
11630    dtype: TensorDtype,
11631    buf_id: u8,
11632    f: F,
11633) -> R {
11634    if dtype == TensorDtype::Q8_2f {
11635        if buf_id == 1 {
11636            PRESCALE_BUF1.with(|b| {
11637                let mut buf = b.borrow_mut();
11638                buf.clear();
11639                buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
11640                f(&buf)
11641            })
11642        } else {
11643            PRESCALE_BUF2.with(|b| {
11644                let mut buf = b.borrow_mut();
11645                buf.clear();
11646                buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
11647                f(&buf)
11648            })
11649        }
11650    } else {
11651        f(x)
11652    }
11653}
11654
11655// ───────────────────── x86-64 AVX2 kernels (roadmap этап 2) ─────────────────────
11656
11657/// AVX2+FMA available? Default ON when the CPU supports both;
11658/// `CMF_AVX2=0` disables (falls back to the autovectorized loops).
11659#[cfg(target_arch = "x86_64")]
11660pub(crate) fn avx2_enabled() -> bool {
11661    use std::sync::OnceLock;
11662    static ON: OnceLock<bool> = OnceLock::new();
11663    *ON.get_or_init(|| {
11664        std::env::var("CMF_AVX2").map(|v| v != "0").unwrap_or(true)
11665            && std::arch::is_x86_feature_detected!("avx2")
11666            && std::arch::is_x86_feature_detected!("fma")
11667    })
11668}
11669
11670/// AVX2 A8W8 allowed? The quantized-activation contract is switched by
11671/// the SAME env as the ARM SDOT path: `CMF_SDOT=0` keeps exact kernels
11672/// (the golden-parity exact gate relies on it) — AVX2 f32 kernels stay
11673/// active either way, they are exact (regrouped sums only).
11674#[cfg(target_arch = "x86_64")]
11675fn avx2_a8w8_enabled() -> bool {
11676    if FLOAT_ACTIVATIONS.get() {
11677        return false;
11678    }
11679    use std::sync::OnceLock;
11680    static ON: OnceLock<bool> = OnceLock::new();
11681    *ON.get_or_init(|| {
11682        avx2_enabled() && std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true)
11683    })
11684}
11685
11686thread_local! {
11687    static FULL_GPU_Q8: std::cell::Cell<bool> = const { std::cell::Cell::new(false) };
11688}
11689
11690/// Match the graph's full-device q8 projection precision on MiMo's host
11691/// tail. Does not enable the GPU or bypass a CPU-only/device-refusal gate.
11692pub(crate) fn enter_full_gpu_q8_scope() -> impl Drop {
11693    struct Restore(bool, std::marker::PhantomData<std::rc::Rc<()>>);
11694    impl Drop for Restore {
11695        fn drop(&mut self) {
11696            FULL_GPU_Q8.set(self.0);
11697        }
11698    }
11699    Restore(FULL_GPU_Q8.replace(true), std::marker::PhantomData)
11700}
11701
11702// Dynamic MiMo experts must not change activation precision when a cache
11703// fill moves them from CPU to GPU. Thread-local: only the cold-expert
11704// dispatch selects float kernels; concurrent pipelines keep their policy.
11705thread_local! {
11706    static FLOAT_ACTIVATIONS: std::cell::Cell<bool> = const { std::cell::Cell::new(false) };
11707}
11708
11709/// Flint's released Q2/Q4 profiles are unusually sensitive to activation
11710/// quantization in the host-side attention path. With the normal A8W8
11711/// default, a short greedy decode can diverge after the first token even
11712/// though the selected-expert Metal path itself is unchanged. Keep the
11713/// reference float-activation contract for Flint unless the operator has
11714/// explicitly selected an `CMF_SDOT` policy.
11715fn flint_exact_cpu_default_enabled() -> bool {
11716    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
11717    *ON.get_or_init(|| std::env::var_os("CMF_SDOT").is_none())
11718}
11719
11720impl QTensor {
11721    #[inline]
11722    fn flint_exact_cpu_default(&self) -> bool {
11723        !FLOAT_ACTIVATIONS.get()
11724            && flint_exact_cpu_default_enabled()
11725            && matches!(
11726                self,
11727                Self::Mapped { model, .. } if model.header.arch.arch_name == "flint"
11728            )
11729    }
11730
11731    #[inline]
11732    fn any_flint_exact_cpu_default(ts: &[&QTensor]) -> bool {
11733        ts.iter().any(|t| t.flint_exact_cpu_default())
11734    }
11735
11736    #[inline]
11737    fn any_flint_exact_cpu_default_pairs(pairs: &[(&QTensor, &QTensor)]) -> bool {
11738        pairs
11739            .iter()
11740            .any(|(gate, up)| gate.flint_exact_cpu_default() || up.flint_exact_cpu_default())
11741    }
11742}
11743
11744/// Run `f` with the exact float-activation contract on this thread: the
11745/// 4-bit and 8-bit kernels skip the int8 activation split (`CMF_SDOT=0`
11746/// semantics, scoped). For parity tests and diagnostics — e.g. an `mxfp4`
11747/// model then computes bit for bit what its dequantized f32 twin computes.
11748pub fn with_float_activations<R>(f: impl FnOnce() -> R) -> R {
11749    float_activations_scope(f)
11750}
11751
11752/// Active on the current thread. `Pool` snapshots this with its job
11753/// descriptor and installs it on persistent workers, so a scoped quantized
11754/// matvec cannot accidentally choose a different kernel on a worker.
11755#[inline]
11756pub(crate) fn float_activations_active() -> bool {
11757    FLOAT_ACTIVATIONS.get()
11758}
11759
11760/// Restores the current thread's float-activation state on drop. The `Rc`
11761/// marker deliberately keeps the guard on the thread that created it: this
11762/// policy is thread-local, not a portable capability.
11763pub(crate) struct FloatActivationsScope {
11764    previous: bool,
11765    _not_send: std::marker::PhantomData<std::rc::Rc<()>>,
11766}
11767
11768impl Drop for FloatActivationsScope {
11769    fn drop(&mut self) {
11770        FLOAT_ACTIVATIONS.set(self.previous);
11771    }
11772}
11773
11774#[inline]
11775pub(crate) fn enter_float_activations_scope() -> FloatActivationsScope {
11776    FloatActivationsScope {
11777        previous: FLOAT_ACTIVATIONS.replace(true),
11778        _not_send: std::marker::PhantomData,
11779    }
11780}
11781
11782/// Run `f` with exact float activations (no A8W8 int8 activation
11783/// quantization in the quantized matvecs). The persistent row pool carries
11784/// this scope into a dispatched job and restores each worker afterwards.
11785/// Public for parity tests that compare the default kernels with the strict
11786/// ones in one process; `CMF_SDOT=0` is the process-wide equivalent.
11787#[doc(hidden)]
11788pub fn float_activations_scope<R>(f: impl FnOnce() -> R) -> R {
11789    let _scope = enter_float_activations_scope();
11790    f()
11791}
11792
11793/// Row-exact batching: while set, the x86 batched kernels (`qmatmat`,
11794/// `q4tp_matmat`) compute every (weight row, token) cell with the
11795/// single-token kernel instead of the blocked 2×4 / 1×4 / 1×8 tiles, so a
11796/// token's result does not depend on the batch it rides in and equals its
11797/// matvec. On ARM `q4tp_matmat` keeps its 1×4 tile but in the matvec's
11798/// reduction order, and no q4tp batch takes the AMX or device GEMM. The
11799/// MiMo speculative verify holds it (`row_exact_scope`) — its accepted
11800/// rows must be the rows plain decode would have produced.
11801// Shared with pool workers, so overlapping requests must keep the mode
11802// enabled until the LAST scope leaves. Saving/restoring a global bool is
11803// incorrect when two threads enter and leave in a non-LIFO order.
11804static ROW_EXACT: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
11805
11806pub(crate) fn row_exact() -> bool {
11807    ROW_EXACT.load(std::sync::atomic::Ordering::Acquire) != 0
11808}
11809
11810fn counted_row_exact_scope<R>(active: &std::sync::atomic::AtomicUsize, f: impl FnOnce() -> R) -> R {
11811    struct Restore<'a>(&'a std::sync::atomic::AtomicUsize);
11812    impl Drop for Restore<'_> {
11813        fn drop(&mut self) {
11814            self.0.fetch_sub(1, std::sync::atomic::Ordering::AcqRel);
11815        }
11816    }
11817    active.fetch_add(1, std::sync::atomic::Ordering::AcqRel);
11818    let _restore = Restore(active);
11819    f()
11820}
11821
11822/// Run `f` with row-exact batching on (also released on unwind).
11823pub(crate) fn row_exact_scope<R>(f: impl FnOnce() -> R) -> R {
11824    counted_row_exact_scope(&ROW_EXACT, f)
11825}
11826
11827/// A8W8 quantized-activation path available on THIS machine? One
11828/// switch across architectures: ARM dotprod (CMF_SDOT) or x86 AVX2
11829/// (CMF_AVX2 + the same CMF_SDOT exact-contract override).
11830#[inline]
11831pub(crate) fn a8w8_enabled() -> bool {
11832    #[cfg(target_arch = "aarch64")]
11833    {
11834        sdot_enabled()
11835    }
11836    #[cfg(target_arch = "x86_64")]
11837    {
11838        avx2_a8w8_enabled()
11839    }
11840    #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
11841    {
11842        false
11843    }
11844}
11845
11846/// int8·int8 dot dispatch: SDOT on ARM; AVX-512 VNNI (vpdpbusd) or AVX2
11847/// maddubs on x86. Callers are gated by `a8w8_enabled()`.
11848#[inline]
11849#[allow(unreachable_code)]
11850fn dot_i8_i8(w: &[u8], xq: &[i8]) -> i32 {
11851    #[cfg(target_arch = "aarch64")]
11852    unsafe {
11853        return dot_i8_sdot(w, xq);
11854    }
11855    #[cfg(target_arch = "x86_64")]
11856    unsafe {
11857        if avx512vnni_enabled() {
11858            return dot_i8_i8_vnni(w, xq);
11859        }
11860        return dot_i8_i8_avx2(w, xq);
11861    }
11862    w.iter()
11863        .zip(xq)
11864        .map(|(&a, &b)| (a as i8) as i32 * b as i32)
11865        .sum()
11866}
11867
11868/// AVX-512 VNNI available? (F+BW+VL+VNNI; `CMF_AVX512=0` falls back to
11869/// AVX2.) VL matters: short 32-byte groups (q4/vbit) ride the 256-bit
11870/// `vpdpbusd` encoding.
11871#[cfg(target_arch = "x86_64")]
11872fn avx512vnni_enabled() -> bool {
11873    use std::sync::OnceLock;
11874    static ON: OnceLock<bool> = OnceLock::new();
11875    *ON.get_or_init(|| {
11876        std::env::var("CMF_AVX512")
11877            .map(|v| v != "0")
11878            .unwrap_or(true)
11879            && std::arch::is_x86_feature_detected!("avx512f")
11880            && std::arch::is_x86_feature_detected!("avx512bw")
11881            && std::arch::is_x86_feature_detected!("avx512vl")
11882            && std::arch::is_x86_feature_detected!("avx512vnni")
11883    })
11884}
11885
11886/// Grouped-codec VNNI arms (the q4t/q4b/q1/q1t tile kernels): default
11887/// ON where AVX-512 VNNI exists (`CMF_VNNI_TILES=0` opt-out). Measured
11888/// on Ryzen 7950X (Zen4, 3 alternating process pairs, blocked GEMM
11889/// 4864×896 b=256): q4t 63→68 GF/s (+8%), q1 53→56 (+6%), q4b 72→75
11890/// (+4%) — consistent, no leg regressed. The tile kernels keep a
11891/// horizontal reduce per 32-weight group, so the `vpdpbusd` saving is
11892/// smaller than the long-dot q8 win (+13%), but it is real and free.
11893#[cfg(target_arch = "x86_64")]
11894fn vnni_tiles_enabled() -> bool {
11895    use std::sync::OnceLock;
11896    static ON: OnceLock<bool> = OnceLock::new();
11897    *ON.get_or_init(|| {
11898        std::env::var("CMF_VNNI_TILES")
11899            .map(|v| v != "0")
11900            .unwrap_or(true)
11901            && avx512vnni_enabled()
11902    })
11903}
11904
11905/// One 256-bit u8×i8 dot → i32 via `vpdpbusd` into a fresh accumulator
11906/// plus the same horizontal reduce the AVX2 kernels use. Products are
11907/// bounded (|w| ≤ 8 or ≤ 1), so maddubs never saturated — the i32 sum
11908/// is bit-identical to the maddubs+madd pair it replaces.
11909#[cfg(target_arch = "x86_64")]
11910#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
11911#[inline]
11912unsafe fn dpbusd_hsum(aw: core::arch::x86_64::__m256i, xs: core::arch::x86_64::__m256i) -> i32 {
11913    // SAFETY: pure register math.
11914    unsafe {
11915        use core::arch::x86_64::*;
11916        let d = _mm256_dpbusd_epi32(_mm256_setzero_si256(), aw, xs);
11917        let hi128 = _mm256_extracti128_si256::<1>(d);
11918        let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
11919        let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
11920        let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
11921        _mm_cvtsi128_si32(s32)
11922    }
11923}
11924
11925/// int8·int8 via AVX-512 VNNI: `vpdpbusd` fuses the maddubs+madd+add
11926/// triple into one u8×i8 dot-accumulate. AVX-512 has no vpsignb, so the
11927/// |w|·sign(x,w) trick becomes |w| × (x negated where w<0) via a mask
11928/// subtract — w==0 lanes contribute 0 through |w|=0 either way.
11929#[cfg(target_arch = "x86_64")]
11930#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
11931unsafe fn dot_i8_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
11932    // SAFETY: callers uphold slice-length contracts (see call sites).
11933    unsafe {
11934        use core::arch::x86_64::*;
11935        let n = w.len();
11936        let mut j = 0usize;
11937        let mut total: i32;
11938        // 4 independent accumulators: vpdpbusd is its own loop-carried
11939        // dependency (~5-cycle latency) — a single-acc loop runs
11940        // latency-bound and LOSES to the AVX2 maddubs kernel, measured
11941        // on Granite Rapids.
11942        {
11943            #[inline(always)]
11944            unsafe fn step(
11945                w: *const u8,
11946                x: *const i8,
11947                acc: core::arch::x86_64::__m512i,
11948            ) -> core::arch::x86_64::__m512i {
11949                unsafe {
11950                    use core::arch::x86_64::*;
11951                    let wv = _mm512_loadu_si512(w as *const _);
11952                    let xv = _mm512_loadu_si512(x as *const _);
11953                    let aw = _mm512_abs_epi8(wv);
11954                    let neg = _mm512_movepi8_mask(wv);
11955                    let sx = _mm512_mask_sub_epi8(xv, neg, _mm512_setzero_si512(), xv);
11956                    _mm512_dpbusd_epi32(acc, aw, sx)
11957                }
11958            }
11959            let (mut a0, mut a1, mut a2, mut a3) = (
11960                _mm512_setzero_si512(),
11961                _mm512_setzero_si512(),
11962                _mm512_setzero_si512(),
11963                _mm512_setzero_si512(),
11964            );
11965            while j + 256 <= n {
11966                a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
11967                a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), a1);
11968                a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), a2);
11969                a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), a3);
11970                j += 256;
11971            }
11972            while j + 64 <= n {
11973                a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
11974                j += 64;
11975            }
11976            let s01 = _mm512_add_epi32(a0, a1);
11977            let s23 = _mm512_add_epi32(a2, a3);
11978            total = _mm512_reduce_add_epi32(_mm512_add_epi32(s01, s23));
11979        }
11980        // 32-wide (q4/vbit groups are exactly 32 bytes).
11981        if j + 32 <= n {
11982            let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
11983            let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
11984            let d = _mm256_dpbusd_epi32(
11985                _mm256_setzero_si256(),
11986                _mm256_abs_epi8(wv),
11987                _mm256_sign_epi8(xv, wv),
11988            );
11989            let hi128 = _mm256_extracti128_si256::<1>(d);
11990            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
11991            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
11992            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
11993            total += _mm_cvtsi128_si32(s32);
11994            j += 32;
11995        }
11996        while j < n {
11997            total += (w[j] as i8) as i32 * xq[j] as i32;
11998            j += 1;
11999        }
12000        total
12001    }
12002}
12003
12004/// i8 row · f32 x via AVX2/FMA (x86 mirror of `dot_i8_f32_neon`).
12005#[cfg(target_arch = "x86_64")]
12006#[target_feature(enable = "avx2,fma")]
12007unsafe fn dot_i8_f32_avx2(w: &[u8], x: &[f32]) -> f32 {
12008    // SAFETY: callers uphold slice-length contracts (see call sites).
12009    unsafe {
12010        use core::arch::x86_64::*;
12011        let n = x.len();
12012        let wp = w.as_ptr();
12013        let xp = x.as_ptr();
12014        let (mut a0, mut a1) = (_mm256_setzero_ps(), _mm256_setzero_ps());
12015        let mut j = 0usize;
12016        while j + 16 <= n {
12017            let wb = _mm_loadu_si128(wp.add(j) as *const __m128i);
12018            let lo = _mm256_cvtepi8_epi32(wb);
12019            let hi = _mm256_cvtepi8_epi32(_mm_srli_si128::<8>(wb));
12020            a0 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(lo), _mm256_loadu_ps(xp.add(j)), a0);
12021            a1 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(hi), _mm256_loadu_ps(xp.add(j + 8)), a1);
12022            j += 16;
12023        }
12024        let acc = _mm256_add_ps(a0, a1);
12025        let hi128 = _mm256_extractf128_ps::<1>(acc);
12026        let s128 = _mm_add_ps(_mm256_castps256_ps128(acc), hi128);
12027        let s64 = _mm_add_ps(s128, _mm_movehl_ps(s128, s128));
12028        let s32 = _mm_add_ss(s64, _mm_shuffle_ps::<1>(s64, s64));
12029        let mut sum = _mm_cvtss_f32(s32);
12030        while j < n {
12031            sum += (*wp.add(j) as i8) as f32 * *xp.add(j);
12032            j += 1;
12033        }
12034        sum
12035    }
12036}
12037
12038/// int8(weight)·int8(activation) → i32 via AVX2 maddubs — the x86
12039/// analogue of the SDOT A8W8 path. `maddubs` takes u8×i8, so the
12040/// standard sign trick applies: |w| × sign(x, w) ≡ w × x per lane.
12041/// Pair saturation is safe: |w|≤128, |x|≤127 → 2·128·127 < 32767.
12042#[cfg(target_arch = "x86_64")]
12043#[target_feature(enable = "avx2")]
12044unsafe fn dot_i8_i8_avx2(w: &[u8], xq: &[i8]) -> i32 {
12045    // SAFETY: callers uphold slice-length contracts (see call sites).
12046    unsafe {
12047        use core::arch::x86_64::*;
12048        let n = w.len();
12049        let ones = _mm256_set1_epi16(1);
12050        let mut acc = _mm256_setzero_si256();
12051        let mut j = 0usize;
12052        while j + 32 <= n {
12053            let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
12054            let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
12055            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
12056            acc = _mm256_add_epi32(acc, _mm256_madd_epi16(p16, ones));
12057            j += 32;
12058        }
12059        let hi128 = _mm256_extracti128_si256::<1>(acc);
12060        let s128 = _mm_add_epi32(_mm256_castsi256_si128(acc), hi128);
12061        let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
12062        let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
12063        let mut s = _mm_cvtsi128_si32(s32);
12064        while j < n {
12065            s += (w[j] as i8) as i32 * xq[j] as i32;
12066            j += 1;
12067        }
12068        s
12069    }
12070}
12071
12072/// smmla 2×4: one instruction covers a 2-row × 2-activation × 8-deep
12073/// tile (32 MACs vs sdot's 16) — the weight pair loads once per 8-k
12074/// slice as a combined 2×8 register and meets two activation pairs.
12075#[cfg(target_arch = "aarch64")]
12076#[target_feature(enable = "neon,i8mm")]
12077unsafe fn dot_i8_smmla_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
12078    // SAFETY: callers uphold slice-length contracts.
12079    unsafe {
12080        use core::arch::aarch64::*;
12081        use core::arch::asm;
12082        let n = w0.len();
12083        let w0p = w0.as_ptr() as *const i8;
12084        let w1p = w1.as_ptr() as *const i8;
12085        // acc01 holds [c(r0,x0) c(r0,x1) c(r1,x0) c(r1,x1)]; acc23 the
12086        // same for x2/x3.
12087        let mut acc01 = vdupq_n_s32(0);
12088        let mut acc23 = vdupq_n_s32(0);
12089        let mut i = 0usize;
12090        while i + 8 <= n {
12091            let wa = vcombine_s8(vld1_s8(w0p.add(i)), vld1_s8(w1p.add(i)));
12092            let xb01 = vcombine_s8(
12093                vld1_s8(xs[0].as_ptr().add(i)),
12094                vld1_s8(xs[1].as_ptr().add(i)),
12095            );
12096            let xb23 = vcombine_s8(
12097                vld1_s8(xs[2].as_ptr().add(i)),
12098                vld1_s8(xs[3].as_ptr().add(i)),
12099            );
12100            asm!(
12101                "smmla {a01:v}.4s, {w:v}.16b, {x01:v}.16b",
12102                "smmla {a23:v}.4s, {w:v}.16b, {x23:v}.16b",
12103                a01 = inout(vreg) acc01, a23 = inout(vreg) acc23,
12104                w = in(vreg) wa, x01 = in(vreg) xb01, x23 = in(vreg) xb23,
12105                options(pure, nomem, nostack),
12106            );
12107            i += 8;
12108        }
12109        let mut out = [[0i32; 4]; 2];
12110        let a01: [i32; 4] = core::mem::transmute(acc01);
12111        let a23: [i32; 4] = core::mem::transmute(acc23);
12112        out[0][0] = a01[0];
12113        out[0][1] = a01[1];
12114        out[1][0] = a01[2];
12115        out[1][1] = a01[3];
12116        out[0][2] = a23[0];
12117        out[0][3] = a23[1];
12118        out[1][2] = a23[2];
12119        out[1][3] = a23[3];
12120        if i < n {
12121            for (k, x) in xs.iter().enumerate() {
12122                for j in i..n {
12123                    out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
12124                    out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
12125                }
12126            }
12127        }
12128        out
12129    }
12130}
12131
12132/// ARM twin of the x86 blocked prefill GEMM: two weight rows stay in
12133/// registers across four activation streams, eight sdot accumulators.
12134/// (The per-row form re-read each W row once per activation.)
12135#[cfg(target_arch = "aarch64")]
12136#[target_feature(enable = "neon,dotprod")]
12137unsafe fn dot_i8_sdot_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
12138    // SAFETY: callers uphold slice-length contracts.
12139    unsafe {
12140        use core::arch::aarch64::*;
12141        use core::arch::asm;
12142        let n = w0.len();
12143        let w0p = w0.as_ptr() as *const i8;
12144        let w1p = w1.as_ptr() as *const i8;
12145        let mut acc = [[vdupq_n_s32(0); 4]; 2];
12146        let mut i = 0usize;
12147        while i + 16 <= n {
12148            let wv0 = vld1q_s8(w0p.add(i));
12149            let wv1 = vld1q_s8(w1p.add(i));
12150            for (k, x) in xs.iter().enumerate() {
12151                let xv = vld1q_s8(x.as_ptr().add(i));
12152                let (mut a0, mut a1) = (acc[0][k], acc[1][k]);
12153                asm!(
12154                    "sdot {a0:v}.4s, {w0:v}.16b, {x:v}.16b",
12155                    "sdot {a1:v}.4s, {w1:v}.16b, {x:v}.16b",
12156                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
12157                    w0 = in(vreg) wv0, w1 = in(vreg) wv1, x = in(vreg) xv,
12158                    options(pure, nomem, nostack),
12159                );
12160                acc[0][k] = a0;
12161                acc[1][k] = a1;
12162            }
12163            i += 16;
12164        }
12165        let mut out = [[0i32; 4]; 2];
12166        for r in 0..2 {
12167            for k in 0..4 {
12168                out[r][k] = vaddvq_s32(acc[r][k]);
12169            }
12170        }
12171        if i < n {
12172            for (k, x) in xs.iter().enumerate() {
12173                for j in i..n {
12174                    out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
12175                    out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
12176                }
12177            }
12178        }
12179        out
12180    }
12181}
12182
12183/// Blocked 2 weight rows × 4 activations for the prefill GEMM
12184/// (roadmap P0: packed panels + multi-row accumulators). The two rows'
12185/// abs() live in registers across all four activation streams; the
12186/// sign-fixup is recomputed per pair (the price of the maddubs trick).
12187/// Returns raw i8·i8 dots; the caller applies scales and outliers.
12188#[cfg(target_arch = "x86_64")]
12189#[target_feature(enable = "avx2")]
12190unsafe fn dot_i8_i8_avx2_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
12191    // SAFETY: callers uphold slice-length contracts.
12192    unsafe {
12193        use core::arch::x86_64::*;
12194        let n = w0.len();
12195        let ones = _mm256_set1_epi16(1);
12196        let mut acc = [[_mm256_setzero_si256(); 4]; 2];
12197        let mut j = 0usize;
12198        while j + 32 <= n {
12199            let wv0 = _mm256_loadu_si256(w0.as_ptr().add(j) as *const __m256i);
12200            let wv1 = _mm256_loadu_si256(w1.as_ptr().add(j) as *const __m256i);
12201            let aw0 = _mm256_abs_epi8(wv0);
12202            let aw1 = _mm256_abs_epi8(wv1);
12203            for (k, x) in xs.iter().enumerate() {
12204                let xv = _mm256_loadu_si256(x.as_ptr().add(j) as *const __m256i);
12205                let p0 = _mm256_maddubs_epi16(aw0, _mm256_sign_epi8(xv, wv0));
12206                acc[0][k] = _mm256_add_epi32(acc[0][k], _mm256_madd_epi16(p0, ones));
12207                let p1 = _mm256_maddubs_epi16(aw1, _mm256_sign_epi8(xv, wv1));
12208                acc[1][k] = _mm256_add_epi32(acc[1][k], _mm256_madd_epi16(p1, ones));
12209            }
12210            j += 32;
12211        }
12212        let mut out = [[0i32; 4]; 2];
12213        for r in 0..2 {
12214            for k in 0..4 {
12215                let a = acc[r][k];
12216                let hi128 = _mm256_extracti128_si256::<1>(a);
12217                let s128 = _mm_add_epi32(_mm256_castsi256_si128(a), hi128);
12218                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
12219                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
12220                out[r][k] = _mm_cvtsi128_si32(s32);
12221            }
12222        }
12223        if j < n {
12224            for (k, x) in xs.iter().enumerate() {
12225                for i in j..n {
12226                    out[0][k] += (w0[i] as i8) as i32 * x[i] as i32;
12227                    out[1][k] += (w1[i] as i8) as i32 * x[i] as i32;
12228                }
12229            }
12230        }
12231        out
12232    }
12233}
12234
12235/// AVX2/VNNI q8 row dot with exact outlier correction (x86 mirror of
12236/// `row_dot_sdot` — same A8W8 contract). With AVX-512 VNNI the row goes
12237/// through the bias trick: Σ(w+128)·x via pure `vpdpbusd` (no per-lane
12238/// sign fixups), corrected by −128·Σx with Σx precomputed per split.
12239#[cfg(target_arch = "x86_64")]
12240#[inline]
12241fn row_dot_avx2(row: &[u8], act: &SplitAct) -> f32 {
12242    let dot = if avx512vnni_enabled() && row.len() >= 64 {
12243        (unsafe { dot_u8p128_i8_vnni(row, &act.xq) }) - 128 * act.xsum
12244    } else {
12245        unsafe { dot_i8_i8_avx2(row, &act.xq) }
12246    };
12247    let mut acc = dot as f32 * act.sx;
12248    for &(j, xv) in &act.outliers {
12249        acc += (row[j] as i8) as f32 * xv;
12250    }
12251    acc
12252}
12253
12254/// Σ (w[i]+128)·x[i] via pure `vpdpbusd` — the caller subtracts
12255/// 128·Σx. Four independent accumulators (dpbusd is ~5-cycle latency;
12256/// a single-acc loop runs latency-bound, measured on Granite Rapids).
12257#[cfg(target_arch = "x86_64")]
12258#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
12259unsafe fn dot_u8p128_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
12260    // SAFETY: callers uphold slice-length contracts (see call sites).
12261    unsafe {
12262        use core::arch::x86_64::*;
12263        let n = w.len();
12264        let flip = _mm512_set1_epi8(-128); // XOR 0x80: i8 w → u8 (w+128)
12265        #[inline(always)]
12266        unsafe fn step(
12267            w: *const u8,
12268            x: *const i8,
12269            flip: core::arch::x86_64::__m512i,
12270            acc: core::arch::x86_64::__m512i,
12271        ) -> core::arch::x86_64::__m512i {
12272            unsafe {
12273                use core::arch::x86_64::*;
12274                let wv = _mm512_xor_si512(_mm512_loadu_si512(w as *const _), flip);
12275                _mm512_dpbusd_epi32(acc, wv, _mm512_loadu_si512(x as *const _))
12276            }
12277        }
12278        let (mut a0, mut a1, mut a2, mut a3) = (
12279            _mm512_setzero_si512(),
12280            _mm512_setzero_si512(),
12281            _mm512_setzero_si512(),
12282            _mm512_setzero_si512(),
12283        );
12284        let mut j = 0usize;
12285        while j + 256 <= n {
12286            a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
12287            a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), flip, a1);
12288            a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), flip, a2);
12289            a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), flip, a3);
12290            j += 256;
12291        }
12292        while j + 64 <= n {
12293            a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
12294            j += 64;
12295        }
12296        let mut total = _mm512_reduce_add_epi32(_mm512_add_epi32(
12297            _mm512_add_epi32(a0, a1),
12298            _mm512_add_epi32(a2, a3),
12299        ));
12300        // Scalar tail: (w as i8) + 128 ≡ (w as u8) ^ 0x80.
12301        while j < n {
12302            total += ((w[j] ^ 0x80) as i32) * xq[j] as i32;
12303            j += 1;
12304        }
12305        total
12306    }
12307}
12308
12309/// One q4 row via AVX2: nibbles → centered i8 (unpacklo/hi restores the
12310/// writer's flat order, same as the NEON vzip pair), maddubs against
12311/// the pre-quantized activation group, × the group's f16 scale. Pair
12312/// saturation safe: |w|≤8, |x|≤127 → 2·8·127 ≪ 32767. Mirror of
12313/// `dot_q4_row_sdot`.
12314#[cfg(target_arch = "x86_64")]
12315#[target_feature(enable = "avx2")]
12316unsafe fn dot_q4_row_avx2(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
12317    // SAFETY: callers uphold slice-length contracts (16 packed bytes and
12318    // 2 scale bytes per group; xq.len() == gpr·GROUP_SIZE).
12319    unsafe {
12320        use core::arch::x86_64::*;
12321        let lomask = _mm_set1_epi8(0x0F);
12322        let eight = _mm256_set1_epi8(8);
12323        let ones = _mm256_set1_epi16(1);
12324        let mut acc = 0f32;
12325        for gi in 0..gpr {
12326            let g = g0 + gi;
12327            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
12328            let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
12329            let lo = _mm_and_si128(b, lomask);
12330            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
12331            let w = _mm256_sub_epi8(
12332                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
12333                eight,
12334            );
12335            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
12336            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
12337            let d = _mm256_madd_epi16(p16, ones);
12338            let hi128 = _mm256_extracti128_si256::<1>(d);
12339            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
12340            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
12341            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
12342            acc += _mm_cvtsi128_si32(s32) as f32 * s;
12343        }
12344        acc
12345    }
12346}
12347
12348/// Two-activation q4 row via AVX2: nibbles unpacked ONCE per group,
12349/// both activations dotted against the same centered i8 register.
12350#[cfg(target_arch = "x86_64")]
12351#[target_feature(enable = "avx2")]
12352unsafe fn dot_q4_row_avx2_2(
12353    packed: &[u8],
12354    scales: &[u8],
12355    g0: usize,
12356    gpr: usize,
12357    xq1: &[i8],
12358    xq2: &[i8],
12359) -> (f32, f32) {
12360    // SAFETY: callers uphold slice-length contracts (see dot_q4_row_avx2).
12361    unsafe {
12362        use core::arch::x86_64::*;
12363        let lomask = _mm_set1_epi8(0x0F);
12364        let eight = _mm256_set1_epi8(8);
12365        let ones = _mm256_set1_epi16(1);
12366        let (mut acc1, mut acc2) = (0f32, 0f32);
12367        #[inline(always)]
12368        unsafe fn hsum(d: core::arch::x86_64::__m256i) -> i32 {
12369            unsafe {
12370                use core::arch::x86_64::*;
12371                let hi128 = _mm256_extracti128_si256::<1>(d);
12372                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
12373                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
12374                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
12375                _mm_cvtsi128_si32(s32)
12376            }
12377        }
12378        for gi in 0..gpr {
12379            let g = g0 + gi;
12380            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
12381            let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
12382            let lo = _mm_and_si128(b, lomask);
12383            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
12384            let w = _mm256_sub_epi8(
12385                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
12386                eight,
12387            );
12388            let aw = _mm256_abs_epi8(w);
12389            let x1 = _mm256_loadu_si256(xq1.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
12390            let x2 = _mm256_loadu_si256(xq2.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
12391            let d1 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x1, w)), ones);
12392            let d2 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x2, w)), ones);
12393            acc1 += hsum(d1) as f32 * s;
12394            acc2 += hsum(d2) as f32 * s;
12395        }
12396        (acc1, acc2)
12397    }
12398}
12399
12400/// One q8 row range via AVX2 (x86 mirror of `q8_range_sdot`).
12401#[cfg(target_arch = "x86_64")]
12402fn q8_range_avx2(
12403    q: &[u8],
12404    row_scale: &[f32],
12405    act: &SplitAct,
12406    cols: usize,
12407    out_addr: SendMut,
12408    start: usize,
12409    end: usize,
12410) {
12411    for o in start..end {
12412        let v = row_dot_avx2(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
12413        // SAFETY: disjoint row ranges per worker.
12414        unsafe { *out_addr.at(o) = v };
12415    }
12416}
12417
12418/// Two-input q8 row range via AVX2 (x86 mirror of `q8_range2_sdot`).
12419#[cfg(target_arch = "x86_64")]
12420#[allow(clippy::too_many_arguments)]
12421fn q8_range2_avx2(
12422    q: &[u8],
12423    row_scale: &[f32],
12424    a1: &SplitAct,
12425    a2: &SplitAct,
12426    cols: usize,
12427    p1: SendMut,
12428    p2: SendMut,
12429    start: usize,
12430    end: usize,
12431) {
12432    for o in start..end {
12433        let row = &q[o * cols..(o + 1) * cols];
12434        // SAFETY: disjoint row ranges per worker.
12435        unsafe {
12436            *p1.at(o) = row_dot_avx2(row, a1) * row_scale[o];
12437            *p2.at(o) = row_dot_avx2(row, a2) * row_scale[o];
12438        }
12439    }
12440}
12441
12442// ───────────────────── A8W8 SDOT path (port of vmfcore, ×1.78 decode) ─────────────────────
12443
12444/// ARMv8.6 i8mm (smmla): 32 int8 MACs per instruction vs sdot's 16 —
12445/// yet MEASURED 2.4× SLOWER than the blocked sdot on Apple silicon
12446/// (108 vs 264 GF/s): the on-the-fly vcombine packing and the two-
12447/// accumulator dependency chain swamp the MAC advantage, and Apple's
12448/// four SIMD pipes already keep sdot fed. OPT-IN (CMF_I8MM=1) for
12449/// field trials on Cortex-A710/X-class parts with two pipes, where the
12450/// balance may differ; a pre-interleaved weight layout (repack infra)
12451/// is the known path if it ever earns its keep.
12452#[cfg(target_arch = "aarch64")]
12453fn i8mm_enabled() -> bool {
12454    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12455    *ON.get_or_init(|| {
12456        std::env::var("CMF_I8MM").map(|v| v == "1").unwrap_or(false)
12457            && std::arch::is_aarch64_feature_detected!("i8mm")
12458    })
12459}
12460
12461/// SDOT enabled? Default ON when the CPU has ARMv8.2 dotprod;
12462/// `CMF_SDOT=0` disables (falls back to i8×f32 NEON).
12463/// (On non-ARM release builds only the test tolerance switch calls it.)
12464#[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
12465fn sdot_enabled() -> bool {
12466    if FLOAT_ACTIVATIONS.get() {
12467        return false;
12468    }
12469    use std::sync::OnceLock;
12470    static ON: OnceLock<bool> = OnceLock::new();
12471    *ON.get_or_init(|| {
12472        let want = std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true);
12473        if !want {
12474            return false;
12475        }
12476
12477        #[cfg(target_arch = "aarch64")]
12478        {
12479            if std::arch::is_aarch64_feature_detected!("dotprod") {
12480                return true;
12481            }
12482            #[cfg(target_os = "android")]
12483            {
12484                if let Ok(cpuinfo) = std::fs::read_to_string("/proc/cpuinfo") {
12485                    if cpuinfo.lines().any(|l| {
12486                        (l.starts_with("Features") || l.starts_with("features"))
12487                            && l.contains("asimddp")
12488                    }) {
12489                        return true;
12490                    }
12491                }
12492            }
12493            false
12494        }
12495        #[cfg(not(target_arch = "aarch64"))]
12496        {
12497            false
12498        }
12499    })
12500}
12501
12502/// Two-field activation split (≡ vmfcore `q8_split_prep`): outlier
12503/// channels (>8·rms) are computed exactly in f32; the bulk (outliers
12504/// zeroed → clean absmax) goes through int8 SDOT. Computed ONCE per
12505/// matvec, shared by all rows/workers.
12506struct SplitAct {
12507    xq: Vec<i8>,
12508    sx: f32,
12509    outliers: Vec<(usize, f32)>,
12510    /// Σ xq — the VNNI bias-trick correction (`(w+128)·x` sums need
12511    /// `−128·Σx`); one i32 per split, computed once per matvec.
12512    #[cfg_attr(not(target_arch = "x86_64"), allow(dead_code))]
12513    xsum: i32,
12514    /// `Σ xq` per 32-element group, already f32 — the q4tp `(w+8)·x − 8·Σx`
12515    /// correction for the VERTICAL accumulator: the f32x8 partials reduce
12516    /// through an 8-lane horizontal sum, and each lane subtracts this once,
12517    /// so the row total carries exactly `8·Σx`. Precomputing it lets the
12518    /// AVX2 kernels drop the per-byte `vpaddb`/`vpabsb`/`vpsignb` (two of
12519    /// the three sat on port 5) and feed the raw unsigned nibbles straight
12520    /// to `vpmaddubsw`. Every value is an integer ≤ 32·127, so both the
12521    /// f32 conversion and each lane's subtraction are exact and the cell
12522    /// equals the old `(nib − 8)` dot bit for bit.
12523    #[cfg_attr(not(target_arch = "x86_64"), allow(dead_code))]
12524    xg8: Vec<f32>,
12525}
12526
12527thread_local! {
12528    /// Recycled xq buffers: split_act runs for every matvec (~200/token)
12529    /// and its hidden-size allocation was steady-state heap churn.
12530    static XQ_FREE: std::cell::RefCell<Vec<Vec<i8>>> =
12531        const { std::cell::RefCell::new(Vec::new()) };
12532    /// Recycled per-group correction buffers (see `SplitAct::xg8`).
12533    static XG8_FREE: std::cell::RefCell<Vec<Vec<f32>>> =
12534        const { std::cell::RefCell::new(Vec::new()) };
12535}
12536
12537/// Pop a recycled `xg8` buffer, or mint one; filled by [`xg8_of`].
12538fn xg8_take() -> Vec<f32> {
12539    let mut v = XG8_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
12540    v.clear();
12541    v
12542}
12543
12544impl Drop for SplitAct {
12545    fn drop(&mut self) {
12546        let buf = std::mem::take(&mut self.xq);
12547        if buf.capacity() > 0 {
12548            XQ_FREE.with(|f| {
12549                let mut f = f.borrow_mut();
12550                if f.len() < 16 {
12551                    f.push(buf);
12552                }
12553            });
12554        }
12555        let buf = std::mem::take(&mut self.xg8);
12556        if buf.capacity() > 0 {
12557            XG8_FREE.with(|f| {
12558                let mut f = f.borrow_mut();
12559                if f.len() < 16 {
12560                    f.push(buf);
12561                }
12562            });
12563        }
12564    }
12565}
12566
12567/// Fill `out` with `Σ xq` per group — see [`SplitAct::xg8`]. A tail
12568/// shorter than one group is dropped (`chunks_exact`); kernels guard on
12569/// `xg8.len() == gpr`.
12570fn xg8_fill(xq: &[i8], out: &mut Vec<f32>) {
12571    out.clear();
12572    out.extend(xq.chunks_exact(GROUP_SIZE).map(|c| {
12573        let mut s = 0i32;
12574        for &v in c {
12575            s += v as i32;
12576        }
12577        s as f32
12578    }));
12579}
12580
12581thread_local! {
12582    /// One scratch row per WORKER, kept for the life of the thread.
12583    ///
12584    /// The kernels take a row of group scales per dispatch, and a fresh
12585    /// `vec![0f32; gpr]` inside the closure is one allocation per worker per
12586    /// dispatch — on the release checkpoint about six thousand a token, a
12587    /// quarter of everything the benchmark counts.
12588    static KROW: UnsafeCell<Vec<f32>> = const { UnsafeCell::new(Vec::new()) };
12589    /// Two scratch rows for a gate/up pair.  Q2TP and Q4TP fused SwiGLU
12590    /// decode each projection's ladder independently, but both ladders have
12591    /// the same lifetime; retaining them per worker removes two heap trips
12592    /// from every dispatched expert pair.
12593    static KROWS: UnsafeCell<[Vec<f32>; 2]> =
12594        const { UnsafeCell::new([Vec::new(), Vec::new()]) };
12595}
12596
12597/// Borrow `n` floats of the calling worker's scratch. Nothing inside a
12598/// kernel body borrows it again, which is what keeps the RefCell honest.
12599#[inline]
12600fn with_krow<R>(n: usize, f: impl FnOnce(&mut [f32]) -> R) -> R {
12601    KROW.with(|s| {
12602        // SAFETY: `KROW` is thread-local and this function does not recurse;
12603        // each worker has exclusive access to its own scratch row.
12604        let b = unsafe { &mut *s.get() };
12605        if b.len() < n {
12606            b.resize(n, 0.0);
12607        }
12608        f(&mut b[..n])
12609    })
12610}
12611
12612/// `t.round().clamp(-127.0, 127.0) as i8`, bit for bit, without the libm
12613/// call. On baseline x86-64 (no SSE4.1 `roundps`) `f32::round` is a
12614/// function call per element, and split_act runs it over every hidden
12615/// state before every matvec: measured 27 us a call on a 2048-wide
12616/// activation on an EPYC 7763 — 5.4 ms of a 55 ms decode token, all of
12617/// it on the caller's thread while thirty workers wait. Clamping first is
12618/// equivalent (round is monotonic and ±127 are integers), and after the
12619/// clamp `t - trunc(t)` is exact, so the half-away-from-zero decision is
12620/// the one `round` makes. NaN clamps to NaN and converts to 0, as before.
12621/// The loop vectorizes (cvttps2dq + compare/select).
12622#[inline(always)]
12623fn q8_round(t: f32) -> i8 {
12624    let t = t.clamp(-127.0, 127.0);
12625    let i = t as i32;
12626    let f = t - i as f32;
12627    let r = if f >= 0.5 {
12628        i + 1
12629    } else if f <= -0.5 {
12630        i - 1
12631    } else {
12632        i
12633    };
12634    r as i8
12635}
12636
12637#[inline]
12638fn with_krows<R>(n: usize, f: impl FnOnce(&mut [f32], &mut [f32]) -> R) -> R {
12639    KROWS.with(|s| {
12640        // SAFETY: KROWS is thread-local and no kernel body recursively calls
12641        // this helper; each participant owns its two vectors exclusively.
12642        let b = unsafe { &mut *s.get() };
12643        for row in b.iter_mut() {
12644            if row.len() < n {
12645                row.resize(n, 0.0);
12646            }
12647        }
12648        let (a, b) = b.split_at_mut(1);
12649        f(&mut a[0][..n], &mut b[0][..n])
12650    })
12651}
12652
12653#[inline]
12654fn silu_mul_limited(mut gate: f32, mut up: f32, limit: f32) -> f32 {
12655    if limit > 0.0 {
12656        up = up.clamp(-limit, limit);
12657        gate = gate.min(limit);
12658    }
12659    gate / (1.0 + (-gate).exp()) * up
12660}
12661
12662fn split_act(x: &[f32]) -> SplitAct {
12663    let _prof = crate::cpuprof::time(crate::cpuprof::Slot::SplitAct);
12664    let n = x.len();
12665    let rms = (x.iter().map(|&v| (v * v) as f64).sum::<f64>() / n.max(1) as f64).sqrt() as f32;
12666    let thr = 8.0 * rms;
12667    // One pass: collect outliers and the bulk absmax (outliers excluded —
12668    // identical to the old zero-then-fold over a copied buffer, minus the
12669    // full-vector copy).
12670    let mut outliers: Vec<(usize, f32)> = Vec::new();
12671    let mut amax = 0f32;
12672    for (j, &v) in x.iter().enumerate() {
12673        let a = v.abs();
12674        if a > thr {
12675            outliers.push((j, v));
12676        } else if a > amax {
12677            amax = a;
12678        }
12679    }
12680    let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
12681    let inv = 1.0 / sx;
12682    crate::cpuprof::count_outliers(outliers.len() as u64);
12683    let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
12684    xq.clear();
12685    xq.reserve(n);
12686    if outliers.is_empty() {
12687        xq.extend(
12688            x.iter()
12689                .map(|&v| q8_round(v * inv)),
12690        );
12691    } else {
12692        // Outlier slots quantize to 0 (their exact term is added later).
12693        xq.extend(x.iter().map(|&v| {
12694            if v.abs() > thr {
12695                0
12696            } else {
12697                q8_round(v * inv)
12698            }
12699        }));
12700    }
12701    let xsum = xq.iter().map(|&v| v as i32).sum();
12702    let mut xg8 = xg8_take();
12703    xg8_fill(&xq, &mut xg8);
12704    SplitAct {
12705        xq,
12706        sx,
12707        outliers,
12708        xsum,
12709        xg8,
12710    }
12711}
12712
12713fn split_act_q8_2f(x: &[f32], col: &[f32]) -> SplitAct {
12714    let _prof = crate::cpuprof::time(crate::cpuprof::Slot::SplitAct);
12715    let n = x.len();
12716    let rms = (x
12717        .iter()
12718        .zip(col)
12719        .map(|(&a, &c)| {
12720            let v = a * c;
12721            (v * v) as f64
12722        })
12723        .sum::<f64>()
12724        / n.max(1) as f64)
12725        .sqrt() as f32;
12726    let thr = 8.0 * rms;
12727
12728    let mut outliers = Vec::new();
12729    let mut amax = 0f32;
12730    for (j, (&a, &c)) in x.iter().zip(col).enumerate() {
12731        let v = a * c;
12732        let s = v.abs();
12733        if s > thr {
12734            outliers.push((j, v));
12735        } else if s > amax {
12736            amax = s;
12737        }
12738    }
12739
12740    let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
12741    let inv = 1.0 / sx;
12742    let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
12743    xq.clear();
12744    xq.reserve(n);
12745    if outliers.is_empty() {
12746        xq.extend(
12747            x.iter()
12748                .zip(col)
12749                .map(|(&a, &c)| q8_round((a * c) * inv)),
12750        );
12751    } else {
12752        xq.extend(x.iter().zip(col).map(|(&a, &c)| {
12753            let v = a * c;
12754            if v.abs() > thr {
12755                0
12756            } else {
12757                q8_round(v * inv)
12758            }
12759        }));
12760    }
12761    let xsum = xq.iter().map(|&v| v as i32).sum();
12762    let mut xg8 = xg8_take();
12763    xg8_fill(&xq, &mut xg8);
12764    SplitAct {
12765        xq,
12766        sx,
12767        outliers,
12768        xsum,
12769        xg8,
12770    }
12771}
12772
12773/// int8(weight)·int8(activation) → i32 via `sdot` (inline asm — the
12774/// vdotq intrinsic is unstable; port of vmfcore `dot_i8_sdot`).
12775#[cfg(target_arch = "aarch64")]
12776#[target_feature(enable = "neon,dotprod")]
12777unsafe fn dot_i8_sdot(w: &[u8], xq: &[i8]) -> i32 {
12778    // SAFETY: callers uphold slice-length contracts (see call sites).
12779    unsafe {
12780        use core::arch::aarch64::*;
12781        use core::arch::asm;
12782        let wp = w.as_ptr() as *const i8;
12783        let n = w.len();
12784        let (mut a0, mut a1, mut a2, mut a3) = (
12785            vdupq_n_s32(0),
12786            vdupq_n_s32(0),
12787            vdupq_n_s32(0),
12788            vdupq_n_s32(0),
12789        );
12790        let mut i = 0;
12791        while i + 64 <= n {
12792            let (w0, x0) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
12793            let (w1, x1) = (vld1q_s8(wp.add(i + 16)), vld1q_s8(xq.as_ptr().add(i + 16)));
12794            let (w2, x2) = (vld1q_s8(wp.add(i + 32)), vld1q_s8(xq.as_ptr().add(i + 32)));
12795            let (w3, x3) = (vld1q_s8(wp.add(i + 48)), vld1q_s8(xq.as_ptr().add(i + 48)));
12796            asm!(
12797                "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
12798                "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
12799                "sdot {a2:v}.4s, {w2:v}.16b, {x2:v}.16b",
12800                "sdot {a3:v}.4s, {w3:v}.16b, {x3:v}.16b",
12801                a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
12802                w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
12803                w2 = in(vreg) w2, x2 = in(vreg) x2, w3 = in(vreg) w3, x3 = in(vreg) x3,
12804                options(pure, nomem, nostack),
12805            );
12806            i += 64;
12807        }
12808        while i + 16 <= n {
12809            let (wv, xv) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
12810            asm!("sdot {a:v}.4s, {w:v}.16b, {x:v}.16b",
12811                 a = inout(vreg) a0, w = in(vreg) wv, x = in(vreg) xv, options(pure, nomem, nostack));
12812            i += 16;
12813        }
12814        let mut s = vaddvq_s32(vaddq_s32(vaddq_s32(a0, a1), vaddq_s32(a2, a3)));
12815        while i < n {
12816            s += (*wp.add(i)) as i32 * xq[i] as i32;
12817            i += 1;
12818        }
12819        s
12820    }
12821}
12822
12823/// Row-blocked SDOT: 4 output rows per pass — the activation chunk is
12824/// loaded once and reused, 4 independent accumulators hide sdot latency
12825/// (port of vmfcore `dot_i8_sdot_4rows`).
12826#[cfg(target_arch = "aarch64")]
12827#[target_feature(enable = "neon,dotprod")]
12828unsafe fn dot_i8_sdot_4rows(w0: &[u8], w1: &[u8], w2: &[u8], w3: &[u8], xq: &[i8]) -> [i32; 4] {
12829    // SAFETY: callers uphold slice-length contracts (see call sites).
12830    unsafe {
12831        use core::arch::aarch64::*;
12832        use core::arch::asm;
12833        let n = xq.len();
12834        let px = xq.as_ptr();
12835        let (p0, p1, p2, p3) = (
12836            w0.as_ptr() as *const i8,
12837            w1.as_ptr() as *const i8,
12838            w2.as_ptr() as *const i8,
12839            w3.as_ptr() as *const i8,
12840        );
12841        let (mut a0, mut a1, mut a2, mut a3) = (
12842            vdupq_n_s32(0),
12843            vdupq_n_s32(0),
12844            vdupq_n_s32(0),
12845            vdupq_n_s32(0),
12846        );
12847        let mut i = 0;
12848        while i + 16 <= n {
12849            let x = vld1q_s8(px.add(i));
12850            let v0 = vld1q_s8(p0.add(i));
12851            let v1 = vld1q_s8(p1.add(i));
12852            let v2 = vld1q_s8(p2.add(i));
12853            let v3 = vld1q_s8(p3.add(i));
12854            asm!(
12855                "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
12856                "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
12857                "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
12858                "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
12859                a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
12860                v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
12861                options(pure, nomem, nostack),
12862            );
12863            i += 16;
12864        }
12865        let mut r = [
12866            vaddvq_s32(a0),
12867            vaddvq_s32(a1),
12868            vaddvq_s32(a2),
12869            vaddvq_s32(a3),
12870        ];
12871        while i < n {
12872            let xi = *px.add(i) as i32;
12873            r[0] += (*p0.add(i)) as i32 * xi;
12874            r[1] += (*p1.add(i)) as i32 * xi;
12875            r[2] += (*p2.add(i)) as i32 * xi;
12876            r[3] += (*p3.add(i)) as i32 * xi;
12877            i += 1;
12878        }
12879        r
12880    }
12881}
12882
12883/// 4 interleaved rows in one pass: the repacked group is [r0[c], r1[c],
12884/// r2[c], r3[c]] per 16-byte chunk, so each iteration reads ONE 64-byte
12885/// line plus the shared activation chunk — a single sequential weight
12886/// stream per worker. Per-row accumulation is the same one-accumulator
12887/// scheme as `dot_i8_sdot_4rows`; integer sums are exact, so outputs
12888/// are bit-identical to the mmap-layout kernel.
12889#[cfg(target_arch = "aarch64")]
12890#[target_feature(enable = "neon,dotprod")]
12891unsafe fn dot_i8_sdot_4rows_il(g: &[u8], xq: &[i8]) -> [i32; 4] {
12892    // SAFETY: callers uphold slice-length contracts (g.len() == 4·n,
12893    // n % 16 == 0 — guaranteed by the repack gate).
12894    unsafe {
12895        use core::arch::aarch64::*;
12896        use core::arch::asm;
12897        let n = xq.len();
12898        let px = xq.as_ptr();
12899        let pg = g.as_ptr() as *const i8;
12900        let (mut a0, mut a1, mut a2, mut a3) = (
12901            vdupq_n_s32(0),
12902            vdupq_n_s32(0),
12903            vdupq_n_s32(0),
12904            vdupq_n_s32(0),
12905        );
12906        let mut i = 0;
12907        while i + 16 <= n {
12908            let x = vld1q_s8(px.add(i));
12909            let base = pg.add(4 * i);
12910            let v0 = vld1q_s8(base);
12911            let v1 = vld1q_s8(base.add(16));
12912            let v2 = vld1q_s8(base.add(32));
12913            let v3 = vld1q_s8(base.add(48));
12914            asm!(
12915                "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
12916                "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
12917                "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
12918                "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
12919                a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
12920                v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
12921                options(pure, nomem, nostack),
12922            );
12923            i += 16;
12924        }
12925        [
12926            vaddvq_s32(a0),
12927            vaddvq_s32(a1),
12928            vaddvq_s32(a2),
12929            vaddvq_s32(a3),
12930        ]
12931    }
12932}
12933
12934/// One q8 row range via SDOT (4-row blocks + tail) — the body of
12935/// `qmatvec`'s hot loop, extracted so multi-matrix jobs can drive the
12936/// SAME kernel for several tensors under one pool dispatch. `rep` — the
12937/// load-time interleaved repack (empty = mmap layout only); rows outside
12938/// full 4-row groups always come from the mmap layout.
12939#[cfg(target_arch = "aarch64")]
12940fn q8_range_sdot(
12941    q: &[u8],
12942    rep: &[u8],
12943    row_scale: &[f32],
12944    act: &SplitAct,
12945    cols: usize,
12946    out_addr: SendMut,
12947    start: usize,
12948    end: usize,
12949) {
12950    let mut o = start;
12951    // Leading rows to the group boundary (repack path only): the pool
12952    // splits row ranges arbitrarily, groups are absolute.
12953    if !rep.is_empty() {
12954        while o < end && o % 4 != 0 {
12955            let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
12956            unsafe { *out_addr.at(o) = v };
12957            o += 1;
12958        }
12959    }
12960    while o + 4 <= end {
12961        let r = if rep.is_empty() {
12962            unsafe {
12963                dot_i8_sdot_4rows(
12964                    &q[o * cols..(o + 1) * cols],
12965                    &q[(o + 1) * cols..(o + 2) * cols],
12966                    &q[(o + 2) * cols..(o + 3) * cols],
12967                    &q[(o + 3) * cols..(o + 4) * cols],
12968                    &act.xq,
12969                )
12970            }
12971        } else {
12972            unsafe { dot_i8_sdot_4rows_il(&rep[o * cols..(o + 4) * cols], &act.xq) }
12973        };
12974        for k in 0..4 {
12975            let mut acc = r[k] as f32 * act.sx;
12976            for &(j, xv) in &act.outliers {
12977                acc += (q[(o + k) * cols + j] as i8) as f32 * xv;
12978            }
12979            // SAFETY: disjoint row ranges per worker.
12980            unsafe { *out_addr.at(o + k) = acc * row_scale[o + k] };
12981        }
12982        o += 4;
12983    }
12984    while o < end {
12985        let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
12986        unsafe { *out_addr.at(o) = v };
12987        o += 1;
12988    }
12989}
12990
12991/// Two-input q8 row range via SDOT — `qmatvec2`'s hot loop, extracted
12992/// for the fused pair multi-matrix job (`matvec2_many`).
12993#[cfg(target_arch = "aarch64")]
12994#[allow(clippy::too_many_arguments)]
12995fn q8_range2_sdot(
12996    q: &[u8],
12997    row_scale: &[f32],
12998    a1: &SplitAct,
12999    a2: &SplitAct,
13000    cols: usize,
13001    p1: SendMut,
13002    p2: SendMut,
13003    start: usize,
13004    end: usize,
13005) {
13006    for o in start..end {
13007        let row = &q[o * cols..(o + 1) * cols];
13008        // SAFETY: disjoint row ranges per worker.
13009        unsafe {
13010            *p1.at(o) = row_dot_sdot(row, a1) * row_scale[o];
13011            *p2.at(o) = row_dot_sdot(row, a2) * row_scale[o];
13012        }
13013    }
13014}
13015
13016/// Two-input q8 row range, f32 kernel (non-SDOT) — same extraction.
13017#[allow(clippy::too_many_arguments)]
13018fn q8_range2_f32(
13019    q: &[u8],
13020    row_scale: &[f32],
13021    x1: &[f32],
13022    x2: &[f32],
13023    cols: usize,
13024    p1: SendMut,
13025    p2: SendMut,
13026    start: usize,
13027    end: usize,
13028) {
13029    for o in start..end {
13030        let row = &q[o * cols..(o + 1) * cols];
13031        // SAFETY: disjoint row ranges per worker.
13032        unsafe {
13033            *p1.at(o) = dot_i8_f32(row, x1) * row_scale[o];
13034            *p2.at(o) = dot_i8_f32(row, x2) * row_scale[o];
13035        }
13036    }
13037}
13038
13039/// Scalar/NEON-f32 q8 row range (non-SDOT platforms) — same extraction.
13040fn q8_range_f32(
13041    q: &[u8],
13042    row_scale: &[f32],
13043    xs: &[f32],
13044    cols: usize,
13045    out_addr: SendMut,
13046    start: usize,
13047    end: usize,
13048) {
13049    for o in start..end {
13050        let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
13051        // SAFETY: disjoint row ranges per worker.
13052        unsafe { *out_addr.at(o) = v };
13053    }
13054}
13055
13056/// One q8 row against a split activation, portable: the per-arch fast
13057/// dots where they exist, the exact scalar loop elsewhere. The scalar
13058/// arm is also the test oracle for both fast arms.
13059#[inline]
13060fn q8_row_dot(row: &[u8], act: &SplitAct) -> f32 {
13061    #[cfg(target_arch = "aarch64")]
13062    return row_dot_sdot(row, act);
13063    #[cfg(target_arch = "x86_64")]
13064    return row_dot_avx2(row, act);
13065    #[allow(unreachable_code)]
13066    q8_row_dot_scalar(row, act)
13067}
13068
13069#[allow(dead_code)]
13070fn q8_row_dot_scalar(row: &[u8], act: &SplitAct) -> f32 {
13071    let mut acc = 0i32;
13072    for (k, &b) in row.iter().enumerate() {
13073        acc += (b as i8) as i32 * act.xq[k] as i32;
13074    }
13075    let mut acc = acc as f32 * act.sx;
13076    for &(j, xv) in &act.outliers {
13077        acc += (row[j] as i8) as f32 * xv;
13078    }
13079    acc
13080}
13081
13082/// SDOT row dot with exact outlier correction:
13083/// `dot = sdot(w, xq)·sx + Σ_outl w[j]·x[j]` (then × row_scale by caller).
13084#[cfg(target_arch = "aarch64")]
13085#[inline]
13086fn row_dot_sdot(row: &[u8], act: &SplitAct) -> f32 {
13087    let mut acc = unsafe { dot_i8_sdot(row, &act.xq) } as f32 * act.sx;
13088    for &(j, xv) in &act.outliers {
13089        acc += (row[j] as i8) as f32 * xv;
13090    }
13091    acc
13092}
13093
13094/// One q4 row via SDOT: each 32-group's nibbles unpack to centered i8
13095/// (nib−8 ∈ [−8,7]), int8×int8 `sdot` against the pre-quantized
13096/// activation group, × the group's f16 scale. Returns Σ_g dot_g·s_g;
13097/// the caller multiplies by the activation scale and adds the exact
13098/// outlier terms (port of vmfcore `dot_q4_block_sdot`, +23% measured).
13099/// Nibble order matches the writer: element 2k = low nibble, 2k+1 = high
13100/// → zip(lo,hi) restores flat order.
13101#[cfg(target_arch = "aarch64")]
13102#[target_feature(enable = "neon,dotprod")]
13103unsafe fn dot_q4_row_sdot(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
13104    // SAFETY: callers uphold slice-length contracts (16 packed bytes and
13105    // 2 scale bytes per group; xq.len() == gpr·GROUP_SIZE).
13106    unsafe {
13107        use core::arch::aarch64::*;
13108        use core::arch::asm;
13109        let lomask = vdupq_n_u8(0x0F);
13110        let eight = vdupq_n_s8(8);
13111        let mut acc = 0f32;
13112        for gi in 0..gpr {
13113            let g = g0 + gi;
13114            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
13115            let b = vld1q_u8(packed.as_ptr().add(g * 16));
13116            let lo = vandq_u8(b, lomask);
13117            let hi = vshrq_n_u8::<4>(b);
13118            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
13119            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
13120            let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
13121            let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
13122            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
13123            asm!(
13124                "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
13125                "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
13126                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
13127                e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
13128                options(pure, nomem, nostack),
13129            );
13130            acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
13131        }
13132        acc
13133    }
13134}
13135
13136/// Two-activation q4 row via SDOT: the nibble unpack (the expensive
13137/// part) happens ONCE per group; both pre-quantized activations are
13138/// dotted against the same centered i8 registers. Per-lane math matches
13139/// `dot_q4_row_sdot` exactly.
13140#[cfg(target_arch = "aarch64")]
13141#[target_feature(enable = "neon,dotprod")]
13142unsafe fn dot_q4_row_sdot2(
13143    packed: &[u8],
13144    scales: &[u8],
13145    g0: usize,
13146    gpr: usize,
13147    xq1: &[i8],
13148    xq2: &[i8],
13149) -> (f32, f32) {
13150    // SAFETY: callers uphold slice-length contracts (16 packed bytes and
13151    // 2 scale bytes per group; xq*.len() == gpr·GROUP_SIZE).
13152    unsafe {
13153        use core::arch::aarch64::*;
13154        use core::arch::asm;
13155        let lomask = vdupq_n_u8(0x0F);
13156        let eight = vdupq_n_s8(8);
13157        let (mut acc1, mut acc2) = (0f32, 0f32);
13158        for gi in 0..gpr {
13159            let g = g0 + gi;
13160            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
13161            let b = vld1q_u8(packed.as_ptr().add(g * 16));
13162            let lo = vandq_u8(b, lomask);
13163            let hi = vshrq_n_u8::<4>(b);
13164            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
13165            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
13166            let x10 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE));
13167            let x11 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE + 16));
13168            let x20 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE));
13169            let x21 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE + 16));
13170            let (mut a0, mut a1, mut b0, mut b1) = (
13171                vdupq_n_s32(0),
13172                vdupq_n_s32(0),
13173                vdupq_n_s32(0),
13174                vdupq_n_s32(0),
13175            );
13176            asm!(
13177                "sdot {a0:v}.4s, {e0:v}.16b, {x10:v}.16b",
13178                "sdot {a1:v}.4s, {e1:v}.16b, {x11:v}.16b",
13179                "sdot {b0:v}.4s, {e0:v}.16b, {x20:v}.16b",
13180                "sdot {b1:v}.4s, {e1:v}.16b, {x21:v}.16b",
13181                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
13182                b0 = inout(vreg) b0, b1 = inout(vreg) b1,
13183                e0 = in(vreg) e0, e1 = in(vreg) e1,
13184                x10 = in(vreg) x10, x11 = in(vreg) x11,
13185                x20 = in(vreg) x20, x21 = in(vreg) x21,
13186                options(pure, nomem, nostack),
13187            );
13188            acc1 += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
13189            acc2 += vaddvq_s32(vaddq_s32(b0, b1)) as f32 * s;
13190        }
13191        (acc1, acc2)
13192    }
13193}
13194
13195// ───────────────────── fused int8 kernels ─────────────────────
13196
13197/// `acc += w · row` where the row is centered i8 — NEON widen+fma on
13198/// aarch64, scalar elsewhere. The KV-cache q8 value path rides on this.
13199#[inline]
13200pub(crate) fn axpy_i8_f32(acc: &mut [f32], row: &[i8], w: f32) {
13201    #[cfg(target_arch = "aarch64")]
13202    unsafe {
13203        return axpy_i8_f32_neon(acc, row, w);
13204    }
13205    #[cfg(target_arch = "x86_64")]
13206    if avx2_enabled() {
13207        return unsafe { axpy_i8_f32_avx2(acc, row, w) };
13208    }
13209    #[allow(unreachable_code)]
13210    {
13211        for (a, &b) in acc.iter_mut().zip(row) {
13212            *a += w * b as f32;
13213        }
13214    }
13215}
13216
13217/// i8→f32 axpy via AVX2/FMA (x86 mirror of `axpy_i8_f32_neon`).
13218#[cfg(target_arch = "x86_64")]
13219#[target_feature(enable = "avx2,fma")]
13220unsafe fn axpy_i8_f32_avx2(acc: &mut [f32], row: &[i8], w: f32) {
13221    // SAFETY: callers uphold slice-length contracts (see call sites).
13222    unsafe {
13223        use core::arch::x86_64::*;
13224        let n = acc.len().min(row.len());
13225        let ap = acc.as_mut_ptr();
13226        let rp = row.as_ptr();
13227        let wv = _mm256_set1_ps(w);
13228        let mut j = 0usize;
13229        while j + 16 <= n {
13230            let rb = _mm_loadu_si128(rp.add(j) as *const __m128i);
13231            let lo = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(rb));
13232            let hi = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(rb)));
13233            let v0 = _mm256_fmadd_ps(wv, lo, _mm256_loadu_ps(ap.add(j)));
13234            let v1 = _mm256_fmadd_ps(wv, hi, _mm256_loadu_ps(ap.add(j + 8)));
13235            _mm256_storeu_ps(ap.add(j), v0);
13236            _mm256_storeu_ps(ap.add(j + 8), v1);
13237            j += 16;
13238        }
13239        while j < n {
13240            *ap.add(j) += w * (*rp.add(j)) as f32;
13241            j += 1;
13242        }
13243    }
13244}
13245
13246#[cfg(target_arch = "aarch64")]
13247#[target_feature(enable = "neon")]
13248unsafe fn axpy_i8_f32_neon(acc: &mut [f32], row: &[i8], w: f32) {
13249    // SAFETY: callers uphold slice-length contracts (see call sites).
13250    unsafe {
13251        use core::arch::aarch64::*;
13252        let n = acc.len().min(row.len());
13253        let ap = acc.as_mut_ptr();
13254        let rp = row.as_ptr();
13255        let wv = vdupq_n_f32(w);
13256        let mut j = 0usize;
13257        while j + 16 <= n {
13258            let rb = vld1q_s8(rp.add(j));
13259            let lo = vmovl_s8(vget_low_s8(rb));
13260            let hi = vmovl_s8(vget_high_s8(rb));
13261            for (off, half) in [(0, lo), (8, hi)] {
13262                let f0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(half)));
13263                let f1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(half)));
13264                let o = j + off;
13265                vst1q_f32(ap.add(o), vfmaq_f32(vld1q_f32(ap.add(o)), wv, f0));
13266                vst1q_f32(ap.add(o + 4), vfmaq_f32(vld1q_f32(ap.add(o + 4)), wv, f1));
13267            }
13268            j += 16;
13269        }
13270        while j < n {
13271            *ap.add(j) += w * (*rp.add(j)) as f32;
13272            j += 1;
13273        }
13274    }
13275}
13276
13277/// i8 row · f32 x. NEON on aarch64 (ported from vmfcore `dot_i8_f32_neon`,
13278/// ≈9× scalar), scalar elsewhere.
13279#[inline]
13280pub(crate) fn dot_i8_f32(w: &[u8], x: &[f32]) -> f32 {
13281    #[cfg(target_arch = "aarch64")]
13282    unsafe {
13283        return dot_i8_f32_neon(w, x);
13284    }
13285    #[cfg(target_arch = "x86_64")]
13286    if avx2_enabled() {
13287        return unsafe { dot_i8_f32_avx2(w, x) };
13288    }
13289    #[allow(unreachable_code)]
13290    {
13291        let mut sum = 0.0f32;
13292        for (j, &b) in w.iter().enumerate() {
13293            sum += (b as i8) as f32 * x[j];
13294        }
13295        sum
13296    }
13297}
13298
13299/// i8 row · (x ⊙ col_field) — the q8_2f row dot with the θ col-field
13300/// folded into the product (no prescaled copy of x). NEON on aarch64,
13301/// scalar elsewhere. Used by the active-neuron path `row_dot`.
13302#[inline]
13303fn dot_i8_col_f32(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
13304    #[cfg(target_arch = "aarch64")]
13305    unsafe {
13306        return dot_i8_col_f32_neon(w, x, col);
13307    }
13308    #[allow(unreachable_code)]
13309    {
13310        let mut sum = 0.0f32;
13311        for (j, &b) in w.iter().enumerate() {
13312            sum += (b as i8) as f32 * x[j] * col[j];
13313        }
13314        sum
13315    }
13316}
13317
13318#[cfg(target_arch = "aarch64")]
13319#[target_feature(enable = "neon")]
13320unsafe fn dot_i8_col_f32_neon(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
13321    // SAFETY: callers uphold slice-length contracts (see call sites).
13322    unsafe {
13323        use core::arch::aarch64::*;
13324        let n = x.len();
13325        let wp = w.as_ptr() as *const i8;
13326        let xp = x.as_ptr();
13327        let cp = col.as_ptr();
13328        let (mut a0, mut a1, mut a2, mut a3) = (
13329            vdupq_n_f32(0.0),
13330            vdupq_n_f32(0.0),
13331            vdupq_n_f32(0.0),
13332            vdupq_n_f32(0.0),
13333        );
13334        let mut j = 0usize;
13335        while j + 16 <= n {
13336            let wb = vld1q_s8(wp.add(j));
13337            let lo = vmovl_s8(vget_low_s8(wb));
13338            let hi = vmovl_s8(vget_high_s8(wb));
13339            let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
13340            let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
13341            let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
13342            let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
13343            a0 = vfmaq_f32(
13344                a0,
13345                w0,
13346                vmulq_f32(vld1q_f32(xp.add(j)), vld1q_f32(cp.add(j))),
13347            );
13348            a1 = vfmaq_f32(
13349                a1,
13350                w1,
13351                vmulq_f32(vld1q_f32(xp.add(j + 4)), vld1q_f32(cp.add(j + 4))),
13352            );
13353            a2 = vfmaq_f32(
13354                a2,
13355                w2,
13356                vmulq_f32(vld1q_f32(xp.add(j + 8)), vld1q_f32(cp.add(j + 8))),
13357            );
13358            a3 = vfmaq_f32(
13359                a3,
13360                w3,
13361                vmulq_f32(vld1q_f32(xp.add(j + 12)), vld1q_f32(cp.add(j + 12))),
13362            );
13363            j += 16;
13364        }
13365        let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
13366        while j < n {
13367            sum += (*wp.add(j)) as f32 * *xp.add(j) * *cp.add(j);
13368            j += 1;
13369        }
13370        sum
13371    }
13372}
13373
13374#[cfg(target_arch = "aarch64")]
13375#[target_feature(enable = "neon")]
13376unsafe fn dot_i8_f32_neon(w: &[u8], x: &[f32]) -> f32 {
13377    // SAFETY: callers uphold slice-length contracts (see call sites).
13378    unsafe {
13379        use core::arch::aarch64::*;
13380        let n = x.len();
13381        let wp = w.as_ptr() as *const i8;
13382        let xp = x.as_ptr();
13383        let (mut a0, mut a1, mut a2, mut a3) = (
13384            vdupq_n_f32(0.0),
13385            vdupq_n_f32(0.0),
13386            vdupq_n_f32(0.0),
13387            vdupq_n_f32(0.0),
13388        );
13389        let mut j = 0usize;
13390        while j + 16 <= n {
13391            let wb = vld1q_s8(wp.add(j));
13392            let lo = vmovl_s8(vget_low_s8(wb));
13393            let hi = vmovl_s8(vget_high_s8(wb));
13394            let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
13395            let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
13396            let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
13397            let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
13398            a0 = vfmaq_f32(a0, w0, vld1q_f32(xp.add(j)));
13399            a1 = vfmaq_f32(a1, w1, vld1q_f32(xp.add(j + 4)));
13400            a2 = vfmaq_f32(a2, w2, vld1q_f32(xp.add(j + 8)));
13401            a3 = vfmaq_f32(a3, w3, vld1q_f32(xp.add(j + 12)));
13402            j += 16;
13403        }
13404        let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
13405        while j < n {
13406            sum += (*wp.add(j)) as f32 * *xp.add(j);
13407            j += 1;
13408        }
13409        sum
13410    }
13411}
13412
13413#[allow(clippy::too_many_arguments)]
13414fn qmatvec(
13415    q: &[u8],
13416    rep: &[u8],
13417    row_scale: &[f32],
13418    x: &[f32],
13419    col_field: &[f32],
13420    dtype: TensorDtype,
13421    rows: usize,
13422    cols: usize,
13423    out: &mut [f32],
13424    pool: Option<&Pool>,
13425) {
13426    debug_assert_eq!(out.len(), rows);
13427    #[cfg(not(target_arch = "aarch64"))]
13428    let _ = rep;
13429
13430    #[cfg(target_arch = "aarch64")]
13431    if sdot_enabled() {
13432        let act = if dtype == TensorDtype::Q8_2f {
13433            split_act_q8_2f(x, col_field)
13434        } else {
13435            split_act(x)
13436        };
13437        let out_addr = SendMut(out.as_mut_ptr());
13438        let run_range = |start: usize, end: usize| {
13439            q8_range_sdot(q, rep, row_scale, &act, cols, out_addr, start, end)
13440        };
13441        match pool {
13442            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
13443            _ => run_range(0, rows),
13444        }
13445        return;
13446    }
13447    // x86 A8W8 via AVX2 maddubs — same quantized-activation contract as
13448    // the SDOT path (CMF_AVX2=0 keeps the exact i8×f32 loop).
13449    #[cfg(target_arch = "x86_64")]
13450    if avx2_a8w8_enabled() {
13451        let act = if dtype == TensorDtype::Q8_2f {
13452            split_act_q8_2f(x, col_field)
13453        } else {
13454            split_act(x)
13455        };
13456        let out_addr = SendMut(out.as_mut_ptr());
13457        let run_range = |start: usize, end: usize| {
13458            q8_range_avx2(q, row_scale, &act, cols, out_addr, start, end)
13459        };
13460        match pool {
13461            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
13462            _ => run_range(0, rows),
13463        }
13464        return;
13465    }
13466
13467    prescale_with(x, col_field, dtype, 1, |xs| {
13468        let out_addr = SendMut(out.as_mut_ptr());
13469        let run_range = move |start: usize, end: usize| {
13470            for o in start..end {
13471                let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
13472                // SAFETY: disjoint row ranges per worker.
13473                unsafe { *out_addr.at(o) = v };
13474            }
13475        };
13476        match pool {
13477            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
13478            _ => run_range(0, rows),
13479        }
13480    });
13481}
13482
13483#[allow(clippy::too_many_arguments)]
13484fn qmatvec2(
13485    q: &[u8],
13486    row_scale: &[f32],
13487    x1: &[f32],
13488    x2: &[f32],
13489    col_field: &[f32],
13490    dtype: TensorDtype,
13491    rows: usize,
13492    cols: usize,
13493    o1: &mut [f32],
13494    o2: &mut [f32],
13495    pool: Option<&Pool>,
13496) {
13497    #[cfg(target_arch = "aarch64")]
13498    if sdot_enabled() {
13499        let a1s = if dtype == TensorDtype::Q8_2f {
13500            split_act_q8_2f(x1, col_field)
13501        } else {
13502            split_act(x1)
13503        };
13504        let a2s = if dtype == TensorDtype::Q8_2f {
13505            split_act_q8_2f(x2, col_field)
13506        } else {
13507            split_act(x2)
13508        };
13509        let p1 = SendMut(o1.as_mut_ptr());
13510        let p2 = SendMut(o2.as_mut_ptr());
13511        let run_range = |start: usize, end: usize| {
13512            q8_range2_sdot(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
13513        };
13514        match pool {
13515            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
13516            _ => run_range(0, rows),
13517        }
13518        return;
13519    }
13520    #[cfg(target_arch = "x86_64")]
13521    if avx2_a8w8_enabled() {
13522        let a1s = if dtype == TensorDtype::Q8_2f {
13523            split_act_q8_2f(x1, col_field)
13524        } else {
13525            split_act(x1)
13526        };
13527        let a2s = if dtype == TensorDtype::Q8_2f {
13528            split_act_q8_2f(x2, col_field)
13529        } else {
13530            split_act(x2)
13531        };
13532        let p1 = SendMut(o1.as_mut_ptr());
13533        let p2 = SendMut(o2.as_mut_ptr());
13534        let run_range = |start: usize, end: usize| {
13535            q8_range2_avx2(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
13536        };
13537        match pool {
13538            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
13539            _ => run_range(0, rows),
13540        }
13541        return;
13542    }
13543
13544    prescale_with(x1, col_field, dtype, 1, |x1s| {
13545        prescale_with(x2, col_field, dtype, 2, |x2s| {
13546            let p1 = SendMut(o1.as_mut_ptr());
13547            let p2 = SendMut(o2.as_mut_ptr());
13548            let run_range = move |start: usize, end: usize| {
13549                for o in start..end {
13550                    let row = &q[o * cols..(o + 1) * cols];
13551                    let s1 = dot_i8_f32(row, x1s) * row_scale[o];
13552                    let s2 = dot_i8_f32(row, x2s) * row_scale[o];
13553                    // SAFETY: disjoint row ranges per worker.
13554                    unsafe {
13555                        *p1.at(o) = s1;
13556                        *p2.at(o) = s2;
13557                    }
13558                }
13559            };
13560            match pool {
13561                Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
13562                _ => run_range(0, rows),
13563            }
13564        });
13565    });
13566}
13567
13568#[derive(Clone, Copy)]
13569struct SendMut(*mut f32);
13570unsafe impl Send for SendMut {}
13571unsafe impl Sync for SendMut {}
13572
13573impl SendMut {
13574    #[inline]
13575    fn at(self, i: usize) -> *mut f32 {
13576        unsafe { self.0.add(i) }
13577    }
13578}
13579
13580#[cfg(test)]
13581mod tests {
13582
13583    /// `q8_round` must be `round().clamp(±127) as i8` bit for bit: every
13584    /// half-integer, their neighbours one ulp either side, the clamp
13585    /// boundary, huge values, infinities and NaN, plus a dense sweep.
13586    #[test]
13587    fn q8_round_is_round_clamp() {
13588        let reference = |t: f32| t.round().clamp(-127.0, 127.0) as i8;
13589        let mut probe = vec![
13590            0.0f32,
13591            -0.0,
13592            f32::NAN,
13593            f32::INFINITY,
13594            f32::NEG_INFINITY,
13595            f32::MAX,
13596            f32::MIN,
13597            1e30,
13598            -1e30,
13599            f32::MIN_POSITIVE,
13600            -f32::MIN_POSITIVE,
13601        ];
13602        for k in -300i32..=300 {
13603            let h = k as f32 * 0.5;
13604            let up = f32::from_bits(h.to_bits() + 1);
13605            let down = f32::from_bits(h.to_bits().wrapping_sub(1));
13606            for t in [h, up, down] {
13607                probe.push(t);
13608                probe.push(-t);
13609            }
13610        }
13611        let mut t = -140.0f32;
13612        while t < 140.0 {
13613            probe.push(t);
13614            t += 0.000_731;
13615        }
13616        for t in probe {
13617            assert_eq!(q8_round(t), reference(t), "t = {t:e} ({:#x})", t.to_bits());
13618        }
13619    }
13620
13621    use super::*;
13622
13623    /// The f32 `matmat` walks the batch four positions a row at a time; every
13624    /// output must still be the scalar-order sum `matvec` gives, bit for bit
13625    /// (a 7-row batch covers the four-wide block and the remainder).
13626    #[test]
13627    fn f32_matmat_equals_matvec_bitwise() {
13628        let (rows, cols, b) = (5usize, 37usize, 7usize);
13629        let w: Vec<f32> = (0..rows * cols).map(|i| ((i * 7919 % 113) as f32 - 56.0) / 37.0).collect();
13630        let xs: Vec<f32> = (0..b * cols).map(|i| ((i * 104_729 % 97) as f32 - 48.0) / 29.0).collect();
13631        let t = QTensor::from_f32(w, rows, cols);
13632        let mut out = vec![0f32; b * rows];
13633        t.matmat(&xs, b, &mut out, None);
13634        for bi in 0..b {
13635            let mut one = vec![0f32; rows];
13636            t.matvec(&xs[bi * cols..(bi + 1) * cols], &mut one, None);
13637            for o in 0..rows {
13638                assert_eq!(out[bi * rows + o].to_bits(), one[o].to_bits(), "bi {bi} row {o}");
13639            }
13640        }
13641    }
13642
13643    #[test]
13644    fn q2tp_i8_dot_matches_exact_on_grid() {
13645        // On-grid activations (±1 → sx=1/127, xq=±127 dequantizes
13646        // exactly, no outliers) must make the integer path agree with
13647        // the exact scalar walk to f32 rounding.
13648        let (rows, cols) = (5, 64);
13649        let gpr = cols / GROUP_SIZE;
13650        // Synthetic codes plane + a flat ladder: scales_into is not under
13651        // test here, so drive dot_q2tp_row_i8 / q2tp_row_exact directly
13652        // with hand-made scales.
13653        let chunks: Vec<u8> = (0..rows * gpr * Q2TP_CHUNK)
13654            .map(|i| (i as u32).wrapping_mul(2654435761) as u8)
13655            .collect();
13656        let scales: Vec<f32> = (0..gpr).map(|g| 0.5 + g as f32 * 0.25).collect();
13657        let x: Vec<f32> = (0..cols)
13658            .map(|i| if i % 3 == 0 { -1.0 } else { 1.0 })
13659            .collect();
13660        let act = split_act(&x);
13661        assert!(
13662            act.outliers.is_empty(),
13663            "on-grid input must have no outliers"
13664        );
13665        let gsum = q1_group_sums(&act.xq, gpr);
13666        for r in 0..rows {
13667            let exact = q2tp_row_exact(&chunks, r, gpr, &x, &scales);
13668            let fast = dot_q2tp_row_i8(&chunks, r, gpr, &act.xq, &gsum, &scales) * act.sx;
13669            assert!(
13670                (exact - fast).abs() <= exact.abs() * 1e-5 + 1e-5,
13671                "row {r}: exact {exact} vs i8 {fast}"
13672            );
13673        }
13674    }
13675
13676    #[cfg(target_arch = "x86_64")]
13677    #[test]
13678    fn q2tp_avx2_dot_matches_scalar_for_random_patterns() {
13679        // Compare the release AVX2 integer dot against the scalar oracle over
13680        // arbitrary packed bytes/activation signs.  This guards the exact
13681        // table-load path used after rejecting a faster-looking decoder whose
13682        // full-checkpoint greedy output drifted.
13683        if !std::arch::is_x86_feature_detected!("avx2") {
13684            return;
13685        }
13686        let mut seed = 0x9e3779b9u32;
13687        let mut next = || {
13688            seed = seed.wrapping_mul(1664525).wrapping_add(1013904223);
13689            seed
13690        };
13691        for _ in 0..20_000 {
13692            let mut ch = [0u8; Q2TP_CHUNK];
13693            let mut x = [0i8; GROUP_SIZE];
13694            for b in &mut ch {
13695                *b = next() as u8;
13696            }
13697            for v in &mut x {
13698                *v = (next() >> 24) as i8;
13699            }
13700            let mut reference = 0i32;
13701            for (k, &b) in ch.iter().enumerate() {
13702                reference += (b & 3) as i32 * x[k * 4] as i32;
13703                reference += ((b >> 2) & 3) as i32 * x[k * 4 + 1] as i32;
13704                reference += ((b >> 4) & 3) as i32 * x[k * 4 + 2] as i32;
13705                reference += ((b >> 6) & 3) as i32 * x[k * 4 + 3] as i32;
13706            }
13707            // SAFETY: guarded by the runtime AVX2 feature check and fixed
13708            // 8-byte/32-byte slice lengths above.
13709            let got = unsafe { q2tp_code_dot_avx2(&ch, &x) };
13710            assert_eq!(got, reference, "packed q2 lane mismatch");
13711        }
13712    }
13713
13714    #[test]
13715    fn q2tp_affine_fuses_half_scale_correction_without_changing_raw_decode() {
13716        let (rows, cols) = (1usize, GROUP_SIZE);
13717        let mut bytes = vec![0u8; Q2TP_CHUNK + 4 + 1];
13718        // Repeating symbols 0,1,2,0 at unit scale.  q2tp's raw B is
13719        // (c-1.5), while the affine Prism operator is (c-1.0).
13720        bytes[..Q2TP_CHUNK].fill(0x24); // codes 0,1,2,0 in LSB-first order
13721        bytes[Q2TP_CHUNK..Q2TP_CHUNK + 2].copy_from_slice(&0u16.to_le_bytes());
13722        bytes[Q2TP_CHUNK + 2..Q2TP_CHUNK + 4].copy_from_slice(&0u16.to_le_bytes());
13723        bytes[Q2TP_CHUNK + 4] = 1; // dtype16 rung 1 = 1.0
13724        let x = vec![1.0f32; cols];
13725        let mut raw = vec![0.0f32; rows];
13726        let mut affine = vec![0.0f32; rows];
13727        q2tp_matvec_for_test(&bytes, &x, rows, cols, &mut raw);
13728        q2tp_affine_matvec_for_test(&bytes, &x, rows, cols, &mut affine);
13729        assert_eq!(raw, vec![-24.0]);
13730        assert_eq!(affine, vec![-8.0]);
13731        assert!((affine[0] - (raw[0] + 0.5 * cols as f32)).abs() < 1e-6);
13732    }
13733
13734    #[test]
13735    fn q8_row_dot_fast_matches_scalar() {
13736        // The per-arch fast dot must agree with the exact scalar oracle
13737        // (same contract the fused q8 FFN arm rides on).
13738        let cols = 96;
13739        let row: Vec<u8> = (0..cols)
13740            .map(|i| ((i * 37 % 251) - 125) as i8 as u8)
13741            .collect();
13742        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.13).sin()).collect();
13743        let act = split_act(&x);
13744        let fast = q8_row_dot(&row, &act);
13745        let scalar = q8_row_dot_scalar(&row, &act);
13746        assert!(
13747            (fast - scalar).abs() <= scalar.abs() * 1e-5 + 1e-5,
13748            "fast {fast} vs scalar {scalar}"
13749        );
13750    }
13751
13752    #[test]
13753    fn f32_matvec_matches_matvec_rows_bitexact() {
13754        let (rows, cols) = (300, 40);
13755        let w: Vec<f32> = (0..rows * cols).map(|i| (i as f32 * 0.017).sin()).collect();
13756        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.05).cos()).collect();
13757        let qt = QTensor::from_f32(w.clone(), rows, cols);
13758
13759        let mut a = vec![0.0f32; rows];
13760        matvec_rows(None, &w, &x, &mut a);
13761        let mut b = vec![0.0f32; rows];
13762        qt.matvec(&x, &mut b, None);
13763        assert_eq!(a, b);
13764    }
13765
13766    #[test]
13767    fn sdot_kernel_exact_on_grid() {
13768        // Activations already on the i8 grid (±1 with amax=1 → sx=1/127,
13769        // xq=±127 dequantizes EXACTLY) → the SDOT path must match the
13770        // exact f32 dot to float rounding. This isolates kernel
13771        // correctness from quantization noise.
13772        eprintln!("sdot_enabled = {}", sdot_enabled());
13773        let (rows, cols) = (9, 80); // odd rows → exercises 4-row + tail
13774        let w: Vec<u8> = (0..rows * cols)
13775            .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
13776            .collect();
13777        let scales: Vec<f32> = (0..rows).map(|o| 0.005 + o as f32 * 0.001).collect();
13778        let x: Vec<f32> = (0..cols)
13779            .map(|i| match i % 3 {
13780                0 => 1.0,
13781                1 => -1.0,
13782                _ => 0.0,
13783            })
13784            .collect();
13785        let mut a = vec![0.0f32; rows];
13786        qmatvec(
13787            &w,
13788            &[],
13789            &scales,
13790            &x,
13791            &[],
13792            TensorDtype::Q8Row,
13793            rows,
13794            cols,
13795            &mut a,
13796            None,
13797        );
13798        for o in 0..rows {
13799            let mut acc = 0.0f32;
13800            for j in 0..cols {
13801                acc += (w[o * cols + j] as i8) as f32 * x[j];
13802            }
13803            let expect = acc * scales[o];
13804            assert!(
13805                (a[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
13806                "row {o}: {} vs {expect}",
13807                a[o]
13808            );
13809        }
13810    }
13811
13812    #[test]
13813    fn q1_tbl_fast_path_matches_reference() {
13814        // gpr = 8 exercises the TBL pair-load fast loop, and the LAST
13815        // row's final 4-tile window trips the 4B-overread guard (the
13816        // payload ends exactly at the last tile) — both paths must
13817        // agree with the dequant reference.
13818        let (rows, cols) = (5, 256);
13819        let gpr = cols / GROUP_SIZE;
13820        let mut bytes = Vec::new();
13821        for t in 0..rows * gpr {
13822            let s = 0.007 + (t % 11) as f32 * 0.004;
13823            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
13824            for j in 0..4 {
13825                bytes.push(((t * 53 + j * 89 + 7) % 249) as u8);
13826            }
13827        }
13828        let x: Vec<f32> = (0..cols)
13829            .map(|i| if (i * 5) % 7 < 3 { 1.0 } else { -1.0 })
13830            .collect();
13831        let mut w = vec![0.0f32; rows * cols];
13832        cortiq_core::quant::dequant_q1(&bytes, &mut w);
13833        let mut got = vec![0.0f32; rows];
13834        q1_matvec(&bytes, &x, rows, cols, &mut got, None);
13835        for o in 0..rows {
13836            let expect: f32 = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
13837            assert!(
13838                (got[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
13839                "row {o}: {} vs {expect}",
13840                got[o]
13841            );
13842        }
13843        // Blocked 1×4 batch (b=5: one quad + remainder) must equal the
13844        // single-matvec path bit-for-bit.
13845        let b = 5usize;
13846        let mut xs_all = Vec::new();
13847        for bi in 0..b {
13848            xs_all.extend(x.iter().map(|v| if bi % 2 == 0 { *v } else { -*v }));
13849        }
13850        let mut mm = vec![0.0f32; b * rows];
13851        q1_matmat(&bytes, &xs_all, b, rows, cols, &mut mm, None);
13852        for bi in 0..b {
13853            let mut single = vec![0.0f32; rows];
13854            q1_matvec(
13855                &bytes,
13856                &xs_all[bi * cols..(bi + 1) * cols],
13857                rows,
13858                cols,
13859                &mut single,
13860                None,
13861            );
13862            assert_eq!(&mm[bi * rows..(bi + 1) * rows], &single[..], "stream {bi}");
13863        }
13864    }
13865
13866    #[test]
13867    fn q1_kernels_match_exact_reference() {
13868        // Synthetic q1 payload: 6-byte tiles [f16 scale][4B bits].
13869        let (rows, cols) = (7, 96);
13870        let gpr = cols / GROUP_SIZE;
13871        let mut bytes = Vec::new();
13872        for t in 0..rows * gpr {
13873            let s = 0.01 + (t % 13) as f32 * 0.003;
13874            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
13875            for j in 0..4 {
13876                bytes.push(((t * 31 + j * 97) % 251) as u8);
13877            }
13878        }
13879        // On-grid activations (±1, amax 1) → the SDOT path is exact.
13880        let x: Vec<f32> = (0..cols)
13881            .map(|i| if i % 3 == 0 { 1.0 } else { -1.0 })
13882            .collect();
13883        // Reference through the core dequant.
13884        let mut w = vec![0.0f32; rows * cols];
13885        cortiq_core::quant::dequant_q1(&bytes, &mut w);
13886        let mut expect = vec![0.0f32; rows];
13887        for o in 0..rows {
13888            expect[o] = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
13889        }
13890        let mut got = vec![0.0f32; rows];
13891        q1_matvec(&bytes, &x, rows, cols, &mut got, None);
13892        for o in 0..rows {
13893            assert!(
13894                (got[o] - expect[o]).abs() < 1e-3 * expect[o].abs().max(1e-3),
13895                "row {o}: {} vs {}",
13896                got[o],
13897                expect[o]
13898            );
13899        }
13900        // Pair and batch paths agree with the single path.
13901        let x2: Vec<f32> = x.iter().map(|v| -v).collect();
13902        let (mut a1, mut a2) = (vec![0.0f32; rows], vec![0.0f32; rows]);
13903        q1_matvec2(&bytes, &x, &x2, rows, cols, &mut a1, &mut a2, None);
13904        assert_eq!(a1, got);
13905        let mut xs = x.clone();
13906        xs.extend_from_slice(&x2);
13907        let mut mm = vec![0.0f32; 2 * rows];
13908        q1_matmat(&bytes, &xs, 2, rows, cols, &mut mm, None);
13909        assert_eq!(&mm[..rows], got.as_slice());
13910        assert_eq!(&mm[rows..], a2.as_slice());
13911    }
13912
13913    /// A Bonsai-shaped q1 payload (one f16 scale per 128-group, stored in
13914    /// all four tiles; pseudo-random signs) and its q1g repack.
13915    fn q1_q1g_pair(rows: usize, cols: usize, salt: usize) -> (Vec<u8>, Vec<u8>) {
13916        let tiles = rows * cols / GROUP_SIZE;
13917        let mut q1 = Vec::with_capacity(tiles * Q1_TILE);
13918        for t in 0..tiles {
13919            let s = 0.004 + ((t / 4 * 37 + salt) % 29) as f32 * 0.0011;
13920            q1.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
13921            let bits = ((t * 2_654_435_761 + salt * 40_503) as u32).rotate_left((t % 31) as u32);
13922            q1.extend_from_slice(&bits.to_le_bytes());
13923        }
13924        let rp = cortiq_core::quant::repack_q1_to_q1g(&q1, rows, cols, 0.0);
13925        assert_eq!(rp.differing, 0, "uniform scales repack exactly");
13926        (q1, rp.data)
13927    }
13928
13929    /// Activations off the int8 grid with two outlier channels, so the
13930    /// A8W8 path quantizes AND takes the exact outlier correction.
13931    fn q1g_test_x(cols: usize, salt: usize) -> Vec<f32> {
13932        let mut x: Vec<f32> = (0..cols)
13933            .map(|i| ((i * 13 + salt * 7) as f32 * 0.071).sin() * 0.9)
13934            .collect();
13935        x[3 % cols] = 41.0;
13936        x[(cols / 2 + 5) % cols] = -37.5;
13937        x
13938    }
13939
13940    /// Every CPU row kernel of q1g returns the q1 kernel's f32 BIT FOR BIT
13941    /// on a uniform-scale repack: A8W8 rows (SIMD of this machine, and the
13942    /// x86 AVX2/VNNI twins directly), the exact scalar rows (CMF_SDOT=0
13943    /// contract), the outlier term, the drivers (matvec, matvec2, matmat
13944    /// with the blocked 1×4 quad + remainder) — including gpr = 4 and the
13945    /// last row, where q1's pair load falls back to its scalar tail.
13946    #[test]
13947    fn q1g_cpu_kernels_equal_q1_bit_for_bit() {
13948        for &(rows, cols) in &[(3usize, 128usize), (5, 256), (9, 2048), (2, 6144)] {
13949            let gpr = cols / GROUP_SIZE;
13950            let (q1, q1g) = q1_q1g_pair(rows, cols, rows + cols);
13951            let (wa, wb) = (BinW::Q1(&q1), BinW::q1g(&q1g, rows, cols));
13952            let x = q1g_test_x(cols, rows);
13953            let act = split_act(&x);
13954            assert!(!act.outliers.is_empty(), "test wants the outlier path");
13955            let gsum = q1_group_sums(&act.xq, gpr);
13956            for r in 0..rows {
13957                let (a, b) = (wa.row_exact(r, gpr, &x), wb.row_exact(r, gpr, &x));
13958                assert_eq!(a.to_bits(), b.to_bits(), "{rows}x{cols} exact row {r}");
13959                let (a, b) = (
13960                    wa.dot_i8(r, gpr, &act.xq, &gsum),
13961                    wb.dot_i8(r, gpr, &act.xq, &gsum),
13962                );
13963                assert_eq!(a.to_bits(), b.to_bits(), "{rows}x{cols} i8 row {r}");
13964                for j in [0, 31, 32, 127, cols - 1] {
13965                    let (a, b) = (wa.outlier(r, gpr, j), wb.outlier(r, gpr, j));
13966                    assert_eq!(
13967                        (a.0.to_bits(), a.1.to_bits()),
13968                        (b.0.to_bits(), b.1.to_bits())
13969                    );
13970                }
13971                #[cfg(target_arch = "x86_64")]
13972                if avx2_enabled() {
13973                    let sc = cortiq_core::quant::q1g_scale_offset(rows, cols);
13974                    // SAFETY: AVX2 checked; payloads sized for (rows, cols).
13975                    unsafe {
13976                        let a = dot_q1_row_avx2(&q1, r, gpr, &act.xq, &gsum);
13977                        let b = dot_q1g_row_avx2(&q1g, sc, r, gpr, &act.xq, &gsum);
13978                        assert_eq!(a.to_bits(), b.to_bits(), "avx2 row {r}");
13979                    }
13980                    if std::is_x86_feature_detected!("avx512vnni")
13981                        && std::is_x86_feature_detected!("avx512bw")
13982                        && std::is_x86_feature_detected!("avx512vl")
13983                    {
13984                        // SAFETY: VNNI family detected just above.
13985                        unsafe {
13986                            let a = dot_q1_row_vnni(&q1, r, gpr, &act.xq, &gsum);
13987                            let b = dot_q1g_row_vnni(&q1g, sc, r, gpr, &act.xq, &gsum);
13988                            assert_eq!(a.to_bits(), b.to_bits(), "vnni row {r}");
13989                        }
13990                    }
13991                }
13992            }
13993            // Drivers: single, pair, batch (b = 5: one blocked quad + one).
13994            let (mut ga, mut gb) = (vec![0f32; rows], vec![0f32; rows]);
13995            bin_matvec(wa, &x, rows, cols, &mut ga, None);
13996            bin_matvec(wb, &x, rows, cols, &mut gb, None);
13997            assert_eq!(ga, gb, "{rows}x{cols} matvec");
13998            let x2 = q1g_test_x(cols, rows + 1);
13999            let (mut a1, mut a2, mut b1, mut b2) = (
14000                vec![0f32; rows],
14001                vec![0f32; rows],
14002                vec![0f32; rows],
14003                vec![0f32; rows],
14004            );
14005            bin_matvec2(wa, &x, &x2, rows, cols, &mut a1, &mut a2, None);
14006            bin_matvec2(wb, &x, &x2, rows, cols, &mut b1, &mut b2, None);
14007            assert_eq!((&a1, &a2), (&b1, &b2), "{rows}x{cols} matvec2");
14008            let b = 5usize;
14009            let xs: Vec<f32> = (0..b).flat_map(|k| q1g_test_x(cols, k * 3 + 2)).collect();
14010            let (mut ma, mut mb) = (vec![0f32; b * rows], vec![0f32; b * rows]);
14011            bin_matmat(wa, &xs, b, rows, cols, &mut ma, None);
14012            bin_matmat(wb, &xs, b, rows, cols, &mut mb, None);
14013            assert_eq!(ma, mb, "{rows}x{cols} matmat");
14014            // The blocked 1×4 kernels directly (the driver only takes them
14015            // under CMF_X86_BLOCKED / SDOT, a process-wide switch).
14016            let acts: Vec<(SplitAct, Vec<i32>)> = (0..4)
14017                .map(|k| {
14018                    let a = split_act(&xs[k * cols..(k + 1) * cols]);
14019                    let g = q1_group_sums(&a.xq, gpr);
14020                    (a, g)
14021                })
14022                .collect();
14023            let xq4: [&[i8]; 4] = std::array::from_fn(|k| acts[k].0.xq.as_slice());
14024            let gs4: [&[i32]; 4] = std::array::from_fn(|k| acts[k].1.as_slice());
14025            for r in 0..rows {
14026                #[cfg(target_arch = "x86_64")]
14027                if avx2_enabled() {
14028                    // SAFETY: AVX2 checked; payloads sized for (rows, cols).
14029                    let (a, b) = unsafe {
14030                        (
14031                            wa.dot_1x4_x86(r, gpr, xq4, gs4),
14032                            wb.dot_1x4_x86(r, gpr, xq4, gs4),
14033                        )
14034                    };
14035                    assert_eq!(a.map(f32::to_bits), b.map(f32::to_bits), "1x4 row {r}");
14036                }
14037                #[cfg(target_arch = "aarch64")]
14038                if sdot_enabled() {
14039                    // SAFETY: SDOT checked; payloads sized for (rows, cols).
14040                    let (a, b) = unsafe {
14041                        (
14042                            wa.dot_1x4_sdot(r, gpr, xq4, gs4),
14043                            wb.dot_1x4_sdot(r, gpr, xq4, gs4),
14044                        )
14045                    };
14046                    assert_eq!(a.map(f32::to_bits), b.map(f32::to_bits), "1x4 row {r}");
14047                }
14048                let _ = (r, &xq4, &gs4);
14049            }
14050        }
14051    }
14052
14053    /// CPU kernel A/B, q1 vs q1g matvec (A8W8 path of this machine, the
14054    /// exact path too), Bonsai-1.7B decode shapes, arms alternated every
14055    /// round, medians printed. `CMF_THREADS` sets the pool (default 8).
14056    /// `cargo test -p cortiq-engine --release --lib q1g_cpu_matvec_speed --
14057    /// --ignored --nocapture`
14058    #[test]
14059    #[ignore]
14060    fn q1g_cpu_matvec_speed() {
14061        use crate::pool::Pool;
14062        let threads = std::env::var("CMF_THREADS")
14063            .ok()
14064            .and_then(|v| v.parse().ok())
14065            .unwrap_or(8usize);
14066        let pool = Pool::new(threads);
14067        for &(name, rows, cols, reps) in &[
14068            ("gate/up_proj", 6144usize, 2048usize, 41usize),
14069            ("down_proj", 2048, 6144, 41),
14070            ("lm_head", 151_669, 2048, 9),
14071        ] {
14072            let (q1, q1g) = q1_q1g_pair(rows, cols, 3);
14073            let (wa, wb) = (BinW::Q1(&q1), BinW::q1g(&q1g, rows, cols));
14074            let x = q1g_test_x(cols, 1);
14075            let mut out = vec![0f32; rows];
14076            let mut time = |w: BinW| {
14077                let t0 = std::time::Instant::now();
14078                bin_matvec(w, &x, rows, cols, &mut out, Some(&pool));
14079                t0.elapsed().as_secs_f64()
14080            };
14081            let _ = (time(wa), time(wb));
14082            let (mut ta, mut tb) = (Vec::new(), Vec::new());
14083            for r in 0..reps {
14084                if r % 2 == 0 {
14085                    ta.push(time(wa));
14086                    tb.push(time(wb));
14087                } else {
14088                    tb.push(time(wb));
14089                    ta.push(time(wa));
14090                }
14091            }
14092            ta.sort_by(|a, b| a.partial_cmp(b).unwrap());
14093            tb.sort_by(|a, b| a.partial_cmp(b).unwrap());
14094            let (ma, mb) = (ta[reps / 2], tb[reps / 2]);
14095            eprintln!(
14096                "cpu x{threads} a8w8={} {name:>12} {rows}x{cols}: q1 {:8.1} us | q1g {:8.1} us | q1g/q1 time {:.3}",
14097                a8w8_enabled(),
14098                ma * 1e6,
14099                mb * 1e6,
14100                mb / ma
14101            );
14102        }
14103    }
14104
14105    /// The mapped QTensor paths a real Bonsai q1g file takes — matvec, the
14106    /// multi-projection job (all-q1g and q1/q1g mixed), the fused gate/up
14107    /// with SiLU, matmat, matvec2 and the embedding row — equal the q1
14108    /// tensor they were repacked from, bit for bit.
14109    #[test]
14110    fn q1g_mapped_tensor_paths_equal_q1() {
14111        use crate::pool::Pool;
14112        use cortiq_core::{CmfHeader, CmfModel, TensorSpec};
14113        let (rows, cols) = (300usize, 256usize);
14114        let (gq1, gq1g) = q1_q1g_pair(rows, cols, 1);
14115        let (uq1, uq1g) = q1_q1g_pair(rows, cols, 2);
14116        let header: CmfHeader = serde_json::from_value(serde_json::json!({
14117            "format": "cmf", "version": cortiq_core::CMF_VERSION,
14118            "arch": {"arch_name": "tiny-q1g", "hidden_size": cols,
14119                "intermediate_size": rows, "num_layers": 1, "num_attention_heads": 2,
14120                "num_kv_heads": 1, "head_dim": 128, "vocab_size": rows,
14121                "layer_types": ["FullAttention"], "rms_norm_eps": 1e-6,
14122                "max_position_embeddings": 8},
14123            "quant_type": "VBIT"
14124        }))
14125        .unwrap();
14126        let spec = |name: &str, dtype, data: &Vec<u8>| TensorSpec {
14127            name: name.into(),
14128            dtype,
14129            shape: vec![rows, cols],
14130            data: data.clone(),
14131        };
14132        let specs = [
14133            spec("g1", TensorDtype::Q1, &gq1),
14134            spec("u1", TensorDtype::Q1, &uq1),
14135            spec("gg", TensorDtype::Q1G, &gq1g),
14136            spec("ug", TensorDtype::Q1G, &uq1g),
14137        ];
14138        let dir = std::env::temp_dir().join(format!("cmf-q1g-mapped-{}", std::process::id()));
14139        std::fs::create_dir_all(&dir).unwrap();
14140        let path = dir.join("m.cmf");
14141        CmfModel::write(&path, &header, &specs, None, None).unwrap();
14142        let model = Arc::new(CmfModel::open(&path).unwrap());
14143        assert!(model.verify().is_empty(), "{:?}", model.verify());
14144        let t = |n: &str| QTensor::from_model(&model, n).unwrap();
14145        let (g1, u1, gg, ug) = (t("g1"), t("u1"), t("gg"), t("ug"));
14146        assert_eq!(gg.model_dtype(), Some(TensorDtype::Q1G));
14147        assert!(gg.is_q1() && gg.graph_weight().map(|w| w.2) == Some(10));
14148        let pool = Pool::new(3);
14149        let x = q1g_test_x(cols, 4);
14150        let run = |a: &QTensor| {
14151            let mut o = vec![0f32; rows];
14152            a.matvec(&x, &mut o, Some(&pool));
14153            o
14154        };
14155        assert_eq!(run(&g1), run(&gg), "matvec");
14156        let many = |a: &QTensor, b: &QTensor| {
14157            let (mut oa, mut ob) = (vec![0f32; rows], vec![0f32; rows]);
14158            QTensor::matvec_many([a, b], &x, [&mut oa, &mut ob], Some(&pool));
14159            (oa, ob)
14160        };
14161        let want = many(&g1, &u1);
14162        assert_eq!(want, many(&gg, &ug), "matvec_many q1g");
14163        assert_eq!(want, many(&g1, &ug), "matvec_many mixed");
14164        assert_eq!(want, (run(&g1), run(&u1)), "matvec_many == singles");
14165        let fused = |g: &QTensor, u: &QTensor| {
14166            let mut o = vec![0f32; rows];
14167            let ok = QTensor::matvec_silu_mul(g, u, &x, &mut o, Some(&pool));
14168            (ok, o)
14169        };
14170        let (ok_a, fa) = fused(&g1, &u1);
14171        let (ok_b, fb) = fused(&gg, &ug);
14172        assert_eq!(ok_a, ok_b, "q1g takes the fused arm whenever q1 does");
14173        assert_eq!(fa, fb, "fused gate/up");
14174        let b = 6usize;
14175        let xs: Vec<f32> = (0..b).flat_map(|k| q1g_test_x(cols, k + 9)).collect();
14176        let mm = |a: &QTensor| {
14177            let mut o = vec![0f32; b * rows];
14178            a.matmat(&xs, b, &mut o, Some(&pool));
14179            o
14180        };
14181        assert_eq!(mm(&g1), mm(&gg), "matmat");
14182        let x2 = q1g_test_x(cols, 77);
14183        let pair = |a: &QTensor| {
14184            let (mut o1, mut o2) = (vec![0f32; rows], vec![0f32; rows]);
14185            a.matvec2(&x, &x2, &mut o1, &mut o2, Some(&pool));
14186            (o1, o2)
14187        };
14188        assert_eq!(pair(&g1), pair(&gg), "matvec2");
14189        for r in [0, 1, rows / 2, rows - 1] {
14190            let (mut ra, mut rb) = (vec![0f32; cols], vec![0f32; cols]);
14191            g1.row_f32(r, &mut ra);
14192            gg.row_f32(r, &mut rb);
14193            assert_eq!(ra, rb, "embedding row {r}");
14194        }
14195        drop((g1, u1, gg, ug, model));
14196        let _ = std::fs::remove_dir_all(&dir);
14197    }
14198
14199    /// A `blr` tensor read from a real CMF file: it maps to the factored
14200    /// arm (not a dense copy), declines every device-graph accessor, and
14201    /// every QTensor entry point a layer may take — matvec, matvec2,
14202    /// matmat (table loop and GEMM sizes), the multi-tensor job mixed with
14203    /// a q8_row sibling, the fused gate/up refusal, rows and columns — equals
14204    /// the dense reference built by the core dequant.
14205    #[test]
14206    fn blr_mapped_tensor_paths_equal_dense() {
14207        use crate::pool::Pool;
14208        use cortiq_core::{CmfHeader, CmfModel, TensorSpec};
14209        let (rows, cols, r) = (96usize, 128usize, 16usize);
14210        let mut rng = crate::blr::tests::Rng(41);
14211        let gate = crate::blr::tests::make(&mut rng, rows, cols, r, 2);
14212        let up = crate::blr::tests::make(&mut rng, rows, cols, r, 1);
14213        // q8_row sibling: [int8 rows·cols][f16 row scale · rows].
14214        let mut q8: Vec<u8> = (0..rows * cols)
14215            .map(|i| (((i * 89) % 253) as i32 - 126) as i8 as u8)
14216            .collect();
14217        for o in 0..rows {
14218            q8.extend_from_slice(&cortiq_core::quant::f32_to_f16(0.002 + o as f32 * 1e-4).to_le_bytes());
14219        }
14220        let header: CmfHeader = serde_json::from_value(serde_json::json!({
14221            "format": "cmf", "version": cortiq_core::CMF_VERSION,
14222            "arch": {"arch_name": "tiny-blr", "hidden_size": cols,
14223                "intermediate_size": rows, "num_layers": 1, "num_attention_heads": 2,
14224                "num_kv_heads": 1, "head_dim": 64, "vocab_size": rows,
14225                "layer_types": ["FullAttention"], "rms_norm_eps": 1e-6,
14226                "max_position_embeddings": 8},
14227            "quant_type": "Q8_2F"
14228        }))
14229        .unwrap();
14230        let specs = [
14231            TensorSpec {
14232                name: "g".into(),
14233                dtype: TensorDtype::BinaryLowRank,
14234                shape: vec![rows, cols],
14235                data: gate.clone(),
14236            },
14237            TensorSpec {
14238                name: "u".into(),
14239                dtype: TensorDtype::BinaryLowRank,
14240                shape: vec![rows, cols],
14241                data: up.clone(),
14242            },
14243            TensorSpec {
14244                name: "q".into(),
14245                dtype: TensorDtype::Q8Row,
14246                shape: vec![rows, cols],
14247                data: q8,
14248            },
14249        ];
14250        let dir = std::env::temp_dir().join(format!("cmf-blr-mapped-{}", std::process::id()));
14251        std::fs::create_dir_all(&dir).unwrap();
14252        let path = dir.join("m.cmf");
14253        CmfModel::write(&path, &header, &specs, None, None).unwrap();
14254        let model = Arc::new(CmfModel::open(&path).unwrap());
14255        assert!(model.verify().is_empty(), "{:?}", model.verify());
14256        let t = |n: &str| QTensor::from_model(&model, n).unwrap();
14257        let (g, u, q) = (t("g"), t("u"), t("q"));
14258        assert!(g.is_blr() && u.is_blr() && !q.is_blr());
14259        assert!(g.graph_weight().is_none() && g.graph_weight_descriptor().is_none());
14260        assert!(g.as_f32().is_none() && g.model_idx().is_none());
14261        assert!(g.mapped_device_gemm().is_none() && g.q1_parts().is_none());
14262        assert!(!g.sparse_col_ok());
14263        assert_eq!((g.rows(), g.cols()), (rows, cols));
14264        let deq = |bytes: &[u8]| {
14265            let mut w = vec![0f32; rows * cols];
14266            cortiq_core::blr::dequant_blr(bytes, rows, cols, &mut w).unwrap();
14267            w
14268        };
14269        let (wg, wu) = (deq(&gate), deq(&up));
14270        let mv = |w: &[f32], x: &[f32]| -> Vec<f32> {
14271            (0..rows)
14272                .map(|i| (0..cols).map(|j| w[i * cols + j] * x[j]).sum())
14273                .collect()
14274        };
14275        let close = |a: &[f32], b: &[f32], what: &str| {
14276            let s = b.iter().map(|v| v * v).sum::<f32>().sqrt().max(1e-30);
14277            let e = a.iter().zip(b).map(|(x, y)| (x - y) * (x - y)).sum::<f32>().sqrt();
14278            assert!(e / s < 1e-5, "{what}: rel {}", e / s);
14279        };
14280        let pool = Pool::new(3);
14281        let x: Vec<f32> = (0..cols).map(|_| rng.f()).collect();
14282        let x2: Vec<f32> = (0..cols).map(|_| rng.f()).collect();
14283        let mut o = vec![0f32; rows];
14284        g.matvec(&x, &mut o, Some(&pool));
14285        close(&o, &mv(&wg, &x), "matvec");
14286        let (mut o1, mut o2) = (vec![0f32; rows], vec![0f32; rows]);
14287        g.matvec2(&x, &x2, &mut o1, &mut o2, Some(&pool));
14288        close(&o1, &mv(&wg, &x), "matvec2 a");
14289        close(&o2, &mv(&wg, &x2), "matvec2 b");
14290        for b in [2usize, 5] {
14291            let xs: Vec<f32> = (0..b * cols).map(|_| rng.f()).collect();
14292            let mut om = vec![0f32; b * rows];
14293            g.matmat(&xs, b, &mut om, Some(&pool));
14294            for t in 0..b {
14295                close(
14296                    &om[t * rows..(t + 1) * rows],
14297                    &mv(&wg, &xs[t * cols..(t + 1) * cols]),
14298                    "matmat",
14299                );
14300            }
14301        }
14302        // Multi-tensor job mixed with a q8_row sibling: per-tensor fallback,
14303        // the q8_row output identical to its own matvec.
14304        let (mut a, mut bq, mut c) = (vec![0f32; rows], vec![0f32; rows], vec![0f32; rows]);
14305        QTensor::matvec_many([&g, &q, &u], &x, [&mut a, &mut bq, &mut c], Some(&pool));
14306        close(&a, &mv(&wg, &x), "matvec_many blr gate");
14307        close(&c, &mv(&wu, &x), "matvec_many blr up");
14308        let mut qo = vec![0f32; rows];
14309        q.matvec(&x, &mut qo, Some(&pool));
14310        assert_eq!(bq, qo, "matvec_many q8_row sibling");
14311        let mut f = vec![0f32; rows];
14312        assert!(!QTensor::matvec_silu_mul(&g, &u, &x, &mut f, Some(&pool)));
14313        let mut row = vec![0f32; cols];
14314        g.row_f32(rows / 3, &mut row);
14315        close(&row, &wg[(rows / 3) * cols..(rows / 3 + 1) * cols], "row_f32");
14316        let mut scratch = vec![0f32; cols];
14317        let d = g.row_dot(7, &x, &mut scratch);
14318        let want: f32 = (0..cols).map(|j| wg[7 * cols + j] * x[j]).sum();
14319        assert!((d - want).abs() < 1e-4 * (1.0 + want.abs()), "row_dot");
14320        let mut col = vec![0f32; rows];
14321        g.add_col_scaled(cols / 2, 1.5, &mut col);
14322        let want: Vec<f32> = (0..rows).map(|i| 1.5 * wg[i * cols + cols / 2]).collect();
14323        close(&col, &want, "add_col_scaled");
14324        drop((g, u, q, model));
14325        let _ = std::fs::remove_dir_all(&dir);
14326    }
14327
14328    #[test]
14329    fn repack_is_bit_identical() {
14330        // The interleaved-repack kernel must produce EXACTLY the same
14331        // bits as the mmap-layout kernel: integer accumulation is order-
14332        // exact, the f32 epilogue is identical. Odd rows exercise the
14333        // tail; direct range calls exercise unaligned pool splits.
14334        let (rows, cols) = (267, 96); // 66 groups + 3 tail rows, cols % 16 == 0
14335        let w: Vec<u8> = (0..rows * cols)
14336            .map(|i| (((i * 89) % 253) as i32 - 126) as i8 as u8)
14337            .collect();
14338        let scales: Vec<f32> = (0..rows).map(|o| 0.003 + o as f32 * 0.0007).collect();
14339        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.37).sin() * 2.0).collect();
14340        let rep = q8_repack_layout(&w, rows, cols);
14341        // Group interleave round-trips.
14342        for g in 0..rows / 4 {
14343            for c in 0..cols / 16 {
14344                for lane in 0..4 {
14345                    assert_eq!(
14346                        &rep[g * 4 * cols + c * 64 + lane * 16
14347                            ..g * 4 * cols + c * 64 + lane * 16 + 16],
14348                        &w[(g * 4 + lane) * cols + c * 16..(g * 4 + lane) * cols + c * 16 + 16],
14349                    );
14350                }
14351            }
14352        }
14353        let mut a = vec![0.0f32; rows];
14354        qmatvec(
14355            &w,
14356            &[],
14357            &scales,
14358            &x,
14359            &[],
14360            TensorDtype::Q8Row,
14361            rows,
14362            cols,
14363            &mut a,
14364            None,
14365        );
14366        let mut b = vec![0.0f32; rows];
14367        qmatvec(
14368            &w,
14369            &rep,
14370            &scales,
14371            &x,
14372            &[],
14373            TensorDtype::Q8Row,
14374            rows,
14375            cols,
14376            &mut b,
14377            None,
14378        );
14379        assert_eq!(a, b, "full-range repack output diverged");
14380
14381        #[cfg(target_arch = "aarch64")]
14382        if sdot_enabled() {
14383            // Unaligned range split (pool workers get arbitrary bounds).
14384            let act = split_act(&x);
14385            let mut c1 = vec![0.0f32; rows];
14386            let mut c2 = vec![0.0f32; rows];
14387            q8_range_sdot(
14388                &w,
14389                &[],
14390                &scales,
14391                &act,
14392                cols,
14393                SendMut(c1.as_mut_ptr()),
14394                3,
14395                rows - 2,
14396            );
14397            q8_range_sdot(
14398                &w,
14399                &rep,
14400                &scales,
14401                &act,
14402                cols,
14403                SendMut(c2.as_mut_ptr()),
14404                3,
14405                rows - 2,
14406            );
14407            assert_eq!(c1, c2, "unaligned-range repack output diverged");
14408        }
14409    }
14410
14411    #[test]
14412    fn sdot_a8w8_noise_is_bounded() {
14413        // Off-grid activations: A8 quantization noise must stay small in
14414        // relative L2 over the whole output (realistic accuracy contract;
14415        // vmfcore measured argmax-identical decode on real models).
14416        let (rows, cols) = (16, 512);
14417        let w: Vec<u8> = (0..rows * cols)
14418            .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
14419            .collect();
14420        let scales = vec![0.01f32; rows];
14421        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.21).sin()).collect();
14422        let mut a = vec![0.0f32; rows];
14423        qmatvec(
14424            &w,
14425            &[],
14426            &scales,
14427            &x,
14428            &[],
14429            TensorDtype::Q8Row,
14430            rows,
14431            cols,
14432            &mut a,
14433            None,
14434        );
14435        let (mut num, mut den) = (0f64, 0f64);
14436        for o in 0..rows {
14437            let mut acc = 0.0f32;
14438            for j in 0..cols {
14439                acc += (w[o * cols + j] as i8) as f32 * x[j];
14440            }
14441            let expect = acc * scales[o];
14442            num += ((a[o] - expect) as f64).powi(2);
14443            den += (expect as f64).powi(2);
14444        }
14445        let rel = (num / den.max(1e-12)).sqrt();
14446        assert!(rel < 0.05, "A8W8 relative L2 error too high: {rel}");
14447    }
14448
14449    #[test]
14450    fn i8_dot_neon_matches_scalar() {
14451        let n = 100;
14452        let w: Vec<u8> = (0..n).map(|i| ((i * 37 + 11) % 251) as u8).collect();
14453        let x: Vec<f32> = (0..n).map(|i| (i as f32 * 0.13).sin()).collect();
14454        let mut scalar = 0.0f32;
14455        for j in 0..n {
14456            scalar += (w[j] as i8) as f32 * x[j];
14457        }
14458        let fast = dot_i8_f32(&w, &x);
14459        assert!((scalar - fast).abs() < 1e-3 * scalar.abs().max(1.0));
14460    }
14461
14462    /// Fused vbit matvec must match full dequant_vbit + dense matvec.
14463    #[test]
14464    fn vbitmatvec_matches_full_dequant() {
14465        let (rows, cols) = (6, 64);
14466        let ng = cols / GROUP_SIZE;
14467        // Hand-craft: bits per row, f16 scales, packed rows.
14468        let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
14469        let mut bytes = bits.clone();
14470        for g in 0..rows * ng {
14471            let s = 0.02 + 0.001 * g as f32;
14472            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
14473        }
14474        for r in 0..rows {
14475            let b = bits[r] as usize;
14476            let (mut acc, mut nb) = (0u64, 0usize);
14477            let mut rowbytes = Vec::new();
14478            for i in 0..cols {
14479                let v = ((i * 7 + r * 13) % (1 << b)) as u64;
14480                acc = (acc << b) | v;
14481                nb += b;
14482                while nb >= 8 {
14483                    nb -= 8;
14484                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
14485                }
14486            }
14487            if nb > 0 {
14488                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
14489            }
14490            bytes.extend_from_slice(&rowbytes);
14491        }
14492        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
14493
14494        let mut reference = vec![0f32; rows * cols];
14495        cortiq_core::quant::dequant_vbit(&bytes, rows, cols, &mut reference).unwrap();
14496        let mut expect = vec![0f32; rows];
14497        for r in 0..rows {
14498            expect[r] = reference[r * cols..(r + 1) * cols]
14499                .iter()
14500                .zip(&x)
14501                .map(|(w, xv)| w * xv)
14502                .sum();
14503        }
14504        let mut got = vec![0f32; rows];
14505        let offsets = vbit_row_offsets(&bytes, rows, cols);
14506        vbitmatvec(&bytes, &offsets, &x, rows, cols, &mut got, None);
14507        // SDOT path quantizes activations to i8 (A8W8): bounded noise,
14508        // same contract as q8 (exact path is pinned by CMF_SDOT=0 in
14509        // the golden-parity gate).
14510        let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
14511        let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
14512        for r in 0..rows {
14513            assert!(
14514                (got[r] - expect[r]).abs() < tol * scale,
14515                "row {r}: {} vs {}",
14516                got[r],
14517                expect[r]
14518            );
14519        }
14520    }
14521
14522    /// Fused q4 matvec must match the reference full-dequant + dense
14523    /// matvec bit-for-bit in structure (same f32 math, group order).
14524    /// vbit matmat: the blocked 1×4 leg must match the per-row path
14525    /// (paired env toggle; larger shape so both code paths engage).
14526    #[test]
14527    #[cfg(target_arch = "x86_64")]
14528    fn vbit_matmat_blocked_matches_per_row() {
14529        let (rows, cols, b) = (64usize, 128usize, 9usize);
14530        let ng = cols / GROUP_SIZE;
14531        let bits: Vec<u8> = (0..rows).map(|r| [3u8, 4, 5, 6][r % 4]).collect();
14532        let mut bytes = bits.clone();
14533        for g in 0..rows * ng {
14534            let sc = 0.02 + 0.0005 * g as f32;
14535            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
14536        }
14537        for r in 0..rows {
14538            let bw = bits[r] as usize;
14539            let (mut acc, mut nb) = (0u64, 0usize);
14540            let mut rowbytes = Vec::new();
14541            for i in 0..cols {
14542                let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
14543                acc = (acc << bw) | v;
14544                nb += bw;
14545                while nb >= 8 {
14546                    nb -= 8;
14547                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
14548                }
14549            }
14550            if nb > 0 {
14551                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
14552            }
14553            bytes.extend_from_slice(&rowbytes);
14554        }
14555        let x: Vec<f32> = (0..b * cols)
14556            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
14557            .collect();
14558        let offsets = vbit_row_offsets(&bytes, rows, cols);
14559        let mut y_a = vec![0f32; b * rows];
14560        let mut y_b = vec![0f32; b * rows];
14561        unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
14562        vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_a, None);
14563        unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
14564        vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_b, None);
14565        unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
14566        let max_d = y_a
14567            .iter()
14568            .zip(&y_b)
14569            .map(|(p, q)| (p - q).abs())
14570            .fold(0.0f32, f32::max);
14571        assert!(max_d < 1e-4, "vbit blocked ≠ per-row: max|Δ| = {max_d}");
14572    }
14573
14574    /// q4t blocked 1×4 (SDOT on ARM, AVX2 on x86) must equal the
14575    /// per-row path exactly: same nibble unpack, same group order,
14576    /// same f32 accumulation — batch == matvec bit-for-bit. b=9 covers
14577    /// two full 1×4 blocks plus a remainder through the single-row
14578    /// kernel. (Both paths produce identical output, so the shared
14579    /// CMF_X86_BLOCKED env var racing with other tests cannot flip
14580    /// the verdict — worst case both sides take the same path.)
14581    #[test]
14582    fn q4t_matmat_blocked_matches_per_row() {
14583        let (rows, cols, b) = (16usize, 64usize, 9usize);
14584        let gpr = cols / GROUP_SIZE;
14585        let mut bytes = vec![0u8; rows * gpr * Q4_TILE];
14586        for r in 0..rows {
14587            for g in 0..gpr {
14588                let t = (r * gpr + g) * Q4_TILE;
14589                let sc = 0.02 + 0.001 * (r * gpr + g) as f32;
14590                bytes[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
14591                for k in 0..16 {
14592                    bytes[t + 2 + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
14593                }
14594            }
14595        }
14596        let x: Vec<f32> = (0..b * cols)
14597            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
14598            .collect();
14599        let mut y_blk = vec![0f32; b * rows];
14600        let mut y_row = vec![0f32; b * rows];
14601        unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
14602        q4t_matmat(&bytes, &x, b, rows, cols, &mut y_blk, None);
14603        unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
14604        q4t_matmat(&bytes, &x, b, rows, cols, &mut y_row, None);
14605        unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
14606        assert_eq!(y_blk, y_row, "q4t blocked 1x4 ≠ per-row");
14607    }
14608
14609    /// The wide-batch Accelerate arm of q4t_matmat vs a brute-force
14610    /// f32 dequant matmul: both are f32 GEMMs, so only reduction
14611    /// order differs — tight tolerance.
14612    /// A synthetic q4tp payload: random nibbles plus a per-row ladder whose
14613    /// span varies row to row, so the codes actually exercise the full 0..31
14614    /// range rather than clustering on one rung.
14615    fn synth_q4tp(rows: usize, cols: usize) -> Vec<u8> {
14616        use cortiq_core::quant::{f32_to_f16, q4tp_code_stride, q4tp_put_code};
14617        let gpr = cols / GROUP_SIZE;
14618        let stride = q4tp_code_stride(gpr);
14619        let (params_off, codes_off, _) = q4tp_sections(rows, cols);
14620        let mut b = vec![0u8; codes_off + rows * stride];
14621        for r in 0..rows {
14622            for g in 0..gpr {
14623                let t = (r * gpr + g) * Q4TP_NIB;
14624                for k in 0..16 {
14625                    b[t + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
14626                }
14627            }
14628            let lo = -6.0 - 0.03 * (r % 17) as f32;
14629            let step = 0.01 + 0.004 * (r % 11) as f32;
14630            let p = params_off + r * 4;
14631            b[p..p + 2].copy_from_slice(&f32_to_f16(lo).to_le_bytes());
14632            b[p + 2..p + 4].copy_from_slice(&f32_to_f16(step).to_le_bytes());
14633            let crow = &mut b[codes_off + r * stride..codes_off + (r + 1) * stride];
14634            for g in 0..gpr {
14635                q4tp_put_code(crow, g, (r * 5 + g * 3) % 32);
14636            }
14637        }
14638        b
14639    }
14640
14641    /// The same weights re-expressed as q4_tiled, so the proven kernel can
14642    /// be the reference: each tile stores the ladder scale its code selects.
14643    /// Only the f16 rounding of that scale separates the two payloads.
14644    fn q4tp_as_q4t(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
14645        let gpr = cols / GROUP_SIZE;
14646        let v = Q4tpView::new(bytes, rows, cols);
14647        let mut out = vec![0u8; rows * gpr * Q4_TILE];
14648        let mut sc = vec![0f32; gpr];
14649        for r in 0..rows {
14650            v.scales_into(r, gpr, &mut sc);
14651            for g in 0..gpr {
14652                let t = (r * gpr + g) * Q4_TILE;
14653                let s = sc[g];
14654                out[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
14655                let src = (r * gpr + g) * Q4TP_NIB;
14656                out[t + 2..t + Q4_TILE].copy_from_slice(&v.nib[src..src + Q4TP_NIB]);
14657            }
14658        }
14659        out
14660    }
14661
14662    /// The exact (`CMF_SDOT=0`) path must reproduce `dequant_q4tp` to f32
14663    /// rounding — that scalar routine is the format's definition, and the
14664    /// kernels re-derive the scale from the ladder independently. Call the
14665    /// row kernel directly: `matmat` picks the int8 arm when a8w8 is on,
14666    /// so routing through it would test the other path by accident.
14667    #[test]
14668    fn q4tp_exact_path_matches_dequant_reference() {
14669        let (rows, cols) = (256usize, 512usize);
14670        let gpr = cols / GROUP_SIZE;
14671        let bytes = synth_q4tp(rows, cols);
14672        let mut w = vec![0f32; rows * cols];
14673        cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
14674
14675        let x: Vec<f32> = (0..cols)
14676            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
14677            .collect();
14678        let v = Q4tpView::new(&bytes, rows, cols);
14679        let mut sc = vec![0f32; gpr];
14680        for r in 0..rows {
14681            v.scales_into(r, gpr, &mut sc);
14682            let got = q4tp_row_exact(v.nib, r, gpr, &x, &sc);
14683            let want: f32 = (0..cols).map(|c| w[r * cols + c] * x[c]).sum();
14684            // These dot products cancel down to ~1e-3 from terms of ~5e-2, so
14685            // the meaningful yardstick is the summed magnitude, not the result:
14686            // against the result any reordering of a 512-term f32 sum "fails".
14687            let mag: f32 = (0..cols).map(|c| (w[r * cols + c] * x[c]).abs()).sum();
14688            assert!(
14689                (got - want).abs() <= 1e-5 * mag,
14690                "row {r}: kernel {got} vs dequant {want}"
14691            );
14692        }
14693    }
14694
14695    /// The int8 (a8w8) path can't be checked against an f32 reference — the
14696    /// activation quantization dominates. Check it against the q4t kernel it
14697    /// was ported from instead, on payloads holding the same weights: that
14698    /// isolates exactly what the port could break (16 B stride, ladder
14699    /// lookup, nibble unpack) from what it deliberately shares.
14700    #[test]
14701    fn q4tp_matvec_matches_the_q4t_kernel_it_was_ported_from() {
14702        let (rows, cols) = (256usize, 512usize);
14703        let bytes = synth_q4tp(rows, cols);
14704        let twin = q4tp_as_q4t(&bytes, rows, cols);
14705        let x: Vec<f32> = (0..cols)
14706            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
14707            .collect();
14708
14709        let mut got = vec![0f32; rows];
14710        q4tp_matvec(&bytes, &x, rows, cols, &mut got, None);
14711        let mut want = vec![0f32; rows];
14712        q4t_matvec(&twin, &x, rows, cols, &mut want, None);
14713
14714        // Scale is f16 in the twin and f32 here, so allow that rounding on
14715        // top of the summed magnitude (same cancellation argument as above).
14716        let mut w = vec![0f32; rows * cols];
14717        cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
14718        for r in 0..rows {
14719            let mag: f32 = (0..cols).map(|c| (w[r * cols + c] * x[c]).abs()).sum();
14720            assert!(
14721                (got[r] - want[r]).abs() <= 1e-3 * mag,
14722                "row {r}: q4tp {} vs q4t {}",
14723                got[r],
14724                want[r]
14725            );
14726        }
14727    }
14728
14729    /// `matmat` carries three arms (Accelerate, blocked int8 1x4, scalar).
14730    /// Batch 5 crosses the blocked kernel's stride, so this exercises the
14731    /// 1x4 path AND its scalar tail in one run — the blocked kernel is new
14732    /// code and its four accumulators are exactly what tends to go wrong.
14733    #[test]
14734    fn q4tp_matmat_matches_the_q4t_kernel_it_was_ported_from() {
14735        let (rows, cols, b) = (256usize, 512usize, 5usize);
14736        let bytes = synth_q4tp(rows, cols);
14737        let twin = q4tp_as_q4t(&bytes, rows, cols);
14738        let xs: Vec<f32> = (0..b * cols)
14739            .map(|i| ((i * 29 + 11) % 89) as f32 / 89.0 - 0.5)
14740            .collect();
14741
14742        let mut got = vec![0f32; b * rows];
14743        q4tp_matmat(&bytes, &xs, b, rows, cols, &mut got, None);
14744        let mut want = vec![0f32; b * rows];
14745        q4t_matmat(&twin, &xs, b, rows, cols, &mut want, None);
14746
14747        let mut w = vec![0f32; rows * cols];
14748        cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
14749        for t in 0..b {
14750            for r in 0..rows {
14751                let mag: f32 = (0..cols)
14752                    .map(|c| (w[r * cols + c] * xs[t * cols + c]).abs())
14753                    .sum();
14754                let (g, wa) = (got[t * rows + r], want[t * rows + r]);
14755                assert!(
14756                    (g - wa).abs() <= 1e-3 * mag,
14757                    "batch {t} row {r}: q4tp {g} vs q4t {wa}"
14758                );
14759            }
14760        }
14761    }
14762
14763    #[test]
14764    fn q4tp_matvec2_matches_the_single_stream_kernel() {
14765        let (rows, cols) = (128usize, 256usize);
14766        let gpr = cols / GROUP_SIZE;
14767        let bytes = synth_q4tp(rows, cols);
14768        let xs: Vec<f32> = (0..2 * cols)
14769            .map(|i| ((i * 29 + 11) % 89) as f32 / 89.0 - 0.5)
14770            .collect();
14771
14772        let (mut o1, mut o2) = (vec![0f32; rows], vec![0f32; rows]);
14773        q4tp_matvec2(
14774            &bytes,
14775            &xs[..cols],
14776            &xs[cols..],
14777            rows,
14778            cols,
14779            &mut o1,
14780            &mut o2,
14781            None,
14782        );
14783
14784        // matvec2 takes the exact path for both streams, so the single-row
14785        // kernel is an exact reference — no tolerance for path differences.
14786        let v = Q4tpView::new(&bytes, rows, cols);
14787        let mut sc = vec![0f32; gpr];
14788        for r in 0..rows {
14789            v.scales_into(r, gpr, &mut sc);
14790            assert_eq!(o1[r], q4tp_row_exact(v.nib, r, gpr, &xs[..cols], &sc));
14791            assert_eq!(o2[r], q4tp_row_exact(v.nib, r, gpr, &xs[cols..], &sc));
14792        }
14793    }
14794
14795    /// q4tp must not COST speed — it exists to save bytes, and a format that
14796    /// trades 7% of a file for a slower model is a bad trade. This guard is
14797    /// here because correctness tests happily passed while `q4tp_matmat` was
14798    /// missing its int8 and Accelerate arms and the model ran 5x slower.
14799    /// Measured on M-series: 0.97-1.04x, i.e. parity (16 B tiles are better
14800    /// aligned than q4t's 18 B, which pays for the scale indirection).
14801    #[test]
14802    fn q4tp_matvec_keeps_pace_with_q4t() {
14803        let (rows, cols) = (4096usize, 3072usize);
14804        let bytes = synth_q4tp(rows, cols);
14805        let twin = q4tp_as_q4t(&bytes, rows, cols);
14806        let x: Vec<f32> = (0..cols).map(|i| (i % 97) as f32 / 97.0 - 0.5).collect();
14807        let mut o = vec![0f32; rows];
14808        let n = 12;
14809        let mut best = (f64::MAX, f64::MAX);
14810        // Interleaved A/B, minimum statistic: this machine throttles, and a
14811        // mean over a thermal ramp reliably indicts whichever ran second.
14812        for _ in 0..3 {
14813            let t0 = std::time::Instant::now();
14814            for _ in 0..n {
14815                q4t_matvec(&twin, &x, rows, cols, &mut o, None);
14816            }
14817            best.0 = best.0.min(t0.elapsed().as_secs_f64());
14818            let t0 = std::time::Instant::now();
14819            for _ in 0..n {
14820                q4tp_matvec(&bytes, &x, rows, cols, &mut o, None);
14821            }
14822            best.1 = best.1.min(t0.elapsed().as_secs_f64());
14823        }
14824        let ratio = best.1 / best.0;
14825        println!(
14826            "q4t {:.3} ms | q4tp {:.3} ms | {ratio:.2}x",
14827            best.0 * 1e3 / n as f64,
14828            best.1 * 1e3 / n as f64
14829        );
14830        assert!(ratio < 2.0, "q4tp matvec {ratio:.2}x slower than q4t");
14831    }
14832
14833    #[cfg(target_os = "macos")]
14834    #[test]
14835    fn q4t_matmat_accel_matches_dequant_reference() {
14836        if !accel_gemm_enabled() {
14837            return; // CMF_ACCEL=0
14838        }
14839        let (rows, cols, b) = (512usize, 1024usize, 8usize); // ≥500K → accel arm
14840        let gpr = cols / GROUP_SIZE;
14841        let mut bytes = vec![0u8; rows * gpr * Q4_TILE];
14842        for r in 0..rows {
14843            for g in 0..gpr {
14844                let t = (r * gpr + g) * Q4_TILE;
14845                let sc = 0.02 + 0.0005 * ((r * gpr + g) % 64) as f32;
14846                bytes[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
14847                for k in 0..16 {
14848                    bytes[t + 2 + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
14849                }
14850            }
14851        }
14852        let x: Vec<f32> = (0..b * cols)
14853            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
14854            .collect();
14855        let mut got = vec![0f32; b * rows];
14856        q4t_matmat(&bytes, &x, b, rows, cols, &mut got, None);
14857        // Brute-force reference off the same tiles.
14858        let mut w = vec![0f32; rows * cols];
14859        for r in 0..rows {
14860            for g in 0..gpr {
14861                let t = (r * gpr + g) * Q4_TILE;
14862                let s = f16_to_f32(u16::from_le_bytes([bytes[t], bytes[t + 1]]));
14863                for (k, &bb) in bytes[t + 2..t + Q4_TILE].iter().enumerate() {
14864                    w[r * cols + g * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s;
14865                    w[r * cols + g * GROUP_SIZE + k * 2 + 1] =
14866                        (((bb >> 4) & 0x0F) as f32 - 8.0) * s;
14867                }
14868            }
14869        }
14870        for bi in 0..b {
14871            for r in 0..rows {
14872                let want: f32 = (0..cols).map(|j| x[bi * cols + j] * w[r * cols + j]).sum();
14873                let d = (got[bi * rows + r] - want).abs();
14874                assert!(
14875                    d <= want.abs().max(1.0) * 1e-4,
14876                    "accel q4t GEMM diverged at ({bi},{r}): {} vs {want}",
14877                    got[bi * rows + r]
14878                );
14879            }
14880        }
14881    }
14882
14883    #[test]
14884    fn q4matvec_matches_full_dequant() {
14885        let (rows, cols) = (8, 64);
14886        let groups = rows * cols / GROUP_SIZE;
14887        // Hand-craft a q4_block blob: nibbles then f16 scales.
14888        let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
14889        for i in 0..groups * 16 {
14890            bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
14891        }
14892        for g in 0..groups {
14893            let s = 0.01 + 0.003 * g as f32;
14894            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
14895        }
14896        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
14897
14898        let mut reference = vec![0.0f32; rows * cols];
14899        cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
14900        let mut expect = vec![0.0f32; rows];
14901        for r in 0..rows {
14902            expect[r] = reference[r * cols..(r + 1) * cols]
14903                .iter()
14904                .zip(&x)
14905                .map(|(w, xv)| w * xv)
14906                .sum();
14907        }
14908
14909        let mut got = vec![0.0f32; rows];
14910        q4matvec(&bytes, &x, rows, cols, &mut got, None);
14911        // SDOT path quantizes activations to i8 (A8W8): bounded noise,
14912        // same contract as q8/vbit (exact path is pinned by CMF_SDOT=0
14913        // in the golden-parity gate).
14914        let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
14915        let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
14916        for r in 0..rows {
14917            assert!(
14918                (got[r] - expect[r]).abs() < tol * scale,
14919                "row {r}: {} vs {}",
14920                got[r],
14921                expect[r]
14922            );
14923        }
14924    }
14925
14926    /// Fused two-input vbit matvec must equal two single matvecs exactly
14927    /// (same per-lane accumulation order on both scalar and SDOT paths).
14928    #[test]
14929    fn vbitmatvec2_equals_two_singles() {
14930        let (rows, cols) = (6, 64);
14931        let ng = cols / GROUP_SIZE;
14932        let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
14933        let mut bytes = bits.clone();
14934        for g in 0..rows * ng {
14935            let s = 0.02 + 0.001 * g as f32;
14936            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
14937        }
14938        for r in 0..rows {
14939            let b = bits[r] as usize;
14940            let (mut acc, mut nb) = (0u64, 0usize);
14941            let mut rowbytes = Vec::new();
14942            for i in 0..cols {
14943                let v = ((i * 7 + r * 13) % (1 << b)) as u64;
14944                acc = (acc << b) | v;
14945                nb += b;
14946                while nb >= 8 {
14947                    nb -= 8;
14948                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
14949                }
14950            }
14951            if nb > 0 {
14952                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
14953            }
14954            bytes.extend_from_slice(&rowbytes);
14955        }
14956        let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
14957        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).cos()).collect();
14958        let offsets = vbit_row_offsets(&bytes, rows, cols);
14959
14960        let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
14961        vbitmatvec(&bytes, &offsets, &x1, rows, cols, &mut a1, None);
14962        vbitmatvec(&bytes, &offsets, &x2, rows, cols, &mut a2, None);
14963        let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
14964        vbitmatvec2(
14965            &bytes, &offsets, &x1, &x2, rows, cols, &mut b1, &mut b2, None,
14966        );
14967        assert_eq!(a1, b1, "fused vbit lane 1 must be bit-identical");
14968        assert_eq!(a2, b2, "fused vbit lane 2 must be bit-identical");
14969    }
14970
14971    /// Fused two-input q4 matvec must equal two single matvecs exactly.
14972    #[test]
14973    fn q4matvec2_equals_two_singles() {
14974        let (rows, cols) = (8, 128);
14975        let groups = rows * cols / GROUP_SIZE;
14976        let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
14977        for i in 0..groups * 16 {
14978            bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
14979        }
14980        for g in 0..groups {
14981            let s = 0.01 + 0.003 * g as f32;
14982            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
14983        }
14984        // Include an outlier channel so the SDOT correction path is
14985        // exercised in the pair kernel too.
14986        let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
14987        x1[9] = 250.0;
14988        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
14989
14990        let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
14991        q4matvec(&bytes, &x1, rows, cols, &mut a1, None);
14992        q4matvec(&bytes, &x2, rows, cols, &mut a2, None);
14993        let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
14994        q4matvec2(&bytes, &x1, &x2, rows, cols, &mut b1, &mut b2, None);
14995        assert_eq!(a1, b1, "fused q4 lane 1 must be bit-identical");
14996        assert_eq!(a2, b2, "fused q4 lane 2 must be bit-identical");
14997    }
14998
14999    /// Multi-matrix job must equal separate matvecs exactly — same
15000    /// kernels, only the dispatch is fused.
15001    #[test]
15002    fn matvec_many_equals_separate_matvecs() {
15003        use crate::pool::Pool;
15004        let (r1, r2, cols) = (300, 200, 64);
15005        let mk = |salt: usize, rows: usize| {
15006            QTensor::from_f32(
15007                (0..rows * cols)
15008                    .map(|i| ((i * 7 + salt) % 97) as f32 / 97.0 - 0.5)
15009                    .collect(),
15010                rows,
15011                cols,
15012            )
15013        };
15014        let (a, b) = (mk(1, r1), mk(5, r2));
15015        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).sin()).collect();
15016        let pool = Pool::new(3);
15017
15018        let (mut ea, mut eb) = (vec![0f32; r1], vec![0f32; r2]);
15019        a.matvec(&x, &mut ea, Some(&pool));
15020        b.matvec(&x, &mut eb, Some(&pool));
15021        let (mut ga, mut gb) = (vec![0f32; r1], vec![0f32; r2]);
15022        QTensor::matvec_many([&a, &b], &x, [&mut ga, &mut gb], Some(&pool));
15023        assert_eq!(ea, ga, "fused multi-matrix lane 1 must be bit-identical");
15024        assert_eq!(eb, gb, "fused multi-matrix lane 2 must be bit-identical");
15025    }
15026
15027    /// The public Q4TP operator must take the real mapped matvec_many arm,
15028    /// rather than the F32 fallback above.  Build a tiny valid CMF so both
15029    /// handles retain their mmap payloads, then compare the fused dispatch
15030    /// with two ordinary mapped matvec calls bit-for-bit.
15031    #[test]
15032    fn q4tp_matvec_many_equals_separate_matvecs() {
15033        use crate::pool::Pool;
15034        use cortiq_core::{CMF_VERSION, CmfHeader, CmfModel, QuantType, TensorSpec};
15035
15036        let (r1, r2, cols) = (300usize, 200usize, 64usize);
15037        let arch: cortiq_core::ModelArch = serde_json::from_value(serde_json::json!({
15038            "arch_name": "tiny-q4tp",
15039            "hidden_size": cols,
15040            "intermediate_size": cols * 2,
15041            "num_layers": 1,
15042            "num_attention_heads": 2,
15043            "num_kv_heads": 1,
15044            "head_dim": 32,
15045            "vocab_size": r1,
15046            "layer_types": ["FullAttention"],
15047            "rms_norm_eps": 1e-6,
15048            "max_position_embeddings": 8,
15049            "linear_conv_kernel_dim": 0,
15050            "linear_num_key_heads": 0,
15051            "linear_num_value_heads": 0
15052        }))
15053        .unwrap();
15054        let header = CmfHeader {
15055            format: "cmf".into(),
15056            version: CMF_VERSION,
15057            arch,
15058            quant_type: QuantType::Q4Block,
15059            provenance: None,
15060            tokenizer_config: None,
15061            section_hashes: None,
15062            skills: Vec::new(),
15063            shard: None,
15064            calibration: None,
15065            routing: None,
15066            genome: None,
15067            lineage: Vec::new(),
15068            router: None,
15069            segments: Vec::new(),
15070        };
15071        let specs = [
15072            TensorSpec {
15073                name: "q".into(),
15074                dtype: TensorDtype::Q4TiledP,
15075                shape: vec![r1, cols],
15076                data: synth_q4tp(r1, cols),
15077            },
15078            TensorSpec {
15079                name: "kv".into(),
15080                dtype: TensorDtype::Q4TiledP,
15081                shape: vec![r2, cols],
15082                data: synth_q4tp(r2, cols),
15083            },
15084        ];
15085        let dir = std::env::temp_dir().join(format!("cmf-q4tp-many-{}", std::process::id()));
15086        std::fs::create_dir_all(&dir).unwrap();
15087        let path = dir.join("m.cmf");
15088        CmfModel::write(&path, &header, &specs, None, None).unwrap();
15089        let model = Arc::new(CmfModel::open(&path).unwrap());
15090        let (a, b) = (
15091            QTensor::from_model(&model, "q").unwrap(),
15092            QTensor::from_model(&model, "kv").unwrap(),
15093        );
15094        assert_eq!(a.model_dtype(), Some(TensorDtype::Q4TiledP));
15095        assert_eq!(b.model_dtype(), Some(TensorDtype::Q4TiledP));
15096        let x: Vec<f32> = (0..cols)
15097            .map(|i| ((i * 17 + 3) % 97) as f32 / 97.0 - 0.5)
15098            .collect();
15099        let pool = Pool::new(3);
15100        let (mut ea, mut eb) = (vec![0.0f32; r1], vec![0.0f32; r2]);
15101        a.matvec(&x, &mut ea, Some(&pool));
15102        b.matvec(&x, &mut eb, Some(&pool));
15103        let (mut ga, mut gb) = (vec![0.0f32; r1], vec![0.0f32; r2]);
15104        QTensor::matvec_many([&a, &b], &x, [&mut ga, &mut gb], Some(&pool));
15105        assert_eq!(ea, ga, "Q4TP fused lane 1 must be bit-identical");
15106        assert_eq!(eb, gb, "Q4TP fused lane 2 must be bit-identical");
15107        let _ = std::fs::remove_dir_all(&dir);
15108    }
15109
15110    /// The MiMo speculative verify's kernels: several tokens' MoE through
15111    /// `moe_gate_up_rows` / `moe_down_rows` (+ the caller's route-order sum)
15112    /// is bit-identical to each token alone through `moe_gate_up_many` /
15113    /// `moe_down_many` (decode), and a row-exact `q4tp_matmat` of five
15114    /// tokens (wide enough for the blocked tiles) equals five matvecs.
15115    #[test]
15116    fn multi_token_moe_rows_equal_single_token_decode() {
15117        use crate::pool::Pool;
15118        use cortiq_core::{CMF_VERSION, CmfHeader, CmfModel, QuantType, TensorSpec};
15119
15120        let (h, inter, ne) = (64usize, 128usize, 3usize);
15121        let arch: cortiq_core::ModelArch = serde_json::from_value(serde_json::json!({
15122            "arch_name": "tiny-q4tp-moe",
15123            "hidden_size": h,
15124            "intermediate_size": inter,
15125            "num_layers": 1,
15126            "num_attention_heads": 2,
15127            "num_kv_heads": 1,
15128            "head_dim": 32,
15129            "vocab_size": 8,
15130            "layer_types": ["FullAttention"],
15131            "rms_norm_eps": 1e-6,
15132            "max_position_embeddings": 8,
15133            "linear_conv_kernel_dim": 0,
15134            "linear_num_key_heads": 0,
15135            "linear_num_value_heads": 0
15136        }))
15137        .unwrap();
15138        let header = CmfHeader {
15139            format: "cmf".into(),
15140            version: CMF_VERSION,
15141            arch,
15142            quant_type: QuantType::Q4Block,
15143            provenance: None,
15144            tokenizer_config: None,
15145            section_hashes: None,
15146            skills: Vec::new(),
15147            shard: None,
15148            calibration: None,
15149            routing: None,
15150            genome: None,
15151            lineage: Vec::new(),
15152            router: None,
15153            segments: Vec::new(),
15154        };
15155        let mut specs = Vec::new();
15156        for e in 0..ne {
15157            for (k, (n, r, c)) in [("g", inter, h), ("u", inter, h), ("d", h, inter)]
15158                .into_iter()
15159                .enumerate()
15160            {
15161                // Distinct experts: perturb only the nibble plane (any byte
15162                // is a valid pair of codes; the ladder stays intact).
15163                let mut data = synth_q4tp(r, c);
15164                for (i, byte) in data[..r * (c / GROUP_SIZE) * Q4TP_NIB]
15165                    .iter_mut()
15166                    .enumerate()
15167                {
15168                    *byte ^= ((i * (e * 3 + k + 1)) % 251) as u8;
15169                }
15170                specs.push(TensorSpec {
15171                    name: format!("{n}{e}"),
15172                    dtype: TensorDtype::Q4TiledP,
15173                    shape: vec![r, c],
15174                    data,
15175                });
15176            }
15177        }
15178        let dir = std::env::temp_dir().join(format!(
15179            "cmf-moe-rows-{}-{}",
15180            std::process::id(),
15181            FLOAT_ACTIVATIONS.get()
15182        ));
15183        std::fs::create_dir_all(&dir).unwrap();
15184        let path = dir.join("m.cmf");
15185        CmfModel::write(&path, &header, &specs, None, None).unwrap();
15186        let model = Arc::new(CmfModel::open(&path).unwrap());
15187        let t = |n: String| QTensor::from_model(&model, &n).unwrap();
15188        let g: Vec<QTensor> = (0..ne).map(|e| t(format!("g{e}"))).collect();
15189        let u: Vec<QTensor> = (0..ne).map(|e| t(format!("u{e}"))).collect();
15190        let d: Vec<QTensor> = (0..ne).map(|e| t(format!("d{e}"))).collect();
15191        let b = 4usize;
15192        let mut xs: Vec<f32> = (0..b * h)
15193            .map(|i| ((i * 31 + 7) % 89) as f32 / 89.0 - 0.5)
15194            .collect();
15195        xs[5] = 9.0; // an activation outlier on token 0
15196        // Token -> (experts in route order, weights).
15197        let routes: Vec<(Vec<usize>, Vec<f32>)> = vec![
15198            (vec![2, 0], vec![0.6, 0.4]),
15199            (vec![0, 1, 2], vec![0.2, 0.5, 0.3]),
15200            (vec![1], vec![1.0]),
15201            (vec![2, 1, 0], vec![0.25, 0.25, 0.5]),
15202        ];
15203        let pool = Pool::new(3);
15204        // Decode reference, token by token.
15205        let mut want = vec![0f32; b * h];
15206        for (tk, (idx, w)) in routes.iter().enumerate() {
15207            let x = &xs[tk * h..(tk + 1) * h];
15208            let pairs: Vec<(&QTensor, &QTensor)> = idx.iter().map(|&e| (&g[e], &u[e])).collect();
15209            let mut gs: Vec<Vec<f32>> = idx.iter().map(|_| vec![0f32; inter]).collect();
15210            assert!(QTensor::moe_gate_up_many(&pairs, x, &mut gs, Some(&pool)));
15211            if FLOAT_ACTIVATIONS.get() {
15212                for (slot, &e) in idx.iter().enumerate() {
15213                    let (mut gate, mut up) = (vec![0.0; inter], vec![0.0; inter]);
15214                    g[e].matvec(x, &mut gate, Some(&pool));
15215                    u[e].matvec(x, &mut up, Some(&pool));
15216                    for (v, u) in gate.iter_mut().zip(up) {
15217                        *v = (*v / (1.0 + (-*v).exp())) * u;
15218                    }
15219                    assert_eq!(gs[slot], gate, "float gate/up must equal ordinary matvecs");
15220                }
15221            }
15222            let downs: Vec<&QTensor> = idx.iter().map(|&e| &d[e]).collect();
15223            assert!(QTensor::moe_down_many(
15224                &downs,
15225                &gs,
15226                w,
15227                &mut want[tk * h..(tk + 1) * h],
15228                Some(&pool)
15229            ));
15230        }
15231        if FLOAT_ACTIVATIONS.get() {
15232            for (tk, (idx, w)) in routes.iter().enumerate() {
15233                let mut scalar = vec![0.0; h];
15234                for (&e, &weight) in idx.iter().zip(w) {
15235                    let (mut gate, mut up, mut down) =
15236                        (vec![0.0; inter], vec![0.0; inter], vec![0.0; h]);
15237                    g[e].matvec(&xs[tk * h..(tk + 1) * h], &mut gate, Some(&pool));
15238                    u[e].matvec(&xs[tk * h..(tk + 1) * h], &mut up, Some(&pool));
15239                    for (v, u) in gate.iter_mut().zip(up) {
15240                        *v = (*v / (1.0 + (-*v).exp())) * u;
15241                    }
15242                    d[e].matvec(&gate, &mut down, Some(&pool));
15243                    for (v, d) in scalar.iter_mut().zip(down) {
15244                        *v += weight * d;
15245                    }
15246                }
15247                assert_eq!(
15248                    &want[tk * h..(tk + 1) * h],
15249                    scalar,
15250                    "float many equals scalar experts"
15251                );
15252            }
15253        }
15254        // All four tokens at once, grouped by expert.
15255        let mut experts: Vec<usize> = Vec::new();
15256        let mut groups: Vec<Vec<usize>> = Vec::new();
15257        for (tk, (idx, _)) in routes.iter().enumerate() {
15258            for &e in idx {
15259                match experts.iter().position(|&x| x == e) {
15260                    Some(k) => groups[k].push(tk),
15261                    None => {
15262                        experts.push(e);
15263                        groups.push(vec![tk]);
15264                    }
15265                }
15266            }
15267        }
15268        let n_pairs: usize = groups.iter().map(|g| g.len()).sum();
15269        let pairs: Vec<(&QTensor, &QTensor)> = experts.iter().map(|&e| (&g[e], &u[e])).collect();
15270        let mut gs: Vec<Vec<f32>> = (0..n_pairs).map(|_| vec![0f32; inter]).collect();
15271        assert!(QTensor::moe_gate_up_rows(
15272            &pairs,
15273            &groups,
15274            &xs,
15275            &mut gs,
15276            Some(&pool)
15277        ));
15278        let downs: Vec<&QTensor> = experts.iter().map(|&e| &d[e]).collect();
15279        let lens: Vec<usize> = groups.iter().map(|g| g.len()).collect();
15280        let mut ds: Vec<Vec<f32>> = (0..n_pairs).map(|_| vec![0f32; h]).collect();
15281        assert!(QTensor::moe_down_rows(
15282            &downs,
15283            &lens,
15284            &gs,
15285            &mut ds,
15286            Some(&pool)
15287        ));
15288        let slot = |tk: usize, e: usize| {
15289            let k = experts.iter().position(|&x| x == e).unwrap();
15290            groups[..k].iter().map(|g| g.len()).sum::<usize>()
15291                + groups[k].iter().position(|&x| x == tk).unwrap()
15292        };
15293        let mut got = vec![0f32; b * h];
15294        for (tk, (idx, w)) in routes.iter().enumerate() {
15295            for i in 0..h {
15296                let mut acc = 0f32;
15297                for (&e, &we) in idx.iter().zip(w) {
15298                    acc += we * ds[slot(tk, e)][i];
15299                }
15300                got[tk * h + i] = acc;
15301            }
15302        }
15303        assert!(want.iter().any(|v| *v != 0.0));
15304        assert_eq!(
15305            want.iter().map(|v| v.to_bits()).collect::<Vec<_>>(),
15306            got.iter().map(|v| v.to_bits()).collect::<Vec<_>>(),
15307            "multi-token MoE must equal decode bit for bit"
15308        );
15309
15310        // Row-exact q4tp_matmat: five tokens (a blocked 1x4 tile + tail
15311        // otherwise) equal five matvecs.
15312        let b5 = 5usize;
15313        let x5: Vec<f32> = (0..b5 * h)
15314            .map(|i| ((i * 13 + 5) % 71) as f32 / 71.0 - 0.5)
15315            .collect();
15316        let mut mm = vec![0f32; b5 * inter];
15317        row_exact_scope(|| g[1].matmat(&x5, b5, &mut mm, Some(&pool)));
15318        for tk in 0..b5 {
15319            let mut mv = vec![0f32; inter];
15320            g[1].matvec(&x5[tk * h..(tk + 1) * h], &mut mv, Some(&pool));
15321            assert_eq!(
15322                mv.iter().map(|v| v.to_bits()).collect::<Vec<_>>(),
15323                mm[tk * inter..(tk + 1) * inter]
15324                    .iter()
15325                    .map(|v| v.to_bits())
15326                    .collect::<Vec<_>>(),
15327                "row-exact matmat token {tk}"
15328            );
15329        }
15330        // Other concurrent tests/requests may still hold the shared mode.
15331        // Nested, overlapping and unwind restoration is checked separately.
15332        let _ = std::fs::remove_dir_all(&dir);
15333    }
15334
15335    /// The row-exact fix must not touch the fast path. Outside the scope
15336    /// `q4tp_matmat` has to produce exactly what it did before, and on ARM
15337    /// "before" is spelled out below: the tuned 1x4 SDOT tile for every
15338    /// four columns and the single-row kernel for the tail. Inside the
15339    /// scope every column equals its token's matvec. `q4tp_matmat_with`
15340    /// takes the mode as an argument, so a concurrent test holding the
15341    /// shared scope cannot flip it under this one.
15342    #[test]
15343    fn q4tp_matmat_fast_path_unchanged_outside_row_exact() {
15344        use crate::pool::Pool;
15345        use std::sync::atomic::Ordering::Relaxed;
15346        let _alt = Q4TP_ALT_TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner());
15347        // The tuned ARM shape, which is also what an unset switch picks
15348        // unless CMF_Q4TP_V1 is exported.
15349        Q4TP_ALT.store(2, Relaxed);
15350        let pool = Pool::new(3);
15351        // Under 500k cells, so macOS keeps the matmat off the AMX; the
15352        // second shape runs across pool workers (rows >= 256) with 32
15353        // groups of accumulation and a tail after two 1x4 tiles.
15354        for &(rows, cols, b) in &[(64usize, 256usize, 7usize), (320, 1024, 9)] {
15355            let bytes = synth_q4tp(rows, cols);
15356            let mut xs: Vec<f32> = (0..b * cols)
15357                .map(|i| ((i * 29 + 11) % 83) as f32 / 83.0 - 0.5)
15358                .collect();
15359            xs[3] = 7.5; // an activation outlier on token 0
15360            let run = |exact: bool| {
15361                let mut out = vec![0f32; b * rows];
15362                q4tp_matmat_with(&bytes, &xs, b, rows, cols, &mut out, Some(&pool), exact);
15363                out
15364            };
15365            let (fast, exact) = (run(false), run(true));
15366            #[cfg(not(target_arch = "aarch64"))]
15367            let _ = fast;
15368            let bits = |v: &[f32]| v.iter().map(|x| x.to_bits()).collect::<Vec<_>>();
15369            let mut matvecs = vec![0f32; b * rows];
15370            for (bi, o) in matvecs.chunks_mut(rows).enumerate() {
15371                q4tp_matvec(
15372                    &bytes,
15373                    &xs[bi * cols..(bi + 1) * cols],
15374                    rows,
15375                    cols,
15376                    o,
15377                    Some(&pool),
15378                );
15379            }
15380            assert!(matvecs.iter().any(|v| *v != 0.0));
15381            assert_eq!(
15382                bits(&exact),
15383                bits(&matvecs),
15384                "{rows}x{cols} b={b}: row-exact matmat must equal per-token matvecs"
15385            );
15386            #[cfg(target_arch = "aarch64")]
15387            {
15388                // The pre-fix ARM loop, cell for cell.
15389                let gpr = cols / GROUP_SIZE;
15390                let v = Q4tpView::new(&bytes, rows, cols);
15391                let mut old = vec![0f32; b * rows];
15392                let mut sc = vec![0f32; gpr];
15393                let a8w8 = a8w8_enabled();
15394                let blocked = sdot_enabled() && blocked_enabled();
15395                let acts: Vec<SplitAct> = (0..b)
15396                    .map(|bi| split_act(&xs[bi * cols..(bi + 1) * cols]))
15397                    .collect();
15398                for r in 0..rows {
15399                    v.scales_into(r, gpr, &mut sc);
15400                    if !a8w8 {
15401                        for bi in 0..b {
15402                            let x = &xs[bi * cols..(bi + 1) * cols];
15403                            old[bi * rows + r] = q4tp_row_exact(v.nib, r, gpr, x, &sc);
15404                        }
15405                        continue;
15406                    }
15407                    let finish = |d: f32, act: &SplitAct| {
15408                        let mut acc = d * act.sx;
15409                        for &(j, xv) in &act.outliers {
15410                            let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
15411                            acc += w * s * xv;
15412                        }
15413                        acc
15414                    };
15415                    let mut bi = 0usize;
15416                    while blocked && bi + 4 <= b {
15417                        let xs4 = [
15418                            acts[bi].xq.as_slice(),
15419                            acts[bi + 1].xq.as_slice(),
15420                            acts[bi + 2].xq.as_slice(),
15421                            acts[bi + 3].xq.as_slice(),
15422                        ];
15423                        let d = unsafe { dot_q4tp_row_1x4_sdot(v.nib, r, gpr, xs4, &sc) };
15424                        for k in 0..4 {
15425                            old[(bi + k) * rows + r] = finish(d[k], &acts[bi + k]);
15426                        }
15427                        bi += 4;
15428                    }
15429                    for (bi, act) in acts.iter().enumerate().skip(bi) {
15430                        let d = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, &sc);
15431                        old[bi * rows + r] = finish(d, act);
15432                    }
15433                }
15434                assert_eq!(
15435                    bits(&fast),
15436                    bits(&old),
15437                    "{rows}x{cols} b={b}: the fast path changed outside row_exact"
15438                );
15439                // And it is still the fast tile that runs: its lane-parallel
15440                // fma sum rounds differently from the matvec somewhere.
15441                if blocked {
15442                    assert_ne!(
15443                        bits(&fast),
15444                        bits(&matvecs),
15445                        "{rows}x{cols} b={b}: the tuned tile no longer runs outside row_exact"
15446                    );
15447                }
15448            }
15449        }
15450        Q4TP_ALT.store(0, Relaxed);
15451    }
15452
15453    #[test]
15454    fn row_exact_scopes_survive_overlap_nesting_and_unwind() {
15455        use std::sync::{Barrier, atomic::{AtomicUsize, Ordering}};
15456        // A private counter makes this restoration test independent of
15457        // numerical tests concurrently using the production counter.
15458        let active = AtomicUsize::new(0);
15459        counted_row_exact_scope(&active, || {
15460            assert_eq!(active.load(Ordering::Acquire), 1);
15461            counted_row_exact_scope(&active, || {
15462                assert_eq!(active.load(Ordering::Acquire), 2);
15463            });
15464            assert_eq!(active.load(Ordering::Acquire), 1);
15465        });
15466        assert_eq!(active.load(Ordering::Acquire), 0);
15467
15468        let both_entered = Barrier::new(2);
15469        let release_last = Barrier::new(2);
15470        std::thread::scope(|s| {
15471            let first = s.spawn(|| counted_row_exact_scope(&active, || {
15472                both_entered.wait();
15473            }));
15474            let last = s.spawn(|| counted_row_exact_scope(&active, || {
15475                both_entered.wait();
15476                release_last.wait();
15477            }));
15478            first.join().unwrap();
15479            let after_first = active.load(Ordering::Acquire);
15480            release_last.wait();
15481            last.join().unwrap();
15482            assert_eq!(after_first, 1, "second request must remain exact");
15483        });
15484        assert_eq!(active.load(Ordering::Acquire), 0);
15485        let panic = std::panic::catch_unwind(|| {
15486            counted_row_exact_scope(&active, || panic!("scope unwind"));
15487        });
15488        assert!(panic.is_err());
15489        assert_eq!(active.load(Ordering::Acquire), 0);
15490    }
15491
15492    #[test]
15493    #[cfg(target_arch = "x86_64")]
15494    fn q4tp_float_avx2_is_bitwise_scalar() {
15495        if !avx2_enabled() {
15496            return;
15497        }
15498        for cols in [32, 64, 96, 2048, 4096] {
15499            let rows = 9;
15500            let bytes = synth_q4tp(rows, cols);
15501            let v = Q4tpView::new(&bytes, rows, cols);
15502            let gpr = cols / GROUP_SIZE;
15503            let mut sc = vec![0.0; gpr];
15504            for seed in 1..=5 {
15505                let xs: Vec<f32> = (0..cols)
15506                    .map(|i| (((i * 104729 + seed * 8191) % 100003) as f32 - 50001.0) / 7919.0)
15507                    .collect();
15508                for r in 0..rows {
15509                    v.scales_into(r, gpr, &mut sc);
15510                    let scalar = q4tp_row_float_scalar(v.nib, r, gpr, &xs, &sc);
15511                    let vector = unsafe { q4tp_row_float_avx2(v.nib, r, gpr, &xs, &sc) };
15512                    assert_eq!(
15513                        scalar.to_bits(),
15514                        vector.to_bits(),
15515                        "cols={cols} row={r} seed={seed}"
15516                    );
15517                }
15518            }
15519        }
15520    }
15521
15522    #[test]
15523    fn multi_token_moe_rows_float_equal_single_token_decode() {
15524        float_activations_scope(multi_token_moe_rows_equal_single_token_decode);
15525    }
15526
15527    #[test]
15528    fn full_gpu_q8_scope_is_nested_and_thread_local() {
15529        assert!(!FULL_GPU_Q8.get());
15530        let before = gpu_split_frac();
15531        {
15532            let _guard = enter_full_gpu_q8_scope();
15533            assert_eq!(gpu_split_frac(), 1.0);
15534            {
15535                let _nested = enter_full_gpu_q8_scope();
15536            }
15537            assert_eq!(gpu_split_frac(), 1.0);
15538            std::thread::spawn(|| assert!(!FULL_GPU_Q8.get())).join().unwrap();
15539        }
15540        assert!(!FULL_GPU_Q8.get());
15541        assert_eq!(gpu_split_frac(), before);
15542    }
15543
15544    #[test]
15545    fn float_activation_scope_is_nested_thread_local_and_unwind_safe() {
15546        assert!(!FLOAT_ACTIVATIONS.get());
15547        let before = a8w8_enabled();
15548        float_activations_scope(|| {
15549            assert!(!a8w8_enabled());
15550            float_activations_scope(|| assert!(!a8w8_enabled()));
15551            assert!(FLOAT_ACTIVATIONS.get());
15552            std::thread::spawn(|| assert!(!FLOAT_ACTIVATIONS.get()))
15553                .join()
15554                .unwrap();
15555        });
15556        assert!(!FLOAT_ACTIVATIONS.get());
15557        assert_eq!(a8w8_enabled(), before);
15558        let _ = std::panic::catch_unwind(|| float_activations_scope(|| panic!("test unwind")));
15559        assert!(!FLOAT_ACTIVATIONS.get());
15560    }
15561
15562    /// Batched q4/vbit matmat must equal per-position matvec calls
15563    /// exactly (the fallback it replaced) — same kernels, same order.
15564    #[test]
15565    fn batched_matmat_equals_per_position_matvec() {
15566        let (rows, cols, b) = (8, 64, 5);
15567        // q4 blob.
15568        let groups = rows * cols / GROUP_SIZE;
15569        let mut q4 = Vec::new();
15570        for i in 0..groups * 16 {
15571            q4.push((((i * 7 + 3) % 256) & 0xFF) as u8);
15572        }
15573        for g in 0..groups {
15574            q4.extend_from_slice(
15575                &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
15576            );
15577        }
15578        // vbit blob (mixed widths incl. 8).
15579        let ng = cols / GROUP_SIZE;
15580        let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4, 5, 3];
15581        let mut vb = bits.clone();
15582        for g in 0..rows * ng {
15583            vb.extend_from_slice(
15584                &cortiq_core::quant::f32_to_f16(0.02 + 0.001 * g as f32).to_le_bytes(),
15585            );
15586        }
15587        for r in 0..rows {
15588            let bw = bits[r] as usize;
15589            let (mut acc, mut nb) = (0u64, 0usize);
15590            let mut rowbytes = Vec::new();
15591            for i in 0..cols {
15592                let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
15593                acc = (acc << bw) | v;
15594                nb += bw;
15595                while nb >= 8 {
15596                    nb -= 8;
15597                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
15598                }
15599            }
15600            if nb > 0 {
15601                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
15602            }
15603            vb.extend_from_slice(&rowbytes);
15604        }
15605        let offsets = vbit_row_offsets(&vb, rows, cols);
15606
15607        let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
15608
15609        // q4: batch vs singles.
15610        let mut got = vec![0f32; b * rows];
15611        q4matmat(&q4, &xs, b, rows, cols, &mut got, None);
15612        for bi in 0..b {
15613            let mut expect = vec![0f32; rows];
15614            q4matvec(
15615                &q4,
15616                &xs[bi * cols..(bi + 1) * cols],
15617                rows,
15618                cols,
15619                &mut expect,
15620                None,
15621            );
15622            assert_eq!(
15623                &got[bi * rows..(bi + 1) * rows],
15624                &expect[..],
15625                "q4 batch pos {bi}"
15626            );
15627        }
15628
15629        // vbit: batch vs singles.
15630        let mut got = vec![0f32; b * rows];
15631        vbitmatmat(&vb, &offsets, &xs, b, rows, cols, &mut got, None);
15632        for bi in 0..b {
15633            let mut expect = vec![0f32; rows];
15634            vbitmatvec(
15635                &vb,
15636                &offsets,
15637                &xs[bi * cols..(bi + 1) * cols],
15638                rows,
15639                cols,
15640                &mut expect,
15641                None,
15642            );
15643            assert_eq!(
15644                &got[bi * rows..(bi + 1) * rows],
15645                &expect[..],
15646                "vbit batch pos {bi}"
15647            );
15648        }
15649    }
15650
15651    /// q4_tiled kernels must produce BIT-identical outputs to the q4
15652    /// split kernels on the same values (same ints, same order — only
15653    /// the byte placement differs).
15654    #[test]
15655    fn q4_tiled_matches_q4_block_bitexact() {
15656        let (rows, cols, b) = (8usize, 128usize, 3usize);
15657        let groups = rows * cols / GROUP_SIZE;
15658        let mut split = Vec::with_capacity(groups * 18);
15659        for i in 0..groups * 16 {
15660            split.push((((i * 7 + 3) % 256) & 0xFF) as u8);
15661        }
15662        for g in 0..groups {
15663            split.extend_from_slice(
15664                &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
15665            );
15666        }
15667        // Re-tile: [scale][nibbles] per group.
15668        let (packed, scales) = split.split_at(groups * 16);
15669        let mut tiled = Vec::with_capacity(groups * Q4_TILE);
15670        for g in 0..groups {
15671            tiled.extend_from_slice(&scales[g * 2..g * 2 + 2]);
15672            tiled.extend_from_slice(&packed[g * 16..(g + 1) * 16]);
15673        }
15674
15675        let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
15676        x1[9] = 250.0; // exercise the outlier path
15677        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
15678
15679        let (mut a, mut t) = (vec![0f32; rows], vec![0f32; rows]);
15680        q4matvec(&split, &x1, rows, cols, &mut a, None);
15681        q4t_matvec(&tiled, &x1, rows, cols, &mut t, None);
15682        assert_eq!(a, t, "q4t matvec must match q4 bit-for-bit");
15683
15684        let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
15685        let (mut t1, mut t2) = (vec![0f32; rows], vec![0f32; rows]);
15686        q4matvec2(&split, &x1, &x2, rows, cols, &mut a1, &mut a2, None);
15687        q4t_matvec2(&tiled, &x1, &x2, rows, cols, &mut t1, &mut t2, None);
15688        assert_eq!(a1, t1);
15689        assert_eq!(a2, t2);
15690
15691        let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
15692        let (mut am, mut tm) = (vec![0f32; b * rows], vec![0f32; b * rows]);
15693        q4matmat(&split, &xs, b, rows, cols, &mut am, None);
15694        q4t_matmat(&tiled, &xs, b, rows, cols, &mut tm, None);
15695        assert_eq!(am, tm, "q4t matmat must match q4 bit-for-bit");
15696    }
15697
15698    /// q4 SDOT outlier correction: a single huge activation channel
15699    /// (>8·rms → outlier, zeroed in xq) must still contribute its EXACT
15700    /// term. On-grid bulk (±1/0 → xq dequantizes exactly) isolates the
15701    /// correction from A8W8 noise. cols must exceed 64: at n=64 the
15702    /// 8·rms threshold equals sqrt(v²+rest) ≥ v, so a single outlier
15703    /// can never qualify (8² = n).
15704    #[test]
15705    fn q4matvec_sdot_outlier_exact() {
15706        let (rows, cols) = (4, 128);
15707        let groups = rows * cols / GROUP_SIZE;
15708        let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
15709        for i in 0..groups * 16 {
15710            bytes.push(((i * 11 + 5) % 256) as u8);
15711        }
15712        for g in 0..groups {
15713            let s = 0.02 + 0.002 * g as f32;
15714            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
15715        }
15716        let mut x: Vec<f32> = (0..cols)
15717            .map(|i| match i % 3 {
15718                0 => 1.0,
15719                1 => -1.0,
15720                _ => 0.0,
15721            })
15722            .collect();
15723        x[17] = 300.0; // ≫ 8·rms → outlier channel
15724
15725        let mut reference = vec![0.0f32; rows * cols];
15726        cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
15727        let mut expect = vec![0.0f32; rows];
15728        for r in 0..rows {
15729            expect[r] = reference[r * cols..(r + 1) * cols]
15730                .iter()
15731                .zip(&x)
15732                .map(|(w, xv)| w * xv)
15733                .sum();
15734        }
15735        let mut got = vec![0.0f32; rows];
15736        q4matvec(&bytes, &x, rows, cols, &mut got, None);
15737        let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
15738        for r in 0..rows {
15739            assert!(
15740                (got[r] - expect[r]).abs() < 2e-3 * scale,
15741                "row {r}: {} vs {} (outlier term must be exact)",
15742                got[r],
15743                expect[r]
15744            );
15745        }
15746    }
15747
15748    /// The fused q1t matvec must equal the reference (dequant_q1t → dot),
15749    /// including the ternary zero level and the binary-searched outlier
15750    /// overlay. Guards the mmap kernel that makes a 12B q1t runnable.
15751    #[test]
15752    fn q1t_matvec_matches_reference() {
15753        use cortiq_core::quant::{dequant_q1t, f32_to_f16};
15754        let (rows, cols) = (3usize, 64usize); // gpr = 2
15755        let gpr = cols / GROUP_SIZE;
15756        let scales = [0.5f32, 0.3, 0.7, 0.2, 0.6, 0.15];
15757        // Overlay (must be sorted by flat index): a few spikes across rows.
15758        let outliers: [(u32, f32); 3] = [(5, 9.0), (70, -4.5), (150, 3.25)];
15759        let is_out = |flat: usize| outliers.iter().any(|&(i, _)| i as usize == flat);
15760        let mut bytes = Vec::new();
15761        for r in 0..rows {
15762            for g in 0..gpr {
15763                bytes.extend_from_slice(&f32_to_f16(scales[r * gpr + g]).to_le_bytes());
15764                let mut c = [0u8; 7];
15765                for k in 0..GROUP_SIZE {
15766                    // Encoder invariant: code 0 at outlier positions.
15767                    let code = if is_out(r * cols + g * GROUP_SIZE + k) {
15768                        0
15769                    } else {
15770                        ((k + r * 3 + g) % 3) as u8 // 0,1,2
15771                    };
15772                    cortiq_core::quant::q1t_pack(&mut c, k, code);
15773                }
15774                bytes.extend_from_slice(&c);
15775            }
15776        }
15777        // Per-row overlay: [u32 row_ptr[rows+1]] then [(u16 col, f16 val)] by
15778        // row (outliers are sorted by flat index → already grouped by row).
15779        let mut row_ptr = vec![0u32; rows + 1];
15780        for &(idx, _) in &outliers {
15781            row_ptr[idx as usize / cols + 1] += 1;
15782        }
15783        for r in 0..rows {
15784            row_ptr[r + 1] += row_ptr[r];
15785        }
15786        for &p in &row_ptr {
15787            bytes.extend_from_slice(&p.to_le_bytes());
15788        }
15789        for &(idx, v) in &outliers {
15790            bytes.extend_from_slice(&((idx as usize % cols) as u16).to_le_bytes());
15791            bytes.extend_from_slice(&f32_to_f16(v).to_le_bytes());
15792        }
15793
15794        let mut refw = vec![0f32; rows * cols];
15795        dequant_q1t(&bytes, rows, cols, &mut refw);
15796        // On-grid activations (±1, amax 1) so the int8 SDOT path reconstructs
15797        // x exactly and matches the f32 reference (same trick as the q1 test).
15798        let x: Vec<f32> = (0..cols)
15799            .map(|j| if j % 3 == 0 { 1.0 } else { -1.0 })
15800            .collect();
15801        let mut expect = vec![0f32; rows];
15802        for r in 0..rows {
15803            let mut a = 0.0f32;
15804            for j in 0..cols {
15805                a += refw[r * cols + j] * x[j];
15806            }
15807            expect[r] = a;
15808        }
15809        let tol = |e: f32| 1e-3 * e.abs().max(1e-3);
15810        let mut got = vec![0f32; rows];
15811        q1t_matvec(&bytes, &x, rows, cols, &mut got, None);
15812        for r in 0..rows {
15813            assert!(
15814                (got[r] - expect[r]).abs() < tol(expect[r]),
15815                "row {r}: {} vs {}",
15816                got[r],
15817                expect[r]
15818            );
15819        }
15820        // matmat (b=2, f32 decode path) must agree too.
15821        let x2: Vec<f32> = x.iter().chain(x.iter()).copied().collect();
15822        let mut gm = vec![0f32; 2 * rows];
15823        q1t_matmat(&bytes, &x2, 2, rows, cols, &mut gm, None);
15824        for r in 0..rows {
15825            assert!((gm[r] - expect[r]).abs() < tol(expect[r]));
15826            assert!((gm[rows + r] - expect[r]).abs() < tol(expect[r]));
15827        }
15828        // Fused pair (q1t_matvec2) must equal two single matvecs
15829        // bit-for-bit: same unpack, same group order, same f32
15830        // accumulation per stream. Distinct x2 exercises both lanes.
15831        let xb: Vec<f32> = (0..cols)
15832            .map(|j| if j % 5 == 0 { -1.0 } else { 1.0 })
15833            .collect();
15834        let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
15835        q1t_matvec(&bytes, &x, rows, cols, &mut s1, None);
15836        q1t_matvec(&bytes, &xb, rows, cols, &mut s2, None);
15837        let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
15838        q1t_matvec2(&bytes, &x, &xb, rows, cols, &mut p1, &mut p2, None);
15839        assert_eq!(p1, s1, "q1t pair lane 1 ≠ single matvec");
15840        assert_eq!(p2, s2, "q1t pair lane 2 ≠ single matvec");
15841    }
15842
15843    /// Pair == 2×matvec with an ODD group count (the kernel's tail
15844    /// group) and no overlay section.
15845    #[test]
15846    fn q1t_matvec2_odd_gpr_matches_singles() {
15847        use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_pack};
15848        let (rows, cols) = (5usize, 96usize); // gpr = 3 → paired + tail
15849        let gpr = cols / GROUP_SIZE;
15850        let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE);
15851        for r in 0..rows {
15852            for g in 0..gpr {
15853                bytes.extend_from_slice(&f32_to_f16(0.1 + 0.05 * (r + g) as f32).to_le_bytes());
15854                let mut c = [0u8; 7];
15855                for k in 0..GROUP_SIZE {
15856                    q1t_pack(&mut c, k, ((k * 7 + r * 5 + g * 3) % 3) as u8);
15857                }
15858                bytes.extend_from_slice(&c);
15859            }
15860        }
15861        let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.31).sin()).collect();
15862        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).cos()).collect();
15863        let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
15864        q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
15865        q1t_matvec(&bytes, &x2, rows, cols, &mut s2, None);
15866        let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
15867        q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
15868        assert_eq!(p1, s1, "odd-gpr pair lane 1 ≠ single");
15869        assert_eq!(p2, s2, "odd-gpr pair lane 2 ≠ single");
15870    }
15871
15872    // Speed A/B: fused pair (one unpack, two streams) vs two single
15873    // matvecs. Single-threaded, FFN-sized, min-of paired in-process.
15874    //   cargo test -p cortiq-engine --release q1t_matvec2_speed -- --ignored --nocapture
15875    #[test]
15876    #[ignore]
15877    fn q1t_matvec2_speed() {
15878        use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_pack};
15879        use std::time::Instant;
15880        let (rows, cols) = (8192usize, 4096usize);
15881        let gpr = cols / GROUP_SIZE;
15882        let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE);
15883        for r in 0..rows {
15884            for g in 0..gpr {
15885                let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
15886                bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
15887                let mut c = [0u8; 7];
15888                for k in 0..GROUP_SIZE {
15889                    q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
15890                }
15891                bytes.extend_from_slice(&c);
15892            }
15893        }
15894        let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.31).sin()).collect();
15895        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).cos()).collect();
15896        let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
15897        let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
15898        // Warm both paths once.
15899        q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
15900        q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
15901        let (mut t_pair, mut t_two) = (f64::MAX, f64::MAX);
15902        for _ in 0..8 {
15903            let t0 = Instant::now();
15904            q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
15905            t_pair = t_pair.min(t0.elapsed().as_secs_f64() * 1000.0);
15906            let t1 = Instant::now();
15907            q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
15908            q1t_matvec(&bytes, &x2, rows, cols, &mut s2, None);
15909            t_two = t_two.min(t1.elapsed().as_secs_f64() * 1000.0);
15910        }
15911        assert_eq!(p1, s1);
15912        assert_eq!(p2, s2);
15913        println!("q1t pair {rows}x{cols}: fused {t_pair:.2} ms | two singles {t_two:.2} ms");
15914    }
15915
15916    // Speed A/B: the base-3-division decode (what the packing commit left in
15917    // place) vs the fused sign-LUT matvec. Both single-threaded, same bytes.
15918    //   cargo test -p cortiq-engine q1t_matvec_speed -- --ignored --nocapture
15919    #[test]
15920    #[ignore]
15921    fn q1t_matvec_speed() {
15922        use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_code, q1t_pack};
15923        use std::time::Instant;
15924        let (rows, cols) = (8192usize, 4096usize); // FFN-sized
15925        let gpr = cols / GROUP_SIZE;
15926        let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE + 16);
15927        for r in 0..rows {
15928            for g in 0..gpr {
15929                let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
15930                bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
15931                let mut c = [0u8; 7];
15932                for k in 0..GROUP_SIZE {
15933                    q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
15934                }
15935                bytes.extend_from_slice(&c);
15936            }
15937        }
15938        let (n, stride) = (rows * cols, 40usize); // ~2.5% outliers, per-row overlay
15939        let mut row_ptr = vec![0u32; rows + 1];
15940        let mut idx = 0usize;
15941        while idx < n {
15942            row_ptr[idx / cols + 1] += 1;
15943            idx += stride;
15944        }
15945        for r in 0..rows {
15946            row_ptr[r + 1] += row_ptr[r];
15947        }
15948        for &p in &row_ptr {
15949            bytes.extend_from_slice(&p.to_le_bytes());
15950        }
15951        let mut idx = 0usize;
15952        while idx < n {
15953            bytes.extend_from_slice(&((idx % cols) as u16).to_le_bytes());
15954            bytes.extend_from_slice(&f32_to_f16((idx % 13) as f32 * 0.1 - 0.6).to_le_bytes());
15955            idx += stride;
15956        }
15957        // On-grid ±1 so the fast path's int8 SDOT is exact vs the f32 "slow"
15958        // reference (the A/B is a timing check; values must still agree).
15959        let x: Vec<f32> = (0..cols)
15960            .map(|j| if j % 3 == 0 { 1.0 } else { -1.0 })
15961            .collect();
15962        let (rp_off, ent_off, has_ov) = q1t_overlay(&bytes, rows * gpr * Q1T_TILE, rows);
15963
15964        // "before": base-3 division decode into a buffer, then dot.
15965        let slow = |out: &mut [f32]| {
15966            let mut buf = vec![0f32; cols];
15967            for r in 0..rows {
15968                for g in 0..gpr {
15969                    let off = (r * gpr + g) * Q1T_TILE;
15970                    let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
15971                    let codes = &bytes[off + 2..off + Q1T_TILE];
15972                    for k in 0..GROUP_SIZE {
15973                        buf[g * GROUP_SIZE + k] = match q1t_code(codes, k) {
15974                            1 => s,
15975                            2 => -s,
15976                            _ => 0.0,
15977                        };
15978                    }
15979                }
15980                out[r] = q1t_row_outlier_correction(&bytes, r, rp_off, ent_off, has_ov, &x)
15981                    + (0..cols).map(|j| buf[j] * x[j]).sum::<f32>();
15982            }
15983        };
15984        let iters = 5;
15985        let mut a = vec![0f32; rows];
15986        slow(&mut a); // warm
15987        let t = Instant::now();
15988        for _ in 0..iters {
15989            slow(&mut a);
15990        }
15991        let slow_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
15992
15993        let mut b = vec![0f32; rows];
15994        q1t_matvec(&bytes, &x, rows, cols, &mut b, None); // warm
15995        let t = Instant::now();
15996        for _ in 0..iters {
15997            q1t_matvec(&bytes, &x, rows, cols, &mut b, None);
15998        }
15999        let fast_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
16000
16001        for r in 0..rows {
16002            assert!((a[r] - b[r]).abs() < 1e-2, "mismatch row {r}");
16003        }
16004        println!(
16005            "q1t matvec {rows}x{cols} (1 thread): div-decode {slow_ms:.2} ms  fused-LUT {fast_ms:.2} ms  => {:.2}x",
16006            slow_ms / fast_ms
16007        );
16008    }
16009}
16010
16011#[cfg(test)]
16012mod gemm_bench {
16013    /// `cargo test -p cortiq-engine --release q4tp_matmat_throughput -- --ignored --nocapture`
16014    /// Times the batched q4tp GEMM at the shapes the image DiT runs
16015    /// (b=296 tokens, 2304 -> 9216), on synthetic bytes: no model, no
16016    /// mmap, no thermal drift over minutes — a kernel change shows up
16017    /// here in seconds where a full render hides it in noise.
16018    ///
16019    /// On macOS add `CMF_ACCEL=0`: this shape is over the 500k-cell mark
16020    /// where the matmat hands off to Accelerate's dequant sgemm, and
16021    /// without the opt-out both rows below measure the AMX, not the
16022    /// kernel under test.
16023    #[test]
16024    #[ignore]
16025    fn q4tp_matmat_throughput() {
16026        let _alt = super::Q4TP_ALT_TEST_LOCK
16027            .lock()
16028            .unwrap_or_else(|e| e.into_inner());
16029        // 296 is a prompt-encode batch; the image DiT runs 2085 at
16030        // 512x512, where the activation panel stops fitting L2 and the
16031        // loop's shape starts to matter more than its instructions.
16032        let b: usize = std::env::var("CMF_BENCH_B")
16033            .ok()
16034            .and_then(|v| v.parse().ok())
16035            .unwrap_or(296);
16036        let (rows, cols) = (9216usize, 2304usize);
16037        let (_, _, _) = (rows, cols, b);
16038        let total =
16039            cortiq_core::quant::expected_nbytes(cortiq_core::TensorDtype::Q4TiledP, &[rows, cols])
16040                .unwrap();
16041        // Random nibbles are fine, but the row params are f16 (lo, step)
16042        // of a geometric ladder: garbage there gives exp2 of a huge
16043        // exponent, the scales come back inf, and the whole bench times
16044        // NaN arithmetic instead of the kernel.
16045        let (params_off, codes_off, _) = cortiq_core::quant::q4tp_sections(rows, cols);
16046        let mut bytes: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
16047        let lo = cortiq_core::quant::f32_to_f16(-4.0);
16048        let step = cortiq_core::quant::f32_to_f16(0.1);
16049        for r in 0..rows {
16050            let o = params_off + r * 4;
16051            bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
16052            bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
16053        }
16054        let _ = codes_off;
16055        let xs: Vec<f32> = (0..b * cols)
16056            .map(|i| ((i % 97) as f32 - 48.0) / 48.0)
16057            .collect();
16058        let mut out = vec![0f32; b * rows];
16059        let pool = crate::pool::Pool::from_env();
16060        // A shared 48-core stand drifts ±25% run to run, which is wider
16061        // than any kernel change worth making. So: alternate the two
16062        // kernels inside one process and keep the BEST time for
16063        // each. Interleaving makes both see the same interference, and a
16064        // minimum is the one statistic another tenant cannot inflate.
16065        super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
16066        let reps: usize = std::env::var("CMF_BENCH_REPS")
16067            .ok()
16068            .and_then(|v| v.parse().ok())
16069            .unwrap_or(10);
16070        let mut best = [f64::MAX; 2];
16071        let mut sums = [0f32; 2];
16072        for _ in 0..reps {
16073            for (k, w) in [(0usize, 1u8), (1usize, 2u8)] {
16074                super::Q4TP_ALT.store(w, std::sync::atomic::Ordering::Relaxed);
16075                let t = std::time::Instant::now();
16076                super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
16077                best[k] = best[k].min(t.elapsed().as_secs_f64());
16078                sums[k] = out.iter().take(64).sum::<f32>();
16079            }
16080        }
16081        let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
16082        for (k, name) in ["previous", "tuned   "].iter().enumerate() {
16083            println!(
16084                "q4tp matmat {rows}x{cols} b={b} {name}: {:.1} ms  {:.1} GFLOP/s  (checksum {:.3})",
16085                best[k] * 1e3,
16086                flops / best[k] / 1e9,
16087                sums[k]
16088            );
16089        }
16090        assert!(
16091            (sums[0] - sums[1]).abs() < 1e-2,
16092            "the tuned kernel changed the result: {} vs {}",
16093            sums[0],
16094            sums[1]
16095        );
16096    }
16097
16098    /// The blocked kernel must agree with the per-column path exactly —
16099    /// same weights, same activation split, only a different instruction
16100    /// mix. Shapes are chosen to hit the awkward cases: a column count
16101    /// that leaves an odd group (the 512-bit kernel does two at a time),
16102    /// and a batch that does not divide by four.
16103    #[test]
16104    fn q4tp_matmat_blocked_matches_scalar() {
16105        use std::sync::atomic::Ordering::Relaxed;
16106        let _alt = super::Q4TP_ALT_TEST_LOCK
16107            .lock()
16108            .unwrap_or_else(|e| e.into_inner());
16109        // The last shape carries the image DiT's column count — 2304, so
16110        // 72 groups of accumulation, which is where a reordered sum can
16111        // actually drift — and runs through the thread pool, since the
16112        // blocked path splits rows across workers. Its row count stays
16113        // under 500k cells on purpose: above that, macOS diverts the whole
16114        // matmat to the Accelerate/AMX dequant sgemm and neither kernel
16115        // here would run.
16116        for &(rows, cols, b) in &[
16117            (64usize, 128usize, 7usize),
16118            (33, 96, 4),
16119            (16, 256, 9),
16120            (192, 2304, 37),
16121        ] {
16122            let total = cortiq_core::quant::expected_nbytes(
16123                cortiq_core::TensorDtype::Q4TiledP,
16124                &[rows, cols],
16125            )
16126            .unwrap();
16127            let (params_off, _, _) = cortiq_core::quant::q4tp_sections(rows, cols);
16128            let mut bytes: Vec<u8> = (0..total).map(|i| (i * 61 % 251) as u8).collect();
16129            let lo = cortiq_core::quant::f32_to_f16(-4.0);
16130            let step = cortiq_core::quant::f32_to_f16(0.1);
16131            for r in 0..rows {
16132                let o = params_off + r * 4;
16133                bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
16134                bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
16135            }
16136            let xs: Vec<f32> = (0..b * cols)
16137                .map(|i| ((i % 89) as f32 - 44.0) / 44.0)
16138                .collect();
16139            let mut got = vec![0f32; b * rows];
16140            let mut want = vec![0f32; b * rows];
16141            let gpr = cols / 32;
16142            let view = super::Q4tpView::new(&bytes, rows, cols);
16143            let pool = crate::pool::Pool::from_env();
16144            super::Q4TP_ALT.store(2, Relaxed);
16145            super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut got, pool.as_deref());
16146            super::Q4TP_ALT.store(1, Relaxed);
16147            super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut want, pool.as_deref());
16148            super::Q4TP_ALT.store(0, Relaxed);
16149            // Measured against the output's scale, not cell by cell: a
16150            // dot product of 2304 terms lands near zero wherever the row
16151            // and the activation nearly cancel, and there a per-cell
16152            // ratio reports 1e-3 for an absolute error of 5e-6 — f32's
16153            // own rounding, reordered. What must stay small is the error
16154            // relative to what the layer actually outputs.
16155            let scale = want.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
16156            let (mut worst, mut at) = (0f32, 0usize);
16157            for (i, (g, w)) in got.iter().zip(&want).enumerate() {
16158                if (g - w).abs() > worst {
16159                    worst = (g - w).abs();
16160                    at = i;
16161                }
16162            }
16163            assert!(
16164                worst <= 1e-4 * scale,
16165                "{rows}x{cols} b={b}: blocked and scalar disagree by {worst:.3e} \
16166                 (scale {scale:.3e}) at cell {at}: {} vs {}",
16167                got[at],
16168                want[at]
16169            );
16170
16171            // "Same speed, no quality loss" is a claim about which answer
16172            // is RIGHT, not about which two agree. Both paths sum the same
16173            // 2304 products in different orders, so f64 decides: the
16174            // blocked kernel keeps sixteen partial sums and folds them at
16175            // the end, which is a shallower addition tree than the
16176            // per-column path's running scalar, and it must not be worse.
16177            let (mut e_blocked, mut e_scalar) = (0f64, 0f64);
16178            for bi in 0..b {
16179                let act = super::split_act(&xs[bi * cols..(bi + 1) * cols]);
16180                for r in 0..rows {
16181                    let mut sc = vec![0f32; gpr];
16182                    view.scales_into(r, gpr, &mut sc);
16183                    let mut exact = 0f64;
16184                    for j in 0..cols {
16185                        let (w, sq) = super::q4tp_outlier(view.nib, r, gpr, j, &sc);
16186                        exact += w as f64 * sq as f64 * act.xq[j] as f64;
16187                    }
16188                    exact *= act.sx as f64;
16189                    for &(j, xv) in &act.outliers {
16190                        let (w, sq) = super::q4tp_outlier(view.nib, r, gpr, j, &sc);
16191                        exact += w as f64 * sq as f64 * xv as f64;
16192                    }
16193                    let i = bi * rows + r;
16194                    e_blocked = e_blocked.max((got[i] as f64 - exact).abs());
16195                    e_scalar = e_scalar.max((want[i] as f64 - exact).abs());
16196                }
16197            }
16198            println!(
16199                "{rows}x{cols} b={b}: worst error vs f64 — blocked {e_blocked:.3e}, \
16200                 per-column {e_scalar:.3e}"
16201            );
16202            // An absolute bar, not a race between the two: at these
16203            // magnitudes both sit in f32's last bits, and on a small shape
16204            // whichever one happens to round the unluckiest cell "wins" by
16205            // a factor the next seed reverses.
16206            assert!(
16207                e_blocked <= 1e-5 * scale as f64 && e_scalar <= 1e-5 * scale as f64,
16208                "{rows}x{cols} b={b}: error against f64 too large — blocked \
16209                 {e_blocked:.3e}, per-column {e_scalar:.3e}, scale {scale:.3e}"
16210            );
16211        }
16212    }
16213
16214    /// What the row-exact contract costs a speculative-verify panel on the
16215    /// host: the same batch through the fast arms, through the row-exact
16216    /// arms, and as one matvec per token (the other way to be exact).
16217    /// Arms alternate inside one process and keep their best time.
16218    /// `CMF_GPU=0 cargo test -p cortiq-engine --release q4tp_matmat_row_exact_cost -- --ignored --nocapture`
16219    #[test]
16220    #[ignore]
16221    fn q4tp_matmat_row_exact_cost() {
16222        let _alt = super::Q4TP_ALT_TEST_LOCK
16223            .lock()
16224            .unwrap_or_else(|e| e.into_inner());
16225        let pool = crate::pool::Pool::from_env();
16226        let reps: usize = std::env::var("CMF_BENCH_REPS")
16227            .ok()
16228            .and_then(|v| v.parse().ok())
16229            .unwrap_or(30);
16230        for &(rows, cols) in &[(2048usize, 4096usize), (4096, 2048), (4096, 4096)] {
16231            let total = cortiq_core::quant::expected_nbytes(
16232                cortiq_core::TensorDtype::Q4TiledP,
16233                &[rows, cols],
16234            )
16235            .unwrap();
16236            let (params_off, _, _) = cortiq_core::quant::q4tp_sections(rows, cols);
16237            let mut bytes: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
16238            let lo = cortiq_core::quant::f32_to_f16(-4.0);
16239            let step = cortiq_core::quant::f32_to_f16(0.1);
16240            for r in 0..rows {
16241                let o = params_off + r * 4;
16242                bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
16243                bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
16244            }
16245            for &b in &[2usize, 4, 5, 8] {
16246                let xs: Vec<f32> = (0..b * cols)
16247                    .map(|i| ((i % 97) as f32 - 48.0) / 48.0)
16248                    .collect();
16249                let mut out = vec![0f32; b * rows];
16250                let mut best = [f64::MAX; 3];
16251                for _ in 0..reps {
16252                    for (k, best_k) in best.iter_mut().enumerate() {
16253                        let t = std::time::Instant::now();
16254                        match k {
16255                            0 | 1 => super::q4tp_matmat_with(
16256                                &bytes,
16257                                &xs,
16258                                b,
16259                                rows,
16260                                cols,
16261                                &mut out,
16262                                pool.as_deref(),
16263                                k == 1,
16264                            ),
16265                            _ => {
16266                                for (bi, o) in out.chunks_mut(rows).enumerate() {
16267                                    super::q4tp_matvec(
16268                                        &bytes,
16269                                        &xs[bi * cols..(bi + 1) * cols],
16270                                        rows,
16271                                        cols,
16272                                        o,
16273                                        pool.as_deref(),
16274                                    );
16275                                }
16276                            }
16277                        }
16278                        *best_k = best_k.min(t.elapsed().as_secs_f64());
16279                    }
16280                }
16281                println!(
16282                    "q4tp {rows}x{cols} b={b}: fast {:.3} ms, row-exact {:.3} ms, \
16283                     {b} matvecs {:.3} ms",
16284                    best[0] * 1e3,
16285                    best[1] * 1e3,
16286                    best[2] * 1e3
16287                );
16288            }
16289        }
16290    }
16291
16292    /// The q4t twin of the throughput bench, same shape and rules, so the
16293    /// two quantisations' batch kernels can be read against each other.
16294    /// `cargo test -p cortiq-engine --release q4t_matmat_throughput -- --ignored --nocapture`
16295    #[test]
16296    #[ignore]
16297    fn q4t_matmat_throughput() {
16298        let (rows, cols, b) = (9216usize, 2304usize, 296usize);
16299        let total =
16300            cortiq_core::quant::expected_nbytes(cortiq_core::TensorDtype::Q4Tiled, &[rows, cols])
16301                .unwrap();
16302        // q4t carries a per-group f16 scale in the tile's first two bytes;
16303        // random bytes there decode to inf and the bench would time NaNs.
16304        let mut bytes: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
16305        let sc = cortiq_core::quant::f32_to_f16(0.02);
16306        for t in bytes.chunks_mut(super::Q4_TILE) {
16307            t[..2].copy_from_slice(&sc.to_le_bytes());
16308        }
16309        let xs: Vec<f32> = (0..b * cols)
16310            .map(|i| ((i % 97) as f32 - 48.0) / 48.0)
16311            .collect();
16312        let mut out = vec![0f32; b * rows];
16313        let pool = crate::pool::Pool::from_env();
16314        super::q4t_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
16315        let reps: usize = std::env::var("CMF_BENCH_REPS")
16316            .ok()
16317            .and_then(|v| v.parse().ok())
16318            .unwrap_or(10);
16319        let mut best = f64::MAX;
16320        for _ in 0..reps {
16321            let t = std::time::Instant::now();
16322            super::q4t_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
16323            best = best.min(t.elapsed().as_secs_f64());
16324        }
16325        let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
16326        println!(
16327            "q4t matmat {rows}x{cols} b={b}: {:.1} ms  {:.1} GFLOP/s  (checksum {:.3})",
16328            best * 1e3,
16329            flops / best / 1e9,
16330            out.iter().take(64).sum::<f32>()
16331        );
16332    }
16333}