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/// Whole q4tp experts on the host's cores (AVX-512 VNNI): the misses of a
28/// card's expert arena.
29#[cfg(target_arch = "x86_64")]
30mod moe_fast;
31
32/// D5 GPU-matmat gate telemetry: large-GEMM calls that reached the gate
33/// vs those the q8_2f device kernel actually consumed (SAM3 profiling).
34pub static MATMAT_GATE: AtomicU64 = AtomicU64::new(0);
35pub static MATMAT_GPU_HIT: AtomicU64 = AtomicU64::new(0);
36
37pub enum QTensor {
38    F32 {
39        data: Vec<f32>,
40        rows: usize,
41        cols: usize,
42    },
43    Mapped {
44        model: Arc<CmfModel>,
45        /// Index into the model's tensor directory.
46        idx: usize,
47        dtype: TensorDtype,
48        rows: usize,
49        cols: usize,
50        /// Per-row scales, dequantized to f32 up front (tiny).
51        row_scale: Vec<f32>,
52        /// q8_2f column field (θ), dequantized up front; empty for q8_row.
53        col_field: Vec<f32>,
54        /// Vbit only: byte offset of each row's packed data within the
55        /// tensor blob (`[rows + 1]`, computed once at load — the per-
56        /// matvec prefix scan over row bit-widths was O(rows) each call).
57        vbit_offsets: Vec<usize>,
58        /// q8-family decode repack (load-time, optional): rows in groups
59        /// of 4, interleaved in 16-byte units — one 64-byte line per
60        /// iteration feeds all 4 sdot lanes, ONE sequential weight
61        /// stream per worker instead of four (this is where llama.cpp's
62        /// repacked Q8 kernels get their bandwidth). Empty = off
63        /// (CMF_REPACK=0, non-SDOT arch, or an ineligible shape). Trades
64        /// an anonymous copy of the quants for mmap pages that go cold.
65        repack: Vec<u8>,
66    },
67    /// Binary low-rank (`blr`): ±1 paths, signs from the mmap, scales as
68    /// f32 (docs/BLR_NVG.md §3). Every device-graph accessor answers `None`
69    /// for it (whole-token graphs decline to per-op/host paths); the CPU
70    /// kernels live in `crate::blr`, the per-op wgpu kernels in
71    /// `gpu_wgpu::blr`.
72    Blr(Box<crate::blr::BlrTensor>),
73}
74
75/// Load-time q8 repack gate (see `Mapped::repack`). OPT-IN
76/// (`CMF_REPACK=1`): the single-stream hypothesis LOST on Apple Silicon
77/// (M4, interleaved A/B: decode 101 vs 94 tok/s — four adjacent row
78/// streams per worker feed the prefetcher MORE memory-level parallelism
79/// than one); kept as an experiment flag for x86, where the tradeoff
80/// may land differently.
81fn repack_enabled() -> bool {
82    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
83    *ON.get_or_init(|| {
84        std::env::var("CMF_REPACK")
85            .map(|v| v == "1")
86            .unwrap_or(cfg!(target_os = "android"))
87    })
88}
89
90/// Interleave q8 rows for the decode kernel: group g holds rows
91/// 4g..4g+4 as [r0[c], r1[c], r2[c], r3[c]] per 16-byte chunk c. Only
92/// full groups are packed — tail rows keep reading the mmap layout.
93fn q8_repack(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
94    #[cfg(target_arch = "aarch64")]
95    let arch_ok = sdot_enabled();
96    #[cfg(not(target_arch = "aarch64"))]
97    let arch_ok = false;
98    if !arch_ok || !repack_enabled() || rows < 256 || cols % 16 != 0 {
99        return Vec::new();
100    }
101    q8_repack_layout(bytes, rows, cols)
102}
103
104/// The pure layout transform behind `q8_repack` (tested directly —
105/// the gate depends on arch and env).
106fn q8_repack_layout(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
107    let groups = rows / 4;
108    let mut rep = vec![0u8; groups * 4 * cols];
109    for g in 0..groups {
110        let dst = &mut rep[g * 4 * cols..(g + 1) * 4 * cols];
111        for c in 0..cols / 16 {
112            for lane in 0..4 {
113                let src = (g * 4 + lane) * cols + c * 16;
114                dst[c * 64 + lane * 16..c * 64 + lane * 16 + 16]
115                    .copy_from_slice(&bytes[src..src + 16]);
116            }
117        }
118    }
119    rep
120}
121
122/// Prefix-sum of vbit row payload offsets (absolute within the tensor
123/// bytes). `offsets[r]..offsets[r+1]` is row r's packed data.
124fn vbit_row_offsets(bytes: &[u8], rows: usize, cols: usize) -> Vec<usize> {
125    let ng = cols / GROUP_SIZE;
126    let bits = &bytes[..rows];
127    let mut offsets = Vec::with_capacity(rows + 1);
128    let mut off = rows + rows * ng * 2;
129    for r in 0..rows {
130        offsets.push(off);
131        off += (cols * bits[r] as usize).div_ceil(8);
132    }
133    offsets.push(off);
134    offsets
135}
136
137/// `CMF_X86_BLOCKED` / `CMF_GPU_LMHEAD` / `CMF_GPU_SPLIT`, read once. They
138/// used to be read from the environment on every large matvec and on every
139/// matmat in six places — microseconds each, but also a knob that could
140/// change under a running process, which is not a thing a kernel choice
141/// should be able to do mid-sequence.
142fn blocked_enabled() -> bool {
143    use std::sync::atomic::Ordering::Relaxed;
144    match BLOCKED_OVERRIDE.load(Relaxed) {
145        1 => false,
146        2 => true,
147        _ => {
148            static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
149            *ON.get_or_init(|| {
150                std::env::var("CMF_X86_BLOCKED")
151                    .map(|v| v != "0")
152                    .unwrap_or(true)
153            })
154        }
155    }
156}
157
158static BLOCKED_OVERRIDE: std::sync::atomic::AtomicU8 = std::sync::atomic::AtomicU8::new(0);
159
160/// Force the blocked GEMM on or off, ignoring the environment; `None`
161/// restores it. For tests that need to run BOTH paths and compare them:
162/// `blocked_enabled` caches its answer for the life of the process, which
163/// is right when the environment is the only input, but leaves a test that
164/// flips `CMF_X86_BLOCKED` between two calls comparing a path against
165/// itself — or against whatever a test running in parallel latched first.
166pub fn set_blocked_override(on: Option<bool>) {
167    let v = match on {
168        None => 0,
169        Some(false) => 1,
170        Some(true) => 2,
171    };
172    BLOCKED_OVERRIDE.store(v, std::sync::atomic::Ordering::Relaxed);
173}
174
175fn gpu_lmhead_enabled() -> bool {
176    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
177    *ON.get_or_init(|| {
178        std::env::var("CMF_GPU_LMHEAD")
179            .map(|v| v != "0")
180            .unwrap_or(true)
181    })
182}
183
184fn gpu_split_frac() -> f32 {
185    // MiMo's banked verification projects every row on the device. Plain
186    // decode must not quantize half of its O/head activation on the CPU.
187    if FULL_GPU_Q8.get() {
188        return 1.0;
189    }
190    static F: std::sync::OnceLock<f32> = std::sync::OnceLock::new();
191    *F.get_or_init(|| {
192        std::env::var("CMF_GPU_SPLIT")
193            .ok()
194            .and_then(|v| v.parse::<f32>().ok())
195            .unwrap_or(0.5)
196            .clamp(0.0, 1.0)
197    })
198}
199
200impl QTensor {
201    pub fn from_f32(data: Vec<f32>, rows: usize, cols: usize) -> Self {
202        debug_assert_eq!(data.len(), rows * cols);
203        Self::F32 { data, rows, cols }
204    }
205
206    /// Wrap a directory tensor without dequantizing the payload.
207    /// Falls back to dequantized f32 for dtypes without a fused kernel.
208    pub fn from_model(model: &Arc<CmfModel>, name: &str) -> Result<Self, String> {
209        // Indexed lookup: the linear directory scan made pipeline build
210        // O(N²) on MoE/skills files with thousands of tensors.
211        let idx = model
212            .tensor_index(name)
213            .ok_or_else(|| format!("tensor '{name}' not found in CMF directory"))?;
214        let entry = &model.tensors[idx];
215        if entry.shape.len() != 2 {
216            return Err(format!("QTensor::from_model needs 2-D, got '{name}'"));
217        }
218        let (rows, cols) = (entry.shape[0], entry.shape[1]);
219        let bytes = model.entry_bytes(entry);
220
221        match entry.dtype {
222            TensorDtype::Q8Row | TensorDtype::Q8_2f => {
223                let n = rows * cols;
224                let scales_off = n;
225                let row_scale: Vec<f32> = (0..rows)
226                    .map(|o| {
227                        f16_to_f32(u16::from_le_bytes([
228                            bytes[scales_off + o * 2],
229                            bytes[scales_off + o * 2 + 1],
230                        ]))
231                    })
232                    .collect();
233                let col_field: Vec<f32> = if entry.dtype == TensorDtype::Q8_2f {
234                    let col_off = n + rows * 2;
235                    (0..cols)
236                        .map(|i| {
237                            f16_to_f32(u16::from_le_bytes([
238                                bytes[col_off + i * 2],
239                                bytes[col_off + i * 2 + 1],
240                            ]))
241                        })
242                        .collect()
243                } else {
244                    Vec::new()
245                };
246                Ok(Self::Mapped {
247                    model: model.clone(),
248                    idx,
249                    dtype: entry.dtype,
250                    rows,
251                    cols,
252                    row_scale,
253                    col_field,
254                    vbit_offsets: Vec::new(),
255                    repack: q8_repack(bytes, rows, cols),
256                })
257            }
258            // vbit: fused kernel unpacks variable-bit rows from mmap.
259            TensorDtype::Vbit if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
260                model: model.clone(),
261                idx,
262                dtype: entry.dtype,
263                rows,
264                cols,
265                row_scale: Vec::new(),
266                col_field: Vec::new(),
267                vbit_offsets: vbit_row_offsets(bytes, rows, cols),
268                repack: Vec::new(),
269            }),
270            // vbit_ro (§4.2): the offset table comes straight from the
271            // file — no load-time prefix scan; kernels are shared with
272            // legacy vbit (they consume absolute offsets either way).
273            TensorDtype::VbitRo if cols % GROUP_SIZE == 0 => {
274                let (_, off_off, packed_off) = cortiq_core::quant::vbit_ro_sections(rows, cols);
275                let offsets: Vec<usize> = (0..=rows)
276                    .map(|r| packed_off + cortiq_core::quant::vbit_ro_offset(bytes, off_off, r))
277                    .collect();
278                Ok(Self::Mapped {
279                    model: model.clone(),
280                    idx,
281                    dtype: entry.dtype,
282                    rows,
283                    cols,
284                    row_scale: Vec::new(),
285                    col_field: Vec::new(),
286                    vbit_offsets: offsets,
287                    repack: Vec::new(),
288                })
289            }
290            // q4_block: fused kernel reads nibbles straight from mmap —
291            // a 14B q4 file no longer explodes into ×8 f32 RAM.
292            // q4_tiled (§4.3): interleaved [scale][nibbles] tiles — one
293            // sequential memory stream (measured ×1.66 ARM / ×1.13 AVX2
294            // at kernel level over the split layout).
295            TensorDtype::Q4Tiled if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
296                model: model.clone(),
297                idx,
298                dtype: entry.dtype,
299                rows,
300                cols,
301                row_scale: Vec::new(),
302                col_field: Vec::new(),
303                vbit_offsets: Vec::new(),
304                repack: Vec::new(),
305            }),
306            // q4tp (§4.10): nibbles from mmap, scale from the row ladder —
307            // 7.3% less file than q4t at the same 4-bit grid.
308            TensorDtype::Q4TiledP if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
309                model: model.clone(),
310                idx,
311                dtype: entry.dtype,
312                rows,
313                cols,
314                row_scale: Vec::new(),
315                col_field: Vec::new(),
316                vbit_offsets: Vec::new(),
317                repack: Vec::new(),
318            }),
319            // q2tp: 2-bit chunks from mmap, scale from the same row ladder.
320            TensorDtype::Q2TiledP if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
321                model: model.clone(),
322                idx,
323                dtype: entry.dtype,
324                rows,
325                cols,
326                row_scale: Vec::new(),
327                col_field: Vec::new(),
328                vbit_offsets: Vec::new(),
329                repack: Vec::new(),
330            }),
331            TensorDtype::Q4Block if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
332                model: model.clone(),
333                idx,
334                dtype: entry.dtype,
335                rows,
336                cols,
337                row_scale: Vec::new(),
338                col_field: Vec::new(),
339                vbit_offsets: Vec::new(),
340                repack: Vec::new(),
341            }),
342            // q1: binary sign-bit tiles from mmap (1-bit-trained models).
343            // q1g: the same signs as planes with one scale per 128-group.
344            TensorDtype::Q1 if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
345                model: model.clone(),
346                idx,
347                dtype: entry.dtype,
348                rows,
349                cols,
350                row_scale: Vec::new(),
351                col_field: Vec::new(),
352                vbit_offsets: Vec::new(),
353                repack: Vec::new(),
354            }),
355            TensorDtype::Q1G if cols % cortiq_core::quant::Q1G_GROUP == 0 => Ok(Self::Mapped {
356                model: model.clone(),
357                idx,
358                dtype: entry.dtype,
359                rows,
360                cols,
361                row_scale: Vec::new(),
362                col_field: Vec::new(),
363                vbit_offsets: Vec::new(),
364                repack: Vec::new(),
365            }),
366            // mxfp4: E2M1 codes + E8M0 scales straight from the mmap (the
367            // native FP4 expert codec; `qtensor::mxfp4`).
368            TensorDtype::Mxfp4 if cols % cortiq_core::mxfp4::GROUP == 0 => Ok(Self::Mapped {
369                model: model.clone(),
370                idx,
371                dtype: entry.dtype,
372                rows,
373                cols,
374                row_scale: Vec::new(),
375                col_field: Vec::new(),
376                vbit_offsets: Vec::new(),
377                repack: Vec::new(),
378            }),
379            // q1t (ternary + outlier overlay): fused per-row dequant kernel
380            // reads straight from mmap — a 12B q1t stays ~its file size in
381            // RAM instead of dequantizing to ~48 GB of f32.
382            TensorDtype::Q1T if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
383                model: model.clone(),
384                idx,
385                dtype: entry.dtype,
386                rows,
387                cols,
388                row_scale: Vec::new(),
389                col_field: Vec::new(),
390                vbit_offsets: Vec::new(),
391                repack: Vec::new(),
392            }),
393            // blr: factored kernels over the packed signs. CMF_BLR_DENSE=1
394            // dequantizes to f32 instead (A/B parity runs).
395            TensorDtype::BinaryLowRank if std::env::var("CMF_BLR_DENSE").as_deref() != Ok("1") => {
396                Ok(Self::Blr(Box::new(crate::blr::BlrTensor::from_model(
397                    model, idx,
398                )?)))
399            }
400            // No fused kernel yet → dequantize once (correct, more RAM).
401            _ => {
402                let mut data = vec![0.0f32; rows * cols];
403                cortiq_core::quant::dequant_tensor(entry, bytes, &mut data)?;
404                Ok(Self::from_f32(data, rows, cols))
405            }
406        }
407    }
408
409    /// Binary low-rank tensor? Fused multi-tensor kernels refuse it and fall
410    /// back to per-tensor calls.
411    pub fn is_blr(&self) -> bool {
412        matches!(self, Self::Blr(_))
413    }
414
415    /// Memory-mapped `mxfp4` tensor? Like `blr`, it has no row-range kernel
416    /// in the fused multi-tensor jobs: those run it per tensor.
417    pub fn is_mxfp4(&self) -> bool {
418        matches!(
419            self,
420            Self::Mapped {
421                dtype: TensorDtype::Mxfp4,
422                ..
423            }
424        )
425    }
426
427    /// Codecs the fused multi-tensor jobs (`matvec_many`, the SwiGLU and
428    /// MoE row kernels) do not carry: they fall back to per-tensor calls.
429    fn per_tensor_only(&self) -> bool {
430        self.is_blr() || self.is_mxfp4()
431    }
432
433    /// (model, idx) of a `blr` tensor, for the wgpu graphs only (their
434    /// `blr` stage kernels, kind [`crate::gpu::GRAPH_KIND_BLR`]). Kept apart
435    /// from [`graph_weight`](Self::graph_weight), which stays None for
436    /// `blr`: every other graph consumer keys its kernels on that answer.
437    /// A file with a Prism contract keeps the per-op path.
438    pub fn graph_blr(&self) -> Option<(&Arc<CmfModel>, usize)> {
439        match self {
440            Self::Blr(t) if !crate::prism::has_contract(&t.model) => Some((&t.model, t.idx)),
441            _ => None,
442        }
443    }
444
445    /// [`graph_weight`](Self::graph_weight), then the Prism descriptor
446    /// form, then a `blr` tensor as kind [`crate::gpu::GRAPH_KIND_BLR`] —
447    /// the wgpu token and batch graph builders' lookup.
448    pub(crate) fn graph_weight_or_blr(
449        &self,
450    ) -> Option<(&std::sync::Arc<CmfModel>, usize, u8, &[f32])> {
451        self.graph_weight()
452            .or_else(|| self.graph_weight_descriptor())
453            .or_else(|| {
454                self.graph_blr()
455                    .map(|(m, i)| (m, i, crate::gpu::GRAPH_KIND_BLR, &[][..]))
456            })
457    }
458
459    /// Binary-weight-mapped tensor, q1 or q1g? (GPU gates: the binary CPU
460    /// kernels are compute-bound, so offload pays at much smaller shapes
461    /// than q8.)
462    pub(crate) fn is_q1(&self) -> bool {
463        matches!(
464            self,
465            Self::Mapped {
466                dtype: TensorDtype::Q1 | TensorDtype::Q1G,
467                ..
468            }
469        )
470    }
471
472    /// The binary-weight view (q1 tiles or q1g planes) the shared CPU row
473    /// drivers take; None for every other storage.
474    fn bin_w(&self) -> Option<BinW<'_>> {
475        match self {
476            Self::Mapped {
477                dtype: TensorDtype::Q1,
478                ..
479            } => Some(BinW::Q1(self.quant_bytes())),
480            Self::Mapped {
481                dtype: TensorDtype::Q1G,
482                rows,
483                cols,
484                ..
485            } => Some(BinW::q1g(self.quant_bytes(), *rows, *cols)),
486            _ => None,
487        }
488    }
489
490    /// Owned-f32 view (data, rows, cols) — the GDN a/b gate projections
491    /// arrive dequantized (force-f16 in the converter → F32 in RAM).
492    pub(crate) fn f32_parts(&self) -> Option<(&[f32], usize, usize)> {
493        match self {
494            Self::F32 { data, rows, cols } => Some((data, *rows, *cols)),
495            _ => None,
496        }
497    }
498
499    /// (directory idx, rows, cols) of a q1-mapped tensor — the
500    /// whole-block GPU path resolves offsets itself.
501    /// (idx, rows, cols) of a mapped tensor the whole-token GPU graph can drive
502    /// — Q1, Q1T or Q4-block (it resolves the offset and picks the kernel by
503    /// dtype). Q4-block lets a precise down_proj/lm_head stay on-device.
504    /// Named `q1_parts` for historical reasons.
505    pub(crate) fn q1_parts(&self) -> Option<(usize, usize, usize)> {
506        if self.has_prism_contract() {
507            return None;
508        }
509        match self {
510            #[cfg(target_os = "macos")]
511            Self::Mapped {
512                dtype: TensorDtype::Q1T,
513                ..
514            } if !crate::gpu::metal_q1t_enabled() => None,
515            Self::Mapped {
516                idx,
517                dtype:
518                    TensorDtype::Q1
519                    | TensorDtype::Q1G
520                    | TensorDtype::Q1T
521                    | TensorDtype::Q4Block
522                    | TensorDtype::Q4Tiled
523                    // Q2TiledP deliberately absent: the Metal graph has no
524                    // q2tp kernel, and advertising it here made the block
525                    // plan truncate mid-run at the first q2tp layer.
526                    | TensorDtype::Q4TiledP
527                    | TensorDtype::Q8Row
528                    | TensorDtype::Q8_2f,
529                rows,
530                cols,
531                ..
532            } => Some((*idx, *rows, *cols)),
533            _ => None,
534        }
535    }
536
537    /// `(directory idx, rows, cols)` for the native Metal token graph.  The
538    /// historical q1 graph gate intentionally refuses every Prism tensor so
539    /// an untransformed q2 payload cannot slip into the resident path.  The
540    /// Metal2 graph is descriptor-aware and admits only the production
541    /// q2tp-affine forward targets; ordinary q1/q4 callers retain the old
542    /// `q1_parts` behaviour.
543    #[cfg(target_os = "macos")]
544    pub(crate) fn metal_graph_parts(&self) -> Option<(usize, usize, usize)> {
545        if let Some((model, idx, kind, _)) = self.graph_weight_descriptor() {
546            let name = &model.tensors[idx].name;
547            let forward = kind == 9 && crate::prism::is_forward_weight(model, name);
548            let affine = kind == 9 && crate::prism::is_affine_target(model, name);
549            if forward && affine {
550                let e = model.tensors.get(idx)?;
551                return Some((idx, *e.shape.first()?, *e.shape.get(1)?));
552            }
553        }
554        self.q1_parts()
555    }
556
557    /// (directory idx, rows, cols) of a q4_tiled mapped tensor. The
558    /// chunk-prefill graph takes it in the same 4-tuple slot as
559    /// `q8_row_parts` with an EMPTY row_scale — q4t carries its scales
560    /// inside the 18-byte tiles, and the empty slice is what tells the
561    /// encoder to reach for the q4t kernels.
562    pub(crate) fn q4t_parts(&self) -> Option<(usize, usize, usize)> {
563        if self.has_prism_contract() {
564            return None;
565        }
566        match self {
567            Self::Mapped {
568                idx,
569                dtype: TensorDtype::Q4Tiled,
570                rows,
571                cols,
572                ..
573            } => Some((*idx, *rows, *cols)),
574            _ => None,
575        }
576    }
577
578    /// (directory idx, rows, cols) of a q4tp mapped tensor. Same empty-scale
579    /// slot as `q4t_parts` in the chunk graph — the encoder tells the two
580    /// apart by the tensor's dtype, not by the slot.
581    pub(crate) fn q4tp_parts(&self) -> Option<(usize, usize, usize)> {
582        if self.has_prism_contract() {
583            return None;
584        }
585        match self {
586            Self::Mapped {
587                idx,
588                dtype: TensorDtype::Q4TiledP,
589                rows,
590                cols,
591                ..
592            } => Some((*idx, *rows, *cols)),
593            _ => None,
594        }
595    }
596
597    /// (directory idx, rows, cols, row_scale) of a plain q8_row mapped
598    /// tensor — the chunk-prefill GPU graph resolves offsets itself.
599    /// q8_2f is excluded on purpose: its column field would need a
600    /// prescale stage on the device.
601    pub(crate) fn q8_row_parts(&self) -> Option<(usize, usize, usize, &[f32])> {
602        if self.has_prism_contract() {
603            return None;
604        }
605        match self {
606            Self::Mapped {
607                idx,
608                dtype: TensorDtype::Q8Row,
609                rows,
610                cols,
611                row_scale,
612                col_field,
613                ..
614            } if col_field.is_empty() => Some((*idx, *rows, *cols, row_scale)),
615            _ => None,
616        }
617    }
618
619    /// (directory idx, rows, cols, row_scale, col_field) of a q8_2f mapped
620    /// tensor — the chunk-prefill graph stages x·col on the device first.
621    pub(crate) fn q8_2f_parts(&self) -> Option<(usize, usize, usize, &[f32], &[f32])> {
622        if self.has_prism_contract() {
623            return None;
624        }
625        match self {
626            Self::Mapped {
627                idx,
628                dtype: TensorDtype::Q8_2f,
629                rows,
630                cols,
631                row_scale,
632                col_field,
633                ..
634            } if col_field.len() == *cols => Some((*idx, *rows, *cols, row_scale, col_field)),
635            _ => None,
636        }
637    }
638
639    /// The layout this tensor is stored in, when it is mapped from a model.
640    /// The frames branch on it — a q2tp gate against a q4tp down is a real
641    /// combination in the 2-bit profile and needs a different kernel.
642    pub fn model_dtype(&self) -> Option<cortiq_core::TensorDtype> {
643        match self {
644            Self::Mapped { dtype, .. } => Some(*dtype),
645            _ => None,
646        }
647    }
648
649    /// The tensor's index in the model directory, when it is mapped from one.
650    /// The GPU frames bind by index rather than by name — a name lookup per
651    /// layer per token is not free, and the index is what the device cache is
652    /// keyed on anyway.
653    pub fn model_idx(&self) -> Option<usize> {
654        match self {
655            Self::Mapped { idx, .. } => Some(*idx),
656            _ => None,
657        }
658    }
659
660    /// The model this tensor is mapped from, when it is mapped at all. The
661    /// GPU frames need the container to reach the bytes; a QTensor already
662    /// holds it, and threading a second handle down every call site to say
663    /// the same thing invites the two to disagree.
664    pub fn model_arc(&self) -> Option<std::sync::Arc<cortiq_core::CmfModel>> {
665        match self {
666            Self::Mapped { model, .. } => Some(model.clone()),
667            _ => None,
668        }
669    }
670
671    /// The mapped byte ranges decoding row `r` reads, as (address, length),
672    /// appended to `out`: what a caller hands the kernel as read-ahead advice
673    /// before a batch of row lookups in a table the page cache does not
674    /// hold. `false`, and nothing appended, for a tensor that is not mapped
675    /// or a codec whose rows have no fixed address.
676    pub fn row_spans(&self, r: usize, out: &mut Vec<(usize, usize)>) -> bool {
677        let Self::Mapped { dtype, .. } = self else {
678            return false;
679        };
680        let Some(spans) = row_span_offsets(*dtype, self.rows(), self.cols(), r) else {
681            return false;
682        };
683        let base = self.quant_bytes().as_ptr() as usize;
684        out.extend(spans.iter().filter(|s| s.1 > 0).map(|&(o, n)| (base + o, n)));
685        true
686    }
687
688    /// Whether this mapped tensor belongs to the Prism/Bonsai transform
689    /// contract.  Device graphs do not carry the descriptor, so callers use
690    /// this conservative predicate to stay on the descriptor-aware CPU path
691    /// instead of silently executing an unrotated matrix.
692    pub(crate) fn has_prism_contract(&self) -> bool {
693        matches!(self, Self::Mapped { model, .. } if crate::prism::has_contract(model))
694    }
695
696    pub fn rows(&self) -> usize {
697        match self {
698            Self::F32 { rows, .. } | Self::Mapped { rows, .. } => *rows,
699            Self::Blr(t) => t.rows,
700        }
701    }
702
703    /// Mapped q4t handle (model + directory index) — the fused GPU FFN
704    /// needs the raw file coordinates of its three projections.
705    pub(crate) fn mapped_q4t(&self) -> Option<(&Arc<CmfModel>, usize)> {
706        if self.has_prism_contract() {
707            return None;
708        }
709        match self {
710            Self::Mapped {
711                model,
712                idx,
713                dtype: TensorDtype::Q4Tiled,
714                ..
715            } => Some((model, *idx)),
716            _ => None,
717        }
718    }
719
720    /// Same slot as `mapped_q4t` for a q4tp tensor — the fused DiT FFN picks
721    /// its kernels by which of the two answers.
722    pub fn mapped_q4tp(&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,
731                ..
732            } => Some((model, *idx)),
733            _ => None,
734        }
735    }
736
737    /// (model, tensor idx) for a mapped weight in ANY codec the fused device
738    /// paths can run — four-bit tiled or either int8 layout.
739    ///
740    /// The fused DiT chains asked for `mapped_q4tp` by name, so an eight-bit
741    /// container never reached them and rendered through per-op GEMMs even
742    /// after those kernels learned its codec. The gate is what the codec has
743    /// a device GEMM for, not which codec it is.
744    pub fn mapped_device_gemm(&self) -> Option<(&Arc<CmfModel>, usize)> {
745        if self.has_prism_contract() {
746            return None;
747        }
748        match self {
749            Self::Mapped {
750                model,
751                idx,
752                dtype: TensorDtype::Q4TiledP | TensorDtype::Q8Row | TensorDtype::Q8_2f,
753                ..
754            } => Some((model, *idx)),
755            _ => None,
756        }
757    }
758
759    /// (model, tensor idx) for a q8_2f mapped weight — the eight-bit slot of
760    /// the MoE graph ladder (int8 body, f16 row scales, f16 column field).
761    pub fn mapped_q8_2f(&self) -> Option<(&Arc<CmfModel>, usize)> {
762        if self.has_prism_contract() {
763            return None;
764        }
765        match self {
766            Self::Mapped {
767                model,
768                idx,
769                dtype: TensorDtype::Q8_2f,
770                ..
771            } => Some((model, *idx)),
772            _ => None,
773        }
774    }
775
776    /// (model, tensor idx) for a q2tp mapped weight — the 2-bit twin of
777    /// `mapped_q4tp`, used by the mixed MoE profile.
778    pub fn mapped_q2tp(&self) -> Option<(&Arc<CmfModel>, usize)> {
779        if self.has_prism_contract() {
780            return None;
781        }
782        match self {
783            Self::Mapped {
784                model,
785                idx,
786                dtype: TensorDtype::Q2TiledP,
787                ..
788            } => Some((model, *idx)),
789            _ => None,
790        }
791    }
792
793    pub fn cols(&self) -> usize {
794        match self {
795            Self::F32 { cols, .. } | Self::Mapped { cols, .. } => *cols,
796            Self::Blr(t) => t.cols,
797        }
798    }
799
800    /// (model, tensor idx) for a q1 mapped weight — the wgpu token graph
801    /// keys its resident VRAM cache by idx. None for any other dtype/kind.
802    pub fn mapped_q1(&self) -> Option<(&std::sync::Arc<CmfModel>, usize)> {
803        if self.has_prism_contract() {
804            return None;
805        }
806        match self {
807            Self::Mapped {
808                model,
809                idx,
810                dtype: TensorDtype::Q1,
811                ..
812            } => Some((model, *idx)),
813            _ => None,
814        }
815    }
816
817    /// (model, idx, kind, row_scale) for a graph-capable mapped weight.
818    /// kind: 0=q8_row (per-row scales), 1=q1, 2=q4_block, 3=q1t
819    /// (tile-embedded, no rs), 5=q4_tiled, 6=q4tp, 7=q8_2f (both scale
820    /// planes live inside the tensor), 9=q2tp, 10=q1g (sign + scale
821    /// planes). None only for `vbit`.
822    ///
823    /// The old comment here claimed q4_block was unhandled while the arm
824    /// right below mapped it, and it named q8_2f as unhandled after that
825    /// stopped being true — a stale comment on this function is how a
826    /// model silently loses the graph, so it is worth keeping honest.
827    pub fn graph_weight(&self) -> Option<(&std::sync::Arc<CmfModel>, usize, u8, &[f32])> {
828        if self.has_prism_contract() {
829            return None;
830        }
831        self.graph_weight_descriptor()
832    }
833
834    /// Descriptor-aware graph handle used only by the Prism token graph.
835    /// Ordinary graph callers continue to use [`graph_weight`] and therefore
836    /// remain fail-closed until they provide the same explicit transform
837    /// contract.
838    pub(crate) fn graph_weight_descriptor(
839        &self,
840    ) -> Option<(&std::sync::Arc<CmfModel>, usize, u8, &[f32])> {
841        match self {
842            Self::Mapped {
843                model,
844                idx,
845                dtype: TensorDtype::Q8Row,
846                row_scale,
847                ..
848            } => Some((model, *idx, 0, row_scale.as_slice())),
849            Self::Mapped {
850                model,
851                idx,
852                dtype: TensorDtype::Q1,
853                ..
854            } => Some((model, *idx, 1, &[])),
855            // Kind 10, never 1: q1g's sign and scale planes are a different
856            // walk from q1's 6-byte tiles (the kind-2/5 mixup taught that).
857            Self::Mapped {
858                model,
859                idx,
860                dtype: TensorDtype::Q1G,
861                ..
862            } => Some((model, *idx, 10, &[])),
863            // Q4Tiled is kind 5, NOT 2: both carried 2 historically, and
864            // the wgpu token graph fed 18B interleaved tiles to the
865            // split-layout q4b kernel — garbage output on q4t models
866            // (caught by an end-to-end answer check on real Vulkan).
867            Self::Mapped {
868                model,
869                idx,
870                dtype: TensorDtype::Q4Tiled,
871                ..
872            } => Some((model, *idx, 5, &[])),
873            // Kind 6, not 5: q4tp's nibble stride and scale planes differ,
874            // and feeding them to the q4t kernel is exactly the mistake that
875            // produced garbage when Q4Tiled shared kind 2 with Q4Block.
876            Self::Mapped {
877                model,
878                idx,
879                dtype: TensorDtype::Q4TiledP,
880                ..
881            } => Some((model, *idx, 6, &[])),
882            Self::Mapped {
883                model,
884                idx,
885                dtype: TensorDtype::Q4Block,
886                ..
887            } => Some((model, *idx, 2, &[])),
888            // q8_2f carries BOTH scale planes after the int8 body (rows
889            // f16, then cols f16), so the graph takes the whole tensor
890            // and the kernel reads them where they lie — no host-side
891            // prescale, which is what the per-op path does instead.
892            Self::Mapped {
893                model,
894                idx,
895                dtype: TensorDtype::Q8_2f,
896                ..
897            } => Some((model, *idx, 7, &[])),
898            Self::Mapped {
899                model,
900                idx,
901                dtype: TensorDtype::Q1T,
902                ..
903            } => Some((model, *idx, 3, &[])),
904            // Kind 9: the 2-bit plane on the q4tp ladder (dense FFN gate/up
905            // of the q2tp profile). Its own kernel — 8 bytes a group where
906            // q4tp has 16, and rung 0 is the exact zero.
907            Self::Mapped {
908                model,
909                idx,
910                dtype: TensorDtype::Q2TiledP,
911                ..
912            } => Some((model, *idx, 9, &[])),
913            _ => None,
914        }
915    }
916
917    /// Dense f32 view — only for owned tensors. Masked/sparse execution
918    /// paths require it; quantized weights don't support masks yet.
919    /// The parts of a `q8_2f` tensor: codes by row, row scales, column field.
920    pub fn q8_2f_codes(&self) -> Option<(&[u8], &[f32], &[f32])> {
921        match self {
922            Self::Mapped { dtype: TensorDtype::Q8_2f, row_scale, col_field, .. } => {
923                Some((self.quant_bytes(), row_scale, col_field))
924            }
925            _ => None,
926        }
927    }
928
929    pub fn as_f32(&self) -> Option<&[f32]> {
930        match self {
931            Self::F32 { data, .. } => Some(data),
932            Self::Mapped { .. } | Self::Blr(_) => None,
933        }
934    }
935
936    /// The tensor's bytes in the file, for a device that takes them as they are.
937    pub fn mapped_bytes(&self) -> Option<&[u8]> {
938        match self {
939            Self::Mapped { .. } => Some(self.quant_bytes()),
940            _ => None,
941        }
942    }
943
944    fn quant_bytes(&self) -> &[u8] {
945        match self {
946            Self::Mapped { model, idx, .. } => model.entry_bytes(&model.tensors[*idx]),
947            Self::Blr(t) => t.bytes(),
948            Self::F32 { .. } => unreachable!("quant_bytes on F32"),
949        }
950    }
951
952    /// Dequantize one row into `dst` (embedding lookup).
953    pub fn row_f32(&self, r: usize, dst: &mut [f32]) {
954        let cols = self.cols();
955        debug_assert_eq!(dst.len(), cols);
956        match self {
957            Self::F32 { data, .. } => dst.copy_from_slice(&data[r * cols..(r + 1) * cols]),
958            Self::Blr(t) => t.row_f32(r, dst),
959            Self::Mapped {
960                model,
961                idx,
962                dtype,
963                row_scale,
964                col_field,
965                vbit_offsets,
966                ..
967            } => {
968                if *dtype == TensorDtype::Q4Tiled {
969                    let bytes = self.quant_bytes();
970                    let gpr = cols / GROUP_SIZE;
971                    for gi in 0..gpr {
972                        let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
973                        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
974                        for (k, &b) in tile[2..].iter().enumerate() {
975                            dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
976                            dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
977                        }
978                    }
979                    if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
980                        crate::prism::inverse_embedding(model, dst);
981                    }
982                    return;
983                }
984                if *dtype == TensorDtype::Q4TiledP {
985                    let bytes = self.quant_bytes();
986                    let gpr = cols / GROUP_SIZE;
987                    let v = Q4tpView::new(bytes, self.rows(), cols);
988                    let mut sc = vec![0f32; gpr];
989                    v.scales_into(r, gpr, &mut sc);
990                    for gi in 0..gpr {
991                        let tile = &v.nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
992                        let s = sc[gi];
993                        for (k, &b) in tile.iter().enumerate() {
994                            dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
995                            dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
996                        }
997                    }
998                    if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
999                        crate::prism::inverse_embedding(model, dst);
1000                    }
1001                    return;
1002                }
1003                if *dtype == TensorDtype::Q2TiledP {
1004                    let bytes = self.quant_bytes();
1005                    let gpr = cols / GROUP_SIZE;
1006                    let v = Q4tpView::new_q2(bytes, self.rows(), cols);
1007                    let mut sc = vec![0f32; gpr];
1008                    v.scales_into(r, gpr, &mut sc);
1009                    for gi in 0..gpr {
1010                        let ch =
1011                            &v.nib[(r * gpr + gi) * Q2TP_CHUNK..(r * gpr + gi + 1) * Q2TP_CHUNK];
1012                        let s = sc[gi];
1013                        for (k, &b) in ch.iter().enumerate() {
1014                            for j in 0..4 {
1015                                let center = if crate::prism::is_affine_target(
1016                                    model,
1017                                    &model.tensors[*idx].name,
1018                                ) {
1019                                    1.0
1020                                } else {
1021                                    1.5
1022                                };
1023                                dst[gi * GROUP_SIZE + k * 4 + j] =
1024                                    (((b >> (2 * j)) & 3) as f32 - center) * s;
1025                            }
1026                        }
1027                    }
1028                    if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
1029                        crate::prism::inverse_embedding(model, dst);
1030                    }
1031                    return;
1032                }
1033                if *dtype == TensorDtype::Q4Block {
1034                    let (packed, scales) = q4_split(self.quant_bytes(), self.rows(), cols);
1035                    let gpr = cols / GROUP_SIZE;
1036                    for gi in 0..gpr {
1037                        let g = r * gpr + gi;
1038                        let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
1039                        for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
1040                            dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
1041                            dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
1042                        }
1043                    }
1044                    if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
1045                        crate::prism::inverse_embedding(model, dst);
1046                    }
1047                    return;
1048                }
1049                if *dtype == TensorDtype::Mxfp4 {
1050                    mxfp4::Mx::new(self.quant_bytes(), self.rows(), cols).row_f32(r, dst);
1051                    if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
1052                        crate::prism::inverse_embedding(model, dst);
1053                    }
1054                    return;
1055                }
1056                if *dtype == TensorDtype::Q1 {
1057                    let bytes = self.quant_bytes();
1058                    let gpr = cols / GROUP_SIZE;
1059                    for gi in 0..gpr {
1060                        let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
1061                        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
1062                        for (j, &b) in tile[2..].iter().enumerate() {
1063                            for k in 0..8 {
1064                                dst[gi * GROUP_SIZE + j * 8 + k] =
1065                                    (((b >> k) & 1) as f32 * 2.0 - 1.0) * s;
1066                            }
1067                        }
1068                    }
1069                    if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
1070                        crate::prism::inverse_embedding(model, dst);
1071                    }
1072                    return;
1073                }
1074                if *dtype == TensorDtype::Q1G {
1075                    // The tied Bonsai embedding is q1g too: one row = cols/8
1076                    // sign bytes + cols/128 scales, same ±s values as q1.
1077                    let bytes = self.quant_bytes();
1078                    let sc_off = cortiq_core::quant::q1g_scale_offset(self.rows(), cols);
1079                    let gpr = cols / GROUP_SIZE;
1080                    for gi in 0..gpr {
1081                        let t = r * gpr + gi;
1082                        let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, t));
1083                        for (j, &b) in bytes[t * 4..t * 4 + 4].iter().enumerate() {
1084                            for k in 0..8 {
1085                                dst[gi * GROUP_SIZE + j * 8 + k] =
1086                                    (((b >> k) & 1) as f32 * 2.0 - 1.0) * s;
1087                            }
1088                        }
1089                    }
1090                    if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
1091                        crate::prism::inverse_embedding(model, dst);
1092                    }
1093                    return;
1094                }
1095                if *dtype == TensorDtype::Q1T {
1096                    let bytes = self.quant_bytes();
1097                    let gpr = cols / GROUP_SIZE;
1098                    let base_len = self.rows() * gpr * cortiq_core::quant::Q1T_TILE;
1099                    for gi in 0..gpr {
1100                        let off = (r * gpr + gi) * cortiq_core::quant::Q1T_TILE;
1101                        let s = cortiq_core::quant::f16_to_f32(u16::from_le_bytes([
1102                            bytes[off],
1103                            bytes[off + 1],
1104                        ]));
1105                        let codes = &bytes[off + 2..off + cortiq_core::quant::Q1T_TILE];
1106                        for k in 0..GROUP_SIZE {
1107                            dst[gi * GROUP_SIZE + k] = match cortiq_core::quant::q1t_code(codes, k)
1108                            {
1109                                1 => s,
1110                                2 => -s,
1111                                _ => 0.0,
1112                            };
1113                        }
1114                    }
1115                    // Overlay
1116                    let rows = self.rows();
1117                    let entries = base_len + (rows + 1) * 4;
1118                    if entries <= bytes.len() {
1119                        let ptrs = &bytes[base_len..base_len + (rows + 1) * 4];
1120                        let r0 = u32::from_le_bytes([
1121                            ptrs[r * 4],
1122                            ptrs[r * 4 + 1],
1123                            ptrs[r * 4 + 2],
1124                            ptrs[r * 4 + 3],
1125                        ]) as usize;
1126                        let r1 = u32::from_le_bytes([
1127                            ptrs[(r + 1) * 4],
1128                            ptrs[(r + 1) * 4 + 1],
1129                            ptrs[(r + 1) * 4 + 2],
1130                            ptrs[(r + 1) * 4 + 3],
1131                        ]) as usize;
1132                        let off = entries + r0 * 4;
1133                        for i in 0..r1 - r0 {
1134                            let item = &bytes[off + i * 4..off + i * 4 + 4];
1135                            let c = u16::from_le_bytes([item[0], item[1]]) as usize;
1136                            let v = cortiq_core::quant::f16_to_f32(u16::from_le_bytes([
1137                                item[2], item[3],
1138                            ]));
1139                            if c < cols {
1140                                dst[c] = v;
1141                            }
1142                        }
1143                    }
1144                    if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
1145                        crate::prism::inverse_embedding(model, dst);
1146                    }
1147                    return;
1148                }
1149                if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1150                    let bytes = self.quant_bytes();
1151                    let rows = self.rows();
1152                    let ng = cols / GROUP_SIZE;
1153                    let bits = &bytes[..rows];
1154                    let sc_off = rows;
1155                    // Precomputed at load — embedding lookup used to scan
1156                    // the bit-widths of every preceding row (O(token_id)).
1157                    let off = vbit_offsets[r];
1158                    let b = bits[r] as usize;
1159                    let l = ((1usize << (b - 1)) - 1) as f32;
1160                    let data = &bytes[off..];
1161                    let (mut acc, mut nbits, mut byte_idx) = (0u64, 0usize, 0usize);
1162                    for (i, d) in dst.iter_mut().enumerate() {
1163                        while nbits < b {
1164                            acc = (acc << 8) | data[byte_idx] as u64;
1165                            byte_idx += 1;
1166                            nbits += 8;
1167                        }
1168                        let u = ((acc >> (nbits - b)) & ((1u64 << b) - 1)) as f32;
1169                        nbits -= b;
1170                        let so = (r * ng + i / GROUP_SIZE) * 2;
1171                        let sv = f16_to_f32(u16::from_le_bytes([
1172                            bytes[sc_off + so],
1173                            bytes[sc_off + so + 1],
1174                        ]));
1175                        *d = (u - l) * sv;
1176                    }
1177                    if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
1178                        crate::prism::inverse_embedding(model, dst);
1179                    }
1180                    return;
1181                }
1182                let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
1183                let s = row_scale[r];
1184                match dtype {
1185                    TensorDtype::Q8Row => {
1186                        for (d, &b) in dst.iter_mut().zip(q) {
1187                            *d = (b as i8) as f32 * s;
1188                        }
1189                    }
1190                    TensorDtype::Q8_2f => {
1191                        for (i, (d, &b)) in dst.iter_mut().zip(q).enumerate() {
1192                            *d = (b as i8) as f32 * s * col_field[i];
1193                        }
1194                    }
1195                    _ => unreachable!(),
1196                }
1197                if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
1198                    crate::prism::inverse_embedding(model, dst);
1199                }
1200            }
1201        }
1202    }
1203
1204    /// Can this tensor's columns be read cheaply (for sparse down_proj)?
1205    /// True for F32/Q8Row/Q8_2f (per-row scale, direct strided access);
1206    /// false for group-packed q4/vbit (column access would unpack whole
1207    /// groups — sparse execution falls back to f32 for those).
1208    pub fn sparse_col_ok(&self) -> bool {
1209        match self {
1210            Self::F32 { .. } => true,
1211            Self::Blr(_) => false,
1212            Self::Mapped { dtype, .. } => {
1213                matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
1214            }
1215        }
1216    }
1217
1218    /// down_proj [hidden, inter]: accumulate `w · col(c)` into `out`
1219    /// [hidden] — reads ONLY column `c` (one neuron) from the mmap,
1220    /// no full-matrix dequant. `out[k] += w · down[k, c]`.
1221    pub fn add_col_scaled(&self, c: usize, w: f32, out: &mut [f32]) {
1222        let inter = self.cols();
1223        let hidden = self.rows();
1224        debug_assert_eq!(out.len(), hidden);
1225        match self {
1226            Self::F32 { data, .. } => {
1227                for (k, o) in out.iter_mut().enumerate() {
1228                    *o += w * data[k * inter + c];
1229                }
1230            }
1231            Self::Blr(t) => t.add_col_scaled(c, w, out),
1232            Self::Mapped {
1233                dtype,
1234                row_scale,
1235                col_field,
1236                ..
1237            } => {
1238                let q = self.quant_bytes();
1239                let colf = if *dtype == TensorDtype::Q8_2f {
1240                    col_field[c]
1241                } else {
1242                    1.0
1243                };
1244                let wc = w * colf;
1245                for (k, o) in out.iter_mut().enumerate() {
1246                    let b = q[k * inter + c] as i8 as f32;
1247                    *o += wc * b * row_scale[k];
1248                }
1249            }
1250        }
1251    }
1252
1253    /// Touch the head of row `r` so the DRAM latency of the next
1254    /// neuron's weights overlaps the current one's arithmetic.
1255    ///
1256    /// Scattered rows are what per-token sparsity reads, and a 2 KB
1257    /// stride is past what the hardware prefetcher follows: without this
1258    /// every row starts with a cold miss that nothing hides. One touch
1259    /// per 512 bytes is enough — the rest of the row is a sequential run
1260    /// the prefetcher does pick up.
1261    #[inline]
1262    pub fn prefetch_row(&self, r: usize) {
1263        let Self::Mapped { dtype, .. } = self else {
1264            return;
1265        };
1266        if !matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f) {
1267            return;
1268        }
1269        let cols = self.cols();
1270        let q = self.quant_bytes();
1271        let (a, b) = (r * cols, (r + 1) * cols);
1272        if b > q.len() {
1273            return;
1274        }
1275        let mut j = a;
1276        while j < b {
1277            unsafe { std::ptr::read_volatile(q.as_ptr().add(j)) };
1278            j += 512;
1279        }
1280    }
1281
1282    /// `out += w · row(r)` — the transposed twin of `add_col_scaled`.
1283    ///
1284    /// A neuron's `down` weights are a COLUMN of `[hidden, inter]`, and a
1285    /// column is strided: reading one costs a cache line per element, so
1286    /// per-neuron dynamic sparsity saves arithmetic and no bytes. Stored
1287    /// transposed (`down_proj.t.weight`, `[inter, hidden]`) the same
1288    /// weights are a contiguous ROW, and this accumulate reads exactly
1289    /// the neurons the token asked for.
1290    pub fn add_row_scaled(&self, r: usize, w: f32, out: &mut [f32], scratch: &mut [f32]) {
1291        let cols = self.cols();
1292        debug_assert_eq!(out.len(), cols);
1293        match self {
1294            Self::F32 { data, .. } => {
1295                let row = &data[r * cols..(r + 1) * cols];
1296                for (o, v) in out.iter_mut().zip(row) {
1297                    *o += w * v;
1298                }
1299            }
1300            Self::Blr(t) => {
1301                let tmp = &mut scratch[..cols];
1302                t.row_f32(r, tmp);
1303                for (o, v) in out.iter_mut().zip(tmp.iter()) {
1304                    *o += w * v;
1305                }
1306            }
1307            Self::Mapped {
1308                dtype,
1309                row_scale,
1310                col_field,
1311                ..
1312            } => match dtype {
1313                TensorDtype::Q8Row => {
1314                    let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
1315                    let ws = w * row_scale[r];
1316                    let row: &[i8] =
1317                        unsafe { std::slice::from_raw_parts(q.as_ptr() as *const i8, q.len()) };
1318                    axpy_i8_f32(out, row, ws);
1319                }
1320                TensorDtype::Q8_2f => {
1321                    let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
1322                    let ws = w * row_scale[r];
1323                    for ((o, b), c) in out.iter_mut().zip(q).zip(col_field) {
1324                        *o += ws * c * (*b as i8 as f32);
1325                    }
1326                }
1327                _ => {
1328                    self.row_f32(r, scratch);
1329                    for (o, v) in out.iter_mut().zip(scratch.iter()) {
1330                        *o += w * v;
1331                    }
1332                }
1333            },
1334        }
1335    }
1336
1337    /// Dot of row `r` with `x` (gate/up active-neuron path). Reads only
1338    /// row `r` from the mmap — no full dequant. q4/vbit dequant the row
1339    /// into `scratch` first (rare for active-FFN weights).
1340    pub fn row_dot(&self, r: usize, x: &[f32], scratch: &mut [f32]) -> f32 {
1341        let cols = self.cols();
1342        match self {
1343            Self::F32 { data, .. } => {
1344                let row = &data[r * cols..(r + 1) * cols];
1345                row.iter().zip(x).map(|(w, v)| w * v).sum()
1346            }
1347            Self::Blr(t) => {
1348                let tmp = &mut scratch[..cols];
1349                t.row_f32(r, tmp);
1350                tmp.iter().zip(x).map(|(w, v)| w * v).sum()
1351            }
1352            Self::Mapped {
1353                model,
1354                idx,
1355                dtype,
1356                row_scale,
1357                col_field,
1358                ..
1359            } => {
1360                let prism_forward =
1361                    crate::prism::is_forward_weight(model, &model.tensors[*idx].name);
1362                if prism_forward {
1363                    let transformed = crate::prism::forward(model, &x[..cols]);
1364                    let gpr = cols / GROUP_SIZE;
1365                    match dtype {
1366                        TensorDtype::Q2TiledP => {
1367                            let v = Q4tpView::new_q2(self.quant_bytes(), self.rows(), cols);
1368                            let mut sc = vec![0f32; gpr];
1369                            v.scales_into(r, gpr, &mut sc);
1370                            if crate::prism::is_affine_target(model, &model.tensors[*idx].name) {
1371                                return q2tp_affine_row_exact(v.nib, r, gpr, &transformed, &sc);
1372                            }
1373                            return q2tp_row_exact(v.nib, r, gpr, &transformed, &sc);
1374                        }
1375                        _ => {
1376                            self.row_f32(r, scratch);
1377                            return scratch.iter().zip(&transformed).map(|(w, v)| w * v).sum();
1378                        }
1379                    }
1380                }
1381                match dtype {
1382                    TensorDtype::Q8Row => {
1383                        let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
1384                        dot_i8_f32(q, x) * row_scale[r]
1385                    }
1386                    TensorDtype::Q8_2f => {
1387                        let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
1388                        dot_i8_col_f32(q, x, col_field) * row_scale[r]
1389                    }
1390                    _ => {
1391                        self.row_f32(r, scratch);
1392                        scratch.iter().zip(x).map(|(w, v)| w * v).sum()
1393                    }
1394                }
1395            }
1396        }
1397    }
1398
1399    /// `out = W · x` (row-major). F32 delegates to the historical
1400    /// bit-exact path; Mapped runs the fused int8 kernel.
1401    pub fn matvec(&self, x: &[f32], out: &mut [f32], pool: Option<&Pool>) {
1402        if self.exact_cpu_default() {
1403            return float_activations_scope(|| self.matvec(x, out, pool));
1404        }
1405        match self {
1406            // NOTE: `out.len()` DRIVES this arm — it computes that many rows,
1407            // and `x.len()` is the stride. A short `out` is legitimate here,
1408            // which is why the check below lives in the Mapped arm only.
1409            Self::F32 { data, .. } => {
1410                if !crate::f32_backend::matvec(data, x, out) {
1411                    matvec_rows(pool, data, x, out);
1412                }
1413            }
1414            Self::Blr(t) => {
1415                if crate::gpu::enabled_here()
1416                    && crate::blr::gpu_wanted()
1417                    && crate::gpu::blr_matmat(t, &x[..t.cols], 1, &mut out[..t.rows])
1418                {
1419                    return;
1420                }
1421                t.matvec(x, out, pool)
1422            }
1423            Self::Mapped {
1424                model,
1425                idx,
1426                dtype,
1427                rows,
1428                cols,
1429                row_scale,
1430                col_field,
1431                vbit_offsets,
1432                repack,
1433            } => {
1434                let _ = (model, idx);
1435                // Every kernel below writes `rows` entries through a raw
1436                // pointer, so a short `out` is an out-of-bounds WRITE, not a
1437                // wrong answer: it scribbles on the allocator's metadata and
1438                // the process aborts much later, somewhere innocent
1439                // (`double free or corruption`, `corrupted double-linked
1440                // list`). The debug_assert two of the kernels carried is
1441                // compiled out of the release — exactly the build where it
1442                // matters. Fail here instead, while the caller is still on
1443                // the stack to be named.
1444                assert!(
1445                    out.len() >= *rows && x.len() >= *cols,
1446                    "matvec {rows}x{cols}: out {} (need {rows}), x {} (need {cols})",
1447                    out.len(),
1448                    x.len(),
1449                );
1450                let prism_forward =
1451                    crate::prism::is_forward_weight(model, &model.tensors[*idx].name);
1452                if *dtype == TensorDtype::Q2TiledP
1453                    && std::env::var("CMF_Q2TP_TRACE").as_deref() == Ok("1")
1454                {
1455                    use std::sync::atomic::{AtomicUsize, Ordering};
1456                    static N: AtomicUsize = AtomicUsize::new(0);
1457                    let n = N.fetch_add(1, Ordering::Relaxed);
1458                    if n < 128 {
1459                        eprintln!(
1460                            "q2tp-dispatch #{n} name={} prism={} rows={} cols={} gpu={} optin={} layer={}",
1461                            model.tensors[*idx].name,
1462                            prism_forward,
1463                            rows,
1464                            cols,
1465                            crate::gpu::enabled_here(),
1466                            crate::gpu::q2tp_gpu_opt_in(),
1467                            crate::gpu::cur_layer(),
1468                        );
1469                    }
1470                }
1471                // Prism stores every manifest-listed forward matrix in the
1472                // signed-Hadamard basis.  The q2tp WGSL path receives that
1473                // transformed vector and an explicit affine bit; codecs
1474                // without a descriptor-aware kernel remain on CPU below.
1475                if prism_forward {
1476                    let transformed = crate::prism::forward(model, &x[..*cols]);
1477                    match dtype {
1478                        TensorDtype::Q4Block => {
1479                            q4matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1480                        }
1481                        TensorDtype::Q4Tiled => {
1482                            q4t_matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1483                        }
1484                        TensorDtype::Q4TiledP => {
1485                            q4tp_matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1486                        }
1487                        TensorDtype::Q2TiledP => {
1488                            let affine =
1489                                crate::prism::is_affine_target(model, &model.tensors[*idx].name);
1490                            if *rows * *cols >= 8_388_608
1491                                && crate::gpu::enabled_here()
1492                                && crate::gpu::q2tp_gpu_opt_in()
1493                            {
1494                                let gpu_ok = if affine {
1495                                    crate::gpu::q2tp_affine_matvec(
1496                                        model,
1497                                        *idx,
1498                                        &transformed,
1499                                        *rows,
1500                                        *cols,
1501                                        out,
1502                                    )
1503                                } else {
1504                                    crate::gpu::q2tp_matvec(
1505                                        model,
1506                                        *idx,
1507                                        &transformed,
1508                                        *rows,
1509                                        *cols,
1510                                        out,
1511                                    )
1512                                };
1513                                if gpu_ok {
1514                                    return;
1515                                }
1516                            }
1517                            if affine {
1518                                q2tp_affine_matvec(
1519                                    self.quant_bytes(),
1520                                    &transformed,
1521                                    *rows,
1522                                    *cols,
1523                                    out,
1524                                    pool,
1525                                )
1526                            } else {
1527                                q2tp_matvec(
1528                                    self.quant_bytes(),
1529                                    &transformed,
1530                                    *rows,
1531                                    *cols,
1532                                    out,
1533                                    pool,
1534                                )
1535                            }
1536                        }
1537                        TensorDtype::Q1 | TensorDtype::Q1G => bin_matvec(
1538                            self.bin_w().expect("binary dtype"),
1539                            &transformed,
1540                            *rows,
1541                            *cols,
1542                            out,
1543                            pool,
1544                        ),
1545                        TensorDtype::Q1T => {
1546                            q1t_matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1547                        }
1548                        TensorDtype::Mxfp4 => {
1549                            mxfp4::matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1550                        }
1551                        TensorDtype::Vbit | TensorDtype::VbitRo => vbitmatvec(
1552                            self.quant_bytes(),
1553                            vbit_offsets,
1554                            &transformed,
1555                            *rows,
1556                            *cols,
1557                            out,
1558                            pool,
1559                        ),
1560                        TensorDtype::Q8Row | TensorDtype::Q8_2f => qmatvec(
1561                            self.quant_bytes(),
1562                            repack,
1563                            row_scale,
1564                            &transformed,
1565                            col_field,
1566                            *dtype,
1567                            *rows,
1568                            *cols,
1569                            out,
1570                            pool,
1571                        ),
1572                        _ => unreachable!("unsupported mapped Prism dtype {dtype:?}"),
1573                    }
1574                    return;
1575                }
1576                if *dtype == TensorDtype::Q4Block {
1577                    // GPU route (wgpu q4b kernel) for large q4_block matvecs —
1578                    // gives NVIDIA/AMD/Intel q4 models a GPU path. Probe keeps
1579                    // the winner; Metal returns false → the CPU kernel below.
1580                    if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1581                        let t0 = std::time::Instant::now();
1582                        match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
1583                            crate::gpu::ProbeArm::Gpu => {
1584                                if crate::gpu::q4b_matvec(model, *idx, x, *rows, *cols, out) {
1585                                    crate::gpu::probe_record(
1586                                        crate::gpu::OpClass::Matvec,
1587                                        true,
1588                                        t0.elapsed(),
1589                                    );
1590                                    return;
1591                                }
1592                            }
1593                            crate::gpu::ProbeArm::CpuTimed => {
1594                                q4matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1595                                crate::gpu::probe_record(
1596                                    crate::gpu::OpClass::Matvec,
1597                                    false,
1598                                    t0.elapsed(),
1599                                );
1600                                return;
1601                            }
1602                            crate::gpu::ProbeArm::Cpu => {}
1603                        }
1604                    }
1605                    q4matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1606                    return;
1607                }
1608                if *dtype == TensorDtype::Q4Tiled {
1609                    // GPU route for large q4t matvecs — the lm_head class,
1610                    // same shape as the q4tp arm below. The probe keeps the
1611                    // winner; a backend without the kernel refuses and the
1612                    // CPU path stays.
1613                    if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1614                        let t0 = std::time::Instant::now();
1615                        let cls = crate::gpu::matvec_class(*rows, *cols);
1616                        match crate::gpu::probe_arm(cls) {
1617                            crate::gpu::ProbeArm::Gpu => {
1618                                if crate::gpu::q4t_matvec(model, *idx, x, *rows, *cols, out) {
1619                                    crate::gpu::probe_record(cls, true, t0.elapsed());
1620                                    return;
1621                                }
1622                            }
1623                            crate::gpu::ProbeArm::CpuTimed => {
1624                                q4t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1625                                crate::gpu::probe_record(cls, false, t0.elapsed());
1626                                return;
1627                            }
1628                            crate::gpu::ProbeArm::Cpu => {}
1629                        }
1630                    }
1631                    q4t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1632                    return;
1633                }
1634                if *dtype == TensorDtype::Q4TiledP {
1635                    // GPU route for large q4tp matvecs — the lm_head class.
1636                    // On a q4tp checkpoint the head is the biggest single
1637                    // host matvec left in the decode step, and the batched
1638                    // kernel at b=1 already exists on both backends. Probe
1639                    // keeps the winner, same as q4_block above.
1640                    if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1641                        let t0 = std::time::Instant::now();
1642                        let cls = crate::gpu::matvec_class(*rows, *cols);
1643                        match crate::gpu::probe_arm(cls) {
1644                            crate::gpu::ProbeArm::Gpu => {
1645                                if crate::gpu::q4tp_matvec(model, *idx, x, *rows, *cols, out) {
1646                                    crate::gpu::probe_record(cls, true, t0.elapsed());
1647                                    return;
1648                                }
1649                            }
1650                            crate::gpu::ProbeArm::CpuTimed => {
1651                                q4tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1652                                crate::gpu::probe_record(cls, false, t0.elapsed());
1653                                return;
1654                            }
1655                            crate::gpu::ProbeArm::Cpu => {}
1656                        }
1657                    }
1658                    q4tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1659                    return;
1660                }
1661                if *dtype == TensorDtype::Q2TiledP {
1662                    q2tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1663                    return;
1664                }
1665                if *dtype == TensorDtype::Mxfp4 {
1666                    self.mxfp4_matvec(x, out, pool);
1667                    return;
1668                }
1669                if *dtype == TensorDtype::Q1 || *dtype == TensorDtype::Q1G {
1670                    // GPU route for large q1 matvecs (out_proj / lm_head
1671                    // class): the CPU q1 kernel is load-port-bound at
1672                    // ~4 GB/s/core, the GPU one is bandwidth-bound — the
1673                    // probe measures both arms and keeps the winner.
1674                    let w = self.bin_w().expect("binary dtype");
1675                    if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1676                        let t0 = std::time::Instant::now();
1677                        let arm = if crate::gpu::q1_force() {
1678                            crate::gpu::ProbeArm::Gpu
1679                        } else {
1680                            crate::gpu::probe_arm(crate::gpu::OpClass::Matvec)
1681                        };
1682                        match arm {
1683                            crate::gpu::ProbeArm::Gpu => {
1684                                let ok = if *dtype == TensorDtype::Q1G {
1685                                    crate::gpu::q1g_matvec(model, *idx, x, *rows, *cols, out)
1686                                } else {
1687                                    crate::gpu::q1_matvec(model, *idx, x, *rows, *cols, out)
1688                                };
1689                                if ok {
1690                                    crate::gpu::probe_record(
1691                                        crate::gpu::OpClass::Matvec,
1692                                        true,
1693                                        t0.elapsed(),
1694                                    );
1695                                    return;
1696                                }
1697                            }
1698                            crate::gpu::ProbeArm::CpuTimed => {
1699                                bin_matvec(w, x, *rows, *cols, out, pool);
1700                                crate::gpu::probe_record(
1701                                    crate::gpu::OpClass::Matvec,
1702                                    false,
1703                                    t0.elapsed(),
1704                                );
1705                                return;
1706                            }
1707                            crate::gpu::ProbeArm::Cpu => {}
1708                        }
1709                    }
1710                    bin_matvec(w, x, *rows, *cols, out, pool);
1711                    return;
1712                }
1713                if *dtype == TensorDtype::Q1T {
1714                    // GPU route for large q1t matvecs: the ternary BASE dot runs
1715                    // on the GPU (load-port-bound on CPU, like q1), then the
1716                    // sparse overlay is added on the CPU. Probe keeps the winner.
1717                    if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1718                        let t0 = std::time::Instant::now();
1719                        match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
1720                            crate::gpu::ProbeArm::Gpu => {
1721                                if crate::gpu::q1t_matvec(model, *idx, x, *rows, *cols, out) {
1722                                    q1t_add_overlay(self.quant_bytes(), x, *rows, *cols, out, pool);
1723                                    crate::gpu::probe_record(
1724                                        crate::gpu::OpClass::Matvec,
1725                                        true,
1726                                        t0.elapsed(),
1727                                    );
1728                                    return;
1729                                }
1730                            }
1731                            crate::gpu::ProbeArm::CpuTimed => {
1732                                q1t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1733                                crate::gpu::probe_record(
1734                                    crate::gpu::OpClass::Matvec,
1735                                    false,
1736                                    t0.elapsed(),
1737                                );
1738                                return;
1739                            }
1740                            crate::gpu::ProbeArm::Cpu => {}
1741                        }
1742                    }
1743                    q1t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1744                    return;
1745                }
1746                if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1747                    vbitmatvec(self.quant_bytes(), vbit_offsets, x, *rows, *cols, out, pool);
1748                    return;
1749                }
1750                let xs = prescale(x, col_field, *dtype);
1751                // D5: large q8 matrices (lm_head-class) — hybrid
1752                // CPU∥GPU: split the rows, both sides compute
1753                // SIMULTANEOUSLY (same math, shared prescale).
1754                // GPU share: CMF_GPU_SPLIT (0..1, default 0.5).
1755                if *rows >= crate::gpu::min_rows()
1756                    && matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
1757                    && gpu_lmhead_enabled()
1758                    && crate::gpu::enabled_here()
1759                {
1760                    // Runtime probe: alternate the hybrid against the
1761                    // pure-CPU matvec, keep whichever is faster HERE.
1762                    let t0 = std::time::Instant::now();
1763                    match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
1764                        crate::gpu::ProbeArm::Gpu => {}
1765                        crate::gpu::ProbeArm::CpuTimed => {
1766                            qmatvec(
1767                                self.quant_bytes(),
1768                                repack,
1769                                row_scale,
1770                                x,
1771                                col_field,
1772                                *dtype,
1773                                *rows,
1774                                *cols,
1775                                out,
1776                                pool,
1777                            );
1778                            crate::gpu::probe_record(
1779                                crate::gpu::OpClass::Matvec,
1780                                false,
1781                                t0.elapsed(),
1782                            );
1783                            return;
1784                        }
1785                        crate::gpu::ProbeArm::Cpu => {
1786                            qmatvec(
1787                                self.quant_bytes(),
1788                                repack,
1789                                row_scale,
1790                                x,
1791                                col_field,
1792                                *dtype,
1793                                *rows,
1794                                *cols,
1795                                out,
1796                                pool,
1797                            );
1798                            return;
1799                        }
1800                    }
1801                    let frac = gpu_split_frac();
1802                    let cpu_rows = ((*rows as f32) * (1.0 - frac)) as usize;
1803                    let (out_cpu, out_gpu) = out.split_at_mut(cpu_rows);
1804                    let bytes = self.quant_bytes();
1805                    let ok = std::thread::scope(|sc| {
1806                        let g = sc.spawn(|| {
1807                            crate::gpu::q8_matvec_range(
1808                                model,
1809                                *idx,
1810                                cpu_rows,
1811                                &row_scale[cpu_rows..],
1812                                &xs,
1813                                *rows - cpu_rows,
1814                                *cols,
1815                                out_gpu,
1816                            )
1817                        });
1818                        if cpu_rows > 0 {
1819                            // Repack prefix covers the full groups of the
1820                            // CPU half (the split starts at row 0).
1821                            let rep_cpu = if repack.is_empty() {
1822                                &[][..]
1823                            } else {
1824                                &repack[..(cpu_rows / 4) * 4 * *cols]
1825                            };
1826                            qmatvec(
1827                                &bytes[..cpu_rows * *cols],
1828                                rep_cpu,
1829                                &row_scale[..cpu_rows],
1830                                x,
1831                                col_field,
1832                                *dtype,
1833                                cpu_rows,
1834                                *cols,
1835                                out_cpu,
1836                                pool,
1837                            );
1838                        }
1839                        g.join().unwrap_or(false)
1840                    });
1841                    if ok {
1842                        crate::gpu::probe_record(crate::gpu::OpClass::Matvec, true, t0.elapsed());
1843                        return;
1844                    }
1845                    // GPU failed — CPU finishes its half (rows rebased —
1846                    // group offsets don't line up, mmap layout only).
1847                    qmatvec(
1848                        &bytes[cpu_rows * *cols..(*rows) * *cols],
1849                        &[],
1850                        &row_scale[cpu_rows..],
1851                        x,
1852                        col_field,
1853                        *dtype,
1854                        *rows - cpu_rows,
1855                        *cols,
1856                        out_gpu,
1857                        pool,
1858                    );
1859                    return;
1860                }
1861                qmatvec(
1862                    self.quant_bytes(),
1863                    repack,
1864                    row_scale,
1865                    x,
1866                    col_field,
1867                    *dtype,
1868                    *rows,
1869                    *cols,
1870                    out,
1871                    pool,
1872                );
1873            }
1874        }
1875    }
1876
1877    /// Fused two-input matvec (MTP verify pair): weights streamed once.
1878    pub fn matvec2(
1879        &self,
1880        x1: &[f32],
1881        x2: &[f32],
1882        o1: &mut [f32],
1883        o2: &mut [f32],
1884        pool: Option<&Pool>,
1885    ) {
1886        if self.exact_cpu_default() {
1887            return float_activations_scope(|| self.matvec2(x1, x2, o1, o2, pool));
1888        }
1889        match self {
1890            Self::F32 { data, .. } => matvec_rows2(pool, data, x1, x2, o1, o2),
1891            Self::Blr(_) => {
1892                self.matvec(x1, o1, pool);
1893                self.matvec(x2, o2, pool);
1894            }
1895            Self::Mapped {
1896                model,
1897                idx,
1898                dtype,
1899                rows,
1900                cols,
1901                row_scale,
1902                col_field,
1903                vbit_offsets,
1904                ..
1905            } => {
1906                if crate::prism::is_forward_weight(model, &model.tensors[*idx].name) {
1907                    let tx1 = crate::prism::forward(model, &x1[..*cols]);
1908                    let tx2 = crate::prism::forward(model, &x2[..*cols]);
1909                    match dtype {
1910                        TensorDtype::Q4Block => {
1911                            q4matvec2(self.quant_bytes(), &tx1, &tx2, *rows, *cols, o1, o2, pool)
1912                        }
1913                        TensorDtype::Q4Tiled => {
1914                            q4t_matvec2(self.quant_bytes(), &tx1, &tx2, *rows, *cols, o1, o2, pool)
1915                        }
1916                        TensorDtype::Q4TiledP => {
1917                            q4tp_matvec2(self.quant_bytes(), &tx1, &tx2, *rows, *cols, o1, o2, pool)
1918                        }
1919                        TensorDtype::Q2TiledP => {
1920                            if crate::prism::is_affine_target(model, &model.tensors[*idx].name) {
1921                                q2tp_affine_matvec2(
1922                                    self.quant_bytes(),
1923                                    &tx1,
1924                                    &tx2,
1925                                    *rows,
1926                                    *cols,
1927                                    o1,
1928                                    o2,
1929                                    pool,
1930                                )
1931                            } else {
1932                                q2tp_matvec2(
1933                                    self.quant_bytes(),
1934                                    &tx1,
1935                                    &tx2,
1936                                    *rows,
1937                                    *cols,
1938                                    o1,
1939                                    o2,
1940                                    pool,
1941                                )
1942                            }
1943                        }
1944                        TensorDtype::Q1 | TensorDtype::Q1G => bin_matvec2(
1945                            self.bin_w().expect("binary dtype"),
1946                            &tx1,
1947                            &tx2,
1948                            *rows,
1949                            *cols,
1950                            o1,
1951                            o2,
1952                            pool,
1953                        ),
1954                        TensorDtype::Q1T => {
1955                            q1t_matvec2(self.quant_bytes(), &tx1, &tx2, *rows, *cols, o1, o2, pool)
1956                        }
1957                        TensorDtype::Mxfp4 => mxfp4::matvec2(
1958                            self.quant_bytes(),
1959                            &tx1,
1960                            &tx2,
1961                            *rows,
1962                            *cols,
1963                            o1,
1964                            o2,
1965                            pool,
1966                        ),
1967                        TensorDtype::Vbit | TensorDtype::VbitRo => vbitmatvec2(
1968                            self.quant_bytes(),
1969                            vbit_offsets,
1970                            &tx1,
1971                            &tx2,
1972                            *rows,
1973                            *cols,
1974                            o1,
1975                            o2,
1976                            pool,
1977                        ),
1978                        TensorDtype::Q8Row | TensorDtype::Q8_2f => qmatvec2(
1979                            self.quant_bytes(),
1980                            row_scale,
1981                            &tx1,
1982                            &tx2,
1983                            col_field,
1984                            *dtype,
1985                            *rows,
1986                            *cols,
1987                            o1,
1988                            o2,
1989                            pool,
1990                        ),
1991                        _ => unreachable!("unsupported mapped Prism dtype {dtype:?}"),
1992                    }
1993                    return;
1994                }
1995                if *dtype == TensorDtype::Q4Block {
1996                    q4matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1997                    return;
1998                }
1999                if *dtype == TensorDtype::Q4Tiled {
2000                    q4t_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
2001                    return;
2002                }
2003                if *dtype == TensorDtype::Q4TiledP {
2004                    q4tp_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
2005                    return;
2006                }
2007                if *dtype == TensorDtype::Q2TiledP {
2008                    q2tp_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
2009                    return;
2010                }
2011                if *dtype == TensorDtype::Mxfp4 {
2012                    mxfp4::matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
2013                    return;
2014                }
2015                if let Some(w) = self.bin_w() {
2016                    bin_matvec2(w, x1, x2, *rows, *cols, o1, o2, pool);
2017                    return;
2018                }
2019                if *dtype == TensorDtype::Q1T {
2020                    // Fused ternary pair: one row pass, the register
2021                    // unpack shared across both streams on ARM. (Q1T
2022                    // lacks a row_scale array — scales live inline in
2023                    // the tiles — so it must not fall through to the
2024                    // q8 qmatvec2 below.)
2025                    q1t_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
2026                    return;
2027                }
2028                if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
2029                    vbitmatvec2(
2030                        self.quant_bytes(),
2031                        vbit_offsets,
2032                        x1,
2033                        x2,
2034                        *rows,
2035                        *cols,
2036                        o1,
2037                        o2,
2038                        pool,
2039                    );
2040                    return;
2041                }
2042                qmatvec2(
2043                    self.quant_bytes(),
2044                    row_scale,
2045                    x1,
2046                    x2,
2047                    col_field,
2048                    *dtype,
2049                    *rows,
2050                    *cols,
2051                    o1,
2052                    o2,
2053                    pool,
2054                );
2055            }
2056        }
2057    }
2058}
2059
2060impl QTensor {
2061    /// Batched matvec (prefill-GEMM): xs — row-major [b, cols],
2062    /// out — row-major [b, rows]. Element-wise semantics are IDENTICAL
2063    /// to b matvec calls (same dot kernels in the same order); the win —
2064    /// the weight row streams from DRAM once per batch, not b times.
2065    /// `(model, index)` when this is a memory-mapped q4tp tensor — the
2066    /// identity a device-resident chain needs to hand `tp_matmat` the
2067    /// weight without going through this struct's own dispatch.
2068    pub fn q4tp_mapped(&self) -> Option<(&std::sync::Arc<CmfModel>, usize)> {
2069        if self.has_prism_contract() {
2070            return None;
2071        }
2072        match self {
2073            Self::Mapped {
2074                model, idx, dtype, ..
2075            } if *dtype == TensorDtype::Q4TiledP => Some((model, *idx)),
2076            _ => None,
2077        }
2078    }
2079
2080    pub fn matmat(&self, xs_all: &[f32], b: usize, out: &mut [f32], pool: Option<&Pool>) {
2081        if self.exact_cpu_default() {
2082            return float_activations_scope(|| self.matmat(xs_all, b, out, pool));
2083        }
2084        let cols = self.cols();
2085        let rows = self.rows();
2086        debug_assert_eq!(xs_all.len(), b * cols);
2087        debug_assert_eq!(out.len(), b * rows);
2088        let _prof = crate::cpuprof::time(crate::cpuprof::Slot::Matmat);
2089        // GPTQ calibration: fold this layer's inputs into its Hessian. Only
2090        // Mapped tensors carry a directory name; the check is a relaxed
2091        // atomic load, free when not calibrating.
2092        if crate::gptq_capture::capturing() {
2093            if let Self::Mapped { model, idx, .. } = self {
2094                crate::gptq_capture::accumulate(&model.tensors[*idx].name, xs_all, b, cols);
2095            }
2096            if let Self::Blr(t) = self {
2097                crate::gptq_capture::accumulate(t.name(), xs_all, b, cols);
2098            }
2099        }
2100        match self {
2101            Self::Blr(t) => {
2102                if crate::gpu::enabled_here()
2103                    && crate::blr::gpu_wanted()
2104                    && crate::gpu::blr_matmat(t, &xs_all[..b * cols], b, &mut out[..b * rows])
2105                {
2106                    return;
2107                }
2108                t.matmat(xs_all, b, out, pool)
2109            }
2110            Self::F32 { data, .. } => {
2111                if crate::f32_backend::matmat(data, xs_all, b, rows, cols, out) {
2112                    return;
2113                }
2114                let out_addr = SendMut(out.as_mut_ptr());
2115                // Over the batch, four positions a row at a time: four
2116                // independent chains, each still summing j = 0..cols in the
2117                // scalar order (bit-identical to the one-at-a-time loop and
2118                // to `matvec_rows`). The batch is the wide axis of a prefill;
2119                // the rows may be few (Spark-X2.5's 16-row g_proj spent
2120                // 17 ms a layer here split over its rows).
2121                let kernel = |start: usize, end: usize, r0: usize, r1: usize| {
2122                    let mut bi = start;
2123                    while bi < end {
2124                        let n = (end - bi).min(4);
2125                        for o in r0..r1 {
2126                            let row = &data[o * cols..(o + 1) * cols];
2127                            if n == 4 {
2128                                let x = |k: usize| &xs_all[(bi + k) * cols..(bi + k + 1) * cols];
2129                                let (x0, x1, x2, x3) = (x(0), x(1), x(2), x(3));
2130                                let (mut a0, mut a1, mut a2, mut a3) = (0f32, 0f32, 0f32, 0f32);
2131                                for j in 0..cols {
2132                                    let w = row[j];
2133                                    a0 += w * x0[j];
2134                                    a1 += w * x1[j];
2135                                    a2 += w * x2[j];
2136                                    a3 += w * x3[j];
2137                                }
2138                                for (k, acc) in [a0, a1, a2, a3].into_iter().enumerate() {
2139                                    unsafe { *out_addr.at((bi + k) * rows + o) = acc };
2140                                }
2141                            } else {
2142                                for k in 0..n {
2143                                    let x = &xs_all[(bi + k) * cols..(bi + k + 1) * cols];
2144                                    let mut acc = 0f32;
2145                                    for j in 0..cols {
2146                                        acc += row[j] * x[j];
2147                                    }
2148                                    unsafe { *out_addr.at((bi + k) * rows + o) = acc };
2149                                }
2150                            }
2151                        }
2152                        bi += n;
2153                    }
2154                };
2155                match pool {
2156                    // A chunk shorter than 256 positions (every prefill
2157                    // chunk on a Mac) left the whole GEMM on one core: an
2158                    // f16 CherryPit-1.2-2B scored 512 tokens in minutes.
2159                    // Split the output rows instead; every element keeps
2160                    // its own chain, so the result is bit-identical.
2161                    Some(pool) if b < 256 && rows >= 256 => {
2162                        pool.run_rows(rows, &|r0, r1| kernel(0, b, r0, r1))
2163                    }
2164                    _ => dispatch_rows(pool, b, &|s, e| kernel(s, e, 0, rows)),
2165                }
2166            }
2167            Self::Mapped {
2168                model,
2169                idx,
2170                dtype,
2171                row_scale,
2172                col_field,
2173                vbit_offsets,
2174                ..
2175            } => {
2176                if crate::prism::is_forward_weight(model, &model.tensors[*idx].name) {
2177                    let mut transformed = Vec::with_capacity(xs_all.len());
2178                    for bi in 0..b {
2179                        transformed.extend_from_slice(&crate::prism::forward(
2180                            model,
2181                            &xs_all[bi * cols..(bi + 1) * cols],
2182                        ));
2183                    }
2184                    match dtype {
2185                        TensorDtype::Q4Block => {
2186                            q4matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
2187                        }
2188                        TensorDtype::Q4Tiled => {
2189                            q4t_matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
2190                        }
2191                        TensorDtype::Q4TiledP => {
2192                            q4tp_matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
2193                        }
2194                        TensorDtype::Q2TiledP => {
2195                            let affine =
2196                                crate::prism::is_affine_target(model, &model.tensors[*idx].name);
2197                            // Affine Prism Q2TP has a descriptor-aware GPU
2198                            // kernel for short/tail batches too.  Unlike the
2199                            // ordinary Q2TP path, don't force b<32 back to a
2200                            // scalar CPU matmat: prefill chunks and the final
2201                            // tail both need to stay on the tested GPU arm.
2202                            let gpu_batch_ok = if affine {
2203                                b >= 2
2204                            } else {
2205                                b >= 32 && b * rows * cols >= 128_000_000
2206                            };
2207                            if gpu_batch_ok
2208                                && cols % 32 == 0
2209                                && crate::gpu::enabled_here()
2210                                && crate::gpu::q2tp_gpu_opt_in()
2211                            {
2212                                let gpu_ok = if affine {
2213                                    crate::gpu::q2tp_affine_matmat(
2214                                        model,
2215                                        *idx,
2216                                        &transformed,
2217                                        b,
2218                                        rows,
2219                                        cols,
2220                                        out,
2221                                    )
2222                                } else {
2223                                    crate::gpu::q2tp_matmat(
2224                                        model,
2225                                        *idx,
2226                                        &transformed,
2227                                        b,
2228                                        rows,
2229                                        cols,
2230                                        out,
2231                                    )
2232                                };
2233                                if gpu_ok {
2234                                    return;
2235                                }
2236                            }
2237                            // A one-token Prism decode is the other short
2238                            // case.  Use the descriptor-aware matvec kernel
2239                            // before falling back to the exact CPU path.
2240                            if affine
2241                                && b == 1
2242                                && cols % 32 == 0
2243                                && crate::gpu::enabled_here()
2244                                && crate::gpu::q2tp_gpu_opt_in()
2245                                && crate::gpu::q2tp_affine_matvec(
2246                                    model,
2247                                    *idx,
2248                                    &transformed[..cols],
2249                                    rows,
2250                                    cols,
2251                                    &mut out[..rows],
2252                                )
2253                            {
2254                                return;
2255                            }
2256                            if affine {
2257                                q2tp_affine_matmat(
2258                                    self.quant_bytes(),
2259                                    &transformed,
2260                                    b,
2261                                    rows,
2262                                    cols,
2263                                    out,
2264                                    pool,
2265                                )
2266                            } else {
2267                                q2tp_matmat(
2268                                    self.quant_bytes(),
2269                                    &transformed,
2270                                    b,
2271                                    rows,
2272                                    cols,
2273                                    out,
2274                                    pool,
2275                                )
2276                            }
2277                        }
2278                        TensorDtype::Q1 | TensorDtype::Q1G => bin_matmat(
2279                            self.bin_w().expect("binary dtype"),
2280                            &transformed,
2281                            b,
2282                            rows,
2283                            cols,
2284                            out,
2285                            pool,
2286                        ),
2287                        TensorDtype::Q1T => {
2288                            q1t_matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
2289                        }
2290                        TensorDtype::Mxfp4 => {
2291                            mxfp4::matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
2292                        }
2293                        TensorDtype::Vbit | TensorDtype::VbitRo => vbitmatmat(
2294                            self.quant_bytes(),
2295                            vbit_offsets,
2296                            &transformed,
2297                            b,
2298                            rows,
2299                            cols,
2300                            out,
2301                            pool,
2302                        ),
2303                        TensorDtype::Q8Row | TensorDtype::Q8_2f => {
2304                            let pre: Vec<std::borrow::Cow<'_, [f32]>> = (0..b)
2305                                .map(|bi| {
2306                                    prescale(
2307                                        &transformed[bi * cols..(bi + 1) * cols],
2308                                        col_field,
2309                                        *dtype,
2310                                    )
2311                                })
2312                                .collect();
2313                            qmatmat(self.quant_bytes(), row_scale, &pre, rows, cols, out, pool)
2314                        }
2315                        _ => unreachable!("unsupported mapped Prism dtype {dtype:?}"),
2316                    }
2317                    return;
2318                }
2319                if *dtype == TensorDtype::Q4Block {
2320                    q4matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2321                    return;
2322                }
2323                if *dtype == TensorDtype::Mxfp4 {
2324                    self.mxfp4_matmat(xs_all, b, out, pool);
2325                    return;
2326                }
2327                if *dtype == TensorDtype::Q4TiledP {
2328                    // GPU batched q4tp GEMM (dequant + f32nt mul_mm on the
2329                    // device); the probe keeps whichever beats the CPU arm.
2330                    // Narrow (prompt-encode) and wide (DiT) batches probe
2331                    // as separate classes — the regimes have opposite
2332                    // winners and one shared verdict locked the wrong arm.
2333                    // Kill switch (gpu::mm_kill): one grossly slow GPU op
2334                    // (a fair-condition op is ≤~100 ms even at 1024px)
2335                    // means the device is contended by another process
2336                    // (e.g. a simulator) — verdicts are per-process, so
2337                    // without the bail the whole render crawls behind
2338                    // someone else's queue.
2339                    // Row-exact batches stay on the host: the device GEMM
2340                    // is an f32 dequant-sgemm, not the host matvec's sum.
2341                    //
2342                    // The native CUDA backend takes every batch of a large
2343                    // matrix, however narrow, and is not arbitrated against
2344                    // the host: its plane is built on the card from the
2345                    // codec's bytes in tens of milliseconds, where a wgpu
2346                    // weight has an upload to amortise and a narrow batch
2347                    // tiles badly.
2348                    if crate::gpu::cuda_active()
2349                        && b >= 2
2350                        && rows * cols >= (1 << 20)
2351                        && cols % 32 == 0
2352                        && !row_exact()
2353                        && crate::gpu::enabled_here()
2354                    {
2355                        if let Self::Mapped { model, idx, .. } = self {
2356                            if crate::gpu::q4tp_matmat(model, *idx, xs_all, b, rows, cols, out) {
2357                                return;
2358                            }
2359                        }
2360                    }
2361                    if b >= 32
2362                        && b * rows * cols >= 128_000_000
2363                        && cols % 32 == 0
2364                        && !row_exact()
2365                        && !crate::gpu::mm_killed()
2366                        && crate::gpu::enabled_here()
2367                    {
2368                        let class = if b >= 128 {
2369                            crate::gpu::OpClass::MatmatWide
2370                        } else {
2371                            crate::gpu::OpClass::Matmat
2372                        };
2373                        if let Self::Mapped { model, idx, .. } = self {
2374                            // In-process A/B (`CMF_MM_AB=1`). Three
2375                            // wall-clock A/Bs on a shared stand disagreed
2376                            // with each other by 25% on the same change,
2377                            // because the machine drifts between processes
2378                            // and interleaving whole renders does not fix
2379                            // that. Here both arms run back to back on the
2380                            // SAME data inside one call, so whatever the
2381                            // machine is doing, it does to both — and the
2382                            // disagreement between their outputs falls out
2383                            // for free. Doubles the work; a diagnostic,
2384                            // not a mode.
2385                            if crate::mm_ab::on() {
2386                                let mut g = vec![0f32; b * rows];
2387                                let t = std::time::Instant::now();
2388                                let took = crate::gpu::q4tp_matmat(
2389                                    model, *idx, xs_all, b, rows, cols, &mut g,
2390                                );
2391                                let dg = t.elapsed();
2392                                let t = std::time::Instant::now();
2393                                q4tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2394                                let dc = t.elapsed();
2395                                crate::mm_ab::record(b, rows, cols, took, dg, dc, &g, out);
2396                                return;
2397                            }
2398                            let t0 = std::time::Instant::now();
2399                            // A cold call takes the device arm: its sample
2400                            // is discarded either way, and the upload is
2401                            // what the next step needs.
2402                            let resident = crate::gpu::weight_is_resident(model, *idx);
2403                            match crate::gpu::probe_arm_cold_prefers_gpu(class, resident) {
2404                                crate::gpu::ProbeArm::Gpu => {
2405                                    if crate::gpu::q4tp_matmat(
2406                                        model, *idx, xs_all, b, rows, cols, out,
2407                                    ) {
2408                                        let el = t0.elapsed();
2409                                        // Work-proportional budget: ~8× the
2410                                        // fair-device estimate (+20 ms slack).
2411                                        // An absolute cap missed the worst
2412                                        // case — contended ops sit at
2413                                        // 100–240 ms each and still bury a
2414                                        // render whose fair op is 3–9 ms.
2415                                        // Cold ops (first PSO build, buffer
2416                                        // alloc) are exempt: a one-off
2417                                        // ~50 ms compile is not contention.
2418                                        let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
2419                                        let budget = std::time::Duration::from_secs_f64(
2420                                            flops / 1.5e12 * 8.0 + 0.020,
2421                                        );
2422                                        crate::gpu::mm_budget_check(
2423                                            "q4tp matmat",
2424                                            el,
2425                                            budget,
2426                                            crate::gpu::probe_was_cold() || !resident,
2427                                        );
2428                                        crate::gpu::probe_record(class, true, el);
2429                                        return;
2430                                    }
2431                                }
2432                                crate::gpu::ProbeArm::CpuTimed => {
2433                                    q4tp_matmat(
2434                                        self.quant_bytes(),
2435                                        xs_all,
2436                                        b,
2437                                        rows,
2438                                        cols,
2439                                        out,
2440                                        pool,
2441                                    );
2442                                    crate::gpu::probe_record(class, false, t0.elapsed());
2443                                    return;
2444                                }
2445                                crate::gpu::ProbeArm::Cpu => {}
2446                            }
2447                        }
2448                    }
2449                    q4tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2450                    return;
2451                }
2452                if *dtype == TensorDtype::Q2TiledP {
2453                    // Same device arm as q4tp, behind the same probe:
2454                    // the planes differ, the dispatch does not. Without
2455                    // this a q2tp file ran its widest projections on the
2456                    // host while the 4-bit one had the card, which is a
2457                    // codec paying for its size twice.
2458                    if b >= 32
2459                        && b * rows * cols >= 128_000_000
2460                        && cols % 32 == 0
2461                        && !crate::gpu::mm_killed()
2462                        && crate::gpu::enabled_here()
2463                    {
2464                        let class = if b >= 128 {
2465                            crate::gpu::OpClass::MatmatWide
2466                        } else {
2467                            crate::gpu::OpClass::Matmat
2468                        };
2469                        if let Self::Mapped { model, idx, .. } = self {
2470                            let t0 = std::time::Instant::now();
2471                            match crate::gpu::probe_arm(class) {
2472                                crate::gpu::ProbeArm::Gpu => {
2473                                    if crate::gpu::q2tp_matmat(
2474                                        model, *idx, xs_all, b, rows, cols, out,
2475                                    ) {
2476                                        crate::gpu::probe_record(class, true, t0.elapsed());
2477                                        return;
2478                                    }
2479                                }
2480                                crate::gpu::ProbeArm::CpuTimed => {
2481                                    q2tp_matmat(
2482                                        self.quant_bytes(),
2483                                        xs_all,
2484                                        b,
2485                                        rows,
2486                                        cols,
2487                                        out,
2488                                        pool,
2489                                    );
2490                                    crate::gpu::probe_record(class, false, t0.elapsed());
2491                                    return;
2492                                }
2493                                crate::gpu::ProbeArm::Cpu => {}
2494                            }
2495                        }
2496                    }
2497                    // Without a host arm a q2tp tensor falls through to
2498                    // the q8 fallback, which reads it at one BYTE per
2499                    // weight — a 2x overrun that killed pool workers
2500                    // mid-prefill while the dispatcher waited forever.
2501                    q2tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2502                    return;
2503                }
2504                if *dtype == TensorDtype::Q4Tiled {
2505                    // GPU batched q4t GEMM (dequant + f32nt mul_mm on the
2506                    // device); the probe keeps whichever beats the CPU arm.
2507                    // Narrow (prompt-encode) and wide (DiT) batches probe
2508                    // as separate classes — the regimes have opposite
2509                    // winners and one shared verdict locked the wrong arm.
2510                    // Kill switch (gpu::mm_kill): one grossly slow GPU op
2511                    // (a fair-condition op is ≤~100 ms even at 1024px)
2512                    // means the device is contended by another process
2513                    // (e.g. a simulator) — verdicts are per-process, so
2514                    // without the bail the whole render crawls behind
2515                    // someone else's queue.
2516                    if b >= 32
2517                        && b * rows * cols >= 128_000_000
2518                        && cols % 32 == 0
2519                        && !crate::gpu::mm_killed()
2520                        && crate::gpu::enabled_here()
2521                    {
2522                        let class = if b >= 128 {
2523                            crate::gpu::OpClass::MatmatWide
2524                        } else {
2525                            crate::gpu::OpClass::Matmat
2526                        };
2527                        if let Self::Mapped { model, idx, .. } = self {
2528                            let t0 = std::time::Instant::now();
2529                            match crate::gpu::probe_arm(class) {
2530                                crate::gpu::ProbeArm::Gpu => {
2531                                    if crate::gpu::q4t_matmat(
2532                                        model, *idx, xs_all, b, rows, cols, out,
2533                                    ) {
2534                                        let el = t0.elapsed();
2535                                        // Work-proportional budget: ~8× the
2536                                        // fair-device estimate (+20 ms slack).
2537                                        // An absolute cap missed the worst
2538                                        // case — contended ops sit at
2539                                        // 100–240 ms each and still bury a
2540                                        // render whose fair op is 3–9 ms.
2541                                        // Cold ops (first PSO build, buffer
2542                                        // alloc) are exempt: a one-off
2543                                        // ~50 ms compile is not contention.
2544                                        let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
2545                                        let budget = std::time::Duration::from_secs_f64(
2546                                            flops / 1.5e12 * 8.0 + 0.020,
2547                                        );
2548                                        crate::gpu::mm_budget_check(
2549                                            "q4t matmat",
2550                                            el,
2551                                            budget,
2552                                            crate::gpu::probe_was_cold(),
2553                                        );
2554                                        crate::gpu::probe_record(class, true, el);
2555                                        return;
2556                                    }
2557                                }
2558                                crate::gpu::ProbeArm::CpuTimed => {
2559                                    q4t_matmat(
2560                                        self.quant_bytes(),
2561                                        xs_all,
2562                                        b,
2563                                        rows,
2564                                        cols,
2565                                        out,
2566                                        pool,
2567                                    );
2568                                    crate::gpu::probe_record(class, false, t0.elapsed());
2569                                    return;
2570                                }
2571                                crate::gpu::ProbeArm::Cpu => {}
2572                            }
2573                        }
2574                    }
2575                    q4t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2576                    return;
2577                }
2578                if *dtype == TensorDtype::Q1 || *dtype == TensorDtype::Q1G {
2579                    // GPU batched q1 GEMM for wide prefill (q1_mul_mm on the
2580                    // device); the probe keeps whichever beats the CPU matmat.
2581                    // q1g has the twin `q1g_mul_mm` (cols % 128 by format).
2582                    let w = self.bin_w().expect("binary dtype");
2583                    if b >= 32
2584                        && b * rows * cols >= 128_000_000
2585                        && cols % 64 == 0
2586                        && crate::gpu::enabled_here()
2587                    {
2588                        if let Self::Mapped { model, idx, .. } = self {
2589                            let t0 = std::time::Instant::now();
2590                            match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
2591                                crate::gpu::ProbeArm::Gpu => {
2592                                    let ok = if *dtype == TensorDtype::Q1G {
2593                                        crate::gpu::q1g_matmat(
2594                                            model, *idx, xs_all, b, rows, cols, out,
2595                                        )
2596                                    } else {
2597                                        crate::gpu::q1_matmat(
2598                                            model, *idx, xs_all, b, rows, cols, out,
2599                                        )
2600                                    };
2601                                    if ok {
2602                                        crate::gpu::probe_record(
2603                                            crate::gpu::OpClass::Matmat,
2604                                            true,
2605                                            t0.elapsed(),
2606                                        );
2607                                        return;
2608                                    }
2609                                }
2610                                crate::gpu::ProbeArm::CpuTimed => {
2611                                    bin_matmat(w, xs_all, b, rows, cols, out, pool);
2612                                    crate::gpu::probe_record(
2613                                        crate::gpu::OpClass::Matmat,
2614                                        false,
2615                                        t0.elapsed(),
2616                                    );
2617                                    return;
2618                                }
2619                                crate::gpu::ProbeArm::Cpu => {}
2620                            }
2621                        }
2622                    }
2623                    bin_matmat(w, xs_all, b, rows, cols, out, pool);
2624                    return;
2625                }
2626                if *dtype == TensorDtype::Q1T {
2627                    // GPU batched GEMM for wide prefill (base + overlay on the
2628                    // device); probe keeps the winner vs the CPU matmat.
2629                    if b >= 32 && b * rows * cols >= 128_000_000 && crate::gpu::enabled_here() {
2630                        if let Self::Mapped { model, idx, .. } = self {
2631                            let t0 = std::time::Instant::now();
2632                            match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
2633                                crate::gpu::ProbeArm::Gpu => {
2634                                    if crate::gpu::q1t_matmat(
2635                                        model, *idx, xs_all, b, rows, cols, out,
2636                                    ) {
2637                                        crate::gpu::probe_record(
2638                                            crate::gpu::OpClass::Matmat,
2639                                            true,
2640                                            t0.elapsed(),
2641                                        );
2642                                        return;
2643                                    }
2644                                }
2645                                crate::gpu::ProbeArm::CpuTimed => {
2646                                    q1t_matmat(
2647                                        self.quant_bytes(),
2648                                        xs_all,
2649                                        b,
2650                                        rows,
2651                                        cols,
2652                                        out,
2653                                        pool,
2654                                    );
2655                                    crate::gpu::probe_record(
2656                                        crate::gpu::OpClass::Matmat,
2657                                        false,
2658                                        t0.elapsed(),
2659                                    );
2660                                    return;
2661                                }
2662                                crate::gpu::ProbeArm::Cpu => {}
2663                            }
2664                        }
2665                    }
2666                    q1t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2667                    return;
2668                }
2669                if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
2670                    vbitmatmat(
2671                        self.quant_bytes(),
2672                        vbit_offsets,
2673                        xs_all,
2674                        b,
2675                        rows,
2676                        cols,
2677                        out,
2678                        pool,
2679                    );
2680                    return;
2681                }
2682                let _pmp = crate::cpuprof::spans::span("mm.prescale");
2683                let pre: Vec<std::borrow::Cow<'_, [f32]>> = (0..b)
2684                    .map(|bi| prescale(&xs_all[bi * cols..(bi + 1) * cols], col_field, *dtype))
2685                    .collect();
2686                drop(_pmp);
2687                // MiMo verification is a 2–4 row decode panel, not a wide
2688                // prompt GEMM. Keep q8 projections on the same device as
2689                // decode; the generic b>=8 gate otherwise silently moves
2690                // every projection back to CPU. The short wgpu matmat uses
2691                // the same 64-lane reduction as its single-token matvec.
2692                if row_exact()
2693                    && (1..=4).contains(&b)
2694                    && matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
2695                    && crate::gpu::enabled_here()
2696                    && crate::gpu::wgpu_active()
2697                {
2698                    let flat: Vec<f32> = pre.iter().flat_map(|v| v.iter().copied()).collect();
2699                    if crate::gpu::q8_matmat(model, *idx, row_scale, &flat, b, rows, cols, out) {
2700                        return;
2701                    }
2702                }
2703                // D5: large prefill-batch GEMMs — on the GPU (threshold by
2704                // work volume: submission carries b×rows×cols MACs).
2705                // Runtime probe: the naive GEMM shader + sync readback
2706                // lose to the CPU GEMM on slow driver stacks — alternate
2707                // both arms and keep the winner.
2708                // CMF_MATMAT_GPU=0: diagnostic off-switch for this D5
2709                // device arm only (default unchanged). The GPU int8 kernel
2710                // and the CPU int8 kernel are not bit-comparable and differ
2711                // ~1e-2/layer (measured on SAM3), which compounds through
2712                // deep stacks — validation harnesses that need a single
2713                // changed block per run turn the others off with this.
2714                if b >= 8
2715                    && b * rows * cols >= 128_000_000
2716                    && crate::gpu::enabled_here()
2717                    && std::env::var("CMF_MATMAT_GPU").as_deref() != Ok("0")
2718                {
2719                    if let Self::Mapped { model, idx, .. } = self {
2720                        MATMAT_GATE.fetch_add(1, Ordering::Relaxed);
2721                        let t0 = std::time::Instant::now();
2722                        match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
2723                            crate::gpu::ProbeArm::Gpu
2724                                if crate::gpu::probe_deciding(crate::gpu::OpClass::Matmat)
2725                                    && !crate::gpu::q8_resident_or_upload(model, *idx) =>
2726                            {
2727                                // Cold weights during probing: the upload
2728                                // has started, the count runs on the CPU —
2729                                // the GPU arm samples on the next touch.
2730                                let q = self.quant_bytes();
2731                                qmatmat(q, row_scale, &pre, rows, cols, out, pool);
2732                                return;
2733                            }
2734                            crate::gpu::ProbeArm::Gpu => {
2735                                // The two-field codec's own entry first, as
2736                                // `device_matmat` takes it: the column field
2737                                // folds into the weight plane and the GEMM
2738                                // runs on the matrix units. The plain int8
2739                                // entry below is the scalar f32 GEMM — on a
2740                                // Spark-X2.5 4B prefill it ran q|k|v|o at
2741                                // host speed.
2742                                if *dtype == TensorDtype::Q8_2f
2743                                    && std::env::var("CMF_Q8_2F_DEV").as_deref() != Ok("0")
2744                                    && crate::gpu::q8_matmat_2f(
2745                                        model, *idx, row_scale, col_field, xs_all, b, rows,
2746                                        cols, out,
2747                                    )
2748                                {
2749                                    MATMAT_GPU_HIT.fetch_add(1, Ordering::Relaxed);
2750                                    crate::gpu::probe_record(
2751                                        crate::gpu::OpClass::Matmat,
2752                                        true,
2753                                        t0.elapsed(),
2754                                    );
2755                                    return;
2756                                }
2757                                let flat: Vec<f32> =
2758                                    pre.iter().flat_map(|v| v.iter().copied()).collect();
2759                                if crate::gpu::q8_matmat(
2760                                    model, *idx, row_scale, &flat, b, rows, cols, out,
2761                                ) {
2762                                    crate::gpu::probe_record(
2763                                        crate::gpu::OpClass::Matmat,
2764                                        true,
2765                                        t0.elapsed(),
2766                                    );
2767                                    return;
2768                                }
2769                            }
2770                            crate::gpu::ProbeArm::CpuTimed => {
2771                                let q = self.quant_bytes();
2772                                qmatmat(q, row_scale, &pre, rows, cols, out, pool);
2773                                crate::gpu::probe_record(
2774                                    crate::gpu::OpClass::Matmat,
2775                                    false,
2776                                    t0.elapsed(),
2777                                );
2778                                return;
2779                            }
2780                            crate::gpu::ProbeArm::Cpu => {}
2781                        }
2782                    }
2783                }
2784                let q = self.quant_bytes();
2785                qmatmat(q, row_scale, &pre, rows, cols, out, pool);
2786            }
2787        }
2788    }
2789}
2790
2791impl QTensor {
2792    /// `mxfp4` matvec: the wgpu kernel for lm_head/expert-class shapes when
2793    /// the probe prefers the device, else the CPU kernels of
2794    /// `qtensor::mxfp4` (always on native Metal, which has no mxfp4 kernel).
2795    fn mxfp4_matvec(&self, x: &[f32], out: &mut [f32], pool: Option<&Pool>) {
2796        let (rows, cols) = (self.rows(), self.cols());
2797        if let Self::Mapped { model, idx, .. } = self {
2798            if rows * cols >= 8_388_608 && crate::gpu::enabled_here() {
2799                let t0 = std::time::Instant::now();
2800                let cls = crate::gpu::matvec_class(rows, cols);
2801                match crate::gpu::probe_arm(cls) {
2802                    crate::gpu::ProbeArm::Gpu => {
2803                        if crate::gpu::mxfp4_matvec(model, *idx, x, rows, cols, out) {
2804                            crate::gpu::probe_record(cls, true, t0.elapsed());
2805                            return;
2806                        }
2807                    }
2808                    crate::gpu::ProbeArm::CpuTimed => {
2809                        mxfp4::matvec(self.quant_bytes(), x, rows, cols, out, pool);
2810                        crate::gpu::probe_record(cls, false, t0.elapsed());
2811                        return;
2812                    }
2813                    crate::gpu::ProbeArm::Cpu => {}
2814                }
2815            }
2816        }
2817        mxfp4::matvec(self.quant_bytes(), x, rows, cols, out, pool);
2818    }
2819
2820    /// `mxfp4` batched matvec (`xs` `[b, cols]` → `out` `[b, rows]`): the
2821    /// device's batched kernel behind the same probe as the other codecs'
2822    /// GEMMs (never for row-exact batches — the device sums in another
2823    /// order than the host matvec), else the CPU kernels.
2824    fn mxfp4_matmat(&self, xs: &[f32], b: usize, out: &mut [f32], pool: Option<&Pool>) {
2825        let (rows, cols) = (self.rows(), self.cols());
2826        if let Self::Mapped { model, idx, .. } = self {
2827            if b >= 2
2828                && b * rows * cols >= 33_554_432
2829                && !row_exact()
2830                && !crate::gpu::mm_killed()
2831                && crate::gpu::enabled_here()
2832            {
2833                let class = if b >= 128 {
2834                    crate::gpu::OpClass::MatmatWide
2835                } else {
2836                    crate::gpu::OpClass::Matmat
2837                };
2838                let t0 = std::time::Instant::now();
2839                match crate::gpu::probe_arm(class) {
2840                    crate::gpu::ProbeArm::Gpu => {
2841                        if crate::gpu::mxfp4_matmat(model, *idx, xs, b, rows, cols, out) {
2842                            crate::gpu::probe_record(class, true, t0.elapsed());
2843                            return;
2844                        }
2845                    }
2846                    crate::gpu::ProbeArm::CpuTimed => {
2847                        mxfp4::matmat(self.quant_bytes(), xs, b, rows, cols, out, pool);
2848                        crate::gpu::probe_record(class, false, t0.elapsed());
2849                        return;
2850                    }
2851                    crate::gpu::ProbeArm::Cpu => {}
2852                }
2853            }
2854        }
2855        mxfp4::matmat(self.quant_bytes(), xs, b, rows, cols, out, pool);
2856    }
2857}
2858
2859impl QTensor {
2860    /// The device GEMM this tensor would take, run once on the caller's
2861    /// data — the startup parity probe's arm, and the one place that knows
2862    /// which entry point each codec has.
2863    ///
2864    /// It exists because the probe used to look for a `q4tp` weight by
2865    /// name AND dtype, and a container packed any other way was declared
2866    /// "host path" for the whole render even though its codec had a device
2867    /// GEMM of its own. A gate that only recognizes one codec is a gate
2868    /// that silently downgrades every other one.
2869    pub fn device_matmat(&self, xs: &[f32], b: usize, out: &mut [f32]) -> bool {
2870        let (rows, cols) = (self.rows(), self.cols());
2871        let Self::Mapped {
2872            model,
2873            idx,
2874            dtype,
2875            row_scale,
2876            col_field,
2877            ..
2878        } = self
2879        else {
2880            return false;
2881        };
2882        if crate::prism::has_contract(model) {
2883            return false;
2884        }
2885        match *dtype {
2886            TensorDtype::Q4TiledP => crate::gpu::q4tp_matmat(model, *idx, xs, b, rows, cols, out),
2887            // The two-field codec folds its column field into the
2888            // activation, which leaves a plain per-row int8 GEMM — the
2889            // same kernel `q8_row` uses, on both backends.
2890            TensorDtype::Q8Row | TensorDtype::Q8_2f => {
2891                // The field belongs to the weight; only a backend that cannot
2892                // apply it there makes a scaled copy of the activation.
2893                if *dtype == TensorDtype::Q8_2f
2894                    && std::env::var("CMF_Q8_2F_DEV").as_deref() != Ok("0")
2895                    && crate::gpu::q8_matmat_2f(
2896                        model, *idx, row_scale, col_field, xs, b, rows, cols, out,
2897                    )
2898                {
2899                    return true;
2900                }
2901                let flat: Vec<f32> = (0..b)
2902                    .flat_map(|bi| {
2903                        prescale(&xs[bi * cols..(bi + 1) * cols], col_field, *dtype).into_owned()
2904                    })
2905                    .collect();
2906                crate::gpu::q8_matmat(model, *idx, row_scale, &flat, b, rows, cols, out)
2907            }
2908            _ => false,
2909        }
2910    }
2911
2912    /// Multi-matrix job (roadmap §3 P0): N tensors sharing one input
2913    /// run under a SINGLE pool dispatch — QKV or gate+up cost one
2914    /// barrier instead of N. Per-row math is the exact same kernel as
2915    /// `matvec` (bit-identical outputs); only the dispatch is fused.
2916    /// Falls back to N sequential matvecs when the set is not a uniform
2917    /// q8-family/F32 group or there is no pool.
2918    pub fn matvec_many<const N: usize>(
2919        ts: [&QTensor; N],
2920        x: &[f32],
2921        mut outs: [&mut [f32]; N],
2922        pool: Option<&Pool>,
2923    ) {
2924        if Self::any_exact_cpu_default(&ts) {
2925            return float_activations_scope(|| Self::matvec_many(ts, x, outs, pool));
2926        }
2927        let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
2928        if ts.iter().any(|t| t.has_prism_contract() || t.per_tensor_only()) {
2929            // The fused range kernels have no transform descriptor.  Let
2930            // each tensor's ordinary matvec dispatch perform the explicit
2931            // signed FWHT (and retain CPU fallback for mixed q2tp/q4tp).
2932            // blr has no row-range kernel either: per-tensor calls.
2933            for (t, o) in ts.iter().zip(outs.iter_mut()) {
2934                t.matvec(x, o, pool);
2935            }
2936            return;
2937        }
2938        let uniform_q8 = ts.iter().all(|t| {
2939            matches!(
2940                t,
2941                Self::Mapped {
2942                    dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f,
2943                    ..
2944                }
2945            )
2946        });
2947        let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
2948        if uniform_f32 && crate::f32_backend::active() {
2949            for (t, o) in ts.iter().zip(outs.iter_mut()) {
2950                t.matvec(x, o, pool);
2951            }
2952            return;
2953        }
2954        let uniform_q4 = ts.iter().all(|t| {
2955            matches!(
2956                t,
2957                Self::Mapped {
2958                    dtype: TensorDtype::Q4Block,
2959                    ..
2960                }
2961            )
2962        });
2963        let uniform_vbit = ts.iter().all(|t| {
2964            matches!(
2965                t,
2966                Self::Mapped {
2967                    dtype: TensorDtype::Vbit | TensorDtype::VbitRo,
2968                    ..
2969                }
2970            )
2971        });
2972        // q1 and q1g may mix freely: the shared activation split and group
2973        // sums are layout-independent, only the row kernel differs.
2974        let uniform_q1 =
2975            ts.iter().all(|t| t.bin_w().is_some()) && ts.iter().all(|t| t.cols() == ts[0].cols());
2976        let uniform_q1t = ts.iter().all(|t| {
2977            matches!(
2978                t,
2979                Self::Mapped {
2980                    dtype: TensorDtype::Q1T,
2981                    ..
2982                }
2983            )
2984        });
2985        // q4tp is the skeleton dtype of the big MoE files, and without an arm
2986        // here every projection that shares an input paid its own pool
2987        // barrier: DeepSeek-V4's attention step alone hands this function
2988        // wq_a, wkv and both compressors' pairs off the same hidden state.
2989        let uniform_q4tp = ts.iter().all(|t| {
2990            matches!(
2991                t,
2992                Self::Mapped {
2993                    dtype: TensorDtype::Q4TiledP,
2994                    ..
2995                }
2996            )
2997        }) && ts
2998            .iter()
2999            .all(|t| t.cols() == ts[0].cols() && t.cols() % GROUP_SIZE == 0);
3000        let Some(pool) = pool else {
3001            for (t, o) in ts.iter().zip(outs.iter_mut()) {
3002                t.matvec(x, o, None);
3003            }
3004            return;
3005        };
3006        if total_rows < 256
3007            || !(uniform_q8
3008                || uniform_f32
3009                || uniform_q4
3010                || uniform_vbit
3011                || uniform_q1
3012                || uniform_q1t
3013                || uniform_q4tp)
3014        {
3015            for (t, o) in ts.iter().zip(outs.iter_mut()) {
3016                t.matvec(x, o, Some(pool));
3017            }
3018            return;
3019        }
3020
3021        if uniform_q4tp {
3022            // Every tensor's rows laid end to end in one virtual row space,
3023            // so the whole set is ONE dispatch. The per-row body is the
3024            // `q4tp_matvec` arm verbatim — same activation split, same
3025            // accumulation order — so the outputs are bit-identical to the
3026            // sequential calls this replaces.
3027            let cols = ts[0].cols();
3028            let gpr = cols / GROUP_SIZE;
3029            let views: [Q4tpView; N] =
3030                std::array::from_fn(|i| Q4tpView::new(ts[i].quant_bytes(), ts[i].rows(), cols));
3031            let rows_of: [usize; N] = std::array::from_fn(|i| ts[i].rows());
3032            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
3033            // flat index -> (which tensor, which of its rows)
3034            let locate = |flat: usize| -> (usize, usize) {
3035                let mut acc = 0;
3036                for (i, &r) in rows_of.iter().enumerate() {
3037                    if flat < acc + r {
3038                        return (i, flat - acc);
3039                    }
3040                    acc += r;
3041                }
3042                (rows_of.len() - 1, 0)
3043            };
3044            let (views, outs_addr) = (&views, &outs_addr);
3045            if a8w8_enabled() {
3046                let act = split_act(x);
3047                let act = &act;
3048                #[cfg(target_arch = "x86_64")]
3049                let lanes = q4tp_lanes();
3050                let run = |start: usize, end: usize| {
3051                    #[cfg(target_arch = "x86_64")]
3052                    if lanes > 1 && act.xg8.len() == gpr {
3053                        // L independent accumulator chains per worker —
3054                        // the plain body's math, lane-interleaved.
3055                        with_krow(gpr * lanes, |sc| {
3056                            // SAFETY: a8w8 implies AVX2; each (tensor, row)
3057                            // pair still belongs to exactly one worker.
3058                            unsafe {
3059                                match lanes {
3060                                    2 => q4tp_many_lanes::<2>(
3061                                        &views[..],
3062                                        &locate,
3063                                        start,
3064                                        end,
3065                                        gpr,
3066                                        &act.xq,
3067                                        &act.xg8,
3068                                        act.sx,
3069                                        &act.outliers,
3070                                        sc,
3071                                        &outs_addr[..],
3072                                    ),
3073                                    8 => q4tp_many_lanes::<8>(
3074                                        &views[..],
3075                                        &locate,
3076                                        start,
3077                                        end,
3078                                        gpr,
3079                                        &act.xq,
3080                                        &act.xg8,
3081                                        act.sx,
3082                                        &act.outliers,
3083                                        sc,
3084                                        &outs_addr[..],
3085                                    ),
3086                                    _ => q4tp_many_lanes::<4>(
3087                                        &views[..],
3088                                        &locate,
3089                                        start,
3090                                        end,
3091                                        gpr,
3092                                        &act.xq,
3093                                        &act.xg8,
3094                                        act.sx,
3095                                        &act.outliers,
3096                                        sc,
3097                                        &outs_addr[..],
3098                                    ),
3099                                }
3100                            }
3101                        });
3102                        return;
3103                    }
3104                    crate::cpuprof::count_plain_rows((end - start) as u64);
3105                    with_krow(gpr, |sc| {
3106                        for flat in start..end {
3107                            let (t, r) = locate(flat);
3108                            let v = &views[t];
3109                            v.scales_into(r, gpr, sc);
3110                            let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, sc) * act.sx;
3111                            for &(j, xv) in &act.outliers {
3112                                let (w, s) = q4tp_outlier(v.nib, r, gpr, j, sc);
3113                                acc += w * s * xv;
3114                            }
3115                            // SAFETY: one worker owns each (tensor, row) pair.
3116                            unsafe { *outs_addr[t].at(r) = acc };
3117                        }
3118                    });
3119                };
3120                pool.run_rows(total_rows, &run);
3121            } else {
3122                let run = |start: usize, end: usize| {
3123                    with_krow(gpr, |sc| {
3124                        for flat in start..end {
3125                            let (t, r) = locate(flat);
3126                            let v = &views[t];
3127                            v.scales_into(r, gpr, sc);
3128                            // SAFETY: one worker owns each (tensor, row) pair.
3129                            unsafe { *outs_addr[t].at(r) = q4tp_row_exact(v.nib, r, gpr, x, sc) };
3130                        }
3131                    });
3132                };
3133                pool.run_rows(total_rows, &run);
3134            }
3135            return;
3136        }
3137
3138        if uniform_q1 {
3139            // One shared activation split + group sums (q1 has no col
3140            // field; the same input feeds every tensor).
3141            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
3142            if a8w8_enabled() {
3143                let act = split_act(x);
3144                let gsum = q1_group_sums(&act.xq, ts[0].cols() / GROUP_SIZE);
3145                let (act, gsum) = (&act, &gsum);
3146                let closures: [_; N] = std::array::from_fn(|i| {
3147                    let (w, gpr, out) = (
3148                        ts[i].bin_w().expect("uniform_q1 checked"),
3149                        ts[i].cols() / GROUP_SIZE,
3150                        outs_addr[i],
3151                    );
3152                    move |s: usize, e: usize| bin_range_a8w8(w, gpr, act, gsum, out, s, e)
3153                });
3154                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3155                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3156                pool.run_many(&parts);
3157            } else {
3158                let closures: [_; N] = std::array::from_fn(|i| {
3159                    let (w, gpr, out) = (
3160                        ts[i].bin_w().expect("uniform_q1 checked"),
3161                        ts[i].cols() / GROUP_SIZE,
3162                        outs_addr[i],
3163                    );
3164                    move |s: usize, e: usize| bin_range_f32(w, gpr, x, out, s, e)
3165                });
3166                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3167                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3168                pool.run_many(&parts);
3169            }
3170            return;
3171        }
3172
3173        if uniform_q1t {
3174            // Q1T batched: one shared activation split + overlay decode,
3175            // all tensors' rows in ONE pool dispatch (saves N−1 dispatches
3176            // and N−1 redundant split_act calls per layer).
3177            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
3178            const TILE: usize = cortiq_core::quant::Q1T_TILE;
3179            if a8w8_enabled() {
3180                let act = split_act(x);
3181                let act = &act;
3182                let x_ref = x;
3183                let closures: [_; N] = std::array::from_fn(|i| {
3184                    let bytes = ts[i].quant_bytes();
3185                    let (rows, cols) = (ts[i].rows(), ts[i].cols());
3186                    let gpr = cols / GROUP_SIZE;
3187                    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
3188                    let out = outs_addr[i];
3189                    move |s: usize, e: usize| {
3190                        q1t_range_a8w8(bytes, gpr, rp_off, ent_off, has_ov, act, x_ref, out, s, e)
3191                    }
3192                });
3193                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3194                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3195                pool.run_many(&parts);
3196            } else {
3197                let x_ref = x;
3198                let closures: [_; N] = std::array::from_fn(|i| {
3199                    let bytes = ts[i].quant_bytes();
3200                    let (rows, cols) = (ts[i].rows(), ts[i].cols());
3201                    let gpr = cols / GROUP_SIZE;
3202                    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
3203                    let out = outs_addr[i];
3204                    move |s: usize, e: usize| {
3205                        q1t_range_f32_batch(bytes, gpr, rp_off, ent_off, has_ov, x_ref, out, s, e)
3206                    }
3207                });
3208                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3209                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3210                pool.run_many(&parts);
3211            }
3212            return;
3213        }
3214
3215        if uniform_q4 || uniform_vbit {
3216            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
3217            // q4/vbit share one activation split — no per-tensor col field.
3218            if a8w8_enabled() {
3219                let act = split_act(x);
3220                let act = &act;
3221                if uniform_q4 {
3222                    let closures: [_; N] = std::array::from_fn(|i| {
3223                        let (packed, scales) =
3224                            q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
3225                        let (gpr, cols, out) =
3226                            (ts[i].cols() / GROUP_SIZE, ts[i].cols(), outs_addr[i]);
3227                        move |s: usize, e: usize| {
3228                            q4_range_a8w8(packed, scales, gpr, cols, act, out, s, e)
3229                        }
3230                    });
3231                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3232                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3233                    pool.run_many(&parts);
3234                } else {
3235                    let closures: [_; N] = std::array::from_fn(|i| {
3236                        let Self::Mapped { vbit_offsets, .. } = ts[i] else {
3237                            unreachable!()
3238                        };
3239                        let (bytes, rows, cols, out) = (
3240                            ts[i].quant_bytes(),
3241                            ts[i].rows(),
3242                            ts[i].cols(),
3243                            outs_addr[i],
3244                        );
3245                        move |s: usize, e: usize| {
3246                            vbit_range_a8w8(bytes, vbit_offsets, x, act, rows, cols, out, s, e)
3247                        }
3248                    });
3249                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3250                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3251                    pool.run_many(&parts);
3252                }
3253                return;
3254            }
3255            if uniform_q4 {
3256                let closures: [_; N] = std::array::from_fn(|i| {
3257                    let (packed, scales) =
3258                        q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
3259                    let (gpr, out) = (ts[i].cols() / GROUP_SIZE, outs_addr[i]);
3260                    move |s: usize, e: usize| q4_range_f32(packed, scales, gpr, x, out, s, e)
3261                });
3262                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3263                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3264                pool.run_many(&parts);
3265            } else {
3266                let closures: [_; N] = std::array::from_fn(|i| {
3267                    let Self::Mapped { vbit_offsets, .. } = ts[i] else {
3268                        unreachable!()
3269                    };
3270                    let (bytes, rows, cols, out) = (
3271                        ts[i].quant_bytes(),
3272                        ts[i].rows(),
3273                        ts[i].cols(),
3274                        outs_addr[i],
3275                    );
3276                    move |s: usize, e: usize| {
3277                        vbit_range_f32(bytes, vbit_offsets, x, rows, cols, out, s, e)
3278                    }
3279                });
3280                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3281                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3282                pool.run_many(&parts);
3283            }
3284            return;
3285        }
3286
3287        if uniform_f32 {
3288            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
3289            let closures: [_; N] = std::array::from_fn(|i| {
3290                let Self::F32 { data, cols, .. } = ts[i] else {
3291                    unreachable!()
3292                };
3293                let out = outs_addr[i];
3294                move |start: usize, end: usize| {
3295                    for o in start..end {
3296                        let row = &data[o * cols..(o + 1) * cols];
3297                        let mut sum = 0.0f32;
3298                        for j in 0..*cols {
3299                            sum += row[j] * x[j];
3300                        }
3301                        // SAFETY: disjoint (tensor, row) cells per worker.
3302                        unsafe { *out.at(o) = sum };
3303                    }
3304                }
3305            });
3306            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3307                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3308            pool.run_many(&parts);
3309            return;
3310        }
3311
3312        // Uniform q8-family: per-tensor prescale (q8_2f col fields
3313        // differ per tensor) + the shared range kernels.
3314        struct Ctx<'a> {
3315            bytes: &'a [u8],
3316            #[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
3317            rep: &'a [u8],
3318            row_scale: &'a [f32],
3319            cols: usize,
3320            xs: std::borrow::Cow<'a, [f32]>,
3321        }
3322        let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
3323            let Self::Mapped {
3324                dtype,
3325                cols,
3326                row_scale,
3327                col_field,
3328                repack,
3329                ..
3330            } = ts[i]
3331            else {
3332                unreachable!()
3333            };
3334            Ctx {
3335                bytes: ts[i].quant_bytes(),
3336                rep: repack,
3337                row_scale,
3338                cols: *cols,
3339                xs: prescale(x, col_field, *dtype),
3340            }
3341        });
3342        let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
3343        #[cfg(target_arch = "aarch64")]
3344        if sdot_enabled() {
3345            let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
3346            let closures: [_; N] = std::array::from_fn(|i| {
3347                let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
3348                move |start: usize, end: usize| {
3349                    q8_range_sdot(c.bytes, c.rep, c.row_scale, act, c.cols, out, start, end)
3350                }
3351            });
3352            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3353                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3354            pool.run_many(&parts);
3355            return;
3356        }
3357        #[cfg(target_arch = "x86_64")]
3358        if avx2_a8w8_enabled() {
3359            let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
3360            let closures: [_; N] = std::array::from_fn(|i| {
3361                let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
3362                move |start: usize, end: usize| {
3363                    q8_range_avx2(c.bytes, c.row_scale, act, c.cols, out, start, end)
3364                }
3365            });
3366            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3367                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3368            pool.run_many(&parts);
3369            return;
3370        }
3371        let closures: [_; N] = std::array::from_fn(|i| {
3372            let (c, out) = (&ctxs[i], outs_addr[i]);
3373            move |start: usize, end: usize| {
3374                q8_range_f32(c.bytes, c.row_scale, &c.xs, c.cols, out, start, end)
3375            }
3376        });
3377        let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3378            std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3379        pool.run_many(&parts);
3380    }
3381}
3382
3383impl QTensor {
3384    /// Pair-input multi-matrix job: N tensors × 2 shared inputs under a
3385    /// single pool dispatch — the MTP/pair decode path publishes one job
3386    /// for Q/K/V (and one for gate+up) instead of one per tensor.
3387    /// Per-row math is exactly `matvec2`'s kernels; bit-identical.
3388    #[allow(clippy::needless_range_loop)]
3389    pub fn matvec2_many<const N: usize>(
3390        ts: [&QTensor; N],
3391        x1: &[f32],
3392        x2: &[f32],
3393        mut o1s: [&mut [f32]; N],
3394        mut o2s: [&mut [f32]; N],
3395        pool: Option<&Pool>,
3396    ) {
3397        if Self::any_exact_cpu_default(&ts) {
3398            return float_activations_scope(|| {
3399                Self::matvec2_many(ts, x1, x2, o1s, o2s, pool)
3400            });
3401        }
3402        let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
3403        if ts.iter().any(|t| t.has_prism_contract() || t.per_tensor_only()) {
3404            for i in 0..N {
3405                ts[i].matvec2(x1, x2, o1s[i], o2s[i], pool);
3406            }
3407            return;
3408        }
3409        let uniform_q8 = ts.iter().all(|t| {
3410            matches!(
3411                t,
3412                Self::Mapped {
3413                    dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f,
3414                    ..
3415                }
3416            )
3417        });
3418        let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
3419        let uniform_q4 = ts.iter().all(|t| {
3420            matches!(
3421                t,
3422                Self::Mapped {
3423                    dtype: TensorDtype::Q4Block,
3424                    ..
3425                }
3426            )
3427        });
3428        let uniform_vbit = ts.iter().all(|t| {
3429            matches!(
3430                t,
3431                Self::Mapped {
3432                    dtype: TensorDtype::Vbit | TensorDtype::VbitRo,
3433                    ..
3434                }
3435            )
3436        });
3437        let fusable = pool.is_some()
3438            && total_rows >= 256
3439            && (uniform_q8 || uniform_f32 || uniform_q4 || uniform_vbit);
3440        if !fusable {
3441            for i in 0..N {
3442                ts[i].matvec2(x1, x2, o1s[i], o2s[i], pool);
3443            }
3444            return;
3445        }
3446        let pool = pool.unwrap();
3447
3448        if uniform_q4 || uniform_vbit {
3449            let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
3450            let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
3451            // q4/vbit share activation splits — no per-tensor col field.
3452            if a8w8_enabled() {
3453                let a1 = split_act(x1);
3454                let a2 = split_act(x2);
3455                let (a1, a2) = (&a1, &a2);
3456                if uniform_q4 {
3457                    let closures: [_; N] = std::array::from_fn(|i| {
3458                        let (packed, scales) =
3459                            q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
3460                        let (gpr, cols, o1, o2) =
3461                            (ts[i].cols() / GROUP_SIZE, ts[i].cols(), p1[i], p2[i]);
3462                        move |s: usize, e: usize| {
3463                            q4_range2_a8w8(packed, scales, gpr, cols, a1, a2, o1, o2, s, e)
3464                        }
3465                    });
3466                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3467                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3468                    pool.run_many(&parts);
3469                } else {
3470                    let closures: [_; N] = std::array::from_fn(|i| {
3471                        let Self::Mapped { vbit_offsets, .. } = ts[i] else {
3472                            unreachable!()
3473                        };
3474                        let (bytes, rows, cols, o1, o2) = (
3475                            ts[i].quant_bytes(),
3476                            ts[i].rows(),
3477                            ts[i].cols(),
3478                            p1[i],
3479                            p2[i],
3480                        );
3481                        move |s: usize, e: usize| {
3482                            vbit_range2_a8w8(
3483                                bytes,
3484                                vbit_offsets,
3485                                x1,
3486                                x2,
3487                                a1,
3488                                a2,
3489                                rows,
3490                                cols,
3491                                o1,
3492                                o2,
3493                                s,
3494                                e,
3495                            )
3496                        }
3497                    });
3498                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3499                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3500                    pool.run_many(&parts);
3501                }
3502                return;
3503            }
3504            if uniform_q4 {
3505                let closures: [_; N] = std::array::from_fn(|i| {
3506                    let (packed, scales) =
3507                        q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
3508                    let (gpr, o1, o2) = (ts[i].cols() / GROUP_SIZE, p1[i], p2[i]);
3509                    move |s: usize, e: usize| {
3510                        q4_range2_f32(packed, scales, gpr, x1, x2, o1, o2, s, e)
3511                    }
3512                });
3513                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3514                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3515                pool.run_many(&parts);
3516            } else {
3517                let closures: [_; N] = std::array::from_fn(|i| {
3518                    let Self::Mapped { vbit_offsets, .. } = ts[i] else {
3519                        unreachable!()
3520                    };
3521                    let (bytes, rows, cols, o1, o2) = (
3522                        ts[i].quant_bytes(),
3523                        ts[i].rows(),
3524                        ts[i].cols(),
3525                        p1[i],
3526                        p2[i],
3527                    );
3528                    move |s: usize, e: usize| {
3529                        vbit_range2_f32(bytes, vbit_offsets, x1, x2, rows, cols, o1, o2, s, e)
3530                    }
3531                });
3532                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3533                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3534                pool.run_many(&parts);
3535            }
3536            return;
3537        }
3538
3539        if uniform_f32 {
3540            let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
3541            let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
3542            let closures: [_; N] = std::array::from_fn(|i| {
3543                let Self::F32 { data, cols, .. } = ts[i] else {
3544                    unreachable!()
3545                };
3546                let (o1, o2) = (p1[i], p2[i]);
3547                move |start: usize, end: usize| {
3548                    for o in start..end {
3549                        let row = &data[o * cols..(o + 1) * cols];
3550                        let (mut s1, mut s2) = (0.0f32, 0.0f32);
3551                        for j in 0..*cols {
3552                            s1 += row[j] * x1[j];
3553                            s2 += row[j] * x2[j];
3554                        }
3555                        // SAFETY: disjoint (tensor, row) cells per worker.
3556                        unsafe {
3557                            *o1.at(o) = s1;
3558                            *o2.at(o) = s2;
3559                        }
3560                    }
3561                }
3562            });
3563            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3564                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3565            pool.run_many(&parts);
3566            return;
3567        }
3568
3569        struct Ctx<'a> {
3570            bytes: &'a [u8],
3571            row_scale: &'a [f32],
3572            cols: usize,
3573            xs1: std::borrow::Cow<'a, [f32]>,
3574            xs2: std::borrow::Cow<'a, [f32]>,
3575        }
3576        let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
3577            let Self::Mapped {
3578                dtype,
3579                cols,
3580                row_scale,
3581                col_field,
3582                ..
3583            } = ts[i]
3584            else {
3585                unreachable!()
3586            };
3587            Ctx {
3588                bytes: ts[i].quant_bytes(),
3589                row_scale,
3590                cols: *cols,
3591                xs1: prescale(x1, col_field, *dtype),
3592                xs2: prescale(x2, col_field, *dtype),
3593            }
3594        });
3595        let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
3596        let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
3597        #[cfg(target_arch = "aarch64")]
3598        if sdot_enabled() {
3599            let acts: [(SplitAct, SplitAct); N] =
3600                std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
3601            let closures: [_; N] = std::array::from_fn(|i| {
3602                let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
3603                move |start: usize, end: usize| {
3604                    q8_range2_sdot(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
3605                }
3606            });
3607            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3608                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3609            pool.run_many(&parts);
3610            return;
3611        }
3612        #[cfg(target_arch = "x86_64")]
3613        if avx2_a8w8_enabled() {
3614            let acts: [(SplitAct, SplitAct); N] =
3615                std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
3616            let closures: [_; N] = std::array::from_fn(|i| {
3617                let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
3618                move |start: usize, end: usize| {
3619                    q8_range2_avx2(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
3620                }
3621            });
3622            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3623                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3624            pool.run_many(&parts);
3625            return;
3626        }
3627        let closures: [_; N] = std::array::from_fn(|i| {
3628            let (c, o1, o2) = (&ctxs[i], p1[i], p2[i]);
3629            move |start: usize, end: usize| {
3630                q8_range2_f32(
3631                    c.bytes,
3632                    c.row_scale,
3633                    &c.xs1,
3634                    &c.xs2,
3635                    c.cols,
3636                    o1,
3637                    o2,
3638                    start,
3639                    end,
3640                )
3641            }
3642        });
3643        let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3644            std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3645        pool.run_many(&parts);
3646    }
3647
3648    /// Fused gate+up matvec with SiLU·mul: for each row r, computes
3649    /// `silu(gate·x) * (up·x)` and writes to `out[r]`. ONE pool dispatch,
3650    /// no intermediate g/u buffers, no separate silu pass. Falls back
3651    /// (returns false) for unsupported dtype combos.
3652    pub fn matvec_silu_mul(
3653        gate: &QTensor,
3654        up: &QTensor,
3655        x: &[f32],
3656        out: &mut [f32],
3657        pool: Option<&Pool>,
3658    ) -> bool {
3659        Self::matvec_silu_mul_limited(gate, up, x, out, 0.0, pool)
3660    }
3661
3662    /// Fused gate+up+SiLU with the GLM asymmetrical clamp.  `limit == 0`
3663    /// preserves the historical unclamped helper; a positive limit clamps
3664    /// `up` to both sides and `gate` only from above, matching the GLM
3665    /// SwiGLU reference.  Keeping the limit in the row kernel avoids the two
3666    /// intermediate vectors and the extra combine pass on the Q2TP experts.
3667    pub fn matvec_silu_mul_limited(
3668        gate: &QTensor,
3669        up: &QTensor,
3670        x: &[f32],
3671        out: &mut [f32],
3672        limit: f32,
3673        pool: Option<&Pool>,
3674    ) -> bool {
3675        if gate.exact_cpu_default() || up.exact_cpu_default() {
3676            return float_activations_scope(|| {
3677                Self::matvec_silu_mul_limited(gate, up, x, out, limit, pool)
3678            });
3679        }
3680        if gate.has_prism_contract()
3681            || up.has_prism_contract()
3682            || gate.per_tensor_only()
3683            || up.per_tensor_only()
3684        {
3685            // The fused gate/up kernels consume x directly.  Prism requires
3686            // a per-matrix signed FWHT, so the caller must use two ordinary
3687            // descriptor-aware matvecs instead of an unrotated fast path.
3688            return false;
3689        }
3690        let inter = gate.rows();
3691        debug_assert_eq!(up.rows(), inter);
3692        debug_assert_eq!(out.len(), inter);
3693        debug_assert_eq!(gate.cols(), up.cols());
3694        if !a8w8_enabled() {
3695            return false;
3696        }
3697        let act = split_act(x);
3698        let act = &act;
3699        let x_ref = x;
3700        let out_addr = SendMut(out.as_mut_ptr());
3701
3702        match (gate, up) {
3703            // Q4Block gate + Q4Block up (most common mobile q4 models)
3704            (
3705                Self::Mapped {
3706                    dtype: TensorDtype::Q4Block,
3707                    ..
3708                },
3709                Self::Mapped {
3710                    dtype: TensorDtype::Q4Block,
3711                    ..
3712                },
3713            ) => {
3714                let (gp, gs) = q4_split(gate.quant_bytes(), gate.rows(), gate.cols());
3715                let (up_p, up_s) = q4_split(up.quant_bytes(), up.rows(), up.cols());
3716                let gpr = gate.cols() / GROUP_SIZE;
3717                let cols = gate.cols();
3718                let run = move |start: usize, end: usize| {
3719                    for r in start..end {
3720                        let mut gv = dot_q4_row_i8(gp, gs, r * gpr, gpr, &act.xq) * act.sx;
3721                        let mut uv = dot_q4_row_i8(up_p, up_s, r * gpr, gpr, &act.xq) * act.sx;
3722                        for &(j, xv) in &act.outliers {
3723                            let flat = r * cols + j;
3724                            let gb = gp[flat / 2];
3725                            let gn = if flat & 1 == 0 { gb & 0x0F } else { gb >> 4 };
3726                            let gsc = f16_to_f32(u16::from_le_bytes([
3727                                gs[(flat / GROUP_SIZE) * 2],
3728                                gs[(flat / GROUP_SIZE) * 2 + 1],
3729                            ]));
3730                            gv += ((gn as i32 - 8) as f32) * gsc * xv;
3731                            let ub = up_p[flat / 2];
3732                            let un = if flat & 1 == 0 { ub & 0x0F } else { ub >> 4 };
3733                            let usc = f16_to_f32(u16::from_le_bytes([
3734                                up_s[(flat / GROUP_SIZE) * 2],
3735                                up_s[(flat / GROUP_SIZE) * 2 + 1],
3736                            ]));
3737                            uv += ((un as i32 - 8) as f32) * usc * xv;
3738                        }
3739                        // SAFETY: disjoint row ranges per worker.
3740                        unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3741                    }
3742                };
3743                dispatch_rows(pool, inter, &run);
3744                true
3745            }
3746            // Q4Tiled gate + Q4Tiled up — one row pass, both tile
3747            // streams sequential, silu·mul fused (same per-row math as
3748            // `q4t_matvec`).
3749            (
3750                Self::Mapped {
3751                    dtype: TensorDtype::Q4Tiled,
3752                    ..
3753                },
3754                Self::Mapped {
3755                    dtype: TensorDtype::Q4Tiled,
3756                    ..
3757                },
3758            ) => {
3759                let g_bytes = gate.quant_bytes();
3760                let u_bytes = up.quant_bytes();
3761                let gpr = gate.cols() / GROUP_SIZE;
3762                let run = move |start: usize, end: usize| {
3763                    for r in start..end {
3764                        let mut gv = dot_q4t_row_i8(g_bytes, r, gpr, &act.xq) * act.sx;
3765                        let mut uv = dot_q4t_row_i8(u_bytes, r, gpr, &act.xq) * act.sx;
3766                        for &(j, xv) in &act.outliers {
3767                            let (w, s) = q4t_outlier(g_bytes, r, gpr, j);
3768                            gv += w * s * xv;
3769                            let (w, s) = q4t_outlier(u_bytes, r, gpr, j);
3770                            uv += w * s * xv;
3771                        }
3772                        // SAFETY: disjoint row ranges per worker.
3773                        unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3774                    }
3775                };
3776                dispatch_rows(pool, inter, &run);
3777                true
3778            }
3779            // Q4TiledP gate + Q4TiledP up — the same fused row pass, with
3780            // each row's two ladders built once and spent on both streams.
3781            (
3782                Self::Mapped {
3783                    dtype: TensorDtype::Q4TiledP,
3784                    ..
3785                },
3786                Self::Mapped {
3787                    dtype: TensorDtype::Q4TiledP,
3788                    ..
3789                },
3790            ) => {
3791                let cols = gate.cols();
3792                let gpr = cols / GROUP_SIZE;
3793                let gv_view = Q4tpView::new(gate.quant_bytes(), inter, cols);
3794                let uv_view = Q4tpView::new(up.quant_bytes(), inter, cols);
3795                #[cfg(target_arch = "x86_64")]
3796                let lanes = q4tp_lanes();
3797                let run = |start: usize, end: usize| {
3798                    #[cfg(target_arch = "x86_64")]
3799                    if lanes > 1 && act.xg8.len() == gpr {
3800                        // L independent accumulator chains per worker —
3801                        // the plain arm's math, lane-interleaved.
3802                        with_krows(gpr * lanes, |gsc, usc| {
3803                            // SAFETY: a8w8 implies AVX2; disjoint row ranges.
3804                            unsafe {
3805                                match lanes {
3806                                    2 => q4tp_gateup_lanes::<2>(
3807                                        &gv_view, &uv_view, start, end, gpr, &act.xq,
3808                                        &act.xg8, act.sx, limit, &act.outliers, gsc, usc,
3809                                        out_addr,
3810                                    ),
3811                                    8 => q4tp_gateup_lanes::<8>(
3812                                        &gv_view, &uv_view, start, end, gpr, &act.xq,
3813                                        &act.xg8, act.sx, limit, &act.outliers, gsc, usc,
3814                                        out_addr,
3815                                    ),
3816                                    _ => q4tp_gateup_lanes::<4>(
3817                                        &gv_view, &uv_view, start, end, gpr, &act.xq,
3818                                        &act.xg8, act.sx, limit, &act.outliers, gsc, usc,
3819                                        out_addr,
3820                                    ),
3821                                }
3822                            }
3823                        });
3824                        return;
3825                    }
3826                    crate::cpuprof::count_plain_rows((end - start) as u64);
3827                    with_krows(gpr, |gsc, usc| {
3828                        for r in start..end {
3829                            gv_view.scales_into(r, gpr, gsc);
3830                            uv_view.scales_into(r, gpr, usc);
3831                            let mut gv =
3832                                dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsc) * act.sx;
3833                            let mut uv =
3834                                dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, usc) * act.sx;
3835                            for &(j, xv) in &act.outliers {
3836                                let (w, s) = q4tp_outlier(gv_view.nib, r, gpr, j, gsc);
3837                                gv += w * s * xv;
3838                                let (w, s) = q4tp_outlier(uv_view.nib, r, gpr, j, usc);
3839                                uv += w * s * xv;
3840                            }
3841                            // SAFETY: disjoint row ranges per worker.
3842                            unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3843                        }
3844                    });
3845                };
3846                dispatch_rows(pool, inter, &run);
3847                true
3848            }
3849            // Q1 gate + Q1 up — one row pass over both sign streams,
3850            // silu·mul fused (the per-row math of `q1_range_a8w8`); the
3851            // activation group sums are shared by both streams. Without
3852            // this arm a q1 dense FFN paid two dispatches + a combine
3853            // loop — the exact barrier this function exists to remove.
3854            // q1g (and a q1/q1g mix) rides the same arm through `BinW`.
3855            (
3856                Self::Mapped {
3857                    dtype: TensorDtype::Q1 | TensorDtype::Q1G,
3858                    ..
3859                },
3860                Self::Mapped {
3861                    dtype: TensorDtype::Q1 | TensorDtype::Q1G,
3862                    ..
3863                },
3864            ) => {
3865                let gw = gate.bin_w().expect("binary dtype matched");
3866                let uw = up.bin_w().expect("binary dtype matched");
3867                let gpr = gate.cols() / GROUP_SIZE;
3868                let gsum = q1_group_sums(&act.xq, gpr);
3869                let gsum = &gsum;
3870                let run = move |start: usize, end: usize| {
3871                    for r in start..end {
3872                        let mut gv = gw.dot_i8(r, gpr, &act.xq, gsum) * act.sx;
3873                        let mut uv = uw.dot_i8(r, gpr, &act.xq, gsum) * act.sx;
3874                        for &(j, xv) in &act.outliers {
3875                            let (w, s) = gw.outlier(r, gpr, j);
3876                            gv += w * s * xv;
3877                            let (w, s) = uw.outlier(r, gpr, j);
3878                            uv += w * s * xv;
3879                        }
3880                        // SAFETY: disjoint row ranges per worker.
3881                        unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3882                    }
3883                };
3884                dispatch_rows(pool, inter, &run);
3885                true
3886            }
3887            // Q2TiledP gate + Q2TiledP up — the 2-bit expert pair (MoE
3888            // FFNs of the W2 class): one row pass, both ladders built
3889            // once, integer code dots with shared group sums.
3890            (
3891                Self::Mapped {
3892                    dtype: TensorDtype::Q2TiledP,
3893                    ..
3894                },
3895                Self::Mapped {
3896                    dtype: TensorDtype::Q2TiledP,
3897                    ..
3898                },
3899            ) => {
3900                let cols = gate.cols();
3901                let gpr = cols / GROUP_SIZE;
3902                let gv_view = Q4tpView::new_q2(gate.quant_bytes(), inter, cols);
3903                let uv_view = Q4tpView::new_q2(up.quant_bytes(), inter, cols);
3904                let gsum = q1_group_sums(&act.xq, gpr);
3905                let gsum = &gsum;
3906                let run = move |start: usize, end: usize| {
3907                    with_krows(gpr, |gsc, usc| {
3908                        for r in start..end {
3909                            gv_view.scales_into(r, gpr, gsc);
3910                            uv_view.scales_into(r, gpr, usc);
3911                            let mut gv =
3912                                dot_q2tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsum, gsc) * act.sx;
3913                            let mut uv =
3914                                dot_q2tp_row_i8(uv_view.nib, r, gpr, &act.xq, gsum, usc) * act.sx;
3915                            for &(j, xv) in &act.outliers {
3916                                let (w, s) = q2tp_outlier(gv_view.nib, r, gpr, j, gsc);
3917                                gv += w * s * xv;
3918                                let (w, s) = q2tp_outlier(uv_view.nib, r, gpr, j, usc);
3919                                uv += w * s * xv;
3920                            }
3921                            // SAFETY: disjoint row ranges per worker.
3922                            unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3923                        }
3924                    });
3925                };
3926                dispatch_rows(pool, inter, &run);
3927                true
3928            }
3929            // Q8Row gate + Q8Row up — one row pass over both i8 streams.
3930            // Q8_2f stays out on purpose: its column field prescales the
3931            // activations PER TENSOR, which breaks this fn's shared
3932            // split_act contract — it keeps the two-dispatch path.
3933            (
3934                Self::Mapped {
3935                    dtype: TensorDtype::Q8Row,
3936                    row_scale: g_rs,
3937                    ..
3938                },
3939                Self::Mapped {
3940                    dtype: TensorDtype::Q8Row,
3941                    row_scale: u_rs,
3942                    ..
3943                },
3944            ) => {
3945                let g_bytes = gate.quant_bytes();
3946                let u_bytes = up.quant_bytes();
3947                let cols = gate.cols();
3948                let run = move |start: usize, end: usize| {
3949                    for r in start..end {
3950                        let gv = q8_row_dot(&g_bytes[r * cols..(r + 1) * cols], act) * g_rs[r];
3951                        let uv = q8_row_dot(&u_bytes[r * cols..(r + 1) * cols], act) * u_rs[r];
3952                        // SAFETY: disjoint row ranges per worker.
3953                        unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3954                    }
3955                };
3956                dispatch_rows(pool, inter, &run);
3957                true
3958            }
3959            // Q1T gate + Q1T up
3960            (
3961                Self::Mapped {
3962                    dtype: TensorDtype::Q1T,
3963                    ..
3964                },
3965                Self::Mapped {
3966                    dtype: TensorDtype::Q1T,
3967                    ..
3968                },
3969            ) => {
3970                const TILE: usize = cortiq_core::quant::Q1T_TILE;
3971                let g_bytes = gate.quant_bytes();
3972                let u_bytes = up.quant_bytes();
3973                let gpr = gate.cols() / GROUP_SIZE;
3974                let (g_rp, g_ent, g_ov) = q1t_overlay(g_bytes, inter * gpr * TILE, inter);
3975                let (u_rp, u_ent, u_ov) = q1t_overlay(u_bytes, inter * gpr * TILE, inter);
3976                let run = move |start: usize, end: usize| {
3977                    for r in start..end {
3978                        let mut gv = q1t_dot_row_i8(g_bytes, r, gpr, &act.xq) * act.sx;
3979                        let mut uv = q1t_dot_row_i8(u_bytes, r, gpr, &act.xq) * act.sx;
3980                        for &(j, xv) in &act.outliers {
3981                            gv += q1t_base_weight(g_bytes, r, gpr, j) * xv;
3982                            uv += q1t_base_weight(u_bytes, r, gpr, j) * xv;
3983                        }
3984                        gv += q1t_row_outlier_correction(g_bytes, r, g_rp, g_ent, g_ov, x_ref);
3985                        uv += q1t_row_outlier_correction(u_bytes, r, u_rp, u_ent, u_ov, x_ref);
3986                        // SAFETY: disjoint row ranges per worker.
3987                        unsafe { *out_addr.at(r) = silu_mul_limited(gv, uv, limit) };
3988                    }
3989                };
3990                dispatch_rows(pool, inter, &run);
3991                true
3992            }
3993            _ => false,
3994        }
3995    }
3996
3997    /// Every routed expert's fused gate/up/SiLU under ONE pool dispatch.
3998    ///
3999    /// The per-expert path pays a pool barrier per expert per stage: at 9
4000    /// experts over 40 layers that is ~720 barriers a token, and a decode
4001    /// profile of Qwen3.6-35B-A3B showed the pool parked in
4002    /// `psynch_cvwait` about twice as long as it spent computing. Laying
4003    /// every expert's rows end-to-end in one virtual row space collapses
4004    /// the stage to a single dispatch. The per-row body is the
4005    /// single-expert q4tp arm verbatim, so outputs are bit-identical.
4006    ///
4007    /// `false` = something is outside the fused q4tp kernel (dtype, shape,
4008    /// or a transformed tensor); the caller walks the ordinary per-expert
4009    /// path. Float activations use the same exact scalar rows, still fused
4010    /// under one pool dispatch.
4011    pub fn moe_gate_up_many(
4012        pairs: &[(&QTensor, &QTensor)],
4013        x: &[f32],
4014        outs: &mut [Vec<f32>],
4015        pool: Option<&Pool>,
4016    ) -> bool {
4017        Self::moe_gate_up_many_limited(pairs, x, outs, 0.0, pool)
4018    }
4019
4020    /// Batched gate/up/SiLU with the optional GLM clamp.  The public legacy
4021    /// helper above keeps its historical unclamped semantics; callers that
4022    /// implement a reference with a positive SwiGLU limit use this variant.
4023    pub fn moe_gate_up_many_limited(
4024        pairs: &[(&QTensor, &QTensor)],
4025        x: &[f32],
4026        outs: &mut [Vec<f32>],
4027        limit: f32,
4028        pool: Option<&Pool>,
4029    ) -> bool {
4030        if Self::any_exact_cpu_default_pairs(pairs) {
4031            return float_activations_scope(|| {
4032                Self::moe_gate_up_many_limited(pairs, x, outs, limit, pool)
4033            });
4034        }
4035        if pairs.is_empty() || pairs.len() != outs.len() {
4036            return false;
4037        }
4038        if pairs.iter().any(|(g, u)| g.per_tensor_only() || u.per_tensor_only()) {
4039            return false;
4040        }
4041        if !a8w8_enabled() {
4042            if limit > 0.0 {
4043                // The exact-row fallback applies no SwiGLU clamp; the
4044                // caller's per-expert path carries it instead.
4045                return false;
4046            }
4047            let groups = vec![vec![0]; pairs.len()];
4048            return Self::moe_gate_up_rows(pairs, &groups, x, outs, pool);
4049        }
4050        let inter = pairs[0].0.rows();
4051        let cols = pairs[0].0.cols();
4052        if cols % GROUP_SIZE != 0 {
4053            return false;
4054        }
4055        let gpr = cols / GROUP_SIZE;
4056        // Uniform layout across every routed pair: q4tp, or the 2-bit
4057        // profile's q2tp gate/up (the W2 class). Mixed sets refuse.
4058        let q2 = matches!(
4059            pairs[0].0,
4060            Self::Mapped {
4061                dtype: TensorDtype::Q2TiledP,
4062                ..
4063            }
4064        );
4065        let want = if q2 {
4066            TensorDtype::Q2TiledP
4067        } else {
4068            TensorDtype::Q4TiledP
4069        };
4070        let mut views = Vec::with_capacity(pairs.len() * 2);
4071        for ((g, u), o) in pairs.iter().zip(outs.iter()) {
4072            let both = matches!(g, Self::Mapped { dtype, .. } if *dtype == want)
4073                && matches!(u, Self::Mapped { dtype, .. } if *dtype == want);
4074            if !both
4075                || g.rows() != inter
4076                || u.rows() != inter
4077                || g.cols() != cols
4078                || u.cols() != cols
4079                || o.len() != inter
4080            {
4081                return false;
4082            }
4083            let mk = if q2 { Q4tpView::new_q2 } else { Q4tpView::new };
4084            views.push(mk(g.quant_bytes(), inter, cols));
4085            views.push(mk(u.quant_bytes(), inter, cols));
4086        }
4087        let act = split_act(x);
4088        let gsum = if q2 {
4089            q1_group_sums(&act.xq, gpr)
4090        } else {
4091            Vec::new()
4092        };
4093        let (act, gsum) = (&act, &gsum);
4094        let ptrs: Vec<SendMut> = outs.iter_mut().map(|o| SendMut(o.as_mut_ptr())).collect();
4095        let (views, ptrs) = (&views, &ptrs);
4096        let run = |start: usize, end: usize| {
4097            with_krows(gpr, |gsc, usc| {
4098                for flat in start..end {
4099                    let (e, r) = (flat / inter, flat % inter);
4100                    let gv_view = &views[e * 2];
4101                    let uv_view = &views[e * 2 + 1];
4102                    gv_view.scales_into(r, gpr, gsc);
4103                    uv_view.scales_into(r, gpr, usc);
4104                    let (mut gv, mut uv) = if q2 {
4105                        (
4106                            dot_q2tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsum, gsc) * act.sx,
4107                            dot_q2tp_row_i8(uv_view.nib, r, gpr, &act.xq, gsum, usc) * act.sx,
4108                        )
4109                    } else {
4110                        (
4111                            dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsc) * act.sx,
4112                            dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, usc) * act.sx,
4113                        )
4114                    };
4115                    for &(j, xv) in &act.outliers {
4116                        let (og, ou) = if q2 {
4117                            (
4118                                q2tp_outlier(gv_view.nib, r, gpr, j, gsc),
4119                                q2tp_outlier(uv_view.nib, r, gpr, j, usc),
4120                            )
4121                        } else {
4122                            (
4123                                q4tp_outlier(gv_view.nib, r, gpr, j, gsc),
4124                                q4tp_outlier(uv_view.nib, r, gpr, j, usc),
4125                            )
4126                        };
4127                        gv += og.0 * og.1 * xv;
4128                        uv += ou.0 * ou.1 * xv;
4129                    }
4130                    // SAFETY: one worker owns each (expert, row) pair.
4131                    unsafe { *ptrs[e].at(r) = silu_mul_limited(gv, uv, limit) };
4132                }
4133            });
4134        };
4135        dispatch_rows(pool, pairs.len() * inter, &run);
4136        true
4137    }
4138
4139    /// Every routed expert's down projection, weighted and summed into
4140    /// `out`, under ONE pool dispatch.
4141    ///
4142    /// Partitioned by OUTPUT row rather than by expert: each row is owned
4143    /// by a single worker, so the experts are summed in the caller's order
4144    /// — the same sequence of f32 adds the serial `out[i] += w·eo[i]` loop
4145    /// performs, hence bit-identical. Partitioning by expert instead would
4146    /// race on the shared accumulator.
4147    pub fn moe_down_many(
4148        downs: &[&QTensor],
4149        gs: &[Vec<f32>],
4150        weights: &[f32],
4151        out: &mut [f32],
4152        pool: Option<&Pool>,
4153    ) -> bool {
4154        if Self::any_exact_cpu_default(downs) {
4155            return float_activations_scope(|| {
4156                Self::moe_down_many(downs, gs, weights, out, pool)
4157            });
4158        }
4159        if downs.is_empty() || downs.len() != gs.len() || downs.len() != weights.len() {
4160            return false;
4161        }
4162        if downs.iter().any(|d| d.per_tensor_only()) {
4163            return false;
4164        }
4165        if !a8w8_enabled() {
4166            let mut terms = vec![vec![0.0; out.len()]; downs.len()];
4167            if !Self::moe_down_rows(downs, &vec![1; downs.len()], gs, &mut terms, pool) {
4168                return false;
4169            }
4170            out.fill(0.0);
4171            for (row, &w) in terms.iter().zip(weights) {
4172                for (o, &v) in out.iter_mut().zip(row) {
4173                    *o += w * v;
4174                }
4175            }
4176            return true;
4177        }
4178        let rows = out.len();
4179        let cols = downs[0].cols();
4180        if cols % GROUP_SIZE != 0 {
4181            return false;
4182        }
4183        let gpr = cols / GROUP_SIZE;
4184        let mut views = Vec::with_capacity(downs.len());
4185        for (d, g) in downs.iter().zip(gs.iter()) {
4186            if !matches!(
4187                d,
4188                Self::Mapped {
4189                    dtype: TensorDtype::Q4TiledP,
4190                    ..
4191                }
4192            ) || d.rows() != rows
4193                || d.cols() != cols
4194                || g.len() != cols
4195            {
4196                return false;
4197            }
4198            views.push(Q4tpView::new(d.quant_bytes(), rows, cols));
4199        }
4200        // One int8 split per expert — the activation vectors differ.
4201        let acts: Vec<SplitAct> = gs.iter().map(|g| split_act(g)).collect();
4202        // Partitioned by OUTPUT row, with the experts folded inside: each
4203        // row is owned by one worker, so they are summed in the caller's
4204        // order — the same f32 sequence the serial `out[i] += w·eo[i]`
4205        // loop produces. Partitioning by expert instead would either race
4206        // on the accumulator or need a scratch plane and a second pass;
4207        // measured, that variant was a wash, so this keeps the simpler
4208        // shape.
4209        let out_addr = SendMut(out.as_mut_ptr());
4210        let (views, acts, weights) = (&views, &acts, &weights);
4211        let run = |start: usize, end: usize| {
4212            with_krow(gpr, |sc| {
4213                for r in start..end {
4214                    let mut acc = 0f32;
4215                    for (e, v) in views.iter().enumerate() {
4216                        v.scales_into(r, gpr, sc);
4217                        let a = &acts[e];
4218                        let mut d = dot_q4tp_row_i8(v.nib, r, gpr, &a.xq, sc) * a.sx;
4219                        for &(j, xv) in &a.outliers {
4220                            let (w, s) = q4tp_outlier(v.nib, r, gpr, j, sc);
4221                            d += w * s * xv;
4222                        }
4223                        acc += weights[e] * d;
4224                    }
4225                    // SAFETY: disjoint row ranges per worker.
4226                    unsafe { *out_addr.at(r) = acc };
4227                }
4228            });
4229        };
4230        dispatch_rows(pool, rows, &run);
4231        true
4232    }
4233
4234    /// `moe_gate_up_many` for SEVERAL tokens at once, decode-exact: expert
4235    /// `e` (`pairs[e]`, q4tp) serves the tokens `groups[e]` (row indices
4236    /// into `xs`, each `cols` wide). Every (expert, token) output is
4237    /// bit-identical to `moe_gate_up_many` run on that token alone — the
4238    /// same int8 activation split, VNNI dots, outlier terms and inline
4239    /// SiLU — while each weight row is read once for all the tokens routed
4240    /// to its expert (the speculative verify's expert sharing). `outs` is
4241    /// flat in (expert, token-of-group) order. False = not covered (not
4242    /// q4tp): the caller takes the per-token path. With float activations,
4243    /// the exact scalar row kernel replaces the int8 dot without changing
4244    /// the shared dispatch or route-order reduction.
4245    pub fn moe_gate_up_rows(
4246        pairs: &[(&QTensor, &QTensor)],
4247        groups: &[Vec<usize>],
4248        xs: &[f32],
4249        outs: &mut [Vec<f32>],
4250        pool: Option<&Pool>,
4251    ) -> bool {
4252        if Self::any_exact_cpu_default_pairs(pairs) {
4253            return float_activations_scope(|| {
4254                Self::moe_gate_up_rows(pairs, groups, xs, outs, pool)
4255            });
4256        }
4257        if pairs.is_empty() || pairs.len() != groups.len() {
4258            return false;
4259        }
4260        if pairs.iter().any(|(g, u)| g.per_tensor_only() || u.per_tensor_only()) {
4261            return false;
4262        }
4263        let inter = pairs[0].0.rows();
4264        let cols = pairs[0].0.cols();
4265        let n_pairs: usize = groups.iter().map(|g| g.len()).sum();
4266        if cols == 0 || cols % GROUP_SIZE != 0 || outs.len() != n_pairs || xs.len() % cols != 0 {
4267            return false;
4268        }
4269        let b = xs.len() / cols;
4270        let gpr = cols / GROUP_SIZE;
4271        let mut views = Vec::with_capacity(pairs.len() * 2);
4272        for (g, u) in pairs {
4273            let q4tp = |t: &QTensor| {
4274                matches!(
4275                    t,
4276                    Self::Mapped {
4277                        dtype: TensorDtype::Q4TiledP,
4278                        ..
4279                    }
4280                )
4281            };
4282            if g.has_prism_contract()
4283                || u.has_prism_contract()
4284                || !q4tp(g)
4285                || !q4tp(u)
4286                || g.rows() != inter
4287                || u.rows() != inter
4288                || g.cols() != cols
4289                || u.cols() != cols
4290            {
4291                return false;
4292            }
4293            views.push(Q4tpView::new(g.quant_bytes(), inter, cols));
4294            views.push(Q4tpView::new(u.quant_bytes(), inter, cols));
4295        }
4296        if outs.iter().any(|o| o.len() != inter) || groups.iter().flatten().any(|&t| t >= b) {
4297            return false;
4298        }
4299        let quantized = a8w8_enabled();
4300        let acts: Vec<SplitAct> = if quantized {
4301            (0..b)
4302                .map(|t| split_act(&xs[t * cols..(t + 1) * cols]))
4303                .collect()
4304        } else {
4305            Vec::new()
4306        };
4307        let mut offs = Vec::with_capacity(groups.len());
4308        let mut o = 0usize;
4309        for g in groups {
4310            offs.push(o);
4311            o += g.len();
4312        }
4313        let ptrs: Vec<SendMut> = outs.iter_mut().map(|o| SendMut(o.as_mut_ptr())).collect();
4314        let (views, ptrs, acts, offs) = (&views, &ptrs, &acts, &offs);
4315        let run = |start: usize, end: usize| {
4316            let (mut gsc, mut usc) = (vec![0f32; gpr], vec![0f32; gpr]);
4317            for flat in start..end {
4318                let (e, r) = (flat / inter, flat % inter);
4319                let (gv_view, uv_view) = (&views[e * 2], &views[e * 2 + 1]);
4320                gv_view.scales_into(r, gpr, &mut gsc);
4321                uv_view.scales_into(r, gpr, &mut usc);
4322                for (k, &t) in groups[e].iter().enumerate() {
4323                    if !quantized {
4324                        let x = &xs[t * cols..(t + 1) * cols];
4325                        let gv = q4tp_row_exact(gv_view.nib, r, gpr, x, &gsc);
4326                        let uv = q4tp_row_exact(uv_view.nib, r, gpr, x, &usc);
4327                        unsafe { *ptrs[offs[e] + k].at(r) = (gv / (1.0 + (-gv).exp())) * uv };
4328                        continue;
4329                    }
4330                    let act = &acts[t];
4331                    let mut gv = dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, &gsc) * act.sx;
4332                    let mut uv = dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, &usc) * act.sx;
4333                    for &(j, xv) in &act.outliers {
4334                        let og = q4tp_outlier(gv_view.nib, r, gpr, j, &gsc);
4335                        let ou = q4tp_outlier(uv_view.nib, r, gpr, j, &usc);
4336                        gv += og.0 * og.1 * xv;
4337                        uv += ou.0 * ou.1 * xv;
4338                    }
4339                    let silu_g = gv / (1.0 + (-gv).exp());
4340                    // SAFETY: one worker owns each (expert, row) cell of
4341                    // every output of the expert's group.
4342                    unsafe { *ptrs[offs[e] + k].at(r) = silu_g * uv };
4343                }
4344            }
4345        };
4346        dispatch_rows(pool, pairs.len() * inter, &run);
4347        true
4348    }
4349
4350    /// The per-(expert, token) down terms `moe_down_many` weights and sums,
4351    /// for SEVERAL tokens: `outs[p][o] = down_e[o] · gs[p]` (int8 split of
4352    /// `gs[p]`, VNNI dot, outlier terms — bit-identical to that kernel's
4353    /// `d`), each down row read once for its expert's whole group. The
4354    /// caller sums `w·d` per token in its route order, which reproduces
4355    /// `moe_down_many`'s f32 sequence exactly. Layout as `moe_gate_up_rows`.
4356    pub fn moe_down_rows(
4357        downs: &[&QTensor],
4358        group_lens: &[usize],
4359        gs: &[Vec<f32>],
4360        outs: &mut [Vec<f32>],
4361        pool: Option<&Pool>,
4362    ) -> bool {
4363        if Self::any_exact_cpu_default(downs) {
4364            return float_activations_scope(|| {
4365                Self::moe_down_rows(downs, group_lens, gs, outs, pool)
4366            });
4367        }
4368        if downs.is_empty() || downs.len() != group_lens.len() {
4369            return false;
4370        }
4371        if downs.iter().any(|d| d.per_tensor_only()) {
4372            return false;
4373        }
4374        let rows = downs[0].rows();
4375        let cols = downs[0].cols();
4376        let n_pairs: usize = group_lens.iter().sum();
4377        if cols == 0 || cols % GROUP_SIZE != 0 || gs.len() != n_pairs || outs.len() != n_pairs {
4378            return false;
4379        }
4380        let gpr = cols / GROUP_SIZE;
4381        let mut views = Vec::with_capacity(downs.len());
4382        for d in downs {
4383            if d.has_prism_contract()
4384                || !matches!(
4385                    d,
4386                    Self::Mapped {
4387                        dtype: TensorDtype::Q4TiledP,
4388                        ..
4389                    }
4390                )
4391                || d.rows() != rows
4392                || d.cols() != cols
4393            {
4394                return false;
4395            }
4396            views.push(Q4tpView::new(d.quant_bytes(), rows, cols));
4397        }
4398        if gs.iter().any(|g| g.len() != cols) || outs.iter().any(|o| o.len() != rows) {
4399            return false;
4400        }
4401        let quantized = a8w8_enabled();
4402        let acts: Vec<SplitAct> = if quantized {
4403            gs.iter().map(|g| split_act(g)).collect()
4404        } else {
4405            Vec::new()
4406        };
4407        let mut offs = Vec::with_capacity(group_lens.len());
4408        let mut o = 0usize;
4409        for &l in group_lens {
4410            offs.push(o);
4411            o += l;
4412        }
4413        let ptrs: Vec<SendMut> = outs.iter_mut().map(|o| SendMut(o.as_mut_ptr())).collect();
4414        let (views, ptrs, acts, offs) = (&views, &ptrs, &acts, &offs);
4415        let run = |start: usize, end: usize| {
4416            let mut sc = vec![0f32; gpr];
4417            for flat in start..end {
4418                let (e, r) = (flat / rows, flat % rows);
4419                let v = &views[e];
4420                v.scales_into(r, gpr, &mut sc);
4421                for k in 0..group_lens[e] {
4422                    if !quantized {
4423                        let d = q4tp_row_exact(v.nib, r, gpr, &gs[offs[e] + k], &sc);
4424                        unsafe { *ptrs[offs[e] + k].at(r) = d };
4425                        continue;
4426                    }
4427                    let a = &acts[offs[e] + k];
4428                    let mut d = dot_q4tp_row_i8(v.nib, r, gpr, &a.xq, &sc) * a.sx;
4429                    for &(j, xv) in &a.outliers {
4430                        let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
4431                        d += w * s * xv;
4432                    }
4433                    // SAFETY: one worker owns each (expert, row) cell.
4434                    unsafe { *ptrs[offs[e] + k].at(r) = d };
4435                }
4436            }
4437        };
4438        dispatch_rows(pool, downs.len() * rows, &run);
4439        true
4440    }
4441
4442    /// Whole experts for SEVERAL tokens in one dispatch: expert `e`
4443    /// (`experts[e]` = its gate, up and down, q4tp) serves the tokens
4444    /// `groups[e]` (row indices into `xs`, each `hidden` wide), and
4445    /// `out[p]` (`hidden` wide, flat in (expert, token-of-group) order) is
4446    /// the term `down_e · (silu(gate_e · x) ⊙ (up_e · x))` the caller weights
4447    /// and sums. The math of `moe_gate_up_rows` followed by `moe_down_rows`
4448    /// (the same int8 activation split, outlier terms and SiLU) with every
4449    /// row's sum taken in sixteen float lanes and the SiLU's exponent a
4450    /// sixteen-lane polynomial, so a value differs from theirs by rounding
4451    /// only; the work is one dispatch and one pass over the weights at the
4452    /// speed of the memory bus (see `qtensor/moe_fast.rs`). False =
4453    /// not covered (no AVX-512 VNNI, not q4tp, float activations, a shape
4454    /// outside 64-wide steps): the caller takes the two row kernels.
4455    pub fn moe_expert_rows_fast(
4456        experts: &[(&QTensor, &QTensor, &QTensor)],
4457        groups: &[Vec<usize>],
4458        xs: &[f32],
4459        out: &mut [f32],
4460        pool: Option<&Pool>,
4461    ) -> bool {
4462        #[cfg(target_arch = "x86_64")]
4463        {
4464            moe_fast::run(experts, groups, xs, out, pool)
4465        }
4466        #[cfg(not(target_arch = "x86_64"))]
4467        {
4468            let _ = (experts, groups, xs, out, pool);
4469            false
4470        }
4471    }
4472
4473    /// Turn the clock of [`QTensor::moe_expert_rows_fast`] on or off; returns what it was.
4474    pub fn moe_expert_rows_fast_prof(on: bool) -> bool {
4475        #[cfg(target_arch = "x86_64")]
4476        {
4477            moe_fast::PROF_ON.swap(on, Ordering::Relaxed)
4478        }
4479        #[cfg(not(target_arch = "x86_64"))]
4480        {
4481            let _ = on;
4482            false
4483        }
4484    }
4485
4486    /// What the clock of [`QTensor::moe_expert_rows_fast`] summed since the last call, as one line; `None`
4487    /// when the kernel did not run.
4488    pub fn moe_expert_rows_fast_prof_take() -> Option<String> {
4489        #[cfg(target_arch = "x86_64")]
4490        {
4491            let p = moe_fast::prof_take();
4492            let (calls, parts, experts) = (p[0], p[1], p[2]);
4493            if calls == 0 || parts == 0 || experts == 0 {
4494                return None;
4495            }
4496            // thread time a phase took of an expert, microseconds
4497            let us = |ns: u64| ns as f64 / 1e3 / experts as f64;
4498            let call = |ns: u64| ns as f64 / 1e3 / calls as f64;
4499            Some(format!(
4500                "{calls} calls of {:.1} threads, {experts} experts, {} pairs, outliers {:.1} a token row and {:.1} a gate row; \
4501                 a call {:.0} us of setup ({:.0} the tokens' split) and {:.0} us of wall; thread time an expert, us: \
4502                 entry {:.0}, gate/up {:.0}, gate rows {:.0}, down {:.0}, waiting {:.0}",
4503                parts as f64 / calls as f64,
4504                p[3],
4505                p[5] as f64 / p[4].max(1) as f64,
4506                p[6] as f64 / p[3].max(1) as f64,
4507                call(p[7]),
4508                call(p[8]),
4509                call(p[9]),
4510                us(p[10]),
4511                us(p[11]),
4512                us(p[12]),
4513                us(p[13]),
4514                us(p[14]),
4515            ))
4516        }
4517        #[cfg(not(target_arch = "x86_64"))]
4518        None
4519    }
4520
4521    /// Set a knob of [`QTensor::moe_expert_rows_fast`] for the calls to come: 0 = the bytes of weights a
4522    /// row prefetches ahead of itself, 1 and 2 = the blocks of rows a core takes at a time of gate/up
4523    /// and of down. For a measurement that compares them in one process; the defaults come from
4524    /// `CMF_QWEN_CPU_PF` and `CMF_QWEN_CPU_CLAIM`. Returns the knob's value before the call.
4525    pub fn moe_expert_rows_fast_tune(knob: usize, value: usize) -> usize {
4526        #[cfg(target_arch = "x86_64")]
4527        {
4528            moe_fast::set_tune(knob, value)
4529        }
4530        #[cfg(not(target_arch = "x86_64"))]
4531        {
4532            let _ = (knob, value);
4533            0
4534        }
4535    }
4536
4537    /// Whether this host has the cores for [`QTensor::moe_expert_rows_fast`] (AVX-512 VNNI with the int8
4538    /// activation path on).
4539    pub fn moe_expert_rows_fast_available() -> bool {
4540        #[cfg(target_arch = "x86_64")]
4541        {
4542            avx512vnni_enabled() && a8w8_enabled()
4543        }
4544        #[cfg(not(target_arch = "x86_64"))]
4545        {
4546            false
4547        }
4548    }
4549}
4550
4551/// Batched q8 kernel: same math as qmatvec, the row makes a single
4552/// pass from memory for the whole batch.
4553/// Accelerate CBLAS — the Apple AMX matrix units, the same engine
4554/// llama.cpp's `-ngl 0` prefill rides via ggml-blas.
4555#[cfg(target_os = "macos")]
4556mod accel_blas {
4557    #[link(name = "Accelerate", kind = "framework")]
4558    unsafe extern "C" {
4559        pub fn cblas_sgemm(
4560            order: i32,
4561            trans_a: i32,
4562            trans_b: i32,
4563            m: i32,
4564            n: i32,
4565            k: i32,
4566            alpha: f32,
4567            a: *const f32,
4568            lda: i32,
4569            b: *const f32,
4570            ldb: i32,
4571            beta: f32,
4572            c: *mut f32,
4573            ldc: i32,
4574        );
4575    }
4576}
4577
4578#[cfg(target_os = "macos")]
4579pub(crate) fn accel_gemm_enabled() -> bool {
4580    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4581    *ON.get_or_init(|| std::env::var("CMF_ACCEL").map(|v| v != "0").unwrap_or(true))
4582}
4583
4584/// Off macOS the "accel" GEMM is the portable NEON micro-kernel below —
4585/// same entry point, so the batched-attention path opens on mobile.
4586#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
4587pub(crate) fn accel_gemm_enabled() -> bool {
4588    true
4589}
4590
4591/// On x86-64 the "accel" GEMM is the AVX2 micro-kernel below, so the
4592/// batched attention opens there too; `CMF_CPU_ATTEND=0` keeps the
4593/// per-position attend, and a CPU without AVX2/FMA never leaves it.
4594#[cfg(all(target_arch = "x86_64", not(target_os = "macos")))]
4595pub(crate) fn accel_gemm_enabled() -> bool {
4596    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4597    *ON.get_or_init(|| {
4598        std::is_x86_feature_detected!("avx2")
4599            && std::is_x86_feature_detected!("fma")
4600            && std::env::var("CMF_CPU_ATTEND").map(|v| v != "0").unwrap_or(true)
4601    })
4602}
4603
4604/// Any other architecture has no CPU micro-GEMM: the per-position attend.
4605#[cfg(not(any(target_os = "macos", target_arch = "aarch64", target_arch = "x86_64")))]
4606#[allow(dead_code)]
4607pub(crate) fn accel_gemm_enabled() -> bool {
4608    false
4609}
4610
4611/// Portable NEON f32 GEMM (row-major, optional Bᵀ): a 4×8 fmla
4612/// micro-kernel with A broadcast against B panels — the mobile stand-in
4613/// for Accelerate in the batched causal attention (QKᵀ and P·V). Not a
4614/// BLAS: shapes here are the attention panels (m ≤ heads·chunk,
4615/// k = head_dim or context), and the goal is removing the per-position
4616/// quadratic wall, not peak GEMM.
4617#[cfg(target_arch = "aarch64")]
4618#[allow(clippy::too_many_arguments)]
4619pub(crate) fn neon_gemm_rm(
4620    m: usize,
4621    n: usize,
4622    k: usize,
4623    alpha: f32,
4624    a: &[f32],
4625    lda: usize,
4626    b_mat: &[f32],
4627    ldb: usize,
4628    b_rows_are_n: bool,
4629    c: &mut [f32],
4630    ldc: usize,
4631) {
4632    debug_assert!(a.len() >= (m - 1) * lda + k);
4633    debug_assert!(c.len() >= (m - 1) * ldc + n);
4634    // SAFETY: bounds asserted above; NEON is baseline on aarch64.
4635    unsafe {
4636        use core::arch::aarch64::*;
4637        let mut i = 0usize;
4638        while i < m {
4639            let mi = (m - i).min(4);
4640            let mut j = 0usize;
4641            while j < n {
4642                let nj = (n - j).min(8);
4643                if mi == 4 && nj == 8 {
4644                    let (mut c0a, mut c0b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
4645                    let (mut c1a, mut c1b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
4646                    let (mut c2a, mut c2b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
4647                    let (mut c3a, mut c3b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
4648                    for p in 0..k {
4649                        let (b0, b1) = if b_rows_are_n {
4650                            // B is [n, k]: column p of Bᵀ = element p of
4651                            // eight consecutive B rows — gathered.
4652                            let base = b_mat.as_ptr().add(j * ldb + p);
4653                            let g = |o: usize| *base.add(o * ldb);
4654                            ([g(0), g(1), g(2), g(3)], [g(4), g(5), g(6), g(7)])
4655                        } else {
4656                            let base = b_mat.as_ptr().add(p * ldb + j);
4657                            (
4658                                [*base, *base.add(1), *base.add(2), *base.add(3)],
4659                                [*base.add(4), *base.add(5), *base.add(6), *base.add(7)],
4660                            )
4661                        };
4662                        let bv0 = vld1q_f32(b0.as_ptr());
4663                        let bv1 = vld1q_f32(b1.as_ptr());
4664                        let a0 = vdupq_n_f32(*a.as_ptr().add(i * lda + p));
4665                        let a1 = vdupq_n_f32(*a.as_ptr().add((i + 1) * lda + p));
4666                        let a2 = vdupq_n_f32(*a.as_ptr().add((i + 2) * lda + p));
4667                        let a3 = vdupq_n_f32(*a.as_ptr().add((i + 3) * lda + p));
4668                        c0a = vfmaq_f32(c0a, a0, bv0);
4669                        c0b = vfmaq_f32(c0b, a0, bv1);
4670                        c1a = vfmaq_f32(c1a, a1, bv0);
4671                        c1b = vfmaq_f32(c1b, a1, bv1);
4672                        c2a = vfmaq_f32(c2a, a2, bv0);
4673                        c2b = vfmaq_f32(c2b, a2, bv1);
4674                        c3a = vfmaq_f32(c3a, a3, bv0);
4675                        c3b = vfmaq_f32(c3b, a3, bv1);
4676                    }
4677                    let al = vdupq_n_f32(alpha);
4678                    for (r, (ca, cb)) in [(c0a, c0b), (c1a, c1b), (c2a, c2b), (c3a, c3b)]
4679                        .iter()
4680                        .enumerate()
4681                    {
4682                        let dst = c.as_mut_ptr().add((i + r) * ldc + j);
4683                        vst1q_f32(dst, vmulq_f32(*ca, al));
4684                        vst1q_f32(dst.add(4), vmulq_f32(*cb, al));
4685                    }
4686                } else {
4687                    for r in 0..mi {
4688                        for q in 0..nj {
4689                            let mut acc = 0f32;
4690                            for p in 0..k {
4691                                let bv = if b_rows_are_n {
4692                                    b_mat[(j + q) * ldb + p]
4693                                } else {
4694                                    b_mat[p * ldb + j + q]
4695                                };
4696                                acc += a[(i + r) * lda + p] * bv;
4697                            }
4698                            c[(i + r) * ldc + j + q] = acc * alpha;
4699                        }
4700                    }
4701                }
4702                j += nj;
4703            }
4704            i += mi;
4705        }
4706    }
4707}
4708
4709/// Off-macOS aarch64: the batched attention rides the NEON micro-GEMM.
4710#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
4711#[allow(clippy::too_many_arguments)]
4712pub(crate) fn sgemm_rm(
4713    m: usize,
4714    n: usize,
4715    k: usize,
4716    alpha: f32,
4717    a: &[f32],
4718    lda: usize,
4719    b_mat: &[f32],
4720    ldb: usize,
4721    b_rows_are_n: bool,
4722    c: &mut [f32],
4723    ldc: usize,
4724) {
4725    neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
4726}
4727
4728/// The portable micro-GEMM of the batched causal attention: NEON on
4729/// aarch64, AVX2 on x86-64. Same contract as `neon_gemm_rm`.
4730#[cfg(target_arch = "aarch64")]
4731#[allow(clippy::too_many_arguments)]
4732pub(crate) fn micro_gemm_rm(
4733    m: usize,
4734    n: usize,
4735    k: usize,
4736    alpha: f32,
4737    a: &[f32],
4738    lda: usize,
4739    b_mat: &[f32],
4740    ldb: usize,
4741    b_rows_are_n: bool,
4742    c: &mut [f32],
4743    ldc: usize,
4744) {
4745    neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
4746}
4747
4748/// The portable micro-GEMM of the batched causal attention (x86-64 arm).
4749#[cfg(target_arch = "x86_64")]
4750#[allow(clippy::too_many_arguments)]
4751pub(crate) fn micro_gemm_rm(
4752    m: usize,
4753    n: usize,
4754    k: usize,
4755    alpha: f32,
4756    a: &[f32],
4757    lda: usize,
4758    b_mat: &[f32],
4759    ldb: usize,
4760    b_rows_are_n: bool,
4761    c: &mut [f32],
4762    ldc: usize,
4763) {
4764    x86_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
4765}
4766
4767/// Off-macOS x86-64: the batched attention rides the AVX2 micro-GEMM.
4768#[cfg(all(target_arch = "x86_64", not(target_os = "macos")))]
4769#[allow(clippy::too_many_arguments)]
4770pub(crate) fn sgemm_rm(
4771    m: usize,
4772    n: usize,
4773    k: usize,
4774    alpha: f32,
4775    a: &[f32],
4776    lda: usize,
4777    b_mat: &[f32],
4778    ldb: usize,
4779    b_rows_are_n: bool,
4780    c: &mut [f32],
4781    ldc: usize,
4782) {
4783    x86_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
4784}
4785
4786/// Portable x86-64 f32 GEMM (row-major, optional Bᵀ): the x86 twin of
4787/// `neon_gemm_rm` for the batched causal attention (Q·Kᵀ and P·V). An
4788/// AVX2+FMA 4×16 micro-kernel with A broadcast against B panels, blocked
4789/// so a 128-column × 256-row slab of B stays in L2 while the rows of A
4790/// stream past it (the P·V multiply has k = context, so without the k
4791/// blocking V would be re-read once per four rows). Same contract:
4792/// C[m,n] = alpha·A[m,k] × B, B stored [k,n] — or [n,k] with
4793/// `b_rows_are_n`, which packs Bᵀ first (the attention callers never
4794/// take that arm: they pack Kᵀ themselves). Not a BLAS: the shapes are
4795/// attention panels, and the goal is the end of the per-position
4796/// quadratic wall on x86, not peak GEMM. Scalar on a CPU without
4797/// AVX2/FMA.
4798#[cfg(target_arch = "x86_64")]
4799#[allow(clippy::too_many_arguments)]
4800pub(crate) fn x86_gemm_rm(
4801    m: usize,
4802    n: usize,
4803    k: usize,
4804    alpha: f32,
4805    a: &[f32],
4806    lda: usize,
4807    b_mat: &[f32],
4808    ldb: usize,
4809    b_rows_are_n: bool,
4810    c: &mut [f32],
4811    ldc: usize,
4812) {
4813    if m == 0 || n == 0 {
4814        return;
4815    }
4816    debug_assert!(a.len() >= (m - 1) * lda + k);
4817    debug_assert!(c.len() >= (m - 1) * ldc + n);
4818    if b_rows_are_n {
4819        debug_assert!(k == 0 || b_mat.len() >= (n - 1) * ldb + k);
4820        let mut bt = vec![0f32; k * n];
4821        for j in 0..n {
4822            for (p, &v) in b_mat[j * ldb..j * ldb + k].iter().enumerate() {
4823                bt[p * n + j] = v;
4824            }
4825        }
4826        return x86_gemm_rm(m, n, k, alpha, a, lda, &bt, n, false, c, ldc);
4827    }
4828    if k == 0 {
4829        for r in 0..m {
4830            c[r * ldc..r * ldc + n].fill(0.0);
4831        }
4832        return;
4833    }
4834    debug_assert!(b_mat.len() >= (k - 1) * ldb + n);
4835    if std::is_x86_feature_detected!("avx2") && std::is_x86_feature_detected!("fma") {
4836        // SAFETY: the bounds hold (asserted above), and both features were
4837        // just detected on this CPU.
4838        unsafe { x86_gemm_avx2(m, n, k, alpha, a, lda, b_mat, ldb, c, ldc) }
4839    } else {
4840        for r in 0..m {
4841            for j in 0..n {
4842                let mut acc = 0f32;
4843                for p in 0..k {
4844                    acc += a[r * lda + p] * b_mat[p * ldb + j];
4845                }
4846                c[r * ldc + j] = acc * alpha;
4847            }
4848        }
4849    }
4850}
4851
4852/// The AVX2+FMA body of `x86_gemm_rm` (B stored [k,n]). Loop nest:
4853/// column block (128) → k block (256) → four rows of A → 16-column
4854/// panel; the first k block starts the accumulators at zero, later ones
4855/// resume from C, the last one applies `alpha`.
4856///
4857/// # Safety
4858/// The caller has checked `avx2` and `fma` and the slice bounds
4859/// (`a` ≥ (m−1)·lda + k, `b` ≥ (k−1)·ldb + n, `c` ≥ (m−1)·ldc + n).
4860#[cfg(target_arch = "x86_64")]
4861#[target_feature(enable = "avx2,fma")]
4862#[allow(clippy::too_many_arguments)]
4863unsafe fn x86_gemm_avx2(
4864    m: usize,
4865    n: usize,
4866    k: usize,
4867    alpha: f32,
4868    a: &[f32],
4869    lda: usize,
4870    b: &[f32],
4871    ldb: usize,
4872    c: &mut [f32],
4873    ldc: usize,
4874) {
4875    use core::arch::x86_64::*;
4876    const NC: usize = 128;
4877    const KC: usize = 256;
4878    unsafe {
4879        let ap = a.as_ptr();
4880        let bp = b.as_ptr();
4881        let cp = c.as_mut_ptr();
4882        let al = _mm256_set1_ps(alpha);
4883        let mut jc = 0usize;
4884        while jc < n {
4885            let jend = (jc + NC).min(n);
4886            let mut pc = 0usize;
4887            while pc < k {
4888                let pend = (pc + KC).min(k);
4889                let (first, last) = (pc == 0, pend == k);
4890                let mut i = 0usize;
4891                while i < m {
4892                    let mi = (m - i).min(4);
4893                    let mut j = jc;
4894                    while j < jend {
4895                        let nj = (jend - j).min(16);
4896                        if mi == 4 && nj == 16 {
4897                            let a0 = ap.add(i * lda);
4898                            let a1 = ap.add((i + 1) * lda);
4899                            let a2 = ap.add((i + 2) * lda);
4900                            let a3 = ap.add((i + 3) * lda);
4901                            let c0 = cp.add(i * ldc + j);
4902                            let c1 = cp.add((i + 1) * ldc + j);
4903                            let c2 = cp.add((i + 2) * ldc + j);
4904                            let c3 = cp.add((i + 3) * ldc + j);
4905                            let (mut c00, mut c01, mut c10, mut c11, mut c20, mut c21, mut c30, mut c31) =
4906                                if first {
4907                                    let z = _mm256_setzero_ps();
4908                                    (z, z, z, z, z, z, z, z)
4909                                } else {
4910                                    (
4911                                        _mm256_loadu_ps(c0),
4912                                        _mm256_loadu_ps(c0.add(8)),
4913                                        _mm256_loadu_ps(c1),
4914                                        _mm256_loadu_ps(c1.add(8)),
4915                                        _mm256_loadu_ps(c2),
4916                                        _mm256_loadu_ps(c2.add(8)),
4917                                        _mm256_loadu_ps(c3),
4918                                        _mm256_loadu_ps(c3.add(8)),
4919                                    )
4920                                };
4921                            for p in pc..pend {
4922                                let bb = bp.add(p * ldb + j);
4923                                let b0 = _mm256_loadu_ps(bb);
4924                                let b1 = _mm256_loadu_ps(bb.add(8));
4925                                let x0 = _mm256_broadcast_ss(&*a0.add(p));
4926                                c00 = _mm256_fmadd_ps(x0, b0, c00);
4927                                c01 = _mm256_fmadd_ps(x0, b1, c01);
4928                                let x1 = _mm256_broadcast_ss(&*a1.add(p));
4929                                c10 = _mm256_fmadd_ps(x1, b0, c10);
4930                                c11 = _mm256_fmadd_ps(x1, b1, c11);
4931                                let x2 = _mm256_broadcast_ss(&*a2.add(p));
4932                                c20 = _mm256_fmadd_ps(x2, b0, c20);
4933                                c21 = _mm256_fmadd_ps(x2, b1, c21);
4934                                let x3 = _mm256_broadcast_ss(&*a3.add(p));
4935                                c30 = _mm256_fmadd_ps(x3, b0, c30);
4936                                c31 = _mm256_fmadd_ps(x3, b1, c31);
4937                            }
4938                            if last {
4939                                c00 = _mm256_mul_ps(c00, al);
4940                                c01 = _mm256_mul_ps(c01, al);
4941                                c10 = _mm256_mul_ps(c10, al);
4942                                c11 = _mm256_mul_ps(c11, al);
4943                                c20 = _mm256_mul_ps(c20, al);
4944                                c21 = _mm256_mul_ps(c21, al);
4945                                c30 = _mm256_mul_ps(c30, al);
4946                                c31 = _mm256_mul_ps(c31, al);
4947                            }
4948                            _mm256_storeu_ps(c0, c00);
4949                            _mm256_storeu_ps(c0.add(8), c01);
4950                            _mm256_storeu_ps(c1, c10);
4951                            _mm256_storeu_ps(c1.add(8), c11);
4952                            _mm256_storeu_ps(c2, c20);
4953                            _mm256_storeu_ps(c2.add(8), c21);
4954                            _mm256_storeu_ps(c3, c30);
4955                            _mm256_storeu_ps(c3.add(8), c31);
4956                        } else {
4957                            // Row-at-a-time tail: fewer than four rows left,
4958                            // or a panel narrower than 16 columns.
4959                            for r in 0..mi {
4960                                let ar = ap.add((i + r) * lda);
4961                                let cr = cp.add((i + r) * ldc + j);
4962                                let mut q = 0usize;
4963                                while q + 8 <= nj {
4964                                    let mut acc = if first {
4965                                        _mm256_setzero_ps()
4966                                    } else {
4967                                        _mm256_loadu_ps(cr.add(q))
4968                                    };
4969                                    for p in pc..pend {
4970                                        acc = _mm256_fmadd_ps(
4971                                            _mm256_broadcast_ss(&*ar.add(p)),
4972                                            _mm256_loadu_ps(bp.add(p * ldb + j + q)),
4973                                            acc,
4974                                        );
4975                                    }
4976                                    if last {
4977                                        acc = _mm256_mul_ps(acc, al);
4978                                    }
4979                                    _mm256_storeu_ps(cr.add(q), acc);
4980                                    q += 8;
4981                                }
4982                                while q < nj {
4983                                    let mut acc = if first { 0.0 } else { *cr.add(q) };
4984                                    for p in pc..pend {
4985                                        acc += *ar.add(p) * *bp.add(p * ldb + j + q);
4986                                    }
4987                                    if last {
4988                                        acc *= alpha;
4989                                    }
4990                                    *cr.add(q) = acc;
4991                                    q += 1;
4992                                }
4993                            }
4994                        }
4995                        j += nj;
4996                    }
4997                    i += mi;
4998                }
4999                pc = pend;
5000            }
5001            jc = jend;
5002        }
5003    }
5004}
5005
5006/// Row-major f32 GEMM, exposed for offline tools (the AWNP pass builds a
5007/// per-layer projection and applies it to every expert; a naive triple loop
5008/// would turn a two-minute job into half an hour).
5009#[allow(clippy::too_many_arguments)]
5010pub fn sgemm_public(
5011    m: usize,
5012    n: usize,
5013    k: usize,
5014    alpha: f32,
5015    a: &[f32],
5016    lda: usize,
5017    b_mat: &[f32],
5018    ldb: usize,
5019    b_rows_are_n: bool,
5020    c: &mut [f32],
5021    ldc: usize,
5022) {
5023    #[cfg(any(target_os = "macos", target_arch = "aarch64"))]
5024    {
5025        sgemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
5026    }
5027    // x86 without Accelerate has no sgemm_rm (the specialized paths there are
5028    // quantized kernels): the packed AVX/FMA/AVX-512 kernels of
5029    // matrixmultiply, single-threaded — offline callers split rows across
5030    // their own threads.
5031    #[cfg(not(any(target_os = "macos", target_arch = "aarch64")))]
5032    {
5033        if m == 0 || n == 0 {
5034            return;
5035        }
5036        assert!(a.len() >= (m - 1) * lda + k && c.len() >= (m - 1) * ldc + n);
5037        let (rsb, csb) = if b_rows_are_n { (1, ldb as isize) } else { (ldb as isize, 1) };
5038        assert!(k == 0 || b_mat.len() > if b_rows_are_n { (n - 1) * ldb + k - 1 } else { (k - 1) * ldb + n - 1 });
5039        // SAFETY: the asserts above bound every element the strides reach.
5040        unsafe {
5041            matrixmultiply::sgemm(
5042                m,
5043                k,
5044                n,
5045                alpha,
5046                a.as_ptr(),
5047                lda as isize,
5048                1,
5049                b_mat.as_ptr(),
5050                rsb,
5051                csb,
5052                0.0,
5053                c.as_mut_ptr(),
5054                ldc as isize,
5055                1,
5056            );
5057        }
5058    }
5059}
5060
5061/// Row-major f32 GEMM on Accelerate: C[m,n] = alpha·A[m,k] × B(ᵀ).
5062/// `b_rows_are_n` = true multiplies by Bᵀ where B is stored [n, k].
5063#[cfg(target_os = "macos")]
5064#[allow(clippy::too_many_arguments)]
5065pub(crate) fn sgemm_rm(
5066    m: usize,
5067    n: usize,
5068    k: usize,
5069    alpha: f32,
5070    a: &[f32],
5071    lda: usize,
5072    b_mat: &[f32],
5073    ldb: usize,
5074    b_rows_are_n: bool,
5075    c: &mut [f32],
5076    ldc: usize,
5077) {
5078    debug_assert!(a.len() >= (m - 1) * lda + k);
5079    debug_assert!(c.len() >= (m - 1) * ldc + n);
5080    // Test hook: route the attention GEMMs through the portable NEON
5081    // micro-kernel ON APPLE SILICON — how the mobile batched attend is
5082    // measured without a phone in the loop. (Intel macOS has no NEON —
5083    // the hook is a no-op there, Accelerate continues below.)
5084    #[cfg(target_arch = "aarch64")]
5085    if std::env::var("CMF_FORCE_NEON_GEMM")
5086        .map(|v| v == "1")
5087        .unwrap_or(false)
5088    {
5089        return neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
5090    }
5091    unsafe {
5092        accel_blas::cblas_sgemm(
5093            101, // RowMajor
5094            111, // NoTrans A
5095            if b_rows_are_n { 112 } else { 111 },
5096            m as i32,
5097            n as i32,
5098            k as i32,
5099            alpha,
5100            a.as_ptr(),
5101            lda as i32,
5102            b_mat.as_ptr(),
5103            ldb as i32,
5104            0.0,
5105            c.as_mut_ptr(),
5106            ldc as i32,
5107        );
5108    }
5109}
5110
5111/// Prefill GEMM through Accelerate (macOS): dequantize q8 rows into
5112/// f32 tiles (scale folded in, pool-parallel) and multiply each tile
5113/// on the AMX with one row-major sgemm. Tiles live in cache, weights
5114/// stream once. Numerics are f32-GEMM (not the int8 dot): prefill
5115/// logits shift within f32 rounding — tolerance-class, like every
5116/// reduction-order change; decode (M=1) never takes this path.
5117#[cfg(target_os = "macos")]
5118fn qmatmat_accel(
5119    q: &[u8],
5120    row_scale: &[f32],
5121    pre: &[std::borrow::Cow<'_, [f32]>],
5122    rows: usize,
5123    cols: usize,
5124    out: &mut [f32],
5125    pool: Option<&Pool>,
5126) {
5127    // NOTE: double-buffering the dequant against the sgemm (a scoped
5128    // thread driving the pool on tile k+1 while the caller multiplies
5129    // tile k) was tried and LOST ~6%: Accelerate's sgemm is itself
5130    // multithreaded, and the dequant workers just steal its cores.
5131    const TR: usize = 2048;
5132    let b = pre.len();
5133    thread_local! {
5134        static XPANEL: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
5135        static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
5136    }
5137    XPANEL.with(|xp| {
5138        WTILE.with(|wt| {
5139            let mut xpanel = xp.borrow_mut();
5140            xpanel.clear();
5141            for x in pre {
5142                xpanel.extend_from_slice(x);
5143            }
5144            let mut wtile = wt.borrow_mut();
5145            wtile.resize(TR * cols, 0.0);
5146            let mut r0 = 0usize;
5147            while r0 < rows {
5148                let tr = TR.min(rows - r0);
5149                // Dequant the tile (scale folded) — pool-parallel.
5150                let wt_addr = SendMut(wtile.as_mut_ptr());
5151                let run = |start: usize, end: usize| {
5152                    for r in start..end {
5153                        let row = &q[(r0 + r) * cols..(r0 + r + 1) * cols];
5154                        let s = row_scale[r0 + r];
5155                        // SAFETY: workers cover disjoint r ranges.
5156                        let dst =
5157                            unsafe { std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols) };
5158                        for (d, &v) in dst.iter_mut().zip(row) {
5159                            *d = (v as i8) as f32 * s;
5160                        }
5161                    }
5162                };
5163                dispatch_rows(pool, tr, &run);
5164                // C[b, tr] (at column r0 of out[b, rows]) = X · Wtileᵀ
5165                unsafe {
5166                    accel_blas::cblas_sgemm(
5167                        101, // RowMajor
5168                        111, // NoTrans A
5169                        112, // Trans B
5170                        b as i32,
5171                        tr as i32,
5172                        cols as i32,
5173                        1.0,
5174                        xpanel.as_ptr(),
5175                        cols as i32,
5176                        wtile.as_ptr(),
5177                        cols as i32,
5178                        0.0,
5179                        out.as_mut_ptr().add(r0),
5180                        rows as i32,
5181                    );
5182                }
5183                r0 += tr;
5184            }
5185        })
5186    });
5187}
5188
5189fn qmatmat(
5190    q: &[u8],
5191    row_scale: &[f32],
5192    pre: &[std::borrow::Cow<'_, [f32]>],
5193    rows: usize,
5194    cols: usize,
5195    out: &mut [f32],
5196    pool: Option<&Pool>,
5197) {
5198    let b = pre.len();
5199    debug_assert_eq!(out.len(), b * rows);
5200    let _p = crate::cpuprof::spans::span("mm.qmatmat");
5201    // Big prefill batches ride the AMX (roadmap PR3): the row×batch
5202    // SDOT loop below peaks near the CPU's dot throughput, an order
5203    // below the matrix units. Small tensors and tiny test models stay
5204    // on the exact integer path.
5205    #[cfg(target_os = "macos")]
5206    if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
5207        qmatmat_accel(q, row_scale, pre, rows, cols, out, pool);
5208        return;
5209    }
5210    #[cfg(target_arch = "aarch64")]
5211    if sdot_enabled() {
5212        let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
5213        let out_addr = SendMut(out.as_mut_ptr());
5214        // Blocked 2×4 (mobile prefill: no AMX to fall back on — this
5215        // path IS the ARM prefill GEMM off Apple silicon).
5216        let blocked_ok = blocked_enabled();
5217        let use_i8mm = i8mm_enabled();
5218        if blocked_ok {
5219            let run = |start: usize, end: usize| {
5220                let mut o = start;
5221                while o < end {
5222                    if o + 2 <= end {
5223                        let r0 = &q[o * cols..(o + 1) * cols];
5224                        let r1 = &q[(o + 1) * cols..(o + 2) * cols];
5225                        let mut bi = 0usize;
5226                        while bi + 4 <= acts.len() {
5227                            let xs = [
5228                                acts[bi].xq.as_slice(),
5229                                acts[bi + 1].xq.as_slice(),
5230                                acts[bi + 2].xq.as_slice(),
5231                                acts[bi + 3].xq.as_slice(),
5232                            ];
5233                            let d = if use_i8mm {
5234                                unsafe { dot_i8_smmla_2x4(r0, r1, xs) }
5235                            } else {
5236                                unsafe { dot_i8_sdot_2x4(r0, r1, xs) }
5237                            };
5238                            for (r, row) in [r0, r1].into_iter().enumerate() {
5239                                for k in 0..4 {
5240                                    let act = &acts[bi + k];
5241                                    let mut v = d[r][k] as f32 * act.sx;
5242                                    for &(j, xv) in &act.outliers {
5243                                        v += (row[j] as i8) as f32 * xv;
5244                                    }
5245                                    unsafe {
5246                                        *out_addr.at((bi + k) * rows + o + r) = v * row_scale[o + r]
5247                                    };
5248                                }
5249                            }
5250                            bi += 4;
5251                        }
5252                        while bi < acts.len() {
5253                            for (r, row) in [r0, r1].into_iter().enumerate() {
5254                                let v = row_dot_sdot(row, &acts[bi]) * row_scale[o + r];
5255                                unsafe { *out_addr.at(bi * rows + o + r) = v };
5256                            }
5257                            bi += 1;
5258                        }
5259                        o += 2;
5260                    } else {
5261                        let row = &q[o * cols..(o + 1) * cols];
5262                        for (bi, act) in acts.iter().enumerate() {
5263                            let v = row_dot_sdot(row, act) * row_scale[o];
5264                            unsafe { *out_addr.at(bi * rows + o) = v };
5265                        }
5266                        o += 1;
5267                    }
5268                }
5269            };
5270            dispatch_rows(pool, rows, &run);
5271            return;
5272        }
5273        let run = |start: usize, end: usize| {
5274            for o in start..end {
5275                let row = &q[o * cols..(o + 1) * cols];
5276                for (bi, act) in acts.iter().enumerate() {
5277                    let v = row_dot_sdot(row, act) * row_scale[o];
5278                    unsafe { *out_addr.at(bi * rows + o) = v };
5279                }
5280            }
5281        };
5282        dispatch_rows(pool, rows, &run);
5283        return;
5284    }
5285    // x86 A8W8 batch. Non-VNNI parts take the BLOCKED 2×4 kernel
5286    // (roadmap P0: two weight rows' abs() stay in registers across four
5287    // activation streams); VNNI machines keep the per-row bias-trick
5288    // dot, which is already throughput-bound there.
5289    #[cfg(target_arch = "x86_64")]
5290    if avx2_a8w8_enabled() {
5291        let acts = {
5292            let _psa = crate::cpuprof::spans::span("mm.split_act");
5293            let a: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
5294            drop(_psa);
5295            a
5296        };
5297        let out_addr = SendMut(out.as_mut_ptr());
5298        // CMF_X86_BLOCKED=0 forces the per-row path (paired in-process
5299        // A/B on noisy shared-vCPU hosts).
5300        let blocked_ok = blocked_enabled();
5301        if !avx512vnni_enabled() && blocked_ok && !row_exact() {
5302            let run = |start: usize, end: usize| {
5303                let mut o = start;
5304                while o < end {
5305                    if o + 2 <= end {
5306                        let r0 = &q[o * cols..(o + 1) * cols];
5307                        let r1 = &q[(o + 1) * cols..(o + 2) * cols];
5308                        let mut bi = 0usize;
5309                        while bi + 4 <= acts.len() {
5310                            let xs = [
5311                                acts[bi].xq.as_slice(),
5312                                acts[bi + 1].xq.as_slice(),
5313                                acts[bi + 2].xq.as_slice(),
5314                                acts[bi + 3].xq.as_slice(),
5315                            ];
5316                            let d = unsafe { dot_i8_i8_avx2_2x4(r0, r1, xs) };
5317                            for (r, row) in [r0, r1].into_iter().enumerate() {
5318                                for k in 0..4 {
5319                                    let act = &acts[bi + k];
5320                                    let mut v = d[r][k] as f32 * act.sx;
5321                                    for &(j, xv) in &act.outliers {
5322                                        v += (row[j] as i8) as f32 * xv;
5323                                    }
5324                                    unsafe {
5325                                        *out_addr.at((bi + k) * rows + o + r) = v * row_scale[o + r]
5326                                    };
5327                                }
5328                            }
5329                            bi += 4;
5330                        }
5331                        while bi < acts.len() {
5332                            for (r, row) in [r0, r1].into_iter().enumerate() {
5333                                let v = row_dot_avx2(row, &acts[bi]) * row_scale[o + r];
5334                                unsafe { *out_addr.at(bi * rows + o + r) = v };
5335                            }
5336                            bi += 1;
5337                        }
5338                        o += 2;
5339                    } else {
5340                        let row = &q[o * cols..(o + 1) * cols];
5341                        for (bi, act) in acts.iter().enumerate() {
5342                            let v = row_dot_avx2(row, act) * row_scale[o];
5343                            unsafe { *out_addr.at(bi * rows + o) = v };
5344                        }
5345                        o += 1;
5346                    }
5347                }
5348            };
5349            dispatch_rows(pool, rows, &run);
5350            return;
5351        }
5352        let run = |start: usize, end: usize| {
5353            for o in start..end {
5354                let row = &q[o * cols..(o + 1) * cols];
5355                for (bi, act) in acts.iter().enumerate() {
5356                    let v = row_dot_avx2(row, act) * row_scale[o];
5357                    unsafe { *out_addr.at(bi * rows + o) = v };
5358                }
5359            }
5360        };
5361        dispatch_rows(pool, rows, &run);
5362        return;
5363    }
5364    let out_addr = SendMut(out.as_mut_ptr());
5365    let run = |start: usize, end: usize| {
5366        for o in start..end {
5367            let row = &q[o * cols..(o + 1) * cols];
5368            for (bi, x) in pre.iter().enumerate() {
5369                let mut acc = 0f32;
5370                for j in 0..cols {
5371                    acc += (row[j] as i8) as f32 * x[j];
5372                }
5373                unsafe { *out_addr.at(bi * rows + o) = acc * row_scale[o] };
5374            }
5375        }
5376    };
5377    dispatch_rows(pool, rows, &run);
5378}
5379
5380/// Split rows across pool workers (shared qmatvec pattern). Self-balancing
5381/// — see `Pool::run_rows` for why a static 1/n split is wrong here.
5382fn dispatch_rows(pool: Option<&Pool>, rows: usize, run: &(dyn Fn(usize, usize) + Sync)) {
5383    match pool {
5384        Some(pool) if rows >= 256 => pool.run_rows(rows, run),
5385        _ => run(0, rows),
5386    }
5387}
5388
5389/// Split a q4_block blob into (packed nibbles, f16 group scales).
5390fn q4_split(bytes: &[u8], rows: usize, cols: usize) -> (&[u8], &[u8]) {
5391    let groups = rows * cols / GROUP_SIZE;
5392    bytes.split_at(groups * 16)
5393}
5394
5395/// SIMD unpack for the dominant vbit width B=4 (94% of rows on the
5396/// log2-shape calibration): 16 packed bytes -> 32 centered i8 values.
5397/// vbit packs MSB-first, so the HIGH nibble is the even element
5398/// (opposite of q4_block's lo-first interleave). Centering is u-7.
5399#[inline]
5400fn vbit_fill4(data: &[u8], buf: &mut [u8]) {
5401    #[cfg(target_arch = "aarch64")]
5402    unsafe {
5403        return vbit_fill4_neon(data, buf);
5404    }
5405    #[cfg(target_arch = "x86_64")]
5406    if avx2_enabled() {
5407        return unsafe { vbit_fill4_avx2(data, buf) };
5408    }
5409    #[allow(unreachable_code)]
5410    for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
5411        let u = unpack8::<4>(&data[blk * 4..]);
5412        for k in 0..8 {
5413            chunk[k] = (u[k] - 7) as i8 as u8;
5414        }
5415    }
5416}
5417
5418#[cfg(target_arch = "aarch64")]
5419#[target_feature(enable = "neon")]
5420unsafe fn vbit_fill4_neon(data: &[u8], buf: &mut [u8]) {
5421    // SAFETY: buf.len() is a multiple of GROUP_SIZE=32; data holds
5422    // buf.len()/2 packed bytes (validated at load).
5423    unsafe {
5424        use core::arch::aarch64::*;
5425        let n = buf.len();
5426        let mask = vdupq_n_u8(0x0F);
5427        let seven = vdupq_n_s8(7);
5428        let mut g = 0usize;
5429        while g * 32 + 32 <= n {
5430            let b = vld1q_u8(data.as_ptr().add(g * 16));
5431            let hi = vshrq_n_u8::<4>(b);
5432            let lo = vandq_u8(b, mask);
5433            let z0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(hi, lo)), seven);
5434            let z1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(hi, lo)), seven);
5435            vst1q_u8(buf.as_mut_ptr().add(g * 32), vreinterpretq_u8_s8(z0));
5436            vst1q_u8(buf.as_mut_ptr().add(g * 32 + 16), vreinterpretq_u8_s8(z1));
5437            g += 1;
5438        }
5439    }
5440}
5441
5442#[cfg(target_arch = "x86_64")]
5443#[target_feature(enable = "avx2")]
5444unsafe fn vbit_fill4_avx2(data: &[u8], buf: &mut [u8]) {
5445    // SAFETY: see vbit_fill4_neon.
5446    unsafe {
5447        use core::arch::x86_64::*;
5448        let n = buf.len();
5449        let mask = _mm_set1_epi8(0x0F);
5450        let seven = _mm256_set1_epi8(7);
5451        let mut g = 0usize;
5452        while g * 32 + 32 <= n {
5453            let b = _mm_loadu_si128(data.as_ptr().add(g * 16) as *const __m128i);
5454            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), mask);
5455            let lo = _mm_and_si128(b, mask);
5456            let z = _mm256_sub_epi8(
5457                _mm256_set_m128i(_mm_unpackhi_epi8(hi, lo), _mm_unpacklo_epi8(hi, lo)),
5458                seven,
5459            );
5460            _mm256_storeu_si256(buf.as_mut_ptr().add(g * 32) as *mut __m256i, z);
5461            g += 1;
5462        }
5463    }
5464}
5465
5466/// Unpack 8 MSB-first B-bit values from exactly B bytes (fixed shifts —
5467/// no serial bit-buffer, auto-vectorizable). Every 32-value group starts
5468/// byte-aligned (32·B/8 is integral for B∈3..8), so groups decompose
5469/// into 4 such blocks.
5470#[inline(always)]
5471fn unpack8<const B: usize>(data: &[u8]) -> [i32; 8] {
5472    let mut acc = 0u64;
5473    for i in 0..B {
5474        acc = (acc << 8) | data[i] as u64;
5475    }
5476    let mask = (1u64 << B) - 1;
5477    let mut out = [0i32; 8];
5478    for (k, o) in out.iter_mut().enumerate() {
5479        *o = ((acc >> ((7 - k) * B)) & mask) as i32;
5480    }
5481    out
5482}
5483
5484/// Fused vbit matvec straight from the mapped bytes (spec §3, P13
5485/// FIG.3): [u8 bits: rows][f16 scales: rows·cols/32][bit-packed rows,
5486/// MSB-first, byte-padded]. Row data offsets are precomputed at load
5487/// (`vbit_row_offsets`) — the per-call prefix scan was O(rows) pure
5488/// overhead on every matvec.
5489#[allow(clippy::too_many_arguments)]
5490fn vbitmatvec(
5491    bytes: &[u8],
5492    offsets: &[usize],
5493    x: &[f32],
5494    rows: usize,
5495    cols: usize,
5496    out: &mut [f32],
5497    pool: Option<&Pool>,
5498) {
5499    debug_assert_eq!(out.len(), rows);
5500    debug_assert_eq!(offsets.len(), rows + 1);
5501
5502    // SDOT path: unpack the row to centered i8 once, then per-group
5503    // int8 dot against the quantized activations — same A8W8 contract
5504    // as q8 (bounded noise; CMF_SDOT=0 keeps the exact scalar path).
5505    if a8w8_enabled() {
5506        let act = split_act(x);
5507        let out_addr = SendMut(out.as_mut_ptr());
5508        let run = move |start: usize, end: usize| {
5509            vbit_range_a8w8(bytes, offsets, x, &act, rows, cols, out_addr, start, end)
5510        };
5511        dispatch_rows(pool, rows, &run);
5512        return;
5513    }
5514
5515    let out_addr = SendMut(out.as_mut_ptr());
5516    let run = move |start: usize, end: usize| {
5517        vbit_range_f32(bytes, offsets, x, rows, cols, out_addr, start, end)
5518    };
5519    dispatch_rows(pool, rows, &run);
5520}
5521
5522/// One vbit row range via the A8W8 int8 path — kernel body of
5523/// `vbitmatvec`, extracted so multi-matrix jobs can drive it for
5524/// several tensors in one dispatch (b=8 rows go exact f32).
5525#[allow(clippy::too_many_arguments)]
5526fn vbit_range_a8w8(
5527    bytes: &[u8],
5528    offsets: &[usize],
5529    x: &[f32],
5530    act: &SplitAct,
5531    rows: usize,
5532    cols: usize,
5533    out: SendMut,
5534    start: usize,
5535    end: usize,
5536) {
5537    let ng = cols / GROUP_SIZE;
5538    let bits = &bytes[..rows];
5539    let sc_off = rows;
5540    let row_dot = |r: usize| -> f32 {
5541        let b = bits[r] as usize;
5542        let l = (1i32 << (b - 1)) - 1;
5543        let mask = (1u64 << b) - 1;
5544        let data = &bytes[offsets[r]..offsets[r + 1]];
5545        if b == 8 {
5546            // u−L reaches 128 → does not fit i8; exact f32 path.
5547            let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
5548            let mut dot = 0f32;
5549            for g in 0..ng {
5550                let so = (r * ng + g) * 2;
5551                let sgf = f16_to_f32(u16::from_le_bytes([
5552                    bytes[sc_off + so],
5553                    bytes[sc_off + so + 1],
5554                ]));
5555                let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
5556                let mut gd = 0f32;
5557                for &xv in xg.iter() {
5558                    if nbits < 8 {
5559                        acc = (acc << 8) | data[idx] as u64;
5560                        idx += 1;
5561                        nbits += 8;
5562                    }
5563                    let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
5564                    nbits -= 8;
5565                    gd += (u - l) as f32 * xv;
5566                }
5567                dot += gd * sgf;
5568            }
5569            return dot;
5570        }
5571        // Per-worker scratch: this closure runs for every row of the
5572        // tensor (lm_head ≈ 150k rows/token) — a heap allocation per
5573        // row was measurable pure overhead.
5574        thread_local! {
5575            static VBIT_SCRATCH: std::cell::RefCell<Vec<u8>> =
5576                const { std::cell::RefCell::new(Vec::new()) };
5577        }
5578        #[inline(always)]
5579        fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
5580            for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
5581                let u = unpack8::<B>(&data[blk * B..]);
5582                for k in 0..8 {
5583                    chunk[k] = (u[k] - l) as i8 as u8;
5584                }
5585            }
5586        }
5587        let _ = mask;
5588        VBIT_SCRATCH.with(|scratch| {
5589            let mut buf = scratch.borrow_mut();
5590            buf.resize(cols, 0);
5591            match b {
5592                3 => fill::<3>(data, l, &mut buf),
5593                4 => vbit_fill4(data, &mut buf),
5594                5 => fill::<5>(data, l, &mut buf),
5595                6 => fill::<6>(data, l, &mut buf),
5596                _ => unreachable!(),
5597            }
5598            let mut dot = 0f32;
5599            for g in 0..ng {
5600                let so = (r * ng + g) * 2;
5601                let s = f16_to_f32(u16::from_le_bytes([
5602                    bytes[sc_off + so],
5603                    bytes[sc_off + so + 1],
5604                ]));
5605                let d = dot_i8_i8(
5606                    &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
5607                    &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
5608                ) as f32
5609                    * act.sx;
5610                dot += d * s;
5611            }
5612            for &(j, xv) in &act.outliers {
5613                let so = (r * ng + j / GROUP_SIZE) * 2;
5614                let s = f16_to_f32(u16::from_le_bytes([
5615                    bytes[sc_off + so],
5616                    bytes[sc_off + so + 1],
5617                ]));
5618                // xq is zeroed at outlier slots — add the exact term.
5619                dot += (buf[j] as i8) as f32 * s * xv;
5620            }
5621            dot
5622        })
5623    };
5624    for r in start..end {
5625        // SAFETY: disjoint row ranges per worker.
5626        unsafe { *out.at(r) = row_dot(r) };
5627    }
5628}
5629
5630/// Exact scalar vbit row range (same extraction, non-SDOT path).
5631#[allow(clippy::too_many_arguments)]
5632fn vbit_range_f32(
5633    bytes: &[u8],
5634    offsets: &[usize],
5635    x: &[f32],
5636    rows: usize,
5637    cols: usize,
5638    out: SendMut,
5639    start: usize,
5640    end: usize,
5641) {
5642    let ng = cols / GROUP_SIZE;
5643    let bits = &bytes[..rows];
5644    let sc_off = rows;
5645    // Per-bit-width specialized inner loops: the compiler unrolls the
5646    // constant shifts (the generic bit-buffer loop was branch-bound —
5647    // 5.6 vs 13.2 tok/s q4 on the 0.8B).
5648    #[inline(always)]
5649    fn dot_row<const B: usize>(
5650        data: &[u8],
5651        bytes: &[u8],
5652        sc_off: usize,
5653        r: usize,
5654        ng: usize,
5655        x: &[f32],
5656    ) -> f32 {
5657        let l = ((1i32 << (B - 1)) - 1) as f32;
5658        let gbytes = GROUP_SIZE * B / 8;
5659        let mut dot = 0f32;
5660        for g in 0..ng {
5661            let so = (r * ng + g) * 2;
5662            let s = f16_to_f32(u16::from_le_bytes([
5663                bytes[sc_off + so],
5664                bytes[sc_off + so + 1],
5665            ]));
5666            let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
5667            let gd0 = &data[g * gbytes..(g + 1) * gbytes];
5668            let mut gd = 0f32;
5669            for blk in 0..GROUP_SIZE / 8 {
5670                let u = unpack8::<B>(&gd0[blk * B..]);
5671                let xb = &xg[blk * 8..blk * 8 + 8];
5672                for k in 0..8 {
5673                    gd += (u[k] as f32 - l) * xb[k];
5674                }
5675            }
5676            dot += gd * s;
5677        }
5678        dot
5679    }
5680    for r in start..end {
5681        let data = &bytes[offsets[r]..offsets[r + 1]];
5682        let v = match bits[r] {
5683            3 => dot_row::<3>(data, bytes, sc_off, r, ng, x),
5684            4 => dot_row::<4>(data, bytes, sc_off, r, ng, x),
5685            5 => dot_row::<5>(data, bytes, sc_off, r, ng, x),
5686            6 => dot_row::<6>(data, bytes, sc_off, r, ng, x),
5687            8 => dot_row::<8>(data, bytes, sc_off, r, ng, x),
5688            b => unreachable!("vbit bit-width {b} (validated at load)"),
5689        };
5690        // SAFETY: disjoint row ranges per worker.
5691        unsafe { *out.at(r) = v };
5692    }
5693}
5694
5695/// Fused two-input vbit matvec: each row is unpacked from the mmap ONCE
5696/// and dotted against BOTH activations (MTP verify / pair prefill used
5697/// to run two full matvecs — double weight traffic and double unpack).
5698/// Per-input math is identical to `vbitmatvec` → same accuracy contract.
5699#[allow(clippy::too_many_arguments)]
5700fn vbitmatvec2(
5701    bytes: &[u8],
5702    offsets: &[usize],
5703    x1: &[f32],
5704    x2: &[f32],
5705    rows: usize,
5706    cols: usize,
5707    o1: &mut [f32],
5708    o2: &mut [f32],
5709    pool: Option<&Pool>,
5710) {
5711    debug_assert_eq!(o1.len(), rows);
5712    debug_assert_eq!(o2.len(), rows);
5713
5714    if a8w8_enabled() {
5715        let a1 = split_act(x1);
5716        let a2 = split_act(x2);
5717        let p1 = SendMut(o1.as_mut_ptr());
5718        let p2 = SendMut(o2.as_mut_ptr());
5719        let run = move |start: usize, end: usize| {
5720            vbit_range2_a8w8(
5721                bytes, offsets, x1, x2, &a1, &a2, rows, cols, p1, p2, start, end,
5722            )
5723        };
5724        dispatch_rows(pool, rows, &run);
5725        return;
5726    }
5727
5728    let p1 = SendMut(o1.as_mut_ptr());
5729    let p2 = SendMut(o2.as_mut_ptr());
5730    let run = move |start: usize, end: usize| {
5731        vbit_range2_f32(bytes, offsets, x1, x2, rows, cols, p1, p2, start, end)
5732    };
5733    dispatch_rows(pool, rows, &run);
5734}
5735
5736/// Two-input vbit row range via the A8W8 int8 path — kernel body of
5737/// `vbitmatvec2`, extracted for pair multi-matrix jobs (b=8 rows go
5738/// exact f32 for both lanes, bits streamed once).
5739#[allow(clippy::too_many_arguments)]
5740fn vbit_range2_a8w8(
5741    bytes: &[u8],
5742    offsets: &[usize],
5743    x1: &[f32],
5744    x2: &[f32],
5745    a1: &SplitAct,
5746    a2: &SplitAct,
5747    rows: usize,
5748    cols: usize,
5749    p1: SendMut,
5750    p2: SendMut,
5751    start: usize,
5752    end: usize,
5753) {
5754    let ng = cols / GROUP_SIZE;
5755    let bits = &bytes[..rows];
5756    let sc_off = rows;
5757    let row_dots = |r: usize| -> (f32, f32) {
5758        let b = bits[r] as usize;
5759        let l = (1i32 << (b - 1)) - 1;
5760        let data = &bytes[offsets[r]..offsets[r + 1]];
5761        if b == 8 {
5762            // u−L reaches 128 → does not fit i8; exact f32 path,
5763            // bits still streamed once for both lanes.
5764            let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
5765            let (mut d1, mut d2) = (0f32, 0f32);
5766            for g in 0..ng {
5767                let so = (r * ng + g) * 2;
5768                let sgf = f16_to_f32(u16::from_le_bytes([
5769                    bytes[sc_off + so],
5770                    bytes[sc_off + so + 1],
5771                ]));
5772                let (mut g1, mut g2) = (0f32, 0f32);
5773                for k in 0..GROUP_SIZE {
5774                    if nbits < 8 {
5775                        acc = (acc << 8) | data[idx] as u64;
5776                        idx += 1;
5777                        nbits += 8;
5778                    }
5779                    let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
5780                    nbits -= 8;
5781                    let w = (u - l) as f32;
5782                    g1 += w * x1[g * GROUP_SIZE + k];
5783                    g2 += w * x2[g * GROUP_SIZE + k];
5784                }
5785                d1 += g1 * sgf;
5786                d2 += g2 * sgf;
5787            }
5788            return (d1, d2);
5789        }
5790        thread_local! {
5791            static VBIT_SCRATCH2: std::cell::RefCell<Vec<u8>> =
5792                const { std::cell::RefCell::new(Vec::new()) };
5793        }
5794        #[inline(always)]
5795        fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
5796            for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
5797                let u = unpack8::<B>(&data[blk * B..]);
5798                for k in 0..8 {
5799                    chunk[k] = (u[k] - l) as i8 as u8;
5800                }
5801            }
5802        }
5803        VBIT_SCRATCH2.with(|scratch| {
5804            let mut buf = scratch.borrow_mut();
5805            buf.resize(cols, 0);
5806            match b {
5807                3 => fill::<3>(data, l, &mut buf),
5808                4 => vbit_fill4(data, &mut buf),
5809                5 => fill::<5>(data, l, &mut buf),
5810                6 => fill::<6>(data, l, &mut buf),
5811                _ => unreachable!(),
5812            }
5813            let (mut d1, mut d2) = (0f32, 0f32);
5814            for g in 0..ng {
5815                let so = (r * ng + g) * 2;
5816                let s = f16_to_f32(u16::from_le_bytes([
5817                    bytes[sc_off + so],
5818                    bytes[sc_off + so + 1],
5819                ]));
5820                let wg = &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
5821                let v1 = dot_i8_i8(wg, &a1.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a1.sx;
5822                let v2 = dot_i8_i8(wg, &a2.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a2.sx;
5823                d1 += v1 * s;
5824                d2 += v2 * s;
5825            }
5826            for &(j, xv) in &a1.outliers {
5827                let so = (r * ng + j / GROUP_SIZE) * 2;
5828                let s = f16_to_f32(u16::from_le_bytes([
5829                    bytes[sc_off + so],
5830                    bytes[sc_off + so + 1],
5831                ]));
5832                d1 += (buf[j] as i8) as f32 * s * xv;
5833            }
5834            for &(j, xv) in &a2.outliers {
5835                let so = (r * ng + j / GROUP_SIZE) * 2;
5836                let s = f16_to_f32(u16::from_le_bytes([
5837                    bytes[sc_off + so],
5838                    bytes[sc_off + so + 1],
5839                ]));
5840                d2 += (buf[j] as i8) as f32 * s * xv;
5841            }
5842            (d1, d2)
5843        })
5844    };
5845    for r in start..end {
5846        let (v1, v2) = row_dots(r);
5847        // SAFETY: disjoint row ranges per worker.
5848        unsafe {
5849            *p1.at(r) = v1;
5850            *p2.at(r) = v2;
5851        }
5852    }
5853}
5854
5855/// Two-input exact scalar vbit row range (same extraction) —
5856/// per-bit-width specialized, two accumulators per row; per-lane
5857/// accumulation order matches `vbitmatvec` exactly.
5858#[allow(clippy::too_many_arguments)]
5859fn vbit_range2_f32(
5860    bytes: &[u8],
5861    offsets: &[usize],
5862    x1: &[f32],
5863    x2: &[f32],
5864    rows: usize,
5865    cols: usize,
5866    p1: SendMut,
5867    p2: SendMut,
5868    start: usize,
5869    end: usize,
5870) {
5871    let ng = cols / GROUP_SIZE;
5872    let bits = &bytes[..rows];
5873    let sc_off = rows;
5874    #[inline(always)]
5875    #[allow(clippy::too_many_arguments)]
5876    fn dot_row2<const B: usize>(
5877        data: &[u8],
5878        bytes: &[u8],
5879        sc_off: usize,
5880        r: usize,
5881        ng: usize,
5882        x1: &[f32],
5883        x2: &[f32],
5884    ) -> (f32, f32) {
5885        let l = ((1i32 << (B - 1)) - 1) as f32;
5886        let gbytes = GROUP_SIZE * B / 8;
5887        let (mut d1, mut d2) = (0f32, 0f32);
5888        for g in 0..ng {
5889            let so = (r * ng + g) * 2;
5890            let s = f16_to_f32(u16::from_le_bytes([
5891                bytes[sc_off + so],
5892                bytes[sc_off + so + 1],
5893            ]));
5894            let x1g = &x1[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
5895            let x2g = &x2[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
5896            let gd0 = &data[g * gbytes..(g + 1) * gbytes];
5897            let (mut g1, mut g2) = (0f32, 0f32);
5898            for blk in 0..GROUP_SIZE / 8 {
5899                let u = unpack8::<B>(&gd0[blk * B..]);
5900                for k in 0..8 {
5901                    let w = u[k] as f32 - l;
5902                    g1 += w * x1g[blk * 8 + k];
5903                    g2 += w * x2g[blk * 8 + k];
5904                }
5905            }
5906            d1 += g1 * s;
5907            d2 += g2 * s;
5908        }
5909        (d1, d2)
5910    }
5911    for r in start..end {
5912        let data = &bytes[offsets[r]..offsets[r + 1]];
5913        let (v1, v2) = match bits[r] {
5914            3 => dot_row2::<3>(data, bytes, sc_off, r, ng, x1, x2),
5915            4 => dot_row2::<4>(data, bytes, sc_off, r, ng, x1, x2),
5916            5 => dot_row2::<5>(data, bytes, sc_off, r, ng, x1, x2),
5917            6 => dot_row2::<6>(data, bytes, sc_off, r, ng, x1, x2),
5918            8 => dot_row2::<8>(data, bytes, sc_off, r, ng, x1, x2),
5919            b => unreachable!("vbit bit-width {b} (validated at load)"),
5920        };
5921        // SAFETY: disjoint row ranges per worker.
5922        unsafe {
5923            *p1.at(r) = v1;
5924            *p2.at(r) = v2;
5925        }
5926    }
5927}
5928
5929// ───────────────────── q4_tiled kernels (§4.3) ─────────────────────
5930
5931/// One q4_tiled row dot on the A8W8 int8 path: per 32-group the tile
5932/// is ONE sequential read — [f16 scale][16B nibbles] — versus the two
5933/// distant streams of the split layout. Values/order identical to the
5934/// split kernels.
5935#[inline]
5936#[allow(unreachable_code)]
5937fn dot_q4t_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
5938    #[cfg(target_arch = "aarch64")]
5939    unsafe {
5940        return dot_q4t_row_sdot(bytes, r, gpr, xq);
5941    }
5942    #[cfg(target_arch = "x86_64")]
5943    unsafe {
5944        if vnni_tiles_enabled() {
5945            return dot_q4t_row_vnni(bytes, r, gpr, xq);
5946        }
5947        return dot_q4t_row_avx2(bytes, r, gpr, xq);
5948    }
5949    let mut acc = 0f32;
5950    for gi in 0..gpr {
5951        let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
5952        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
5953        let mut d = 0i32;
5954        for (k, &b) in tile[2..].iter().enumerate() {
5955            d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
5956                + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
5957        }
5958        acc += d as f32 * s;
5959    }
5960    acc
5961}
5962
5963#[cfg(target_arch = "aarch64")]
5964#[target_feature(enable = "neon,dotprod")]
5965unsafe fn dot_q4t_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
5966    // SAFETY: callers uphold slice-length contracts (18B tile per group,
5967    // xq.len() == gpr·GROUP_SIZE).
5968    unsafe {
5969        use core::arch::aarch64::*;
5970        use core::arch::asm;
5971        let lomask = vdupq_n_u8(0x0F);
5972        let eight = vdupq_n_s8(8);
5973        let mut acc = 0f32;
5974        for gi in 0..gpr {
5975            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
5976            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5977            let b = vld1q_u8(t.add(2));
5978            let lo = vandq_u8(b, lomask);
5979            let hi = vshrq_n_u8::<4>(b);
5980            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
5981            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
5982            let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
5983            let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
5984            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5985            asm!(
5986                "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
5987                "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
5988                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5989                e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
5990                options(pure, nomem, nostack),
5991            );
5992            acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
5993        }
5994        acc
5995    }
5996}
5997
5998#[cfg(target_arch = "x86_64")]
5999#[target_feature(enable = "avx2")]
6000unsafe fn dot_q4t_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
6001    // SAFETY: see dot_q4t_row_sdot.
6002    unsafe {
6003        use core::arch::x86_64::*;
6004        let lomask = _mm_set1_epi8(0x0F);
6005        let eight = _mm256_set1_epi8(8);
6006        let ones = _mm256_set1_epi16(1);
6007        let mut acc = 0f32;
6008        for gi in 0..gpr {
6009            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
6010            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
6011            let b = _mm_loadu_si128(t.add(2) as *const __m128i);
6012            let lo = _mm_and_si128(b, lomask);
6013            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
6014            let w = _mm256_sub_epi8(
6015                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
6016                eight,
6017            );
6018            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6019            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
6020            let d = _mm256_madd_epi16(p16, ones);
6021            let hi128 = _mm256_extracti128_si256::<1>(d);
6022            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
6023            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
6024            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
6025            acc += _mm_cvtsi128_si32(s32) as f32 * s;
6026        }
6027        acc
6028    }
6029}
6030
6031/// VNNI twin of `dot_q4t_row_avx2`: same unpack, `vpdpbusd` replaces
6032/// the maddubs+madd pair (see `dpbusd_hsum` — sums are bit-identical).
6033/// 256-bit VL encoding, so the VEX `vpsignb` stays usable.
6034#[cfg(target_arch = "x86_64")]
6035#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
6036unsafe fn dot_q4t_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
6037    // SAFETY: see dot_q4t_row_sdot.
6038    unsafe {
6039        use core::arch::x86_64::*;
6040        let lomask = _mm_set1_epi8(0x0F);
6041        let eight = _mm256_set1_epi8(8);
6042        let mut acc = 0f32;
6043        for gi in 0..gpr {
6044            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
6045            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
6046            let b = _mm_loadu_si128(t.add(2) as *const __m128i);
6047            let lo = _mm_and_si128(b, lomask);
6048            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
6049            let w = _mm256_sub_epi8(
6050                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
6051                eight,
6052            );
6053            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6054            let d = dpbusd_hsum(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
6055            acc += d as f32 * s;
6056        }
6057        acc
6058    }
6059}
6060
6061/// One q4_tiled row against FOUR activation streams: the nibble unpack
6062/// and abs() happen once per group instead of once per (group,
6063/// activation) — the unpack is the dominant per-element cost of the
6064/// tiled format (roadmap P0 portable blocking, q4t leg).
6065#[cfg(target_arch = "x86_64")]
6066// `fma` is NOT implied by `avx2`: without it LLVM lowers _mm256_fmadd_ps
6067// to a libm call per lane — measured 2x slower than the reduction this
6068// kernel replaces. The runtime gate (`avx2_enabled`) already requires
6069// both features, so declaring it here is safe.
6070#[target_feature(enable = "avx2,fma")]
6071unsafe fn dot_q4t_row_1x4_avx2(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
6072    // SAFETY: callers uphold the 18B-tile and xq-length contracts.
6073    unsafe {
6074        use core::arch::x86_64::*;
6075        let lomask = _mm_set1_epi8(0x0F);
6076        let eight = _mm256_set1_epi8(8);
6077        let ones = _mm256_set1_epi16(1);
6078        // One f32 accumulator VECTOR per activation, reduced once at the
6079        // end. Folding each group's i32 lanes to a scalar inside the loop
6080        // costs an extracti128 + three shift/add + a movd — a cross-lane
6081        // dependency chain per (group, activation), 288 of them per row at
6082        // cols=2304. The per-group scale is what forces a float
6083        // accumulator; it does not force a horizontal sum.
6084        //
6085        // The four accumulators are NAMED, not an array: as `[__m256; 4]`
6086        // indexed by a loop variable LLVM keeps them in memory and every
6087        // group pays four 32-byte loads and stores. That alone made this
6088        // kernel 2x SLOWER than the per-group reduction it replaces
6089        // (measured on the EPYC box: 150 s vs 71 s for two 256² steps).
6090        let mut f0 = _mm256_setzero_ps();
6091        let mut f1 = _mm256_setzero_ps();
6092        let mut f2 = _mm256_setzero_ps();
6093        let mut f3 = _mm256_setzero_ps();
6094        for gi in 0..gpr {
6095            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
6096            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
6097            let sv = _mm256_set1_ps(s);
6098            let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
6099            let lo = _mm_and_si128(bb, lomask);
6100            let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
6101            let w = _mm256_sub_epi8(
6102                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
6103                eight,
6104            );
6105            let aw = _mm256_abs_epi8(w);
6106            let off = gi * GROUP_SIZE;
6107            let dot = |xq: &[i8]| {
6108                let x = _mm256_loadu_si256(xq.as_ptr().add(off) as *const __m256i);
6109                let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
6110                _mm256_cvtepi32_ps(_mm256_madd_epi16(p16, ones))
6111            };
6112            f0 = _mm256_fmadd_ps(dot(xs[0]), sv, f0);
6113            f1 = _mm256_fmadd_ps(dot(xs[1]), sv, f1);
6114            f2 = _mm256_fmadd_ps(dot(xs[2]), sv, f2);
6115            f3 = _mm256_fmadd_ps(dot(xs[3]), sv, f3);
6116        }
6117        [
6118            hsum256_ps(f0),
6119            hsum256_ps(f1),
6120            hsum256_ps(f2),
6121            hsum256_ps(f3),
6122        ]
6123    }
6124}
6125
6126/// Horizontal sum of eight f32 lanes — the one cross-lane reduction the
6127/// blocked kernels pay, once per row instead of once per group.
6128#[cfg(target_arch = "x86_64")]
6129#[target_feature(enable = "avx2")]
6130#[inline]
6131unsafe fn hsum256_ps(v: core::arch::x86_64::__m256) -> f32 {
6132    // SAFETY: pure register arithmetic on the caller's vector.
6133    unsafe {
6134        use core::arch::x86_64::*;
6135        let hi = _mm256_extractf128_ps::<1>(v);
6136        let s = _mm_add_ps(_mm256_castps256_ps128(v), hi);
6137        let s = _mm_add_ps(s, _mm_movehl_ps(s, s));
6138        let s = _mm_add_ss(s, _mm_shuffle_ps::<0x55>(s, s));
6139        _mm_cvtss_f32(s)
6140    }
6141}
6142
6143/// VNNI twin of `dot_q4t_row_1x4_avx2` (see `dpbusd_hsum`).
6144#[cfg(target_arch = "x86_64")]
6145#[target_feature(enable = "avx2,fma,avx512f,avx512bw,avx512vl,avx512vnni")]
6146unsafe fn dot_q4t_row_1x4_vnni(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
6147    // SAFETY: callers uphold the 18B-tile and xq-length contracts.
6148    unsafe {
6149        use core::arch::x86_64::*;
6150        let lomask = _mm_set1_epi8(0x0F);
6151        let eight = _mm256_set1_epi8(8);
6152        // Same shape as the AVX2 twin: accumulate in f32 vectors and pay
6153        // one cross-lane reduction per row, not per (group, activation).
6154        let mut f0 = _mm256_setzero_ps();
6155        let mut f1 = _mm256_setzero_ps();
6156        let mut f2 = _mm256_setzero_ps();
6157        let mut f3 = _mm256_setzero_ps();
6158        for gi in 0..gpr {
6159            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
6160            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
6161            let sv = _mm256_set1_ps(s);
6162            let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
6163            let lo = _mm_and_si128(bb, lomask);
6164            let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
6165            let w = _mm256_sub_epi8(
6166                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
6167                eight,
6168            );
6169            let aw = _mm256_abs_epi8(w);
6170            let off = gi * GROUP_SIZE;
6171            let dot = |xq: &[i8]| {
6172                let x = _mm256_loadu_si256(xq.as_ptr().add(off) as *const __m256i);
6173                _mm256_cvtepi32_ps(_mm256_dpbusd_epi32(
6174                    _mm256_setzero_si256(),
6175                    aw,
6176                    _mm256_sign_epi8(x, w),
6177                ))
6178            };
6179            f0 = _mm256_fmadd_ps(dot(xs[0]), sv, f0);
6180            f1 = _mm256_fmadd_ps(dot(xs[1]), sv, f1);
6181            f2 = _mm256_fmadd_ps(dot(xs[2]), sv, f2);
6182            f3 = _mm256_fmadd_ps(dot(xs[3]), sv, f3);
6183        }
6184        let acc = [
6185            hsum256_ps(f0),
6186            hsum256_ps(f1),
6187            hsum256_ps(f2),
6188            hsum256_ps(f3),
6189        ];
6190        acc
6191    }
6192}
6193
6194/// ARM twin of `dot_q4t_row_1x4_avx2`: one nibble unpack per group
6195/// serves FOUR activation streams. Per stream the group order and f32
6196/// accumulation match `dot_q4t_row_sdot` exactly — batch == matvec
6197/// bit-for-bit.
6198#[cfg(target_arch = "aarch64")]
6199#[target_feature(enable = "neon,dotprod")]
6200unsafe fn dot_q4t_row_1x4_sdot(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
6201    // SAFETY: callers uphold the 18B-tile and xq-length contracts.
6202    unsafe {
6203        use core::arch::aarch64::*;
6204        use core::arch::asm;
6205        let lomask = vdupq_n_u8(0x0F);
6206        let eight = vdupq_n_s8(8);
6207        let mut acc = [0f32; 4];
6208        for gi in 0..gpr {
6209            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
6210            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
6211            let b = vld1q_u8(t.add(2));
6212            let lo = vandq_u8(b, lomask);
6213            let hi = vshrq_n_u8::<4>(b);
6214            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
6215            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
6216            for (k, xq) in xs.iter().enumerate() {
6217                let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
6218                let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
6219                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6220                asm!(
6221                    "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
6222                    "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
6223                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6224                    e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
6225                    options(pure, nomem, nostack),
6226                );
6227                acc[k] += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
6228            }
6229        }
6230        acc
6231    }
6232}
6233
6234/// Exact-term correction for A8W8 outliers on a tiled row.
6235#[inline]
6236fn q4t_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
6237    let gi = j / GROUP_SIZE;
6238    let k = j % GROUP_SIZE;
6239    let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
6240    let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
6241    let byte = tile[2 + k / 2];
6242    let nib = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
6243    ((nib as i32 - 8) as f32, s)
6244}
6245
6246/// Exact scalar q4_tiled row (CMF_SDOT=0 contract) — same pairwise
6247/// accumulation shape as `q4_range_f32`.
6248#[inline]
6249fn q4t_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
6250    let mut acc = 0f32;
6251    for gi in 0..gpr {
6252        let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
6253        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
6254        let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
6255        let mut ga = 0f32;
6256        for (k, &b) in tile[2..].iter().enumerate() {
6257            ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
6258                + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
6259        }
6260        acc += ga * s;
6261    }
6262    acc
6263}
6264
6265/// Split view of a `q4tp` payload. The three planes are resolved once per
6266/// matvec instead of per row — `q4tp_sections` is cheap, but doing it inside
6267/// the row loop would put a division on the hot path for nothing.
6268/// Where row `r` of a coded matrix lives inside its bytes, as up to three
6269/// (offset, length) ranges (a zero length marks an unused one). A q4t row is
6270/// one run of tiles; a q4tp row is its nibbles, its two ladder parameters
6271/// and its scale codes, in three sections of the tensor.
6272fn row_span_offsets(
6273    dtype: TensorDtype,
6274    rows: usize,
6275    cols: usize,
6276    r: usize,
6277) -> Option<[(usize, usize); 3]> {
6278    if r >= rows || cols % GROUP_SIZE != 0 {
6279        return None;
6280    }
6281    let gpr = cols / GROUP_SIZE;
6282    match dtype {
6283        TensorDtype::Q4Tiled => Some([(r * gpr * Q4_TILE, gpr * Q4_TILE), (0, 0), (0, 0)]),
6284        TensorDtype::Q4TiledP => {
6285            let (params, codes, stride) = q4tp_sections(rows, cols);
6286            Some([
6287                (r * gpr * Q4TP_NIB, gpr * Q4TP_NIB),
6288                (params + r * 4, 4),
6289                (codes + r * stride, stride),
6290            ])
6291        }
6292        _ => None,
6293    }
6294}
6295
6296struct Q4tpView<'a> {
6297    nib: &'a [u8],
6298    params: &'a [u8],
6299    codes: &'a [u8],
6300    stride: usize,
6301    /// q2tp reads the ladder with rung 0 = exact zero.
6302    zero_rung: bool,
6303}
6304
6305impl<'a> Q4tpView<'a> {
6306    fn new(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
6307        let (params_off, codes_off, stride) = q4tp_sections(rows, cols);
6308        Self {
6309            nib: &bytes[..params_off],
6310            params: &bytes[params_off..codes_off],
6311            codes: &bytes[codes_off..],
6312            stride,
6313            zero_rung: false,
6314        }
6315    }
6316
6317    /// The q2tp view: identical params/codes planes, 8 B weight chunks.
6318    fn new_q2(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
6319        let (params_off, codes_off, stride) = q2tp_sections(rows, cols);
6320        Self {
6321            nib: &bytes[..params_off],
6322            params: &bytes[params_off..codes_off],
6323            codes: &bytes[codes_off..],
6324            stride,
6325            zero_rung: true,
6326        }
6327    }
6328
6329    /// Expand row `r`'s per-tile scales into `out` (length `gpr`).
6330    ///
6331    /// Doing this once per row — rather than decoding a 5-bit code inside the
6332    /// tile loop — is what makes the format free at runtime. Random access to
6333    /// a packed 5-bit field costs a division, two bounds checks and a branch;
6334    /// the tile's actual work is two `sdot`s, so per-tile decoding dominated
6335    /// the kernel and cost 5x (measured: 1.4 vs 6.9 tok/s on Nanbeige-3B).
6336    /// Walking the plane sequentially with a bit accumulator is ~3 ops.
6337    /// Eight 5-bit codes are exactly five bytes, so a whole group of
6338    /// eight decodes from one little-endian word at fixed shifts. The
6339    /// bit-accumulator this replaces carried a data-dependent `while
6340    /// have < 5` refill whose branch sat in the innermost loop of every
6341    /// q4tp row; a decode profile put this function above the dot
6342    /// products it feeds. Same bitstream, same codes — just no branch
6343    /// and eight independent extractions.
6344    /// `avx2_enabled()` and `ladder_4` are compiled only for x86-64, so a
6345    /// plain `cfg!(target_arch = ...)` guard is NOT enough: `cfg!` is a
6346    /// runtime `false` on other targets, yet both arms still get compiled
6347    /// and the x86-only symbols do not exist there at all (aarch64 build:
6348    /// E0425). This associated fn IS a compile-time `cfg` — the whole
6349    /// branch, symbol references included, is dropped on non-x86 targets,
6350    /// and the non-x86 arm falls straight through to the scalar ladder.
6351    #[cfg(target_arch = "x86_64")]
6352    #[inline]
6353    fn ladder_via_avx2(&self) -> bool {
6354        avx2_enabled()
6355    }
6356    #[cfg(not(target_arch = "x86_64"))]
6357    #[inline]
6358    fn ladder_via_avx2(&self) -> bool {
6359        false
6360    }
6361
6362    #[inline]
6363    fn scales_into(&self, r: usize, gpr: usize, out: &mut [f32]) {
6364        let tab = if self.zero_rung {
6365            q2tp_ladder(self.params, r)
6366        } else if self.ladder_via_avx2() {
6367            // One lane of the 4-wide vector ladder (the other three are
6368            // duplicates) — every q4tp scale in the process then comes
6369            // from the SAME ladder bits, whichever call site decoded it.
6370            // `ladder_via_avx2()` is cfg-gated, so on non-x86 targets
6371            // this arm is never taken and never even compiled with the
6372            // x86-only symbols in scope.
6373            #[cfg(target_arch = "x86_64")]
6374            let tab = {
6375                let p = &self.params[r * 4..r * 4 + 4];
6376                let lo = f16_to_f32(u16::from_le_bytes([p[0], p[1]]));
6377                let st = f16_to_f32(u16::from_le_bytes([p[2], p[3]]));
6378                let mut tabc = [0f32; 128];
6379                unsafe { ladder_4(&[lo; 4], &[st; 4], &mut tabc) };
6380                let mut t = [0f32; 32];
6381                for (g, o) in t.iter_mut().enumerate() {
6382                    *o = tabc[g * 4];
6383                }
6384                t
6385            };
6386            // On a non-x86 target the arm is unreachable (the `cfg!`
6387            // above is false), so it decodes through the scalar ladder —
6388            // the very path this branch would have fallen back to.
6389            #[cfg(not(target_arch = "x86_64"))]
6390            let tab = q4tp_ladder(self.params, r);
6391            tab
6392        } else {
6393            q4tp_ladder(self.params, r)
6394        };
6395        let codes = &self.codes[r * self.stride..(r + 1) * self.stride];
6396        let out = &mut out[..gpr];
6397        let mut chunks = out.chunks_exact_mut(8);
6398        let mut ci = 0usize;
6399        for c in &mut chunks {
6400            let w = u64::from(codes[ci])
6401                | u64::from(codes[ci + 1]) << 8
6402                | u64::from(codes[ci + 2]) << 16
6403                | u64::from(codes[ci + 3]) << 24
6404                | u64::from(codes[ci + 4]) << 32;
6405            for (k, o) in c.iter_mut().enumerate() {
6406                *o = tab[((w >> (5 * k)) & 31) as usize];
6407            }
6408            ci += 5;
6409        }
6410        // Fewer than eight codes left: the shared total accessor, which
6411        // tolerates a 5-bit field whose spill byte is past the stride.
6412        let tail = &codes[ci..];
6413        for (k, o) in chunks.into_remainder().iter_mut().enumerate() {
6414            *o = tab[q4tp_code(tail, k)];
6415        }
6416    }
6417}
6418
6419#[inline]
6420fn dot_q4tp_row_i8(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
6421    #[cfg(target_arch = "aarch64")]
6422    unsafe {
6423        return dot_q4tp_row_sdot(nib, r, gpr, xq, scales);
6424    }
6425    #[cfg(target_arch = "x86_64")]
6426    unsafe {
6427        if vnni_tiles_enabled() {
6428            return dot_q4tp_row_vnni(nib, r, gpr, xq, scales);
6429        }
6430        return dot_q4tp_row_avx2(nib, r, gpr, xq, scales);
6431    }
6432    #[allow(unreachable_code)]
6433    {
6434        let mut acc = 0f32;
6435        for gi in 0..gpr {
6436            let tile = &nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
6437            let s = scales[gi];
6438            let mut d = 0i32;
6439            for (k, &b) in tile.iter().enumerate() {
6440                d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
6441                    + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
6442            }
6443            acc += d as f32 * s;
6444        }
6445        acc
6446    }
6447}
6448
6449/// q4tp twin of `dot_q4t_row_sdot`: identical nibble math, but the tile
6450/// stride is 16 B (no inline scale) and the scale is a ladder lookup.
6451#[cfg(target_arch = "aarch64")]
6452#[target_feature(enable = "neon,dotprod")]
6453unsafe fn dot_q4tp_row_sdot(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
6454    // SAFETY: callers uphold slice-length contracts (16B tile per group,
6455    // xq.len() == gpr·GROUP_SIZE, codes covering gpr 5-bit fields).
6456    unsafe {
6457        use core::arch::aarch64::*;
6458        use core::arch::asm;
6459        let lomask = vdupq_n_u8(0x0F);
6460        let eight = vdupq_n_s8(8);
6461        let mut acc = 0f32;
6462        for gi in 0..gpr {
6463            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
6464            let s = *scales.get_unchecked(gi);
6465            let b = vld1q_u8(t);
6466            let lo = vandq_u8(b, lomask);
6467            let hi = vshrq_n_u8::<4>(b);
6468            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
6469            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
6470            let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
6471            let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
6472            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6473            asm!(
6474                "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
6475                "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
6476                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6477                e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
6478                options(pure, nomem, nostack),
6479            );
6480            acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
6481        }
6482        acc
6483    }
6484}
6485
6486#[cfg(target_arch = "x86_64")]
6487#[target_feature(enable = "avx2")]
6488unsafe fn dot_q4tp_row_avx2(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
6489    // SAFETY: see dot_q4tp_row_sdot.
6490    unsafe {
6491        use core::arch::x86_64::*;
6492        let lomask = _mm_set1_epi8(0x0F);
6493        let eight = _mm256_set1_epi8(8);
6494        let ones = _mm256_set1_epi16(1);
6495        let mut acc = 0f32;
6496        for gi in 0..gpr {
6497            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
6498            let s = *scales.get_unchecked(gi);
6499            let b = _mm_loadu_si128(t as *const __m128i);
6500            let lo = _mm_and_si128(b, lomask);
6501            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
6502            let w = _mm256_sub_epi8(
6503                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
6504                eight,
6505            );
6506            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6507            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
6508            let d = _mm256_madd_epi16(p16, ones);
6509            let hi128 = _mm256_extracti128_si256::<1>(d);
6510            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
6511            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
6512            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
6513            acc += _mm_cvtsi128_si32(s32) as f32 * s;
6514        }
6515        acc
6516    }
6517}
6518
6519/// VNNI twin of `dot_q4tp_row_avx2` (see `dot_q4t_row_vnni` for why the
6520/// 256-bit VL encoding is the one to use here).
6521#[cfg(target_arch = "x86_64")]
6522#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
6523unsafe fn dot_q4tp_row_vnni(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
6524    // SAFETY: see dot_q4tp_row_sdot.
6525    unsafe {
6526        use core::arch::x86_64::*;
6527        let lomask = _mm_set1_epi8(0x0F);
6528        let eight = _mm256_set1_epi8(8);
6529        let mut acc = 0f32;
6530        for gi in 0..gpr {
6531            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
6532            let s = *scales.get_unchecked(gi);
6533            let b = _mm_loadu_si128(t as *const __m128i);
6534            let lo = _mm_and_si128(b, lomask);
6535            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
6536            let w = _mm256_sub_epi8(
6537                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
6538                eight,
6539            );
6540            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6541            acc += dpbusd_hsum(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w)) as f32 * s;
6542        }
6543        acc
6544    }
6545}
6546
6547/// `2^x` for a 4-lane vector (range-reduced degree-6 polynomial).
6548/// ≤1 ulp vs libm `exp2f` — f32's own precision rounds the polynomial
6549/// home. Used only by [`q4tp_scales_lanes`]; see the ladder note there.
6550///
6551/// # Safety
6552/// AVX2.
6553#[cfg(target_arch = "x86_64")]
6554#[target_feature(enable = "avx2")]
6555#[inline]
6556unsafe fn exp2_128(x: core::arch::x86_64::__m128) -> core::arch::x86_64::__m128 {
6557    unsafe {
6558        use core::arch::x86_64::*;
6559        // floor via trunc with a below-integer correction.
6560        let tr = _mm_cvttps_epi32(x);
6561        // Compare against cvt(tr), NOT a bit-cast: the bit pattern of a
6562        // negative i32 is a NaN and cmplt against it is always false.
6563        let below = _mm_and_si128(
6564            _mm_castps_si128(_mm_cmplt_ps(x, _mm_cvtepi32_ps(tr))),
6565            _mm_set1_epi32(1),
6566        );
6567        let n = _mm_sub_epi32(tr, below); // floor(x)
6568        let r = _mm_sub_ps(x, _mm_cvtepi32_ps(n)); // x - floor(x) ∈ [0,1)
6569        // 2^r = exp(r·ln2) ≈ Σ aₖrᵏ, aₖ = ln2ᵏ/k!, k = 1..9 (Horner).
6570        // At r ∈ [0,1) the a₁₀ term is 7.1e-9 — below f32's own ulp.
6571        // The coefficients are carried at full decimal precision on
6572        // purpose: truncating them to an f32 spelling (what
6573        // `excessive_precision` suggests) would round the constants before
6574        // they reach the AVX registers. `approx_constant` fires on a₁ =
6575        // LN_2, but the table must stay internally consistent, so both
6576        // lints are allowed for this table only.
6577        #[allow(clippy::excessive_precision, clippy::approx_constant)]
6578        let a = [
6579            6.931_471_805_599_453e-1,
6580            2.402_265_069_591_007e-1,
6581            5.550_410_866_482_158e-2,
6582            9.618_129_107_628_477e-3,
6583            1.333_355_814_642_844e-3,
6584            1.540_353_039_338_161e-4,
6585            1.525_273_380_405_984e-5,
6586            1.321_548_679_014_431e-6,
6587            1.017_808_600_923_969_6e-7,
6588        ];
6589        let mut p = _mm_set1_ps(a[8]);
6590        for &c in a[..8].iter().rev() {
6591            p = _mm_add_ps(_mm_mul_ps(p, r), _mm_set1_ps(c));
6592        }
6593        p = _mm_add_ps(_mm_mul_ps(p, r), _mm_set1_ps(1.0));
6594        // Combine: p ∈ [1,2) carries the mantissa; n the exponent.
6595        // Underflow (n < −126) → 0; our args are f16-derived, so this
6596        // only matters for pathological rows.
6597        // bits(p) already carries exponent 127; adding n<<23 (mod 2^32,
6598        // correct for negative n too) retargets it to 127+n.
6599        let ncl = _mm_min_epi32(
6600            _mm_max_epi32(n, _mm_set1_epi32(-126)),
6601            _mm_set1_epi32(127),
6602        );
6603        let mut bits = _mm_add_epi32(_mm_castps_si128(p), _mm_slli_epi32(ncl, 23));
6604        let over = _mm_cmpgt_epi32(n, _mm_set1_epi32(127));
6605        bits = _mm_or_si128(
6606            _mm_andnot_si128(over, bits),
6607            _mm_and_si128(over, _mm_set1_epi32(0x007F_8000)),
6608        );
6609        let zero = _mm_cmplt_ps(x, _mm_set1_ps(-149.0));
6610        _mm_andnot_ps(zero, _mm_castsi128_ps(bits))
6611    }
6612}
6613
6614/// Horizontal helper: the four lanes of `v` as `[f32; 4]`.
6615#[cfg(target_arch = "x86_64")]
6616#[inline]
6617unsafe fn lanes4(v: core::arch::x86_64::__m128) -> [f32; 4] {
6618    use core::arch::x86_64::*;
6619    unsafe {
6620        [
6621            _mm_cvtss_f32(v),
6622            _mm_cvtss_f32(_mm_shuffle_ps(v, v, 0x55)),
6623            _mm_cvtss_f32(_mm_shuffle_ps(v, v, 0xAA)),
6624            _mm_cvtss_f32(_mm_shuffle_ps(v, v, 0xFF)),
6625        ]
6626    }
6627}
6628
6629/// FOUR rows' 32-run q4tp ladders at once into `[group * 4 + lane]`.
6630/// The chain is `t = exp2(lo); t *= exp2(st)` — per lane this is the
6631/// scalar `q4tp_ladder` recurrence run as `vmulps`, bit-identical; only
6632/// `exp2` is replaced (≤1 ulp, see [`exp2_128`]).
6633///
6634/// # Safety
6635/// AVX2.
6636#[cfg(target_arch = "x86_64")]
6637#[target_feature(enable = "avx2")]
6638#[inline]
6639unsafe fn ladder_4(lo4: &[f32; 4], st4: &[f32; 4], tabc: &mut [f32; 128]) {
6640    unsafe {
6641        use core::arch::x86_64::*;
6642        let ratio = exp2_128(_mm_loadu_ps(st4.as_ptr()));
6643        let mut t = exp2_128(_mm_loadu_ps(lo4.as_ptr()));
6644        tabc[0..4].copy_from_slice(&lanes4(t));
6645        for c in 1..32 {
6646            t = _mm_mul_ps(t, ratio);
6647            tabc[c * 4..c * 4 + 4].copy_from_slice(&lanes4(t));
6648        }
6649    }
6650}
6651
6652/// Decode FOUR rows' q4tp group scales at once. The scalar
6653/// [`Q4tpView::scales_into`] costs ~364 cycles a row — two libm `exp2`
6654/// calls plus the 31-step serial `t *= ratio` ladder chain — and the
6655/// ladder is 17 % of every decode token. Here the chains run as one
6656/// 4-lane vector: `t[c] = t[c-1] * ratio` is `vmulps`, bit-identical to
6657/// the scalar f32 chain (and to `q4tp_ladder`'s definition), so the only
6658/// bit change is `exp2` (≤1 ulp, see [`exp2_128`]); the dequant reference
6659/// shares this function. Output layout is `[row][group]`.
6660///
6661/// Only for `Q4TiledP` views (`!zero_rung`); rows need not be
6662/// contiguous. Falls back per incomplete tail chunk.
6663///
6664/// # Safety
6665/// AVX2; `out` must hold `rows.len() * gpr` floats.
6666#[cfg(target_arch = "x86_64")]
6667#[target_feature(enable = "avx2")]
6668unsafe fn q4tp_scales_lanes(
6669    vs: &[&Q4tpView<'_>],
6670    rows: &[usize],
6671    gpr: usize,
6672    out: &mut [f32],
6673) {
6674    use core::arch::x86_64::*;
6675    debug_assert_eq!(vs.len(), rows.len());
6676    debug_assert!(out.len() >= rows.len() * gpr);
6677    let mut done = 0usize;
6678    for (ci, chunk) in rows.chunks(4).enumerate() {
6679        let vchunk = &vs[ci * 4..ci * 4 + chunk.len()];
6680        if chunk.len() < 4 || vchunk.iter().any(|v| v.zero_rung) {
6681            for (l, &r) in chunk.iter().enumerate() {
6682                vchunk[l].scales_into(r, gpr, &mut out[(done + l) * gpr..(done + l + 1) * gpr]);
6683            }
6684            done += chunk.len();
6685            continue;
6686        }
6687        let (r0, r1, r2, r3) = (chunk[0], chunk[1], chunk[2], chunk[3]);
6688        let v0 = vchunk[0];
6689        let v1 = vchunk[1];
6690        let v2 = vchunk[2];
6691        let v3 = vchunk[3];
6692        // Gather the four (lo, st) f16 pairs — 4 B per row.
6693        let mut lo = [0f32; 4];
6694        let mut st = [0f32; 4];
6695        for (l, r) in [r0, r1, r2, r3].iter().enumerate() {
6696            let o = *r * 4;
6697            let vv = vchunk[l];
6698            lo[l] = cortiq_core::quant::f16_to_f32(u16::from_le_bytes([
6699                vv.params[o],
6700                vv.params[o + 1],
6701            ]));
6702            st[l] = cortiq_core::quant::f16_to_f32(u16::from_le_bytes([
6703                vv.params[o + 2],
6704                vv.params[o + 3],
6705            ]));
6706        }
6707        let lo_v = _mm_loadu_ps(lo.as_ptr());
6708        let st_v = _mm_loadu_ps(st.as_ptr());
6709        let ratio = exp2_128(st_v);
6710        let mut t = exp2_128(lo_v);
6711        // tab layout [group * 4 + lane] — contiguous stores per step.
6712        let mut tabc = [0f32; 32 * 4];
6713        tabc[0..4].copy_from_slice(&[
6714            _mm_cvtss_f32(t),
6715            _mm_cvtss_f32(_mm_shuffle_ps(t, t, 0x55)),
6716            _mm_cvtss_f32(_mm_shuffle_ps(t, t, 0xAA)),
6717            _mm_cvtss_f32(_mm_shuffle_ps(t, t, 0xFF)),
6718        ]);
6719        for c in 1..32 {
6720            t = _mm_mul_ps(t, ratio);
6721            tabc[c * 4..c * 4 + 4].copy_from_slice(&[
6722                _mm_cvtss_f32(t),
6723                _mm_cvtss_f32(_mm_shuffle_ps(t, t, 0x55)),
6724                _mm_cvtss_f32(_mm_shuffle_ps(t, t, 0xAA)),
6725                _mm_cvtss_f32(_mm_shuffle_ps(t, t, 0xFF)),
6726            ]);
6727        }
6728        // Scalar code walks (per row), reading tabc[group * 4 + lane].
6729        let vs4 = [v0, v1, v2, v3];
6730        for (l, &r) in [r0, r1, r2, r3].iter().enumerate() {
6731            let v = vs4[l];
6732            let codes = &v.codes[r * v.stride..(r + 1) * v.stride];
6733            let dst = &mut out[(done + l) * gpr..(done + l + 1) * gpr];
6734            let mut chunks = dst.chunks_exact_mut(8);
6735            let mut ci = 0usize;
6736            for c in &mut chunks {
6737                let w = u64::from(codes[ci])
6738                    | u64::from(codes[ci + 1]) << 8
6739                    | u64::from(codes[ci + 2]) << 16
6740                    | u64::from(codes[ci + 3]) << 24
6741                    | u64::from(codes[ci + 4]) << 32;
6742                for (k, o) in c.iter_mut().enumerate() {
6743                    *o = tabc[(((w >> (5 * k)) & 31) as usize) * 4 + l];
6744                }
6745                ci += 5;
6746            }
6747            let tail = &codes[ci..];
6748            for (k, o) in chunks.into_remainder().iter_mut().enumerate() {
6749                *o = tabc[q4tp_code(tail, k) * 4 + l];
6750            }
6751        }
6752        done += 4;
6753    }
6754}
6755
6756/// Decode-lane count for the q4tp row loops: how many rows one worker
6757/// steps through with independent `acc += d * scale` chains in ONE body.
6758/// The per-group AVX2 dot is a ~17-cycle dependent chain per 32 weights;
6759/// a single row serializes on it, overlapping rows reaches the uop-issue
6760/// bound instead. Each row keeps its own group order, so every output is
6761/// bit-identical to the plain row loop. `CMF_Q4TP_LANES=1` restores it;
6762/// 2/4/8 are the supported lane counts (default 4).
6763fn q4tp_lanes() -> usize {
6764    #[cfg(not(target_arch = "x86_64"))]
6765    {
6766        1
6767    }
6768    #[cfg(target_arch = "x86_64")]
6769    {
6770        static ON: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
6771        *ON.get_or_init(|| {
6772            let n: usize = std::env::var("CMF_Q4TP_LANES")
6773                .ok()
6774                .and_then(|v| v.parse().ok())
6775                .unwrap_or(4);
6776            match n {
6777                1 => 1,
6778                2 => 2,
6779                8 => 8,
6780                _ => 4,
6781            }
6782        })
6783    }
6784}
6785
6786/// One q4tp group tile unpacked to interleaved weights in 0..=16 — the
6787/// raw nibbles, WITHOUT the −8 shift. The correction travels in
6788/// `SplitAct::xg8` instead (each vertical lane subtracts `Σx`, and the
6789/// 8-lane reduction makes that `8·Σx` overall), which lets `vpmaddubsw`
6790/// consume the unsigned codes directly and deletes the per-byte `vpaddb`
6791/// + `vpabsb`/`vpsignb` — two of them port-5 ops in the old row
6792/// kernel's innermost loop.
6793///
6794/// # Safety
6795/// AVX2; `tile` must point at 16 valid weight bytes.
6796#[cfg(target_arch = "x86_64")]
6797#[target_feature(enable = "avx2")]
6798#[inline]
6799unsafe fn q4tp_unpack_w(tile: *const u8) -> core::arch::x86_64::__m256i {
6800    unsafe {
6801        use core::arch::x86_64::*;
6802        let lomask = _mm_set1_epi8(0x0F);
6803        let b = _mm_loadu_si128(tile as *const __m128i);
6804        let lo = _mm_and_si128(b, lomask);
6805        let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
6806        _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi))
6807    }
6808}
6809
6810/// Σ over 32 columns of `w·x` as i32x8 lane partials: `vpmaddubsw`
6811/// (|w| ≤ 15, |x| ≤ 127 — pair sums ≤ 3810, no saturation) then `vpmaddwd`.
6812/// Exact in any association.
6813///
6814/// # Safety
6815/// AVX2; `xp` must point at 32 valid i8 activation bytes.
6816#[cfg(target_arch = "x86_64")]
6817#[target_feature(enable = "avx2")]
6818#[inline]
6819unsafe fn q4tp_madd_wx(
6820    w: core::arch::x86_64::__m256i,
6821    xp: *const i8,
6822) -> core::arch::x86_64::__m256i {
6823    unsafe {
6824        use core::arch::x86_64::*;
6825        let ones = _mm256_set1_epi16(1);
6826        let x = _mm256_loadu_si256(xp as *const __m256i);
6827        _mm256_madd_epi16(_mm256_maddubs_epi16(w, x), ones)
6828    }
6829}
6830
6831/// THE shared end-of-row reduction of the vertical f32 accumulator.
6832/// Every a8w8 q4tp path — the decode lane drivers and the row-exact
6833/// matmat cells — must reduce through this one function, or the
6834/// matvec↔matmat bitwise parity test breaks.
6835#[cfg(target_arch = "x86_64")]
6836#[target_feature(enable = "avx2")]
6837#[inline]
6838unsafe fn q4tp_hsum256(acc: core::arch::x86_64::__m256) -> f32 {
6839    unsafe {
6840        use core::arch::x86_64::*;
6841        let a = _mm_add_ps(_mm256_castps256_ps128(acc), _mm256_extractf128_ps::<1>(acc));
6842        let b = _mm_add_ps(a, _mm_movehl_ps(a, a));
6843        let c = _mm_add_ss(b, _mm_shuffle_ps(b, b, 0x55));
6844        _mm_cvtss_f32(c)
6845    }
6846}
6847
6848/// The `− 8·Σx` term of a q4tp a8w8 row: `8 · Σ_g Σx_g · scale_g`, an f32
6849/// dot over the groups. Every vertical-cell path subtracts it after its
6850/// horizontal sum through this one function, so their bits agree.
6851#[cfg(target_arch = "x86_64")]
6852#[inline]
6853fn q4tp_xg8_corr(xg8: &[f32], sc: &[f32]) -> f32 {
6854    let n = sc.len().min(xg8.len());
6855    8.0 * crate::attention::dot_f32(&xg8[..n], &sc[..n])
6856}
6857
6858/// One (row, activation) cell computed the vertical way: per group,
6859/// `Σ(w+8)·x as f32 · scale` fused into lane-wise f32x8 accumulators, one
6860/// horizontal reduction, then the `8·Σx` term (`q4tp_xg8_corr`). Used by the lane drivers' tails
6861/// and by the row-exact matmat column loop so both land on identical bits.
6862///
6863/// # Safety
6864/// AVX2; `gpr` groups must be covered by `xq`/`sc`/`xg8`.
6865#[cfg(target_arch = "x86_64")]
6866#[target_feature(enable = "avx2,fma")]
6867#[inline]
6868unsafe fn q4tp_cell_dot_v(
6869    nib: *const u8,
6870    r: usize,
6871    gpr: usize,
6872    xp: *const i8,
6873    sc: *const f32,
6874    xg8: *const f32,
6875) -> f32 {
6876    unsafe {
6877        use core::arch::x86_64::*;
6878        let mut acc = _mm256_setzero_ps();
6879        for gi in 0..gpr {
6880            let w = q4tp_unpack_w(nib.add((r * gpr + gi) * Q4TP_NIB));
6881            let d = q4tp_madd_wx(w, xp.add(gi * GROUP_SIZE));
6882            acc = _mm256_fmadd_ps(_mm256_cvtepi32_ps(d), _mm256_set1_ps(*sc.add(gi)), acc);
6883        }
6884        q4tp_hsum256(acc)
6885            - q4tp_xg8_corr(
6886                std::slice::from_raw_parts(xg8, gpr),
6887                std::slice::from_raw_parts(sc, gpr),
6888            )
6889    }
6890}
6891
6892/// Row `r` of a row-exact q4tp matmat against one activation panel (up to
6893/// four columns at a time): each weight group is unpacked ONCE and fed to
6894/// four independent column accumulators — the cell math of
6895/// [`q4tp_cell_dot_v`], gi-ascending, so every cell equals the decode
6896/// matvec's bits while the columns overlap their latency chains and the
6897/// unpack is shared. This is the AVX2 twin of the avx512-only blocked
6898/// kernel: without it, Broadwell prefill re-decodes every row once per
6899/// column (7 tok/s measured against q8_2f's 24).
6900///
6901/// # Safety
6902/// AVX2 (caller gates on `a8w8_enabled()` + `xg8.len() == gpr`);
6903/// `out` cells for this row are this worker's alone.
6904#[cfg(target_arch = "x86_64")]
6905#[target_feature(enable = "avx2,fma")]
6906unsafe fn q4tp_matmat_row_cols_v(
6907    v: &Q4tpView<'_>,
6908    r: usize,
6909    gpr: usize,
6910    acts: &[SplitAct],
6911    sc: &[f32],
6912    abase: usize,
6913    rows: usize,
6914    out: SendMut,
6915) {
6916    use core::arch::x86_64::*;
6917    debug_assert_eq!(sc.len(), gpr);
6918    let n = acts.len();
6919    let mut bi = 0usize;
6920    while bi + 4 <= n {
6921        let a: [&SplitAct; 4] = [&acts[bi], &acts[bi + 1], &acts[bi + 2], &acts[bi + 3]];
6922        let mut accs = [_mm256_setzero_ps(); 4];
6923        for gi in 0..gpr {
6924            let w = q4tp_unpack_w(v.nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB));
6925            let sb = _mm256_set1_ps(*sc.get_unchecked(gi));
6926            for (k, act) in a.iter().enumerate() {
6927                let xp = act.xq.as_ptr().add(gi * GROUP_SIZE);
6928                let d = q4tp_madd_wx(w, xp);
6929                accs[k] = _mm256_fmadd_ps(_mm256_cvtepi32_ps(d), sb, accs[k]);
6930            }
6931        }
6932        for k in 0..4 {
6933            let act = a[k];
6934            let mut acc = (q4tp_hsum256(accs[k]) - q4tp_xg8_corr(&act.xg8, sc)) * act.sx;
6935            for &(j, xv) in &act.outliers {
6936                let (w, s) = q4tp_outlier(v.nib, r, gpr, j, sc);
6937                acc += w * s * xv;
6938            }
6939            // SAFETY: disjoint (bi, r) cells per worker.
6940            unsafe { *out.at((abase + bi + k) * rows + r) = acc };
6941        }
6942        bi += 4;
6943    }
6944    // Remainder columns (1..=3): one vertical cell each.
6945    while bi < n {
6946        let act = &acts[bi];
6947        let mut acc = q4tp_cell_dot_v(
6948            v.nib.as_ptr(),
6949            r,
6950            gpr,
6951            act.xq.as_ptr(),
6952            sc.as_ptr(),
6953            act.xg8.as_ptr(),
6954        ) * act.sx;
6955        for &(j, xv) in &act.outliers {
6956            let (w, s) = q4tp_outlier(v.nib, r, gpr, j, sc);
6957            acc += w * s * xv;
6958        }
6959        // SAFETY: disjoint (bi, r) cells per worker.
6960        unsafe { *out.at((abase + bi) * rows + r) = acc };
6961        bi += 1;
6962    }
6963}
6964
6965/// Columns per activation panel of the q4tp batch: `CMF_Q4TP_PANEL`, else
6966/// 256 — and on x86-64 as many as fill half the core's L2 (`cols` int8
6967/// bytes a column), a multiple of four from 8 to 256. A panel of 256
6968/// columns of a 2560-wide row is 655 KB: on a 512 KB L2 (Zen 1/2) every
6969/// weight row streamed the whole panel from L3 again.
6970fn q4tp_panel_cols(cols: usize) -> usize {
6971    static ENV: std::sync::OnceLock<Option<usize>> = std::sync::OnceLock::new();
6972    if let Some(v) = *ENV.get_or_init(|| {
6973        std::env::var("CMF_Q4TP_PANEL")
6974            .ok()
6975            .and_then(|v| v.parse().ok())
6976            .filter(|v: &usize| *v > 0)
6977    }) {
6978        return v;
6979    }
6980    #[cfg(target_arch = "x86_64")]
6981    {
6982        ((l2_bytes() / 2) / cols.max(1)).clamp(8, 256) & !3
6983    }
6984    #[cfg(not(target_arch = "x86_64"))]
6985    {
6986        let _ = cols;
6987        256
6988    }
6989}
6990
6991/// The per-core L2 size (Linux sysfs; 1 MB when unknown).
6992#[cfg(target_arch = "x86_64")]
6993fn l2_bytes() -> usize {
6994    static L2: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
6995    *L2.get_or_init(|| {
6996        let s = std::fs::read_to_string("/sys/devices/system/cpu/cpu0/cache/index2/size").unwrap_or_default();
6997        let s = s.trim();
6998        let (num, mul) = match s.strip_suffix('K') {
6999            Some(n) => (n, 1024),
7000            None => match s.strip_suffix('M') {
7001                Some(n) => (n, 1024 * 1024),
7002                None => (s, 1),
7003            },
7004        };
7005        num.parse::<usize>().ok().map(|n| n * mul).filter(|&b| b >= 64 * 1024).unwrap_or(1 << 20)
7006    })
7007}
7008
7009/// `CMF_Q4TP_FMA=0` keeps the row-exact vertical cell
7010/// (`q4tp_matmat_row_cols_v`) for a batch on an AVX2 CPU (A/B only).
7011#[cfg(target_arch = "x86_64")]
7012fn q4tp_fma_x86() -> bool {
7013    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
7014    *ON.get_or_init(|| {
7015        std::is_x86_feature_detected!("avx2")
7016            && std::is_x86_feature_detected!("fma")
7017            && std::env::var("CMF_Q4TP_FMA").map(|v| v != "0").unwrap_or(true)
7018    })
7019}
7020
7021/// Two weight rows (`r`, `r + 1`) against a panel of activation columns,
7022/// the tuned AVX2 shape for a batch that need not be row-exact: four
7023/// columns at a time, so each group's tile is unpacked once for four
7024/// columns and each column's 32 bytes are loaded once for two rows (eight
7025/// lane accumulators, the register file's worth). Per group and cell the
7026/// loop is `vpmaddubsw`, `vpmaddwd`, a convert and one FMA by the group
7027/// scale — the vertical cell's `− 8·Σx` correction leaves the loop and is
7028/// taken once per cell as `8 · Σ_g Σx_g · scale_g` (an f32 dot over the
7029/// groups). The same products as the vertical cell, summed in another
7030/// order: tolerance-class against the matvec, like the ARM tuned shape.
7031/// Columns past the last group of four take the vertical cell, row by row.
7032///
7033/// # Safety
7034/// AVX2 and FMA (`q4tp_fma_x86`); every column carries `xg8` of `gpr`
7035/// groups; `out` cells (`(abase + k)·rows + r`, `r + 1`) disjoint per worker.
7036#[cfg(target_arch = "x86_64")]
7037#[target_feature(enable = "avx2,fma")]
7038#[allow(clippy::too_many_arguments)]
7039unsafe fn q4tp_matmat_2rows_fma(
7040    v: &Q4tpView<'_>,
7041    r: usize,
7042    gpr: usize,
7043    acts: &[SplitAct],
7044    sc0: &[f32],
7045    sc1: &[f32],
7046    abase: usize,
7047    rows: usize,
7048    out: SendMut,
7049) {
7050    use core::arch::x86_64::*;
7051    debug_assert!(sc0.len() == gpr && sc1.len() == gpr);
7052    let n = acts.len();
7053    let nib = v.nib.as_ptr();
7054    let ones = _mm256_set1_epi16(1);
7055    let mut bi = 0usize;
7056    unsafe {
7057        while bi + 4 <= n {
7058            let a: [&SplitAct; 4] = [&acts[bi], &acts[bi + 1], &acts[bi + 2], &acts[bi + 3]];
7059            let xp: [*const i8; 4] = [
7060                a[0].xq.as_ptr(),
7061                a[1].xq.as_ptr(),
7062                a[2].xq.as_ptr(),
7063                a[3].xq.as_ptr(),
7064            ];
7065            let mut c0 = [_mm256_setzero_ps(); 4];
7066            let mut c1 = [_mm256_setzero_ps(); 4];
7067            for gi in 0..gpr {
7068                let w0 = q4tp_unpack_w(nib.add((r * gpr + gi) * Q4TP_NIB));
7069                let w1 = q4tp_unpack_w(nib.add(((r + 1) * gpr + gi) * Q4TP_NIB));
7070                let s0 = _mm256_set1_ps(*sc0.get_unchecked(gi));
7071                let s1 = _mm256_set1_ps(*sc1.get_unchecked(gi));
7072                for k in 0..4 {
7073                    let x = _mm256_loadu_si256(xp[k].add(gi * GROUP_SIZE) as *const __m256i);
7074                    let d0 = _mm256_madd_epi16(_mm256_maddubs_epi16(w0, x), ones);
7075                    let d1 = _mm256_madd_epi16(_mm256_maddubs_epi16(w1, x), ones);
7076                    c0[k] = _mm256_fmadd_ps(_mm256_cvtepi32_ps(d0), s0, c0[k]);
7077                    c1[k] = _mm256_fmadd_ps(_mm256_cvtepi32_ps(d1), s1, c1[k]);
7078                }
7079            }
7080            for (row, cr, scr) in [(r, &c0, sc0), (r + 1, &c1, sc1)] {
7081                for k in 0..4 {
7082                    let act = a[k];
7083                    let mut acc = (q4tp_hsum256(cr[k]) - q4tp_xg8_corr(&act.xg8, scr)) * act.sx;
7084                    for &(j, xv) in &act.outliers {
7085                        let (w, s) = q4tp_outlier(v.nib, row, gpr, j, scr);
7086                        acc += w * s * xv;
7087                    }
7088                    // SAFETY: disjoint (bi, r) cells per worker.
7089                    *out.at((abase + bi + k) * rows + row) = acc;
7090                }
7091            }
7092            bi += 4;
7093        }
7094        if bi < n {
7095            q4tp_matmat_row_cols_v(v, r, gpr, &acts[bi..], sc0, abase + bi, rows, out);
7096            q4tp_matmat_row_cols_v(v, r + 1, gpr, &acts[bi..], sc1, abase + bi, rows, out);
7097        }
7098    }
7099}
7100
7101/// Rows `start..end` of a single-output q4tp a8w8 matvec, `L` independent
7102/// accumulator chains wide. Per row: scales decoded once, groups stepped
7103/// in ascending `gi` with `acc += d as f32 * scale` (the exact plain-row
7104/// order), then the `* sx` scale and the outlier corrections in list
7105/// order — the plain row loop's arithmetic, line for line.
7106///
7107/// # Safety
7108/// AVX2 (callers gate on `a8w8_enabled()`); `out` disjoint per row.
7109#[cfg(target_arch = "x86_64")]
7110#[target_feature(enable = "avx2,fma")]
7111unsafe fn q4tp_matvec_lanes<const L: usize>(
7112    v: &Q4tpView<'_>,
7113    start: usize,
7114    end: usize,
7115    gpr: usize,
7116    xq: &[i8],
7117    xg8: &[f32],
7118    sx: f32,
7119    outliers: &[(usize, f32)],
7120    sc: &mut [f32],
7121    out: SendMut,
7122) {
7123    use core::arch::x86_64::*;
7124    debug_assert!(L >= 2 && sc.len() >= L * gpr && xg8.len() == gpr);
7125    crate::cpuprof::count_lane_rows((end - start) as u64);
7126    let mut base = start;
7127    while base + L <= end {
7128        let _t0 = crate::cpuprof::time(crate::cpuprof::Slot::GuScales);
7129        let mut rows = [0usize; L];
7130        for l in 0..L {
7131            rows[l] = base + l;
7132        }
7133        let v4: [&Q4tpView; L] = [v; L];
7134        q4tp_scales_lanes(&v4, &rows, gpr, sc);
7135        drop(_t0);
7136        let _t1 = crate::cpuprof::time(crate::cpuprof::Slot::GuDots);
7137        let mut accs = [_mm256_setzero_ps(); L];
7138        for gi in 0..gpr {
7139            let xp = xq.as_ptr().add(gi * GROUP_SIZE);
7140            for l in 0..L {
7141                let w = q4tp_unpack_w(v.nib.as_ptr().add((rows[l] * gpr + gi) * Q4TP_NIB));
7142                let d = q4tp_madd_wx(w, xp);
7143                let sb = _mm256_set1_ps(*sc.get_unchecked(l * gpr + gi));
7144                accs[l] = _mm256_fmadd_ps(_mm256_cvtepi32_ps(d), sb, accs[l]);
7145            }
7146        }
7147        drop(_t1);
7148        let _t2 = crate::cpuprof::time(crate::cpuprof::Slot::GuPost);
7149        for l in 0..L {
7150            let r = rows[l];
7151            let scl = &sc[l * gpr..(l + 1) * gpr];
7152            let mut acc = (q4tp_hsum256(accs[l]) - q4tp_xg8_corr(xg8, scl)) * sx;
7153            for &(j, xv) in outliers {
7154                let (w, s) = q4tp_outlier(v.nib, r, gpr, j, scl);
7155                acc += w * s * xv;
7156            }
7157            // SAFETY: one lane owns each row.
7158            unsafe { *out.at(r) = acc };
7159        }
7160        drop(_t2);
7161        base += L;
7162    }
7163    // Partial lane batch (and any tail): one vertical cell per row —
7164    // the same reduction shape as the lanes above and as the row-exact
7165    // matmat, so bits agree across every path.
7166    for r in base..end {
7167        let scl = &mut sc[..gpr];
7168        v.scales_into(r, gpr, scl);
7169        let scl: &[f32] = scl;
7170        let mut acc = q4tp_cell_dot_v(v.nib.as_ptr(), r, gpr, xq.as_ptr(), scl.as_ptr(), xg8.as_ptr())
7171            * sx;
7172        for &(j, xv) in outliers {
7173            let (w, s) = q4tp_outlier(v.nib, r, gpr, j, scl);
7174            acc += w * s * xv;
7175        }
7176        // SAFETY: disjoint row ranges per worker.
7177        unsafe { *out.at(r) = acc };
7178    }
7179}
7180
7181/// Rows `start..end` of the fused SwiGLU gate+up q4tp arm, `L` chains
7182/// wide: both projections' scales decoded per lane, both group loops
7183/// stepped in the plain arm's per-row order, then `silu_mul_limited`
7184/// over the outlier-corrected pair — bit-identical to the plain arm.
7185///
7186/// # Safety
7187/// AVX2 (callers gate on `a8w8_enabled()`); `out` disjoint per row.
7188#[cfg(target_arch = "x86_64")]
7189#[target_feature(enable = "avx2,fma")]
7190unsafe fn q4tp_gateup_lanes<const L: usize>(
7191    gv: &Q4tpView<'_>,
7192    uv: &Q4tpView<'_>,
7193    start: usize,
7194    end: usize,
7195    gpr: usize,
7196    xq: &[i8],
7197    xg8: &[f32],
7198    sx: f32,
7199    limit: f32,
7200    outliers: &[(usize, f32)],
7201    gsc: &mut [f32],
7202    usc: &mut [f32],
7203    out: SendMut,
7204) {
7205    use core::arch::x86_64::*;
7206    debug_assert!(
7207        L >= 2 && gsc.len() >= L * gpr && usc.len() >= L * gpr && xg8.len() == gpr
7208    );
7209    crate::cpuprof::count_lane_rows((end - start) as u64);
7210    let mut base = start;
7211    while base + L <= end {
7212        let _t0 = crate::cpuprof::time(crate::cpuprof::Slot::GuScales);
7213        let mut rows = [0usize; L];
7214        for l in 0..L {
7215            rows[l] = base + l;
7216        }
7217        let gv4: [&Q4tpView; L] = [gv; L];
7218        let uv4: [&Q4tpView; L] = [uv; L];
7219        q4tp_scales_lanes(&gv4, &rows, gpr, gsc);
7220        q4tp_scales_lanes(&uv4, &rows, gpr, usc);
7221        drop(_t0);
7222        let _t1 = crate::cpuprof::time(crate::cpuprof::Slot::GuDots);
7223        let mut ag = [_mm256_setzero_ps(); L];
7224        let mut au = [_mm256_setzero_ps(); L];
7225        for gi in 0..gpr {
7226            let xp = xq.as_ptr().add(gi * GROUP_SIZE);
7227            for l in 0..L {
7228                let w = q4tp_unpack_w(gv.nib.as_ptr().add((rows[l] * gpr + gi) * Q4TP_NIB));
7229                let d = q4tp_madd_wx(w, xp);
7230                let sb = _mm256_set1_ps(*gsc.get_unchecked(l * gpr + gi));
7231                ag[l] = _mm256_fmadd_ps(_mm256_cvtepi32_ps(d), sb, ag[l]);
7232            }
7233            for l in 0..L {
7234                let w = q4tp_unpack_w(uv.nib.as_ptr().add((rows[l] * gpr + gi) * Q4TP_NIB));
7235                let d = q4tp_madd_wx(w, xp);
7236                let sb = _mm256_set1_ps(*usc.get_unchecked(l * gpr + gi));
7237                au[l] = _mm256_fmadd_ps(_mm256_cvtepi32_ps(d), sb, au[l]);
7238            }
7239        }
7240        drop(_t1);
7241        let _t2 = crate::cpuprof::time(crate::cpuprof::Slot::GuPost);
7242        for l in 0..L {
7243            let r = rows[l];
7244            let gs = &gsc[l * gpr..(l + 1) * gpr];
7245            let us = &usc[l * gpr..(l + 1) * gpr];
7246            let mut g = (q4tp_hsum256(ag[l]) - q4tp_xg8_corr(xg8, gs)) * sx;
7247            let mut u = (q4tp_hsum256(au[l]) - q4tp_xg8_corr(xg8, us)) * sx;
7248            for &(j, xv) in outliers {
7249                let (w, s) = q4tp_outlier(gv.nib, r, gpr, j, gs);
7250                g += w * s * xv;
7251                let (w, s) = q4tp_outlier(uv.nib, r, gpr, j, us);
7252                u += w * s * xv;
7253            }
7254            // SAFETY: one lane owns each row.
7255            unsafe { *out.at(r) = silu_mul_limited(g, u, limit) };
7256        }
7257        drop(_t2);
7258        base += L;
7259    }
7260    // Partial lane batch (and any tail): one vertical cell per stream.
7261    for r in base..end {
7262        let g = &mut gsc[..gpr];
7263        let u = &mut usc[..gpr];
7264        gv.scales_into(r, gpr, g);
7265        uv.scales_into(r, gpr, u);
7266        let (g, u): (&[f32], &[f32]) = (g, u);
7267        let mut gv_ =
7268            q4tp_cell_dot_v(gv.nib.as_ptr(), r, gpr, xq.as_ptr(), g.as_ptr(), xg8.as_ptr()) * sx;
7269        let mut uv_ =
7270            q4tp_cell_dot_v(uv.nib.as_ptr(), r, gpr, xq.as_ptr(), u.as_ptr(), xg8.as_ptr()) * sx;
7271        for &(j, xv) in outliers {
7272            let (w, s) = q4tp_outlier(gv.nib, r, gpr, j, g);
7273            gv_ += w * s * xv;
7274            let (w, s) = q4tp_outlier(uv.nib, r, gpr, j, u);
7275            uv_ += w * s * xv;
7276        }
7277        // SAFETY: disjoint row ranges per worker.
7278        unsafe { *out.at(r) = silu_mul_limited(gv_, uv_, limit) };
7279    }
7280}
7281
7282/// Flat-space sibling of [`q4tp_matvec_lanes`] for `matvec_many`: the
7283/// virtual row space spans several tensors, so each lane resolves its
7284/// (tensor, row) pair through the caller's `locate`, then runs the plain
7285/// per-row math of the fused q4tp arm.
7286///
7287/// # Safety
7288/// AVX2 (callers gate on `a8w8_enabled()`); `outs` disjoint per row.
7289#[cfg(target_arch = "x86_64")]
7290#[target_feature(enable = "avx2,fma")]
7291unsafe fn q4tp_many_lanes<const L: usize>(
7292    views: &[Q4tpView<'_>],
7293    locate: &(dyn Fn(usize) -> (usize, usize)),
7294    start: usize,
7295    end: usize,
7296    gpr: usize,
7297    xq: &[i8],
7298    xg8: &[f32],
7299    sx: f32,
7300    outliers: &[(usize, f32)],
7301    sc: &mut [f32],
7302    outs: &[SendMut],
7303) {
7304    use core::arch::x86_64::*;
7305    debug_assert!(L >= 2 && sc.len() >= L * gpr && xg8.len() == gpr);
7306    crate::cpuprof::count_lane_rows((end - start) as u64);
7307    let mut base = start;
7308    while base + L <= end {
7309        let _t0 = crate::cpuprof::time(crate::cpuprof::Slot::GuScales);
7310        let mut tk = [0usize; L];
7311        let mut rows = [0usize; L];
7312        for l in 0..L {
7313            let (t, r) = locate(base + l);
7314            tk[l] = t;
7315            rows[l] = r;
7316        }
7317        let v4: [&Q4tpView; L] = std::array::from_fn(|l| &views[tk[l]]);
7318        q4tp_scales_lanes(&v4, &rows, gpr, sc);
7319        drop(_t0);
7320        let _t1 = crate::cpuprof::time(crate::cpuprof::Slot::GuDots);
7321        let mut accs = [_mm256_setzero_ps(); L];
7322        for gi in 0..gpr {
7323            let xp = xq.as_ptr().add(gi * GROUP_SIZE);
7324            for l in 0..L {
7325                let w = q4tp_unpack_w(
7326                    views[tk[l]].nib.as_ptr().add((rows[l] * gpr + gi) * Q4TP_NIB),
7327                );
7328                let d = q4tp_madd_wx(w, xp);
7329                let sb = _mm256_set1_ps(*sc.get_unchecked(l * gpr + gi));
7330                accs[l] = _mm256_fmadd_ps(_mm256_cvtepi32_ps(d), sb, accs[l]);
7331            }
7332        }
7333        drop(_t1);
7334        let _t2 = crate::cpuprof::time(crate::cpuprof::Slot::GuPost);
7335        for l in 0..L {
7336            let (t, r) = (tk[l], rows[l]);
7337            let scl = &sc[l * gpr..(l + 1) * gpr];
7338            let mut acc = (q4tp_hsum256(accs[l]) - q4tp_xg8_corr(xg8, scl)) * sx;
7339            for &(j, xv) in outliers {
7340                let (w, s) = q4tp_outlier(views[t].nib, r, gpr, j, scl);
7341                acc += w * s * xv;
7342            }
7343            // SAFETY: one lane owns each (tensor, row) pair.
7344            unsafe { *outs[t].at(r) = acc };
7345        }
7346        drop(_t2);
7347        base += L;
7348    }
7349    // Partial lane batch (and any tail): one vertical cell per row.
7350    for flat in base..end {
7351        let (t, r) = locate(flat);
7352        let v = &views[t];
7353        let scl = &mut sc[..gpr];
7354        v.scales_into(r, gpr, scl);
7355        let scl: &[f32] = scl;
7356        let mut acc =
7357            q4tp_cell_dot_v(v.nib.as_ptr(), r, gpr, xq.as_ptr(), scl.as_ptr(), xg8.as_ptr()) * sx;
7358        for &(j, xv) in outliers {
7359            let (w, s) = q4tp_outlier(v.nib, r, gpr, j, scl);
7360            acc += w * s * xv;
7361        }
7362        // SAFETY: one worker owns each (tensor, row) pair.
7363        unsafe { *outs[t].at(r) = acc };
7364    }
7365}
7366
7367/// Exact scalar q4tp row — the `CMF_SDOT=0` contract, same pairwise
7368/// accumulation shape as `q4t_row_exact`.
7369#[inline]
7370fn q4tp_row_exact(nib: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
7371    #[cfg(target_arch = "x86_64")]
7372    if avx2_enabled() {
7373        // Keep the scalar pair/group reduction order, not a vector sum.
7374        return unsafe { q4tp_row_float_avx2(nib, r, gpr, x, scales) };
7375    }
7376    q4tp_row_float_scalar(nib, r, gpr, x, scales)
7377}
7378
7379#[inline]
7380fn q4tp_row_float_scalar(nib: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
7381    let mut acc = 0f32;
7382    for gi in 0..gpr {
7383        let tile = &nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
7384        let s = scales[gi];
7385        let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
7386        let mut ga = 0f32;
7387        for (k, &b) in tile.iter().enumerate() {
7388            ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
7389                + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
7390        }
7391        acc += ga * s;
7392    }
7393    acc
7394}
7395
7396/// Vectorize unpack, conversion and multiplication, but preserve every
7397/// pair addition and the scalar accumulation order. No activation rounding
7398/// or FMA: bit-identical to the float scalar row, including its group scale.
7399#[cfg(target_arch = "x86_64")]
7400#[target_feature(enable = "avx2")]
7401unsafe fn q4tp_row_float_avx2(nib: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
7402    // SAFETY: caller checks AVX2 and provides the same complete 32-element
7403    // groups as the scalar row. Loads/stores are explicitly unaligned.
7404    unsafe {
7405        use core::arch::x86_64::*;
7406        let mask = _mm_set1_epi8(15);
7407        let eight = _mm_set1_epi8(8);
7408        let order = _mm256_setr_epi32(0, 1, 4, 5, 2, 3, 6, 7);
7409        let mut acc = 0.0f32;
7410        for gi in 0..gpr {
7411            let packed = _mm_loadu_si128(nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB).cast());
7412            let lo = _mm_and_si128(packed, mask);
7413            let hi = _mm_and_si128(_mm_srli_epi16::<4>(packed), mask);
7414            let w0 = _mm_sub_epi8(_mm_unpacklo_epi8(lo, hi), eight);
7415            let w1 = _mm_sub_epi8(_mm_unpackhi_epi8(lo, hi), eight);
7416            let xp = x.as_ptr().add(gi * GROUP_SIZE);
7417            let a = _mm256_mul_ps(
7418                _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(w0)),
7419                _mm256_loadu_ps(xp),
7420            );
7421            let b = _mm256_mul_ps(
7422                _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(w0))),
7423                _mm256_loadu_ps(xp.add(8)),
7424            );
7425            let c = _mm256_mul_ps(
7426                _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(w1)),
7427                _mm256_loadu_ps(xp.add(16)),
7428            );
7429            let d = _mm256_mul_ps(
7430                _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(w1))),
7431                _mm256_loadu_ps(xp.add(24)),
7432            );
7433            let mut pairs = [0.0f32; 16];
7434            _mm256_storeu_ps(
7435                pairs.as_mut_ptr(),
7436                _mm256_permutevar8x32_ps(_mm256_hadd_ps(a, b), order),
7437            );
7438            _mm256_storeu_ps(
7439                pairs.as_mut_ptr().add(8),
7440                _mm256_permutevar8x32_ps(_mm256_hadd_ps(c, d), order),
7441            );
7442            let mut ga = 0.0f32;
7443            for v in pairs {
7444                ga += v;
7445            }
7446            acc += ga * scales[gi];
7447        }
7448        acc
7449    }
7450}
7451
7452/// Single weight of a q4tp tensor — the a8w8 outlier path, which restores
7453/// activation outliers at full precision after the int8 pass.
7454#[inline]
7455fn q4tp_outlier(nib: &[u8], r: usize, gpr: usize, j: usize, scales: &[f32]) -> (f32, f32) {
7456    let (gi, k) = (j / GROUP_SIZE, j % GROUP_SIZE);
7457    let byte = nib[(r * gpr + gi) * Q4TP_NIB + k / 2];
7458    let n = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
7459    ((n as i32 - 8) as f32, scales[gi])
7460}
7461
7462/// Fused q4tp matvec (dispatch mirrors `q4t_matvec`).
7463fn q4tp_matvec(
7464    bytes: &[u8],
7465    x: &[f32],
7466    rows: usize,
7467    cols: usize,
7468    out: &mut [f32],
7469    pool: Option<&Pool>,
7470) {
7471    debug_assert_eq!(out.len(), rows);
7472    let gpr = cols / GROUP_SIZE;
7473    let v = Q4tpView::new(bytes, rows, cols);
7474    let out_addr = SendMut(out.as_mut_ptr());
7475    if a8w8_enabled() {
7476        let act = split_act(x);
7477        #[cfg(target_arch = "x86_64")]
7478        let lanes = q4tp_lanes();
7479        let run = |start: usize, end: usize| {
7480            #[cfg(target_arch = "x86_64")]
7481            if lanes > 1 && act.xg8.len() == gpr {
7482                // L independent accumulator chains per worker — the plain
7483                // row body's math, lane-interleaved.
7484                with_krow(gpr * lanes, |sc| {
7485                    // SAFETY: a8w8 implies AVX2; disjoint row ranges.
7486                    unsafe {
7487                        match lanes {
7488                            2 => q4tp_matvec_lanes::<2>(
7489                                &v, start, end, gpr, &act.xq, &act.xg8, act.sx,
7490                                &act.outliers, sc, out_addr,
7491                            ),
7492                            8 => q4tp_matvec_lanes::<8>(
7493                                &v, start, end, gpr, &act.xq, &act.xg8, act.sx,
7494                                &act.outliers, sc, out_addr,
7495                            ),
7496                            _ => q4tp_matvec_lanes::<4>(
7497                                &v, start, end, gpr, &act.xq, &act.xg8, act.sx,
7498                                &act.outliers, sc, out_addr,
7499                            ),
7500                        }
7501                    }
7502                });
7503                return;
7504            }
7505            crate::cpuprof::count_plain_rows((end - start) as u64);
7506            // One scratch row of scales per worker — borrowed, not minted.
7507            with_krow(gpr, |sc| {
7508                for r in start..end {
7509                    v.scales_into(r, gpr, sc);
7510                    let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, sc) * act.sx;
7511                    for &(j, xv) in &act.outliers {
7512                        let (w, s) = q4tp_outlier(v.nib, r, gpr, j, sc);
7513                        acc += w * s * xv;
7514                    }
7515                    // SAFETY: disjoint row ranges per worker.
7516                    unsafe { *out_addr.at(r) = acc }
7517                }
7518            })
7519        };
7520        dispatch_rows(pool, rows, &run);
7521        return;
7522    }
7523    let run = |start: usize, end: usize| {
7524        with_krow(gpr, |sc| {
7525            for r in start..end {
7526                v.scales_into(r, gpr, sc);
7527                // SAFETY: disjoint row ranges per worker.
7528                unsafe { *out_addr.at(r) = q4tp_row_exact(v.nib, r, gpr, x, sc) };
7529            }
7530        })
7531    };
7532    dispatch_rows(pool, rows, &run);
7533}
7534
7535/// Fused two-input q4tp matvec — the SwiGLU gate/up pair. Weights and the
7536/// row ladder are read once and spent on both activation streams.
7537#[allow(clippy::too_many_arguments)]
7538fn q4tp_matvec2(
7539    bytes: &[u8],
7540    x1: &[f32],
7541    x2: &[f32],
7542    rows: usize,
7543    cols: usize,
7544    o1: &mut [f32],
7545    o2: &mut [f32],
7546    pool: Option<&Pool>,
7547) {
7548    let gpr = cols / GROUP_SIZE;
7549    let v = Q4tpView::new(bytes, rows, cols);
7550    let (p1, p2) = (SendMut(o1.as_mut_ptr()), SendMut(o2.as_mut_ptr()));
7551    let run = |start: usize, end: usize| {
7552        let mut sc = vec![0f32; gpr];
7553        for r in start..end {
7554            v.scales_into(r, gpr, &mut sc);
7555            // SAFETY: disjoint row ranges per worker.
7556            unsafe {
7557                *p1.at(r) = q4tp_row_exact(v.nib, r, gpr, x1, &sc);
7558                *p2.at(r) = q4tp_row_exact(v.nib, r, gpr, x2, &sc);
7559            }
7560        }
7561    };
7562    dispatch_rows(pool, rows, &run);
7563}
7564
7565/// One q2tp outlier weight at column `j` of row `r`: the 2-bit code and
7566/// its group scale, mirrored on `q4tp_outlier`.
7567#[inline]
7568fn q2tp_outlier(chunks: &[u8], r: usize, gpr: usize, j: usize, scales: &[f32]) -> (f32, f32) {
7569    let (gi, k) = (j / GROUP_SIZE, j % GROUP_SIZE);
7570    let byte = chunks[(r * gpr + gi) * Q2TP_CHUNK + k / 4];
7571    let c = (byte >> (2 * (k % 4))) & 3;
7572    (c as f32 - 1.5, scales[gi])
7573}
7574
7575#[cfg(target_arch = "x86_64")]
7576const Q2TP_DECODE_U32: [u32; 256] = {
7577    let mut tab = [0u32; 256];
7578    let mut b = 0usize;
7579    while b < 256 {
7580        tab[b] = ((b as u32) & 3)
7581            | ((((b as u32) >> 2) & 3) << 8)
7582            | ((((b as u32) >> 4) & 3) << 16)
7583            | ((((b as u32) >> 6) & 3) << 24);
7584        b += 1;
7585    }
7586    tab
7587};
7588
7589/// Eight packed q2tp bytes against 32 signed activation bytes. `maddubs`
7590/// exactly computes unsigned 2-bit code × signed i8; its pair sums cannot
7591/// saturate (2 × 3 × 127 < i16::MAX), and the second madd widens to i32.
7592#[cfg(target_arch = "x86_64")]
7593#[target_feature(enable = "avx2")]
7594unsafe fn q2tp_code_dot_avx2(ch: &[u8], x: &[i8]) -> i32 {
7595    use core::arch::x86_64::*;
7596    debug_assert!(ch.len() >= Q2TP_CHUNK && x.len() >= GROUP_SIZE);
7597    let codes = _mm256_setr_epi32(
7598        Q2TP_DECODE_U32[ch[0] as usize] as i32,
7599        Q2TP_DECODE_U32[ch[1] as usize] as i32,
7600        Q2TP_DECODE_U32[ch[2] as usize] as i32,
7601        Q2TP_DECODE_U32[ch[3] as usize] as i32,
7602        Q2TP_DECODE_U32[ch[4] as usize] as i32,
7603        Q2TP_DECODE_U32[ch[5] as usize] as i32,
7604        Q2TP_DECODE_U32[ch[6] as usize] as i32,
7605        Q2TP_DECODE_U32[ch[7] as usize] as i32,
7606    );
7607    let xv = unsafe { _mm256_loadu_si256(x.as_ptr().cast()) };
7608    let pair = _mm256_maddubs_epi16(codes, xv);
7609    let quad = _mm256_madd_epi16(pair, _mm256_set1_epi16(1));
7610    let sum128 = _mm_add_epi32(
7611        _mm256_castsi256_si128(quad),
7612        _mm256_extracti128_si256(quad, 1),
7613    );
7614    let sum64 = _mm_hadd_epi32(sum128, sum128);
7615    _mm_cvtsi128_si32(_mm_hadd_epi32(sum64, sum64))
7616}
7617
7618#[cfg(target_arch = "x86_64")]
7619#[inline]
7620fn q2tp_avx2_enabled() -> bool {
7621    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
7622    *ON.get_or_init(|| std::arch::is_x86_feature_detected!("avx2"))
7623}
7624
7625/// Integer dot of one q2tp row against pre-quantized activations:
7626/// Σ_g s_g · (Σ c·xq − 1.5·Σ xq). The half-integer grid (c − 1.5)
7627/// becomes exact integer math through the group sums — the same trick
7628/// every a8w8 kernel in this file rides. The codes decode into a
7629/// 32-byte scratch in natural order and the dot itself is the shared
7630/// SDOT primitive; elsewhere a scalar integer loop.
7631#[inline]
7632fn dot_q2tp_row_i8(
7633    chunks: &[u8],
7634    r: usize,
7635    gpr: usize,
7636    xq: &[i8],
7637    gsum: &[i32],
7638    scales: &[f32],
7639) -> f32 {
7640    let mut acc = 0f32;
7641    let base = r * gpr * Q2TP_CHUNK;
7642    #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
7643    let mut codes = [0i8; GROUP_SIZE];
7644    #[cfg(target_arch = "x86_64")]
7645    // Cache the process-wide AVX2 decision; probing it for every 32-weight
7646    // group adds a branch to the hottest Q2TP row loop while retaining the
7647    // established table-load AVX2 arithmetic (and without coupling this path
7648    // to the separate FMA-gated A8W8 switch).
7649    let avx2 = q2tp_avx2_enabled();
7650    for gi in 0..gpr {
7651        let ch = &chunks[base + gi * Q2TP_CHUNK..base + (gi + 1) * Q2TP_CHUNK];
7652        let xg = &xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
7653        #[cfg(target_arch = "aarch64")]
7654        // NEON: the byte's four 2-bit fields land in four lane vectors
7655        // (shift+mask), vld4 de-interleaves xq to match (xj[k] =
7656        // xq[4k+j]), widening MACs accumulate exactly in i32. A scalar
7657        // decode here cost as much as the dot it fed — the profile put
7658        // it at the top of the whole W2 decode.
7659        let dot = unsafe {
7660            use core::arch::aarch64::*;
7661            let b = vld1_u8(ch.as_ptr());
7662            let three = vdup_n_u8(3);
7663            let c0 = vreinterpret_s8_u8(vand_u8(b, three));
7664            let c1 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 2), three));
7665            let c2 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 4), three));
7666            let c3 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 6), three));
7667            let x4 = vld4_s8(xg.as_ptr());
7668            let mut acc4 = vdupq_n_s32(0);
7669            acc4 = vpadalq_s16(acc4, vmull_s8(c0, x4.0));
7670            acc4 = vpadalq_s16(acc4, vmull_s8(c1, x4.1));
7671            acc4 = vpadalq_s16(acc4, vmull_s8(c2, x4.2));
7672            acc4 = vpadalq_s16(acc4, vmull_s8(c3, x4.3));
7673            vaddvq_s32(acc4)
7674        };
7675        #[cfg(target_arch = "x86_64")]
7676        let dot: i32 = if avx2 {
7677            // SAFETY: the runtime feature check gates the target-feature body;
7678            // the group slices above are exactly 8 and 32 bytes long.
7679            unsafe { q2tp_code_dot_avx2(ch, xg) }
7680        } else {
7681            ch.iter()
7682                .enumerate()
7683                .map(|(k, &b)| {
7684                    ((b & 3) as i32) * xg[k * 4] as i32
7685                        + (((b >> 2) & 3) as i32) * xg[k * 4 + 1] as i32
7686                        + (((b >> 4) & 3) as i32) * xg[k * 4 + 2] as i32
7687                        + (((b >> 6) & 3) as i32) * xg[k * 4 + 3] as i32
7688                })
7689                .sum()
7690        };
7691        #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
7692        let dot: i32 = {
7693            for (k, &b) in ch.iter().enumerate() {
7694                codes[k * 4] = (b & 3) as i8;
7695                codes[k * 4 + 1] = ((b >> 2) & 3) as i8;
7696                codes[k * 4 + 2] = ((b >> 4) & 3) as i8;
7697                codes[k * 4 + 3] = ((b >> 6) & 3) as i8;
7698            }
7699            codes
7700                .iter()
7701                .zip(xg)
7702                .map(|(&c, &x)| c as i32 * x as i32)
7703                .sum()
7704        };
7705        acc += scales[gi] * (dot as f32 - 1.5 * gsum[gi] as f32);
7706    }
7707    acc
7708}
7709
7710/// Exact f32 dot of one q2tp row: 2-bit fields LSB-first, (c − 1.5)·s.
7711/// Scalar on purpose — the 2-bit class targets the GPU graph; the CPU
7712/// path exists for parity gates and small-machine fallback.
7713fn q2tp_row_exact(chunks: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
7714    q2tp_row_exact_center(chunks, r, gpr, x, scales, 1.5)
7715}
7716
7717/// Fused Prism affine row: the derived correction is applied inside the
7718/// decoded code, avoiding a second accumulated dot and avoiding cancellation
7719/// between `B=(c-1.5)s` and `+.5s` for long 5120/17408 rows.
7720#[inline]
7721fn q2tp_affine_row_exact(chunks: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
7722    q2tp_row_exact_center(chunks, r, gpr, x, scales, 1.0)
7723}
7724
7725#[inline]
7726fn q2tp_row_exact_center(
7727    chunks: &[u8],
7728    r: usize,
7729    gpr: usize,
7730    x: &[f32],
7731    scales: &[f32],
7732    center: f32,
7733) -> f32 {
7734    let mut acc = 0f32;
7735    for gi in 0..gpr {
7736        let ch = &chunks[(r * gpr + gi) * Q2TP_CHUNK..(r * gpr + gi + 1) * Q2TP_CHUNK];
7737        let s = scales[gi];
7738        let xb = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
7739        let mut g = 0f32;
7740        for (k, &b) in ch.iter().enumerate() {
7741            g += ((b & 3) as f32 - center) * xb[k * 4]
7742                + (((b >> 2) & 3) as f32 - center) * xb[k * 4 + 1]
7743                + (((b >> 4) & 3) as f32 - center) * xb[k * 4 + 2]
7744                + (((b >> 6) & 3) as f32 - center) * xb[k * 4 + 3];
7745        }
7746        acc += s * g;
7747    }
7748    acc
7749}
7750
7751fn q2tp_matvec(
7752    bytes: &[u8],
7753    x: &[f32],
7754    rows: usize,
7755    cols: usize,
7756    out: &mut [f32],
7757    pool: Option<&Pool>,
7758) {
7759    q2tp_matvec_mode(bytes, x, rows, cols, out, pool, false);
7760}
7761
7762fn q2tp_affine_matvec(
7763    bytes: &[u8],
7764    x: &[f32],
7765    rows: usize,
7766    cols: usize,
7767    out: &mut [f32],
7768    pool: Option<&Pool>,
7769) {
7770    q2tp_matvec_mode(bytes, x, rows, cols, out, pool, true);
7771}
7772
7773fn q2tp_matvec_mode(
7774    bytes: &[u8],
7775    x: &[f32],
7776    rows: usize,
7777    cols: usize,
7778    out: &mut [f32],
7779    pool: Option<&Pool>,
7780    affine: bool,
7781) {
7782    debug_assert_eq!(out.len(), rows);
7783    let gpr = cols / GROUP_SIZE;
7784    let v = Q4tpView::new_q2(bytes, rows, cols);
7785    let out_addr = SendMut(out.as_mut_ptr());
7786    // a8w8 fast path (CMF_SDOT=0 keeps the exact scalar walk): integer
7787    // code dots + group sums, exact outlier correction — the same
7788    // contract as every sibling kernel; measured 2-bit rows were the
7789    // only scalar holdout in the family.
7790    if !affine && a8w8_enabled() {
7791        let act = split_act(x);
7792        let gsum = q1_group_sums(&act.xq, gpr);
7793        let (act, gsum) = (&act, &gsum);
7794        let run = move |start: usize, end: usize| {
7795            with_krow(gpr, |sc| {
7796                for r in start..end {
7797                    v.scales_into(r, gpr, sc);
7798                    let mut acc = dot_q2tp_row_i8(v.nib, r, gpr, &act.xq, gsum, sc) * act.sx;
7799                    for &(j, xv) in &act.outliers {
7800                        let (w, s) = q2tp_outlier(v.nib, r, gpr, j, sc);
7801                        acc += w * s * xv;
7802                    }
7803                    // SAFETY: disjoint row ranges per worker.
7804                    unsafe { *out_addr.at(r) = acc };
7805                }
7806            })
7807        };
7808        dispatch_rows(pool, rows, &run);
7809        return;
7810    }
7811    let run = |start: usize, end: usize| {
7812        with_krow(gpr, |sc| {
7813            for r in start..end {
7814                v.scales_into(r, gpr, sc);
7815                // SAFETY: disjoint row ranges per worker.
7816                unsafe {
7817                    *out_addr.at(r) = if affine {
7818                        q2tp_affine_row_exact(v.nib, r, gpr, x, sc)
7819                    } else {
7820                        q2tp_row_exact(v.nib, r, gpr, x, sc)
7821                    }
7822                };
7823            }
7824        })
7825    };
7826    dispatch_rows(pool, rows, &run);
7827}
7828
7829/// Fused two-input q2tp matvec — the SwiGLU gate/up pair.
7830#[allow(clippy::too_many_arguments)]
7831fn q2tp_matvec2(
7832    bytes: &[u8],
7833    x1: &[f32],
7834    x2: &[f32],
7835    rows: usize,
7836    cols: usize,
7837    o1: &mut [f32],
7838    o2: &mut [f32],
7839    pool: Option<&Pool>,
7840) {
7841    q2tp_matvec2_mode(bytes, x1, x2, rows, cols, o1, o2, pool, false);
7842}
7843
7844#[allow(clippy::too_many_arguments)]
7845fn q2tp_affine_matvec2(
7846    bytes: &[u8],
7847    x1: &[f32],
7848    x2: &[f32],
7849    rows: usize,
7850    cols: usize,
7851    o1: &mut [f32],
7852    o2: &mut [f32],
7853    pool: Option<&Pool>,
7854) {
7855    q2tp_matvec2_mode(bytes, x1, x2, rows, cols, o1, o2, pool, true);
7856}
7857
7858#[allow(clippy::too_many_arguments)]
7859fn q2tp_matvec2_mode(
7860    bytes: &[u8],
7861    x1: &[f32],
7862    x2: &[f32],
7863    rows: usize,
7864    cols: usize,
7865    o1: &mut [f32],
7866    o2: &mut [f32],
7867    pool: Option<&Pool>,
7868    affine: bool,
7869) {
7870    let gpr = cols / GROUP_SIZE;
7871    let v = Q4tpView::new_q2(bytes, rows, cols);
7872    let (p1, p2) = (SendMut(o1.as_mut_ptr()), SendMut(o2.as_mut_ptr()));
7873    let run = |start: usize, end: usize| {
7874        let mut sc = vec![0f32; gpr];
7875        for r in start..end {
7876            v.scales_into(r, gpr, &mut sc);
7877            // SAFETY: disjoint row ranges per worker.
7878            unsafe {
7879                *p1.at(r) = if affine {
7880                    q2tp_affine_row_exact(v.nib, r, gpr, x1, &sc)
7881                } else {
7882                    q2tp_row_exact(v.nib, r, gpr, x1, &sc)
7883                };
7884                *p2.at(r) = if affine {
7885                    q2tp_affine_row_exact(v.nib, r, gpr, x2, &sc)
7886                } else {
7887                    q2tp_row_exact(v.nib, r, gpr, x2, &sc)
7888                };
7889            }
7890        }
7891    };
7892    dispatch_rows(pool, rows, &run);
7893}
7894
7895/// Batched q2tp matmat: scalar row kernel over every batch column. CPU
7896/// prefill only — decode rides the graph, so plain and correct beats
7897/// clever here.
7898/// Test doors into the host 2-bit kernels: the stand's heap corruption
7899/// pointed at down-shaped tensors, and the private fns need a way to be
7900/// held to a reference without a model file around them.
7901pub fn q2tp_matvec_for_test(bytes: &[u8], x: &[f32], rows: usize, cols: usize, out: &mut [f32]) {
7902    // The facade IS the reference: encoder oracles hold requant output
7903    // to the exact scalar walk. The production dispatch may take the i8
7904    // fast path, whose error scale is the ACTIVATIONS' — a different
7905    // claim than the encoder correctness these tests pin.
7906    let gpr = cols / GROUP_SIZE;
7907    let v = Q4tpView::new_q2(bytes, rows, cols);
7908    with_krow(gpr, |sc| {
7909        for r in 0..rows {
7910            v.scales_into(r, gpr, sc);
7911            out[r] = q2tp_row_exact(v.nib, r, gpr, x, sc);
7912        }
7913    });
7914}
7915
7916/// Test door for the descriptor-specific fused affine decode.  Production
7917/// callers select this through a validated Prism header, never by dtype alone.
7918pub fn q2tp_affine_matvec_for_test(
7919    bytes: &[u8],
7920    x: &[f32],
7921    rows: usize,
7922    cols: usize,
7923    out: &mut [f32],
7924) {
7925    q2tp_affine_matvec(bytes, x, rows, cols, out, None);
7926}
7927
7928pub fn q2tp_matmat_for_test(
7929    bytes: &[u8],
7930    xs_all: &[f32],
7931    b: usize,
7932    rows: usize,
7933    cols: usize,
7934    out: &mut [f32],
7935) {
7936    q2tp_matmat(bytes, xs_all, b, rows, cols, out, None);
7937}
7938
7939fn q2tp_matmat(
7940    bytes: &[u8],
7941    xs_all: &[f32],
7942    b: usize,
7943    rows: usize,
7944    cols: usize,
7945    out: &mut [f32],
7946    pool: Option<&Pool>,
7947) {
7948    q2tp_matmat_mode(bytes, xs_all, b, rows, cols, out, pool, false);
7949}
7950
7951fn q2tp_affine_matmat(
7952    bytes: &[u8],
7953    xs_all: &[f32],
7954    b: usize,
7955    rows: usize,
7956    cols: usize,
7957    out: &mut [f32],
7958    pool: Option<&Pool>,
7959) {
7960    q2tp_matmat_mode(bytes, xs_all, b, rows, cols, out, pool, true);
7961}
7962
7963fn q2tp_matmat_mode(
7964    bytes: &[u8],
7965    xs_all: &[f32],
7966    b: usize,
7967    rows: usize,
7968    cols: usize,
7969    out: &mut [f32],
7970    pool: Option<&Pool>,
7971    affine: bool,
7972) {
7973    debug_assert_eq!(out.len(), b * rows);
7974    let gpr = cols / GROUP_SIZE;
7975    let v = Q4tpView::new_q2(bytes, rows, cols);
7976    let out_addr = SendMut(out.as_mut_ptr());
7977    let run = |start: usize, end: usize| {
7978        let mut sc = vec![0f32; gpr];
7979        for r in start..end {
7980            v.scales_into(r, gpr, &mut sc);
7981            for bi in 0..b {
7982                let x = &xs_all[bi * cols..(bi + 1) * cols];
7983                // SAFETY: disjoint row ranges per worker.
7984                unsafe {
7985                    *out_addr.at(bi * rows + r) = if affine {
7986                        q2tp_affine_row_exact(v.nib, r, gpr, x, &sc)
7987                    } else {
7988                        q2tp_row_exact(v.nib, r, gpr, x, &sc)
7989                    }
7990                };
7991            }
7992        }
7993    };
7994    dispatch_rows(pool, rows, &run);
7995}
7996
7997/// The pre-vectorised shape, kept for A/B (`CMF_Q4TP_V1=1`): the
7998/// horizontal add lands once per group per column instead of once per
7999/// row. Same weights, same activations — only the reduction differs.
8000/// It is also the row-exact batch kernel: per group and column it forms
8001/// `int dot as f32 * scale` and adds it to a scalar running sum, which is
8002/// `dot_q4tp_row_sdot` step for step (Rust never contracts to an fma),
8003/// so each column is bit-identical to that column's matvec.
8004#[cfg(target_arch = "aarch64")]
8005#[target_feature(enable = "neon,dotprod")]
8006unsafe fn dot_q4tp_row_1x4_sdot_v1(
8007    nib: &[u8],
8008    r: usize,
8009    gpr: usize,
8010    xs: [&[i8]; 4],
8011    scales: &[f32],
8012) -> [f32; 4] {
8013    unsafe {
8014        use core::arch::aarch64::*;
8015        use core::arch::asm;
8016        let lomask = vdupq_n_u8(0x0F);
8017        let eight = vdupq_n_s8(8);
8018        let (mut f0, mut f1, mut f2, mut f3) = (0f32, 0f32, 0f32, 0f32);
8019        for gi in 0..gpr {
8020            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
8021            let s = *scales.get_unchecked(gi);
8022            let bb = vld1q_u8(t);
8023            let lo = vandq_u8(bb, lomask);
8024            let hi = vshrq_n_u8::<4>(bb);
8025            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
8026            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
8027            let mut d = [0f32; 4];
8028            for (k, dk) in d.iter_mut().enumerate() {
8029                let x0 = vld1q_s8(xs[k].as_ptr().add(gi * GROUP_SIZE));
8030                let x1 = vld1q_s8(xs[k].as_ptr().add(gi * GROUP_SIZE + 16));
8031                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
8032                asm!(
8033                    "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
8034                    "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
8035                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
8036                    e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
8037                    options(pure, nomem, nostack),
8038                );
8039                *dk = vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
8040            }
8041            f0 += d[0];
8042            f1 += d[1];
8043            f2 += d[2];
8044            f3 += d[3];
8045        }
8046        [f0, f1, f2, f3]
8047    }
8048}
8049
8050/// Which q4tp batch kernel to run: 1 = the previous one, 2 = the tuned
8051/// one, 0 = decide from the CPU. An atomic rather than a `OnceLock` so a
8052/// benchmark can alternate the two inside one process, where the machine's
8053/// mood — a shared box drifts ±25% between runs — is the same for both.
8054/// What the two mean is per-architecture: on x86 the blocked AVX-512 path
8055/// against the per-column one, on ARM the two reduction shapes.
8056#[allow(dead_code)]
8057static Q4TP_ALT: std::sync::atomic::AtomicU8 = std::sync::atomic::AtomicU8::new(0);
8058
8059/// Tests that store `Q4TP_ALT` hold this, so one test's kernel pick does
8060/// not leak into another's bit-exact comparison running in parallel.
8061#[cfg(test)]
8062static Q4TP_ALT_TEST_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
8063
8064/// Blocking pays on x86 only with 512-bit VNNI. With AVX2 alone, four
8065/// columns sharing an unpack still measured slower than the per-column
8066/// path (23.2 ms against 19.4 on a 48-thread EPYC), because that path
8067/// already dequantizes the row once — so the blocked kernel bought a
8068/// second unpack-free pass at the price of half the vector width.
8069#[cfg(target_arch = "x86_64")]
8070fn q4tp_blocked_x86() -> bool {
8071    match Q4TP_ALT.load(std::sync::atomic::Ordering::Relaxed) {
8072        1 => false,
8073        // A forced ON still asks the CPU. The switch exists so a bench can
8074        // pick a kernel, not so it can promise instructions the machine
8075        // does not have — CI caught that as a SIGILL on a runner without
8076        // AVX-512, where the parity test had turned the path on by hand.
8077        2 => avx512vnni_enabled(),
8078        // Deliberately not cached back into the switch: both gates below
8079        // hold their own `OnceLock`, and latching their answer here would
8080        // make a test's override outlive the test that set it.
8081        _ => blocked_enabled() && avx512vnni_enabled(),
8082    }
8083}
8084
8085/// `CMF_Q4TP_V1=1` picks the old reduction shape (A/B only).
8086#[cfg(target_arch = "aarch64")]
8087#[allow(dead_code)]
8088fn q4tp_v1() -> bool {
8089    match Q4TP_ALT.load(std::sync::atomic::Ordering::Relaxed) {
8090        1 => true,
8091        2 => false,
8092        _ => {
8093            static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8094            *ON.get_or_init(|| std::env::var("CMF_Q4TP_V1").is_ok_and(|v| v != "0"))
8095        }
8096    }
8097}
8098
8099/// Two weight rows against eight columns. The activation load is the
8100/// same for both rows, so it is paid once for twice the arithmetic, and
8101/// sixteen accumulator chains run where eight did — which is what a kernel
8102/// retiring 0.29 instructions a cycle is short of. Register pressure is
8103/// the limit: sixteen `zmm` accumulators, two weight tiles, one
8104/// activation, of thirty-two.
8105///
8106/// Four rows by four columns spends the same sixteen accumulators the
8107/// other way and measured worse — 1488 GFLOP/s against 1644 — so the
8108/// unpack, which four rows pay twice as often, costs more than the extra
8109/// sharing of one activation load buys.
8110#[cfg(target_arch = "x86_64")]
8111#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
8112unsafe fn dot_q4tp_2x8_avx512(
8113    nib: &[u8],
8114    r0: usize,
8115    gpr: usize,
8116    xs: [&[i8]; 8],
8117    sc0: &[f32],
8118    sc1: &[f32],
8119) -> [[f32; 8]; 2] {
8120    // SAFETY: as dot_q4tp_row_1x8_avx512, two adjacent rows at once; the
8121    // caller guarantees r0 + 1 < rows and the ISA.
8122    unsafe {
8123        use core::arch::x86_64::*;
8124        let lomask = _mm256_set1_epi8(0x0F);
8125        let eight = _mm256_set1_epi8(8);
8126        let zero = _mm512_setzero_si512();
8127        let mut v0 = [_mm512_setzero_ps(); 8];
8128        let mut v1 = [_mm512_setzero_ps(); 8];
8129        let pairs = gpr / 2;
8130        let unpack = |r: usize, gi: usize| -> (__m512i, __mmask64) {
8131            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
8132            let bb = _mm256_loadu_si256(t as *const __m256i);
8133            let lo = _mm256_and_si256(bb, lomask);
8134            let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
8135            let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
8136            let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
8137            let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
8138            let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
8139            (_mm512_abs_epi8(w), _mm512_movepi8_mask(w))
8140        };
8141        for gp in 0..pairs {
8142            let gi = gp * 2;
8143            let (wa0, neg0) = unpack(r0, gi);
8144            let (wa1, neg1) = unpack(r0 + 1, gi);
8145            let off = gi * GROUP_SIZE;
8146            let sv = |sc: &[f32]| {
8147                _mm512_insertf32x8::<1>(
8148                    _mm512_castps256_ps512(_mm256_set1_ps(*sc.get_unchecked(gi))),
8149                    _mm256_set1_ps(*sc.get_unchecked(gi + 1)),
8150                )
8151            };
8152            let s0 = sv(sc0);
8153            let s1 = sv(sc1);
8154            for k in 0..8 {
8155                let xv = _mm512_loadu_si512(xs[k].as_ptr().add(off) as *const __m512i);
8156                let d0 = _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(
8157                    zero,
8158                    wa0,
8159                    _mm512_mask_sub_epi8(xv, neg0, zero, xv),
8160                ));
8161                let d1 = _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(
8162                    zero,
8163                    wa1,
8164                    _mm512_mask_sub_epi8(xv, neg1, zero, xv),
8165                ));
8166                v0[k] = _mm512_fmadd_ps(d0, s0, v0[k]);
8167                v1[k] = _mm512_fmadd_ps(d1, s1, v1[k]);
8168            }
8169        }
8170        let mut acc = [[0f32; 8]; 2];
8171        for k in 0..8 {
8172            acc[0][k] = _mm512_reduce_add_ps(v0[k]);
8173            acc[1][k] = _mm512_reduce_add_ps(v1[k]);
8174        }
8175        if gpr % 2 == 1 {
8176            let off = (gpr - 1) * GROUP_SIZE;
8177            for j in off..off + GROUP_SIZE {
8178                let (w0, sa) = q4tp_outlier(nib, r0, gpr, j, sc0);
8179                let (w1, sb) = q4tp_outlier(nib, r0 + 1, gpr, j, sc1);
8180                for k in 0..8 {
8181                    let x = *xs[k].get_unchecked(j) as f32;
8182                    acc[0][k] += w0 * sa * x;
8183                    acc[1][k] += w1 * sb * x;
8184                }
8185            }
8186        }
8187        acc
8188    }
8189}
8190
8191/// The same, eight columns at a time. One unpack then feeds twice as many
8192/// activation streams, so a wide batch reads the weight tile half as
8193/// often; the price is eight accumulators live at once. Measured 9.0 ->
8194/// 8.3 ms at 9216x2304, b=296 on a 48-thread EPYC 9B45.
8195#[cfg(target_arch = "x86_64")]
8196#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
8197unsafe fn dot_q4tp_row_1x8_avx512(
8198    nib: &[u8],
8199    r: usize,
8200    gpr: usize,
8201    xs: [&[i8]; 8],
8202    scales: &[f32],
8203) -> [f32; 8] {
8204    // SAFETY: as dot_q4tp_row_1x4_avx2; caller guarantees the ISA.
8205    unsafe {
8206        use core::arch::x86_64::*;
8207        let lomask = _mm256_set1_epi8(0x0F);
8208        let eight = _mm256_set1_epi8(8);
8209        let zero = _mm512_setzero_si512();
8210        let (mut v0, mut v1, mut v2, mut v3) = (
8211            _mm512_setzero_ps(),
8212            _mm512_setzero_ps(),
8213            _mm512_setzero_ps(),
8214            _mm512_setzero_ps(),
8215        );
8216        let (mut v4, mut v5, mut v6, mut v7) = (
8217            _mm512_setzero_ps(),
8218            _mm512_setzero_ps(),
8219            _mm512_setzero_ps(),
8220            _mm512_setzero_ps(),
8221        );
8222        let pairs = gpr / 2;
8223        for gp in 0..pairs {
8224            let gi = gp * 2;
8225            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
8226            let bb = _mm256_loadu_si256(t as *const __m256i);
8227            let lo = _mm256_and_si256(bb, lomask);
8228            let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
8229            // `unpack` works per 128-bit lane, so the halves come out as
8230            // [A.lo, B.lo] and [A.hi, B.hi]; the shuffle reorders the four
8231            // 128-bit lanes into the weights' natural order, which is what
8232            // the straight activation load expects.
8233            let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
8234            let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
8235            let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
8236            let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
8237            let wabs = _mm512_abs_epi8(w);
8238            let neg = _mm512_movepi8_mask(w);
8239            let off = gi * GROUP_SIZE;
8240            let sv = _mm512_insertf32x8::<1>(
8241                _mm512_castps256_ps512(_mm256_set1_ps(*scales.get_unchecked(gi))),
8242                _mm256_set1_ps(*scales.get_unchecked(gi + 1)),
8243            );
8244            let dot = |x: &[i8]| -> __m512 {
8245                let xv = _mm512_loadu_si512(x.as_ptr().add(off) as *const __m512i);
8246                let sx = _mm512_mask_sub_epi8(xv, neg, zero, xv);
8247                _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(zero, wabs, sx))
8248            };
8249            v0 = _mm512_fmadd_ps(dot(xs[0]), sv, v0);
8250            v1 = _mm512_fmadd_ps(dot(xs[1]), sv, v1);
8251            v2 = _mm512_fmadd_ps(dot(xs[2]), sv, v2);
8252            v3 = _mm512_fmadd_ps(dot(xs[3]), sv, v3);
8253            v4 = _mm512_fmadd_ps(dot(xs[4]), sv, v4);
8254            v5 = _mm512_fmadd_ps(dot(xs[5]), sv, v5);
8255            v6 = _mm512_fmadd_ps(dot(xs[6]), sv, v6);
8256            v7 = _mm512_fmadd_ps(dot(xs[7]), sv, v7);
8257        }
8258        let mut acc = [
8259            _mm512_reduce_add_ps(v0),
8260            _mm512_reduce_add_ps(v1),
8261            _mm512_reduce_add_ps(v2),
8262            _mm512_reduce_add_ps(v3),
8263            _mm512_reduce_add_ps(v4),
8264            _mm512_reduce_add_ps(v5),
8265            _mm512_reduce_add_ps(v6),
8266            _mm512_reduce_add_ps(v7),
8267        ];
8268        // An odd group count leaves one group over; the narrow kernel
8269        // finishes it rather than the tail being a special case here.
8270        if gpr % 2 == 1 {
8271            let off = (gpr - 1) * GROUP_SIZE;
8272            for j in off..off + GROUP_SIZE {
8273                let (w, s) = q4tp_outlier(nib, r, gpr, j, scales);
8274                let ws = w * s;
8275                for k in 0..8 {
8276                    acc[k] += ws * *xs[k].get_unchecked(j) as f32;
8277                }
8278            }
8279        }
8280        acc
8281    }
8282}
8283
8284/// The same four columns, 512 bits wide. Two groups (64 weights) ride one
8285/// unpack and one `vpdpbusd`, where AVX2 needs two unpacks and four
8286/// `maddubs`/`madd` pairs — about 2.3x fewer instructions for the same
8287/// arithmetic. The two groups carry different scales, so the fma takes a
8288/// vector whose halves hold each group's scale rather than a broadcast.
8289///
8290/// There is no 512-bit `vpsignb`, so the activation's sign is applied by
8291/// negating under a mask taken from the weight's sign bits. That mask is
8292/// per-tile, so it is hoisted out of the column loop and the per-column
8293/// cost stays exactly one instruction, as with `sign_epi8`. Weights of
8294/// zero are not zeroed by the mask trick and do not need to be: their
8295/// magnitude is zero, so the product is.
8296#[cfg(target_arch = "x86_64")]
8297#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
8298unsafe fn dot_q4tp_row_1x4_avx512(
8299    nib: &[u8],
8300    r: usize,
8301    gpr: usize,
8302    xs: [&[i8]; 4],
8303    scales: &[f32],
8304) -> [f32; 4] {
8305    // SAFETY: as dot_q4tp_row_1x4_avx2; caller guarantees the ISA.
8306    unsafe {
8307        use core::arch::x86_64::*;
8308        let lomask = _mm256_set1_epi8(0x0F);
8309        let eight = _mm256_set1_epi8(8);
8310        let zero = _mm512_setzero_si512();
8311        let (mut v0, mut v1, mut v2, mut v3) = (
8312            _mm512_setzero_ps(),
8313            _mm512_setzero_ps(),
8314            _mm512_setzero_ps(),
8315            _mm512_setzero_ps(),
8316        );
8317        let pairs = gpr / 2;
8318        for gp in 0..pairs {
8319            let gi = gp * 2;
8320            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
8321            let bb = _mm256_loadu_si256(t as *const __m256i);
8322            let lo = _mm256_and_si256(bb, lomask);
8323            let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
8324            // `unpack` works per 128-bit lane, so the halves come out as
8325            // [A.lo, B.lo] and [A.hi, B.hi]; the shuffle reorders the four
8326            // 128-bit lanes into the weights' natural order, which is what
8327            // the straight activation load expects.
8328            let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
8329            let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
8330            let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
8331            let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
8332            let wabs = _mm512_abs_epi8(w);
8333            let neg = _mm512_movepi8_mask(w);
8334            let off = gi * GROUP_SIZE;
8335            let sv = _mm512_insertf32x8::<1>(
8336                _mm512_castps256_ps512(_mm256_set1_ps(*scales.get_unchecked(gi))),
8337                _mm256_set1_ps(*scales.get_unchecked(gi + 1)),
8338            );
8339            let dot = |x: &[i8]| -> __m512 {
8340                let xv = _mm512_loadu_si512(x.as_ptr().add(off) as *const __m512i);
8341                let sx = _mm512_mask_sub_epi8(xv, neg, zero, xv);
8342                _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(zero, wabs, sx))
8343            };
8344            v0 = _mm512_fmadd_ps(dot(xs[0]), sv, v0);
8345            v1 = _mm512_fmadd_ps(dot(xs[1]), sv, v1);
8346            v2 = _mm512_fmadd_ps(dot(xs[2]), sv, v2);
8347            v3 = _mm512_fmadd_ps(dot(xs[3]), sv, v3);
8348        }
8349        let mut acc = [
8350            _mm512_reduce_add_ps(v0),
8351            _mm512_reduce_add_ps(v1),
8352            _mm512_reduce_add_ps(v2),
8353            _mm512_reduce_add_ps(v3),
8354        ];
8355        // An odd group count leaves one group over; the narrow kernel
8356        // finishes it rather than the tail being a special case here.
8357        if gpr % 2 == 1 {
8358            let off = (gpr - 1) * GROUP_SIZE;
8359            for j in off..off + GROUP_SIZE {
8360                let (w, s) = q4tp_outlier(nib, r, gpr, j, scales);
8361                let ws = w * s;
8362                for k in 0..4 {
8363                    acc[k] += ws * *xs[k].get_unchecked(j) as f32;
8364                }
8365            }
8366        }
8367        acc
8368    }
8369}
8370
8371/// Four batch columns against one q4tp row: the tile is unpacked ONCE and
8372/// spent on four activation streams, which is where a prefill batch stops
8373/// being weight-bandwidth-bound. Twin of `dot_q4t_row_1x4_sdot`.
8374#[cfg(target_arch = "aarch64")]
8375#[target_feature(enable = "neon,dotprod")]
8376unsafe fn dot_q4tp_row_1x4_sdot(
8377    nib: &[u8],
8378    r: usize,
8379    gpr: usize,
8380    xs: [&[i8]; 4],
8381    scales: &[f32],
8382) -> [f32; 4] {
8383    // SAFETY: see dot_q4tp_row_sdot; every xs[k] is gpr·GROUP_SIZE long.
8384    unsafe {
8385        use core::arch::aarch64::*;
8386        use core::arch::asm;
8387        let lomask = vdupq_n_u8(0x0F);
8388        let eight = vdupq_n_s8(8);
8389        // Named accumulators, NOT an array indexed by a loop variable: the
8390        // latter does not stay in registers (the same defect cost 2x in the
8391        // AVX2 q4t kernel and again in WGSL).
8392        //
8393        // They are VECTORS, and the horizontal add happens once at the end
8394        // instead of once per group per column. `vaddvq` is a cross-lane
8395        // reduction — with 72 groups and four columns the old shape paid
8396        // 288 of them per row, each one a dependency stall the pipeline
8397        // cannot hide, to save four float adds. The group's scale now
8398        // rides an fma into the lane accumulators, so the arithmetic per
8399        // group is one convert and one fma. Summation order changes (the
8400        // lanes carry independent partial sums), which is the same
8401        // round-off class the SDOT path already lives in — the strict
8402        // kernel (`CMF_SDOT=0`, what `cortiq ppl` runs) is unchanged and
8403        // stays the reference.
8404        let (mut v0, mut v1, mut v2, mut v3) = (
8405            vdupq_n_f32(0.0),
8406            vdupq_n_f32(0.0),
8407            vdupq_n_f32(0.0),
8408            vdupq_n_f32(0.0),
8409        );
8410        for gi in 0..gpr {
8411            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
8412            let s = *scales.get_unchecked(gi);
8413            let bb = vld1q_u8(t);
8414            let lo = vandq_u8(bb, lomask);
8415            let hi = vshrq_n_u8::<4>(bb);
8416            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
8417            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
8418            let off = gi * GROUP_SIZE;
8419            let dot4 = |x: &[i8]| -> int32x4_t {
8420                let x0 = vld1q_s8(x.as_ptr().add(off));
8421                let x1 = vld1q_s8(x.as_ptr().add(off + 16));
8422                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
8423                asm!(
8424                    "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
8425                    "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
8426                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
8427                    e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
8428                    options(pure, nomem, nostack),
8429                );
8430                vaddq_s32(a0, a1)
8431            };
8432            v0 = vfmaq_n_f32(v0, vcvtq_f32_s32(dot4(xs[0])), s);
8433            v1 = vfmaq_n_f32(v1, vcvtq_f32_s32(dot4(xs[1])), s);
8434            v2 = vfmaq_n_f32(v2, vcvtq_f32_s32(dot4(xs[2])), s);
8435            v3 = vfmaq_n_f32(v3, vcvtq_f32_s32(dot4(xs[3])), s);
8436        }
8437        [
8438            vaddvq_f32(v0),
8439            vaddvq_f32(v1),
8440            vaddvq_f32(v2),
8441            vaddvq_f32(v3),
8442        ]
8443    }
8444}
8445
8446/// Fused q4tp matmat — the same three arms `q4t_matmat` has. Shipping only
8447/// the scalar one made Nanbeige-3B decode at 1.2 tok/s against q4t's 5.9:
8448/// the format was fine, the missing arms were the whole regression.
8449///
8450/// Under `row_exact()` every cell is summed in `q4tp_matvec`'s order on
8451/// every architecture; the mode is read once, so one call never mixes
8452/// the two contracts when a concurrent scope opens or closes mid-call.
8453fn q4tp_matmat(
8454    bytes: &[u8],
8455    xs_all: &[f32],
8456    b: usize,
8457    rows: usize,
8458    cols: usize,
8459    out: &mut [f32],
8460    pool: Option<&Pool>,
8461) {
8462    q4tp_matmat_with(bytes, xs_all, b, rows, cols, out, pool, row_exact())
8463}
8464
8465/// `q4tp_matmat` with the row-exact mode passed in rather than read from
8466/// the shared counter: `exact` = each cell equals its token's matvec bit
8467/// for bit, otherwise the fast blocked / AMX arms are free to reorder.
8468#[allow(clippy::too_many_arguments)]
8469fn q4tp_matmat_with(
8470    bytes: &[u8],
8471    xs_all: &[f32],
8472    b: usize,
8473    rows: usize,
8474    cols: usize,
8475    out: &mut [f32],
8476    pool: Option<&Pool>,
8477    exact: bool,
8478) {
8479    debug_assert_eq!(out.len(), b * rows);
8480    let gpr = cols / GROUP_SIZE;
8481    let v = Q4tpView::new(bytes, rows, cols);
8482
8483    // Wide batches ride the AMX through a dequant-tile sgemm, as in q4t.
8484    // An f32 GEMM over dequantized weights is not the matvec's int8 sum,
8485    // so a row-exact batch never takes it.
8486    #[cfg(target_os = "macos")]
8487    if !exact && b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
8488        dequant_matmat_accel(
8489            &|r, dst| {
8490                let mut sc = [0f32; 32];
8491                let mut scv;
8492                let s: &[f32] = if gpr <= 32 {
8493                    v.scales_into(r, gpr, &mut sc);
8494                    &sc[..gpr]
8495                } else {
8496                    scv = vec![0f32; gpr];
8497                    v.scales_into(r, gpr, &mut scv);
8498                    &scv
8499                };
8500                for gi in 0..gpr {
8501                    let tile = &v.nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
8502                    for (k, &bb) in tile.iter().enumerate() {
8503                        dst[gi * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s[gi];
8504                        dst[gi * GROUP_SIZE + k * 2 + 1] =
8505                            (((bb >> 4) & 0x0F) as f32 - 8.0) * s[gi];
8506                    }
8507                }
8508            },
8509            xs_all,
8510            b,
8511            rows,
8512            cols,
8513            out,
8514            pool,
8515        );
8516        return;
8517    }
8518
8519    let out_addr = SendMut(out.as_mut_ptr());
8520    if a8w8_enabled() {
8521        let acts = split_acts_pooled(xs_all, b, cols, pool);
8522        let acts = &acts;
8523        // ARM stays blocked under `exact` too: the 1x4 kernel then runs in
8524        // its v1 shape, whose per-group `int dot as f32 * scale` and scalar
8525        // running sum are exactly `dot_q4tp_row_sdot`'s, so the tile is
8526        // still unpacked once for four columns and each column equals its
8527        // matvec. The tuned shape (fma into lane partials, one horizontal
8528        // add per row) is 1-16 ulp off the matvec and is kept for `!exact`.
8529        #[cfg(target_arch = "aarch64")]
8530        let blocked_ok = sdot_enabled() && blocked_enabled();
8531        // x86 gets the same blocking: one tile unpack spent on four
8532        // columns. Without it every column re-decoded the row, which is
8533        // why a 48-core EPYC measured a sixth of an M4's per-core rate.
8534        // The gate is `avx2_enabled`, as in q4t — `sdot_enabled` answers
8535        // for ARM's dotprod and is hard-wired false everywhere else, so
8536        // asking it here left the whole blocked path unreachable on x86.
8537        #[cfg(target_arch = "x86_64")]
8538        let blocked_ok = q4tp_blocked_x86() && !exact;
8539        #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
8540        let blocked_ok = {
8541            let _ = exact;
8542            false
8543        };
8544        // Columns are swept in panels that fit L2. Without this a
8545        // row-pair walks every activation in the batch — 4.8 MB at
8546        // 512x512 — and does it again for the next pair, so the whole
8547        // batch streams out of the shared cache once per row. Measured
8548        // 800 GB/s of it, flat across batch sizes, which is the signature
8549        // of a loop bound by traffic rather than by arithmetic. A panel of
8550        // 256 columns is 590 KB beside 221 KB of this worker's weights:
8551        // both stay resident and the batch crosses L3 once instead of
8552        // once per row.
8553        let panel_cols = q4tp_panel_cols(cols);
8554        let run = |start: usize, end: usize| {
8555            for abase in (0..acts.len()).step_by(panel_cols) {
8556                let alen = (acts.len() - abase).min(panel_cols);
8557                let mut sc = vec![0f32; gpr];
8558                #[cfg(target_arch = "x86_64")]
8559                let mut r_lo = start;
8560                #[cfg(target_arch = "x86_64")]
8561                if blocked_ok && alen >= 8 {
8562                    let mut sc1 = vec![0f32; gpr];
8563                    while r_lo + 2 <= end {
8564                        v.scales_into(r_lo, gpr, &mut sc);
8565                        v.scales_into(r_lo + 1, gpr, &mut sc1);
8566                        let mut bi = 0usize;
8567                        while bi + 8 <= alen {
8568                            let xs = [
8569                                acts[abase + bi].xq.as_slice(),
8570                                acts[abase + bi + 1].xq.as_slice(),
8571                                acts[abase + bi + 2].xq.as_slice(),
8572                                acts[abase + bi + 3].xq.as_slice(),
8573                                acts[abase + bi + 4].xq.as_slice(),
8574                                acts[abase + bi + 5].xq.as_slice(),
8575                                acts[abase + bi + 6].xq.as_slice(),
8576                                acts[abase + bi + 7].xq.as_slice(),
8577                            ];
8578                            let d = unsafe { dot_q4tp_2x8_avx512(v.nib, r_lo, gpr, xs, &sc, &sc1) };
8579                            for (row, dr, scr) in [(r_lo, &d[0], &sc), (r_lo + 1, &d[1], &sc1)] {
8580                                for k in 0..8 {
8581                                    let act = &acts[abase + bi + k];
8582                                    let mut acc = dr[k] * act.sx;
8583                                    for &(j, xv) in &act.outliers {
8584                                        let (w, s) = q4tp_outlier(v.nib, row, gpr, j, scr);
8585                                        acc += w * s * xv;
8586                                    }
8587                                    // SAFETY: disjoint (bi, r) cells per worker.
8588                                    unsafe { *out_addr.at((abase + bi + k) * rows + row) = acc };
8589                                }
8590                            }
8591                            bi += 8;
8592                        }
8593                        // Columns past the last group of eight, both rows —
8594                        // the same single-row kernel the tail below uses.
8595                        for row in [r_lo, r_lo + 1] {
8596                            let scr: &[f32] = if row == r_lo { &sc } else { &sc1 };
8597                            for b2 in bi..alen {
8598                                let act = &acts[abase + b2];
8599                                let xs4 = [
8600                                    act.xq.as_slice(),
8601                                    act.xq.as_slice(),
8602                                    act.xq.as_slice(),
8603                                    act.xq.as_slice(),
8604                                ];
8605                                let d =
8606                                    unsafe { dot_q4tp_row_1x4_avx512(v.nib, row, gpr, xs4, scr) };
8607                                let mut acc = d[0] * act.sx;
8608                                for &(j, xv) in &act.outliers {
8609                                    let (w, s) = q4tp_outlier(v.nib, row, gpr, j, scr);
8610                                    acc += w * s * xv;
8611                                }
8612                                // SAFETY: disjoint (bi, r) cells per worker.
8613                                unsafe { *out_addr.at((abase + b2) * rows + row) = acc };
8614                            }
8615                        }
8616                        r_lo += 2;
8617                    }
8618                }
8619                // AVX2 without VNNI, a batch that need not be row-exact:
8620                // two rows against four columns with FMA
8621                // (`q4tp_matmat_2rows_fma`); an odd last row falls to the
8622                // vertical cell below.
8623                #[cfg(target_arch = "x86_64")]
8624                if !blocked_ok
8625                    && !exact
8626                    && alen >= 4
8627                    && q4tp_fma_x86()
8628                    && acts[abase..abase + alen].iter().all(|a| a.xg8.len() == gpr)
8629                {
8630                    let mut sc1 = vec![0f32; gpr];
8631                    while r_lo + 2 <= end {
8632                        v.scales_into(r_lo, gpr, &mut sc);
8633                        v.scales_into(r_lo + 1, gpr, &mut sc1);
8634                        // SAFETY: AVX2+FMA checked; cells disjoint per
8635                        // (worker, bi, r).
8636                        unsafe {
8637                            q4tp_matmat_2rows_fma(
8638                                &v,
8639                                r_lo,
8640                                gpr,
8641                                &acts[abase..abase + alen],
8642                                &sc,
8643                                &sc1,
8644                                abase,
8645                                rows,
8646                                out_addr,
8647                            )
8648                        };
8649                        r_lo += 2;
8650                    }
8651                }
8652                #[cfg(target_arch = "x86_64")]
8653                let row_start = r_lo;
8654                #[cfg(not(target_arch = "x86_64"))]
8655                let row_start = start;
8656                for r in row_start..end {
8657                    v.scales_into(r, gpr, &mut sc);
8658                    let mut bi = 0usize;
8659                    #[cfg(target_arch = "x86_64")]
8660                    if blocked_ok {
8661                        while bi + 8 <= alen {
8662                            let xs = [
8663                                acts[abase + bi].xq.as_slice(),
8664                                acts[abase + bi + 1].xq.as_slice(),
8665                                acts[abase + bi + 2].xq.as_slice(),
8666                                acts[abase + bi + 3].xq.as_slice(),
8667                                acts[abase + bi + 4].xq.as_slice(),
8668                                acts[abase + bi + 5].xq.as_slice(),
8669                                acts[abase + bi + 6].xq.as_slice(),
8670                                acts[abase + bi + 7].xq.as_slice(),
8671                            ];
8672                            let d = unsafe { dot_q4tp_row_1x8_avx512(v.nib, r, gpr, xs, &sc) };
8673                            for k in 0..8 {
8674                                let act = &acts[abase + bi + k];
8675                                let mut acc = d[k] * act.sx;
8676                                for &(j, xv) in &act.outliers {
8677                                    let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
8678                                    acc += w * s * xv;
8679                                }
8680                                // SAFETY: disjoint (bi, r) cells per worker.
8681                                unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
8682                            }
8683                            bi += 8;
8684                        }
8685                        while bi + 4 <= alen {
8686                            let xs = [
8687                                acts[abase + bi].xq.as_slice(),
8688                                acts[abase + bi + 1].xq.as_slice(),
8689                                acts[abase + bi + 2].xq.as_slice(),
8690                                acts[abase + bi + 3].xq.as_slice(),
8691                            ];
8692                            let d = unsafe { dot_q4tp_row_1x4_avx512(v.nib, r, gpr, xs, &sc) };
8693                            for k in 0..4 {
8694                                let act = &acts[abase + bi + k];
8695                                let mut acc = d[k] * act.sx;
8696                                for &(j, xv) in &act.outliers {
8697                                    let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
8698                                    acc += w * s * xv;
8699                                }
8700                                // SAFETY: disjoint (bi, r) cells per worker.
8701                                unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
8702                            }
8703                            bi += 4;
8704                        }
8705                    }
8706                    #[cfg(target_arch = "aarch64")]
8707                    if blocked_ok {
8708                        while bi + 4 <= alen {
8709                            let xs = [
8710                                acts[abase + bi].xq.as_slice(),
8711                                acts[abase + bi + 1].xq.as_slice(),
8712                                acts[abase + bi + 2].xq.as_slice(),
8713                                acts[abase + bi + 3].xq.as_slice(),
8714                            ];
8715                            let d = unsafe {
8716                                if exact || q4tp_v1() {
8717                                    dot_q4tp_row_1x4_sdot_v1(v.nib, r, gpr, xs, &sc)
8718                                } else {
8719                                    dot_q4tp_row_1x4_sdot(v.nib, r, gpr, xs, &sc)
8720                                }
8721                            };
8722                            for k in 0..4 {
8723                                let act = &acts[abase + bi + k];
8724                                let mut acc = d[k] * act.sx;
8725                                for &(j, xv) in &act.outliers {
8726                                    let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
8727                                    acc += w * s * xv;
8728                                }
8729                                // SAFETY: disjoint (bi, r) cells per worker.
8730                                unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
8731                            }
8732                            bi += 4;
8733                        }
8734                    }
8735                    let _ = blocked_ok;
8736                    // AVX2: the vertical column panel shares each group's
8737                    // unpack across four columns and keeps the row-exact
8738                    // cell bits (see `q4tp_matmat_row_cols_v`). Falls
8739                    // through to the single-column loop only when a column
8740                    // lacks its per-group correction.
8741                    #[cfg(target_arch = "x86_64")]
8742                    if acts[abase..abase + alen]
8743                        .iter()
8744                        .all(|a| a.xg8.len() == gpr)
8745                    {
8746                        // SAFETY: a8w8 implies AVX2; cells stay disjoint
8747                        // per (worker, bi, r).
8748                        unsafe {
8749                            q4tp_matmat_row_cols_v(
8750                                &v,
8751                                r,
8752                                gpr,
8753                                &acts[abase + bi..abase + alen],
8754                                &sc,
8755                                abase + bi,
8756                                rows,
8757                                out_addr,
8758                            )
8759                        };
8760                        bi = alen;
8761                    }
8762                    while bi < alen {
8763                        let act = &acts[abase + bi];
8764                        let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, &sc) * act.sx;
8765                        for &(j, xv) in &act.outliers {
8766                            let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
8767                            acc += w * s * xv;
8768                        }
8769                        // SAFETY: disjoint (bi, r) cells per worker range.
8770                        unsafe { *out_addr.at((abase + bi) * rows + r) = acc };
8771                        bi += 1;
8772                    }
8773                }
8774            }
8775        };
8776        dispatch_rows(pool, rows, &run);
8777        return;
8778    }
8779
8780    let run = |start: usize, end: usize| {
8781        let mut sc = vec![0f32; gpr];
8782        for r in start..end {
8783            v.scales_into(r, gpr, &mut sc);
8784            for bi in 0..b {
8785                let x = &xs_all[bi * cols..(bi + 1) * cols];
8786                // SAFETY: disjoint (bi, r) cells per worker range.
8787                unsafe { *out_addr.at(bi * rows + r) = q4tp_row_exact(v.nib, r, gpr, x, &sc) };
8788            }
8789        }
8790    };
8791    dispatch_rows(pool, rows, &run);
8792}
8793
8794/// Fused q4_tiled matvec (dispatch mirrors `q4matvec`).
8795fn q4t_matvec(
8796    bytes: &[u8],
8797    x: &[f32],
8798    rows: usize,
8799    cols: usize,
8800    out: &mut [f32],
8801    pool: Option<&Pool>,
8802) {
8803    debug_assert_eq!(out.len(), rows);
8804    let gpr = cols / GROUP_SIZE;
8805    let out_addr = SendMut(out.as_mut_ptr());
8806    if a8w8_enabled() {
8807        let act = split_act(x);
8808        let run = move |start: usize, end: usize| {
8809            for r in start..end {
8810                let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
8811                for &(j, xv) in &act.outliers {
8812                    let (w, s) = q4t_outlier(bytes, r, gpr, j);
8813                    acc += w * s * xv;
8814                }
8815                // SAFETY: disjoint row ranges per worker.
8816                unsafe { *out_addr.at(r) = acc };
8817            }
8818        };
8819        dispatch_rows(pool, rows, &run);
8820        return;
8821    }
8822    let run = move |start: usize, end: usize| {
8823        for r in start..end {
8824            // SAFETY: disjoint row ranges per worker.
8825            unsafe { *out_addr.at(r) = q4t_row_exact(bytes, r, gpr, x) };
8826        }
8827    };
8828    dispatch_rows(pool, rows, &run);
8829}
8830
8831/// Fused two-input q4_tiled matvec (weights read once per pair).
8832#[allow(clippy::too_many_arguments)]
8833fn q4t_matvec2(
8834    bytes: &[u8],
8835    x1: &[f32],
8836    x2: &[f32],
8837    rows: usize,
8838    cols: usize,
8839    o1: &mut [f32],
8840    o2: &mut [f32],
8841    pool: Option<&Pool>,
8842) {
8843    let gpr = cols / GROUP_SIZE;
8844    let p1 = SendMut(o1.as_mut_ptr());
8845    let p2 = SendMut(o2.as_mut_ptr());
8846    if a8w8_enabled() {
8847        let a1 = split_act(x1);
8848        let a2 = split_act(x2);
8849        let run = move |start: usize, end: usize| {
8850            for r in start..end {
8851                let mut v1 = dot_q4t_row_i8(bytes, r, gpr, &a1.xq) * a1.sx;
8852                let mut v2 = dot_q4t_row_i8(bytes, r, gpr, &a2.xq) * a2.sx;
8853                for &(j, xv) in &a1.outliers {
8854                    let (w, s) = q4t_outlier(bytes, r, gpr, j);
8855                    v1 += w * s * xv;
8856                }
8857                for &(j, xv) in &a2.outliers {
8858                    let (w, s) = q4t_outlier(bytes, r, gpr, j);
8859                    v2 += w * s * xv;
8860                }
8861                // SAFETY: disjoint row ranges per worker.
8862                unsafe {
8863                    *p1.at(r) = v1;
8864                    *p2.at(r) = v2;
8865                }
8866            }
8867        };
8868        dispatch_rows(pool, rows, &run);
8869        return;
8870    }
8871    let run = move |start: usize, end: usize| {
8872        for r in start..end {
8873            // SAFETY: disjoint row ranges per worker.
8874            unsafe {
8875                *p1.at(r) = q4t_row_exact(bytes, r, gpr, x1);
8876                *p2.at(r) = q4t_row_exact(bytes, r, gpr, x2);
8877            }
8878        }
8879    };
8880    dispatch_rows(pool, rows, &run);
8881}
8882
8883/// Batched q4_tiled matmat: each row's tiles stream once per microbatch.
8884#[allow(clippy::too_many_arguments)]
8885/// Prefill GEMM through Accelerate for group-quantized codecs: a
8886/// caller-supplied row dequantizer fills f32 tiles (pool-parallel) and
8887/// each tile rides the AMX with one sgemm — the generic sibling of
8888/// `qmatmat_accel` (q8). Numerics are f32-GEMM (tolerance class);
8889/// decode (b=1) never takes this path.
8890#[cfg(target_os = "macos")]
8891fn dequant_matmat_accel(
8892    dequant_row: &(dyn Fn(usize, &mut [f32]) + Sync),
8893    xs_all: &[f32],
8894    b: usize,
8895    rows: usize,
8896    cols: usize,
8897    out: &mut [f32],
8898    pool: Option<&Pool>,
8899) {
8900    const TR: usize = 2048;
8901    thread_local! {
8902        static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
8903    }
8904    WTILE.with(|wt| {
8905        let mut wtile = wt.borrow_mut();
8906        wtile.resize(TR * cols, 0.0);
8907        let mut r0 = 0usize;
8908        while r0 < rows {
8909            let tr = TR.min(rows - r0);
8910            let wt_addr = SendMut(wtile.as_mut_ptr());
8911            let run = |start: usize, end: usize| {
8912                for r in start..end {
8913                    // SAFETY: workers cover disjoint r ranges.
8914                    let dst = unsafe { std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols) };
8915                    dequant_row(r0 + r, dst);
8916                }
8917            };
8918            dispatch_rows(pool, tr, &run);
8919            unsafe {
8920                accel_blas::cblas_sgemm(
8921                    101, // RowMajor
8922                    111, // NoTrans A
8923                    112, // Trans B
8924                    b as i32,
8925                    tr as i32,
8926                    cols as i32,
8927                    1.0,
8928                    xs_all.as_ptr(),
8929                    cols as i32,
8930                    wtile.as_ptr(),
8931                    cols as i32,
8932                    0.0,
8933                    out.as_mut_ptr().add(r0),
8934                    rows as i32,
8935                );
8936            }
8937            r0 += tr;
8938        }
8939    });
8940}
8941
8942fn q4t_matmat(
8943    bytes: &[u8],
8944    xs_all: &[f32],
8945    b: usize,
8946    rows: usize,
8947    cols: usize,
8948    out: &mut [f32],
8949    pool: Option<&Pool>,
8950) {
8951    debug_assert_eq!(out.len(), b * rows);
8952    let gpr = cols / GROUP_SIZE;
8953    // Wide batches ride the AMX like q8's qmatmat: on Apple silicon
8954    // the dequant-tile sgemm is an order above the SDOT row loop for
8955    // prefill shapes (imagegen DiT forwards are exactly this).
8956    #[cfg(target_os = "macos")]
8957    if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
8958        dequant_matmat_accel(
8959            &|r, dst| {
8960                for gi in 0..gpr {
8961                    let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
8962                    let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
8963                    for (k, &bb) in tile[2..].iter().enumerate() {
8964                        dst[gi * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s;
8965                        dst[gi * GROUP_SIZE + k * 2 + 1] = (((bb >> 4) & 0x0F) as f32 - 8.0) * s;
8966                    }
8967                }
8968            },
8969            xs_all,
8970            b,
8971            rows,
8972            cols,
8973            out,
8974            pool,
8975        );
8976        return;
8977    }
8978    let out_addr = SendMut(out.as_mut_ptr());
8979    if a8w8_enabled() {
8980        let acts: Vec<SplitAct> = (0..b)
8981            .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
8982            .collect();
8983        let acts = &acts;
8984        #[cfg(target_arch = "x86_64")]
8985        let blocked_ok = avx2_enabled() && blocked_enabled();
8986        #[cfg(target_arch = "aarch64")]
8987        let blocked_ok = sdot_enabled() && blocked_enabled();
8988        #[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))]
8989        let blocked_ok = false;
8990        let run = move |start: usize, end: usize| {
8991            for r in start..end {
8992                let mut bi = 0usize;
8993                #[cfg(target_arch = "aarch64")]
8994                if blocked_ok {
8995                    while bi + 4 <= acts.len() {
8996                        let xs = [
8997                            acts[bi].xq.as_slice(),
8998                            acts[bi + 1].xq.as_slice(),
8999                            acts[bi + 2].xq.as_slice(),
9000                            acts[bi + 3].xq.as_slice(),
9001                        ];
9002                        let d = unsafe { dot_q4t_row_1x4_sdot(bytes, r, gpr, xs) };
9003                        for k in 0..4 {
9004                            let act = &acts[bi + k];
9005                            let mut acc = d[k] * act.sx;
9006                            for &(j, xv) in &act.outliers {
9007                                let (w, sc) = q4t_outlier(bytes, r, gpr, j);
9008                                acc += w * sc * xv;
9009                            }
9010                            // SAFETY: disjoint (bi, r) cells per worker.
9011                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
9012                        }
9013                        bi += 4;
9014                    }
9015                }
9016                #[cfg(target_arch = "x86_64")]
9017                if blocked_ok {
9018                    while bi + 4 <= acts.len() {
9019                        let xs = [
9020                            acts[bi].xq.as_slice(),
9021                            acts[bi + 1].xq.as_slice(),
9022                            acts[bi + 2].xq.as_slice(),
9023                            acts[bi + 3].xq.as_slice(),
9024                        ];
9025                        let d = unsafe {
9026                            if vnni_tiles_enabled() {
9027                                dot_q4t_row_1x4_vnni(bytes, r, gpr, xs)
9028                            } else {
9029                                dot_q4t_row_1x4_avx2(bytes, r, gpr, xs)
9030                            }
9031                        };
9032                        for k in 0..4 {
9033                            let act = &acts[bi + k];
9034                            let mut acc = d[k] * act.sx;
9035                            for &(j, xv) in &act.outliers {
9036                                let (w, sc) = q4t_outlier(bytes, r, gpr, j);
9037                                acc += w * sc * xv;
9038                            }
9039                            // SAFETY: disjoint (bi, r) cells per worker.
9040                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
9041                        }
9042                        bi += 4;
9043                    }
9044                }
9045                let _ = blocked_ok;
9046                while bi < acts.len() {
9047                    let act = &acts[bi];
9048                    let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
9049                    for &(j, xv) in &act.outliers {
9050                        let (w, s) = q4t_outlier(bytes, r, gpr, j);
9051                        acc += w * s * xv;
9052                    }
9053                    // SAFETY: disjoint (bi, r) cells per worker range.
9054                    unsafe { *out_addr.at(bi * rows + r) = acc };
9055                    bi += 1;
9056                }
9057            }
9058        };
9059        dispatch_rows(pool, rows, &run);
9060        return;
9061    }
9062    let run = move |start: usize, end: usize| {
9063        for r in start..end {
9064            for bi in 0..b {
9065                let x = &xs_all[bi * cols..(bi + 1) * cols];
9066                // SAFETY: disjoint (bi, r) cells per worker range.
9067                unsafe { *out_addr.at(bi * rows + r) = q4t_row_exact(bytes, r, gpr, x) };
9068            }
9069        }
9070    };
9071    dispatch_rows(pool, rows, &run);
9072}
9073
9074// ── q1 (dtype 12): binary weights, [f16 scale][4B sign bits] per
9075// 32-group tile. The kernel family mirrors q4_tiled: one sequential
9076// stream of 6-byte tiles, per-tile integer dot × scale, exact outlier
9077// correction (A8W8 contract), exact scalar path under CMF_SDOT=0. ──
9078
9079/// Per-32-group sums of the quantized activation — the ±1 identity's
9080/// shared half: `dot = −2·sdot(mask, x) − gsum[g]`, computed ONCE per
9081/// matvec and reused by every row.
9082fn q1_group_sums(xq: &[i8], gpr: usize) -> Vec<i32> {
9083    (0..gpr)
9084        .map(|gi| {
9085            xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE]
9086                .iter()
9087                .map(|&v| v as i32)
9088                .sum()
9089        })
9090        .collect()
9091}
9092
9093/// One q1 row via the A8W8 int8 path — mask-SDOT on ARM (no ±1
9094/// expansion at all), scalar bit loop elsewhere (AVX2 queued with the
9095/// x86 pass).
9096#[inline]
9097#[allow(unreachable_code)]
9098/// AVX2 q1 row via the same ±1 identity as the ARM sdot kernel: the
9099/// sign bits expand to a {0, −1} byte mask through shuffle+cmpeq, the
9100/// masked activation sums through maddubs(1, x&mask), and
9101/// `dot = −(2·masked_sum + Σx_group)` — bit-identical integer math.
9102#[cfg(target_arch = "x86_64")]
9103#[target_feature(enable = "avx2")]
9104unsafe fn dot_q1_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
9105    // SAFETY: callers uphold the 6B-tile and xq/gsum length contracts.
9106    unsafe {
9107        use core::arch::x86_64::*;
9108        // Byte j of the mask must replicate bits-byte j/8.
9109        let expand = _mm256_setr_epi8(
9110            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,
9111            3, 3, 3,
9112        );
9113        let bitsel = _mm256_setr_epi8(
9114            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
9115            -128, 1, 2, 4, 8, 16, 32, 64, -128,
9116        );
9117        let ones8 = _mm256_set1_epi8(1);
9118        let ones16 = _mm256_set1_epi16(1);
9119        let mut acc = 0f32;
9120        for gi in 0..gpr {
9121            let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
9122            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
9123            let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
9124            let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
9125            let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
9126            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9127            let sel = _mm256_and_si256(x, mask);
9128            // Σ of selected i8 lanes: maddubs(1u8, sel_i8) pairs → madd.
9129            let p16 = _mm256_maddubs_epi16(ones8, sel);
9130            let d32 = _mm256_madd_epi16(p16, ones16);
9131            let hi128 = _mm256_extracti128_si256::<1>(d32);
9132            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
9133            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
9134            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
9135            let msum = _mm_cvtsi128_si32(s32);
9136            // The and-select keeps x UN-negated (unlike ARM's −1-mask
9137            // sdot): d = Σ_set − Σ_unset = 2·Σ_set − Σ_all.
9138            let d = 2 * msum - gsum[gi];
9139            acc += d as f32 * s;
9140        }
9141        acc
9142    }
9143}
9144
9145/// VNNI twin of `dot_q1_row_avx2`: the masked-select sum goes through
9146/// one `vpdpbusd(1u8, sel)` (see `dpbusd_hsum` — bit-identical).
9147#[cfg(target_arch = "x86_64")]
9148#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
9149unsafe fn dot_q1_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
9150    // SAFETY: callers uphold the 6B-tile and xq/gsum length contracts.
9151    unsafe {
9152        use core::arch::x86_64::*;
9153        let expand = _mm256_setr_epi8(
9154            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,
9155            3, 3, 3,
9156        );
9157        let bitsel = _mm256_setr_epi8(
9158            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
9159            -128, 1, 2, 4, 8, 16, 32, 64, -128,
9160        );
9161        let ones8 = _mm256_set1_epi8(1);
9162        let mut acc = 0f32;
9163        for gi in 0..gpr {
9164            let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
9165            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
9166            let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
9167            let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
9168            let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
9169            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9170            let msum = dpbusd_hsum(ones8, _mm256_and_si256(x, mask));
9171            let d = 2 * msum - gsum[gi];
9172            acc += d as f32 * s;
9173        }
9174        acc
9175    }
9176}
9177
9178/// VNNI twin of `dot_q1_row_1x4_avx2` (see `dpbusd_hsum`).
9179#[cfg(target_arch = "x86_64")]
9180#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
9181unsafe fn dot_q1_row_1x4_vnni(
9182    bytes: &[u8],
9183    r: usize,
9184    gpr: usize,
9185    xs: [&[i8]; 4],
9186    gsums: [&[i32]; 4],
9187) -> [f32; 4] {
9188    // SAFETY: callers uphold the 6B-tile and xq/gsum length contracts.
9189    unsafe {
9190        use core::arch::x86_64::*;
9191        let expand = _mm256_setr_epi8(
9192            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,
9193            3, 3, 3,
9194        );
9195        let bitsel = _mm256_setr_epi8(
9196            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
9197            -128, 1, 2, 4, 8, 16, 32, 64, -128,
9198        );
9199        let ones8 = _mm256_set1_epi8(1);
9200        let mut acc = [0f32; 4];
9201        for gi in 0..gpr {
9202            let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
9203            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
9204            let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
9205            let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
9206            let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
9207            for (k, xq) in xs.iter().enumerate() {
9208                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9209                let msum = dpbusd_hsum(ones8, _mm256_and_si256(x, mask));
9210                let d = 2 * msum - gsums[k][gi];
9211                acc[k] += d as f32 * s;
9212            }
9213        }
9214        acc
9215    }
9216}
9217
9218/// The blocked 1×4 flavor: the expanded bit mask serves four activation
9219/// streams per group (mask build once, four select+reduce chains).
9220#[cfg(target_arch = "x86_64")]
9221#[target_feature(enable = "avx2")]
9222unsafe fn dot_q1_row_1x4_avx2(
9223    bytes: &[u8],
9224    r: usize,
9225    gpr: usize,
9226    xs: [&[i8]; 4],
9227    gsums: [&[i32]; 4],
9228) -> [f32; 4] {
9229    // SAFETY: callers uphold the 6B-tile and xq/gsum length contracts.
9230    unsafe {
9231        use core::arch::x86_64::*;
9232        let expand = _mm256_setr_epi8(
9233            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,
9234            3, 3, 3,
9235        );
9236        let bitsel = _mm256_setr_epi8(
9237            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
9238            -128, 1, 2, 4, 8, 16, 32, 64, -128,
9239        );
9240        let ones8 = _mm256_set1_epi8(1);
9241        let ones16 = _mm256_set1_epi16(1);
9242        let mut acc = [0f32; 4];
9243        for gi in 0..gpr {
9244            let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
9245            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
9246            let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
9247            let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
9248            let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
9249            for (k, xq) in xs.iter().enumerate() {
9250                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9251                let sel = _mm256_and_si256(x, mask);
9252                let p16 = _mm256_maddubs_epi16(ones8, sel);
9253                let d32 = _mm256_madd_epi16(p16, ones16);
9254                let hi128 = _mm256_extracti128_si256::<1>(d32);
9255                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
9256                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
9257                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
9258                let msum = _mm_cvtsi128_si32(s32);
9259                let d = 2 * msum - gsums[k][gi];
9260                acc[k] += d as f32 * s;
9261            }
9262        }
9263        acc
9264    }
9265}
9266
9267#[allow(unreachable_code)]
9268fn dot_q1_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
9269    #[cfg(target_arch = "aarch64")]
9270    unsafe {
9271        return dot_q1_row_sdot(bytes, r, gpr, xq, gsum);
9272    }
9273    #[cfg(target_arch = "x86_64")]
9274    if avx2_enabled() {
9275        unsafe {
9276            if vnni_tiles_enabled() {
9277                return dot_q1_row_vnni(bytes, r, gpr, xq, gsum);
9278            }
9279            return dot_q1_row_avx2(bytes, r, gpr, xq, gsum);
9280        }
9281    }
9282    let _ = gsum;
9283    let mut acc = 0f32;
9284    for gi in 0..gpr {
9285        let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
9286        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
9287        let mut d = 0i32;
9288        for (j, &b) in tile[2..].iter().enumerate() {
9289            for k in 0..8 {
9290                let w = ((b >> k) & 1) as i32 * 2 - 1;
9291                d += w * xq[gi * GROUP_SIZE + j * 8 + k] as i32;
9292            }
9293        }
9294        acc += d as f32 * s;
9295    }
9296    acc
9297}
9298
9299/// SDOT q1 row via the ±1 identity: the vtst mask (0xFF where the bit
9300/// is set, i.e. −1 as i8) feeds `sdot` DIRECTLY — no expansion to ±1
9301/// lanes at all — and `dot = −(2·sdot(mask, x) + Σx_group)`, with the
9302/// per-group activation sums shared across every row of the matvec.
9303/// Four tiles (128 weights) per iteration: integer dots reduce through
9304/// a vpaddq tree into ONE i32x4 that meets its four scales in a single
9305/// fused f32 multiply-add. Integer math throughout — bit-identical to
9306/// the scalar ±1 reference.
9307#[cfg(target_arch = "aarch64")]
9308#[target_feature(enable = "neon,dotprod")]
9309unsafe fn dot_q1_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
9310    // SAFETY: callers uphold slice-length contracts (6B tile per group,
9311    // xq.len() == gpr·GROUP_SIZE, gsum.len() == gpr).
9312    unsafe {
9313        use core::arch::aarch64::*;
9314        use core::arch::asm;
9315        const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
9316        let m = vld1q_u8(MASKS.as_ptr());
9317        // One tile's −Σ_set(x) as an UNREDUCED i32x4 (two mask-sdots).
9318        macro_rules! tile_dot {
9319            ($t:expr, $x:expr) => {{
9320                let v0 = vcombine_u8(vdup_n_u8(*$t.add(2)), vdup_n_u8(*$t.add(3)));
9321                let v1 = vcombine_u8(vdup_n_u8(*$t.add(4)), vdup_n_u8(*$t.add(5)));
9322                let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
9323                let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
9324                let x0 = vld1q_s8($x);
9325                let x1 = vld1q_s8($x.add(16));
9326                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
9327                asm!(
9328                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
9329                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
9330                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
9331                    w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
9332                    options(pure, nomem, nostack),
9333                );
9334                vaddq_s32(a0, a1)
9335            }};
9336        }
9337        // TBL unpack over PAIR loads: one vld1q covers two 6B tiles
9338        // ([s s b b b b][s s b b b b] + 4B slack), TBL replicates each
9339        // bit-byte across 8 lanes for vtst, and the four scales gather
9340        // through tbl2 into one fcvtl — the 16 ld1r broadcast loads and
9341        // 4 branchy software f16 conversions per 128 weights (the
9342        // measured load-port wall of this kernel) become 2 vector
9343        // loads + 9 table lookups. Integer math order is unchanged —
9344        // bit-identical results (FCVTL is exact on every f16).
9345        const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
9346        const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
9347        const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
9348        const IW11: [u8; 16] = [
9349            10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
9350        ];
9351        const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
9352        let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
9353        let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
9354        let isc = vld1_u8(ISC.as_ptr());
9355        // One tile's −Σ_set(x) from a TBL-unpacked pair load.
9356        macro_rules! tile_dot_tbl {
9357            ($ld:expr, $i0:expr, $i1:expr, $x:expr) => {{
9358                let w0 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i0), m));
9359                let w1 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i1), m));
9360                let x0 = vld1q_s8($x);
9361                let x1 = vld1q_s8($x.add(16));
9362                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
9363                asm!(
9364                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
9365                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
9366                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
9367                    w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
9368                    options(pure, nomem, nostack),
9369                );
9370                vaddq_s32(a0, a1)
9371            }};
9372        }
9373        let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
9374        let row_base = r * gpr * Q1_TILE;
9375        let abs_end = bytes.len();
9376        let xp = xq.as_ptr();
9377        let gp = gsum.as_ptr();
9378        let mut accv = vdupq_n_f32(0.0);
9379        let mut gi = 0;
9380        // The second pair load reads 4B past tile gi+3 — stay inside
9381        // the payload slice (only the file's final tiles fall back).
9382        while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
9383            let t0 = base.add(gi * Q1_TILE);
9384            let ld_a = vld1q_u8(t0);
9385            let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
9386            let d0 = tile_dot_tbl!(ld_a, iw00, iw01, xp.add(gi * GROUP_SIZE));
9387            let d1 = tile_dot_tbl!(ld_a, iw10, iw11, xp.add((gi + 1) * GROUP_SIZE));
9388            let d2 = tile_dot_tbl!(ld_b, iw00, iw01, xp.add((gi + 2) * GROUP_SIZE));
9389            let d3 = tile_dot_tbl!(ld_b, iw10, iw11, xp.add((gi + 3) * GROUP_SIZE));
9390            // [−Σ0, −Σ1, −Σ2, −Σ3] → dots = −(2·Σset_neg + gsum)
9391            let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
9392            let g = vld1q_s32(gp.add(gi));
9393            let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
9394            let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
9395            let scf: float32x4_t;
9396            asm!(
9397                "fcvtl {o:v}.4s, {i:v}.4h",
9398                o = out(vreg) scf, i = in(vreg) sc16,
9399                options(pure, nomem, nostack),
9400            );
9401            accv = vfmaq_f32(accv, vcvtq_f32_s32(dots), scf);
9402            gi += 4;
9403        }
9404        let mut acc = vaddvq_f32(accv);
9405        while gi < gpr {
9406            let t = base.add(gi * Q1_TILE);
9407            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
9408            let d = vaddvq_s32(tile_dot!(t, xp.add(gi * GROUP_SIZE)));
9409            acc += (-(2 * d + *gp.add(gi))) as f32 * s;
9410            gi += 1;
9411        }
9412        acc
9413    }
9414}
9415
9416/// Blocked q1 1×4: one TBL unpack of the tile pair serves FOUR
9417/// activation streams (prefill amortization — the same idea as the
9418/// AVX2 twin; per stream the group order, fma order and tail match the
9419/// single-row kernel exactly, so batch == matvec bit-for-bit).
9420#[cfg(target_arch = "aarch64")]
9421#[target_feature(enable = "neon,dotprod")]
9422unsafe fn dot_q1_row_1x4_sdot(
9423    bytes: &[u8],
9424    r: usize,
9425    gpr: usize,
9426    xs: [&[i8]; 4],
9427    gs: [&[i32]; 4],
9428) -> [f32; 4] {
9429    // SAFETY: same slice-length contracts as `dot_q1_row_sdot`, ×4.
9430    unsafe {
9431        use core::arch::aarch64::*;
9432        use core::arch::asm;
9433        const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
9434        const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
9435        const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
9436        const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
9437        const IW11: [u8; 16] = [
9438            10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
9439        ];
9440        const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
9441        let m = vld1q_u8(MASKS.as_ptr());
9442        let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
9443        let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
9444        let isc = vld1_u8(ISC.as_ptr());
9445        macro_rules! sdot2 {
9446            ($w0:expr, $w1:expr, $x:expr) => {{
9447                let x0 = vld1q_s8($x);
9448                let x1 = vld1q_s8($x.add(16));
9449                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
9450                asm!(
9451                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
9452                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
9453                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
9454                    w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
9455                    options(pure, nomem, nostack),
9456                );
9457                vaddq_s32(a0, a1)
9458            }};
9459        }
9460        let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
9461        let row_base = r * gpr * Q1_TILE;
9462        let abs_end = bytes.len();
9463        let mut accv = [vdupq_n_f32(0.0); 4];
9464        let mut gi = 0;
9465        while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
9466            let t0 = base.add(gi * Q1_TILE);
9467            let ld_a = vld1q_u8(t0);
9468            let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
9469            // Unpack ONCE — eight ±mask vectors serve all four streams.
9470            let w00 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw00), m));
9471            let w01 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw01), m));
9472            let w10 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw10), m));
9473            let w11 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw11), m));
9474            let w20 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw00), m));
9475            let w21 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw01), m));
9476            let w30 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw10), m));
9477            let w31 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw11), m));
9478            let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
9479            let scf: float32x4_t;
9480            asm!(
9481                "fcvtl {o:v}.4s, {i:v}.4h",
9482                o = out(vreg) scf, i = in(vreg) sc16,
9483                options(pure, nomem, nostack),
9484            );
9485            for k in 0..4 {
9486                let xp = xs[k].as_ptr();
9487                let d0 = sdot2!(w00, w01, xp.add(gi * GROUP_SIZE));
9488                let d1 = sdot2!(w10, w11, xp.add((gi + 1) * GROUP_SIZE));
9489                let d2 = sdot2!(w20, w21, xp.add((gi + 2) * GROUP_SIZE));
9490                let d3 = sdot2!(w30, w31, xp.add((gi + 3) * GROUP_SIZE));
9491                let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
9492                let g = vld1q_s32(gs[k].as_ptr().add(gi));
9493                let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
9494                accv[k] = vfmaq_f32(accv[k], vcvtq_f32_s32(dots), scf);
9495            }
9496            gi += 4;
9497        }
9498        let mut acc = [
9499            vaddvq_f32(accv[0]),
9500            vaddvq_f32(accv[1]),
9501            vaddvq_f32(accv[2]),
9502            vaddvq_f32(accv[3]),
9503        ];
9504        while gi < gpr {
9505            let t = base.add(gi * Q1_TILE);
9506            let sc = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
9507            let v0 = vcombine_u8(vdup_n_u8(*t.add(2)), vdup_n_u8(*t.add(3)));
9508            let v1 = vcombine_u8(vdup_n_u8(*t.add(4)), vdup_n_u8(*t.add(5)));
9509            let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
9510            let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
9511            for k in 0..4 {
9512                let d = vaddvq_s32(sdot2!(w0, w1, xs[k].as_ptr().add(gi * GROUP_SIZE)));
9513                acc[k] += (-(2 * d + *gs[k].as_ptr().add(gi))) as f32 * sc;
9514            }
9515            gi += 1;
9516        }
9517        acc
9518    }
9519}
9520
9521/// (weight ±1, scale) of one q1 element — the exact outlier term.
9522#[inline]
9523fn q1_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
9524    let gi = j / GROUP_SIZE;
9525    let k = j % GROUP_SIZE;
9526    let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
9527    let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
9528    let bit = (tile[2 + k / 8] >> (k % 8)) & 1;
9529    ((bit as i32 * 2 - 1) as f32, s)
9530}
9531
9532/// Exact scalar q1 row (CMF_SDOT=0 contract).
9533#[inline]
9534fn q1_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
9535    let mut acc = 0f32;
9536    for gi in 0..gpr {
9537        let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
9538        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
9539        let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
9540        let mut ga = 0f32;
9541        for (j, &b) in tile[2..].iter().enumerate() {
9542            for k in 0..8 {
9543                ga += (((b >> k) & 1) as f32 * 2.0 - 1.0) * xg[j * 8 + k];
9544            }
9545        }
9546        acc += ga * s;
9547    }
9548    acc
9549}
9550
9551/// One binary-weight row range via A8W8 (the body of `bin_matvec`'s hot
9552/// loop, extracted so multi-matrix jobs drive the same kernel).
9553#[allow(clippy::too_many_arguments)]
9554fn bin_range_a8w8(
9555    w: BinW,
9556    gpr: usize,
9557    act: &SplitAct,
9558    gsum: &[i32],
9559    out: SendMut,
9560    start: usize,
9561    end: usize,
9562) {
9563    for r in start..end {
9564        let mut acc = w.dot_i8(r, gpr, &act.xq, gsum) * act.sx;
9565        for &(j, xv) in &act.outliers {
9566            let (wv, s) = w.outlier(r, gpr, j);
9567            acc += wv * s * xv;
9568        }
9569        // SAFETY: disjoint row ranges per worker.
9570        unsafe { *out.at(r) = acc };
9571    }
9572}
9573
9574/// Exact-scalar binary-weight row range (CMF_SDOT=0 contract).
9575fn bin_range_f32(w: BinW, gpr: usize, x: &[f32], out: SendMut, start: usize, end: usize) {
9576    for r in start..end {
9577        // SAFETY: disjoint row ranges per worker.
9578        unsafe { *out.at(r) = w.row_exact(r, gpr, x) };
9579    }
9580}
9581
9582/// A binary (±s) weight in either of its layouts: `q1` 6-byte tiles, or
9583/// `q1g` planes with one scale per 128-group. The row drivers — matvec,
9584/// matvec2, matmat, the multi-projection job and the fused gate/up pair —
9585/// are written once against this; only the per-row kernels differ, and
9586/// the match costs one predictable branch per row. `gpr` is always the
9587/// row's count of 32-tiles (cols / 32): the activation group sums and the
9588/// tile walk are shared by both layouts.
9589#[derive(Clone, Copy)]
9590pub(crate) enum BinW<'a> {
9591    Q1(&'a [u8]),
9592    Q1g {
9593        bytes: &'a [u8],
9594        /// Byte offset of the f16 scale plane (= rows · cols/8).
9595        sc_off: usize,
9596        /// Row count: the q1 kernels treat the payload's final row
9597        /// specially (see `dot_q1g_row_sdot`), and q1g mirrors that.
9598        rows: usize,
9599    },
9600}
9601
9602impl<'a> BinW<'a> {
9603    fn q1g(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
9604        Self::Q1g {
9605            bytes,
9606            sc_off: cortiq_core::quant::q1g_scale_offset(rows, cols),
9607            rows,
9608        }
9609    }
9610
9611    #[inline]
9612    fn dot_i8(self, r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
9613        match self {
9614            Self::Q1(b) => dot_q1_row_i8(b, r, gpr, xq, gsum),
9615            Self::Q1g {
9616                bytes,
9617                sc_off,
9618                rows,
9619            } => dot_q1g_row_i8(bytes, sc_off, rows, r, gpr, xq, gsum),
9620        }
9621    }
9622
9623    #[inline]
9624    fn row_exact(self, r: usize, gpr: usize, x: &[f32]) -> f32 {
9625        match self {
9626            Self::Q1(b) => q1_row_exact(b, r, gpr, x),
9627            Self::Q1g { bytes, sc_off, .. } => q1g_row_exact(bytes, sc_off, r, gpr, x),
9628        }
9629    }
9630
9631    #[inline]
9632    fn outlier(self, r: usize, gpr: usize, j: usize) -> (f32, f32) {
9633        match self {
9634            Self::Q1(b) => q1_outlier(b, r, gpr, j),
9635            Self::Q1g { bytes, sc_off, .. } => q1g_outlier(bytes, sc_off, r, gpr, j),
9636        }
9637    }
9638
9639    /// Blocked 1×4 row (one unpack, four activation streams).
9640    ///
9641    /// # Safety
9642    /// The SDOT/dotprod target feature must be present (`sdot_enabled`).
9643    #[cfg(target_arch = "aarch64")]
9644    #[inline]
9645    unsafe fn dot_1x4_sdot(
9646        self,
9647        r: usize,
9648        gpr: usize,
9649        xs: [&[i8]; 4],
9650        gs: [&[i32]; 4],
9651    ) -> [f32; 4] {
9652        // SAFETY: forwarded caller contract.
9653        unsafe {
9654            match self {
9655                Self::Q1(b) => dot_q1_row_1x4_sdot(b, r, gpr, xs, gs),
9656                Self::Q1g {
9657                    bytes,
9658                    sc_off,
9659                    rows,
9660                } => dot_q1g_row_1x4_sdot(bytes, sc_off, rows, r, gpr, xs, gs),
9661            }
9662        }
9663    }
9664
9665    /// Blocked 1×4 row on x86 (VNNI when the tile switch is on).
9666    ///
9667    /// # Safety
9668    /// AVX2 must be present (`avx2_enabled`).
9669    #[cfg(target_arch = "x86_64")]
9670    #[inline]
9671    unsafe fn dot_1x4_x86(self, r: usize, gpr: usize, xs: [&[i8]; 4], gs: [&[i32]; 4]) -> [f32; 4] {
9672        // SAFETY: forwarded caller contract.
9673        unsafe {
9674            match (self, vnni_tiles_enabled()) {
9675                (Self::Q1(b), true) => dot_q1_row_1x4_vnni(b, r, gpr, xs, gs),
9676                (Self::Q1(b), false) => dot_q1_row_1x4_avx2(b, r, gpr, xs, gs),
9677                (Self::Q1g { bytes, sc_off, .. }, true) => {
9678                    dot_q1g_row_1x4_vnni(bytes, sc_off, r, gpr, xs, gs)
9679                }
9680                (Self::Q1g { bytes, sc_off, .. }, false) => {
9681                    dot_q1g_row_1x4_avx2(bytes, sc_off, r, gpr, xs, gs)
9682                }
9683            }
9684        }
9685    }
9686}
9687
9688// ── q1g (dtype 20): binary weights, ONE f16 scale per 128-group, in sign
9689// and scale PLANES (`[bits: rows·cols/8][f16: rows·cols/128]`). Same
9690// kernel family as q1 and — deliberately — the same float accumulation
9691// order (one `acc += d·s` per 32-tile, the group's scale reused four
9692// times), so a q1g repack of a q1 tensor whose tile scales agree gives
9693// the same f32 outputs bit for bit on every path. The win is the stream:
9694// 18 bytes per 128 weights instead of 24, the bits 16-byte aligned. ──
9695
9696/// Sign bits of q1g tile `t` (flat 32-tile index `r·gpr + gi`).
9697#[inline]
9698fn q1g_bits(bytes: &[u8], t: usize) -> u32 {
9699    u32::from_le_bytes([
9700        bytes[t * 4],
9701        bytes[t * 4 + 1],
9702        bytes[t * 4 + 2],
9703        bytes[t * 4 + 3],
9704    ])
9705}
9706
9707/// Raw f16 bits of the scale of the 128-group holding q1g tile `t`.
9708#[inline]
9709fn q1g_scale_bits(bytes: &[u8], sc_off: usize, t: usize) -> u16 {
9710    let o = sc_off + (t / 4) * 2;
9711    u16::from_le_bytes([bytes[o], bytes[o + 1]])
9712}
9713
9714/// One q1g row via the A8W8 int8 path (dispatch twin of `dot_q1_row_i8`).
9715#[allow(unreachable_code)]
9716#[allow(clippy::too_many_arguments)]
9717fn dot_q1g_row_i8(
9718    bytes: &[u8],
9719    sc_off: usize,
9720    rows: usize,
9721    r: usize,
9722    gpr: usize,
9723    xq: &[i8],
9724    gsum: &[i32],
9725) -> f32 {
9726    #[cfg(target_arch = "aarch64")]
9727    unsafe {
9728        return dot_q1g_row_sdot(bytes, sc_off, rows, r, gpr, xq, gsum);
9729    }
9730    #[cfg(target_arch = "x86_64")]
9731    if avx2_enabled() {
9732        unsafe {
9733            if vnni_tiles_enabled() {
9734                return dot_q1g_row_vnni(bytes, sc_off, r, gpr, xq, gsum);
9735            }
9736            return dot_q1g_row_avx2(bytes, sc_off, r, gpr, xq, gsum);
9737        }
9738    }
9739    let _ = (gsum, rows);
9740    let mut acc = 0f32;
9741    for gi in 0..gpr {
9742        let t = r * gpr + gi;
9743        let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, t));
9744        let bits = q1g_bits(bytes, t);
9745        let mut d = 0i32;
9746        for k in 0..GROUP_SIZE {
9747            let w = ((bits >> k) & 1) as i32 * 2 - 1;
9748            d += w * xq[gi * GROUP_SIZE + k] as i32;
9749        }
9750        acc += d as f32 * s;
9751    }
9752    acc
9753}
9754
9755/// Four tiles' unreduced `i32×8` partial sums → `[Σ0, Σ1, Σ2, Σ3]` in one
9756/// hadd tree: 3 hadds + 1 add for four tiles where the per-tile reduce of
9757/// the q1 kernels costs 7 ops each. Integer math, so any order gives the
9758/// q1 kernel's per-tile sums exactly.
9759#[cfg(target_arch = "x86_64")]
9760#[target_feature(enable = "avx2")]
9761#[inline]
9762unsafe fn q1g_hsum4(
9763    d0: core::arch::x86_64::__m256i,
9764    d1: core::arch::x86_64::__m256i,
9765    d2: core::arch::x86_64::__m256i,
9766    d3: core::arch::x86_64::__m256i,
9767) -> core::arch::x86_64::__m128i {
9768    // SAFETY: pure register math.
9769    unsafe {
9770        use core::arch::x86_64::*;
9771        let h = _mm256_hadd_epi32(_mm256_hadd_epi32(d0, d1), _mm256_hadd_epi32(d2, d3));
9772        _mm_add_epi32(_mm256_castsi256_si128(h), _mm256_extracti128_si256::<1>(h))
9773    }
9774}
9775
9776/// One 128-group's epilogue, in q1's float order: `d = 2·Σset − Σx` per
9777/// tile (exact i32 → exact f32), then `acc += d·s` tile by tile — four
9778/// separately rounded multiply/add pairs, the q1 kernels' sequence.
9779#[cfg(target_arch = "x86_64")]
9780#[target_feature(enable = "avx2")]
9781#[inline]
9782unsafe fn q1g_group_acc(
9783    acc: f32,
9784    msum4: core::arch::x86_64::__m128i,
9785    gsum4: &[i32],
9786    s: f32,
9787) -> f32 {
9788    // SAFETY: gsum4 holds the group's four tile sums (callers slice it).
9789    unsafe {
9790        use core::arch::x86_64::*;
9791        let g = _mm_loadu_si128(gsum4.as_ptr() as *const __m128i);
9792        let d = _mm_cvtepi32_ps(_mm_sub_epi32(_mm_slli_epi32::<1>(msum4), g));
9793        let mut f = [0f32; 4];
9794        _mm_storeu_ps(f.as_mut_ptr(), d);
9795        let mut acc = acc;
9796        for v in f {
9797            acc += v * s;
9798        }
9799        acc
9800    }
9801}
9802
9803/// AVX2 q1g row: `dot_q1_row_avx2`'s mask/select/maddubs per tile, but the
9804/// four tiles of a 128-group reduce together (`q1g_hsum4`) and the scale
9805/// converts once per group — bit-identical results (integer sums, then
9806/// q1's per-tile float sequence), fewer horizontal ops and f16 decodes.
9807#[cfg(target_arch = "x86_64")]
9808#[target_feature(enable = "avx2")]
9809unsafe fn dot_q1g_row_avx2(
9810    bytes: &[u8],
9811    sc_off: usize,
9812    r: usize,
9813    gpr: usize,
9814    xq: &[i8],
9815    gsum: &[i32],
9816) -> f32 {
9817    // SAFETY: callers uphold the plane-length and xq/gsum contracts
9818    // (gpr % 4 == 0, xq.len() == gpr·32, gsum.len() == gpr).
9819    unsafe {
9820        use core::arch::x86_64::*;
9821        let expand = _mm256_setr_epi8(
9822            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,
9823            3, 3, 3,
9824        );
9825        let bitsel = _mm256_setr_epi8(
9826            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
9827            -128, 1, 2, 4, 8, 16, 32, 64, -128,
9828        );
9829        let ones8 = _mm256_set1_epi8(1);
9830        let ones16 = _mm256_set1_epi16(1);
9831        let bp = bytes.as_ptr().add(r * gpr * 4);
9832        let mut acc = 0f32;
9833        let mut gi = 0;
9834        while gi < gpr {
9835            let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
9836            let mut d = [_mm256_setzero_si256(); 4];
9837            for (k, dk) in d.iter_mut().enumerate() {
9838                let t = gi + k;
9839                let bits = (bp.add(t * 4) as *const u32).read_unaligned();
9840                let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
9841                let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
9842                let x = _mm256_loadu_si256(xq.as_ptr().add(t * GROUP_SIZE) as *const __m256i);
9843                let sel = _mm256_and_si256(x, mask);
9844                *dk = _mm256_madd_epi16(_mm256_maddubs_epi16(ones8, sel), ones16);
9845            }
9846            let m4 = q1g_hsum4(d[0], d[1], d[2], d[3]);
9847            acc = q1g_group_acc(acc, m4, &gsum[gi..gi + 4], s);
9848            gi += 4;
9849        }
9850        acc
9851    }
9852}
9853
9854/// VNNI twin of `dot_q1g_row_avx2`: `vpdpbusd(1u8, sel)` per tile (the
9855/// same i32 lanes as maddubs+madd — bounded products), same group reduce.
9856#[cfg(target_arch = "x86_64")]
9857#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
9858unsafe fn dot_q1g_row_vnni(
9859    bytes: &[u8],
9860    sc_off: usize,
9861    r: usize,
9862    gpr: usize,
9863    xq: &[i8],
9864    gsum: &[i32],
9865) -> f32 {
9866    // SAFETY: callers uphold the plane-length and xq/gsum contracts.
9867    unsafe {
9868        use core::arch::x86_64::*;
9869        let expand = _mm256_setr_epi8(
9870            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,
9871            3, 3, 3,
9872        );
9873        let bitsel = _mm256_setr_epi8(
9874            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
9875            -128, 1, 2, 4, 8, 16, 32, 64, -128,
9876        );
9877        let ones8 = _mm256_set1_epi8(1);
9878        let bp = bytes.as_ptr().add(r * gpr * 4);
9879        let mut acc = 0f32;
9880        let mut gi = 0;
9881        while gi < gpr {
9882            let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
9883            let mut d = [_mm256_setzero_si256(); 4];
9884            for (k, dk) in d.iter_mut().enumerate() {
9885                let t = gi + k;
9886                let bits = (bp.add(t * 4) as *const u32).read_unaligned();
9887                let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
9888                let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
9889                let x = _mm256_loadu_si256(xq.as_ptr().add(t * GROUP_SIZE) as *const __m256i);
9890                *dk = _mm256_dpbusd_epi32(_mm256_setzero_si256(), ones8, _mm256_and_si256(x, mask));
9891            }
9892            let m4 = q1g_hsum4(d[0], d[1], d[2], d[3]);
9893            acc = q1g_group_acc(acc, m4, &gsum[gi..gi + 4], s);
9894            gi += 4;
9895        }
9896        acc
9897    }
9898}
9899
9900/// Blocked 1×4 AVX2 q1g row: each tile's mask is built once and serves
9901/// four activation streams; per stream, the group reduce and float order
9902/// of `dot_q1g_row_avx2` (so batch == matvec bit for bit).
9903#[cfg(target_arch = "x86_64")]
9904#[target_feature(enable = "avx2")]
9905unsafe fn dot_q1g_row_1x4_avx2(
9906    bytes: &[u8],
9907    sc_off: usize,
9908    r: usize,
9909    gpr: usize,
9910    xs: [&[i8]; 4],
9911    gsums: [&[i32]; 4],
9912) -> [f32; 4] {
9913    // SAFETY: callers uphold the plane-length and xq/gsum contracts.
9914    unsafe {
9915        use core::arch::x86_64::*;
9916        let expand = _mm256_setr_epi8(
9917            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,
9918            3, 3, 3,
9919        );
9920        let bitsel = _mm256_setr_epi8(
9921            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
9922            -128, 1, 2, 4, 8, 16, 32, 64, -128,
9923        );
9924        let ones8 = _mm256_set1_epi8(1);
9925        let ones16 = _mm256_set1_epi16(1);
9926        let bp = bytes.as_ptr().add(r * gpr * 4);
9927        let mut acc = [0f32; 4];
9928        let mut gi = 0;
9929        while gi < gpr {
9930            let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
9931            // d[stream][tile]
9932            let mut d = [[_mm256_setzero_si256(); 4]; 4];
9933            for t in 0..4 {
9934                let bits = (bp.add((gi + t) * 4) as *const u32).read_unaligned();
9935                let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
9936                let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
9937                for (k, xq) in xs.iter().enumerate() {
9938                    let x = _mm256_loadu_si256(
9939                        xq.as_ptr().add((gi + t) * GROUP_SIZE) as *const __m256i
9940                    );
9941                    let sel = _mm256_and_si256(x, mask);
9942                    d[k][t] = _mm256_madd_epi16(_mm256_maddubs_epi16(ones8, sel), ones16);
9943                }
9944            }
9945            for k in 0..4 {
9946                let m4 = q1g_hsum4(d[k][0], d[k][1], d[k][2], d[k][3]);
9947                acc[k] = q1g_group_acc(acc[k], m4, &gsums[k][gi..gi + 4], s);
9948            }
9949            gi += 4;
9950        }
9951        acc
9952    }
9953}
9954
9955/// VNNI twin of `dot_q1g_row_1x4_avx2`.
9956#[cfg(target_arch = "x86_64")]
9957#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
9958unsafe fn dot_q1g_row_1x4_vnni(
9959    bytes: &[u8],
9960    sc_off: usize,
9961    r: usize,
9962    gpr: usize,
9963    xs: [&[i8]; 4],
9964    gsums: [&[i32]; 4],
9965) -> [f32; 4] {
9966    // SAFETY: callers uphold the plane-length and xq/gsum contracts.
9967    unsafe {
9968        use core::arch::x86_64::*;
9969        let expand = _mm256_setr_epi8(
9970            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,
9971            3, 3, 3,
9972        );
9973        let bitsel = _mm256_setr_epi8(
9974            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
9975            -128, 1, 2, 4, 8, 16, 32, 64, -128,
9976        );
9977        let ones8 = _mm256_set1_epi8(1);
9978        let bp = bytes.as_ptr().add(r * gpr * 4);
9979        let mut acc = [0f32; 4];
9980        let mut gi = 0;
9981        while gi < gpr {
9982            let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
9983            let mut d = [[_mm256_setzero_si256(); 4]; 4];
9984            for t in 0..4 {
9985                let bits = (bp.add((gi + t) * 4) as *const u32).read_unaligned();
9986                let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
9987                let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
9988                for (k, xq) in xs.iter().enumerate() {
9989                    let x = _mm256_loadu_si256(
9990                        xq.as_ptr().add((gi + t) * GROUP_SIZE) as *const __m256i
9991                    );
9992                    d[k][t] = _mm256_dpbusd_epi32(
9993                        _mm256_setzero_si256(),
9994                        ones8,
9995                        _mm256_and_si256(x, mask),
9996                    );
9997                }
9998            }
9999            for k in 0..4 {
10000                let m4 = q1g_hsum4(d[k][0], d[k][1], d[k][2], d[k][3]);
10001                acc[k] = q1g_group_acc(acc[k], m4, &gsums[k][gi..gi + 4], s);
10002            }
10003            gi += 4;
10004        }
10005        acc
10006    }
10007}
10008
10009/// SDOT q1g row: the `dot_q1_row_sdot` math on the plane layout. One
10010/// aligned 16-byte load carries a whole 128-group's signs (q1 needs two
10011/// pair loads and a TBL gather of four scales for the same weights), TBL
10012/// replicates each sign byte across 8 lanes for vtst, and the group's ONE
10013/// scale converts once (FCVTL, exact) and broadcasts. Lane order, the
10014/// vpaddq tree, the fma into the four lane accumulators and the tail are
10015/// the q1 kernel's — including its quirk: q1's pair load reads 4 bytes
10016/// past the 4th tile, so the payload's final four tiles (last row) take
10017/// the scalar tail there. q1g has no overread, but takes the same tail on
10018/// the last row to stay bit-identical to q1.
10019#[cfg(target_arch = "aarch64")]
10020#[target_feature(enable = "neon,dotprod")]
10021#[allow(clippy::too_many_arguments)]
10022unsafe fn dot_q1g_row_sdot(
10023    bytes: &[u8],
10024    sc_off: usize,
10025    rows: usize,
10026    r: usize,
10027    gpr: usize,
10028    xq: &[i8],
10029    gsum: &[i32],
10030) -> f32 {
10031    // SAFETY: callers uphold the plane-length and xq/gsum contracts
10032    // (gpr % 4 == 0, xq.len() == gpr·32, gsum.len() == gpr).
10033    unsafe {
10034        use core::arch::aarch64::*;
10035        use core::arch::asm;
10036        const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
10037        // Tile t of the group = sign bytes 4t..4t+4 of the 16-byte load.
10038        const IW: [[u8; 16]; 8] = [
10039            [0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1],
10040            [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3],
10041            [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5],
10042            [6, 6, 6, 6, 6, 6, 6, 6, 7, 7, 7, 7, 7, 7, 7, 7],
10043            [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9],
10044            [
10045                10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
10046            ],
10047            [
10048                12, 12, 12, 12, 12, 12, 12, 12, 13, 13, 13, 13, 13, 13, 13, 13,
10049            ],
10050            [
10051                14, 14, 14, 14, 14, 14, 14, 14, 15, 15, 15, 15, 15, 15, 15, 15,
10052            ],
10053        ];
10054        let m = vld1q_u8(MASKS.as_ptr());
10055        let iw: [uint8x16_t; 8] = core::array::from_fn(|i| vld1q_u8(IW[i].as_ptr()));
10056        macro_rules! tile_dot {
10057            ($w0:expr, $w1:expr, $x:expr) => {{
10058                let x0 = vld1q_s8($x);
10059                let x1 = vld1q_s8($x.add(16));
10060                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
10061                asm!(
10062                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
10063                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
10064                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
10065                    w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
10066                    options(pure, nomem, nostack),
10067                );
10068                vaddq_s32(a0, a1)
10069            }};
10070        }
10071        macro_rules! tbl_mask {
10072            ($ld:expr, $i:expr) => {
10073                vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, iw[$i]), m))
10074            };
10075        }
10076        let bp = bytes.as_ptr().add(r * gpr * 4);
10077        let xp = xq.as_ptr();
10078        let gp = gsum.as_ptr();
10079        let vec_end = if r + 1 == rows { gpr - 4 } else { gpr };
10080        let mut accv = vdupq_n_f32(0.0);
10081        let mut gi = 0;
10082        while gi + 4 <= vec_end {
10083            let ld = vld1q_u8(bp.add(gi * 4));
10084            let d0 = tile_dot!(tbl_mask!(ld, 0), tbl_mask!(ld, 1), xp.add(gi * GROUP_SIZE));
10085            let d1 = tile_dot!(
10086                tbl_mask!(ld, 2),
10087                tbl_mask!(ld, 3),
10088                xp.add((gi + 1) * GROUP_SIZE)
10089            );
10090            let d2 = tile_dot!(
10091                tbl_mask!(ld, 4),
10092                tbl_mask!(ld, 5),
10093                xp.add((gi + 2) * GROUP_SIZE)
10094            );
10095            let d3 = tile_dot!(
10096                tbl_mask!(ld, 6),
10097                tbl_mask!(ld, 7),
10098                xp.add((gi + 3) * GROUP_SIZE)
10099            );
10100            let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
10101            let g = vld1q_s32(gp.add(gi));
10102            let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
10103            let h = vdup_n_u16(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
10104            let scf: float32x4_t;
10105            asm!(
10106                "fcvtl {o:v}.4s, {i:v}.4h",
10107                o = out(vreg) scf, i = in(vreg) h,
10108                options(pure, nomem, nostack),
10109            );
10110            accv = vfmaq_f32(accv, vcvtq_f32_s32(dots), scf);
10111            gi += 4;
10112        }
10113        let mut acc = vaddvq_f32(accv);
10114        while gi < gpr {
10115            let t = bp.add(gi * 4);
10116            let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
10117            let v0 = vcombine_u8(vdup_n_u8(*t), vdup_n_u8(*t.add(1)));
10118            let v1 = vcombine_u8(vdup_n_u8(*t.add(2)), vdup_n_u8(*t.add(3)));
10119            let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
10120            let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
10121            let d = vaddvq_s32(tile_dot!(w0, w1, xp.add(gi * GROUP_SIZE)));
10122            acc += (-(2 * d + *gp.add(gi))) as f32 * s;
10123            gi += 1;
10124        }
10125        acc
10126    }
10127}
10128
10129/// Blocked q1g 1×4: one unpack of the group serves four activation
10130/// streams; per stream bit-identical to `dot_q1g_row_sdot` (and so to the
10131/// q1 1×4 kernel, same last-row tail).
10132#[cfg(target_arch = "aarch64")]
10133#[target_feature(enable = "neon,dotprod")]
10134#[allow(clippy::too_many_arguments)]
10135unsafe fn dot_q1g_row_1x4_sdot(
10136    bytes: &[u8],
10137    sc_off: usize,
10138    rows: usize,
10139    r: usize,
10140    gpr: usize,
10141    xs: [&[i8]; 4],
10142    gs: [&[i32]; 4],
10143) -> [f32; 4] {
10144    // SAFETY: same contracts as `dot_q1g_row_sdot`, ×4.
10145    unsafe {
10146        use core::arch::aarch64::*;
10147        use core::arch::asm;
10148        const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
10149        const IW: [[u8; 16]; 8] = [
10150            [0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1],
10151            [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3],
10152            [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5],
10153            [6, 6, 6, 6, 6, 6, 6, 6, 7, 7, 7, 7, 7, 7, 7, 7],
10154            [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9],
10155            [
10156                10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
10157            ],
10158            [
10159                12, 12, 12, 12, 12, 12, 12, 12, 13, 13, 13, 13, 13, 13, 13, 13,
10160            ],
10161            [
10162                14, 14, 14, 14, 14, 14, 14, 14, 15, 15, 15, 15, 15, 15, 15, 15,
10163            ],
10164        ];
10165        let m = vld1q_u8(MASKS.as_ptr());
10166        let iw: [uint8x16_t; 8] = core::array::from_fn(|i| vld1q_u8(IW[i].as_ptr()));
10167        macro_rules! sdot2 {
10168            ($w0:expr, $w1:expr, $x:expr) => {{
10169                let x0 = vld1q_s8($x);
10170                let x1 = vld1q_s8($x.add(16));
10171                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
10172                asm!(
10173                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
10174                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
10175                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
10176                    w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
10177                    options(pure, nomem, nostack),
10178                );
10179                vaddq_s32(a0, a1)
10180            }};
10181        }
10182        let bp = bytes.as_ptr().add(r * gpr * 4);
10183        let vec_end = if r + 1 == rows { gpr - 4 } else { gpr };
10184        let mut accv = [vdupq_n_f32(0.0); 4];
10185        let mut gi = 0;
10186        while gi + 4 <= vec_end {
10187            let ld = vld1q_u8(bp.add(gi * 4));
10188            // Unpack ONCE — eight ±mask vectors serve all four streams.
10189            let w: [int8x16_t; 8] =
10190                core::array::from_fn(|i| vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld, iw[i]), m)));
10191            let h = vdup_n_u16(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
10192            let scf: float32x4_t;
10193            asm!(
10194                "fcvtl {o:v}.4s, {i:v}.4h",
10195                o = out(vreg) scf, i = in(vreg) h,
10196                options(pure, nomem, nostack),
10197            );
10198            for k in 0..4 {
10199                let xp = xs[k].as_ptr();
10200                let d0 = sdot2!(w[0], w[1], xp.add(gi * GROUP_SIZE));
10201                let d1 = sdot2!(w[2], w[3], xp.add((gi + 1) * GROUP_SIZE));
10202                let d2 = sdot2!(w[4], w[5], xp.add((gi + 2) * GROUP_SIZE));
10203                let d3 = sdot2!(w[6], w[7], xp.add((gi + 3) * GROUP_SIZE));
10204                let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
10205                let g = vld1q_s32(gs[k].as_ptr().add(gi));
10206                let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
10207                accv[k] = vfmaq_f32(accv[k], vcvtq_f32_s32(dots), scf);
10208            }
10209            gi += 4;
10210        }
10211        let mut acc = [
10212            vaddvq_f32(accv[0]),
10213            vaddvq_f32(accv[1]),
10214            vaddvq_f32(accv[2]),
10215            vaddvq_f32(accv[3]),
10216        ];
10217        while gi < gpr {
10218            let t = bp.add(gi * 4);
10219            let sc = f16_to_f32(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
10220            let v0 = vcombine_u8(vdup_n_u8(*t), vdup_n_u8(*t.add(1)));
10221            let v1 = vcombine_u8(vdup_n_u8(*t.add(2)), vdup_n_u8(*t.add(3)));
10222            let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
10223            let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
10224            for k in 0..4 {
10225                let d = vaddvq_s32(sdot2!(w0, w1, xs[k].as_ptr().add(gi * GROUP_SIZE)));
10226                acc[k] += (-(2 * d + *gs[k].as_ptr().add(gi))) as f32 * sc;
10227            }
10228            gi += 1;
10229        }
10230        acc
10231    }
10232}
10233
10234/// (weight ±1, scale) of one q1g element — the exact outlier term.
10235#[inline]
10236fn q1g_outlier(bytes: &[u8], sc_off: usize, r: usize, gpr: usize, j: usize) -> (f32, f32) {
10237    let t = r * gpr + j / GROUP_SIZE;
10238    let k = j % GROUP_SIZE;
10239    let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, t));
10240    let bit = (bytes[t * 4 + k / 8] >> (k % 8)) & 1;
10241    ((bit as i32 * 2 - 1) as f32, s)
10242}
10243
10244/// Exact scalar q1g row (CMF_SDOT=0 contract) — `q1_row_exact`'s order.
10245#[inline]
10246fn q1g_row_exact(bytes: &[u8], sc_off: usize, r: usize, gpr: usize, x: &[f32]) -> f32 {
10247    let mut acc = 0f32;
10248    for gi in 0..gpr {
10249        let t = r * gpr + gi;
10250        let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, t));
10251        let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
10252        let mut ga = 0f32;
10253        for (j, &b) in bytes[t * 4..t * 4 + 4].iter().enumerate() {
10254            for k in 0..8 {
10255                ga += (((b >> k) & 1) as f32 * 2.0 - 1.0) * xg[j * 8 + k];
10256            }
10257        }
10258        acc += ga * s;
10259    }
10260    acc
10261}
10262
10263/// q1t per-row overlay locator. After the base (`base_len`) come
10264/// `[u32 row_ptr[rows+1]]` then `[(u16 col, f16 val)]` grouped by row (row
10265/// `r`'s entries are `[row_ptr[r], row_ptr[r+1])`). Returns
10266/// `(row_ptr offset, entries offset, present)`.
10267fn q1t_overlay(bytes: &[u8], base_len: usize, rows: usize) -> (usize, usize, bool) {
10268    let entries = base_len + (rows + 1) * 4;
10269    (base_len, entries, entries <= bytes.len())
10270}
10271
10272/// Read `row_ptr[r]` from the overlay's prefix-sum table.
10273#[inline]
10274fn q1t_rowptr(bytes: &[u8], rp_off: usize, r: usize) -> usize {
10275    let o = rp_off + r * 4;
10276    u32::from_le_bytes([bytes[o], bytes[o + 1], bytes[o + 2], bytes[o + 3]]) as usize
10277}
10278
10279/// Byte → the 5 ternary signs it packs `{−1,0,+1}` as f32, precomputed so
10280/// decoding a q1t code is a table load, not the base-3 divide/modulo per
10281/// weight (division is ~20–40× the cost of a load). Built at compile time.
10282const SIGN5: [[f32; 5]; 256] = {
10283    let mut lut = [[0.0f32; 5]; 256];
10284    let pow3 = [1u16, 3, 9, 27, 81];
10285    let mut byte = 0usize;
10286    while byte < 256 {
10287        let mut i = 0usize;
10288        while i < 5 {
10289            let code = (byte as u16 / pow3[i]) % 3;
10290            lut[byte][i] = if code == 1 {
10291                1.0
10292            } else if code == 2 {
10293                -1.0
10294            } else {
10295                0.0
10296            };
10297            i += 1;
10298        }
10299        byte += 1;
10300    }
10301    lut
10302};
10303
10304/// Same table, as i8 signs — the operand for the int8 SDOT base kernel.
10305const SIGN5_I8: [[i8; 5]; 256] = {
10306    let mut lut = [[0i8; 5]; 256];
10307    let pow3 = [1u16, 3, 9, 27, 81];
10308    let mut byte = 0usize;
10309    while byte < 256 {
10310        let mut i = 0usize;
10311        while i < 5 {
10312            let code = (byte as u16 / pow3[i]) % 3;
10313            lut[byte][i] = if code == 1 {
10314                1
10315            } else if code == 2 {
10316                -1
10317            } else {
10318                0
10319            };
10320            i += 1;
10321        }
10322        byte += 1;
10323    }
10324    lut
10325};
10326
10327/// The same 5 i8 signs packed into a u64 (`[s0 s1 s2 s3 s4 0 0 0]`, LE) so the
10328/// group unpack is 7 unaligned u64 stores at offsets 0,5,10,…,30 instead of
10329/// six 5-byte copies + LUT indexing — each store's trailing zeros are fixed by
10330/// the next store, and the last one runs 6 B past the 32nd weight (the unpack
10331/// buffer is padded to 40). This is the decode/prefill hot inner op.
10332const SIGN5_U64: [u64; 256] = {
10333    let mut lut = [0u64; 256];
10334    let pow3 = [1u16, 3, 9, 27, 81];
10335    let mut byte = 0usize;
10336    while byte < 256 {
10337        let mut v = 0u64;
10338        let mut i = 0usize;
10339        while i < 5 {
10340            let code = (byte as u16 / pow3[i]) % 3;
10341            let s: u8 = if code == 1 {
10342                1
10343            } else if code == 2 {
10344                0xFF
10345            } else {
10346                0
10347            };
10348            v |= (s as u64) << (i * 8);
10349            i += 1;
10350        }
10351        lut[byte] = v;
10352        byte += 1;
10353    }
10354    lut
10355};
10356
10357/// Ternary base weight at `(row r, col j)` = `sign(code)·s_group`. Used to add
10358/// back activation-outlier columns, whose `x` was zeroed for the int8 bulk dot
10359/// (`split_act`). At a weight-outlier position the code is 0, so this is 0 and
10360/// the overlay correction owns that column — no double counting.
10361#[inline]
10362fn q1t_base_weight(bytes: &[u8], r: usize, gpr: usize, j: usize) -> f32 {
10363    const TILE: usize = cortiq_core::quant::Q1T_TILE;
10364    let off = (r * gpr + j / GROUP_SIZE) * TILE;
10365    let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
10366    let within = j % GROUP_SIZE;
10367    SIGN5[bytes[off + 2 + within / 5] as usize][within % 5] * s
10368}
10369
10370/// One 32-group int8 dot via two SDOTs. Bit-exact vs the scalar i8 sum
10371/// (integer accumulation is order-independent).
10372#[cfg(target_arch = "aarch64")]
10373#[target_feature(enable = "neon,dotprod")]
10374#[inline]
10375unsafe fn sdot32_i8(w: *const i8, x: *const i8) -> i32 {
10376    // SAFETY: caller guarantees 32 readable i8 at each pointer.
10377    unsafe {
10378        use core::arch::aarch64::*;
10379        use core::arch::asm;
10380        let w0 = vld1q_s8(w);
10381        let w1 = vld1q_s8(w.add(16));
10382        let x0 = vld1q_s8(x);
10383        let x1 = vld1q_s8(x.add(16));
10384        let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
10385        asm!(
10386            "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
10387            "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
10388            a0 = inout(vreg) a0, a1 = inout(vreg) a1,
10389            w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
10390            options(pure, nomem, nostack),
10391        );
10392        vaddvq_s32(vaddq_s32(a0, a1))
10393    }
10394}
10395
10396/// One 32-group int8 dot via AVX2: signed·signed as `maddubs(|w|, sign(x,w))`
10397/// then `madd` and a horizontal reduce (the same idiom as `dot_q4t_row_avx2`).
10398#[cfg(target_arch = "x86_64")]
10399#[target_feature(enable = "avx2")]
10400#[inline]
10401unsafe fn i8dot32_avx2(w: *const i8, x: *const i8) -> i32 {
10402    // SAFETY: caller guarantees 32 readable i8 at each pointer.
10403    unsafe {
10404        use core::arch::x86_64::*;
10405        let wv = _mm256_loadu_si256(w as *const __m256i);
10406        let xv = _mm256_loadu_si256(x as *const __m256i);
10407        let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
10408        let d = _mm256_madd_epi16(p16, _mm256_set1_epi16(1));
10409        let hi128 = _mm256_extracti128_si256::<1>(d);
10410        let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
10411        let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
10412        let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
10413        _mm_cvtsi128_si32(s32)
10414    }
10415}
10416
10417/// Unpack one q1t group's base-3 codes into 32 i8 signs via 7 unaligned u64
10418/// stores (see `SIGN5_U64`). `dst` MUST have ≥ 40 bytes: the 7th store writes
10419/// `dst[30..38]`. Stores go in order so each one's trailing zeros are
10420/// overwritten by the next; the final 6 padding bytes are unused by the dot.
10421#[inline]
10422fn q1t_unpack_group_i8(codes: *const u8, dst: &mut [i8]) {
10423    debug_assert!(dst.len() >= 40);
10424    // SAFETY: codes points at 7 readable bytes; dst has ≥ 40 bytes so every
10425    // 8-byte store at offset bi*5 (bi ≤ 6 → ≤ 30) stays in bounds.
10426    unsafe {
10427        let p = dst.as_mut_ptr();
10428        for bi in 0..7 {
10429            core::ptr::write_unaligned(
10430                p.add(bi * 5) as *mut u64,
10431                SIGN5_U64[*codes.add(bi) as usize],
10432            );
10433        }
10434    }
10435}
10436
10437/// One 32-group int8 dot, arch-dispatched (the matmat inner loop, where the
10438/// row's signs are unpacked once and dotted against every batch input).
10439/// Callers are gated by `a8w8_enabled()`, so the target-feature arms are
10440/// reachable; the scalar arm is a non-SIMD-arch fallback.
10441#[inline]
10442fn q1t_i8dot32(w: *const i8, x: *const i8) -> i32 {
10443    #[cfg(target_arch = "aarch64")]
10444    unsafe {
10445        return sdot32_i8(w, x);
10446    }
10447    #[cfg(target_arch = "x86_64")]
10448    unsafe {
10449        return i8dot32_avx2(w, x);
10450    }
10451    #[allow(unreachable_code)]
10452    unsafe {
10453        let mut s = 0i32;
10454        for k in 0..GROUP_SIZE {
10455            s += *w.add(k) as i32 * *x.add(k) as i32;
10456        }
10457        s
10458    }
10459}
10460
10461#[inline]
10462unsafe fn q1t_unpack_reg_u64s(codes: *const u8) -> (u64, u64, u64, u64) {
10463    let (s0, s1, s2, s3, s4, s5, s6) = unsafe {
10464        (
10465            SIGN5_U64[*codes as usize],
10466            SIGN5_U64[*codes.add(1) as usize],
10467            SIGN5_U64[*codes.add(2) as usize],
10468            SIGN5_U64[*codes.add(3) as usize],
10469            SIGN5_U64[*codes.add(4) as usize],
10470            SIGN5_U64[*codes.add(5) as usize],
10471            SIGN5_U64[*codes.add(6) as usize],
10472        )
10473    };
10474
10475    let u0 = s0 | (s1 << 40);
10476    let u1 = (s1 >> 24) | (s2 << 16) | (s3 << 56);
10477    let u2 = (s3 >> 8) | (s4 << 32);
10478    let u3 = (s4 >> 32) | (s5 << 8) | (s6 << 48);
10479
10480    (u0, u1, u2, u3)
10481}
10482
10483/// One q1t row's int8 base dot: `Σ_group s·dot(signs, xq)` (before the shared
10484/// `sx`). Direct register unpacking (zero stack stores/loads, no STLF stalls).
10485/// ARM SDOT.
10486#[cfg(target_arch = "aarch64")]
10487#[target_feature(enable = "neon,dotprod")]
10488unsafe fn q1t_dot_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
10489    use core::arch::aarch64::*;
10490    use core::arch::asm;
10491    unsafe {
10492        const TILE: usize = cortiq_core::quant::Q1T_TILE;
10493        let mut acc = 0f32;
10494        let bytes_ptr = bytes.as_ptr();
10495        let xq_ptr = xq.as_ptr();
10496        let row_off = r * gpr * TILE;
10497
10498        let gpr2 = gpr & !1;
10499        let mut gi = 0;
10500        while gi < gpr2 {
10501            let off0 = row_off + gi * TILE;
10502            let off1 = off0 + TILE;
10503            let s0 = f16_to_f32(u16::from_le_bytes([
10504                *bytes_ptr.add(off0),
10505                *bytes_ptr.add(off0 + 1),
10506            ]));
10507            let s1 = f16_to_f32(u16::from_le_bytes([
10508                *bytes_ptr.add(off1),
10509                *bytes_ptr.add(off1 + 1),
10510            ]));
10511
10512            let (u0_0, u1_0, u2_0, u3_0) = q1t_unpack_reg_u64s(bytes_ptr.add(off0 + 2));
10513            let (u0_1, u1_1, u2_1, u3_1) = q1t_unpack_reg_u64s(bytes_ptr.add(off1 + 2));
10514
10515            let w0_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_0), vcreate_u64(u1_0)));
10516            let w1_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_0), vcreate_u64(u3_0)));
10517            let w0_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_1), vcreate_u64(u1_1)));
10518            let w1_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_1), vcreate_u64(u3_1)));
10519
10520            let x0_0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE));
10521            let x1_0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE + 16));
10522            let x0_1 = vld1q_s8(xq_ptr.add((gi + 1) * GROUP_SIZE));
10523            let x1_1 = vld1q_s8(xq_ptr.add((gi + 1) * GROUP_SIZE + 16));
10524
10525            let (mut a0_0, mut a1_0) = (vdupq_n_s32(0), vdupq_n_s32(0));
10526            let (mut a0_1, mut a1_1) = (vdupq_n_s32(0), vdupq_n_s32(0));
10527            asm!(
10528                "sdot {a0_0:v}.4s, {w0_0:v}.16b, {x0_0:v}.16b",
10529                "sdot {a1_0:v}.4s, {w1_0:v}.16b, {x1_0:v}.16b",
10530                "sdot {a0_1:v}.4s, {w0_1:v}.16b, {x0_1:v}.16b",
10531                "sdot {a1_1:v}.4s, {w1_1:v}.16b, {x1_1:v}.16b",
10532                a0_0 = inout(vreg) a0_0, a1_0 = inout(vreg) a1_0,
10533                a0_1 = inout(vreg) a0_1, a1_1 = inout(vreg) a1_1,
10534                w0_0 = in(vreg) w0_0, x0_0 = in(vreg) x0_0, w1_0 = in(vreg) w1_0, x1_0 = in(vreg) x1_0,
10535                w0_1 = in(vreg) w0_1, x0_1 = in(vreg) x0_1, w1_1 = in(vreg) w1_1, x1_1 = in(vreg) x1_1,
10536                options(pure, nomem, nostack),
10537            );
10538            let d0 = vaddvq_s32(vaddq_s32(a0_0, a1_0));
10539            let d1 = vaddvq_s32(vaddq_s32(a0_1, a1_1));
10540            acc += d0 as f32 * s0 + d1 as f32 * s1;
10541            gi += 2;
10542        }
10543
10544        if gi < gpr {
10545            let off = row_off + gi * TILE;
10546            let s = f16_to_f32(u16::from_le_bytes([
10547                *bytes_ptr.add(off),
10548                *bytes_ptr.add(off + 1),
10549            ]));
10550            let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
10551            let w0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0), vcreate_u64(u1)));
10552            let w1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2), vcreate_u64(u3)));
10553            let x0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE));
10554            let x1 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE + 16));
10555            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
10556            asm!(
10557                "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
10558                "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
10559                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
10560                w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
10561                options(pure, nomem, nostack),
10562            );
10563            let d = vaddvq_s32(vaddq_s32(a0, a1));
10564            acc += d as f32 * s;
10565        }
10566        acc
10567    }
10568}
10569
10570/// x86 AVX2 mirror of `q1t_dot_row_sdot` (maddubs int8 dot per group).
10571#[cfg(target_arch = "x86_64")]
10572#[target_feature(enable = "avx2")]
10573unsafe fn q1t_dot_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
10574    use core::arch::x86_64::*;
10575    unsafe {
10576        const TILE: usize = cortiq_core::quant::Q1T_TILE;
10577        let mut acc = 0f32;
10578        let bytes_ptr = bytes.as_ptr();
10579        let xq_ptr = xq.as_ptr();
10580        let row_off = r * gpr * TILE;
10581
10582        let ones = _mm256_set1_epi16(1);
10583        for gi in 0..gpr {
10584            let off = row_off + gi * TILE;
10585            let s = f16_to_f32(u16::from_le_bytes([
10586                *bytes_ptr.add(off),
10587                *bytes_ptr.add(off + 1),
10588            ]));
10589            let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
10590            let wv = _mm256_set_epi64x(u3 as i64, u2 as i64, u1 as i64, u0 as i64);
10591            let xv = _mm256_loadu_si256(xq_ptr.add(gi * GROUP_SIZE) as *const __m256i);
10592            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
10593            let d256 = _mm256_madd_epi16(p16, ones);
10594            let d128 = _mm_add_epi32(
10595                _mm256_castsi256_si128(d256),
10596                _mm256_extracti128_si256(d256, 1),
10597            );
10598            let d64 = _mm_add_epi32(d128, _mm_shuffle_epi32(d128, 0xee));
10599            let d32 = _mm_cvtsi128_si32(_mm_add_epi32(d64, _mm_shuffle_epi32(d64, 0x55)));
10600            acc += d32 as f32 * s;
10601        }
10602        acc
10603    }
10604}
10605
10606/// VNNI twin of `q1t_dot_row_avx2` (see `dpbusd_hsum`).
10607#[cfg(target_arch = "x86_64")]
10608#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
10609unsafe fn q1t_dot_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
10610    use core::arch::x86_64::*;
10611    // SAFETY: same tile/xq contracts as `q1t_dot_row_avx2`.
10612    unsafe {
10613        const TILE: usize = cortiq_core::quant::Q1T_TILE;
10614        let mut acc = 0f32;
10615        let bytes_ptr = bytes.as_ptr();
10616        let xq_ptr = xq.as_ptr();
10617        let row_off = r * gpr * TILE;
10618        for gi in 0..gpr {
10619            let off = row_off + gi * TILE;
10620            let s = f16_to_f32(u16::from_le_bytes([
10621                *bytes_ptr.add(off),
10622                *bytes_ptr.add(off + 1),
10623            ]));
10624            let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
10625            let wv = _mm256_set_epi64x(u3 as i64, u2 as i64, u1 as i64, u0 as i64);
10626            let xv = _mm256_loadu_si256(xq_ptr.add(gi * GROUP_SIZE) as *const __m256i);
10627            let d = dpbusd_hsum(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
10628            acc += d as f32 * s;
10629        }
10630        acc
10631    }
10632}
10633
10634/// Per-row int8 base dot, dispatched once per row (matvec decode hot path).
10635/// Callers are gated by `a8w8_enabled()`, so the target-feature kernels are
10636/// reachable.
10637#[inline]
10638fn q1t_dot_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
10639    #[cfg(target_arch = "aarch64")]
10640    unsafe {
10641        return q1t_dot_row_sdot(bytes, r, gpr, xq);
10642    }
10643    #[cfg(target_arch = "x86_64")]
10644    unsafe {
10645        if vnni_tiles_enabled() {
10646            return q1t_dot_row_vnni(bytes, r, gpr, xq);
10647        }
10648        return q1t_dot_row_avx2(bytes, r, gpr, xq);
10649    }
10650    #[allow(unreachable_code)]
10651    {
10652        const TILE: usize = cortiq_core::quant::Q1T_TILE;
10653        let mut acc = 0f32;
10654        let mut sg = [0i8; GROUP_SIZE + 8]; // +8 slack for the u64-store unpack
10655        for gi in 0..gpr {
10656            let off = (r * gpr + gi) * TILE;
10657            let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
10658            q1t_unpack_group_i8(bytes.as_ptr().wrapping_add(off + 2), &mut sg);
10659            let mut d = 0i32;
10660            for k in 0..GROUP_SIZE {
10661                d += sg[k] as i32 * xq[gi * GROUP_SIZE + k] as i32;
10662            }
10663            acc += d as f32 * s;
10664        }
10665        acc
10666    }
10667}
10668
10669/// Σ over a row's outliers of `value·x[col]` — the correction that adds the
10670/// overlay's exact weights on top of the base dot. INVARIANT: the encoder
10671/// writes ternary code 0 at every outlier position (`quantize_q1t`), so the
10672/// base contributes nothing there and this is a plain `value·x`, not
10673/// `(value − base)·x` — no scattered per-outlier scale read. Row `r`'s entries
10674/// are the contiguous slice `[row_ptr[r], row_ptr[r+1])`, so no binary search.
10675fn q1t_row_outlier_correction(
10676    bytes: &[u8],
10677    r: usize,
10678    rp_off: usize,
10679    entries_off: usize,
10680    has_ov: bool,
10681    x: &[f32],
10682) -> f32 {
10683    if !has_ov {
10684        return 0.0;
10685    }
10686    let (c0, c1) = (
10687        q1t_rowptr(bytes, rp_off, r),
10688        q1t_rowptr(bytes, rp_off, r + 1),
10689    );
10690    let mut corr = 0f32;
10691    for p in c0..c1 {
10692        let e = entries_off + p * 4;
10693        let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
10694        let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
10695        corr += val * x[col];
10696    }
10697    corr
10698}
10699
10700/// Dequantize one q1t row into `buf[..cols]` via the sign LUT (no division),
10701/// then apply the row's outliers (its `[row_ptr[r], row_ptr[r+1])` slice).
10702/// Used by the batched (prefill) path where the decode amortizes over the batch.
10703fn q1t_dequant_row(
10704    bytes: &[u8],
10705    r: usize,
10706    gpr: usize,
10707    rp_off: usize,
10708    entries_off: usize,
10709    has_ov: bool,
10710    buf: &mut [f32],
10711) {
10712    const TILE: usize = cortiq_core::quant::Q1T_TILE;
10713    for g in 0..gpr {
10714        let off = (r * gpr + g) * TILE;
10715        let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
10716        let codes = &bytes[off + 2..off + TILE];
10717        let bc = g * GROUP_SIZE;
10718        // 6 full bytes (30 codes) + a 7th byte holding the last 2.
10719        for bi in 0..6 {
10720            let lut = &SIGN5[codes[bi] as usize];
10721            let d = &mut buf[bc + bi * 5..bc + bi * 5 + 5];
10722            for i in 0..5 {
10723                d[i] = lut[i] * s;
10724            }
10725        }
10726        let lut = &SIGN5[codes[6] as usize];
10727        buf[bc + 30] = lut[0] * s;
10728        buf[bc + 31] = lut[1] * s;
10729    }
10730    if !has_ov {
10731        return;
10732    }
10733    let (c0, c1) = (
10734        q1t_rowptr(bytes, rp_off, r),
10735        q1t_rowptr(bytes, rp_off, r + 1),
10736    );
10737    for p in c0..c1 {
10738        let e = entries_off + p * 4;
10739        let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
10740        buf[col] = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
10741    }
10742}
10743
10744/// Add the sparse outlier overlay onto a base dot already in `out` (the GPU
10745/// computes the ternary base; the overlay stays on the CPU — its entries are
10746/// few and its per-row gather doesn't vectorize on the GPU). Row-parallel.
10747fn q1t_add_overlay(
10748    bytes: &[u8],
10749    x: &[f32],
10750    rows: usize,
10751    cols: usize,
10752    out: &mut [f32],
10753    pool: Option<&Pool>,
10754) {
10755    const TILE: usize = cortiq_core::quant::Q1T_TILE;
10756    let gpr = cols / GROUP_SIZE;
10757    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
10758    if !has_ov {
10759        return;
10760    }
10761    let out_addr = SendMut(out.as_mut_ptr());
10762    let run = move |start: usize, end: usize| {
10763        for r in start..end {
10764            let corr = q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
10765            // SAFETY: disjoint rows; add onto the base the GPU already wrote.
10766            unsafe { *out_addr.at(r) += corr };
10767        }
10768    };
10769    dispatch_rows(pool, rows, &run);
10770}
10771
10772/// Q1T row range via the A8W8 int8 path — shared activation split,
10773/// per-row: base SDOT dot + outlier correction + overlay.
10774#[allow(clippy::too_many_arguments)]
10775fn q1t_range_a8w8(
10776    bytes: &[u8],
10777    gpr: usize,
10778    rp_off: usize,
10779    ent_off: usize,
10780    has_ov: bool,
10781    act: &SplitAct,
10782    x: &[f32],
10783    out: SendMut,
10784    start: usize,
10785    end: usize,
10786) {
10787    for r in start..end {
10788        let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
10789        for &(j, xv) in &act.outliers {
10790            acc += q1t_base_weight(bytes, r, gpr, j) * xv;
10791        }
10792        acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
10793        // SAFETY: disjoint row ranges per worker.
10794        unsafe { *out.at(r) = acc };
10795    }
10796}
10797
10798/// Q1T row range via the f32 path (no SDOT) — for matvec_many batched
10799/// dispatch when a8w8 is unavailable.
10800#[allow(clippy::too_many_arguments)]
10801fn q1t_range_f32_batch(
10802    bytes: &[u8],
10803    gpr: usize,
10804    rp_off: usize,
10805    ent_off: usize,
10806    has_ov: bool,
10807    x: &[f32],
10808    out: SendMut,
10809    start: usize,
10810    end: usize,
10811) {
10812    const TILE: usize = cortiq_core::quant::Q1T_TILE;
10813    let mut sg = [0f32; GROUP_SIZE];
10814    for r in start..end {
10815        let mut acc = 0f32;
10816        for g in 0..gpr {
10817            let off = (r * gpr + g) * TILE;
10818            let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
10819            let codes = &bytes[off + 2..off + TILE];
10820            let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
10821            for bi in 0..6 {
10822                sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
10823            }
10824            let lut = &SIGN5[codes[6] as usize];
10825            sg[30] = lut[0];
10826            sg[31] = lut[1];
10827            let mut gsum = 0f32;
10828            for k in 0..GROUP_SIZE {
10829                gsum += sg[k] * xg[k];
10830            }
10831            acc += s * gsum;
10832        }
10833        acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
10834        // SAFETY: disjoint row ranges per worker.
10835        unsafe { *out.at(r) = acc };
10836    }
10837}
10838
10839/// Ternary (q1t) matvec — decode+dot straight from mmap, one group at a time:
10840/// no per-ROW buffer, no division (the sign LUT), and a tiny per-group sign
10841/// buffer so the 32-wide dot vectorizes. This is the decode hot path.
10842fn q1t_matvec(
10843    bytes: &[u8],
10844    x: &[f32],
10845    rows: usize,
10846    cols: usize,
10847    out: &mut [f32],
10848    pool: Option<&Pool>,
10849) {
10850    debug_assert_eq!(out.len(), rows);
10851    const TILE: usize = cortiq_core::quant::Q1T_TILE;
10852    let gpr = cols / GROUP_SIZE;
10853    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
10854    let out_addr = SendMut(out.as_mut_ptr());
10855    // int8 SDOT base dot (ARM dotprod): ~4× the f32 arithmetic. x → i8 once
10856    // (`split_act`), activation outliers added back exactly in f32, weight
10857    // overlay on top. ARM SDOT / x86 AVX2; CMF_SDOT=0 keeps the exact f32 path.
10858    if a8w8_enabled() {
10859        let act = split_act(x);
10860        let act = &act;
10861        let run = move |start: usize, end: usize| {
10862            for r in start..end {
10863                let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
10864                for &(j, xv) in &act.outliers {
10865                    acc += q1t_base_weight(bytes, r, gpr, j) * xv;
10866                }
10867                acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
10868                // SAFETY: disjoint row ranges per worker.
10869                unsafe { *out_addr.at(r) = acc };
10870            }
10871        };
10872        dispatch_rows(pool, rows, &run);
10873        return;
10874    }
10875    let run = move |start: usize, end: usize| {
10876        // Per-group signs, unpacked contiguously so the dot below is a clean
10877        // 32-wide reduction the autovectorizer turns into f32x4 FMAs — the
10878        // 5-values-per-byte base-3 layout won't SIMD in place.
10879        let mut sg = [0f32; GROUP_SIZE];
10880        for r in start..end {
10881            let mut acc = 0f32;
10882            for g in 0..gpr {
10883                let off = (r * gpr + g) * TILE;
10884                let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
10885                let codes = &bytes[off + 2..off + TILE];
10886                let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
10887                for bi in 0..6 {
10888                    sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
10889                }
10890                let lut = &SIGN5[codes[6] as usize];
10891                sg[30] = lut[0];
10892                sg[31] = lut[1];
10893                let mut gsum = 0f32;
10894                for k in 0..GROUP_SIZE {
10895                    gsum += sg[k] * xg[k];
10896                }
10897                acc += s * gsum;
10898            }
10899            acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
10900            unsafe { *out_addr.at(r) = acc };
10901        }
10902    };
10903    dispatch_rows(pool, rows, &run);
10904}
10905
10906/// Fused-pair twin of `q1t_dot_row_sdot`: ONE register unpack of the
10907/// ternary codes serves BOTH activation streams (the unpack chain is
10908/// the dominant per-row cost — MTP verify pairs paid it twice). Per
10909/// stream the group order and f32 accumulation match the single-row
10910/// kernel exactly, so pair == 2×matvec bit-for-bit.
10911#[cfg(target_arch = "aarch64")]
10912#[target_feature(enable = "neon,dotprod")]
10913unsafe fn q1t_dot_row_sdot2(bytes: &[u8], r: usize, gpr: usize, xa: &[i8], xb: &[i8]) -> [f32; 2] {
10914    use core::arch::aarch64::*;
10915    use core::arch::asm;
10916    // SAFETY: same slice-length contracts as `q1t_dot_row_sdot`, ×2.
10917    unsafe {
10918        const TILE: usize = cortiq_core::quant::Q1T_TILE;
10919        let bytes_ptr = bytes.as_ptr();
10920        let row_off = r * gpr * TILE;
10921        let xp = [xa.as_ptr(), xb.as_ptr()];
10922        let mut acc = [0f32; 2];
10923        macro_rules! sdot2 {
10924            ($w0:expr, $w1:expr, $x:expr) => {{
10925                let x0 = vld1q_s8($x);
10926                let x1 = vld1q_s8($x.add(16));
10927                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
10928                asm!(
10929                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
10930                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
10931                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
10932                    w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
10933                    options(pure, nomem, nostack),
10934                );
10935                vaddvq_s32(vaddq_s32(a0, a1))
10936            }};
10937        }
10938        let gpr2 = gpr & !1;
10939        let mut gi = 0;
10940        while gi < gpr2 {
10941            let off0 = row_off + gi * TILE;
10942            let off1 = off0 + TILE;
10943            let s0 = f16_to_f32(u16::from_le_bytes([
10944                *bytes_ptr.add(off0),
10945                *bytes_ptr.add(off0 + 1),
10946            ]));
10947            let s1 = f16_to_f32(u16::from_le_bytes([
10948                *bytes_ptr.add(off1),
10949                *bytes_ptr.add(off1 + 1),
10950            ]));
10951            let (u0_0, u1_0, u2_0, u3_0) = q1t_unpack_reg_u64s(bytes_ptr.add(off0 + 2));
10952            let (u0_1, u1_1, u2_1, u3_1) = q1t_unpack_reg_u64s(bytes_ptr.add(off1 + 2));
10953            let w0_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_0), vcreate_u64(u1_0)));
10954            let w1_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_0), vcreate_u64(u3_0)));
10955            let w0_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_1), vcreate_u64(u1_1)));
10956            let w1_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_1), vcreate_u64(u3_1)));
10957            for k in 0..2 {
10958                let d0 = sdot2!(w0_0, w1_0, xp[k].add(gi * GROUP_SIZE));
10959                let d1 = sdot2!(w0_1, w1_1, xp[k].add((gi + 1) * GROUP_SIZE));
10960                acc[k] += d0 as f32 * s0 + d1 as f32 * s1;
10961            }
10962            gi += 2;
10963        }
10964        if gi < gpr {
10965            let off = row_off + gi * TILE;
10966            let s = f16_to_f32(u16::from_le_bytes([
10967                *bytes_ptr.add(off),
10968                *bytes_ptr.add(off + 1),
10969            ]));
10970            let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
10971            let w0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0), vcreate_u64(u1)));
10972            let w1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2), vcreate_u64(u3)));
10973            for k in 0..2 {
10974                let d = sdot2!(w0, w1, xp[k].add(gi * GROUP_SIZE));
10975                acc[k] += d as f32 * s;
10976            }
10977        }
10978        acc
10979    }
10980}
10981
10982/// Fused Q1T pair matvec: ONE pass over the rows serves both
10983/// activation streams — on ARM the ternary register unpack happens
10984/// once per tile pair (`q1t_dot_row_sdot2`); elsewhere the second dot
10985/// rides the row's L1-warm tile bytes. Per stream the math matches
10986/// `q1t_matvec` exactly.
10987fn q1t_matvec2(
10988    bytes: &[u8],
10989    x1: &[f32],
10990    x2: &[f32],
10991    rows: usize,
10992    cols: usize,
10993    o1: &mut [f32],
10994    o2: &mut [f32],
10995    pool: Option<&Pool>,
10996) {
10997    debug_assert_eq!(o1.len(), rows);
10998    debug_assert_eq!(o2.len(), rows);
10999    const TILE: usize = cortiq_core::quant::Q1T_TILE;
11000    let gpr = cols / GROUP_SIZE;
11001    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
11002    let out1 = SendMut(o1.as_mut_ptr());
11003    let out2 = SendMut(o2.as_mut_ptr());
11004    if a8w8_enabled() {
11005        let a1 = split_act(x1);
11006        let a2 = split_act(x2);
11007        let (a1, a2) = (&a1, &a2);
11008        let run = move |start: usize, end: usize| {
11009            for r in start..end {
11010                #[cfg(target_arch = "aarch64")]
11011                // a8w8 on aarch64 ⇔ sdot_enabled(), so the kernel's
11012                // target features are present.
11013                let ds = unsafe { q1t_dot_row_sdot2(bytes, r, gpr, &a1.xq, &a2.xq) };
11014                #[cfg(not(target_arch = "aarch64"))]
11015                let ds = [
11016                    q1t_dot_row_i8(bytes, r, gpr, &a1.xq),
11017                    q1t_dot_row_i8(bytes, r, gpr, &a2.xq),
11018                ];
11019                let mut acc1 = ds[0] * a1.sx;
11020                for &(j, xv) in &a1.outliers {
11021                    acc1 += q1t_base_weight(bytes, r, gpr, j) * xv;
11022                }
11023                acc1 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x1);
11024                let mut acc2 = ds[1] * a2.sx;
11025                for &(j, xv) in &a2.outliers {
11026                    acc2 += q1t_base_weight(bytes, r, gpr, j) * xv;
11027                }
11028                acc2 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x2);
11029                // SAFETY: disjoint row ranges per worker.
11030                unsafe {
11031                    *out1.at(r) = acc1;
11032                    *out2.at(r) = acc2;
11033                }
11034            }
11035        };
11036        dispatch_rows(pool, rows, &run);
11037        return;
11038    }
11039    let run = move |start: usize, end: usize| {
11040        // Exact path (CMF_SDOT=0): unpack the sign LUT once per group,
11041        // dot both streams — same op order per stream as `q1t_matvec`.
11042        let mut sg = [0f32; GROUP_SIZE];
11043        for r in start..end {
11044            let mut acc1 = 0f32;
11045            let mut acc2 = 0f32;
11046            for g in 0..gpr {
11047                let off = (r * gpr + g) * TILE;
11048                let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
11049                let codes = &bytes[off + 2..off + TILE];
11050                for bi in 0..6 {
11051                    sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
11052                }
11053                let lut = &SIGN5[codes[6] as usize];
11054                sg[30] = lut[0];
11055                sg[31] = lut[1];
11056                let xg1 = &x1[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
11057                let xg2 = &x2[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
11058                let mut gsum1 = 0f32;
11059                for k in 0..GROUP_SIZE {
11060                    gsum1 += sg[k] * xg1[k];
11061                }
11062                acc1 += s * gsum1;
11063                let mut gsum2 = 0f32;
11064                for k in 0..GROUP_SIZE {
11065                    gsum2 += sg[k] * xg2[k];
11066                }
11067                acc2 += s * gsum2;
11068            }
11069            acc1 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x1);
11070            acc2 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x2);
11071            // SAFETY: disjoint row ranges per worker.
11072            unsafe {
11073                *out1.at(r) = acc1;
11074                *out2.at(r) = acc2;
11075            }
11076        }
11077    };
11078    dispatch_rows(pool, rows, &run);
11079}
11080
11081/// Ternary (q1t) matmat (prefill) — dequant each row once, dot the whole
11082/// batch against it (amortizes the per-row decode).
11083fn q1t_matmat(
11084    bytes: &[u8],
11085    xs: &[f32],
11086    b: usize,
11087    rows: usize,
11088    cols: usize,
11089    out: &mut [f32],
11090    pool: Option<&Pool>,
11091) {
11092    debug_assert_eq!(out.len(), b * rows);
11093    const TILE: usize = cortiq_core::quant::Q1T_TILE;
11094    let gpr = cols / GROUP_SIZE;
11095    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
11096    let out_addr = SendMut(out.as_mut_ptr());
11097    // int8 prefill (ARM SDOT / x86 AVX2): quantize the B inputs once, unpack
11098    // each weight row's signs to i8 ONCE, then int8-dot against every input —
11099    // the row sign-decode amortizes over the whole batch. CMF_SDOT=0 → f32.
11100    if a8w8_enabled() {
11101        let acts: Vec<SplitAct> = (0..b)
11102            .map(|bi| split_act(&xs[bi * cols..(bi + 1) * cols]))
11103            .collect();
11104        let acts = &acts;
11105        let run = move |start: usize, end: usize| {
11106            let mut sg = vec![0i8; cols + 8]; // row signs, i8 (+8 unpack slack)
11107            let mut sc = vec![0f32; gpr]; // per-group scales
11108            let mut accs = vec![0f32; b]; // per-batch accumulators, reused per row
11109            for r in start..end {
11110                for g in 0..gpr {
11111                    let off = (r * gpr + g) * TILE;
11112                    sc[g] = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
11113                    q1t_unpack_group_i8(
11114                        bytes.as_ptr().wrapping_add(off + 2),
11115                        &mut sg[g * GROUP_SIZE..],
11116                    );
11117                }
11118                for bi in 0..b {
11119                    let act = &acts[bi];
11120                    let mut isum = 0f32;
11121                    for g in 0..gpr {
11122                        let d = q1t_i8dot32(
11123                            sg.as_ptr().wrapping_add(g * GROUP_SIZE),
11124                            act.xq.as_ptr().wrapping_add(g * GROUP_SIZE),
11125                        );
11126                        isum += d as f32 * sc[g];
11127                    }
11128                    let mut acc = isum * act.sx;
11129                    for &(j, xv) in &act.outliers {
11130                        acc += q1t_base_weight(bytes, r, gpr, j) * xv;
11131                    }
11132                    accs[bi] = acc;
11133                }
11134                // Overlay ONCE per row for the whole batch: read each (col, val)
11135                // from mmap a single time (was b× — the re-read dominated prefill)
11136                // and fan it out over the batch via the cached inputs.
11137                if has_ov {
11138                    let (c0, c1) = (
11139                        q1t_rowptr(bytes, rp_off, r),
11140                        q1t_rowptr(bytes, rp_off, r + 1),
11141                    );
11142                    for p in c0..c1 {
11143                        let e = ent_off + p * 4;
11144                        let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
11145                        let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
11146                        for bi in 0..b {
11147                            accs[bi] += val * xs[bi * cols + col];
11148                        }
11149                    }
11150                }
11151                for bi in 0..b {
11152                    unsafe { *out_addr.at(bi * rows + r) = accs[bi] };
11153                }
11154            }
11155        };
11156        dispatch_rows(pool, rows, &run);
11157        return;
11158    }
11159    let run = move |start: usize, end: usize| {
11160        let mut buf = vec![0f32; cols];
11161        for r in start..end {
11162            q1t_dequant_row(bytes, r, gpr, rp_off, ent_off, has_ov, &mut buf);
11163            for bi in 0..b {
11164                let xr = &xs[bi * cols..(bi + 1) * cols];
11165                let mut acc = 0f32;
11166                for j in 0..cols {
11167                    acc += buf[j] * xr[j];
11168                }
11169                unsafe { *out_addr.at(bi * rows + r) = acc };
11170            }
11171        }
11172    };
11173    dispatch_rows(pool, rows, &run);
11174}
11175
11176/// q1-tile entry points of the binary drivers (the tests' handles; the
11177/// runtime goes through `QTensor::bin_w`).
11178#[cfg(test)]
11179fn q1_matvec(
11180    bytes: &[u8],
11181    x: &[f32],
11182    rows: usize,
11183    cols: usize,
11184    out: &mut [f32],
11185    pool: Option<&Pool>,
11186) {
11187    bin_matvec(BinW::Q1(bytes), x, rows, cols, out, pool)
11188}
11189
11190/// Binary-weight matvec (q1 or q1g): A8W8 rows with the shared group
11191/// sums, or the exact scalar walk under CMF_SDOT=0.
11192fn bin_matvec(w: BinW, x: &[f32], rows: usize, cols: usize, out: &mut [f32], pool: Option<&Pool>) {
11193    debug_assert_eq!(out.len(), rows);
11194    let gpr = cols / GROUP_SIZE;
11195    let out_addr = SendMut(out.as_mut_ptr());
11196    if a8w8_enabled() {
11197        let act = split_act(x);
11198        let gsum = q1_group_sums(&act.xq, gpr);
11199        let (act, gsum) = (&act, &gsum);
11200        let run =
11201            move |start: usize, end: usize| bin_range_a8w8(w, gpr, act, gsum, out_addr, start, end);
11202        dispatch_rows(pool, rows, &run);
11203        return;
11204    }
11205    let run = move |start: usize, end: usize| bin_range_f32(w, gpr, x, out_addr, start, end);
11206    dispatch_rows(pool, rows, &run);
11207}
11208
11209/// Fused two-input q1 matvec (weights read once per pair).
11210#[cfg(test)]
11211#[allow(clippy::too_many_arguments)]
11212fn q1_matvec2(
11213    bytes: &[u8],
11214    x1: &[f32],
11215    x2: &[f32],
11216    rows: usize,
11217    cols: usize,
11218    o1: &mut [f32],
11219    o2: &mut [f32],
11220    pool: Option<&Pool>,
11221) {
11222    bin_matvec2(BinW::Q1(bytes), x1, x2, rows, cols, o1, o2, pool)
11223}
11224
11225/// Fused two-input binary-weight matvec (q1 or q1g).
11226#[allow(clippy::too_many_arguments)]
11227fn bin_matvec2(
11228    w: BinW,
11229    x1: &[f32],
11230    x2: &[f32],
11231    rows: usize,
11232    cols: usize,
11233    o1: &mut [f32],
11234    o2: &mut [f32],
11235    pool: Option<&Pool>,
11236) {
11237    let gpr = cols / GROUP_SIZE;
11238    let p1 = SendMut(o1.as_mut_ptr());
11239    let p2 = SendMut(o2.as_mut_ptr());
11240    if a8w8_enabled() {
11241        let a1 = split_act(x1);
11242        let a2 = split_act(x2);
11243        let g1 = q1_group_sums(&a1.xq, gpr);
11244        let g2 = q1_group_sums(&a2.xq, gpr);
11245        let (a1, a2, g1, g2) = (&a1, &a2, &g1, &g2);
11246        let run = move |start: usize, end: usize| {
11247            for r in start..end {
11248                let mut v1 = w.dot_i8(r, gpr, &a1.xq, g1) * a1.sx;
11249                let mut v2 = w.dot_i8(r, gpr, &a2.xq, g2) * a2.sx;
11250                for &(j, xv) in &a1.outliers {
11251                    let (wv, s) = w.outlier(r, gpr, j);
11252                    v1 += wv * s * xv;
11253                }
11254                for &(j, xv) in &a2.outliers {
11255                    let (wv, s) = w.outlier(r, gpr, j);
11256                    v2 += wv * s * xv;
11257                }
11258                // SAFETY: disjoint row ranges per worker.
11259                unsafe {
11260                    *p1.at(r) = v1;
11261                    *p2.at(r) = v2;
11262                }
11263            }
11264        };
11265        dispatch_rows(pool, rows, &run);
11266        return;
11267    }
11268    let run = move |start: usize, end: usize| {
11269        for r in start..end {
11270            // SAFETY: disjoint row ranges per worker.
11271            unsafe {
11272                *p1.at(r) = w.row_exact(r, gpr, x1);
11273                *p2.at(r) = w.row_exact(r, gpr, x2);
11274            }
11275        }
11276    };
11277    dispatch_rows(pool, rows, &run);
11278}
11279
11280/// Batched q1 matmat: each row's tiles stream once per microbatch.
11281#[cfg(test)]
11282#[allow(clippy::too_many_arguments)]
11283fn q1_matmat(
11284    bytes: &[u8],
11285    xs_all: &[f32],
11286    b: usize,
11287    rows: usize,
11288    cols: usize,
11289    out: &mut [f32],
11290    pool: Option<&Pool>,
11291) {
11292    bin_matmat(BinW::Q1(bytes), xs_all, b, rows, cols, out, pool)
11293}
11294
11295/// Batched binary-weight matmat (q1 or q1g).
11296#[allow(clippy::too_many_arguments)]
11297fn bin_matmat(
11298    w: BinW,
11299    xs_all: &[f32],
11300    b: usize,
11301    rows: usize,
11302    cols: usize,
11303    out: &mut [f32],
11304    pool: Option<&Pool>,
11305) {
11306    debug_assert_eq!(out.len(), b * rows);
11307    let gpr = cols / GROUP_SIZE;
11308    let out_addr = SendMut(out.as_mut_ptr());
11309    if a8w8_enabled() {
11310        let acts: Vec<(SplitAct, Vec<i32>)> = (0..b)
11311            .map(|bi| {
11312                let act = split_act(&xs_all[bi * cols..(bi + 1) * cols]);
11313                let gsum = q1_group_sums(&act.xq, gpr);
11314                (act, gsum)
11315            })
11316            .collect();
11317        let acts = &acts;
11318        #[cfg(target_arch = "x86_64")]
11319        let blocked_ok = avx2_enabled() && blocked_enabled();
11320        #[cfg(target_arch = "aarch64")]
11321        let blocked_ok = sdot_enabled() && blocked_enabled();
11322        let run = move |start: usize, end: usize| {
11323            for r in start..end {
11324                let mut bi = 0usize;
11325                // Blocked 1×4: the unpacked bit mask serves four
11326                // activation streams per group.
11327                #[cfg(target_arch = "aarch64")]
11328                if blocked_ok {
11329                    while bi + 4 <= acts.len() {
11330                        let xs = [
11331                            acts[bi].0.xq.as_slice(),
11332                            acts[bi + 1].0.xq.as_slice(),
11333                            acts[bi + 2].0.xq.as_slice(),
11334                            acts[bi + 3].0.xq.as_slice(),
11335                        ];
11336                        let gs = [
11337                            acts[bi].1.as_slice(),
11338                            acts[bi + 1].1.as_slice(),
11339                            acts[bi + 2].1.as_slice(),
11340                            acts[bi + 3].1.as_slice(),
11341                        ];
11342                        // SAFETY: blocked_ok checked SDOT.
11343                        let d = unsafe { w.dot_1x4_sdot(r, gpr, xs, gs) };
11344                        for k in 0..4 {
11345                            let (act, _) = &acts[bi + k];
11346                            let mut acc = d[k] * act.sx;
11347                            for &(j, xv) in &act.outliers {
11348                                let (wv, sc) = w.outlier(r, gpr, j);
11349                                acc += wv * sc * xv;
11350                            }
11351                            // SAFETY: disjoint (bi, r) cells per worker.
11352                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
11353                        }
11354                        bi += 4;
11355                    }
11356                }
11357                #[cfg(target_arch = "x86_64")]
11358                if blocked_ok {
11359                    while bi + 4 <= acts.len() {
11360                        let xs = [
11361                            acts[bi].0.xq.as_slice(),
11362                            acts[bi + 1].0.xq.as_slice(),
11363                            acts[bi + 2].0.xq.as_slice(),
11364                            acts[bi + 3].0.xq.as_slice(),
11365                        ];
11366                        let gs = [
11367                            acts[bi].1.as_slice(),
11368                            acts[bi + 1].1.as_slice(),
11369                            acts[bi + 2].1.as_slice(),
11370                            acts[bi + 3].1.as_slice(),
11371                        ];
11372                        // SAFETY: blocked_ok checked AVX2.
11373                        let d = unsafe { w.dot_1x4_x86(r, gpr, xs, gs) };
11374                        for k in 0..4 {
11375                            let (act, _) = &acts[bi + k];
11376                            let mut acc = d[k] * act.sx;
11377                            for &(j, xv) in &act.outliers {
11378                                let (wv, sc) = w.outlier(r, gpr, j);
11379                                acc += wv * sc * xv;
11380                            }
11381                            // SAFETY: disjoint (bi, r) cells per worker.
11382                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
11383                        }
11384                        bi += 4;
11385                    }
11386                }
11387                while bi < acts.len() {
11388                    let (act, gsum) = &acts[bi];
11389                    let mut acc = w.dot_i8(r, gpr, &act.xq, gsum) * act.sx;
11390                    for &(j, xv) in &act.outliers {
11391                        let (wv, s) = w.outlier(r, gpr, j);
11392                        acc += wv * s * xv;
11393                    }
11394                    // SAFETY: disjoint (bi, r) cells per worker range.
11395                    unsafe { *out_addr.at(bi * rows + r) = acc };
11396                    bi += 1;
11397                }
11398            }
11399        };
11400        dispatch_rows(pool, rows, &run);
11401        return;
11402    }
11403    let run = move |start: usize, end: usize| {
11404        for r in start..end {
11405            for bi in 0..b {
11406                let x = &xs_all[bi * cols..(bi + 1) * cols];
11407                // SAFETY: disjoint (bi, r) cells per worker range.
11408                unsafe { *out_addr.at(bi * rows + r) = w.row_exact(r, gpr, x) };
11409            }
11410        }
11411    };
11412    dispatch_rows(pool, rows, &run);
11413}
11414
11415/// Fused q4_block matvec straight from the mapped bytes. SDOT path when
11416/// dotprod is available (port of vmfcore `dot_q4_block_sdot`, measured
11417/// +23% on q4 decode): nibbles → centered i8, int8×int8 `sdot` per
11418/// 32-group, exact outlier correction — the same A8W8 contract as q8.
11419/// `CMF_SDOT=0` keeps the exact scalar path.
11420fn q4matvec(
11421    bytes: &[u8],
11422    x: &[f32],
11423    rows: usize,
11424    cols: usize,
11425    out: &mut [f32],
11426    pool: Option<&Pool>,
11427) {
11428    debug_assert_eq!(out.len(), rows);
11429    let (packed, scales) = q4_split(bytes, rows, cols);
11430    let gpr = cols / GROUP_SIZE;
11431    let out_addr = SendMut(out.as_mut_ptr());
11432
11433    if a8w8_enabled() {
11434        let act = split_act(x);
11435        let run = move |start: usize, end: usize| {
11436            q4_range_a8w8(packed, scales, gpr, cols, &act, out_addr, start, end)
11437        };
11438        dispatch_rows(pool, rows, &run);
11439        return;
11440    }
11441
11442    let run =
11443        move |start: usize, end: usize| q4_range_f32(packed, scales, gpr, x, out_addr, start, end);
11444    dispatch_rows(pool, rows, &run);
11445}
11446
11447/// One q4 row via the A8W8 int8 path — SDOT on ARM, AVX2 maddubs on
11448/// x86 (scalar fallback is unreachable: callers gate on a8w8_enabled).
11449#[inline]
11450#[allow(unreachable_code)]
11451/// One UNPACKED q4 row (centered i8 in `buf`) against four activation
11452/// streams: the 32-byte weight chunk and its abs() load once per group,
11453/// the per-group f16 scale decodes once — four maddubs+reduce chains
11454/// instead of four full (load, abs, dot) rounds.
11455#[cfg(target_arch = "x86_64")]
11456#[target_feature(enable = "avx2")]
11457unsafe fn dot_q4b_row_1x4_avx2(
11458    buf: &[u8],
11459    scales: &[u8],
11460    g0: usize,
11461    gpr: usize,
11462    xs: [&[i8]; 4],
11463) -> [f32; 4] {
11464    // SAFETY: callers uphold buffer contracts (buf.len() == gpr·32).
11465    unsafe {
11466        use core::arch::x86_64::*;
11467        let ones = _mm256_set1_epi16(1);
11468        let mut acc = [0f32; 4];
11469        for gi in 0..gpr {
11470            let s = f16_to_f32(u16::from_le_bytes([
11471                scales[(g0 + gi) * 2],
11472                scales[(g0 + gi) * 2 + 1],
11473            ]));
11474            let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
11475            let aw = _mm256_abs_epi8(w);
11476            for (k, xq) in xs.iter().enumerate() {
11477                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
11478                let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
11479                let d = _mm256_madd_epi16(p16, ones);
11480                let hi128 = _mm256_extracti128_si256::<1>(d);
11481                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
11482                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
11483                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
11484                acc[k] += _mm_cvtsi128_si32(s32) as f32 * s;
11485            }
11486        }
11487        acc
11488    }
11489}
11490
11491/// VNNI twin of `dot_q4b_row_1x4_avx2` (see `dpbusd_hsum`).
11492#[cfg(target_arch = "x86_64")]
11493#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
11494unsafe fn dot_q4b_row_1x4_vnni(
11495    buf: &[u8],
11496    scales: &[u8],
11497    g0: usize,
11498    gpr: usize,
11499    xs: [&[i8]; 4],
11500) -> [f32; 4] {
11501    // SAFETY: callers uphold buffer contracts (buf.len() == gpr·32).
11502    unsafe {
11503        use core::arch::x86_64::*;
11504        let mut acc = [0f32; 4];
11505        for gi in 0..gpr {
11506            let s = f16_to_f32(u16::from_le_bytes([
11507                scales[(g0 + gi) * 2],
11508                scales[(g0 + gi) * 2 + 1],
11509            ]));
11510            let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
11511            let aw = _mm256_abs_epi8(w);
11512            for (k, xq) in xs.iter().enumerate() {
11513                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
11514                let d = dpbusd_hsum(aw, _mm256_sign_epi8(x, w));
11515                acc[k] += d as f32 * s;
11516            }
11517        }
11518        acc
11519    }
11520}
11521
11522/// The vbit flavor of the blocked 1×4: the per-activation A8W8 scale
11523/// folds in PER GROUP as `(d·sx)·s` — bit-matching the single-matvec
11524/// accumulation order (the q4_block flavor applies sx once at the end,
11525/// matching ITS single path; the two conventions are historical and
11526/// each blocked leg must mirror its own).
11527#[cfg(target_arch = "x86_64")]
11528#[target_feature(enable = "avx2")]
11529unsafe fn dot_q4b_row_1x4_sx_avx2(
11530    buf: &[u8],
11531    scales: &[u8],
11532    g0: usize,
11533    gpr: usize,
11534    xs: [&[i8]; 4],
11535    sxs: [f32; 4],
11536) -> [f32; 4] {
11537    // SAFETY: callers uphold buffer contracts (buf.len() == gpr·32).
11538    unsafe {
11539        use core::arch::x86_64::*;
11540        let ones = _mm256_set1_epi16(1);
11541        let mut acc = [0f32; 4];
11542        for gi in 0..gpr {
11543            let s = f16_to_f32(u16::from_le_bytes([
11544                scales[(g0 + gi) * 2],
11545                scales[(g0 + gi) * 2 + 1],
11546            ]));
11547            let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
11548            let aw = _mm256_abs_epi8(w);
11549            for (k, xq) in xs.iter().enumerate() {
11550                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
11551                let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
11552                let d = _mm256_madd_epi16(p16, ones);
11553                let hi128 = _mm256_extracti128_si256::<1>(d);
11554                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
11555                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
11556                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
11557                acc[k] += (_mm_cvtsi128_si32(s32) as f32 * sxs[k]) * s;
11558            }
11559        }
11560        acc
11561    }
11562}
11563
11564/// VNNI twin of `dot_q4b_row_1x4_sx_avx2` (see `dpbusd_hsum`; the
11565/// per-group `(d·sx)·s` fold mirrors the vbit single path).
11566#[cfg(target_arch = "x86_64")]
11567#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
11568unsafe fn dot_q4b_row_1x4_sx_vnni(
11569    buf: &[u8],
11570    scales: &[u8],
11571    g0: usize,
11572    gpr: usize,
11573    xs: [&[i8]; 4],
11574    sxs: [f32; 4],
11575) -> [f32; 4] {
11576    // SAFETY: callers uphold buffer contracts (buf.len() == gpr·32).
11577    unsafe {
11578        use core::arch::x86_64::*;
11579        let mut acc = [0f32; 4];
11580        for gi in 0..gpr {
11581            let s = f16_to_f32(u16::from_le_bytes([
11582                scales[(g0 + gi) * 2],
11583                scales[(g0 + gi) * 2 + 1],
11584            ]));
11585            let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
11586            let aw = _mm256_abs_epi8(w);
11587            for (k, xq) in xs.iter().enumerate() {
11588                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
11589                let d = dpbusd_hsum(aw, _mm256_sign_epi8(x, w));
11590                acc[k] += (d as f32 * sxs[k]) * s;
11591            }
11592        }
11593        acc
11594    }
11595}
11596
11597#[allow(unreachable_code)]
11598fn dot_q4_row_i8(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
11599    #[cfg(target_arch = "aarch64")]
11600    unsafe {
11601        return dot_q4_row_sdot(packed, scales, g0, gpr, xq);
11602    }
11603    #[cfg(target_arch = "x86_64")]
11604    unsafe {
11605        return dot_q4_row_avx2(packed, scales, g0, gpr, xq);
11606    }
11607    let mut acc = 0f32;
11608    for gi in 0..gpr {
11609        let g = g0 + gi;
11610        let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
11611        let mut d = 0i32;
11612        for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
11613            d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
11614                + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
11615        }
11616        acc += d as f32 * s;
11617    }
11618    acc
11619}
11620
11621/// Two-activation q4 row via the A8W8 int8 path (see `dot_q4_row_i8`).
11622#[inline]
11623#[allow(unreachable_code)]
11624fn dot_q4_row_i8_2(
11625    packed: &[u8],
11626    scales: &[u8],
11627    g0: usize,
11628    gpr: usize,
11629    xq1: &[i8],
11630    xq2: &[i8],
11631) -> (f32, f32) {
11632    #[cfg(target_arch = "aarch64")]
11633    unsafe {
11634        return dot_q4_row_sdot2(packed, scales, g0, gpr, xq1, xq2);
11635    }
11636    #[cfg(target_arch = "x86_64")]
11637    unsafe {
11638        return dot_q4_row_avx2_2(packed, scales, g0, gpr, xq1, xq2);
11639    }
11640    (
11641        dot_q4_row_i8(packed, scales, g0, gpr, xq1),
11642        dot_q4_row_i8(packed, scales, g0, gpr, xq2),
11643    )
11644}
11645
11646/// One q4 row range via SDOT (kernel body of `q4matvec`, extracted so
11647/// multi-matrix jobs can drive it for several tensors in one dispatch).
11648#[allow(clippy::too_many_arguments)]
11649fn q4_range_a8w8(
11650    packed: &[u8],
11651    scales: &[u8],
11652    gpr: usize,
11653    cols: usize,
11654    act: &SplitAct,
11655    out: SendMut,
11656    start: usize,
11657    end: usize,
11658) {
11659    for r in start..end {
11660        let mut acc = dot_q4_row_i8(packed, scales, r * gpr, gpr, &act.xq) * act.sx;
11661        // xq is zeroed at outlier slots — add the exact terms.
11662        for &(j, xv) in &act.outliers {
11663            let flat = r * cols + j;
11664            let byte = packed[flat / 2];
11665            let nib = if flat & 1 == 0 {
11666                byte & 0x0F
11667            } else {
11668                byte >> 4
11669            };
11670            let s = f16_to_f32(u16::from_le_bytes([
11671                scales[(flat / GROUP_SIZE) * 2],
11672                scales[(flat / GROUP_SIZE) * 2 + 1],
11673            ]));
11674            acc += ((nib as i32 - 8) as f32) * s * xv;
11675        }
11676        // SAFETY: disjoint row ranges per worker.
11677        unsafe { *out.at(r) = acc };
11678    }
11679}
11680
11681/// Two-input q4 row range via the A8W8 int8 path — kernel body of
11682/// `q4matvec2`, extracted for pair multi-matrix jobs.
11683#[allow(clippy::too_many_arguments)]
11684fn q4_range2_a8w8(
11685    packed: &[u8],
11686    scales: &[u8],
11687    gpr: usize,
11688    cols: usize,
11689    a1: &SplitAct,
11690    a2: &SplitAct,
11691    p1: SendMut,
11692    p2: SendMut,
11693    start: usize,
11694    end: usize,
11695) {
11696    for r in start..end {
11697        let (s1, s2) = dot_q4_row_i8_2(packed, scales, r * gpr, gpr, &a1.xq, &a2.xq);
11698        let mut acc1 = s1 * a1.sx;
11699        let mut acc2 = s2 * a2.sx;
11700        // xq is zeroed at outlier slots — add the exact terms.
11701        let fix = |outliers: &[(usize, f32)], acc: &mut f32| {
11702            for &(j, xv) in outliers {
11703                let flat = r * cols + j;
11704                let byte = packed[flat / 2];
11705                let nib = if flat & 1 == 0 {
11706                    byte & 0x0F
11707                } else {
11708                    byte >> 4
11709                };
11710                let s = f16_to_f32(u16::from_le_bytes([
11711                    scales[(flat / GROUP_SIZE) * 2],
11712                    scales[(flat / GROUP_SIZE) * 2 + 1],
11713                ]));
11714                *acc += ((nib as i32 - 8) as f32) * s * xv;
11715            }
11716        };
11717        fix(&a1.outliers, &mut acc1);
11718        fix(&a2.outliers, &mut acc2);
11719        // SAFETY: disjoint row ranges per worker.
11720        unsafe {
11721            *p1.at(r) = acc1;
11722            *p2.at(r) = acc2;
11723        }
11724    }
11725}
11726
11727/// Exact scalar q4 row range (same extraction, non-SDOT path).
11728fn q4_range_f32(
11729    packed: &[u8],
11730    scales: &[u8],
11731    gpr: usize,
11732    x: &[f32],
11733    out: SendMut,
11734    start: usize,
11735    end: usize,
11736) {
11737    for r in start..end {
11738        let mut acc = 0f32;
11739        for gi in 0..gpr {
11740            let g = r * gpr + gi;
11741            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
11742            let pk = &packed[g * 16..(g + 1) * 16];
11743            let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
11744            let mut ga = 0f32;
11745            for (k, &b) in pk.iter().enumerate() {
11746                ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
11747                    + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
11748            }
11749            acc += ga * s;
11750        }
11751        // SAFETY: disjoint row ranges per worker.
11752        unsafe { *out.at(r) = acc };
11753    }
11754}
11755
11756/// Fused two-input q4 matvec: nibbles are unpacked ONCE per group and
11757/// dotted against both activations (was: two full matvecs — double
11758/// weight traffic). Per-lane math matches `q4matvec` exactly.
11759#[allow(clippy::too_many_arguments)]
11760fn q4matvec2(
11761    bytes: &[u8],
11762    x1: &[f32],
11763    x2: &[f32],
11764    rows: usize,
11765    cols: usize,
11766    o1: &mut [f32],
11767    o2: &mut [f32],
11768    pool: Option<&Pool>,
11769) {
11770    debug_assert_eq!(o1.len(), rows);
11771    debug_assert_eq!(o2.len(), rows);
11772    let (packed, scales) = q4_split(bytes, rows, cols);
11773    let gpr = cols / GROUP_SIZE;
11774
11775    if a8w8_enabled() {
11776        let a1 = split_act(x1);
11777        let a2 = split_act(x2);
11778        let p1 = SendMut(o1.as_mut_ptr());
11779        let p2 = SendMut(o2.as_mut_ptr());
11780        let run = move |start: usize, end: usize| {
11781            q4_range2_a8w8(packed, scales, gpr, cols, &a1, &a2, p1, p2, start, end)
11782        };
11783        dispatch_rows(pool, rows, &run);
11784        return;
11785    }
11786
11787    let p1 = SendMut(o1.as_mut_ptr());
11788    let p2 = SendMut(o2.as_mut_ptr());
11789    let run = move |start: usize, end: usize| {
11790        q4_range2_f32(packed, scales, gpr, x1, x2, p1, p2, start, end)
11791    };
11792    dispatch_rows(pool, rows, &run);
11793}
11794
11795/// Two-input exact scalar q4 row range (same extraction).
11796#[allow(clippy::too_many_arguments)]
11797fn q4_range2_f32(
11798    packed: &[u8],
11799    scales: &[u8],
11800    gpr: usize,
11801    x1: &[f32],
11802    x2: &[f32],
11803    p1: SendMut,
11804    p2: SendMut,
11805    start: usize,
11806    end: usize,
11807) {
11808    for r in start..end {
11809        let (mut acc1, mut acc2) = (0f32, 0f32);
11810        for gi in 0..gpr {
11811            let g = r * gpr + gi;
11812            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
11813            let pk = &packed[g * 16..(g + 1) * 16];
11814            let x1g = &x1[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
11815            let x2g = &x2[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
11816            let (mut g1, mut g2) = (0f32, 0f32);
11817            for (k, &b) in pk.iter().enumerate() {
11818                let wl = (b & 0x0F) as f32 - 8.0;
11819                let wh = ((b >> 4) & 0x0F) as f32 - 8.0;
11820                g1 += wl * x1g[k * 2] + wh * x1g[k * 2 + 1];
11821                g2 += wl * x2g[k * 2] + wh * x2g[k * 2 + 1];
11822            }
11823            acc1 += g1 * s;
11824            acc2 += g2 * s;
11825        }
11826        // SAFETY: disjoint row ranges per worker.
11827        unsafe {
11828            *p1.at(r) = acc1;
11829            *p2.at(r) = acc2;
11830        }
11831    }
11832}
11833
11834thread_local! {
11835    /// Per-worker decoded-row scratch for the batched q4/vbit kernels
11836    /// (centered i8 for SDOT, f32 for the exact/scalar paths).
11837    static ROW_I8: std::cell::RefCell<Vec<u8>> = const { std::cell::RefCell::new(Vec::new()) };
11838    static ROW_F32: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
11839}
11840
11841/// Batched q4 matmat: each weight row is unpacked from the mmap ONCE
11842/// and dotted against ALL b activations (prefill used to fall back to b
11843/// full matvecs — b× weight traffic and b× nibble decode). Per-position
11844/// math matches `q4matvec` exactly: same group order, same accumulation.
11845/// `out` is row-major [b, rows] like `qmatmat`.
11846#[allow(clippy::too_many_arguments)]
11847fn q4matmat(
11848    bytes: &[u8],
11849    xs_all: &[f32],
11850    b: usize,
11851    rows: usize,
11852    cols: usize,
11853    out: &mut [f32],
11854    pool: Option<&Pool>,
11855) {
11856    debug_assert_eq!(xs_all.len(), b * cols);
11857    debug_assert_eq!(out.len(), b * rows);
11858    let (packed, scales) = q4_split(bytes, rows, cols);
11859    let gpr = cols / GROUP_SIZE;
11860    let gscale = |g: usize| f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
11861
11862    if a8w8_enabled() {
11863        let acts: Vec<SplitAct> = (0..b)
11864            .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
11865            .collect();
11866        let acts = &acts;
11867        let out_addr = SendMut(out.as_mut_ptr());
11868        let run = move |start: usize, end: usize| {
11869            ROW_I8.with(|rb| {
11870                let mut buf = rb.borrow_mut();
11871                buf.resize(cols, 0);
11872                for r in start..end {
11873                    // Unpack the row's nibbles to centered i8 once
11874                    // (element 2k = low nibble, 2k+1 = high — flat order,
11875                    // same as dot_q4_row_sdot's zip).
11876                    for gi in 0..gpr {
11877                        let g = r * gpr + gi;
11878                        for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
11879                            buf[gi * GROUP_SIZE + k * 2] = ((bt & 0x0F) as i32 - 8) as i8 as u8;
11880                            buf[gi * GROUP_SIZE + k * 2 + 1] =
11881                                (((bt >> 4) & 0x0F) as i32 - 8) as i8 as u8;
11882                        }
11883                    }
11884                    let mut bi = 0usize;
11885                    #[cfg(target_arch = "x86_64")]
11886                    if avx2_enabled() && blocked_enabled() {
11887                        while bi + 4 <= acts.len() {
11888                            let xs = [
11889                                acts[bi].xq.as_slice(),
11890                                acts[bi + 1].xq.as_slice(),
11891                                acts[bi + 2].xq.as_slice(),
11892                                acts[bi + 3].xq.as_slice(),
11893                            ];
11894                            let d = unsafe {
11895                                if vnni_tiles_enabled() {
11896                                    dot_q4b_row_1x4_vnni(&buf, scales, r * gpr, gpr, xs)
11897                                } else {
11898                                    dot_q4b_row_1x4_avx2(&buf, scales, r * gpr, gpr, xs)
11899                                }
11900                            };
11901                            for k in 0..4 {
11902                                let act = &acts[bi + k];
11903                                let mut acc = d[k] * act.sx;
11904                                for &(j, xv) in &act.outliers {
11905                                    acc += (buf[j] as i8) as f32
11906                                        * gscale((r * cols + j) / GROUP_SIZE)
11907                                        * xv;
11908                                }
11909                                // SAFETY: disjoint (bi, r) cells per worker.
11910                                unsafe { *out_addr.at((bi + k) * rows + r) = acc };
11911                            }
11912                            bi += 4;
11913                        }
11914                    }
11915                    while bi < acts.len() {
11916                        let act = &acts[bi];
11917                        let mut acc = 0f32;
11918                        for gi in 0..gpr {
11919                            let d = dot_i8_i8(
11920                                &buf[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
11921                                &act.xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
11922                            );
11923                            acc += d as f32 * gscale(r * gpr + gi);
11924                        }
11925                        acc *= act.sx;
11926                        // xq is zeroed at outlier slots — exact terms.
11927                        for &(j, xv) in &act.outliers {
11928                            acc += (buf[j] as i8) as f32 * gscale((r * cols + j) / GROUP_SIZE) * xv;
11929                        }
11930                        // SAFETY: disjoint (bi, r) cells per worker row range.
11931                        unsafe { *out_addr.at(bi * rows + r) = acc };
11932                        bi += 1;
11933                    }
11934                }
11935            })
11936        };
11937        dispatch_rows(pool, rows, &run);
11938        return;
11939    }
11940
11941    let out_addr = SendMut(out.as_mut_ptr());
11942    let run = move |start: usize, end: usize| {
11943        ROW_F32.with(|rb| {
11944            let mut buf = rb.borrow_mut();
11945            buf.resize(cols, 0.0);
11946            for r in start..end {
11947                // Decode raw (nib − 8) values once; scales stay per-group
11948                // so the accumulation order matches q4matvec bit-for-bit.
11949                for gi in 0..gpr {
11950                    let g = r * gpr + gi;
11951                    for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
11952                        buf[gi * GROUP_SIZE + k * 2] = (bt & 0x0F) as f32 - 8.0;
11953                        buf[gi * GROUP_SIZE + k * 2 + 1] = ((bt >> 4) & 0x0F) as f32 - 8.0;
11954                    }
11955                }
11956                for bi in 0..b {
11957                    let x = &xs_all[bi * cols..(bi + 1) * cols];
11958                    let mut acc = 0f32;
11959                    for gi in 0..gpr {
11960                        let mut ga = 0f32;
11961                        // Pairwise (lo + hi) addition, matching
11962                        // q4matvec's `ga += lo·x + hi·x` shape exactly —
11963                        // a flat one-per-element loop rounds differently
11964                        // and broke bit-parity on the scalar (x86) path.
11965                        for k in 0..GROUP_SIZE / 2 {
11966                            let e = gi * GROUP_SIZE + k * 2;
11967                            ga += buf[e] * x[e] + buf[e + 1] * x[e + 1];
11968                        }
11969                        acc += ga * gscale(r * gpr + gi);
11970                    }
11971                    // SAFETY: disjoint (bi, r) cells per worker row range.
11972                    unsafe { *out_addr.at(bi * rows + r) = acc };
11973                }
11974            }
11975        })
11976    };
11977    dispatch_rows(pool, rows, &run);
11978}
11979
11980/// Batched vbit matmat: each variable-bit row is decoded from the mmap
11981/// ONCE for the whole microbatch. Same per-position math as
11982/// `vbitmatvec` (SDOT A8W8 with exact outliers / exact f32 for b=8 rows
11983/// and the scalar path).
11984#[allow(clippy::too_many_arguments)]
11985fn vbitmatmat(
11986    bytes: &[u8],
11987    offsets: &[usize],
11988    xs_all: &[f32],
11989    b: usize,
11990    rows: usize,
11991    cols: usize,
11992    out: &mut [f32],
11993    pool: Option<&Pool>,
11994) {
11995    debug_assert_eq!(xs_all.len(), b * cols);
11996    debug_assert_eq!(out.len(), b * rows);
11997    debug_assert_eq!(offsets.len(), rows + 1);
11998    let ng = cols / GROUP_SIZE;
11999    let bits = &bytes[..rows];
12000    let sc_off = rows;
12001    let gscale = |r: usize, g: usize| {
12002        let so = (r * ng + g) * 2;
12003        f16_to_f32(u16::from_le_bytes([
12004            bytes[sc_off + so],
12005            bytes[sc_off + so + 1],
12006        ]))
12007    };
12008
12009    // Decode row r's raw (u − L) values into `dst` (f32, unscaled).
12010    let decode_f32 = |r: usize, dst: &mut [f32]| {
12011        let bw = bits[r] as usize;
12012        let l = ((1i32 << (bw - 1)) - 1) as f32;
12013        let data = &bytes[offsets[r]..offsets[r + 1]];
12014        let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
12015        for d in dst.iter_mut() {
12016            while nbits < bw {
12017                acc = (acc << 8) | data[idx] as u64;
12018                idx += 1;
12019                nbits += 8;
12020            }
12021            let u = ((acc >> (nbits - bw)) & ((1u64 << bw) - 1)) as f32;
12022            nbits -= bw;
12023            *d = u - l;
12024        }
12025    };
12026
12027    if a8w8_enabled() {
12028        let acts: Vec<SplitAct> = (0..b)
12029            .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
12030            .collect();
12031        let acts = &acts;
12032        let out_addr = SendMut(out.as_mut_ptr());
12033        let run = move |start: usize, end: usize| {
12034            for r in start..end {
12035                let bw = bits[r] as usize;
12036                if bw == 8 {
12037                    // u−L reaches 128 → no i8 path; decode once, exact
12038                    // f32 dots for every position (same as vbitmatvec).
12039                    ROW_F32.with(|rb| {
12040                        let mut buf = rb.borrow_mut();
12041                        buf.resize(cols, 0.0);
12042                        decode_f32(r, &mut buf);
12043                        for bi in 0..b {
12044                            let x = &xs_all[bi * cols..(bi + 1) * cols];
12045                            let mut dot = 0f32;
12046                            for g in 0..ng {
12047                                let mut gd = 0f32;
12048                                for k in 0..GROUP_SIZE {
12049                                    gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
12050                                }
12051                                dot += gd * gscale(r, g);
12052                            }
12053                            // SAFETY: disjoint (bi, r) cells per worker range.
12054                            unsafe { *out_addr.at(bi * rows + r) = dot };
12055                        }
12056                    });
12057                    continue;
12058                }
12059                let l = (1i32 << (bw - 1)) - 1;
12060                let data = &bytes[offsets[r]..offsets[r + 1]];
12061                ROW_I8.with(|rb| {
12062                    let mut buf = rb.borrow_mut();
12063                    buf.resize(cols, 0);
12064                    #[inline(always)]
12065                    fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
12066                        for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
12067                            let u = unpack8::<B>(&data[blk * B..]);
12068                            for k in 0..8 {
12069                                chunk[k] = (u[k] - l) as i8 as u8;
12070                            }
12071                        }
12072                    }
12073                    match bw {
12074                        3 => fill::<3>(data, l, &mut buf),
12075                        4 => vbit_fill4(data, &mut buf),
12076                        5 => fill::<5>(data, l, &mut buf),
12077                        6 => fill::<6>(data, l, &mut buf),
12078                        _ => unreachable!("vbit bit-width {bw} (validated at load)"),
12079                    }
12080                    let mut bi = 0usize;
12081                    // The vbit scale table shares q4_block's layout
12082                    // (contiguous f16 per (row·ng + g)), so the same
12083                    // blocked 1×4 kernel serves the decoded row.
12084                    #[cfg(target_arch = "x86_64")]
12085                    if avx2_enabled() && blocked_enabled() {
12086                        while bi + 4 <= acts.len() {
12087                            let xs = [
12088                                acts[bi].xq.as_slice(),
12089                                acts[bi + 1].xq.as_slice(),
12090                                acts[bi + 2].xq.as_slice(),
12091                                acts[bi + 3].xq.as_slice(),
12092                            ];
12093                            let sxs = [
12094                                acts[bi].sx,
12095                                acts[bi + 1].sx,
12096                                acts[bi + 2].sx,
12097                                acts[bi + 3].sx,
12098                            ];
12099                            let d = unsafe {
12100                                if vnni_tiles_enabled() {
12101                                    dot_q4b_row_1x4_sx_vnni(
12102                                        &buf,
12103                                        &bytes[sc_off..],
12104                                        r * ng,
12105                                        ng,
12106                                        xs,
12107                                        sxs,
12108                                    )
12109                                } else {
12110                                    dot_q4b_row_1x4_sx_avx2(
12111                                        &buf,
12112                                        &bytes[sc_off..],
12113                                        r * ng,
12114                                        ng,
12115                                        xs,
12116                                        sxs,
12117                                    )
12118                                }
12119                            };
12120                            for k in 0..4 {
12121                                let act = &acts[bi + k];
12122                                let mut dot = d[k];
12123                                for &(j, xv) in &act.outliers {
12124                                    dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
12125                                }
12126                                // SAFETY: disjoint (bi, r) cells per worker.
12127                                unsafe { *out_addr.at((bi + k) * rows + r) = dot };
12128                            }
12129                            bi += 4;
12130                        }
12131                    }
12132                    while bi < acts.len() {
12133                        let act = &acts[bi];
12134                        let mut dot = 0f32;
12135                        for g in 0..ng {
12136                            let d = dot_i8_i8(
12137                                &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
12138                                &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
12139                            ) as f32
12140                                * act.sx;
12141                            dot += d * gscale(r, g);
12142                        }
12143                        for &(j, xv) in &act.outliers {
12144                            dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
12145                        }
12146                        // SAFETY: disjoint (bi, r) cells per worker range.
12147                        unsafe { *out_addr.at(bi * rows + r) = dot };
12148                        bi += 1;
12149                    }
12150                });
12151            }
12152        };
12153        dispatch_rows(pool, rows, &run);
12154        return;
12155    }
12156
12157    let out_addr = SendMut(out.as_mut_ptr());
12158    let run = move |start: usize, end: usize| {
12159        ROW_F32.with(|rb| {
12160            let mut buf = rb.borrow_mut();
12161            buf.resize(cols, 0.0);
12162            for r in start..end {
12163                decode_f32(r, &mut buf);
12164                for bi in 0..b {
12165                    let x = &xs_all[bi * cols..(bi + 1) * cols];
12166                    let mut dot = 0f32;
12167                    for g in 0..ng {
12168                        let mut gd = 0f32;
12169                        for k in 0..GROUP_SIZE {
12170                            gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
12171                        }
12172                        dot += gd * gscale(r, g);
12173                    }
12174                    // SAFETY: disjoint (bi, r) cells per worker range.
12175                    unsafe { *out_addr.at(bi * rows + r) = dot };
12176                }
12177            }
12178        })
12179    };
12180    dispatch_rows(pool, rows, &run);
12181}
12182
12183/// Build a GPU batch job for a q8-family mapped tensor (primary
12184/// shard): prescaled input + directory coordinates. None → not
12185/// GPU-eligible, caller stays on the CPU.
12186pub(crate) fn gpu_batch_job<'a>(
12187    t: &'a QTensor,
12188    x: &[f32],
12189) -> Option<(std::sync::Arc<CmfModel>, crate::gpu::BatchJob<'a>)> {
12190    match t {
12191        QTensor::Mapped {
12192            model,
12193            idx,
12194            dtype: dt @ (TensorDtype::Q8Row | TensorDtype::Q8_2f),
12195            rows,
12196            cols,
12197            row_scale,
12198            col_field,
12199            ..
12200        } => Some((
12201            model.clone(),
12202            crate::gpu::BatchJob {
12203                idx: *idx,
12204                rows: *rows,
12205                cols: *cols,
12206                row_scale,
12207                xs: prescale(x, col_field, *dt).into_owned(),
12208                layout: crate::gpu::BatchLayout::Q8,
12209            },
12210        )),
12211        // q1: raw f32 activations, tile-embedded scales.
12212        QTensor::Mapped {
12213            model,
12214            idx,
12215            dtype: TensorDtype::Q1,
12216            rows,
12217            cols,
12218            ..
12219        } => Some((
12220            model.clone(),
12221            crate::gpu::BatchJob {
12222                idx: *idx,
12223                rows: *rows,
12224                cols: *cols,
12225                row_scale: &[],
12226                xs: x.to_vec(),
12227                layout: crate::gpu::BatchLayout::Q1,
12228            },
12229        )),
12230        // q1g: same raw-f32 contract, sign + scale planes.
12231        QTensor::Mapped {
12232            model,
12233            idx,
12234            dtype: TensorDtype::Q1G,
12235            rows,
12236            cols,
12237            ..
12238        } => Some((
12239            model.clone(),
12240            crate::gpu::BatchJob {
12241                idx: *idx,
12242                rows: *rows,
12243                cols: *cols,
12244                row_scale: &[],
12245                xs: x.to_vec(),
12246                layout: crate::gpu::BatchLayout::Q1g,
12247            },
12248        )),
12249        // q4_tiled / q4tp: raw f32 activations; the scales live in the
12250        // payload (inline tiles / row ladder), so row_scale stays empty.
12251        // The GDN projection batch already runs these layouts on Metal —
12252        // this arm lets the attention QKV batch reach the same kernels.
12253        QTensor::Mapped {
12254            model,
12255            idx,
12256            dtype: dt @ (TensorDtype::Q4Tiled | TensorDtype::Q4TiledP),
12257            rows,
12258            cols,
12259            ..
12260        } => Some((
12261            model.clone(),
12262            crate::gpu::BatchJob {
12263                idx: *idx,
12264                rows: *rows,
12265                cols: *cols,
12266                row_scale: &[],
12267                xs: x.to_vec(),
12268                layout: if *dt == TensorDtype::Q4Tiled {
12269                    crate::gpu::BatchLayout::Q4t
12270                } else {
12271                    crate::gpu::BatchLayout::Q4tp
12272                },
12273            },
12274        )),
12275        _ => None,
12276    }
12277}
12278
12279thread_local! {
12280    static PRESCALE_BUF1: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
12281    static PRESCALE_BUF2: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
12282}
12283
12284pub(crate) fn prescale<'a>(
12285    x: &'a [f32],
12286    col_field: &[f32],
12287    dtype: TensorDtype,
12288) -> std::borrow::Cow<'a, [f32]> {
12289    if dtype == TensorDtype::Q8_2f {
12290        x.iter().zip(col_field).map(|(a, c)| a * c).collect()
12291    } else {
12292        std::borrow::Cow::Borrowed(x)
12293    }
12294}
12295
12296/// θ col-field fold for q8_2f activations. Borrowed pass-through for
12297/// every other dtype, using thread-local buffers to eliminate per-matvec allocations.
12298pub(crate) fn prescale_with<R, F: FnOnce(&[f32]) -> R>(
12299    x: &[f32],
12300    col_field: &[f32],
12301    dtype: TensorDtype,
12302    buf_id: u8,
12303    f: F,
12304) -> R {
12305    if dtype == TensorDtype::Q8_2f {
12306        if buf_id == 1 {
12307            PRESCALE_BUF1.with(|b| {
12308                let mut buf = b.borrow_mut();
12309                buf.clear();
12310                buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
12311                f(&buf)
12312            })
12313        } else {
12314            PRESCALE_BUF2.with(|b| {
12315                let mut buf = b.borrow_mut();
12316                buf.clear();
12317                buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
12318                f(&buf)
12319            })
12320        }
12321    } else {
12322        f(x)
12323    }
12324}
12325
12326// ───────────────────── x86-64 AVX2 kernels (roadmap этап 2) ─────────────────────
12327
12328/// AVX2+FMA available? Default ON when the CPU supports both;
12329/// `CMF_AVX2=0` disables (falls back to the autovectorized loops).
12330#[cfg(target_arch = "x86_64")]
12331pub(crate) fn avx2_enabled() -> bool {
12332    use std::sync::OnceLock;
12333    static ON: OnceLock<bool> = OnceLock::new();
12334    *ON.get_or_init(|| {
12335        std::env::var("CMF_AVX2").map(|v| v != "0").unwrap_or(true)
12336            && std::arch::is_x86_feature_detected!("avx2")
12337            && std::arch::is_x86_feature_detected!("fma")
12338    })
12339}
12340
12341/// AVX2 A8W8 allowed? The quantized-activation contract is switched by
12342/// the SAME env as the ARM SDOT path: `CMF_SDOT=0` keeps exact kernels
12343/// (the golden-parity exact gate relies on it) — AVX2 f32 kernels stay
12344/// active either way, they are exact (regrouped sums only).
12345#[cfg(target_arch = "x86_64")]
12346fn avx2_a8w8_enabled() -> bool {
12347    if FLOAT_ACTIVATIONS.get() {
12348        return false;
12349    }
12350    use std::sync::OnceLock;
12351    static ON: OnceLock<bool> = OnceLock::new();
12352    *ON.get_or_init(|| {
12353        avx2_enabled() && std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true)
12354    })
12355}
12356
12357thread_local! {
12358    static FULL_GPU_Q8: std::cell::Cell<bool> = const { std::cell::Cell::new(false) };
12359}
12360
12361/// Match the graph's full-device q8 projection precision on MiMo's host
12362/// tail. Does not enable the GPU or bypass a CPU-only/device-refusal gate.
12363pub(crate) fn enter_full_gpu_q8_scope() -> impl Drop {
12364    struct Restore(bool, std::marker::PhantomData<std::rc::Rc<()>>);
12365    impl Drop for Restore {
12366        fn drop(&mut self) {
12367            FULL_GPU_Q8.set(self.0);
12368        }
12369    }
12370    Restore(FULL_GPU_Q8.replace(true), std::marker::PhantomData)
12371}
12372
12373// Dynamic MiMo experts must not change activation precision when a cache
12374// fill moves them from CPU to GPU. Thread-local: only the cold-expert
12375// dispatch selects float kernels; concurrent pipelines keep their policy.
12376thread_local! {
12377    static FLOAT_ACTIVATIONS: std::cell::Cell<bool> = const { std::cell::Cell::new(false) };
12378}
12379
12380/// Flint's released Q2/Q4 profiles are unusually sensitive to activation
12381/// quantization in the host-side attention path. With the normal A8W8
12382/// default, a short greedy decode can diverge after the first token even
12383/// though the selected-expert Metal path itself is unchanged. An Embryo file
12384/// carries a router fitted on the exact φ, which A8W8 moves by 4e-3, and its
12385/// CPU decode costs the same without it. Keep the reference float-activation
12386/// contract for both unless the operator has explicitly selected an
12387/// `CMF_SDOT` policy.
12388fn exact_cpu_default_enabled() -> bool {
12389    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12390    *ON.get_or_init(|| std::env::var_os("CMF_SDOT").is_none())
12391}
12392
12393impl QTensor {
12394    #[inline]
12395    fn exact_cpu_default(&self) -> bool {
12396        !FLOAT_ACTIVATIONS.get()
12397            && exact_cpu_default_enabled()
12398            && matches!(
12399                self,
12400                Self::Mapped { model, dtype, .. } if matches!(model.header.arch.arch_name.as_str(), "flint" | "cortiq_embryo")
12401                    // Dory on q8 runs as fast on float activations and exact; on
12402                    // q4tp int8 is 4.7× on an M4 for 0.1 % of perplexity.
12403                    || (model.header.arch.arch_name == "dory" && *dtype == TensorDtype::Q8_2f)
12404            )
12405    }
12406
12407    #[inline]
12408    fn any_exact_cpu_default(ts: &[&QTensor]) -> bool {
12409        ts.iter().any(|t| t.exact_cpu_default())
12410    }
12411
12412    #[inline]
12413    fn any_exact_cpu_default_pairs(pairs: &[(&QTensor, &QTensor)]) -> bool {
12414        pairs
12415            .iter()
12416            .any(|(gate, up)| gate.exact_cpu_default() || up.exact_cpu_default())
12417    }
12418}
12419
12420/// Run `f` with the exact float-activation contract on this thread: the
12421/// 4-bit and 8-bit kernels skip the int8 activation split (`CMF_SDOT=0`
12422/// semantics, scoped). For parity tests and diagnostics — e.g. an `mxfp4`
12423/// model then computes bit for bit what its dequantized f32 twin computes.
12424pub fn with_float_activations<R>(f: impl FnOnce() -> R) -> R {
12425    float_activations_scope(f)
12426}
12427
12428/// Active on the current thread. `Pool` snapshots this with its job
12429/// descriptor and installs it on persistent workers, so a scoped quantized
12430/// matvec cannot accidentally choose a different kernel on a worker.
12431#[inline]
12432pub(crate) fn float_activations_active() -> bool {
12433    FLOAT_ACTIVATIONS.get()
12434}
12435
12436/// Restores the current thread's float-activation state on drop. The `Rc`
12437/// marker deliberately keeps the guard on the thread that created it: this
12438/// policy is thread-local, not a portable capability.
12439pub(crate) struct FloatActivationsScope {
12440    previous: bool,
12441    _not_send: std::marker::PhantomData<std::rc::Rc<()>>,
12442}
12443
12444impl Drop for FloatActivationsScope {
12445    fn drop(&mut self) {
12446        FLOAT_ACTIVATIONS.set(self.previous);
12447    }
12448}
12449
12450#[inline]
12451pub(crate) fn enter_float_activations_scope() -> FloatActivationsScope {
12452    FloatActivationsScope {
12453        previous: FLOAT_ACTIVATIONS.replace(true),
12454        _not_send: std::marker::PhantomData,
12455    }
12456}
12457
12458/// Run `f` with exact float activations (no A8W8 int8 activation
12459/// quantization in the quantized matvecs). The persistent row pool carries
12460/// this scope into a dispatched job and restores each worker afterwards.
12461/// Public for parity tests that compare the default kernels with the strict
12462/// ones in one process; `CMF_SDOT=0` is the process-wide equivalent.
12463#[doc(hidden)]
12464pub fn float_activations_scope<R>(f: impl FnOnce() -> R) -> R {
12465    let _scope = enter_float_activations_scope();
12466    f()
12467}
12468
12469/// Row-exact batching: while set, the x86 batched kernels (`qmatmat`,
12470/// `q4tp_matmat`) compute every (weight row, token) cell with the
12471/// single-token kernel instead of the blocked 2×4 / 1×4 / 1×8 tiles, so a
12472/// token's result does not depend on the batch it rides in and equals its
12473/// matvec. On ARM `q4tp_matmat` keeps its 1×4 tile but in the matvec's
12474/// reduction order, and no q4tp batch takes the AMX or device GEMM. The
12475/// MiMo speculative verify holds it (`row_exact_scope`) — its accepted
12476/// rows must be the rows plain decode would have produced.
12477// Shared with pool workers, so overlapping requests must keep the mode
12478// enabled until the LAST scope leaves. Saving/restoring a global bool is
12479// incorrect when two threads enter and leave in a non-LIFO order.
12480static ROW_EXACT: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
12481
12482pub(crate) fn row_exact() -> bool {
12483    ROW_EXACT.load(std::sync::atomic::Ordering::Acquire) != 0
12484}
12485
12486fn counted_row_exact_scope<R>(active: &std::sync::atomic::AtomicUsize, f: impl FnOnce() -> R) -> R {
12487    struct Restore<'a>(&'a std::sync::atomic::AtomicUsize);
12488    impl Drop for Restore<'_> {
12489        fn drop(&mut self) {
12490            self.0.fetch_sub(1, std::sync::atomic::Ordering::AcqRel);
12491        }
12492    }
12493    active.fetch_add(1, std::sync::atomic::Ordering::AcqRel);
12494    let _restore = Restore(active);
12495    f()
12496}
12497
12498/// Run `f` with row-exact batching on (also released on unwind).
12499pub(crate) fn row_exact_scope<R>(f: impl FnOnce() -> R) -> R {
12500    counted_row_exact_scope(&ROW_EXACT, f)
12501}
12502
12503/// A8W8 quantized-activation path available on THIS machine? One
12504/// switch across architectures: ARM dotprod (CMF_SDOT) or x86 AVX2
12505/// (CMF_AVX2 + the same CMF_SDOT exact-contract override).
12506#[inline]
12507pub(crate) fn a8w8_enabled() -> bool {
12508    #[cfg(target_arch = "aarch64")]
12509    {
12510        sdot_enabled()
12511    }
12512    #[cfg(target_arch = "x86_64")]
12513    {
12514        avx2_a8w8_enabled()
12515    }
12516    #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
12517    {
12518        false
12519    }
12520}
12521
12522/// int8·int8 dot dispatch: SDOT on ARM; AVX-512 VNNI (vpdpbusd) or AVX2
12523/// maddubs on x86. Callers are gated by `a8w8_enabled()`.
12524#[inline]
12525#[allow(unreachable_code)]
12526fn dot_i8_i8(w: &[u8], xq: &[i8]) -> i32 {
12527    #[cfg(target_arch = "aarch64")]
12528    unsafe {
12529        return dot_i8_sdot(w, xq);
12530    }
12531    #[cfg(target_arch = "x86_64")]
12532    unsafe {
12533        if avx512vnni_enabled() {
12534            return dot_i8_i8_vnni(w, xq);
12535        }
12536        return dot_i8_i8_avx2(w, xq);
12537    }
12538    w.iter()
12539        .zip(xq)
12540        .map(|(&a, &b)| (a as i8) as i32 * b as i32)
12541        .sum()
12542}
12543
12544/// AVX-512 VNNI available? (F+BW+VL+VNNI; `CMF_AVX512=0` falls back to
12545/// AVX2.) VL matters: short 32-byte groups (q4/vbit) ride the 256-bit
12546/// `vpdpbusd` encoding.
12547#[cfg(target_arch = "x86_64")]
12548fn avx512vnni_enabled() -> bool {
12549    use std::sync::OnceLock;
12550    static ON: OnceLock<bool> = OnceLock::new();
12551    *ON.get_or_init(|| {
12552        std::env::var("CMF_AVX512")
12553            .map(|v| v != "0")
12554            .unwrap_or(true)
12555            && std::arch::is_x86_feature_detected!("avx512f")
12556            && std::arch::is_x86_feature_detected!("avx512bw")
12557            && std::arch::is_x86_feature_detected!("avx512vl")
12558            && std::arch::is_x86_feature_detected!("avx512vnni")
12559    })
12560}
12561
12562/// Grouped-codec VNNI arms (the q4t/q4b/q1/q1t tile kernels): default
12563/// ON where AVX-512 VNNI exists (`CMF_VNNI_TILES=0` opt-out). Measured
12564/// on Ryzen 7950X (Zen4, 3 alternating process pairs, blocked GEMM
12565/// 4864×896 b=256): q4t 63→68 GF/s (+8%), q1 53→56 (+6%), q4b 72→75
12566/// (+4%) — consistent, no leg regressed. The tile kernels keep a
12567/// horizontal reduce per 32-weight group, so the `vpdpbusd` saving is
12568/// smaller than the long-dot q8 win (+13%), but it is real and free.
12569#[cfg(target_arch = "x86_64")]
12570fn vnni_tiles_enabled() -> bool {
12571    use std::sync::OnceLock;
12572    static ON: OnceLock<bool> = OnceLock::new();
12573    *ON.get_or_init(|| {
12574        std::env::var("CMF_VNNI_TILES")
12575            .map(|v| v != "0")
12576            .unwrap_or(true)
12577            && avx512vnni_enabled()
12578    })
12579}
12580
12581/// One 256-bit u8×i8 dot → i32 via `vpdpbusd` into a fresh accumulator
12582/// plus the same horizontal reduce the AVX2 kernels use. Products are
12583/// bounded (|w| ≤ 8 or ≤ 1), so maddubs never saturated — the i32 sum
12584/// is bit-identical to the maddubs+madd pair it replaces.
12585#[cfg(target_arch = "x86_64")]
12586#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
12587#[inline]
12588unsafe fn dpbusd_hsum(aw: core::arch::x86_64::__m256i, xs: core::arch::x86_64::__m256i) -> i32 {
12589    // SAFETY: pure register math.
12590    unsafe {
12591        use core::arch::x86_64::*;
12592        let d = _mm256_dpbusd_epi32(_mm256_setzero_si256(), aw, xs);
12593        let hi128 = _mm256_extracti128_si256::<1>(d);
12594        let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
12595        let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
12596        let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
12597        _mm_cvtsi128_si32(s32)
12598    }
12599}
12600
12601/// int8·int8 via AVX-512 VNNI: `vpdpbusd` fuses the maddubs+madd+add
12602/// triple into one u8×i8 dot-accumulate. AVX-512 has no vpsignb, so the
12603/// |w|·sign(x,w) trick becomes |w| × (x negated where w<0) via a mask
12604/// subtract — w==0 lanes contribute 0 through |w|=0 either way.
12605#[cfg(target_arch = "x86_64")]
12606#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
12607unsafe fn dot_i8_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
12608    // SAFETY: callers uphold slice-length contracts (see call sites).
12609    unsafe {
12610        use core::arch::x86_64::*;
12611        let n = w.len();
12612        let mut j = 0usize;
12613        let mut total: i32;
12614        // 4 independent accumulators: vpdpbusd is its own loop-carried
12615        // dependency (~5-cycle latency) — a single-acc loop runs
12616        // latency-bound and LOSES to the AVX2 maddubs kernel, measured
12617        // on Granite Rapids.
12618        {
12619            #[inline(always)]
12620            unsafe fn step(
12621                w: *const u8,
12622                x: *const i8,
12623                acc: core::arch::x86_64::__m512i,
12624            ) -> core::arch::x86_64::__m512i {
12625                unsafe {
12626                    use core::arch::x86_64::*;
12627                    let wv = _mm512_loadu_si512(w as *const _);
12628                    let xv = _mm512_loadu_si512(x as *const _);
12629                    let aw = _mm512_abs_epi8(wv);
12630                    let neg = _mm512_movepi8_mask(wv);
12631                    let sx = _mm512_mask_sub_epi8(xv, neg, _mm512_setzero_si512(), xv);
12632                    _mm512_dpbusd_epi32(acc, aw, sx)
12633                }
12634            }
12635            let (mut a0, mut a1, mut a2, mut a3) = (
12636                _mm512_setzero_si512(),
12637                _mm512_setzero_si512(),
12638                _mm512_setzero_si512(),
12639                _mm512_setzero_si512(),
12640            );
12641            while j + 256 <= n {
12642                a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
12643                a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), a1);
12644                a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), a2);
12645                a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), a3);
12646                j += 256;
12647            }
12648            while j + 64 <= n {
12649                a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
12650                j += 64;
12651            }
12652            let s01 = _mm512_add_epi32(a0, a1);
12653            let s23 = _mm512_add_epi32(a2, a3);
12654            total = _mm512_reduce_add_epi32(_mm512_add_epi32(s01, s23));
12655        }
12656        // 32-wide (q4/vbit groups are exactly 32 bytes).
12657        if j + 32 <= n {
12658            let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
12659            let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
12660            let d = _mm256_dpbusd_epi32(
12661                _mm256_setzero_si256(),
12662                _mm256_abs_epi8(wv),
12663                _mm256_sign_epi8(xv, wv),
12664            );
12665            let hi128 = _mm256_extracti128_si256::<1>(d);
12666            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
12667            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
12668            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
12669            total += _mm_cvtsi128_si32(s32);
12670            j += 32;
12671        }
12672        while j < n {
12673            total += (w[j] as i8) as i32 * xq[j] as i32;
12674            j += 1;
12675        }
12676        total
12677    }
12678}
12679
12680/// i8 row · f32 x via AVX2/FMA (x86 mirror of `dot_i8_f32_neon`).
12681#[cfg(target_arch = "x86_64")]
12682#[target_feature(enable = "avx2,fma")]
12683unsafe fn dot_i8_f32_avx2(w: &[u8], x: &[f32]) -> f32 {
12684    // SAFETY: callers uphold slice-length contracts (see call sites).
12685    unsafe {
12686        use core::arch::x86_64::*;
12687        let n = x.len();
12688        let wp = w.as_ptr();
12689        let xp = x.as_ptr();
12690        let (mut a0, mut a1) = (_mm256_setzero_ps(), _mm256_setzero_ps());
12691        let mut j = 0usize;
12692        while j + 16 <= n {
12693            let wb = _mm_loadu_si128(wp.add(j) as *const __m128i);
12694            let lo = _mm256_cvtepi8_epi32(wb);
12695            let hi = _mm256_cvtepi8_epi32(_mm_srli_si128::<8>(wb));
12696            a0 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(lo), _mm256_loadu_ps(xp.add(j)), a0);
12697            a1 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(hi), _mm256_loadu_ps(xp.add(j + 8)), a1);
12698            j += 16;
12699        }
12700        let acc = _mm256_add_ps(a0, a1);
12701        let hi128 = _mm256_extractf128_ps::<1>(acc);
12702        let s128 = _mm_add_ps(_mm256_castps256_ps128(acc), hi128);
12703        let s64 = _mm_add_ps(s128, _mm_movehl_ps(s128, s128));
12704        let s32 = _mm_add_ss(s64, _mm_shuffle_ps::<1>(s64, s64));
12705        let mut sum = _mm_cvtss_f32(s32);
12706        while j < n {
12707            sum += (*wp.add(j) as i8) as f32 * *xp.add(j);
12708            j += 1;
12709        }
12710        sum
12711    }
12712}
12713
12714/// int8(weight)·int8(activation) → i32 via AVX2 maddubs — the x86
12715/// analogue of the SDOT A8W8 path. `maddubs` takes u8×i8, so the
12716/// standard sign trick applies: |w| × sign(x, w) ≡ w × x per lane.
12717/// Pair saturation is safe: |w|≤128, |x|≤127 → 2·128·127 < 32767.
12718#[cfg(target_arch = "x86_64")]
12719#[target_feature(enable = "avx2")]
12720unsafe fn dot_i8_i8_avx2(w: &[u8], xq: &[i8]) -> i32 {
12721    // SAFETY: callers uphold slice-length contracts (see call sites).
12722    unsafe {
12723        use core::arch::x86_64::*;
12724        let n = w.len();
12725        let ones = _mm256_set1_epi16(1);
12726        let mut acc = _mm256_setzero_si256();
12727        let mut j = 0usize;
12728        while j + 32 <= n {
12729            let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
12730            let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
12731            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
12732            acc = _mm256_add_epi32(acc, _mm256_madd_epi16(p16, ones));
12733            j += 32;
12734        }
12735        let hi128 = _mm256_extracti128_si256::<1>(acc);
12736        let s128 = _mm_add_epi32(_mm256_castsi256_si128(acc), hi128);
12737        let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
12738        let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
12739        let mut s = _mm_cvtsi128_si32(s32);
12740        while j < n {
12741            s += (w[j] as i8) as i32 * xq[j] as i32;
12742            j += 1;
12743        }
12744        s
12745    }
12746}
12747
12748/// smmla 2×4: one instruction covers a 2-row × 2-activation × 8-deep
12749/// tile (32 MACs vs sdot's 16) — the weight pair loads once per 8-k
12750/// slice as a combined 2×8 register and meets two activation pairs.
12751#[cfg(target_arch = "aarch64")]
12752#[target_feature(enable = "neon,i8mm")]
12753unsafe fn dot_i8_smmla_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
12754    // SAFETY: callers uphold slice-length contracts.
12755    unsafe {
12756        use core::arch::aarch64::*;
12757        use core::arch::asm;
12758        let n = w0.len();
12759        let w0p = w0.as_ptr() as *const i8;
12760        let w1p = w1.as_ptr() as *const i8;
12761        // acc01 holds [c(r0,x0) c(r0,x1) c(r1,x0) c(r1,x1)]; acc23 the
12762        // same for x2/x3.
12763        let mut acc01 = vdupq_n_s32(0);
12764        let mut acc23 = vdupq_n_s32(0);
12765        let mut i = 0usize;
12766        while i + 8 <= n {
12767            let wa = vcombine_s8(vld1_s8(w0p.add(i)), vld1_s8(w1p.add(i)));
12768            let xb01 = vcombine_s8(
12769                vld1_s8(xs[0].as_ptr().add(i)),
12770                vld1_s8(xs[1].as_ptr().add(i)),
12771            );
12772            let xb23 = vcombine_s8(
12773                vld1_s8(xs[2].as_ptr().add(i)),
12774                vld1_s8(xs[3].as_ptr().add(i)),
12775            );
12776            asm!(
12777                "smmla {a01:v}.4s, {w:v}.16b, {x01:v}.16b",
12778                "smmla {a23:v}.4s, {w:v}.16b, {x23:v}.16b",
12779                a01 = inout(vreg) acc01, a23 = inout(vreg) acc23,
12780                w = in(vreg) wa, x01 = in(vreg) xb01, x23 = in(vreg) xb23,
12781                options(pure, nomem, nostack),
12782            );
12783            i += 8;
12784        }
12785        let mut out = [[0i32; 4]; 2];
12786        let a01: [i32; 4] = core::mem::transmute(acc01);
12787        let a23: [i32; 4] = core::mem::transmute(acc23);
12788        out[0][0] = a01[0];
12789        out[0][1] = a01[1];
12790        out[1][0] = a01[2];
12791        out[1][1] = a01[3];
12792        out[0][2] = a23[0];
12793        out[0][3] = a23[1];
12794        out[1][2] = a23[2];
12795        out[1][3] = a23[3];
12796        if i < n {
12797            for (k, x) in xs.iter().enumerate() {
12798                for j in i..n {
12799                    out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
12800                    out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
12801                }
12802            }
12803        }
12804        out
12805    }
12806}
12807
12808/// ARM twin of the x86 blocked prefill GEMM: two weight rows stay in
12809/// registers across four activation streams, eight sdot accumulators.
12810/// (The per-row form re-read each W row once per activation.)
12811#[cfg(target_arch = "aarch64")]
12812#[target_feature(enable = "neon,dotprod")]
12813unsafe fn dot_i8_sdot_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
12814    // SAFETY: callers uphold slice-length contracts.
12815    unsafe {
12816        use core::arch::aarch64::*;
12817        use core::arch::asm;
12818        let n = w0.len();
12819        let w0p = w0.as_ptr() as *const i8;
12820        let w1p = w1.as_ptr() as *const i8;
12821        let mut acc = [[vdupq_n_s32(0); 4]; 2];
12822        let mut i = 0usize;
12823        while i + 16 <= n {
12824            let wv0 = vld1q_s8(w0p.add(i));
12825            let wv1 = vld1q_s8(w1p.add(i));
12826            for (k, x) in xs.iter().enumerate() {
12827                let xv = vld1q_s8(x.as_ptr().add(i));
12828                let (mut a0, mut a1) = (acc[0][k], acc[1][k]);
12829                asm!(
12830                    "sdot {a0:v}.4s, {w0:v}.16b, {x:v}.16b",
12831                    "sdot {a1:v}.4s, {w1:v}.16b, {x:v}.16b",
12832                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
12833                    w0 = in(vreg) wv0, w1 = in(vreg) wv1, x = in(vreg) xv,
12834                    options(pure, nomem, nostack),
12835                );
12836                acc[0][k] = a0;
12837                acc[1][k] = a1;
12838            }
12839            i += 16;
12840        }
12841        let mut out = [[0i32; 4]; 2];
12842        for r in 0..2 {
12843            for k in 0..4 {
12844                out[r][k] = vaddvq_s32(acc[r][k]);
12845            }
12846        }
12847        if i < n {
12848            for (k, x) in xs.iter().enumerate() {
12849                for j in i..n {
12850                    out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
12851                    out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
12852                }
12853            }
12854        }
12855        out
12856    }
12857}
12858
12859/// Blocked 2 weight rows × 4 activations for the prefill GEMM
12860/// (roadmap P0: packed panels + multi-row accumulators). The two rows'
12861/// abs() live in registers across all four activation streams; the
12862/// sign-fixup is recomputed per pair (the price of the maddubs trick).
12863/// Returns raw i8·i8 dots; the caller applies scales and outliers.
12864#[cfg(target_arch = "x86_64")]
12865#[target_feature(enable = "avx2")]
12866unsafe fn dot_i8_i8_avx2_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
12867    // SAFETY: callers uphold slice-length contracts.
12868    unsafe {
12869        use core::arch::x86_64::*;
12870        let n = w0.len();
12871        let ones = _mm256_set1_epi16(1);
12872        let mut acc = [[_mm256_setzero_si256(); 4]; 2];
12873        let mut j = 0usize;
12874        while j + 32 <= n {
12875            let wv0 = _mm256_loadu_si256(w0.as_ptr().add(j) as *const __m256i);
12876            let wv1 = _mm256_loadu_si256(w1.as_ptr().add(j) as *const __m256i);
12877            let aw0 = _mm256_abs_epi8(wv0);
12878            let aw1 = _mm256_abs_epi8(wv1);
12879            for (k, x) in xs.iter().enumerate() {
12880                let xv = _mm256_loadu_si256(x.as_ptr().add(j) as *const __m256i);
12881                let p0 = _mm256_maddubs_epi16(aw0, _mm256_sign_epi8(xv, wv0));
12882                acc[0][k] = _mm256_add_epi32(acc[0][k], _mm256_madd_epi16(p0, ones));
12883                let p1 = _mm256_maddubs_epi16(aw1, _mm256_sign_epi8(xv, wv1));
12884                acc[1][k] = _mm256_add_epi32(acc[1][k], _mm256_madd_epi16(p1, ones));
12885            }
12886            j += 32;
12887        }
12888        let mut out = [[0i32; 4]; 2];
12889        for r in 0..2 {
12890            for k in 0..4 {
12891                let a = acc[r][k];
12892                let hi128 = _mm256_extracti128_si256::<1>(a);
12893                let s128 = _mm_add_epi32(_mm256_castsi256_si128(a), hi128);
12894                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
12895                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
12896                out[r][k] = _mm_cvtsi128_si32(s32);
12897            }
12898        }
12899        if j < n {
12900            for (k, x) in xs.iter().enumerate() {
12901                for i in j..n {
12902                    out[0][k] += (w0[i] as i8) as i32 * x[i] as i32;
12903                    out[1][k] += (w1[i] as i8) as i32 * x[i] as i32;
12904                }
12905            }
12906        }
12907        out
12908    }
12909}
12910
12911/// AVX2/VNNI q8 row dot with exact outlier correction (x86 mirror of
12912/// `row_dot_sdot` — same A8W8 contract). With AVX-512 VNNI the row goes
12913/// through the bias trick: Σ(w+128)·x via pure `vpdpbusd` (no per-lane
12914/// sign fixups), corrected by −128·Σx with Σx precomputed per split.
12915#[cfg(target_arch = "x86_64")]
12916#[inline]
12917fn row_dot_avx2(row: &[u8], act: &SplitAct) -> f32 {
12918    let dot = if avx512vnni_enabled() && row.len() >= 64 {
12919        (unsafe { dot_u8p128_i8_vnni(row, &act.xq) }) - 128 * act.xsum
12920    } else {
12921        unsafe { dot_i8_i8_avx2(row, &act.xq) }
12922    };
12923    let mut acc = dot as f32 * act.sx;
12924    for &(j, xv) in &act.outliers {
12925        acc += (row[j] as i8) as f32 * xv;
12926    }
12927    acc
12928}
12929
12930/// Σ (w[i]+128)·x[i] via pure `vpdpbusd` — the caller subtracts
12931/// 128·Σx. Four independent accumulators (dpbusd is ~5-cycle latency;
12932/// a single-acc loop runs latency-bound, measured on Granite Rapids).
12933#[cfg(target_arch = "x86_64")]
12934#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
12935unsafe fn dot_u8p128_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
12936    // SAFETY: callers uphold slice-length contracts (see call sites).
12937    unsafe {
12938        use core::arch::x86_64::*;
12939        let n = w.len();
12940        let flip = _mm512_set1_epi8(-128); // XOR 0x80: i8 w → u8 (w+128)
12941        #[inline(always)]
12942        unsafe fn step(
12943            w: *const u8,
12944            x: *const i8,
12945            flip: core::arch::x86_64::__m512i,
12946            acc: core::arch::x86_64::__m512i,
12947        ) -> core::arch::x86_64::__m512i {
12948            unsafe {
12949                use core::arch::x86_64::*;
12950                let wv = _mm512_xor_si512(_mm512_loadu_si512(w as *const _), flip);
12951                _mm512_dpbusd_epi32(acc, wv, _mm512_loadu_si512(x as *const _))
12952            }
12953        }
12954        let (mut a0, mut a1, mut a2, mut a3) = (
12955            _mm512_setzero_si512(),
12956            _mm512_setzero_si512(),
12957            _mm512_setzero_si512(),
12958            _mm512_setzero_si512(),
12959        );
12960        let mut j = 0usize;
12961        while j + 256 <= n {
12962            a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
12963            a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), flip, a1);
12964            a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), flip, a2);
12965            a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), flip, a3);
12966            j += 256;
12967        }
12968        while j + 64 <= n {
12969            a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
12970            j += 64;
12971        }
12972        let mut total = _mm512_reduce_add_epi32(_mm512_add_epi32(
12973            _mm512_add_epi32(a0, a1),
12974            _mm512_add_epi32(a2, a3),
12975        ));
12976        // Scalar tail: (w as i8) + 128 ≡ (w as u8) ^ 0x80.
12977        while j < n {
12978            total += ((w[j] ^ 0x80) as i32) * xq[j] as i32;
12979            j += 1;
12980        }
12981        total
12982    }
12983}
12984
12985/// One q4 row via AVX2: nibbles → centered i8 (unpacklo/hi restores the
12986/// writer's flat order, same as the NEON vzip pair), maddubs against
12987/// the pre-quantized activation group, × the group's f16 scale. Pair
12988/// saturation safe: |w|≤8, |x|≤127 → 2·8·127 ≪ 32767. Mirror of
12989/// `dot_q4_row_sdot`.
12990#[cfg(target_arch = "x86_64")]
12991#[target_feature(enable = "avx2")]
12992unsafe fn dot_q4_row_avx2(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
12993    // SAFETY: callers uphold slice-length contracts (16 packed bytes and
12994    // 2 scale bytes per group; xq.len() == gpr·GROUP_SIZE).
12995    unsafe {
12996        use core::arch::x86_64::*;
12997        let lomask = _mm_set1_epi8(0x0F);
12998        let eight = _mm256_set1_epi8(8);
12999        let ones = _mm256_set1_epi16(1);
13000        let mut acc = 0f32;
13001        for gi in 0..gpr {
13002            let g = g0 + gi;
13003            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
13004            let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
13005            let lo = _mm_and_si128(b, lomask);
13006            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
13007            let w = _mm256_sub_epi8(
13008                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
13009                eight,
13010            );
13011            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
13012            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
13013            let d = _mm256_madd_epi16(p16, ones);
13014            let hi128 = _mm256_extracti128_si256::<1>(d);
13015            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
13016            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
13017            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
13018            acc += _mm_cvtsi128_si32(s32) as f32 * s;
13019        }
13020        acc
13021    }
13022}
13023
13024/// Two-activation q4 row via AVX2: nibbles unpacked ONCE per group,
13025/// both activations dotted against the same centered i8 register.
13026#[cfg(target_arch = "x86_64")]
13027#[target_feature(enable = "avx2")]
13028unsafe fn dot_q4_row_avx2_2(
13029    packed: &[u8],
13030    scales: &[u8],
13031    g0: usize,
13032    gpr: usize,
13033    xq1: &[i8],
13034    xq2: &[i8],
13035) -> (f32, f32) {
13036    // SAFETY: callers uphold slice-length contracts (see dot_q4_row_avx2).
13037    unsafe {
13038        use core::arch::x86_64::*;
13039        let lomask = _mm_set1_epi8(0x0F);
13040        let eight = _mm256_set1_epi8(8);
13041        let ones = _mm256_set1_epi16(1);
13042        let (mut acc1, mut acc2) = (0f32, 0f32);
13043        #[inline(always)]
13044        unsafe fn hsum(d: core::arch::x86_64::__m256i) -> i32 {
13045            unsafe {
13046                use core::arch::x86_64::*;
13047                let hi128 = _mm256_extracti128_si256::<1>(d);
13048                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
13049                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
13050                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
13051                _mm_cvtsi128_si32(s32)
13052            }
13053        }
13054        for gi in 0..gpr {
13055            let g = g0 + gi;
13056            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
13057            let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
13058            let lo = _mm_and_si128(b, lomask);
13059            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
13060            let w = _mm256_sub_epi8(
13061                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
13062                eight,
13063            );
13064            let aw = _mm256_abs_epi8(w);
13065            let x1 = _mm256_loadu_si256(xq1.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
13066            let x2 = _mm256_loadu_si256(xq2.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
13067            let d1 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x1, w)), ones);
13068            let d2 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x2, w)), ones);
13069            acc1 += hsum(d1) as f32 * s;
13070            acc2 += hsum(d2) as f32 * s;
13071        }
13072        (acc1, acc2)
13073    }
13074}
13075
13076/// One q8 row range via AVX2 (x86 mirror of `q8_range_sdot`).
13077#[cfg(target_arch = "x86_64")]
13078fn q8_range_avx2(
13079    q: &[u8],
13080    row_scale: &[f32],
13081    act: &SplitAct,
13082    cols: usize,
13083    out_addr: SendMut,
13084    start: usize,
13085    end: usize,
13086) {
13087    for o in start..end {
13088        let v = row_dot_avx2(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
13089        // SAFETY: disjoint row ranges per worker.
13090        unsafe { *out_addr.at(o) = v };
13091    }
13092}
13093
13094/// Two-input q8 row range via AVX2 (x86 mirror of `q8_range2_sdot`).
13095#[cfg(target_arch = "x86_64")]
13096#[allow(clippy::too_many_arguments)]
13097fn q8_range2_avx2(
13098    q: &[u8],
13099    row_scale: &[f32],
13100    a1: &SplitAct,
13101    a2: &SplitAct,
13102    cols: usize,
13103    p1: SendMut,
13104    p2: SendMut,
13105    start: usize,
13106    end: usize,
13107) {
13108    for o in start..end {
13109        let row = &q[o * cols..(o + 1) * cols];
13110        // SAFETY: disjoint row ranges per worker.
13111        unsafe {
13112            *p1.at(o) = row_dot_avx2(row, a1) * row_scale[o];
13113            *p2.at(o) = row_dot_avx2(row, a2) * row_scale[o];
13114        }
13115    }
13116}
13117
13118// ───────────────────── A8W8 SDOT path (port of vmfcore, ×1.78 decode) ─────────────────────
13119
13120/// ARMv8.6 i8mm (smmla): 32 int8 MACs per instruction vs sdot's 16 —
13121/// yet MEASURED 2.4× SLOWER than the blocked sdot on Apple silicon
13122/// (108 vs 264 GF/s): the on-the-fly vcombine packing and the two-
13123/// accumulator dependency chain swamp the MAC advantage, and Apple's
13124/// four SIMD pipes already keep sdot fed. OPT-IN (CMF_I8MM=1) for
13125/// field trials on Cortex-A710/X-class parts with two pipes, where the
13126/// balance may differ; a pre-interleaved weight layout (repack infra)
13127/// is the known path if it ever earns its keep.
13128#[cfg(target_arch = "aarch64")]
13129fn i8mm_enabled() -> bool {
13130    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
13131    *ON.get_or_init(|| {
13132        std::env::var("CMF_I8MM").map(|v| v == "1").unwrap_or(false)
13133            && std::arch::is_aarch64_feature_detected!("i8mm")
13134    })
13135}
13136
13137/// SDOT enabled? Default ON when the CPU has ARMv8.2 dotprod;
13138/// `CMF_SDOT=0` disables (falls back to i8×f32 NEON).
13139/// (On non-ARM release builds only the test tolerance switch calls it.)
13140#[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
13141fn sdot_enabled() -> bool {
13142    if FLOAT_ACTIVATIONS.get() {
13143        return false;
13144    }
13145    use std::sync::OnceLock;
13146    static ON: OnceLock<bool> = OnceLock::new();
13147    *ON.get_or_init(|| {
13148        let want = std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true);
13149        if !want {
13150            return false;
13151        }
13152
13153        #[cfg(target_arch = "aarch64")]
13154        {
13155            if std::arch::is_aarch64_feature_detected!("dotprod") {
13156                return true;
13157            }
13158            #[cfg(target_os = "android")]
13159            {
13160                if let Ok(cpuinfo) = std::fs::read_to_string("/proc/cpuinfo") {
13161                    if cpuinfo.lines().any(|l| {
13162                        (l.starts_with("Features") || l.starts_with("features"))
13163                            && l.contains("asimddp")
13164                    }) {
13165                        return true;
13166                    }
13167                }
13168            }
13169            false
13170        }
13171        #[cfg(not(target_arch = "aarch64"))]
13172        {
13173            false
13174        }
13175    })
13176}
13177
13178/// Two-field activation split (≡ vmfcore `q8_split_prep`): outlier
13179/// channels (>8·rms) are computed exactly in f32; the bulk (outliers
13180/// zeroed → clean absmax) goes through int8 SDOT. Computed ONCE per
13181/// matvec, shared by all rows/workers.
13182struct SplitAct {
13183    xq: Vec<i8>,
13184    sx: f32,
13185    outliers: Vec<(usize, f32)>,
13186    /// Σ xq — the VNNI bias-trick correction (`(w+128)·x` sums need
13187    /// `−128·Σx`); one i32 per split, computed once per matvec.
13188    #[cfg_attr(not(target_arch = "x86_64"), allow(dead_code))]
13189    xsum: i32,
13190    /// `Σ xq` per 32-element group, already f32 — the q4tp `(w+8)·x − 8·Σx`
13191    /// correction for the VERTICAL accumulator: the f32x8 partials reduce
13192    /// through an 8-lane horizontal sum, and each lane subtracts this once,
13193    /// so the row total carries exactly `8·Σx`. Precomputing it lets the
13194    /// AVX2 kernels drop the per-byte `vpaddb`/`vpabsb`/`vpsignb` (two of
13195    /// the three sat on port 5) and feed the raw unsigned nibbles straight
13196    /// to `vpmaddubsw`. Every value is an integer ≤ 32·127, so both the
13197    /// f32 conversion and each lane's subtraction are exact and the cell
13198    /// equals the old `(nib − 8)` dot bit for bit.
13199    #[cfg_attr(not(target_arch = "x86_64"), allow(dead_code))]
13200    xg8: Vec<f32>,
13201}
13202
13203thread_local! {
13204    /// Recycled xq buffers: split_act runs for every matvec (~200/token)
13205    /// and its hidden-size allocation was steady-state heap churn.
13206    static XQ_FREE: std::cell::RefCell<Vec<Vec<i8>>> =
13207        const { std::cell::RefCell::new(Vec::new()) };
13208    /// Recycled per-group correction buffers (see `SplitAct::xg8`).
13209    static XG8_FREE: std::cell::RefCell<Vec<Vec<f32>>> =
13210        const { std::cell::RefCell::new(Vec::new()) };
13211}
13212
13213/// Pop a recycled `xg8` buffer, or mint one; filled by [`xg8_of`].
13214fn xg8_take() -> Vec<f32> {
13215    let mut v = XG8_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
13216    v.clear();
13217    v
13218}
13219
13220impl Drop for SplitAct {
13221    fn drop(&mut self) {
13222        let buf = std::mem::take(&mut self.xq);
13223        if buf.capacity() > 0 {
13224            XQ_FREE.with(|f| {
13225                let mut f = f.borrow_mut();
13226                if f.len() < 16 {
13227                    f.push(buf);
13228                }
13229            });
13230        }
13231        let buf = std::mem::take(&mut self.xg8);
13232        if buf.capacity() > 0 {
13233            XG8_FREE.with(|f| {
13234                let mut f = f.borrow_mut();
13235                if f.len() < 16 {
13236                    f.push(buf);
13237                }
13238            });
13239        }
13240    }
13241}
13242
13243/// Fill `out` with `Σ xq` per group — see [`SplitAct::xg8`]. A tail
13244/// shorter than one group is dropped (`chunks_exact`); kernels guard on
13245/// `xg8.len() == gpr`.
13246fn xg8_fill(xq: &[i8], out: &mut Vec<f32>) {
13247    out.clear();
13248    out.extend(xq.chunks_exact(GROUP_SIZE).map(|c| {
13249        let mut s = 0i32;
13250        for &v in c {
13251            s += v as i32;
13252        }
13253        s as f32
13254    }));
13255}
13256
13257thread_local! {
13258    /// One scratch row per WORKER, kept for the life of the thread.
13259    ///
13260    /// The kernels take a row of group scales per dispatch, and a fresh
13261    /// `vec![0f32; gpr]` inside the closure is one allocation per worker per
13262    /// dispatch — on the release checkpoint about six thousand a token, a
13263    /// quarter of everything the benchmark counts.
13264    static KROW: UnsafeCell<Vec<f32>> = const { UnsafeCell::new(Vec::new()) };
13265    /// Two scratch rows for a gate/up pair.  Q2TP and Q4TP fused SwiGLU
13266    /// decode each projection's ladder independently, but both ladders have
13267    /// the same lifetime; retaining them per worker removes two heap trips
13268    /// from every dispatched expert pair.
13269    static KROWS: UnsafeCell<[Vec<f32>; 2]> =
13270        const { UnsafeCell::new([Vec::new(), Vec::new()]) };
13271}
13272
13273/// Borrow `n` floats of the calling worker's scratch. Nothing inside a
13274/// kernel body borrows it again, which is what keeps the RefCell honest.
13275#[inline]
13276fn with_krow<R>(n: usize, f: impl FnOnce(&mut [f32]) -> R) -> R {
13277    KROW.with(|s| {
13278        // SAFETY: `KROW` is thread-local and this function does not recurse;
13279        // each worker has exclusive access to its own scratch row.
13280        let b = unsafe { &mut *s.get() };
13281        if b.len() < n {
13282            b.resize(n, 0.0);
13283        }
13284        f(&mut b[..n])
13285    })
13286}
13287
13288/// `t.round().clamp(-127.0, 127.0) as i8`, bit for bit, without the libm
13289/// call. On baseline x86-64 (no SSE4.1 `roundps`) `f32::round` is a
13290/// function call per element, and split_act runs it over every hidden
13291/// state before every matvec: measured 27 us a call on a 2048-wide
13292/// activation on an EPYC 7763 — 5.4 ms of a 55 ms decode token, all of
13293/// it on the caller's thread while thirty workers wait. Clamping first is
13294/// equivalent (round is monotonic and ±127 are integers), and after the
13295/// clamp `t - trunc(t)` is exact, so the half-away-from-zero decision is
13296/// the one `round` makes. NaN clamps to NaN and converts to 0, as before.
13297/// The loop vectorizes (cvttps2dq + compare/select).
13298#[inline(always)]
13299fn q8_round(t: f32) -> i8 {
13300    let t = t.clamp(-127.0, 127.0);
13301    let i = t as i32;
13302    let f = t - i as f32;
13303    let r = if f >= 0.5 {
13304        i + 1
13305    } else if f <= -0.5 {
13306        i - 1
13307    } else {
13308        i
13309    };
13310    r as i8
13311}
13312
13313#[inline]
13314fn with_krows<R>(n: usize, f: impl FnOnce(&mut [f32], &mut [f32]) -> R) -> R {
13315    KROWS.with(|s| {
13316        // SAFETY: KROWS is thread-local and no kernel body recursively calls
13317        // this helper; each participant owns its two vectors exclusively.
13318        let b = unsafe { &mut *s.get() };
13319        for row in b.iter_mut() {
13320            if row.len() < n {
13321                row.resize(n, 0.0);
13322            }
13323        }
13324        let (a, b) = b.split_at_mut(1);
13325        f(&mut a[0][..n], &mut b[0][..n])
13326    })
13327}
13328
13329#[inline]
13330fn silu_mul_limited(mut gate: f32, mut up: f32, limit: f32) -> f32 {
13331    if limit > 0.0 {
13332        up = up.clamp(-limit, limit);
13333        gate = gate.min(limit);
13334    }
13335    gate / (1.0 + (-gate).exp()) * up
13336}
13337
13338/// `split_act` of every column of a batch, the columns spread over the
13339/// pool from 16 on. Each column's split is its own (the same bits as the
13340/// serial loop); serially it was 25 µs a column, 6 ms of a Spark-4B
13341/// prefill token on an 8-thread EPYC while the workers idled.
13342fn split_acts_pooled(xs_all: &[f32], b: usize, cols: usize, pool: Option<&Pool>) -> Vec<SplitAct> {
13343    match pool {
13344        Some(p) if b >= 16 => {
13345            let slots: Vec<std::sync::Mutex<Option<SplitAct>>> =
13346                (0..b).map(|_| std::sync::Mutex::new(None)).collect();
13347            p.run_rows(b, &|start, end| {
13348                for bi in start..end {
13349                    let a = split_act(&xs_all[bi * cols..(bi + 1) * cols]);
13350                    *slots[bi].lock().unwrap_or_else(|e| e.into_inner()) = Some(a);
13351                }
13352            });
13353            slots
13354                .into_iter()
13355                .map(|m| {
13356                    m.into_inner()
13357                        .unwrap_or_else(|e| e.into_inner())
13358                        .expect("every column split")
13359                })
13360                .collect()
13361        }
13362        _ => (0..b)
13363            .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
13364            .collect(),
13365    }
13366}
13367
13368fn split_act(x: &[f32]) -> SplitAct {
13369    let _prof = crate::cpuprof::time(crate::cpuprof::Slot::SplitAct);
13370    let n = x.len();
13371    let rms = (x.iter().map(|&v| (v * v) as f64).sum::<f64>() / n.max(1) as f64).sqrt() as f32;
13372    let thr = 8.0 * rms;
13373    // One pass: collect outliers and the bulk absmax (outliers excluded —
13374    // identical to the old zero-then-fold over a copied buffer, minus the
13375    // full-vector copy).
13376    let mut outliers: Vec<(usize, f32)> = Vec::new();
13377    let mut amax = 0f32;
13378    for (j, &v) in x.iter().enumerate() {
13379        let a = v.abs();
13380        if a > thr {
13381            outliers.push((j, v));
13382        } else if a > amax {
13383            amax = a;
13384        }
13385    }
13386    let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
13387    let inv = 1.0 / sx;
13388    crate::cpuprof::count_outliers(outliers.len() as u64);
13389    let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
13390    xq.clear();
13391    xq.reserve(n);
13392    if outliers.is_empty() {
13393        xq.extend(
13394            x.iter()
13395                .map(|&v| q8_round(v * inv)),
13396        );
13397    } else {
13398        // Outlier slots quantize to 0 (their exact term is added later).
13399        xq.extend(x.iter().map(|&v| {
13400            if v.abs() > thr {
13401                0
13402            } else {
13403                q8_round(v * inv)
13404            }
13405        }));
13406    }
13407    let xsum = xq.iter().map(|&v| v as i32).sum();
13408    let mut xg8 = xg8_take();
13409    xg8_fill(&xq, &mut xg8);
13410    SplitAct {
13411        xq,
13412        sx,
13413        outliers,
13414        xsum,
13415        xg8,
13416    }
13417}
13418
13419fn split_act_q8_2f(x: &[f32], col: &[f32]) -> SplitAct {
13420    let _prof = crate::cpuprof::time(crate::cpuprof::Slot::SplitAct);
13421    let n = x.len();
13422    let rms = (x
13423        .iter()
13424        .zip(col)
13425        .map(|(&a, &c)| {
13426            let v = a * c;
13427            (v * v) as f64
13428        })
13429        .sum::<f64>()
13430        / n.max(1) as f64)
13431        .sqrt() as f32;
13432    let thr = 8.0 * rms;
13433
13434    let mut outliers = Vec::new();
13435    let mut amax = 0f32;
13436    for (j, (&a, &c)) in x.iter().zip(col).enumerate() {
13437        let v = a * c;
13438        let s = v.abs();
13439        if s > thr {
13440            outliers.push((j, v));
13441        } else if s > amax {
13442            amax = s;
13443        }
13444    }
13445
13446    let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
13447    let inv = 1.0 / sx;
13448    let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
13449    xq.clear();
13450    xq.reserve(n);
13451    if outliers.is_empty() {
13452        xq.extend(
13453            x.iter()
13454                .zip(col)
13455                .map(|(&a, &c)| q8_round((a * c) * inv)),
13456        );
13457    } else {
13458        xq.extend(x.iter().zip(col).map(|(&a, &c)| {
13459            let v = a * c;
13460            if v.abs() > thr {
13461                0
13462            } else {
13463                q8_round(v * inv)
13464            }
13465        }));
13466    }
13467    let xsum = xq.iter().map(|&v| v as i32).sum();
13468    let mut xg8 = xg8_take();
13469    xg8_fill(&xq, &mut xg8);
13470    SplitAct {
13471        xq,
13472        sx,
13473        outliers,
13474        xsum,
13475        xg8,
13476    }
13477}
13478
13479/// int8(weight)·int8(activation) → i32 via `sdot` (inline asm — the
13480/// vdotq intrinsic is unstable; port of vmfcore `dot_i8_sdot`).
13481#[cfg(target_arch = "aarch64")]
13482#[target_feature(enable = "neon,dotprod")]
13483unsafe fn dot_i8_sdot(w: &[u8], xq: &[i8]) -> i32 {
13484    // SAFETY: callers uphold slice-length contracts (see call sites).
13485    unsafe {
13486        use core::arch::aarch64::*;
13487        use core::arch::asm;
13488        let wp = w.as_ptr() as *const i8;
13489        let n = w.len();
13490        let (mut a0, mut a1, mut a2, mut a3) = (
13491            vdupq_n_s32(0),
13492            vdupq_n_s32(0),
13493            vdupq_n_s32(0),
13494            vdupq_n_s32(0),
13495        );
13496        let mut i = 0;
13497        while i + 64 <= n {
13498            let (w0, x0) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
13499            let (w1, x1) = (vld1q_s8(wp.add(i + 16)), vld1q_s8(xq.as_ptr().add(i + 16)));
13500            let (w2, x2) = (vld1q_s8(wp.add(i + 32)), vld1q_s8(xq.as_ptr().add(i + 32)));
13501            let (w3, x3) = (vld1q_s8(wp.add(i + 48)), vld1q_s8(xq.as_ptr().add(i + 48)));
13502            asm!(
13503                "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
13504                "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
13505                "sdot {a2:v}.4s, {w2:v}.16b, {x2:v}.16b",
13506                "sdot {a3:v}.4s, {w3:v}.16b, {x3:v}.16b",
13507                a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
13508                w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
13509                w2 = in(vreg) w2, x2 = in(vreg) x2, w3 = in(vreg) w3, x3 = in(vreg) x3,
13510                options(pure, nomem, nostack),
13511            );
13512            i += 64;
13513        }
13514        while i + 16 <= n {
13515            let (wv, xv) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
13516            asm!("sdot {a:v}.4s, {w:v}.16b, {x:v}.16b",
13517                 a = inout(vreg) a0, w = in(vreg) wv, x = in(vreg) xv, options(pure, nomem, nostack));
13518            i += 16;
13519        }
13520        let mut s = vaddvq_s32(vaddq_s32(vaddq_s32(a0, a1), vaddq_s32(a2, a3)));
13521        while i < n {
13522            s += (*wp.add(i)) as i32 * xq[i] as i32;
13523            i += 1;
13524        }
13525        s
13526    }
13527}
13528
13529/// Row-blocked SDOT: 4 output rows per pass — the activation chunk is
13530/// loaded once and reused, 4 independent accumulators hide sdot latency
13531/// (port of vmfcore `dot_i8_sdot_4rows`).
13532#[cfg(target_arch = "aarch64")]
13533#[target_feature(enable = "neon,dotprod")]
13534unsafe fn dot_i8_sdot_4rows(w0: &[u8], w1: &[u8], w2: &[u8], w3: &[u8], xq: &[i8]) -> [i32; 4] {
13535    // SAFETY: callers uphold slice-length contracts (see call sites).
13536    unsafe {
13537        use core::arch::aarch64::*;
13538        use core::arch::asm;
13539        let n = xq.len();
13540        let px = xq.as_ptr();
13541        let (p0, p1, p2, p3) = (
13542            w0.as_ptr() as *const i8,
13543            w1.as_ptr() as *const i8,
13544            w2.as_ptr() as *const i8,
13545            w3.as_ptr() as *const i8,
13546        );
13547        let (mut a0, mut a1, mut a2, mut a3) = (
13548            vdupq_n_s32(0),
13549            vdupq_n_s32(0),
13550            vdupq_n_s32(0),
13551            vdupq_n_s32(0),
13552        );
13553        let mut i = 0;
13554        while i + 16 <= n {
13555            let x = vld1q_s8(px.add(i));
13556            let v0 = vld1q_s8(p0.add(i));
13557            let v1 = vld1q_s8(p1.add(i));
13558            let v2 = vld1q_s8(p2.add(i));
13559            let v3 = vld1q_s8(p3.add(i));
13560            asm!(
13561                "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
13562                "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
13563                "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
13564                "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
13565                a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
13566                v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
13567                options(pure, nomem, nostack),
13568            );
13569            i += 16;
13570        }
13571        let mut r = [
13572            vaddvq_s32(a0),
13573            vaddvq_s32(a1),
13574            vaddvq_s32(a2),
13575            vaddvq_s32(a3),
13576        ];
13577        while i < n {
13578            let xi = *px.add(i) as i32;
13579            r[0] += (*p0.add(i)) as i32 * xi;
13580            r[1] += (*p1.add(i)) as i32 * xi;
13581            r[2] += (*p2.add(i)) as i32 * xi;
13582            r[3] += (*p3.add(i)) as i32 * xi;
13583            i += 1;
13584        }
13585        r
13586    }
13587}
13588
13589/// 4 interleaved rows in one pass: the repacked group is [r0[c], r1[c],
13590/// r2[c], r3[c]] per 16-byte chunk, so each iteration reads ONE 64-byte
13591/// line plus the shared activation chunk — a single sequential weight
13592/// stream per worker. Per-row accumulation is the same one-accumulator
13593/// scheme as `dot_i8_sdot_4rows`; integer sums are exact, so outputs
13594/// are bit-identical to the mmap-layout kernel.
13595#[cfg(target_arch = "aarch64")]
13596#[target_feature(enable = "neon,dotprod")]
13597unsafe fn dot_i8_sdot_4rows_il(g: &[u8], xq: &[i8]) -> [i32; 4] {
13598    // SAFETY: callers uphold slice-length contracts (g.len() == 4·n,
13599    // n % 16 == 0 — guaranteed by the repack gate).
13600    unsafe {
13601        use core::arch::aarch64::*;
13602        use core::arch::asm;
13603        let n = xq.len();
13604        let px = xq.as_ptr();
13605        let pg = g.as_ptr() as *const i8;
13606        let (mut a0, mut a1, mut a2, mut a3) = (
13607            vdupq_n_s32(0),
13608            vdupq_n_s32(0),
13609            vdupq_n_s32(0),
13610            vdupq_n_s32(0),
13611        );
13612        let mut i = 0;
13613        while i + 16 <= n {
13614            let x = vld1q_s8(px.add(i));
13615            let base = pg.add(4 * i);
13616            let v0 = vld1q_s8(base);
13617            let v1 = vld1q_s8(base.add(16));
13618            let v2 = vld1q_s8(base.add(32));
13619            let v3 = vld1q_s8(base.add(48));
13620            asm!(
13621                "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
13622                "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
13623                "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
13624                "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
13625                a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
13626                v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
13627                options(pure, nomem, nostack),
13628            );
13629            i += 16;
13630        }
13631        [
13632            vaddvq_s32(a0),
13633            vaddvq_s32(a1),
13634            vaddvq_s32(a2),
13635            vaddvq_s32(a3),
13636        ]
13637    }
13638}
13639
13640/// One q8 row range via SDOT (4-row blocks + tail) — the body of
13641/// `qmatvec`'s hot loop, extracted so multi-matrix jobs can drive the
13642/// SAME kernel for several tensors under one pool dispatch. `rep` — the
13643/// load-time interleaved repack (empty = mmap layout only); rows outside
13644/// full 4-row groups always come from the mmap layout.
13645#[cfg(target_arch = "aarch64")]
13646fn q8_range_sdot(
13647    q: &[u8],
13648    rep: &[u8],
13649    row_scale: &[f32],
13650    act: &SplitAct,
13651    cols: usize,
13652    out_addr: SendMut,
13653    start: usize,
13654    end: usize,
13655) {
13656    let mut o = start;
13657    // Leading rows to the group boundary (repack path only): the pool
13658    // splits row ranges arbitrarily, groups are absolute.
13659    if !rep.is_empty() {
13660        while o < end && o % 4 != 0 {
13661            let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
13662            unsafe { *out_addr.at(o) = v };
13663            o += 1;
13664        }
13665    }
13666    while o + 4 <= end {
13667        let r = if rep.is_empty() {
13668            unsafe {
13669                dot_i8_sdot_4rows(
13670                    &q[o * cols..(o + 1) * cols],
13671                    &q[(o + 1) * cols..(o + 2) * cols],
13672                    &q[(o + 2) * cols..(o + 3) * cols],
13673                    &q[(o + 3) * cols..(o + 4) * cols],
13674                    &act.xq,
13675                )
13676            }
13677        } else {
13678            unsafe { dot_i8_sdot_4rows_il(&rep[o * cols..(o + 4) * cols], &act.xq) }
13679        };
13680        for k in 0..4 {
13681            let mut acc = r[k] as f32 * act.sx;
13682            for &(j, xv) in &act.outliers {
13683                acc += (q[(o + k) * cols + j] as i8) as f32 * xv;
13684            }
13685            // SAFETY: disjoint row ranges per worker.
13686            unsafe { *out_addr.at(o + k) = acc * row_scale[o + k] };
13687        }
13688        o += 4;
13689    }
13690    while o < end {
13691        let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
13692        unsafe { *out_addr.at(o) = v };
13693        o += 1;
13694    }
13695}
13696
13697/// Two-input q8 row range via SDOT — `qmatvec2`'s hot loop, extracted
13698/// for the fused pair multi-matrix job (`matvec2_many`).
13699#[cfg(target_arch = "aarch64")]
13700#[allow(clippy::too_many_arguments)]
13701fn q8_range2_sdot(
13702    q: &[u8],
13703    row_scale: &[f32],
13704    a1: &SplitAct,
13705    a2: &SplitAct,
13706    cols: usize,
13707    p1: SendMut,
13708    p2: SendMut,
13709    start: usize,
13710    end: usize,
13711) {
13712    for o in start..end {
13713        let row = &q[o * cols..(o + 1) * cols];
13714        // SAFETY: disjoint row ranges per worker.
13715        unsafe {
13716            *p1.at(o) = row_dot_sdot(row, a1) * row_scale[o];
13717            *p2.at(o) = row_dot_sdot(row, a2) * row_scale[o];
13718        }
13719    }
13720}
13721
13722/// Two-input q8 row range, f32 kernel (non-SDOT) — same extraction.
13723#[allow(clippy::too_many_arguments)]
13724fn q8_range2_f32(
13725    q: &[u8],
13726    row_scale: &[f32],
13727    x1: &[f32],
13728    x2: &[f32],
13729    cols: usize,
13730    p1: SendMut,
13731    p2: SendMut,
13732    start: usize,
13733    end: usize,
13734) {
13735    for o in start..end {
13736        let row = &q[o * cols..(o + 1) * cols];
13737        // SAFETY: disjoint row ranges per worker.
13738        unsafe {
13739            *p1.at(o) = dot_i8_f32(row, x1) * row_scale[o];
13740            *p2.at(o) = dot_i8_f32(row, x2) * row_scale[o];
13741        }
13742    }
13743}
13744
13745/// Scalar/NEON-f32 q8 row range (non-SDOT platforms) — same extraction.
13746fn q8_range_f32(
13747    q: &[u8],
13748    row_scale: &[f32],
13749    xs: &[f32],
13750    cols: usize,
13751    out_addr: SendMut,
13752    start: usize,
13753    end: usize,
13754) {
13755    for o in start..end {
13756        let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
13757        // SAFETY: disjoint row ranges per worker.
13758        unsafe { *out_addr.at(o) = v };
13759    }
13760}
13761
13762/// One q8 row against a split activation, portable: the per-arch fast
13763/// dots where they exist, the exact scalar loop elsewhere. The scalar
13764/// arm is also the test oracle for both fast arms.
13765#[inline]
13766fn q8_row_dot(row: &[u8], act: &SplitAct) -> f32 {
13767    #[cfg(target_arch = "aarch64")]
13768    return row_dot_sdot(row, act);
13769    #[cfg(target_arch = "x86_64")]
13770    return row_dot_avx2(row, act);
13771    #[allow(unreachable_code)]
13772    q8_row_dot_scalar(row, act)
13773}
13774
13775#[allow(dead_code)]
13776fn q8_row_dot_scalar(row: &[u8], act: &SplitAct) -> f32 {
13777    let mut acc = 0i32;
13778    for (k, &b) in row.iter().enumerate() {
13779        acc += (b as i8) as i32 * act.xq[k] as i32;
13780    }
13781    let mut acc = acc as f32 * act.sx;
13782    for &(j, xv) in &act.outliers {
13783        acc += (row[j] as i8) as f32 * xv;
13784    }
13785    acc
13786}
13787
13788/// SDOT row dot with exact outlier correction:
13789/// `dot = sdot(w, xq)·sx + Σ_outl w[j]·x[j]` (then × row_scale by caller).
13790#[cfg(target_arch = "aarch64")]
13791#[inline]
13792fn row_dot_sdot(row: &[u8], act: &SplitAct) -> f32 {
13793    let mut acc = unsafe { dot_i8_sdot(row, &act.xq) } as f32 * act.sx;
13794    for &(j, xv) in &act.outliers {
13795        acc += (row[j] as i8) as f32 * xv;
13796    }
13797    acc
13798}
13799
13800/// One q4 row via SDOT: each 32-group's nibbles unpack to centered i8
13801/// (nib−8 ∈ [−8,7]), int8×int8 `sdot` against the pre-quantized
13802/// activation group, × the group's f16 scale. Returns Σ_g dot_g·s_g;
13803/// the caller multiplies by the activation scale and adds the exact
13804/// outlier terms (port of vmfcore `dot_q4_block_sdot`, +23% measured).
13805/// Nibble order matches the writer: element 2k = low nibble, 2k+1 = high
13806/// → zip(lo,hi) restores flat order.
13807#[cfg(target_arch = "aarch64")]
13808#[target_feature(enable = "neon,dotprod")]
13809unsafe fn dot_q4_row_sdot(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
13810    // SAFETY: callers uphold slice-length contracts (16 packed bytes and
13811    // 2 scale bytes per group; xq.len() == gpr·GROUP_SIZE).
13812    unsafe {
13813        use core::arch::aarch64::*;
13814        use core::arch::asm;
13815        let lomask = vdupq_n_u8(0x0F);
13816        let eight = vdupq_n_s8(8);
13817        let mut acc = 0f32;
13818        for gi in 0..gpr {
13819            let g = g0 + gi;
13820            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
13821            let b = vld1q_u8(packed.as_ptr().add(g * 16));
13822            let lo = vandq_u8(b, lomask);
13823            let hi = vshrq_n_u8::<4>(b);
13824            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
13825            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
13826            let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
13827            let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
13828            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
13829            asm!(
13830                "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
13831                "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
13832                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
13833                e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
13834                options(pure, nomem, nostack),
13835            );
13836            acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
13837        }
13838        acc
13839    }
13840}
13841
13842/// Two-activation q4 row via SDOT: the nibble unpack (the expensive
13843/// part) happens ONCE per group; both pre-quantized activations are
13844/// dotted against the same centered i8 registers. Per-lane math matches
13845/// `dot_q4_row_sdot` exactly.
13846#[cfg(target_arch = "aarch64")]
13847#[target_feature(enable = "neon,dotprod")]
13848unsafe fn dot_q4_row_sdot2(
13849    packed: &[u8],
13850    scales: &[u8],
13851    g0: usize,
13852    gpr: usize,
13853    xq1: &[i8],
13854    xq2: &[i8],
13855) -> (f32, f32) {
13856    // SAFETY: callers uphold slice-length contracts (16 packed bytes and
13857    // 2 scale bytes per group; xq*.len() == gpr·GROUP_SIZE).
13858    unsafe {
13859        use core::arch::aarch64::*;
13860        use core::arch::asm;
13861        let lomask = vdupq_n_u8(0x0F);
13862        let eight = vdupq_n_s8(8);
13863        let (mut acc1, mut acc2) = (0f32, 0f32);
13864        for gi in 0..gpr {
13865            let g = g0 + gi;
13866            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
13867            let b = vld1q_u8(packed.as_ptr().add(g * 16));
13868            let lo = vandq_u8(b, lomask);
13869            let hi = vshrq_n_u8::<4>(b);
13870            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
13871            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
13872            let x10 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE));
13873            let x11 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE + 16));
13874            let x20 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE));
13875            let x21 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE + 16));
13876            let (mut a0, mut a1, mut b0, mut b1) = (
13877                vdupq_n_s32(0),
13878                vdupq_n_s32(0),
13879                vdupq_n_s32(0),
13880                vdupq_n_s32(0),
13881            );
13882            asm!(
13883                "sdot {a0:v}.4s, {e0:v}.16b, {x10:v}.16b",
13884                "sdot {a1:v}.4s, {e1:v}.16b, {x11:v}.16b",
13885                "sdot {b0:v}.4s, {e0:v}.16b, {x20:v}.16b",
13886                "sdot {b1:v}.4s, {e1:v}.16b, {x21:v}.16b",
13887                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
13888                b0 = inout(vreg) b0, b1 = inout(vreg) b1,
13889                e0 = in(vreg) e0, e1 = in(vreg) e1,
13890                x10 = in(vreg) x10, x11 = in(vreg) x11,
13891                x20 = in(vreg) x20, x21 = in(vreg) x21,
13892                options(pure, nomem, nostack),
13893            );
13894            acc1 += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
13895            acc2 += vaddvq_s32(vaddq_s32(b0, b1)) as f32 * s;
13896        }
13897        (acc1, acc2)
13898    }
13899}
13900
13901// ───────────────────── fused int8 kernels ─────────────────────
13902
13903/// `acc += w · row` where the row is centered i8 — NEON widen+fma on
13904/// aarch64, scalar elsewhere. The KV-cache q8 value path rides on this.
13905#[inline]
13906pub(crate) fn axpy_i8_f32(acc: &mut [f32], row: &[i8], w: f32) {
13907    #[cfg(target_arch = "aarch64")]
13908    unsafe {
13909        return axpy_i8_f32_neon(acc, row, w);
13910    }
13911    #[cfg(target_arch = "x86_64")]
13912    if avx2_enabled() {
13913        return unsafe { axpy_i8_f32_avx2(acc, row, w) };
13914    }
13915    #[allow(unreachable_code)]
13916    {
13917        for (a, &b) in acc.iter_mut().zip(row) {
13918            *a += w * b as f32;
13919        }
13920    }
13921}
13922
13923/// i8→f32 axpy via AVX2/FMA (x86 mirror of `axpy_i8_f32_neon`).
13924#[cfg(target_arch = "x86_64")]
13925#[target_feature(enable = "avx2,fma")]
13926unsafe fn axpy_i8_f32_avx2(acc: &mut [f32], row: &[i8], w: f32) {
13927    // SAFETY: callers uphold slice-length contracts (see call sites).
13928    unsafe {
13929        use core::arch::x86_64::*;
13930        let n = acc.len().min(row.len());
13931        let ap = acc.as_mut_ptr();
13932        let rp = row.as_ptr();
13933        let wv = _mm256_set1_ps(w);
13934        let mut j = 0usize;
13935        while j + 16 <= n {
13936            let rb = _mm_loadu_si128(rp.add(j) as *const __m128i);
13937            let lo = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(rb));
13938            let hi = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(rb)));
13939            let v0 = _mm256_fmadd_ps(wv, lo, _mm256_loadu_ps(ap.add(j)));
13940            let v1 = _mm256_fmadd_ps(wv, hi, _mm256_loadu_ps(ap.add(j + 8)));
13941            _mm256_storeu_ps(ap.add(j), v0);
13942            _mm256_storeu_ps(ap.add(j + 8), v1);
13943            j += 16;
13944        }
13945        while j < n {
13946            *ap.add(j) += w * (*rp.add(j)) as f32;
13947            j += 1;
13948        }
13949    }
13950}
13951
13952#[cfg(target_arch = "aarch64")]
13953#[target_feature(enable = "neon")]
13954unsafe fn axpy_i8_f32_neon(acc: &mut [f32], row: &[i8], w: f32) {
13955    // SAFETY: callers uphold slice-length contracts (see call sites).
13956    unsafe {
13957        use core::arch::aarch64::*;
13958        let n = acc.len().min(row.len());
13959        let ap = acc.as_mut_ptr();
13960        let rp = row.as_ptr();
13961        let wv = vdupq_n_f32(w);
13962        let mut j = 0usize;
13963        while j + 16 <= n {
13964            let rb = vld1q_s8(rp.add(j));
13965            let lo = vmovl_s8(vget_low_s8(rb));
13966            let hi = vmovl_s8(vget_high_s8(rb));
13967            for (off, half) in [(0, lo), (8, hi)] {
13968                let f0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(half)));
13969                let f1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(half)));
13970                let o = j + off;
13971                vst1q_f32(ap.add(o), vfmaq_f32(vld1q_f32(ap.add(o)), wv, f0));
13972                vst1q_f32(ap.add(o + 4), vfmaq_f32(vld1q_f32(ap.add(o + 4)), wv, f1));
13973            }
13974            j += 16;
13975        }
13976        while j < n {
13977            *ap.add(j) += w * (*rp.add(j)) as f32;
13978            j += 1;
13979        }
13980    }
13981}
13982
13983/// i8 row · f32 x. NEON on aarch64 (ported from vmfcore `dot_i8_f32_neon`,
13984/// ≈9× scalar), scalar elsewhere.
13985#[inline]
13986pub(crate) fn dot_i8_f32(w: &[u8], x: &[f32]) -> f32 {
13987    #[cfg(target_arch = "aarch64")]
13988    unsafe {
13989        return dot_i8_f32_neon(w, x);
13990    }
13991    #[cfg(target_arch = "x86_64")]
13992    if avx2_enabled() {
13993        return unsafe { dot_i8_f32_avx2(w, x) };
13994    }
13995    #[allow(unreachable_code)]
13996    {
13997        let mut sum = 0.0f32;
13998        for (j, &b) in w.iter().enumerate() {
13999            sum += (b as i8) as f32 * x[j];
14000        }
14001        sum
14002    }
14003}
14004
14005/// i8 row · (x ⊙ col_field) — the q8_2f row dot with the θ col-field
14006/// folded into the product (no prescaled copy of x). NEON on aarch64,
14007/// scalar elsewhere. Used by the active-neuron path `row_dot`.
14008#[inline]
14009fn dot_i8_col_f32(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
14010    #[cfg(target_arch = "aarch64")]
14011    unsafe {
14012        return dot_i8_col_f32_neon(w, x, col);
14013    }
14014    #[allow(unreachable_code)]
14015    {
14016        let mut sum = 0.0f32;
14017        for (j, &b) in w.iter().enumerate() {
14018            sum += (b as i8) as f32 * x[j] * col[j];
14019        }
14020        sum
14021    }
14022}
14023
14024#[cfg(target_arch = "aarch64")]
14025#[target_feature(enable = "neon")]
14026unsafe fn dot_i8_col_f32_neon(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
14027    // SAFETY: callers uphold slice-length contracts (see call sites).
14028    unsafe {
14029        use core::arch::aarch64::*;
14030        let n = x.len();
14031        let wp = w.as_ptr() as *const i8;
14032        let xp = x.as_ptr();
14033        let cp = col.as_ptr();
14034        let (mut a0, mut a1, mut a2, mut a3) = (
14035            vdupq_n_f32(0.0),
14036            vdupq_n_f32(0.0),
14037            vdupq_n_f32(0.0),
14038            vdupq_n_f32(0.0),
14039        );
14040        let mut j = 0usize;
14041        while j + 16 <= n {
14042            let wb = vld1q_s8(wp.add(j));
14043            let lo = vmovl_s8(vget_low_s8(wb));
14044            let hi = vmovl_s8(vget_high_s8(wb));
14045            let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
14046            let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
14047            let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
14048            let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
14049            a0 = vfmaq_f32(
14050                a0,
14051                w0,
14052                vmulq_f32(vld1q_f32(xp.add(j)), vld1q_f32(cp.add(j))),
14053            );
14054            a1 = vfmaq_f32(
14055                a1,
14056                w1,
14057                vmulq_f32(vld1q_f32(xp.add(j + 4)), vld1q_f32(cp.add(j + 4))),
14058            );
14059            a2 = vfmaq_f32(
14060                a2,
14061                w2,
14062                vmulq_f32(vld1q_f32(xp.add(j + 8)), vld1q_f32(cp.add(j + 8))),
14063            );
14064            a3 = vfmaq_f32(
14065                a3,
14066                w3,
14067                vmulq_f32(vld1q_f32(xp.add(j + 12)), vld1q_f32(cp.add(j + 12))),
14068            );
14069            j += 16;
14070        }
14071        let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
14072        while j < n {
14073            sum += (*wp.add(j)) as f32 * *xp.add(j) * *cp.add(j);
14074            j += 1;
14075        }
14076        sum
14077    }
14078}
14079
14080#[cfg(target_arch = "aarch64")]
14081#[target_feature(enable = "neon")]
14082unsafe fn dot_i8_f32_neon(w: &[u8], x: &[f32]) -> f32 {
14083    // SAFETY: callers uphold slice-length contracts (see call sites).
14084    unsafe {
14085        use core::arch::aarch64::*;
14086        let n = x.len();
14087        let wp = w.as_ptr() as *const i8;
14088        let xp = x.as_ptr();
14089        let (mut a0, mut a1, mut a2, mut a3) = (
14090            vdupq_n_f32(0.0),
14091            vdupq_n_f32(0.0),
14092            vdupq_n_f32(0.0),
14093            vdupq_n_f32(0.0),
14094        );
14095        let mut j = 0usize;
14096        while j + 16 <= n {
14097            let wb = vld1q_s8(wp.add(j));
14098            let lo = vmovl_s8(vget_low_s8(wb));
14099            let hi = vmovl_s8(vget_high_s8(wb));
14100            let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
14101            let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
14102            let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
14103            let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
14104            a0 = vfmaq_f32(a0, w0, vld1q_f32(xp.add(j)));
14105            a1 = vfmaq_f32(a1, w1, vld1q_f32(xp.add(j + 4)));
14106            a2 = vfmaq_f32(a2, w2, vld1q_f32(xp.add(j + 8)));
14107            a3 = vfmaq_f32(a3, w3, vld1q_f32(xp.add(j + 12)));
14108            j += 16;
14109        }
14110        let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
14111        while j < n {
14112            sum += (*wp.add(j)) as f32 * *xp.add(j);
14113            j += 1;
14114        }
14115        sum
14116    }
14117}
14118
14119#[allow(clippy::too_many_arguments)]
14120fn qmatvec(
14121    q: &[u8],
14122    rep: &[u8],
14123    row_scale: &[f32],
14124    x: &[f32],
14125    col_field: &[f32],
14126    dtype: TensorDtype,
14127    rows: usize,
14128    cols: usize,
14129    out: &mut [f32],
14130    pool: Option<&Pool>,
14131) {
14132    debug_assert_eq!(out.len(), rows);
14133    #[cfg(not(target_arch = "aarch64"))]
14134    let _ = rep;
14135
14136    #[cfg(target_arch = "aarch64")]
14137    if sdot_enabled() {
14138        let act = if dtype == TensorDtype::Q8_2f {
14139            split_act_q8_2f(x, col_field)
14140        } else {
14141            split_act(x)
14142        };
14143        let out_addr = SendMut(out.as_mut_ptr());
14144        let run_range = |start: usize, end: usize| {
14145            q8_range_sdot(q, rep, row_scale, &act, cols, out_addr, start, end)
14146        };
14147        match pool {
14148            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
14149            _ => run_range(0, rows),
14150        }
14151        return;
14152    }
14153    // x86 A8W8 via AVX2 maddubs — same quantized-activation contract as
14154    // the SDOT path (CMF_AVX2=0 keeps the exact i8×f32 loop).
14155    #[cfg(target_arch = "x86_64")]
14156    if avx2_a8w8_enabled() {
14157        let act = if dtype == TensorDtype::Q8_2f {
14158            split_act_q8_2f(x, col_field)
14159        } else {
14160            split_act(x)
14161        };
14162        let out_addr = SendMut(out.as_mut_ptr());
14163        let run_range = |start: usize, end: usize| {
14164            q8_range_avx2(q, row_scale, &act, cols, out_addr, start, end)
14165        };
14166        match pool {
14167            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
14168            _ => run_range(0, rows),
14169        }
14170        return;
14171    }
14172
14173    prescale_with(x, col_field, dtype, 1, |xs| {
14174        let out_addr = SendMut(out.as_mut_ptr());
14175        let run_range = move |start: usize, end: usize| {
14176            for o in start..end {
14177                let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
14178                // SAFETY: disjoint row ranges per worker.
14179                unsafe { *out_addr.at(o) = v };
14180            }
14181        };
14182        match pool {
14183            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
14184            _ => run_range(0, rows),
14185        }
14186    });
14187}
14188
14189#[allow(clippy::too_many_arguments)]
14190fn qmatvec2(
14191    q: &[u8],
14192    row_scale: &[f32],
14193    x1: &[f32],
14194    x2: &[f32],
14195    col_field: &[f32],
14196    dtype: TensorDtype,
14197    rows: usize,
14198    cols: usize,
14199    o1: &mut [f32],
14200    o2: &mut [f32],
14201    pool: Option<&Pool>,
14202) {
14203    #[cfg(target_arch = "aarch64")]
14204    if sdot_enabled() {
14205        let a1s = if dtype == TensorDtype::Q8_2f {
14206            split_act_q8_2f(x1, col_field)
14207        } else {
14208            split_act(x1)
14209        };
14210        let a2s = if dtype == TensorDtype::Q8_2f {
14211            split_act_q8_2f(x2, col_field)
14212        } else {
14213            split_act(x2)
14214        };
14215        let p1 = SendMut(o1.as_mut_ptr());
14216        let p2 = SendMut(o2.as_mut_ptr());
14217        let run_range = |start: usize, end: usize| {
14218            q8_range2_sdot(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
14219        };
14220        match pool {
14221            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
14222            _ => run_range(0, rows),
14223        }
14224        return;
14225    }
14226    #[cfg(target_arch = "x86_64")]
14227    if avx2_a8w8_enabled() {
14228        let a1s = if dtype == TensorDtype::Q8_2f {
14229            split_act_q8_2f(x1, col_field)
14230        } else {
14231            split_act(x1)
14232        };
14233        let a2s = if dtype == TensorDtype::Q8_2f {
14234            split_act_q8_2f(x2, col_field)
14235        } else {
14236            split_act(x2)
14237        };
14238        let p1 = SendMut(o1.as_mut_ptr());
14239        let p2 = SendMut(o2.as_mut_ptr());
14240        let run_range = |start: usize, end: usize| {
14241            q8_range2_avx2(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
14242        };
14243        match pool {
14244            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
14245            _ => run_range(0, rows),
14246        }
14247        return;
14248    }
14249
14250    prescale_with(x1, col_field, dtype, 1, |x1s| {
14251        prescale_with(x2, col_field, dtype, 2, |x2s| {
14252            let p1 = SendMut(o1.as_mut_ptr());
14253            let p2 = SendMut(o2.as_mut_ptr());
14254            let run_range = move |start: usize, end: usize| {
14255                for o in start..end {
14256                    let row = &q[o * cols..(o + 1) * cols];
14257                    let s1 = dot_i8_f32(row, x1s) * row_scale[o];
14258                    let s2 = dot_i8_f32(row, x2s) * row_scale[o];
14259                    // SAFETY: disjoint row ranges per worker.
14260                    unsafe {
14261                        *p1.at(o) = s1;
14262                        *p2.at(o) = s2;
14263                    }
14264                }
14265            };
14266            match pool {
14267                Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
14268                _ => run_range(0, rows),
14269            }
14270        });
14271    });
14272}
14273
14274#[derive(Clone, Copy)]
14275struct SendMut(*mut f32);
14276unsafe impl Send for SendMut {}
14277unsafe impl Sync for SendMut {}
14278
14279impl SendMut {
14280    #[inline]
14281    fn at(self, i: usize) -> *mut f32 {
14282        unsafe { self.0.add(i) }
14283    }
14284}
14285
14286#[cfg(test)]
14287mod tests {
14288
14289    /// `q8_round` must be `round().clamp(±127) as i8` bit for bit: every
14290    /// half-integer, their neighbours one ulp either side, the clamp
14291    /// boundary, huge values, infinities and NaN, plus a dense sweep.
14292    #[test]
14293    fn q8_round_is_round_clamp() {
14294        let reference = |t: f32| t.round().clamp(-127.0, 127.0) as i8;
14295        let mut probe = vec![
14296            0.0f32,
14297            -0.0,
14298            f32::NAN,
14299            f32::INFINITY,
14300            f32::NEG_INFINITY,
14301            f32::MAX,
14302            f32::MIN,
14303            1e30,
14304            -1e30,
14305            f32::MIN_POSITIVE,
14306            -f32::MIN_POSITIVE,
14307        ];
14308        for k in -300i32..=300 {
14309            let h = k as f32 * 0.5;
14310            let up = f32::from_bits(h.to_bits() + 1);
14311            let down = f32::from_bits(h.to_bits().wrapping_sub(1));
14312            for t in [h, up, down] {
14313                probe.push(t);
14314                probe.push(-t);
14315            }
14316        }
14317        let mut t = -140.0f32;
14318        while t < 140.0 {
14319            probe.push(t);
14320            t += 0.000_731;
14321        }
14322        for t in probe {
14323            assert_eq!(q8_round(t), reference(t), "t = {t:e} ({:#x})", t.to_bits());
14324        }
14325    }
14326
14327    use super::*;
14328
14329    /// The f32 `matmat` walks the batch four positions a row at a time; every
14330    /// output must still be the scalar-order sum `matvec` gives, bit for bit
14331    /// (a 7-row batch covers the four-wide block and the remainder).
14332    #[test]
14333    fn f32_matmat_equals_matvec_bitwise() {
14334        let (rows, cols, b) = (5usize, 37usize, 7usize);
14335        let w: Vec<f32> = (0..rows * cols).map(|i| ((i * 7919 % 113) as f32 - 56.0) / 37.0).collect();
14336        let xs: Vec<f32> = (0..b * cols).map(|i| ((i * 104_729 % 97) as f32 - 48.0) / 29.0).collect();
14337        let t = QTensor::from_f32(w, rows, cols);
14338        let mut out = vec![0f32; b * rows];
14339        t.matmat(&xs, b, &mut out, None);
14340        for bi in 0..b {
14341            let mut one = vec![0f32; rows];
14342            t.matvec(&xs[bi * cols..(bi + 1) * cols], &mut one, None);
14343            for o in 0..rows {
14344                assert_eq!(out[bi * rows + o].to_bits(), one[o].to_bits(), "bi {bi} row {o}");
14345            }
14346        }
14347    }
14348
14349    /// A short batch on a pool splits the f32 GEMM over its output rows
14350    /// (it used to run on one core); the sums must not change by a bit.
14351    #[test]
14352    fn f32_matmat_row_split_on_pool_is_bitwise() {
14353        let (rows, cols, b) = (300usize, 37usize, 7usize);
14354        let w: Vec<f32> = (0..rows * cols).map(|i| ((i * 7919 % 113) as f32 - 56.0) / 37.0).collect();
14355        let xs: Vec<f32> = (0..b * cols).map(|i| ((i * 104_729 % 97) as f32 - 48.0) / 29.0).collect();
14356        let t = QTensor::from_f32(w, rows, cols);
14357        let mut alone = vec![0f32; b * rows];
14358        t.matmat(&xs, b, &mut alone, None);
14359        let pool = Pool::new(3);
14360        let mut split = vec![0f32; b * rows];
14361        t.matmat(&xs, b, &mut split, Some(&pool));
14362        let bits = |v: &[f32]| v.iter().map(|x| x.to_bits()).collect::<Vec<_>>();
14363        assert_eq!(bits(&alone), bits(&split));
14364    }
14365
14366    #[test]
14367    fn q2tp_i8_dot_matches_exact_on_grid() {
14368        // On-grid activations (±1 → sx=1/127, xq=±127 dequantizes
14369        // exactly, no outliers) must make the integer path agree with
14370        // the exact scalar walk to f32 rounding.
14371        let (rows, cols) = (5, 64);
14372        let gpr = cols / GROUP_SIZE;
14373        // Synthetic codes plane + a flat ladder: scales_into is not under
14374        // test here, so drive dot_q2tp_row_i8 / q2tp_row_exact directly
14375        // with hand-made scales.
14376        let chunks: Vec<u8> = (0..rows * gpr * Q2TP_CHUNK)
14377            .map(|i| (i as u32).wrapping_mul(2654435761) as u8)
14378            .collect();
14379        let scales: Vec<f32> = (0..gpr).map(|g| 0.5 + g as f32 * 0.25).collect();
14380        let x: Vec<f32> = (0..cols)
14381            .map(|i| if i % 3 == 0 { -1.0 } else { 1.0 })
14382            .collect();
14383        let act = split_act(&x);
14384        assert!(
14385            act.outliers.is_empty(),
14386            "on-grid input must have no outliers"
14387        );
14388        let gsum = q1_group_sums(&act.xq, gpr);
14389        for r in 0..rows {
14390            let exact = q2tp_row_exact(&chunks, r, gpr, &x, &scales);
14391            let fast = dot_q2tp_row_i8(&chunks, r, gpr, &act.xq, &gsum, &scales) * act.sx;
14392            assert!(
14393                (exact - fast).abs() <= exact.abs() * 1e-5 + 1e-5,
14394                "row {r}: exact {exact} vs i8 {fast}"
14395            );
14396        }
14397    }
14398
14399    #[cfg(target_arch = "x86_64")]
14400    #[test]
14401    fn q2tp_avx2_dot_matches_scalar_for_random_patterns() {
14402        // Compare the release AVX2 integer dot against the scalar oracle over
14403        // arbitrary packed bytes/activation signs.  This guards the exact
14404        // table-load path used after rejecting a faster-looking decoder whose
14405        // full-checkpoint greedy output drifted.
14406        if !std::arch::is_x86_feature_detected!("avx2") {
14407            return;
14408        }
14409        let mut seed = 0x9e3779b9u32;
14410        let mut next = || {
14411            seed = seed.wrapping_mul(1664525).wrapping_add(1013904223);
14412            seed
14413        };
14414        for _ in 0..20_000 {
14415            let mut ch = [0u8; Q2TP_CHUNK];
14416            let mut x = [0i8; GROUP_SIZE];
14417            for b in &mut ch {
14418                *b = next() as u8;
14419            }
14420            for v in &mut x {
14421                *v = (next() >> 24) as i8;
14422            }
14423            let mut reference = 0i32;
14424            for (k, &b) in ch.iter().enumerate() {
14425                reference += (b & 3) as i32 * x[k * 4] as i32;
14426                reference += ((b >> 2) & 3) as i32 * x[k * 4 + 1] as i32;
14427                reference += ((b >> 4) & 3) as i32 * x[k * 4 + 2] as i32;
14428                reference += ((b >> 6) & 3) as i32 * x[k * 4 + 3] as i32;
14429            }
14430            // SAFETY: guarded by the runtime AVX2 feature check and fixed
14431            // 8-byte/32-byte slice lengths above.
14432            let got = unsafe { q2tp_code_dot_avx2(&ch, &x) };
14433            assert_eq!(got, reference, "packed q2 lane mismatch");
14434        }
14435    }
14436
14437    #[test]
14438    fn q2tp_affine_fuses_half_scale_correction_without_changing_raw_decode() {
14439        let (rows, cols) = (1usize, GROUP_SIZE);
14440        let mut bytes = vec![0u8; Q2TP_CHUNK + 4 + 1];
14441        // Repeating symbols 0,1,2,0 at unit scale.  q2tp's raw B is
14442        // (c-1.5), while the affine Prism operator is (c-1.0).
14443        bytes[..Q2TP_CHUNK].fill(0x24); // codes 0,1,2,0 in LSB-first order
14444        bytes[Q2TP_CHUNK..Q2TP_CHUNK + 2].copy_from_slice(&0u16.to_le_bytes());
14445        bytes[Q2TP_CHUNK + 2..Q2TP_CHUNK + 4].copy_from_slice(&0u16.to_le_bytes());
14446        bytes[Q2TP_CHUNK + 4] = 1; // dtype16 rung 1 = 1.0
14447        let x = vec![1.0f32; cols];
14448        let mut raw = vec![0.0f32; rows];
14449        let mut affine = vec![0.0f32; rows];
14450        q2tp_matvec_for_test(&bytes, &x, rows, cols, &mut raw);
14451        q2tp_affine_matvec_for_test(&bytes, &x, rows, cols, &mut affine);
14452        assert_eq!(raw, vec![-24.0]);
14453        assert_eq!(affine, vec![-8.0]);
14454        assert!((affine[0] - (raw[0] + 0.5 * cols as f32)).abs() < 1e-6);
14455    }
14456
14457    #[test]
14458    fn q8_row_dot_fast_matches_scalar() {
14459        // The per-arch fast dot must agree with the exact scalar oracle
14460        // (same contract the fused q8 FFN arm rides on).
14461        let cols = 96;
14462        let row: Vec<u8> = (0..cols)
14463            .map(|i| ((i * 37 % 251) - 125) as i8 as u8)
14464            .collect();
14465        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.13).sin()).collect();
14466        let act = split_act(&x);
14467        let fast = q8_row_dot(&row, &act);
14468        let scalar = q8_row_dot_scalar(&row, &act);
14469        assert!(
14470            (fast - scalar).abs() <= scalar.abs() * 1e-5 + 1e-5,
14471            "fast {fast} vs scalar {scalar}"
14472        );
14473    }
14474
14475    #[test]
14476    fn f32_matvec_matches_matvec_rows_bitexact() {
14477        let (rows, cols) = (300, 40);
14478        let w: Vec<f32> = (0..rows * cols).map(|i| (i as f32 * 0.017).sin()).collect();
14479        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.05).cos()).collect();
14480        let qt = QTensor::from_f32(w.clone(), rows, cols);
14481
14482        let mut a = vec![0.0f32; rows];
14483        matvec_rows(None, &w, &x, &mut a);
14484        let mut b = vec![0.0f32; rows];
14485        qt.matvec(&x, &mut b, None);
14486        assert_eq!(a, b);
14487    }
14488
14489    #[test]
14490    fn sdot_kernel_exact_on_grid() {
14491        // Activations already on the i8 grid (±1 with amax=1 → sx=1/127,
14492        // xq=±127 dequantizes EXACTLY) → the SDOT path must match the
14493        // exact f32 dot to float rounding. This isolates kernel
14494        // correctness from quantization noise.
14495        eprintln!("sdot_enabled = {}", sdot_enabled());
14496        let (rows, cols) = (9, 80); // odd rows → exercises 4-row + tail
14497        let w: Vec<u8> = (0..rows * cols)
14498            .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
14499            .collect();
14500        let scales: Vec<f32> = (0..rows).map(|o| 0.005 + o as f32 * 0.001).collect();
14501        let x: Vec<f32> = (0..cols)
14502            .map(|i| match i % 3 {
14503                0 => 1.0,
14504                1 => -1.0,
14505                _ => 0.0,
14506            })
14507            .collect();
14508        let mut a = vec![0.0f32; rows];
14509        qmatvec(
14510            &w,
14511            &[],
14512            &scales,
14513            &x,
14514            &[],
14515            TensorDtype::Q8Row,
14516            rows,
14517            cols,
14518            &mut a,
14519            None,
14520        );
14521        for o in 0..rows {
14522            let mut acc = 0.0f32;
14523            for j in 0..cols {
14524                acc += (w[o * cols + j] as i8) as f32 * x[j];
14525            }
14526            let expect = acc * scales[o];
14527            assert!(
14528                (a[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
14529                "row {o}: {} vs {expect}",
14530                a[o]
14531            );
14532        }
14533    }
14534
14535    #[test]
14536    fn q1_tbl_fast_path_matches_reference() {
14537        // gpr = 8 exercises the TBL pair-load fast loop, and the LAST
14538        // row's final 4-tile window trips the 4B-overread guard (the
14539        // payload ends exactly at the last tile) — both paths must
14540        // agree with the dequant reference.
14541        let (rows, cols) = (5, 256);
14542        let gpr = cols / GROUP_SIZE;
14543        let mut bytes = Vec::new();
14544        for t in 0..rows * gpr {
14545            let s = 0.007 + (t % 11) as f32 * 0.004;
14546            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
14547            for j in 0..4 {
14548                bytes.push(((t * 53 + j * 89 + 7) % 249) as u8);
14549            }
14550        }
14551        let x: Vec<f32> = (0..cols)
14552            .map(|i| if (i * 5) % 7 < 3 { 1.0 } else { -1.0 })
14553            .collect();
14554        let mut w = vec![0.0f32; rows * cols];
14555        cortiq_core::quant::dequant_q1(&bytes, &mut w);
14556        let mut got = vec![0.0f32; rows];
14557        q1_matvec(&bytes, &x, rows, cols, &mut got, None);
14558        for o in 0..rows {
14559            let expect: f32 = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
14560            assert!(
14561                (got[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
14562                "row {o}: {} vs {expect}",
14563                got[o]
14564            );
14565        }
14566        // Blocked 1×4 batch (b=5: one quad + remainder) must equal the
14567        // single-matvec path bit-for-bit.
14568        let b = 5usize;
14569        let mut xs_all = Vec::new();
14570        for bi in 0..b {
14571            xs_all.extend(x.iter().map(|v| if bi % 2 == 0 { *v } else { -*v }));
14572        }
14573        let mut mm = vec![0.0f32; b * rows];
14574        q1_matmat(&bytes, &xs_all, b, rows, cols, &mut mm, None);
14575        for bi in 0..b {
14576            let mut single = vec![0.0f32; rows];
14577            q1_matvec(
14578                &bytes,
14579                &xs_all[bi * cols..(bi + 1) * cols],
14580                rows,
14581                cols,
14582                &mut single,
14583                None,
14584            );
14585            assert_eq!(&mm[bi * rows..(bi + 1) * rows], &single[..], "stream {bi}");
14586        }
14587    }
14588
14589    #[test]
14590    fn q1_kernels_match_exact_reference() {
14591        // Synthetic q1 payload: 6-byte tiles [f16 scale][4B bits].
14592        let (rows, cols) = (7, 96);
14593        let gpr = cols / GROUP_SIZE;
14594        let mut bytes = Vec::new();
14595        for t in 0..rows * gpr {
14596            let s = 0.01 + (t % 13) as f32 * 0.003;
14597            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
14598            for j in 0..4 {
14599                bytes.push(((t * 31 + j * 97) % 251) as u8);
14600            }
14601        }
14602        // On-grid activations (±1, amax 1) → the SDOT path is exact.
14603        let x: Vec<f32> = (0..cols)
14604            .map(|i| if i % 3 == 0 { 1.0 } else { -1.0 })
14605            .collect();
14606        // Reference through the core dequant.
14607        let mut w = vec![0.0f32; rows * cols];
14608        cortiq_core::quant::dequant_q1(&bytes, &mut w);
14609        let mut expect = vec![0.0f32; rows];
14610        for o in 0..rows {
14611            expect[o] = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
14612        }
14613        let mut got = vec![0.0f32; rows];
14614        q1_matvec(&bytes, &x, rows, cols, &mut got, None);
14615        for o in 0..rows {
14616            assert!(
14617                (got[o] - expect[o]).abs() < 1e-3 * expect[o].abs().max(1e-3),
14618                "row {o}: {} vs {}",
14619                got[o],
14620                expect[o]
14621            );
14622        }
14623        // Pair and batch paths agree with the single path.
14624        let x2: Vec<f32> = x.iter().map(|v| -v).collect();
14625        let (mut a1, mut a2) = (vec![0.0f32; rows], vec![0.0f32; rows]);
14626        q1_matvec2(&bytes, &x, &x2, rows, cols, &mut a1, &mut a2, None);
14627        assert_eq!(a1, got);
14628        let mut xs = x.clone();
14629        xs.extend_from_slice(&x2);
14630        let mut mm = vec![0.0f32; 2 * rows];
14631        q1_matmat(&bytes, &xs, 2, rows, cols, &mut mm, None);
14632        assert_eq!(&mm[..rows], got.as_slice());
14633        assert_eq!(&mm[rows..], a2.as_slice());
14634    }
14635
14636    /// A Bonsai-shaped q1 payload (one f16 scale per 128-group, stored in
14637    /// all four tiles; pseudo-random signs) and its q1g repack.
14638    fn q1_q1g_pair(rows: usize, cols: usize, salt: usize) -> (Vec<u8>, Vec<u8>) {
14639        let tiles = rows * cols / GROUP_SIZE;
14640        let mut q1 = Vec::with_capacity(tiles * Q1_TILE);
14641        for t in 0..tiles {
14642            let s = 0.004 + ((t / 4 * 37 + salt) % 29) as f32 * 0.0011;
14643            q1.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
14644            let bits = ((t * 2_654_435_761 + salt * 40_503) as u32).rotate_left((t % 31) as u32);
14645            q1.extend_from_slice(&bits.to_le_bytes());
14646        }
14647        let rp = cortiq_core::quant::repack_q1_to_q1g(&q1, rows, cols, 0.0);
14648        assert_eq!(rp.differing, 0, "uniform scales repack exactly");
14649        (q1, rp.data)
14650    }
14651
14652    /// Activations off the int8 grid with two outlier channels, so the
14653    /// A8W8 path quantizes AND takes the exact outlier correction.
14654    fn q1g_test_x(cols: usize, salt: usize) -> Vec<f32> {
14655        let mut x: Vec<f32> = (0..cols)
14656            .map(|i| ((i * 13 + salt * 7) as f32 * 0.071).sin() * 0.9)
14657            .collect();
14658        x[3 % cols] = 41.0;
14659        x[(cols / 2 + 5) % cols] = -37.5;
14660        x
14661    }
14662
14663    /// Every CPU row kernel of q1g returns the q1 kernel's f32 BIT FOR BIT
14664    /// on a uniform-scale repack: A8W8 rows (SIMD of this machine, and the
14665    /// x86 AVX2/VNNI twins directly), the exact scalar rows (CMF_SDOT=0
14666    /// contract), the outlier term, the drivers (matvec, matvec2, matmat
14667    /// with the blocked 1×4 quad + remainder) — including gpr = 4 and the
14668    /// last row, where q1's pair load falls back to its scalar tail.
14669    #[test]
14670    fn q1g_cpu_kernels_equal_q1_bit_for_bit() {
14671        for &(rows, cols) in &[(3usize, 128usize), (5, 256), (9, 2048), (2, 6144)] {
14672            let gpr = cols / GROUP_SIZE;
14673            let (q1, q1g) = q1_q1g_pair(rows, cols, rows + cols);
14674            let (wa, wb) = (BinW::Q1(&q1), BinW::q1g(&q1g, rows, cols));
14675            let x = q1g_test_x(cols, rows);
14676            let act = split_act(&x);
14677            assert!(!act.outliers.is_empty(), "test wants the outlier path");
14678            let gsum = q1_group_sums(&act.xq, gpr);
14679            for r in 0..rows {
14680                let (a, b) = (wa.row_exact(r, gpr, &x), wb.row_exact(r, gpr, &x));
14681                assert_eq!(a.to_bits(), b.to_bits(), "{rows}x{cols} exact row {r}");
14682                let (a, b) = (
14683                    wa.dot_i8(r, gpr, &act.xq, &gsum),
14684                    wb.dot_i8(r, gpr, &act.xq, &gsum),
14685                );
14686                assert_eq!(a.to_bits(), b.to_bits(), "{rows}x{cols} i8 row {r}");
14687                for j in [0, 31, 32, 127, cols - 1] {
14688                    let (a, b) = (wa.outlier(r, gpr, j), wb.outlier(r, gpr, j));
14689                    assert_eq!(
14690                        (a.0.to_bits(), a.1.to_bits()),
14691                        (b.0.to_bits(), b.1.to_bits())
14692                    );
14693                }
14694                #[cfg(target_arch = "x86_64")]
14695                if avx2_enabled() {
14696                    let sc = cortiq_core::quant::q1g_scale_offset(rows, cols);
14697                    // SAFETY: AVX2 checked; payloads sized for (rows, cols).
14698                    unsafe {
14699                        let a = dot_q1_row_avx2(&q1, r, gpr, &act.xq, &gsum);
14700                        let b = dot_q1g_row_avx2(&q1g, sc, r, gpr, &act.xq, &gsum);
14701                        assert_eq!(a.to_bits(), b.to_bits(), "avx2 row {r}");
14702                    }
14703                    if std::is_x86_feature_detected!("avx512vnni")
14704                        && std::is_x86_feature_detected!("avx512bw")
14705                        && std::is_x86_feature_detected!("avx512vl")
14706                    {
14707                        // SAFETY: VNNI family detected just above.
14708                        unsafe {
14709                            let a = dot_q1_row_vnni(&q1, r, gpr, &act.xq, &gsum);
14710                            let b = dot_q1g_row_vnni(&q1g, sc, r, gpr, &act.xq, &gsum);
14711                            assert_eq!(a.to_bits(), b.to_bits(), "vnni row {r}");
14712                        }
14713                    }
14714                }
14715            }
14716            // Drivers: single, pair, batch (b = 5: one blocked quad + one).
14717            let (mut ga, mut gb) = (vec![0f32; rows], vec![0f32; rows]);
14718            bin_matvec(wa, &x, rows, cols, &mut ga, None);
14719            bin_matvec(wb, &x, rows, cols, &mut gb, None);
14720            assert_eq!(ga, gb, "{rows}x{cols} matvec");
14721            let x2 = q1g_test_x(cols, rows + 1);
14722            let (mut a1, mut a2, mut b1, mut b2) = (
14723                vec![0f32; rows],
14724                vec![0f32; rows],
14725                vec![0f32; rows],
14726                vec![0f32; rows],
14727            );
14728            bin_matvec2(wa, &x, &x2, rows, cols, &mut a1, &mut a2, None);
14729            bin_matvec2(wb, &x, &x2, rows, cols, &mut b1, &mut b2, None);
14730            assert_eq!((&a1, &a2), (&b1, &b2), "{rows}x{cols} matvec2");
14731            let b = 5usize;
14732            let xs: Vec<f32> = (0..b).flat_map(|k| q1g_test_x(cols, k * 3 + 2)).collect();
14733            let (mut ma, mut mb) = (vec![0f32; b * rows], vec![0f32; b * rows]);
14734            bin_matmat(wa, &xs, b, rows, cols, &mut ma, None);
14735            bin_matmat(wb, &xs, b, rows, cols, &mut mb, None);
14736            assert_eq!(ma, mb, "{rows}x{cols} matmat");
14737            // The blocked 1×4 kernels directly (the driver only takes them
14738            // under CMF_X86_BLOCKED / SDOT, a process-wide switch).
14739            let acts: Vec<(SplitAct, Vec<i32>)> = (0..4)
14740                .map(|k| {
14741                    let a = split_act(&xs[k * cols..(k + 1) * cols]);
14742                    let g = q1_group_sums(&a.xq, gpr);
14743                    (a, g)
14744                })
14745                .collect();
14746            let xq4: [&[i8]; 4] = std::array::from_fn(|k| acts[k].0.xq.as_slice());
14747            let gs4: [&[i32]; 4] = std::array::from_fn(|k| acts[k].1.as_slice());
14748            for r in 0..rows {
14749                #[cfg(target_arch = "x86_64")]
14750                if avx2_enabled() {
14751                    // SAFETY: AVX2 checked; payloads sized for (rows, cols).
14752                    let (a, b) = unsafe {
14753                        (
14754                            wa.dot_1x4_x86(r, gpr, xq4, gs4),
14755                            wb.dot_1x4_x86(r, gpr, xq4, gs4),
14756                        )
14757                    };
14758                    assert_eq!(a.map(f32::to_bits), b.map(f32::to_bits), "1x4 row {r}");
14759                }
14760                #[cfg(target_arch = "aarch64")]
14761                if sdot_enabled() {
14762                    // SAFETY: SDOT checked; payloads sized for (rows, cols).
14763                    let (a, b) = unsafe {
14764                        (
14765                            wa.dot_1x4_sdot(r, gpr, xq4, gs4),
14766                            wb.dot_1x4_sdot(r, gpr, xq4, gs4),
14767                        )
14768                    };
14769                    assert_eq!(a.map(f32::to_bits), b.map(f32::to_bits), "1x4 row {r}");
14770                }
14771                let _ = (r, &xq4, &gs4);
14772            }
14773        }
14774    }
14775
14776    /// CPU kernel A/B, q1 vs q1g matvec (A8W8 path of this machine, the
14777    /// exact path too), Bonsai-1.7B decode shapes, arms alternated every
14778    /// round, medians printed. `CMF_THREADS` sets the pool (default 8).
14779    /// `cargo test -p cortiq-engine --release --lib q1g_cpu_matvec_speed --
14780    /// --ignored --nocapture`
14781    #[test]
14782    #[ignore]
14783    fn q1g_cpu_matvec_speed() {
14784        use crate::pool::Pool;
14785        let threads = std::env::var("CMF_THREADS")
14786            .ok()
14787            .and_then(|v| v.parse().ok())
14788            .unwrap_or(8usize);
14789        let pool = Pool::new(threads);
14790        for &(name, rows, cols, reps) in &[
14791            ("gate/up_proj", 6144usize, 2048usize, 41usize),
14792            ("down_proj", 2048, 6144, 41),
14793            ("lm_head", 151_669, 2048, 9),
14794        ] {
14795            let (q1, q1g) = q1_q1g_pair(rows, cols, 3);
14796            let (wa, wb) = (BinW::Q1(&q1), BinW::q1g(&q1g, rows, cols));
14797            let x = q1g_test_x(cols, 1);
14798            let mut out = vec![0f32; rows];
14799            let mut time = |w: BinW| {
14800                let t0 = std::time::Instant::now();
14801                bin_matvec(w, &x, rows, cols, &mut out, Some(&pool));
14802                t0.elapsed().as_secs_f64()
14803            };
14804            let _ = (time(wa), time(wb));
14805            let (mut ta, mut tb) = (Vec::new(), Vec::new());
14806            for r in 0..reps {
14807                if r % 2 == 0 {
14808                    ta.push(time(wa));
14809                    tb.push(time(wb));
14810                } else {
14811                    tb.push(time(wb));
14812                    ta.push(time(wa));
14813                }
14814            }
14815            ta.sort_by(|a, b| a.partial_cmp(b).unwrap());
14816            tb.sort_by(|a, b| a.partial_cmp(b).unwrap());
14817            let (ma, mb) = (ta[reps / 2], tb[reps / 2]);
14818            eprintln!(
14819                "cpu x{threads} a8w8={} {name:>12} {rows}x{cols}: q1 {:8.1} us | q1g {:8.1} us | q1g/q1 time {:.3}",
14820                a8w8_enabled(),
14821                ma * 1e6,
14822                mb * 1e6,
14823                mb / ma
14824            );
14825        }
14826    }
14827
14828    /// The mapped QTensor paths a real Bonsai q1g file takes — matvec, the
14829    /// multi-projection job (all-q1g and q1/q1g mixed), the fused gate/up
14830    /// with SiLU, matmat, matvec2 and the embedding row — equal the q1
14831    /// tensor they were repacked from, bit for bit.
14832    #[test]
14833    fn q1g_mapped_tensor_paths_equal_q1() {
14834        use crate::pool::Pool;
14835        use cortiq_core::{CmfHeader, CmfModel, TensorSpec};
14836        let (rows, cols) = (300usize, 256usize);
14837        let (gq1, gq1g) = q1_q1g_pair(rows, cols, 1);
14838        let (uq1, uq1g) = q1_q1g_pair(rows, cols, 2);
14839        let header: CmfHeader = serde_json::from_value(serde_json::json!({
14840            "format": "cmf", "version": cortiq_core::CMF_VERSION,
14841            "arch": {"arch_name": "tiny-q1g", "hidden_size": cols,
14842                "intermediate_size": rows, "num_layers": 1, "num_attention_heads": 2,
14843                "num_kv_heads": 1, "head_dim": 128, "vocab_size": rows,
14844                "layer_types": ["FullAttention"], "rms_norm_eps": 1e-6,
14845                "max_position_embeddings": 8},
14846            "quant_type": "VBIT"
14847        }))
14848        .unwrap();
14849        let spec = |name: &str, dtype, data: &Vec<u8>| TensorSpec {
14850            name: name.into(),
14851            dtype,
14852            shape: vec![rows, cols],
14853            data: data.clone(),
14854        };
14855        let specs = [
14856            spec("g1", TensorDtype::Q1, &gq1),
14857            spec("u1", TensorDtype::Q1, &uq1),
14858            spec("gg", TensorDtype::Q1G, &gq1g),
14859            spec("ug", TensorDtype::Q1G, &uq1g),
14860        ];
14861        let dir = std::env::temp_dir().join(format!("cmf-q1g-mapped-{}", std::process::id()));
14862        std::fs::create_dir_all(&dir).unwrap();
14863        let path = dir.join("m.cmf");
14864        CmfModel::write(&path, &header, &specs, None, None).unwrap();
14865        let model = Arc::new(CmfModel::open(&path).unwrap());
14866        assert!(model.verify().is_empty(), "{:?}", model.verify());
14867        let t = |n: &str| QTensor::from_model(&model, n).unwrap();
14868        let (g1, u1, gg, ug) = (t("g1"), t("u1"), t("gg"), t("ug"));
14869        assert_eq!(gg.model_dtype(), Some(TensorDtype::Q1G));
14870        assert!(gg.is_q1() && gg.graph_weight().map(|w| w.2) == Some(10));
14871        let pool = Pool::new(3);
14872        let x = q1g_test_x(cols, 4);
14873        let run = |a: &QTensor| {
14874            let mut o = vec![0f32; rows];
14875            a.matvec(&x, &mut o, Some(&pool));
14876            o
14877        };
14878        assert_eq!(run(&g1), run(&gg), "matvec");
14879        let many = |a: &QTensor, b: &QTensor| {
14880            let (mut oa, mut ob) = (vec![0f32; rows], vec![0f32; rows]);
14881            QTensor::matvec_many([a, b], &x, [&mut oa, &mut ob], Some(&pool));
14882            (oa, ob)
14883        };
14884        let want = many(&g1, &u1);
14885        assert_eq!(want, many(&gg, &ug), "matvec_many q1g");
14886        assert_eq!(want, many(&g1, &ug), "matvec_many mixed");
14887        assert_eq!(want, (run(&g1), run(&u1)), "matvec_many == singles");
14888        let fused = |g: &QTensor, u: &QTensor| {
14889            let mut o = vec![0f32; rows];
14890            let ok = QTensor::matvec_silu_mul(g, u, &x, &mut o, Some(&pool));
14891            (ok, o)
14892        };
14893        let (ok_a, fa) = fused(&g1, &u1);
14894        let (ok_b, fb) = fused(&gg, &ug);
14895        assert_eq!(ok_a, ok_b, "q1g takes the fused arm whenever q1 does");
14896        assert_eq!(fa, fb, "fused gate/up");
14897        let b = 6usize;
14898        let xs: Vec<f32> = (0..b).flat_map(|k| q1g_test_x(cols, k + 9)).collect();
14899        let mm = |a: &QTensor| {
14900            let mut o = vec![0f32; b * rows];
14901            a.matmat(&xs, b, &mut o, Some(&pool));
14902            o
14903        };
14904        assert_eq!(mm(&g1), mm(&gg), "matmat");
14905        let x2 = q1g_test_x(cols, 77);
14906        let pair = |a: &QTensor| {
14907            let (mut o1, mut o2) = (vec![0f32; rows], vec![0f32; rows]);
14908            a.matvec2(&x, &x2, &mut o1, &mut o2, Some(&pool));
14909            (o1, o2)
14910        };
14911        assert_eq!(pair(&g1), pair(&gg), "matvec2");
14912        for r in [0, 1, rows / 2, rows - 1] {
14913            let (mut ra, mut rb) = (vec![0f32; cols], vec![0f32; cols]);
14914            g1.row_f32(r, &mut ra);
14915            gg.row_f32(r, &mut rb);
14916            assert_eq!(ra, rb, "embedding row {r}");
14917        }
14918        drop((g1, u1, gg, ug, model));
14919        let _ = std::fs::remove_dir_all(&dir);
14920    }
14921
14922    /// A `blr` tensor read from a real CMF file: it maps to the factored
14923    /// arm (not a dense copy), declines every device-graph accessor, and
14924    /// every QTensor entry point a layer may take — matvec, matvec2,
14925    /// matmat (table loop and GEMM sizes), the multi-tensor job mixed with
14926    /// a q8_row sibling, the fused gate/up refusal, rows and columns — equals
14927    /// the dense reference built by the core dequant.
14928    #[test]
14929    fn blr_mapped_tensor_paths_equal_dense() {
14930        use crate::pool::Pool;
14931        use cortiq_core::{CmfHeader, CmfModel, TensorSpec};
14932        let (rows, cols, r) = (96usize, 128usize, 16usize);
14933        let mut rng = crate::blr::tests::Rng(41);
14934        let gate = crate::blr::tests::make(&mut rng, rows, cols, r, 2);
14935        let up = crate::blr::tests::make(&mut rng, rows, cols, r, 1);
14936        // q8_row sibling: [int8 rows·cols][f16 row scale · rows].
14937        let mut q8: Vec<u8> = (0..rows * cols)
14938            .map(|i| (((i * 89) % 253) as i32 - 126) as i8 as u8)
14939            .collect();
14940        for o in 0..rows {
14941            q8.extend_from_slice(&cortiq_core::quant::f32_to_f16(0.002 + o as f32 * 1e-4).to_le_bytes());
14942        }
14943        let header: CmfHeader = serde_json::from_value(serde_json::json!({
14944            "format": "cmf", "version": cortiq_core::CMF_VERSION,
14945            "arch": {"arch_name": "tiny-blr", "hidden_size": cols,
14946                "intermediate_size": rows, "num_layers": 1, "num_attention_heads": 2,
14947                "num_kv_heads": 1, "head_dim": 64, "vocab_size": rows,
14948                "layer_types": ["FullAttention"], "rms_norm_eps": 1e-6,
14949                "max_position_embeddings": 8},
14950            "quant_type": "Q8_2F"
14951        }))
14952        .unwrap();
14953        let specs = [
14954            TensorSpec {
14955                name: "g".into(),
14956                dtype: TensorDtype::BinaryLowRank,
14957                shape: vec![rows, cols],
14958                data: gate.clone(),
14959            },
14960            TensorSpec {
14961                name: "u".into(),
14962                dtype: TensorDtype::BinaryLowRank,
14963                shape: vec![rows, cols],
14964                data: up.clone(),
14965            },
14966            TensorSpec {
14967                name: "q".into(),
14968                dtype: TensorDtype::Q8Row,
14969                shape: vec![rows, cols],
14970                data: q8,
14971            },
14972        ];
14973        let dir = std::env::temp_dir().join(format!("cmf-blr-mapped-{}", std::process::id()));
14974        std::fs::create_dir_all(&dir).unwrap();
14975        let path = dir.join("m.cmf");
14976        CmfModel::write(&path, &header, &specs, None, None).unwrap();
14977        let model = Arc::new(CmfModel::open(&path).unwrap());
14978        assert!(model.verify().is_empty(), "{:?}", model.verify());
14979        let t = |n: &str| QTensor::from_model(&model, n).unwrap();
14980        let (g, u, q) = (t("g"), t("u"), t("q"));
14981        assert!(g.is_blr() && u.is_blr() && !q.is_blr());
14982        assert!(g.graph_weight().is_none() && g.graph_weight_descriptor().is_none());
14983        assert!(g.as_f32().is_none() && g.model_idx().is_none());
14984        assert!(g.mapped_device_gemm().is_none() && g.q1_parts().is_none());
14985        assert!(!g.sparse_col_ok());
14986        assert_eq!((g.rows(), g.cols()), (rows, cols));
14987        let deq = |bytes: &[u8]| {
14988            let mut w = vec![0f32; rows * cols];
14989            cortiq_core::blr::dequant_blr(bytes, rows, cols, &mut w).unwrap();
14990            w
14991        };
14992        let (wg, wu) = (deq(&gate), deq(&up));
14993        let mv = |w: &[f32], x: &[f32]| -> Vec<f32> {
14994            (0..rows)
14995                .map(|i| (0..cols).map(|j| w[i * cols + j] * x[j]).sum())
14996                .collect()
14997        };
14998        let close = |a: &[f32], b: &[f32], what: &str| {
14999            let s = b.iter().map(|v| v * v).sum::<f32>().sqrt().max(1e-30);
15000            let e = a.iter().zip(b).map(|(x, y)| (x - y) * (x - y)).sum::<f32>().sqrt();
15001            assert!(e / s < 1e-5, "{what}: rel {}", e / s);
15002        };
15003        let pool = Pool::new(3);
15004        let x: Vec<f32> = (0..cols).map(|_| rng.f()).collect();
15005        let x2: Vec<f32> = (0..cols).map(|_| rng.f()).collect();
15006        let mut o = vec![0f32; rows];
15007        g.matvec(&x, &mut o, Some(&pool));
15008        close(&o, &mv(&wg, &x), "matvec");
15009        let (mut o1, mut o2) = (vec![0f32; rows], vec![0f32; rows]);
15010        g.matvec2(&x, &x2, &mut o1, &mut o2, Some(&pool));
15011        close(&o1, &mv(&wg, &x), "matvec2 a");
15012        close(&o2, &mv(&wg, &x2), "matvec2 b");
15013        for b in [2usize, 5] {
15014            let xs: Vec<f32> = (0..b * cols).map(|_| rng.f()).collect();
15015            let mut om = vec![0f32; b * rows];
15016            g.matmat(&xs, b, &mut om, Some(&pool));
15017            for t in 0..b {
15018                close(
15019                    &om[t * rows..(t + 1) * rows],
15020                    &mv(&wg, &xs[t * cols..(t + 1) * cols]),
15021                    "matmat",
15022                );
15023            }
15024        }
15025        // Multi-tensor job mixed with a q8_row sibling: per-tensor fallback,
15026        // the q8_row output identical to its own matvec.
15027        let (mut a, mut bq, mut c) = (vec![0f32; rows], vec![0f32; rows], vec![0f32; rows]);
15028        QTensor::matvec_many([&g, &q, &u], &x, [&mut a, &mut bq, &mut c], Some(&pool));
15029        close(&a, &mv(&wg, &x), "matvec_many blr gate");
15030        close(&c, &mv(&wu, &x), "matvec_many blr up");
15031        let mut qo = vec![0f32; rows];
15032        q.matvec(&x, &mut qo, Some(&pool));
15033        assert_eq!(bq, qo, "matvec_many q8_row sibling");
15034        let mut f = vec![0f32; rows];
15035        assert!(!QTensor::matvec_silu_mul(&g, &u, &x, &mut f, Some(&pool)));
15036        let mut row = vec![0f32; cols];
15037        g.row_f32(rows / 3, &mut row);
15038        close(&row, &wg[(rows / 3) * cols..(rows / 3 + 1) * cols], "row_f32");
15039        let mut scratch = vec![0f32; cols];
15040        let d = g.row_dot(7, &x, &mut scratch);
15041        let want: f32 = (0..cols).map(|j| wg[7 * cols + j] * x[j]).sum();
15042        assert!((d - want).abs() < 1e-4 * (1.0 + want.abs()), "row_dot");
15043        let mut col = vec![0f32; rows];
15044        g.add_col_scaled(cols / 2, 1.5, &mut col);
15045        let want: Vec<f32> = (0..rows).map(|i| 1.5 * wg[i * cols + cols / 2]).collect();
15046        close(&col, &want, "add_col_scaled");
15047        drop((g, u, q, model));
15048        let _ = std::fs::remove_dir_all(&dir);
15049    }
15050
15051    #[test]
15052    fn repack_is_bit_identical() {
15053        // The interleaved-repack kernel must produce EXACTLY the same
15054        // bits as the mmap-layout kernel: integer accumulation is order-
15055        // exact, the f32 epilogue is identical. Odd rows exercise the
15056        // tail; direct range calls exercise unaligned pool splits.
15057        let (rows, cols) = (267, 96); // 66 groups + 3 tail rows, cols % 16 == 0
15058        let w: Vec<u8> = (0..rows * cols)
15059            .map(|i| (((i * 89) % 253) as i32 - 126) as i8 as u8)
15060            .collect();
15061        let scales: Vec<f32> = (0..rows).map(|o| 0.003 + o as f32 * 0.0007).collect();
15062        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.37).sin() * 2.0).collect();
15063        let rep = q8_repack_layout(&w, rows, cols);
15064        // Group interleave round-trips.
15065        for g in 0..rows / 4 {
15066            for c in 0..cols / 16 {
15067                for lane in 0..4 {
15068                    assert_eq!(
15069                        &rep[g * 4 * cols + c * 64 + lane * 16
15070                            ..g * 4 * cols + c * 64 + lane * 16 + 16],
15071                        &w[(g * 4 + lane) * cols + c * 16..(g * 4 + lane) * cols + c * 16 + 16],
15072                    );
15073                }
15074            }
15075        }
15076        let mut a = vec![0.0f32; rows];
15077        qmatvec(
15078            &w,
15079            &[],
15080            &scales,
15081            &x,
15082            &[],
15083            TensorDtype::Q8Row,
15084            rows,
15085            cols,
15086            &mut a,
15087            None,
15088        );
15089        let mut b = vec![0.0f32; rows];
15090        qmatvec(
15091            &w,
15092            &rep,
15093            &scales,
15094            &x,
15095            &[],
15096            TensorDtype::Q8Row,
15097            rows,
15098            cols,
15099            &mut b,
15100            None,
15101        );
15102        assert_eq!(a, b, "full-range repack output diverged");
15103
15104        #[cfg(target_arch = "aarch64")]
15105        if sdot_enabled() {
15106            // Unaligned range split (pool workers get arbitrary bounds).
15107            let act = split_act(&x);
15108            let mut c1 = vec![0.0f32; rows];
15109            let mut c2 = vec![0.0f32; rows];
15110            q8_range_sdot(
15111                &w,
15112                &[],
15113                &scales,
15114                &act,
15115                cols,
15116                SendMut(c1.as_mut_ptr()),
15117                3,
15118                rows - 2,
15119            );
15120            q8_range_sdot(
15121                &w,
15122                &rep,
15123                &scales,
15124                &act,
15125                cols,
15126                SendMut(c2.as_mut_ptr()),
15127                3,
15128                rows - 2,
15129            );
15130            assert_eq!(c1, c2, "unaligned-range repack output diverged");
15131        }
15132    }
15133
15134    #[test]
15135    fn sdot_a8w8_noise_is_bounded() {
15136        // Off-grid activations: A8 quantization noise must stay small in
15137        // relative L2 over the whole output (realistic accuracy contract;
15138        // vmfcore measured argmax-identical decode on real models).
15139        let (rows, cols) = (16, 512);
15140        let w: Vec<u8> = (0..rows * cols)
15141            .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
15142            .collect();
15143        let scales = vec![0.01f32; rows];
15144        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.21).sin()).collect();
15145        let mut a = vec![0.0f32; rows];
15146        qmatvec(
15147            &w,
15148            &[],
15149            &scales,
15150            &x,
15151            &[],
15152            TensorDtype::Q8Row,
15153            rows,
15154            cols,
15155            &mut a,
15156            None,
15157        );
15158        let (mut num, mut den) = (0f64, 0f64);
15159        for o in 0..rows {
15160            let mut acc = 0.0f32;
15161            for j in 0..cols {
15162                acc += (w[o * cols + j] as i8) as f32 * x[j];
15163            }
15164            let expect = acc * scales[o];
15165            num += ((a[o] - expect) as f64).powi(2);
15166            den += (expect as f64).powi(2);
15167        }
15168        let rel = (num / den.max(1e-12)).sqrt();
15169        assert!(rel < 0.05, "A8W8 relative L2 error too high: {rel}");
15170    }
15171
15172    #[test]
15173    fn i8_dot_neon_matches_scalar() {
15174        let n = 100;
15175        let w: Vec<u8> = (0..n).map(|i| ((i * 37 + 11) % 251) as u8).collect();
15176        let x: Vec<f32> = (0..n).map(|i| (i as f32 * 0.13).sin()).collect();
15177        let mut scalar = 0.0f32;
15178        for j in 0..n {
15179            scalar += (w[j] as i8) as f32 * x[j];
15180        }
15181        let fast = dot_i8_f32(&w, &x);
15182        assert!((scalar - fast).abs() < 1e-3 * scalar.abs().max(1.0));
15183    }
15184
15185    /// Fused vbit matvec must match full dequant_vbit + dense matvec.
15186    #[test]
15187    fn vbitmatvec_matches_full_dequant() {
15188        let (rows, cols) = (6, 64);
15189        let ng = cols / GROUP_SIZE;
15190        // Hand-craft: bits per row, f16 scales, packed rows.
15191        let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
15192        let mut bytes = bits.clone();
15193        for g in 0..rows * ng {
15194            let s = 0.02 + 0.001 * g as f32;
15195            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
15196        }
15197        for r in 0..rows {
15198            let b = bits[r] as usize;
15199            let (mut acc, mut nb) = (0u64, 0usize);
15200            let mut rowbytes = Vec::new();
15201            for i in 0..cols {
15202                let v = ((i * 7 + r * 13) % (1 << b)) as u64;
15203                acc = (acc << b) | v;
15204                nb += b;
15205                while nb >= 8 {
15206                    nb -= 8;
15207                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
15208                }
15209            }
15210            if nb > 0 {
15211                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
15212            }
15213            bytes.extend_from_slice(&rowbytes);
15214        }
15215        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
15216
15217        let mut reference = vec![0f32; rows * cols];
15218        cortiq_core::quant::dequant_vbit(&bytes, rows, cols, &mut reference).unwrap();
15219        let mut expect = vec![0f32; rows];
15220        for r in 0..rows {
15221            expect[r] = reference[r * cols..(r + 1) * cols]
15222                .iter()
15223                .zip(&x)
15224                .map(|(w, xv)| w * xv)
15225                .sum();
15226        }
15227        let mut got = vec![0f32; rows];
15228        let offsets = vbit_row_offsets(&bytes, rows, cols);
15229        vbitmatvec(&bytes, &offsets, &x, rows, cols, &mut got, None);
15230        // SDOT path quantizes activations to i8 (A8W8): bounded noise,
15231        // same contract as q8 (exact path is pinned by CMF_SDOT=0 in
15232        // the golden-parity gate).
15233        let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
15234        let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
15235        for r in 0..rows {
15236            assert!(
15237                (got[r] - expect[r]).abs() < tol * scale,
15238                "row {r}: {} vs {}",
15239                got[r],
15240                expect[r]
15241            );
15242        }
15243    }
15244
15245    /// Fused q4 matvec must match the reference full-dequant + dense
15246    /// matvec bit-for-bit in structure (same f32 math, group order).
15247    /// vbit matmat: the blocked 1×4 leg must match the per-row path
15248    /// (paired env toggle; larger shape so both code paths engage).
15249    #[test]
15250    #[cfg(target_arch = "x86_64")]
15251    fn vbit_matmat_blocked_matches_per_row() {
15252        let (rows, cols, b) = (64usize, 128usize, 9usize);
15253        let ng = cols / GROUP_SIZE;
15254        let bits: Vec<u8> = (0..rows).map(|r| [3u8, 4, 5, 6][r % 4]).collect();
15255        let mut bytes = bits.clone();
15256        for g in 0..rows * ng {
15257            let sc = 0.02 + 0.0005 * g as f32;
15258            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
15259        }
15260        for r in 0..rows {
15261            let bw = bits[r] as usize;
15262            let (mut acc, mut nb) = (0u64, 0usize);
15263            let mut rowbytes = Vec::new();
15264            for i in 0..cols {
15265                let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
15266                acc = (acc << bw) | v;
15267                nb += bw;
15268                while nb >= 8 {
15269                    nb -= 8;
15270                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
15271                }
15272            }
15273            if nb > 0 {
15274                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
15275            }
15276            bytes.extend_from_slice(&rowbytes);
15277        }
15278        let x: Vec<f32> = (0..b * cols)
15279            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
15280            .collect();
15281        let offsets = vbit_row_offsets(&bytes, rows, cols);
15282        let mut y_a = vec![0f32; b * rows];
15283        let mut y_b = vec![0f32; b * rows];
15284        unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
15285        vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_a, None);
15286        unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
15287        vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_b, None);
15288        unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
15289        let max_d = y_a
15290            .iter()
15291            .zip(&y_b)
15292            .map(|(p, q)| (p - q).abs())
15293            .fold(0.0f32, f32::max);
15294        assert!(max_d < 1e-4, "vbit blocked ≠ per-row: max|Δ| = {max_d}");
15295    }
15296
15297    /// q4t blocked 1×4 (SDOT on ARM, AVX2 on x86) must equal the
15298    /// per-row path exactly: same nibble unpack, same group order,
15299    /// same f32 accumulation — batch == matvec bit-for-bit. b=9 covers
15300    /// two full 1×4 blocks plus a remainder through the single-row
15301    /// kernel. (Both paths produce identical output, so the shared
15302    /// CMF_X86_BLOCKED env var racing with other tests cannot flip
15303    /// the verdict — worst case both sides take the same path.)
15304    #[test]
15305    fn q4t_matmat_blocked_matches_per_row() {
15306        let (rows, cols, b) = (16usize, 64usize, 9usize);
15307        let gpr = cols / GROUP_SIZE;
15308        let mut bytes = vec![0u8; rows * gpr * Q4_TILE];
15309        for r in 0..rows {
15310            for g in 0..gpr {
15311                let t = (r * gpr + g) * Q4_TILE;
15312                let sc = 0.02 + 0.001 * (r * gpr + g) as f32;
15313                bytes[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
15314                for k in 0..16 {
15315                    bytes[t + 2 + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
15316                }
15317            }
15318        }
15319        let x: Vec<f32> = (0..b * cols)
15320            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
15321            .collect();
15322        let mut y_blk = vec![0f32; b * rows];
15323        let mut y_row = vec![0f32; b * rows];
15324        unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
15325        q4t_matmat(&bytes, &x, b, rows, cols, &mut y_blk, None);
15326        unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
15327        q4t_matmat(&bytes, &x, b, rows, cols, &mut y_row, None);
15328        unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
15329        assert_eq!(y_blk, y_row, "q4t blocked 1x4 ≠ per-row");
15330    }
15331
15332    /// The wide-batch Accelerate arm of q4t_matmat vs a brute-force
15333    /// f32 dequant matmul: both are f32 GEMMs, so only reduction
15334    /// order differs — tight tolerance.
15335    /// A synthetic q4tp payload: random nibbles plus a per-row ladder whose
15336    /// span varies row to row, so the codes actually exercise the full 0..31
15337    /// range rather than clustering on one rung.
15338    fn synth_q4tp(rows: usize, cols: usize) -> Vec<u8> {
15339        use cortiq_core::quant::{f32_to_f16, q4tp_code_stride, q4tp_put_code};
15340        let gpr = cols / GROUP_SIZE;
15341        let stride = q4tp_code_stride(gpr);
15342        let (params_off, codes_off, _) = q4tp_sections(rows, cols);
15343        let mut b = vec![0u8; codes_off + rows * stride];
15344        for r in 0..rows {
15345            for g in 0..gpr {
15346                let t = (r * gpr + g) * Q4TP_NIB;
15347                for k in 0..16 {
15348                    b[t + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
15349                }
15350            }
15351            let lo = -6.0 - 0.03 * (r % 17) as f32;
15352            let step = 0.01 + 0.004 * (r % 11) as f32;
15353            let p = params_off + r * 4;
15354            b[p..p + 2].copy_from_slice(&f32_to_f16(lo).to_le_bytes());
15355            b[p + 2..p + 4].copy_from_slice(&f32_to_f16(step).to_le_bytes());
15356            let crow = &mut b[codes_off + r * stride..codes_off + (r + 1) * stride];
15357            for g in 0..gpr {
15358                q4tp_put_code(crow, g, (r * 5 + g * 3) % 32);
15359            }
15360        }
15361        b
15362    }
15363
15364    /// The same weights re-expressed as q4_tiled, so the proven kernel can
15365    /// be the reference: each tile stores the ladder scale its code selects.
15366    /// Only the f16 rounding of that scale separates the two payloads.
15367    fn q4tp_as_q4t(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
15368        let gpr = cols / GROUP_SIZE;
15369        let v = Q4tpView::new(bytes, rows, cols);
15370        let mut out = vec![0u8; rows * gpr * Q4_TILE];
15371        let mut sc = vec![0f32; gpr];
15372        for r in 0..rows {
15373            v.scales_into(r, gpr, &mut sc);
15374            for g in 0..gpr {
15375                let t = (r * gpr + g) * Q4_TILE;
15376                let s = sc[g];
15377                out[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
15378                let src = (r * gpr + g) * Q4TP_NIB;
15379                out[t + 2..t + Q4_TILE].copy_from_slice(&v.nib[src..src + Q4TP_NIB]);
15380            }
15381        }
15382        out
15383    }
15384
15385    /// The exact (`CMF_SDOT=0`) path must reproduce `dequant_q4tp` to f32
15386    /// rounding — that scalar routine is the format's definition, and the
15387    /// kernels re-derive the scale from the ladder independently. Call the
15388    /// row kernel directly: `matmat` picks the int8 arm when a8w8 is on,
15389    /// so routing through it would test the other path by accident.
15390    #[test]
15391    fn q4tp_exact_path_matches_dequant_reference() {
15392        let (rows, cols) = (256usize, 512usize);
15393        let gpr = cols / GROUP_SIZE;
15394        let bytes = synth_q4tp(rows, cols);
15395        let mut w = vec![0f32; rows * cols];
15396        cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
15397
15398        let x: Vec<f32> = (0..cols)
15399            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
15400            .collect();
15401        let v = Q4tpView::new(&bytes, rows, cols);
15402        let mut sc = vec![0f32; gpr];
15403        for r in 0..rows {
15404            v.scales_into(r, gpr, &mut sc);
15405            let got = q4tp_row_exact(v.nib, r, gpr, &x, &sc);
15406            let want: f32 = (0..cols).map(|c| w[r * cols + c] * x[c]).sum();
15407            // These dot products cancel down to ~1e-3 from terms of ~5e-2, so
15408            // the meaningful yardstick is the summed magnitude, not the result:
15409            // against the result any reordering of a 512-term f32 sum "fails".
15410            let mag: f32 = (0..cols).map(|c| (w[r * cols + c] * x[c]).abs()).sum();
15411            assert!(
15412                (got - want).abs() <= 1e-5 * mag,
15413                "row {r}: kernel {got} vs dequant {want}"
15414            );
15415        }
15416    }
15417
15418    /// The int8 (a8w8) path can't be checked against an f32 reference — the
15419    /// activation quantization dominates. Check it against the q4t kernel it
15420    /// was ported from instead, on payloads holding the same weights: that
15421    /// isolates exactly what the port could break (16 B stride, ladder
15422    /// lookup, nibble unpack) from what it deliberately shares.
15423    #[test]
15424    fn q4tp_matvec_matches_the_q4t_kernel_it_was_ported_from() {
15425        let (rows, cols) = (256usize, 512usize);
15426        let bytes = synth_q4tp(rows, cols);
15427        let twin = q4tp_as_q4t(&bytes, rows, cols);
15428        let x: Vec<f32> = (0..cols)
15429            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
15430            .collect();
15431
15432        let mut got = vec![0f32; rows];
15433        q4tp_matvec(&bytes, &x, rows, cols, &mut got, None);
15434        let mut want = vec![0f32; rows];
15435        q4t_matvec(&twin, &x, rows, cols, &mut want, None);
15436
15437        // Scale is f16 in the twin and f32 here, so allow that rounding on
15438        // top of the summed magnitude (same cancellation argument as above).
15439        let mut w = vec![0f32; rows * cols];
15440        cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
15441        for r in 0..rows {
15442            let mag: f32 = (0..cols).map(|c| (w[r * cols + c] * x[c]).abs()).sum();
15443            assert!(
15444                (got[r] - want[r]).abs() <= 1e-3 * mag,
15445                "row {r}: q4tp {} vs q4t {}",
15446                got[r],
15447                want[r]
15448            );
15449        }
15450    }
15451
15452    /// `matmat` carries three arms (Accelerate, blocked int8 1x4, scalar).
15453    /// Batch 5 crosses the blocked kernel's stride, so this exercises the
15454    /// 1x4 path AND its scalar tail in one run — the blocked kernel is new
15455    /// code and its four accumulators are exactly what tends to go wrong.
15456    #[test]
15457    fn q4tp_matmat_matches_the_q4t_kernel_it_was_ported_from() {
15458        let (rows, cols, b) = (256usize, 512usize, 5usize);
15459        let bytes = synth_q4tp(rows, cols);
15460        let twin = q4tp_as_q4t(&bytes, rows, cols);
15461        let xs: Vec<f32> = (0..b * cols)
15462            .map(|i| ((i * 29 + 11) % 89) as f32 / 89.0 - 0.5)
15463            .collect();
15464
15465        let mut got = vec![0f32; b * rows];
15466        q4tp_matmat(&bytes, &xs, b, rows, cols, &mut got, None);
15467        let mut want = vec![0f32; b * rows];
15468        q4t_matmat(&twin, &xs, b, rows, cols, &mut want, None);
15469
15470        let mut w = vec![0f32; rows * cols];
15471        cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
15472        for t in 0..b {
15473            for r in 0..rows {
15474                let mag: f32 = (0..cols)
15475                    .map(|c| (w[r * cols + c] * xs[t * cols + c]).abs())
15476                    .sum();
15477                let (g, wa) = (got[t * rows + r], want[t * rows + r]);
15478                assert!(
15479                    (g - wa).abs() <= 1e-3 * mag,
15480                    "batch {t} row {r}: q4tp {g} vs q4t {wa}"
15481                );
15482            }
15483        }
15484    }
15485
15486    #[test]
15487    fn q4tp_matvec2_matches_the_single_stream_kernel() {
15488        let (rows, cols) = (128usize, 256usize);
15489        let gpr = cols / GROUP_SIZE;
15490        let bytes = synth_q4tp(rows, cols);
15491        let xs: Vec<f32> = (0..2 * cols)
15492            .map(|i| ((i * 29 + 11) % 89) as f32 / 89.0 - 0.5)
15493            .collect();
15494
15495        let (mut o1, mut o2) = (vec![0f32; rows], vec![0f32; rows]);
15496        q4tp_matvec2(
15497            &bytes,
15498            &xs[..cols],
15499            &xs[cols..],
15500            rows,
15501            cols,
15502            &mut o1,
15503            &mut o2,
15504            None,
15505        );
15506
15507        // matvec2 takes the exact path for both streams, so the single-row
15508        // kernel is an exact reference — no tolerance for path differences.
15509        let v = Q4tpView::new(&bytes, rows, cols);
15510        let mut sc = vec![0f32; gpr];
15511        for r in 0..rows {
15512            v.scales_into(r, gpr, &mut sc);
15513            assert_eq!(o1[r], q4tp_row_exact(v.nib, r, gpr, &xs[..cols], &sc));
15514            assert_eq!(o2[r], q4tp_row_exact(v.nib, r, gpr, &xs[cols..], &sc));
15515        }
15516    }
15517
15518    /// q4tp must not COST speed — it exists to save bytes, and a format that
15519    /// trades 7% of a file for a slower model is a bad trade. This guard is
15520    /// here because correctness tests happily passed while `q4tp_matmat` was
15521    /// missing its int8 and Accelerate arms and the model ran 5x slower.
15522    /// Measured on M-series: 0.97-1.04x, i.e. parity (16 B tiles are better
15523    /// aligned than q4t's 18 B, which pays for the scale indirection).
15524    #[test]
15525    fn q4tp_matvec_keeps_pace_with_q4t() {
15526        let (rows, cols) = (4096usize, 3072usize);
15527        let bytes = synth_q4tp(rows, cols);
15528        let twin = q4tp_as_q4t(&bytes, rows, cols);
15529        let x: Vec<f32> = (0..cols).map(|i| (i % 97) as f32 / 97.0 - 0.5).collect();
15530        let mut o = vec![0f32; rows];
15531        let n = 12;
15532        let mut best = (f64::MAX, f64::MAX);
15533        // Interleaved A/B, minimum statistic: this machine throttles, and a
15534        // mean over a thermal ramp reliably indicts whichever ran second.
15535        for _ in 0..3 {
15536            let t0 = std::time::Instant::now();
15537            for _ in 0..n {
15538                q4t_matvec(&twin, &x, rows, cols, &mut o, None);
15539            }
15540            best.0 = best.0.min(t0.elapsed().as_secs_f64());
15541            let t0 = std::time::Instant::now();
15542            for _ in 0..n {
15543                q4tp_matvec(&bytes, &x, rows, cols, &mut o, None);
15544            }
15545            best.1 = best.1.min(t0.elapsed().as_secs_f64());
15546        }
15547        let ratio = best.1 / best.0;
15548        println!(
15549            "q4t {:.3} ms | q4tp {:.3} ms | {ratio:.2}x",
15550            best.0 * 1e3 / n as f64,
15551            best.1 * 1e3 / n as f64
15552        );
15553        assert!(ratio < 2.0, "q4tp matvec {ratio:.2}x slower than q4t");
15554    }
15555
15556    #[cfg(target_os = "macos")]
15557    #[test]
15558    fn q4t_matmat_accel_matches_dequant_reference() {
15559        if !accel_gemm_enabled() {
15560            return; // CMF_ACCEL=0
15561        }
15562        let (rows, cols, b) = (512usize, 1024usize, 8usize); // ≥500K → accel arm
15563        let gpr = cols / GROUP_SIZE;
15564        let mut bytes = vec![0u8; rows * gpr * Q4_TILE];
15565        for r in 0..rows {
15566            for g in 0..gpr {
15567                let t = (r * gpr + g) * Q4_TILE;
15568                let sc = 0.02 + 0.0005 * ((r * gpr + g) % 64) as f32;
15569                bytes[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
15570                for k in 0..16 {
15571                    bytes[t + 2 + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
15572                }
15573            }
15574        }
15575        let x: Vec<f32> = (0..b * cols)
15576            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
15577            .collect();
15578        let mut got = vec![0f32; b * rows];
15579        q4t_matmat(&bytes, &x, b, rows, cols, &mut got, None);
15580        // Brute-force reference off the same tiles.
15581        let mut w = vec![0f32; rows * cols];
15582        for r in 0..rows {
15583            for g in 0..gpr {
15584                let t = (r * gpr + g) * Q4_TILE;
15585                let s = f16_to_f32(u16::from_le_bytes([bytes[t], bytes[t + 1]]));
15586                for (k, &bb) in bytes[t + 2..t + Q4_TILE].iter().enumerate() {
15587                    w[r * cols + g * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s;
15588                    w[r * cols + g * GROUP_SIZE + k * 2 + 1] =
15589                        (((bb >> 4) & 0x0F) as f32 - 8.0) * s;
15590                }
15591            }
15592        }
15593        for bi in 0..b {
15594            for r in 0..rows {
15595                let want: f32 = (0..cols).map(|j| x[bi * cols + j] * w[r * cols + j]).sum();
15596                let d = (got[bi * rows + r] - want).abs();
15597                assert!(
15598                    d <= want.abs().max(1.0) * 1e-4,
15599                    "accel q4t GEMM diverged at ({bi},{r}): {} vs {want}",
15600                    got[bi * rows + r]
15601                );
15602            }
15603        }
15604    }
15605
15606    #[test]
15607    fn q4matvec_matches_full_dequant() {
15608        let (rows, cols) = (8, 64);
15609        let groups = rows * cols / GROUP_SIZE;
15610        // Hand-craft a q4_block blob: nibbles then f16 scales.
15611        let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
15612        for i in 0..groups * 16 {
15613            bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
15614        }
15615        for g in 0..groups {
15616            let s = 0.01 + 0.003 * g as f32;
15617            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
15618        }
15619        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
15620
15621        let mut reference = vec![0.0f32; rows * cols];
15622        cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
15623        let mut expect = vec![0.0f32; rows];
15624        for r in 0..rows {
15625            expect[r] = reference[r * cols..(r + 1) * cols]
15626                .iter()
15627                .zip(&x)
15628                .map(|(w, xv)| w * xv)
15629                .sum();
15630        }
15631
15632        let mut got = vec![0.0f32; rows];
15633        q4matvec(&bytes, &x, rows, cols, &mut got, None);
15634        // SDOT path quantizes activations to i8 (A8W8): bounded noise,
15635        // same contract as q8/vbit (exact path is pinned by CMF_SDOT=0
15636        // in the golden-parity gate).
15637        let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
15638        let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
15639        for r in 0..rows {
15640            assert!(
15641                (got[r] - expect[r]).abs() < tol * scale,
15642                "row {r}: {} vs {}",
15643                got[r],
15644                expect[r]
15645            );
15646        }
15647    }
15648
15649    /// Fused two-input vbit matvec must equal two single matvecs exactly
15650    /// (same per-lane accumulation order on both scalar and SDOT paths).
15651    #[test]
15652    fn vbitmatvec2_equals_two_singles() {
15653        let (rows, cols) = (6, 64);
15654        let ng = cols / GROUP_SIZE;
15655        let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
15656        let mut bytes = bits.clone();
15657        for g in 0..rows * ng {
15658            let s = 0.02 + 0.001 * g as f32;
15659            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
15660        }
15661        for r in 0..rows {
15662            let b = bits[r] as usize;
15663            let (mut acc, mut nb) = (0u64, 0usize);
15664            let mut rowbytes = Vec::new();
15665            for i in 0..cols {
15666                let v = ((i * 7 + r * 13) % (1 << b)) as u64;
15667                acc = (acc << b) | v;
15668                nb += b;
15669                while nb >= 8 {
15670                    nb -= 8;
15671                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
15672                }
15673            }
15674            if nb > 0 {
15675                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
15676            }
15677            bytes.extend_from_slice(&rowbytes);
15678        }
15679        let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
15680        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).cos()).collect();
15681        let offsets = vbit_row_offsets(&bytes, rows, cols);
15682
15683        let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
15684        vbitmatvec(&bytes, &offsets, &x1, rows, cols, &mut a1, None);
15685        vbitmatvec(&bytes, &offsets, &x2, rows, cols, &mut a2, None);
15686        let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
15687        vbitmatvec2(
15688            &bytes, &offsets, &x1, &x2, rows, cols, &mut b1, &mut b2, None,
15689        );
15690        assert_eq!(a1, b1, "fused vbit lane 1 must be bit-identical");
15691        assert_eq!(a2, b2, "fused vbit lane 2 must be bit-identical");
15692    }
15693
15694    /// Fused two-input q4 matvec must equal two single matvecs exactly.
15695    #[test]
15696    fn q4matvec2_equals_two_singles() {
15697        let (rows, cols) = (8, 128);
15698        let groups = rows * cols / GROUP_SIZE;
15699        let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
15700        for i in 0..groups * 16 {
15701            bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
15702        }
15703        for g in 0..groups {
15704            let s = 0.01 + 0.003 * g as f32;
15705            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
15706        }
15707        // Include an outlier channel so the SDOT correction path is
15708        // exercised in the pair kernel too.
15709        let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
15710        x1[9] = 250.0;
15711        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
15712
15713        let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
15714        q4matvec(&bytes, &x1, rows, cols, &mut a1, None);
15715        q4matvec(&bytes, &x2, rows, cols, &mut a2, None);
15716        let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
15717        q4matvec2(&bytes, &x1, &x2, rows, cols, &mut b1, &mut b2, None);
15718        assert_eq!(a1, b1, "fused q4 lane 1 must be bit-identical");
15719        assert_eq!(a2, b2, "fused q4 lane 2 must be bit-identical");
15720    }
15721
15722    /// Multi-matrix job must equal separate matvecs exactly — same
15723    /// kernels, only the dispatch is fused.
15724    #[test]
15725    fn matvec_many_equals_separate_matvecs() {
15726        use crate::pool::Pool;
15727        let (r1, r2, cols) = (300, 200, 64);
15728        let mk = |salt: usize, rows: usize| {
15729            QTensor::from_f32(
15730                (0..rows * cols)
15731                    .map(|i| ((i * 7 + salt) % 97) as f32 / 97.0 - 0.5)
15732                    .collect(),
15733                rows,
15734                cols,
15735            )
15736        };
15737        let (a, b) = (mk(1, r1), mk(5, r2));
15738        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).sin()).collect();
15739        let pool = Pool::new(3);
15740
15741        let (mut ea, mut eb) = (vec![0f32; r1], vec![0f32; r2]);
15742        a.matvec(&x, &mut ea, Some(&pool));
15743        b.matvec(&x, &mut eb, Some(&pool));
15744        let (mut ga, mut gb) = (vec![0f32; r1], vec![0f32; r2]);
15745        QTensor::matvec_many([&a, &b], &x, [&mut ga, &mut gb], Some(&pool));
15746        assert_eq!(ea, ga, "fused multi-matrix lane 1 must be bit-identical");
15747        assert_eq!(eb, gb, "fused multi-matrix lane 2 must be bit-identical");
15748    }
15749
15750    /// The public Q4TP operator must take the real mapped matvec_many arm,
15751    /// rather than the F32 fallback above.  Build a tiny valid CMF so both
15752    /// handles retain their mmap payloads, then compare the fused dispatch
15753    /// with two ordinary mapped matvec calls bit-for-bit.
15754    #[test]
15755    fn q4tp_matvec_many_equals_separate_matvecs() {
15756        use crate::pool::Pool;
15757        use cortiq_core::{CMF_VERSION, CmfHeader, CmfModel, QuantType, TensorSpec};
15758
15759        let (r1, r2, cols) = (300usize, 200usize, 64usize);
15760        let arch: cortiq_core::ModelArch = serde_json::from_value(serde_json::json!({
15761            "arch_name": "tiny-q4tp",
15762            "hidden_size": cols,
15763            "intermediate_size": cols * 2,
15764            "num_layers": 1,
15765            "num_attention_heads": 2,
15766            "num_kv_heads": 1,
15767            "head_dim": 32,
15768            "vocab_size": r1,
15769            "layer_types": ["FullAttention"],
15770            "rms_norm_eps": 1e-6,
15771            "max_position_embeddings": 8,
15772            "linear_conv_kernel_dim": 0,
15773            "linear_num_key_heads": 0,
15774            "linear_num_value_heads": 0
15775        }))
15776        .unwrap();
15777        let header = CmfHeader {
15778            format: "cmf".into(),
15779            version: CMF_VERSION,
15780            arch,
15781            quant_type: QuantType::Q4Block,
15782            provenance: None,
15783            tokenizer_config: None,
15784            section_hashes: None,
15785            skills: Vec::new(),
15786            shard: None,
15787            calibration: None,
15788            routing: None,
15789            genome: None,
15790            lineage: Vec::new(),
15791            router: None,
15792            segments: Vec::new(),
15793        };
15794        let specs = [
15795            TensorSpec {
15796                name: "q".into(),
15797                dtype: TensorDtype::Q4TiledP,
15798                shape: vec![r1, cols],
15799                data: synth_q4tp(r1, cols),
15800            },
15801            TensorSpec {
15802                name: "kv".into(),
15803                dtype: TensorDtype::Q4TiledP,
15804                shape: vec![r2, cols],
15805                data: synth_q4tp(r2, cols),
15806            },
15807        ];
15808        let dir = std::env::temp_dir().join(format!("cmf-q4tp-many-{}", std::process::id()));
15809        std::fs::create_dir_all(&dir).unwrap();
15810        let path = dir.join("m.cmf");
15811        CmfModel::write(&path, &header, &specs, None, None).unwrap();
15812        let model = Arc::new(CmfModel::open(&path).unwrap());
15813        let (a, b) = (
15814            QTensor::from_model(&model, "q").unwrap(),
15815            QTensor::from_model(&model, "kv").unwrap(),
15816        );
15817        assert_eq!(a.model_dtype(), Some(TensorDtype::Q4TiledP));
15818        assert_eq!(b.model_dtype(), Some(TensorDtype::Q4TiledP));
15819        let x: Vec<f32> = (0..cols)
15820            .map(|i| ((i * 17 + 3) % 97) as f32 / 97.0 - 0.5)
15821            .collect();
15822        let pool = Pool::new(3);
15823        let (mut ea, mut eb) = (vec![0.0f32; r1], vec![0.0f32; r2]);
15824        a.matvec(&x, &mut ea, Some(&pool));
15825        b.matvec(&x, &mut eb, Some(&pool));
15826        let (mut ga, mut gb) = (vec![0.0f32; r1], vec![0.0f32; r2]);
15827        QTensor::matvec_many([&a, &b], &x, [&mut ga, &mut gb], Some(&pool));
15828        assert_eq!(ea, ga, "Q4TP fused lane 1 must be bit-identical");
15829        assert_eq!(eb, gb, "Q4TP fused lane 2 must be bit-identical");
15830        let _ = std::fs::remove_dir_all(&dir);
15831    }
15832
15833    /// The fast expert kernel (`moe_fast`): its scales are bit for bit
15834    /// those of `scales_into`, whole experts equal `moe_gate_up_rows` +
15835    /// `moe_down_rows` up to the rounding of a sixteen-lane sum, with
15836    /// activation outliers on both stages, empty groups, groups past four
15837    /// tokens, and the same bits with and without a pool.
15838    #[cfg(target_arch = "x86_64")]
15839    #[test]
15840    fn fast_expert_rows_match_the_row_kernels() {
15841        use crate::pool::Pool;
15842        use cortiq_core::{CMF_VERSION, CmfHeader, CmfModel, QuantType, TensorSpec};
15843
15844        // The scales, every row width the code decode has a path for.
15845        for (rows, cols) in [(32usize, 2560usize), (16, 640), (48, 192), (16, 64)] {
15846            let bytes = synth_q4tp(rows, cols);
15847            let v = Q4tpView::new(&bytes, rows, cols);
15848            let gpr = cols / GROUP_SIZE;
15849            for r0 in (0..rows).step_by(16) {
15850                let Some(got) = moe_fast::block_scales(&v, r0, gpr) else {
15851                    eprintln!("fast_expert_rows: no AVX-512 VNNI here, skipped");
15852                    return;
15853                };
15854                for i in 0..16 {
15855                    let mut want = vec![0f32; gpr];
15856                    v.scales_into(r0 + i, gpr, &mut want);
15857                    assert_eq!(
15858                        want.iter().map(|x| x.to_bits()).collect::<Vec<_>>(),
15859                        got[i * gpr..(i + 1) * gpr].iter().map(|x| x.to_bits()).collect::<Vec<_>>(),
15860                        "scales of row {} ({rows}x{cols})",
15861                        r0 + i
15862                    );
15863                }
15864            }
15865        }
15866
15867        let (h, inter, ne) = (256usize, 640usize, 5usize);
15868        let arch: cortiq_core::ModelArch = serde_json::from_value(serde_json::json!({
15869            "arch_name": "tiny-q4tp-moe",
15870            "hidden_size": h,
15871            "intermediate_size": inter,
15872            "num_layers": 1,
15873            "num_attention_heads": h / 32,
15874            "num_kv_heads": 1,
15875            "head_dim": 32,
15876            "vocab_size": 8,
15877            "layer_types": ["FullAttention"],
15878            "rms_norm_eps": 1e-6,
15879            "max_position_embeddings": 8,
15880            "linear_conv_kernel_dim": 0,
15881            "linear_num_key_heads": 0,
15882            "linear_num_value_heads": 0
15883        }))
15884        .unwrap();
15885        let header = CmfHeader {
15886            format: "cmf".into(),
15887            version: CMF_VERSION,
15888            arch,
15889            quant_type: QuantType::Q4Block,
15890            provenance: None,
15891            tokenizer_config: None,
15892            section_hashes: None,
15893            skills: Vec::new(),
15894            shard: None,
15895            calibration: None,
15896            routing: None,
15897            genome: None,
15898            lineage: Vec::new(),
15899            router: None,
15900            segments: Vec::new(),
15901        };
15902        let mut specs = Vec::new();
15903        for e in 0..ne {
15904            for (k, (n, r, c)) in [("g", inter, h), ("u", inter, h), ("d", h, inter)]
15905                .into_iter()
15906                .enumerate()
15907            {
15908                let mut data = synth_q4tp(r, c);
15909                for (i, byte) in data[..r * (c / GROUP_SIZE) * Q4TP_NIB]
15910                    .iter_mut()
15911                    .enumerate()
15912                {
15913                    *byte ^= ((i * (e * 3 + k + 1)) % 251) as u8;
15914                }
15915                specs.push(TensorSpec {
15916                    name: format!("{n}{e}"),
15917                    dtype: TensorDtype::Q4TiledP,
15918                    shape: vec![r, c],
15919                    data,
15920                });
15921            }
15922        }
15923        let dir = std::env::temp_dir().join(format!("cmf-moe-fast-{}", std::process::id()));
15924        std::fs::create_dir_all(&dir).unwrap();
15925        let path = dir.join("m.cmf");
15926        CmfModel::write(&path, &header, &specs, None, None).unwrap();
15927        let model = Arc::new(CmfModel::open(&path).unwrap());
15928        let t = |n: String| QTensor::from_model(&model, &n).unwrap();
15929        let g: Vec<QTensor> = (0..ne).map(|e| t(format!("g{e}"))).collect();
15930        let u: Vec<QTensor> = (0..ne).map(|e| t(format!("u{e}"))).collect();
15931        let d: Vec<QTensor> = (0..ne).map(|e| t(format!("d{e}"))).collect();
15932        let b = 6usize;
15933        let mut xs: Vec<f32> = (0..b * h)
15934            .map(|i| ((i * 31 + 7) % 89) as f32 / 89.0 - 0.5)
15935            .collect();
15936        xs[5] = 9.0; // an activation outlier on token 0
15937        // Expert -> its tokens: five (a full chunk and one more), one, none, three, two.
15938        let groups: Vec<Vec<usize>> = vec![vec![0, 1, 2, 3, 4], vec![5], vec![], vec![2, 0, 5], vec![1, 3]];
15939        let lens: Vec<usize> = groups.iter().map(|v| v.len()).collect();
15940        let n_pairs: usize = lens.iter().sum();
15941        let pool = Pool::new(3);
15942        let close = |want: &[f32], got: &[f32], rel: f32, what: &str| {
15943            assert_eq!(want.len(), got.len(), "{what}: length");
15944            let top = want.iter().fold(0f32, |m, v| m.max(v.abs()));
15945            assert!(top > 0.0, "{what}: an all-zero reference proves nothing");
15946            for (i, (w, v)) in want.iter().zip(got).enumerate() {
15947                assert!((w - v).abs() <= rel * top, "{what}: [{i}] {w} vs {v} (scale {top})");
15948            }
15949        };
15950
15951        // The reference: the two row kernels.
15952        let pairs: Vec<(&QTensor, &QTensor)> = (0..ne).map(|e| (&g[e], &u[e])).collect();
15953        let mut gs: Vec<Vec<f32>> = (0..n_pairs).map(|_| vec![0f32; inter]).collect();
15954        assert!(QTensor::moe_gate_up_rows(&pairs, &groups, &xs, &mut gs, Some(&pool)));
15955        let downs: Vec<&QTensor> = d.iter().collect();
15956        let mut ds: Vec<Vec<f32>> = (0..n_pairs).map(|_| vec![0f32; h]).collect();
15957        assert!(QTensor::moe_down_rows(&downs, &lens, &gs, &mut ds, Some(&pool)));
15958
15959        let triples: Vec<(&QTensor, &QTensor, &QTensor)> = (0..ne).map(|e| (&g[e], &u[e], &d[e])).collect();
15960        let mut out = vec![0f32; n_pairs * h];
15961        let mut gs_fast = Vec::new();
15962        if !moe_fast::run_stages(&triples, &groups, &xs, &mut out, Some(&pool), None, Some(&mut gs_fast)) {
15963            eprintln!("fast_expert_rows: the int8 activation kernels are off here, skipped");
15964            let _ = std::fs::remove_dir_all(&dir);
15965            return;
15966        }
15967        close(&gs.concat(), &gs_fast, 1e-4, "gate and up");
15968        close(&ds.concat(), &out, 1e-2, "whole experts");
15969
15970        // The down stage alone on the reference's gate rows, one of them
15971        // with an outlier (the gate rows of a real expert have them).
15972        let mut gs_out = gs.clone();
15973        let top = gs.iter().flatten().fold(0f32, |m, v| m.max(v.abs()));
15974        gs_out[0][3] = 40.0 * top;
15975        let mut ds_out: Vec<Vec<f32>> = (0..n_pairs).map(|_| vec![0f32; h]).collect();
15976        assert!(QTensor::moe_down_rows(&downs, &lens, &gs_out, &mut ds_out, Some(&pool)));
15977        let mut out2 = vec![0f32; n_pairs * h];
15978        assert!(moe_fast::run_stages(&triples, &groups, &xs, &mut out2, Some(&pool), Some(&gs_out.concat()), None));
15979        close(&ds_out.concat(), &out2, 1e-4, "down");
15980
15981        // The same bits on the calling thread alone, and a refusal the caller can fall back from.
15982        let mut out3 = vec![0f32; n_pairs * h];
15983        assert!(QTensor::moe_expert_rows_fast(&triples, &groups, &xs, &mut out3, None));
15984        assert_eq!(
15985            out.iter().map(|x| x.to_bits()).collect::<Vec<_>>(),
15986            out3.iter().map(|x| x.to_bits()).collect::<Vec<_>>(),
15987            "pooled and serial runs must agree bit for bit"
15988        );
15989        // A pool on two NUMA nodes gives each side half of every matrix's rows: the same bits again.
15990        pool.mark_two_nodes(true);
15991        let mut out4 = vec![0f32; n_pairs * h];
15992        assert!(QTensor::moe_expert_rows_fast(&triples, &groups, &xs, &mut out4, Some(&pool)));
15993        let mut out5 = vec![0f32; n_pairs * h];
15994        assert!(moe_fast::run_stages(&triples, &groups, &xs, &mut out5, Some(&pool), Some(&gs_out.concat()), None));
15995        pool.mark_two_nodes(false);
15996        assert_eq!(
15997            out.iter().map(|x| x.to_bits()).collect::<Vec<_>>(),
15998            out4.iter().map(|x| x.to_bits()).collect::<Vec<_>>(),
15999            "the rows split over two sides must agree bit for bit"
16000        );
16001        assert_eq!(
16002            out2.iter().map(|x| x.to_bits()).collect::<Vec<_>>(),
16003            out5.iter().map(|x| x.to_bits()).collect::<Vec<_>>(),
16004            "the down rows split over two sides must agree bit for bit"
16005        );
16006        let mut short = vec![0f32; n_pairs * h - 1];
16007        assert!(!QTensor::moe_expert_rows_fast(&triples, &groups, &xs, &mut short, Some(&pool)));
16008        let _ = std::fs::remove_dir_all(&dir);
16009    }
16010
16011    /// The MiMo speculative verify's kernels: several tokens' MoE through
16012    /// `moe_gate_up_rows` / `moe_down_rows` (+ the caller's route-order sum)
16013    /// is bit-identical to each token alone through `moe_gate_up_many` /
16014    /// `moe_down_many` (decode), and a row-exact `q4tp_matmat` of five
16015    /// tokens (wide enough for the blocked tiles) equals five matvecs.
16016    #[test]
16017    fn multi_token_moe_rows_equal_single_token_decode() {
16018        use crate::pool::Pool;
16019        use cortiq_core::{CMF_VERSION, CmfHeader, CmfModel, QuantType, TensorSpec};
16020
16021        let (h, inter, ne) = (64usize, 128usize, 3usize);
16022        let arch: cortiq_core::ModelArch = serde_json::from_value(serde_json::json!({
16023            "arch_name": "tiny-q4tp-moe",
16024            "hidden_size": h,
16025            "intermediate_size": inter,
16026            "num_layers": 1,
16027            "num_attention_heads": 2,
16028            "num_kv_heads": 1,
16029            "head_dim": 32,
16030            "vocab_size": 8,
16031            "layer_types": ["FullAttention"],
16032            "rms_norm_eps": 1e-6,
16033            "max_position_embeddings": 8,
16034            "linear_conv_kernel_dim": 0,
16035            "linear_num_key_heads": 0,
16036            "linear_num_value_heads": 0
16037        }))
16038        .unwrap();
16039        let header = CmfHeader {
16040            format: "cmf".into(),
16041            version: CMF_VERSION,
16042            arch,
16043            quant_type: QuantType::Q4Block,
16044            provenance: None,
16045            tokenizer_config: None,
16046            section_hashes: None,
16047            skills: Vec::new(),
16048            shard: None,
16049            calibration: None,
16050            routing: None,
16051            genome: None,
16052            lineage: Vec::new(),
16053            router: None,
16054            segments: Vec::new(),
16055        };
16056        let mut specs = Vec::new();
16057        for e in 0..ne {
16058            for (k, (n, r, c)) in [("g", inter, h), ("u", inter, h), ("d", h, inter)]
16059                .into_iter()
16060                .enumerate()
16061            {
16062                // Distinct experts: perturb only the nibble plane (any byte
16063                // is a valid pair of codes; the ladder stays intact).
16064                let mut data = synth_q4tp(r, c);
16065                for (i, byte) in data[..r * (c / GROUP_SIZE) * Q4TP_NIB]
16066                    .iter_mut()
16067                    .enumerate()
16068                {
16069                    *byte ^= ((i * (e * 3 + k + 1)) % 251) as u8;
16070                }
16071                specs.push(TensorSpec {
16072                    name: format!("{n}{e}"),
16073                    dtype: TensorDtype::Q4TiledP,
16074                    shape: vec![r, c],
16075                    data,
16076                });
16077            }
16078        }
16079        let dir = std::env::temp_dir().join(format!(
16080            "cmf-moe-rows-{}-{}",
16081            std::process::id(),
16082            FLOAT_ACTIVATIONS.get()
16083        ));
16084        std::fs::create_dir_all(&dir).unwrap();
16085        let path = dir.join("m.cmf");
16086        CmfModel::write(&path, &header, &specs, None, None).unwrap();
16087        let model = Arc::new(CmfModel::open(&path).unwrap());
16088        let t = |n: String| QTensor::from_model(&model, &n).unwrap();
16089        let g: Vec<QTensor> = (0..ne).map(|e| t(format!("g{e}"))).collect();
16090        let u: Vec<QTensor> = (0..ne).map(|e| t(format!("u{e}"))).collect();
16091        let d: Vec<QTensor> = (0..ne).map(|e| t(format!("d{e}"))).collect();
16092        let b = 4usize;
16093        let mut xs: Vec<f32> = (0..b * h)
16094            .map(|i| ((i * 31 + 7) % 89) as f32 / 89.0 - 0.5)
16095            .collect();
16096        xs[5] = 9.0; // an activation outlier on token 0
16097        // Token -> (experts in route order, weights).
16098        let routes: Vec<(Vec<usize>, Vec<f32>)> = vec![
16099            (vec![2, 0], vec![0.6, 0.4]),
16100            (vec![0, 1, 2], vec![0.2, 0.5, 0.3]),
16101            (vec![1], vec![1.0]),
16102            (vec![2, 1, 0], vec![0.25, 0.25, 0.5]),
16103        ];
16104        let pool = Pool::new(3);
16105        // Decode reference, token by token.
16106        let mut want = vec![0f32; b * h];
16107        for (tk, (idx, w)) in routes.iter().enumerate() {
16108            let x = &xs[tk * h..(tk + 1) * h];
16109            let pairs: Vec<(&QTensor, &QTensor)> = idx.iter().map(|&e| (&g[e], &u[e])).collect();
16110            let mut gs: Vec<Vec<f32>> = idx.iter().map(|_| vec![0f32; inter]).collect();
16111            assert!(QTensor::moe_gate_up_many(&pairs, x, &mut gs, Some(&pool)));
16112            if FLOAT_ACTIVATIONS.get() {
16113                for (slot, &e) in idx.iter().enumerate() {
16114                    let (mut gate, mut up) = (vec![0.0; inter], vec![0.0; inter]);
16115                    g[e].matvec(x, &mut gate, Some(&pool));
16116                    u[e].matvec(x, &mut up, Some(&pool));
16117                    for (v, u) in gate.iter_mut().zip(up) {
16118                        *v = (*v / (1.0 + (-*v).exp())) * u;
16119                    }
16120                    assert_eq!(gs[slot], gate, "float gate/up must equal ordinary matvecs");
16121                }
16122            }
16123            let downs: Vec<&QTensor> = idx.iter().map(|&e| &d[e]).collect();
16124            assert!(QTensor::moe_down_many(
16125                &downs,
16126                &gs,
16127                w,
16128                &mut want[tk * h..(tk + 1) * h],
16129                Some(&pool)
16130            ));
16131        }
16132        if FLOAT_ACTIVATIONS.get() {
16133            for (tk, (idx, w)) in routes.iter().enumerate() {
16134                let mut scalar = vec![0.0; h];
16135                for (&e, &weight) in idx.iter().zip(w) {
16136                    let (mut gate, mut up, mut down) =
16137                        (vec![0.0; inter], vec![0.0; inter], vec![0.0; h]);
16138                    g[e].matvec(&xs[tk * h..(tk + 1) * h], &mut gate, Some(&pool));
16139                    u[e].matvec(&xs[tk * h..(tk + 1) * h], &mut up, Some(&pool));
16140                    for (v, u) in gate.iter_mut().zip(up) {
16141                        *v = (*v / (1.0 + (-*v).exp())) * u;
16142                    }
16143                    d[e].matvec(&gate, &mut down, Some(&pool));
16144                    for (v, d) in scalar.iter_mut().zip(down) {
16145                        *v += weight * d;
16146                    }
16147                }
16148                assert_eq!(
16149                    &want[tk * h..(tk + 1) * h],
16150                    scalar,
16151                    "float many equals scalar experts"
16152                );
16153            }
16154        }
16155        // All four tokens at once, grouped by expert.
16156        let mut experts: Vec<usize> = Vec::new();
16157        let mut groups: Vec<Vec<usize>> = Vec::new();
16158        for (tk, (idx, _)) in routes.iter().enumerate() {
16159            for &e in idx {
16160                match experts.iter().position(|&x| x == e) {
16161                    Some(k) => groups[k].push(tk),
16162                    None => {
16163                        experts.push(e);
16164                        groups.push(vec![tk]);
16165                    }
16166                }
16167            }
16168        }
16169        let n_pairs: usize = groups.iter().map(|g| g.len()).sum();
16170        let pairs: Vec<(&QTensor, &QTensor)> = experts.iter().map(|&e| (&g[e], &u[e])).collect();
16171        let mut gs: Vec<Vec<f32>> = (0..n_pairs).map(|_| vec![0f32; inter]).collect();
16172        assert!(QTensor::moe_gate_up_rows(
16173            &pairs,
16174            &groups,
16175            &xs,
16176            &mut gs,
16177            Some(&pool)
16178        ));
16179        let downs: Vec<&QTensor> = experts.iter().map(|&e| &d[e]).collect();
16180        let lens: Vec<usize> = groups.iter().map(|g| g.len()).collect();
16181        let mut ds: Vec<Vec<f32>> = (0..n_pairs).map(|_| vec![0f32; h]).collect();
16182        assert!(QTensor::moe_down_rows(
16183            &downs,
16184            &lens,
16185            &gs,
16186            &mut ds,
16187            Some(&pool)
16188        ));
16189        let slot = |tk: usize, e: usize| {
16190            let k = experts.iter().position(|&x| x == e).unwrap();
16191            groups[..k].iter().map(|g| g.len()).sum::<usize>()
16192                + groups[k].iter().position(|&x| x == tk).unwrap()
16193        };
16194        let mut got = vec![0f32; b * h];
16195        for (tk, (idx, w)) in routes.iter().enumerate() {
16196            for i in 0..h {
16197                let mut acc = 0f32;
16198                for (&e, &we) in idx.iter().zip(w) {
16199                    acc += we * ds[slot(tk, e)][i];
16200                }
16201                got[tk * h + i] = acc;
16202            }
16203        }
16204        assert!(want.iter().any(|v| *v != 0.0));
16205        assert_eq!(
16206            want.iter().map(|v| v.to_bits()).collect::<Vec<_>>(),
16207            got.iter().map(|v| v.to_bits()).collect::<Vec<_>>(),
16208            "multi-token MoE must equal decode bit for bit"
16209        );
16210
16211        // Row-exact q4tp_matmat: five tokens (a blocked 1x4 tile + tail
16212        // otherwise) equal five matvecs.
16213        let b5 = 5usize;
16214        let x5: Vec<f32> = (0..b5 * h)
16215            .map(|i| ((i * 13 + 5) % 71) as f32 / 71.0 - 0.5)
16216            .collect();
16217        let mut mm = vec![0f32; b5 * inter];
16218        row_exact_scope(|| g[1].matmat(&x5, b5, &mut mm, Some(&pool)));
16219        for tk in 0..b5 {
16220            let mut mv = vec![0f32; inter];
16221            g[1].matvec(&x5[tk * h..(tk + 1) * h], &mut mv, Some(&pool));
16222            assert_eq!(
16223                mv.iter().map(|v| v.to_bits()).collect::<Vec<_>>(),
16224                mm[tk * inter..(tk + 1) * inter]
16225                    .iter()
16226                    .map(|v| v.to_bits())
16227                    .collect::<Vec<_>>(),
16228                "row-exact matmat token {tk}"
16229            );
16230        }
16231        // Other concurrent tests/requests may still hold the shared mode.
16232        // Nested, overlapping and unwind restoration is checked separately.
16233        let _ = std::fs::remove_dir_all(&dir);
16234    }
16235
16236    /// The row-exact fix must not touch the fast path. Outside the scope
16237    /// `q4tp_matmat` has to produce exactly what it did before, and on ARM
16238    /// "before" is spelled out below: the tuned 1x4 SDOT tile for every
16239    /// four columns and the single-row kernel for the tail. Inside the
16240    /// scope every column equals its token's matvec. `q4tp_matmat_with`
16241    /// takes the mode as an argument, so a concurrent test holding the
16242    /// shared scope cannot flip it under this one.
16243    #[test]
16244    fn q4tp_matmat_fast_path_unchanged_outside_row_exact() {
16245        use crate::pool::Pool;
16246        use std::sync::atomic::Ordering::Relaxed;
16247        let _alt = Q4TP_ALT_TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner());
16248        // The tuned ARM shape, which is also what an unset switch picks
16249        // unless CMF_Q4TP_V1 is exported.
16250        Q4TP_ALT.store(2, Relaxed);
16251        let pool = Pool::new(3);
16252        // Under 500k cells, so macOS keeps the matmat off the AMX; the
16253        // second shape runs across pool workers (rows >= 256) with 32
16254        // groups of accumulation and a tail after two 1x4 tiles.
16255        for &(rows, cols, b) in &[(64usize, 256usize, 7usize), (320, 1024, 9)] {
16256            let bytes = synth_q4tp(rows, cols);
16257            let mut xs: Vec<f32> = (0..b * cols)
16258                .map(|i| ((i * 29 + 11) % 83) as f32 / 83.0 - 0.5)
16259                .collect();
16260            xs[3] = 7.5; // an activation outlier on token 0
16261            let run = |exact: bool| {
16262                let mut out = vec![0f32; b * rows];
16263                q4tp_matmat_with(&bytes, &xs, b, rows, cols, &mut out, Some(&pool), exact);
16264                out
16265            };
16266            let (fast, exact) = (run(false), run(true));
16267            #[cfg(not(target_arch = "aarch64"))]
16268            let _ = fast;
16269            let bits = |v: &[f32]| v.iter().map(|x| x.to_bits()).collect::<Vec<_>>();
16270            let mut matvecs = vec![0f32; b * rows];
16271            for (bi, o) in matvecs.chunks_mut(rows).enumerate() {
16272                q4tp_matvec(
16273                    &bytes,
16274                    &xs[bi * cols..(bi + 1) * cols],
16275                    rows,
16276                    cols,
16277                    o,
16278                    Some(&pool),
16279                );
16280            }
16281            assert!(matvecs.iter().any(|v| *v != 0.0));
16282            assert_eq!(
16283                bits(&exact),
16284                bits(&matvecs),
16285                "{rows}x{cols} b={b}: row-exact matmat must equal per-token matvecs"
16286            );
16287            #[cfg(target_arch = "aarch64")]
16288            {
16289                // The pre-fix ARM loop, cell for cell.
16290                let gpr = cols / GROUP_SIZE;
16291                let v = Q4tpView::new(&bytes, rows, cols);
16292                let mut old = vec![0f32; b * rows];
16293                let mut sc = vec![0f32; gpr];
16294                let a8w8 = a8w8_enabled();
16295                let blocked = sdot_enabled() && blocked_enabled();
16296                let acts: Vec<SplitAct> = (0..b)
16297                    .map(|bi| split_act(&xs[bi * cols..(bi + 1) * cols]))
16298                    .collect();
16299                for r in 0..rows {
16300                    v.scales_into(r, gpr, &mut sc);
16301                    if !a8w8 {
16302                        for bi in 0..b {
16303                            let x = &xs[bi * cols..(bi + 1) * cols];
16304                            old[bi * rows + r] = q4tp_row_exact(v.nib, r, gpr, x, &sc);
16305                        }
16306                        continue;
16307                    }
16308                    let finish = |d: f32, act: &SplitAct| {
16309                        let mut acc = d * act.sx;
16310                        for &(j, xv) in &act.outliers {
16311                            let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
16312                            acc += w * s * xv;
16313                        }
16314                        acc
16315                    };
16316                    let mut bi = 0usize;
16317                    while blocked && bi + 4 <= b {
16318                        let xs4 = [
16319                            acts[bi].xq.as_slice(),
16320                            acts[bi + 1].xq.as_slice(),
16321                            acts[bi + 2].xq.as_slice(),
16322                            acts[bi + 3].xq.as_slice(),
16323                        ];
16324                        let d = unsafe { dot_q4tp_row_1x4_sdot(v.nib, r, gpr, xs4, &sc) };
16325                        for k in 0..4 {
16326                            old[(bi + k) * rows + r] = finish(d[k], &acts[bi + k]);
16327                        }
16328                        bi += 4;
16329                    }
16330                    for (bi, act) in acts.iter().enumerate().skip(bi) {
16331                        let d = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, &sc);
16332                        old[bi * rows + r] = finish(d, act);
16333                    }
16334                }
16335                assert_eq!(
16336                    bits(&fast),
16337                    bits(&old),
16338                    "{rows}x{cols} b={b}: the fast path changed outside row_exact"
16339                );
16340                // And it is still the fast tile that runs: its lane-parallel
16341                // fma sum rounds differently from the matvec somewhere.
16342                if blocked {
16343                    assert_ne!(
16344                        bits(&fast),
16345                        bits(&matvecs),
16346                        "{rows}x{cols} b={b}: the tuned tile no longer runs outside row_exact"
16347                    );
16348                }
16349            }
16350        }
16351        Q4TP_ALT.store(0, Relaxed);
16352    }
16353
16354    #[test]
16355    fn row_exact_scopes_survive_overlap_nesting_and_unwind() {
16356        use std::sync::{Barrier, atomic::{AtomicUsize, Ordering}};
16357        // A private counter makes this restoration test independent of
16358        // numerical tests concurrently using the production counter.
16359        let active = AtomicUsize::new(0);
16360        counted_row_exact_scope(&active, || {
16361            assert_eq!(active.load(Ordering::Acquire), 1);
16362            counted_row_exact_scope(&active, || {
16363                assert_eq!(active.load(Ordering::Acquire), 2);
16364            });
16365            assert_eq!(active.load(Ordering::Acquire), 1);
16366        });
16367        assert_eq!(active.load(Ordering::Acquire), 0);
16368
16369        let both_entered = Barrier::new(2);
16370        let release_last = Barrier::new(2);
16371        std::thread::scope(|s| {
16372            let first = s.spawn(|| counted_row_exact_scope(&active, || {
16373                both_entered.wait();
16374            }));
16375            let last = s.spawn(|| counted_row_exact_scope(&active, || {
16376                both_entered.wait();
16377                release_last.wait();
16378            }));
16379            first.join().unwrap();
16380            let after_first = active.load(Ordering::Acquire);
16381            release_last.wait();
16382            last.join().unwrap();
16383            assert_eq!(after_first, 1, "second request must remain exact");
16384        });
16385        assert_eq!(active.load(Ordering::Acquire), 0);
16386        let panic = std::panic::catch_unwind(|| {
16387            counted_row_exact_scope(&active, || panic!("scope unwind"));
16388        });
16389        assert!(panic.is_err());
16390        assert_eq!(active.load(Ordering::Acquire), 0);
16391    }
16392
16393    #[test]
16394    #[cfg(target_arch = "x86_64")]
16395    fn q4tp_float_avx2_is_bitwise_scalar() {
16396        if !avx2_enabled() {
16397            return;
16398        }
16399        for cols in [32, 64, 96, 2048, 4096] {
16400            let rows = 9;
16401            let bytes = synth_q4tp(rows, cols);
16402            let v = Q4tpView::new(&bytes, rows, cols);
16403            let gpr = cols / GROUP_SIZE;
16404            let mut sc = vec![0.0; gpr];
16405            for seed in 1..=5 {
16406                let xs: Vec<f32> = (0..cols)
16407                    .map(|i| (((i * 104729 + seed * 8191) % 100003) as f32 - 50001.0) / 7919.0)
16408                    .collect();
16409                for r in 0..rows {
16410                    v.scales_into(r, gpr, &mut sc);
16411                    let scalar = q4tp_row_float_scalar(v.nib, r, gpr, &xs, &sc);
16412                    let vector = unsafe { q4tp_row_float_avx2(v.nib, r, gpr, &xs, &sc) };
16413                    assert_eq!(
16414                        scalar.to_bits(),
16415                        vector.to_bits(),
16416                        "cols={cols} row={r} seed={seed}"
16417                    );
16418                }
16419            }
16420        }
16421    }
16422
16423    #[test]
16424    fn multi_token_moe_rows_float_equal_single_token_decode() {
16425        float_activations_scope(multi_token_moe_rows_equal_single_token_decode);
16426    }
16427
16428    #[test]
16429    fn full_gpu_q8_scope_is_nested_and_thread_local() {
16430        assert!(!FULL_GPU_Q8.get());
16431        let before = gpu_split_frac();
16432        {
16433            let _guard = enter_full_gpu_q8_scope();
16434            assert_eq!(gpu_split_frac(), 1.0);
16435            {
16436                let _nested = enter_full_gpu_q8_scope();
16437            }
16438            assert_eq!(gpu_split_frac(), 1.0);
16439            std::thread::spawn(|| assert!(!FULL_GPU_Q8.get())).join().unwrap();
16440        }
16441        assert!(!FULL_GPU_Q8.get());
16442        assert_eq!(gpu_split_frac(), before);
16443    }
16444
16445    #[test]
16446    fn float_activation_scope_is_nested_thread_local_and_unwind_safe() {
16447        assert!(!FLOAT_ACTIVATIONS.get());
16448        let before = a8w8_enabled();
16449        float_activations_scope(|| {
16450            assert!(!a8w8_enabled());
16451            float_activations_scope(|| assert!(!a8w8_enabled()));
16452            assert!(FLOAT_ACTIVATIONS.get());
16453            std::thread::spawn(|| assert!(!FLOAT_ACTIVATIONS.get()))
16454                .join()
16455                .unwrap();
16456        });
16457        assert!(!FLOAT_ACTIVATIONS.get());
16458        assert_eq!(a8w8_enabled(), before);
16459        let _ = std::panic::catch_unwind(|| float_activations_scope(|| panic!("test unwind")));
16460        assert!(!FLOAT_ACTIVATIONS.get());
16461    }
16462
16463    /// Batched q4/vbit matmat must equal per-position matvec calls
16464    /// exactly (the fallback it replaced) — same kernels, same order.
16465    #[test]
16466    fn batched_matmat_equals_per_position_matvec() {
16467        let (rows, cols, b) = (8, 64, 5);
16468        // q4 blob.
16469        let groups = rows * cols / GROUP_SIZE;
16470        let mut q4 = Vec::new();
16471        for i in 0..groups * 16 {
16472            q4.push((((i * 7 + 3) % 256) & 0xFF) as u8);
16473        }
16474        for g in 0..groups {
16475            q4.extend_from_slice(
16476                &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
16477            );
16478        }
16479        // vbit blob (mixed widths incl. 8).
16480        let ng = cols / GROUP_SIZE;
16481        let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4, 5, 3];
16482        let mut vb = bits.clone();
16483        for g in 0..rows * ng {
16484            vb.extend_from_slice(
16485                &cortiq_core::quant::f32_to_f16(0.02 + 0.001 * g as f32).to_le_bytes(),
16486            );
16487        }
16488        for r in 0..rows {
16489            let bw = bits[r] as usize;
16490            let (mut acc, mut nb) = (0u64, 0usize);
16491            let mut rowbytes = Vec::new();
16492            for i in 0..cols {
16493                let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
16494                acc = (acc << bw) | v;
16495                nb += bw;
16496                while nb >= 8 {
16497                    nb -= 8;
16498                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
16499                }
16500            }
16501            if nb > 0 {
16502                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
16503            }
16504            vb.extend_from_slice(&rowbytes);
16505        }
16506        let offsets = vbit_row_offsets(&vb, rows, cols);
16507
16508        let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
16509
16510        // q4: batch vs singles.
16511        let mut got = vec![0f32; b * rows];
16512        q4matmat(&q4, &xs, b, rows, cols, &mut got, None);
16513        for bi in 0..b {
16514            let mut expect = vec![0f32; rows];
16515            q4matvec(
16516                &q4,
16517                &xs[bi * cols..(bi + 1) * cols],
16518                rows,
16519                cols,
16520                &mut expect,
16521                None,
16522            );
16523            assert_eq!(
16524                &got[bi * rows..(bi + 1) * rows],
16525                &expect[..],
16526                "q4 batch pos {bi}"
16527            );
16528        }
16529
16530        // vbit: batch vs singles.
16531        let mut got = vec![0f32; b * rows];
16532        vbitmatmat(&vb, &offsets, &xs, b, rows, cols, &mut got, None);
16533        for bi in 0..b {
16534            let mut expect = vec![0f32; rows];
16535            vbitmatvec(
16536                &vb,
16537                &offsets,
16538                &xs[bi * cols..(bi + 1) * cols],
16539                rows,
16540                cols,
16541                &mut expect,
16542                None,
16543            );
16544            assert_eq!(
16545                &got[bi * rows..(bi + 1) * rows],
16546                &expect[..],
16547                "vbit batch pos {bi}"
16548            );
16549        }
16550    }
16551
16552    /// q4_tiled kernels must produce BIT-identical outputs to the q4
16553    /// split kernels on the same values (same ints, same order — only
16554    /// the byte placement differs).
16555    #[test]
16556    fn q4_tiled_matches_q4_block_bitexact() {
16557        let (rows, cols, b) = (8usize, 128usize, 3usize);
16558        let groups = rows * cols / GROUP_SIZE;
16559        let mut split = Vec::with_capacity(groups * 18);
16560        for i in 0..groups * 16 {
16561            split.push((((i * 7 + 3) % 256) & 0xFF) as u8);
16562        }
16563        for g in 0..groups {
16564            split.extend_from_slice(
16565                &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
16566            );
16567        }
16568        // Re-tile: [scale][nibbles] per group.
16569        let (packed, scales) = split.split_at(groups * 16);
16570        let mut tiled = Vec::with_capacity(groups * Q4_TILE);
16571        for g in 0..groups {
16572            tiled.extend_from_slice(&scales[g * 2..g * 2 + 2]);
16573            tiled.extend_from_slice(&packed[g * 16..(g + 1) * 16]);
16574        }
16575
16576        let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
16577        x1[9] = 250.0; // exercise the outlier path
16578        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
16579
16580        let (mut a, mut t) = (vec![0f32; rows], vec![0f32; rows]);
16581        q4matvec(&split, &x1, rows, cols, &mut a, None);
16582        q4t_matvec(&tiled, &x1, rows, cols, &mut t, None);
16583        assert_eq!(a, t, "q4t matvec must match q4 bit-for-bit");
16584
16585        let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
16586        let (mut t1, mut t2) = (vec![0f32; rows], vec![0f32; rows]);
16587        q4matvec2(&split, &x1, &x2, rows, cols, &mut a1, &mut a2, None);
16588        q4t_matvec2(&tiled, &x1, &x2, rows, cols, &mut t1, &mut t2, None);
16589        assert_eq!(a1, t1);
16590        assert_eq!(a2, t2);
16591
16592        let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
16593        let (mut am, mut tm) = (vec![0f32; b * rows], vec![0f32; b * rows]);
16594        q4matmat(&split, &xs, b, rows, cols, &mut am, None);
16595        q4t_matmat(&tiled, &xs, b, rows, cols, &mut tm, None);
16596        assert_eq!(am, tm, "q4t matmat must match q4 bit-for-bit");
16597    }
16598
16599    /// q4 SDOT outlier correction: a single huge activation channel
16600    /// (>8·rms → outlier, zeroed in xq) must still contribute its EXACT
16601    /// term. On-grid bulk (±1/0 → xq dequantizes exactly) isolates the
16602    /// correction from A8W8 noise. cols must exceed 64: at n=64 the
16603    /// 8·rms threshold equals sqrt(v²+rest) ≥ v, so a single outlier
16604    /// can never qualify (8² = n).
16605    #[test]
16606    fn q4matvec_sdot_outlier_exact() {
16607        let (rows, cols) = (4, 128);
16608        let groups = rows * cols / GROUP_SIZE;
16609        let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
16610        for i in 0..groups * 16 {
16611            bytes.push(((i * 11 + 5) % 256) as u8);
16612        }
16613        for g in 0..groups {
16614            let s = 0.02 + 0.002 * g as f32;
16615            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
16616        }
16617        let mut x: Vec<f32> = (0..cols)
16618            .map(|i| match i % 3 {
16619                0 => 1.0,
16620                1 => -1.0,
16621                _ => 0.0,
16622            })
16623            .collect();
16624        x[17] = 300.0; // ≫ 8·rms → outlier channel
16625
16626        let mut reference = vec![0.0f32; rows * cols];
16627        cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
16628        let mut expect = vec![0.0f32; rows];
16629        for r in 0..rows {
16630            expect[r] = reference[r * cols..(r + 1) * cols]
16631                .iter()
16632                .zip(&x)
16633                .map(|(w, xv)| w * xv)
16634                .sum();
16635        }
16636        let mut got = vec![0.0f32; rows];
16637        q4matvec(&bytes, &x, rows, cols, &mut got, None);
16638        let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
16639        for r in 0..rows {
16640            assert!(
16641                (got[r] - expect[r]).abs() < 2e-3 * scale,
16642                "row {r}: {} vs {} (outlier term must be exact)",
16643                got[r],
16644                expect[r]
16645            );
16646        }
16647    }
16648
16649    /// The fused q1t matvec must equal the reference (dequant_q1t → dot),
16650    /// including the ternary zero level and the binary-searched outlier
16651    /// overlay. Guards the mmap kernel that makes a 12B q1t runnable.
16652    #[test]
16653    fn q1t_matvec_matches_reference() {
16654        use cortiq_core::quant::{dequant_q1t, f32_to_f16};
16655        let (rows, cols) = (3usize, 64usize); // gpr = 2
16656        let gpr = cols / GROUP_SIZE;
16657        let scales = [0.5f32, 0.3, 0.7, 0.2, 0.6, 0.15];
16658        // Overlay (must be sorted by flat index): a few spikes across rows.
16659        let outliers: [(u32, f32); 3] = [(5, 9.0), (70, -4.5), (150, 3.25)];
16660        let is_out = |flat: usize| outliers.iter().any(|&(i, _)| i as usize == flat);
16661        let mut bytes = Vec::new();
16662        for r in 0..rows {
16663            for g in 0..gpr {
16664                bytes.extend_from_slice(&f32_to_f16(scales[r * gpr + g]).to_le_bytes());
16665                let mut c = [0u8; 7];
16666                for k in 0..GROUP_SIZE {
16667                    // Encoder invariant: code 0 at outlier positions.
16668                    let code = if is_out(r * cols + g * GROUP_SIZE + k) {
16669                        0
16670                    } else {
16671                        ((k + r * 3 + g) % 3) as u8 // 0,1,2
16672                    };
16673                    cortiq_core::quant::q1t_pack(&mut c, k, code);
16674                }
16675                bytes.extend_from_slice(&c);
16676            }
16677        }
16678        // Per-row overlay: [u32 row_ptr[rows+1]] then [(u16 col, f16 val)] by
16679        // row (outliers are sorted by flat index → already grouped by row).
16680        let mut row_ptr = vec![0u32; rows + 1];
16681        for &(idx, _) in &outliers {
16682            row_ptr[idx as usize / cols + 1] += 1;
16683        }
16684        for r in 0..rows {
16685            row_ptr[r + 1] += row_ptr[r];
16686        }
16687        for &p in &row_ptr {
16688            bytes.extend_from_slice(&p.to_le_bytes());
16689        }
16690        for &(idx, v) in &outliers {
16691            bytes.extend_from_slice(&((idx as usize % cols) as u16).to_le_bytes());
16692            bytes.extend_from_slice(&f32_to_f16(v).to_le_bytes());
16693        }
16694
16695        let mut refw = vec![0f32; rows * cols];
16696        dequant_q1t(&bytes, rows, cols, &mut refw);
16697        // On-grid activations (±1, amax 1) so the int8 SDOT path reconstructs
16698        // x exactly and matches the f32 reference (same trick as the q1 test).
16699        let x: Vec<f32> = (0..cols)
16700            .map(|j| if j % 3 == 0 { 1.0 } else { -1.0 })
16701            .collect();
16702        let mut expect = vec![0f32; rows];
16703        for r in 0..rows {
16704            let mut a = 0.0f32;
16705            for j in 0..cols {
16706                a += refw[r * cols + j] * x[j];
16707            }
16708            expect[r] = a;
16709        }
16710        let tol = |e: f32| 1e-3 * e.abs().max(1e-3);
16711        let mut got = vec![0f32; rows];
16712        q1t_matvec(&bytes, &x, rows, cols, &mut got, None);
16713        for r in 0..rows {
16714            assert!(
16715                (got[r] - expect[r]).abs() < tol(expect[r]),
16716                "row {r}: {} vs {}",
16717                got[r],
16718                expect[r]
16719            );
16720        }
16721        // matmat (b=2, f32 decode path) must agree too.
16722        let x2: Vec<f32> = x.iter().chain(x.iter()).copied().collect();
16723        let mut gm = vec![0f32; 2 * rows];
16724        q1t_matmat(&bytes, &x2, 2, rows, cols, &mut gm, None);
16725        for r in 0..rows {
16726            assert!((gm[r] - expect[r]).abs() < tol(expect[r]));
16727            assert!((gm[rows + r] - expect[r]).abs() < tol(expect[r]));
16728        }
16729        // Fused pair (q1t_matvec2) must equal two single matvecs
16730        // bit-for-bit: same unpack, same group order, same f32
16731        // accumulation per stream. Distinct x2 exercises both lanes.
16732        let xb: Vec<f32> = (0..cols)
16733            .map(|j| if j % 5 == 0 { -1.0 } else { 1.0 })
16734            .collect();
16735        let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
16736        q1t_matvec(&bytes, &x, rows, cols, &mut s1, None);
16737        q1t_matvec(&bytes, &xb, rows, cols, &mut s2, None);
16738        let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
16739        q1t_matvec2(&bytes, &x, &xb, rows, cols, &mut p1, &mut p2, None);
16740        assert_eq!(p1, s1, "q1t pair lane 1 ≠ single matvec");
16741        assert_eq!(p2, s2, "q1t pair lane 2 ≠ single matvec");
16742    }
16743
16744    /// Pair == 2×matvec with an ODD group count (the kernel's tail
16745    /// group) and no overlay section.
16746    #[test]
16747    fn q1t_matvec2_odd_gpr_matches_singles() {
16748        use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_pack};
16749        let (rows, cols) = (5usize, 96usize); // gpr = 3 → paired + tail
16750        let gpr = cols / GROUP_SIZE;
16751        let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE);
16752        for r in 0..rows {
16753            for g in 0..gpr {
16754                bytes.extend_from_slice(&f32_to_f16(0.1 + 0.05 * (r + g) as f32).to_le_bytes());
16755                let mut c = [0u8; 7];
16756                for k in 0..GROUP_SIZE {
16757                    q1t_pack(&mut c, k, ((k * 7 + r * 5 + g * 3) % 3) as u8);
16758                }
16759                bytes.extend_from_slice(&c);
16760            }
16761        }
16762        let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.31).sin()).collect();
16763        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).cos()).collect();
16764        let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
16765        q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
16766        q1t_matvec(&bytes, &x2, rows, cols, &mut s2, None);
16767        let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
16768        q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
16769        assert_eq!(p1, s1, "odd-gpr pair lane 1 ≠ single");
16770        assert_eq!(p2, s2, "odd-gpr pair lane 2 ≠ single");
16771    }
16772
16773    // Speed A/B: fused pair (one unpack, two streams) vs two single
16774    // matvecs. Single-threaded, FFN-sized, min-of paired in-process.
16775    //   cargo test -p cortiq-engine --release q1t_matvec2_speed -- --ignored --nocapture
16776    #[test]
16777    #[ignore]
16778    fn q1t_matvec2_speed() {
16779        use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_pack};
16780        use std::time::Instant;
16781        let (rows, cols) = (8192usize, 4096usize);
16782        let gpr = cols / GROUP_SIZE;
16783        let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE);
16784        for r in 0..rows {
16785            for g in 0..gpr {
16786                let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
16787                bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
16788                let mut c = [0u8; 7];
16789                for k in 0..GROUP_SIZE {
16790                    q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
16791                }
16792                bytes.extend_from_slice(&c);
16793            }
16794        }
16795        let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.31).sin()).collect();
16796        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).cos()).collect();
16797        let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
16798        let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
16799        // Warm both paths once.
16800        q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
16801        q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
16802        let (mut t_pair, mut t_two) = (f64::MAX, f64::MAX);
16803        for _ in 0..8 {
16804            let t0 = Instant::now();
16805            q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
16806            t_pair = t_pair.min(t0.elapsed().as_secs_f64() * 1000.0);
16807            let t1 = Instant::now();
16808            q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
16809            q1t_matvec(&bytes, &x2, rows, cols, &mut s2, None);
16810            t_two = t_two.min(t1.elapsed().as_secs_f64() * 1000.0);
16811        }
16812        assert_eq!(p1, s1);
16813        assert_eq!(p2, s2);
16814        println!("q1t pair {rows}x{cols}: fused {t_pair:.2} ms | two singles {t_two:.2} ms");
16815    }
16816
16817    // Speed A/B: the base-3-division decode (what the packing commit left in
16818    // place) vs the fused sign-LUT matvec. Both single-threaded, same bytes.
16819    //   cargo test -p cortiq-engine q1t_matvec_speed -- --ignored --nocapture
16820    #[test]
16821    #[ignore]
16822    fn q1t_matvec_speed() {
16823        use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_code, q1t_pack};
16824        use std::time::Instant;
16825        let (rows, cols) = (8192usize, 4096usize); // FFN-sized
16826        let gpr = cols / GROUP_SIZE;
16827        let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE + 16);
16828        for r in 0..rows {
16829            for g in 0..gpr {
16830                let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
16831                bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
16832                let mut c = [0u8; 7];
16833                for k in 0..GROUP_SIZE {
16834                    q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
16835                }
16836                bytes.extend_from_slice(&c);
16837            }
16838        }
16839        let (n, stride) = (rows * cols, 40usize); // ~2.5% outliers, per-row overlay
16840        let mut row_ptr = vec![0u32; rows + 1];
16841        let mut idx = 0usize;
16842        while idx < n {
16843            row_ptr[idx / cols + 1] += 1;
16844            idx += stride;
16845        }
16846        for r in 0..rows {
16847            row_ptr[r + 1] += row_ptr[r];
16848        }
16849        for &p in &row_ptr {
16850            bytes.extend_from_slice(&p.to_le_bytes());
16851        }
16852        let mut idx = 0usize;
16853        while idx < n {
16854            bytes.extend_from_slice(&((idx % cols) as u16).to_le_bytes());
16855            bytes.extend_from_slice(&f32_to_f16((idx % 13) as f32 * 0.1 - 0.6).to_le_bytes());
16856            idx += stride;
16857        }
16858        // On-grid ±1 so the fast path's int8 SDOT is exact vs the f32 "slow"
16859        // reference (the A/B is a timing check; values must still agree).
16860        let x: Vec<f32> = (0..cols)
16861            .map(|j| if j % 3 == 0 { 1.0 } else { -1.0 })
16862            .collect();
16863        let (rp_off, ent_off, has_ov) = q1t_overlay(&bytes, rows * gpr * Q1T_TILE, rows);
16864
16865        // "before": base-3 division decode into a buffer, then dot.
16866        let slow = |out: &mut [f32]| {
16867            let mut buf = vec![0f32; cols];
16868            for r in 0..rows {
16869                for g in 0..gpr {
16870                    let off = (r * gpr + g) * Q1T_TILE;
16871                    let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
16872                    let codes = &bytes[off + 2..off + Q1T_TILE];
16873                    for k in 0..GROUP_SIZE {
16874                        buf[g * GROUP_SIZE + k] = match q1t_code(codes, k) {
16875                            1 => s,
16876                            2 => -s,
16877                            _ => 0.0,
16878                        };
16879                    }
16880                }
16881                out[r] = q1t_row_outlier_correction(&bytes, r, rp_off, ent_off, has_ov, &x)
16882                    + (0..cols).map(|j| buf[j] * x[j]).sum::<f32>();
16883            }
16884        };
16885        let iters = 5;
16886        let mut a = vec![0f32; rows];
16887        slow(&mut a); // warm
16888        let t = Instant::now();
16889        for _ in 0..iters {
16890            slow(&mut a);
16891        }
16892        let slow_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
16893
16894        let mut b = vec![0f32; rows];
16895        q1t_matvec(&bytes, &x, rows, cols, &mut b, None); // warm
16896        let t = Instant::now();
16897        for _ in 0..iters {
16898            q1t_matvec(&bytes, &x, rows, cols, &mut b, None);
16899        }
16900        let fast_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
16901
16902        for r in 0..rows {
16903            assert!((a[r] - b[r]).abs() < 1e-2, "mismatch row {r}");
16904        }
16905        println!(
16906            "q1t matvec {rows}x{cols} (1 thread): div-decode {slow_ms:.2} ms  fused-LUT {fast_ms:.2} ms  => {:.2}x",
16907            slow_ms / fast_ms
16908        );
16909    }
16910}
16911
16912#[cfg(test)]
16913mod row_span_tests {
16914    use cortiq_core::TensorDtype;
16915
16916    /// The three ranges `row_span_offsets` names are exactly what decoding a
16917    /// q4tp row reads: with every other byte of the tensor overwritten the
16918    /// row's nibbles and scales come back unchanged.
16919    #[test]
16920    fn q4tp_row_spans_cover_the_row() {
16921        let (rows, cols) = (37usize, 160usize);
16922        let gpr = cols / 32;
16923        let total = cortiq_core::quant::expected_nbytes(TensorDtype::Q4TiledP, &[rows, cols]).unwrap();
16924        let (params_off, _, _) = cortiq_core::quant::q4tp_sections(rows, cols);
16925        let mut bytes: Vec<u8> = (0..total).map(|i| (i * 61 % 251) as u8).collect();
16926        let lo = cortiq_core::quant::f32_to_f16(-4.0);
16927        let step = cortiq_core::quant::f32_to_f16(0.1);
16928        for r in 0..rows {
16929            let o = params_off + r * 4;
16930            bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
16931            bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
16932        }
16933        for r in [0usize, 1, 17, rows - 1] {
16934            let spans = super::row_span_offsets(TensorDtype::Q4TiledP, rows, cols, r).unwrap();
16935            assert_eq!(spans.iter().map(|s| s.1).sum::<usize>(), total / rows);
16936            let mut keep = vec![false; total];
16937            for &(o, n) in &spans {
16938                keep[o..o + n].fill(true);
16939            }
16940            let other: Vec<u8> =
16941                bytes.iter().zip(&keep).map(|(&b, &k)| if k { b } else { 0x5A }).collect();
16942            let (a, b) = (super::Q4tpView::new(&bytes, rows, cols), super::Q4tpView::new(&other, rows, cols));
16943            let (mut sa, mut sb) = (vec![0f32; gpr], vec![0f32; gpr]);
16944            a.scales_into(r, gpr, &mut sa);
16945            b.scales_into(r, gpr, &mut sb);
16946            assert_eq!(sa, sb, "row {r} scales");
16947            let nib = r * gpr * super::Q4TP_NIB..(r + 1) * gpr * super::Q4TP_NIB;
16948            assert_eq!(a.nib[nib.clone()], b.nib[nib], "row {r} nibbles");
16949        }
16950        assert!(super::row_span_offsets(TensorDtype::Q4TiledP, rows, cols, rows).is_none());
16951        assert!(super::row_span_offsets(TensorDtype::Q8_2f, rows, cols, 0).is_none());
16952        let t = super::row_span_offsets(TensorDtype::Q4Tiled, rows, cols, 3).unwrap();
16953        assert_eq!(t[0], (3 * gpr * super::Q4_TILE, gpr * super::Q4_TILE));
16954    }
16955}
16956
16957#[cfg(test)]
16958mod gemm_bench {
16959    /// `cargo test -p cortiq-engine --release q4tp_matmat_throughput -- --ignored --nocapture`
16960    /// Times the batched q4tp GEMM at the shapes the image DiT runs
16961    /// (b=296 tokens, 2304 -> 9216), on synthetic bytes: no model, no
16962    /// mmap, no thermal drift over minutes — a kernel change shows up
16963    /// here in seconds where a full render hides it in noise.
16964    ///
16965    /// On macOS add `CMF_ACCEL=0`: this shape is over the 500k-cell mark
16966    /// where the matmat hands off to Accelerate's dequant sgemm, and
16967    /// without the opt-out both rows below measure the AMX, not the
16968    /// kernel under test.
16969    #[test]
16970    #[ignore]
16971    fn q4tp_matmat_throughput() {
16972        let _alt = super::Q4TP_ALT_TEST_LOCK
16973            .lock()
16974            .unwrap_or_else(|e| e.into_inner());
16975        // 296 is a prompt-encode batch; the image DiT runs 2085 at
16976        // 512x512, where the activation panel stops fitting L2 and the
16977        // loop's shape starts to matter more than its instructions.
16978        let b: usize = std::env::var("CMF_BENCH_B")
16979            .ok()
16980            .and_then(|v| v.parse().ok())
16981            .unwrap_or(296);
16982        let (rows, cols) = (9216usize, 2304usize);
16983        let (_, _, _) = (rows, cols, b);
16984        let total =
16985            cortiq_core::quant::expected_nbytes(cortiq_core::TensorDtype::Q4TiledP, &[rows, cols])
16986                .unwrap();
16987        // Random nibbles are fine, but the row params are f16 (lo, step)
16988        // of a geometric ladder: garbage there gives exp2 of a huge
16989        // exponent, the scales come back inf, and the whole bench times
16990        // NaN arithmetic instead of the kernel.
16991        let (params_off, codes_off, _) = cortiq_core::quant::q4tp_sections(rows, cols);
16992        let mut bytes: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
16993        let lo = cortiq_core::quant::f32_to_f16(-4.0);
16994        let step = cortiq_core::quant::f32_to_f16(0.1);
16995        for r in 0..rows {
16996            let o = params_off + r * 4;
16997            bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
16998            bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
16999        }
17000        let _ = codes_off;
17001        let xs: Vec<f32> = (0..b * cols)
17002            .map(|i| ((i % 97) as f32 - 48.0) / 48.0)
17003            .collect();
17004        let mut out = vec![0f32; b * rows];
17005        let pool = crate::pool::Pool::from_env();
17006        // A shared 48-core stand drifts ±25% run to run, which is wider
17007        // than any kernel change worth making. So: alternate the two
17008        // kernels inside one process and keep the BEST time for
17009        // each. Interleaving makes both see the same interference, and a
17010        // minimum is the one statistic another tenant cannot inflate.
17011        super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
17012        let reps: usize = std::env::var("CMF_BENCH_REPS")
17013            .ok()
17014            .and_then(|v| v.parse().ok())
17015            .unwrap_or(10);
17016        let mut best = [f64::MAX; 2];
17017        let mut sums = [0f32; 2];
17018        for _ in 0..reps {
17019            for (k, w) in [(0usize, 1u8), (1usize, 2u8)] {
17020                super::Q4TP_ALT.store(w, std::sync::atomic::Ordering::Relaxed);
17021                let t = std::time::Instant::now();
17022                super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
17023                best[k] = best[k].min(t.elapsed().as_secs_f64());
17024                sums[k] = out.iter().take(64).sum::<f32>();
17025            }
17026        }
17027        let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
17028        for (k, name) in ["previous", "tuned   "].iter().enumerate() {
17029            println!(
17030                "q4tp matmat {rows}x{cols} b={b} {name}: {:.1} ms  {:.1} GFLOP/s  (checksum {:.3})",
17031                best[k] * 1e3,
17032                flops / best[k] / 1e9,
17033                sums[k]
17034            );
17035        }
17036        assert!(
17037            (sums[0] - sums[1]).abs() < 1e-2,
17038            "the tuned kernel changed the result: {} vs {}",
17039            sums[0],
17040            sums[1]
17041        );
17042    }
17043
17044    /// The blocked kernel must agree with the per-column path exactly —
17045    /// same weights, same activation split, only a different instruction
17046    /// mix. Shapes are chosen to hit the awkward cases: a column count
17047    /// that leaves an odd group (the 512-bit kernel does two at a time),
17048    /// and a batch that does not divide by four.
17049    #[test]
17050    fn q4tp_matmat_blocked_matches_scalar() {
17051        use std::sync::atomic::Ordering::Relaxed;
17052        let _alt = super::Q4TP_ALT_TEST_LOCK
17053            .lock()
17054            .unwrap_or_else(|e| e.into_inner());
17055        // The last shape carries the image DiT's column count — 2304, so
17056        // 72 groups of accumulation, which is where a reordered sum can
17057        // actually drift — and runs through the thread pool, since the
17058        // blocked path splits rows across workers. Its row count stays
17059        // under 500k cells on purpose: above that, macOS diverts the whole
17060        // matmat to the Accelerate/AMX dequant sgemm and neither kernel
17061        // here would run.
17062        for &(rows, cols, b) in &[
17063            (64usize, 128usize, 7usize),
17064            (33, 96, 4),
17065            (16, 256, 9),
17066            (192, 2304, 37),
17067        ] {
17068            let total = cortiq_core::quant::expected_nbytes(
17069                cortiq_core::TensorDtype::Q4TiledP,
17070                &[rows, cols],
17071            )
17072            .unwrap();
17073            let (params_off, _, _) = cortiq_core::quant::q4tp_sections(rows, cols);
17074            let mut bytes: Vec<u8> = (0..total).map(|i| (i * 61 % 251) as u8).collect();
17075            let lo = cortiq_core::quant::f32_to_f16(-4.0);
17076            let step = cortiq_core::quant::f32_to_f16(0.1);
17077            for r in 0..rows {
17078                let o = params_off + r * 4;
17079                bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
17080                bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
17081            }
17082            let xs: Vec<f32> = (0..b * cols)
17083                .map(|i| ((i % 89) as f32 - 44.0) / 44.0)
17084                .collect();
17085            let mut got = vec![0f32; b * rows];
17086            let mut want = vec![0f32; b * rows];
17087            let gpr = cols / 32;
17088            let view = super::Q4tpView::new(&bytes, rows, cols);
17089            let pool = crate::pool::Pool::from_env();
17090            super::Q4TP_ALT.store(2, Relaxed);
17091            super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut got, pool.as_deref());
17092            super::Q4TP_ALT.store(1, Relaxed);
17093            super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut want, pool.as_deref());
17094            super::Q4TP_ALT.store(0, Relaxed);
17095            // Measured against the output's scale, not cell by cell: a
17096            // dot product of 2304 terms lands near zero wherever the row
17097            // and the activation nearly cancel, and there a per-cell
17098            // ratio reports 1e-3 for an absolute error of 5e-6 — f32's
17099            // own rounding, reordered. What must stay small is the error
17100            // relative to what the layer actually outputs.
17101            let scale = want.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
17102            let (mut worst, mut at) = (0f32, 0usize);
17103            for (i, (g, w)) in got.iter().zip(&want).enumerate() {
17104                if (g - w).abs() > worst {
17105                    worst = (g - w).abs();
17106                    at = i;
17107                }
17108            }
17109            assert!(
17110                worst <= 1e-4 * scale,
17111                "{rows}x{cols} b={b}: blocked and scalar disagree by {worst:.3e} \
17112                 (scale {scale:.3e}) at cell {at}: {} vs {}",
17113                got[at],
17114                want[at]
17115            );
17116
17117            // "Same speed, no quality loss" is a claim about which answer
17118            // is RIGHT, not about which two agree. Both paths sum the same
17119            // 2304 products in different orders, so f64 decides: the
17120            // blocked kernel keeps sixteen partial sums and folds them at
17121            // the end, which is a shallower addition tree than the
17122            // per-column path's running scalar, and it must not be worse.
17123            let (mut e_blocked, mut e_scalar) = (0f64, 0f64);
17124            for bi in 0..b {
17125                let act = super::split_act(&xs[bi * cols..(bi + 1) * cols]);
17126                for r in 0..rows {
17127                    let mut sc = vec![0f32; gpr];
17128                    view.scales_into(r, gpr, &mut sc);
17129                    let mut exact = 0f64;
17130                    for j in 0..cols {
17131                        let (w, sq) = super::q4tp_outlier(view.nib, r, gpr, j, &sc);
17132                        exact += w as f64 * sq as f64 * act.xq[j] as f64;
17133                    }
17134                    exact *= act.sx as f64;
17135                    for &(j, xv) in &act.outliers {
17136                        let (w, sq) = super::q4tp_outlier(view.nib, r, gpr, j, &sc);
17137                        exact += w as f64 * sq as f64 * xv as f64;
17138                    }
17139                    let i = bi * rows + r;
17140                    e_blocked = e_blocked.max((got[i] as f64 - exact).abs());
17141                    e_scalar = e_scalar.max((want[i] as f64 - exact).abs());
17142                }
17143            }
17144            println!(
17145                "{rows}x{cols} b={b}: worst error vs f64 — blocked {e_blocked:.3e}, \
17146                 per-column {e_scalar:.3e}"
17147            );
17148            // An absolute bar, not a race between the two: at these
17149            // magnitudes both sit in f32's last bits, and on a small shape
17150            // whichever one happens to round the unluckiest cell "wins" by
17151            // a factor the next seed reverses.
17152            assert!(
17153                e_blocked <= 1e-5 * scale as f64 && e_scalar <= 1e-5 * scale as f64,
17154                "{rows}x{cols} b={b}: error against f64 too large — blocked \
17155                 {e_blocked:.3e}, per-column {e_scalar:.3e}, scale {scale:.3e}"
17156            );
17157        }
17158    }
17159
17160    /// What the row-exact contract costs a speculative-verify panel on the
17161    /// host: the same batch through the fast arms, through the row-exact
17162    /// arms, and as one matvec per token (the other way to be exact).
17163    /// Arms alternate inside one process and keep their best time.
17164    /// `CMF_GPU=0 cargo test -p cortiq-engine --release q4tp_matmat_row_exact_cost -- --ignored --nocapture`
17165    #[test]
17166    #[ignore]
17167    fn q4tp_matmat_row_exact_cost() {
17168        let _alt = super::Q4TP_ALT_TEST_LOCK
17169            .lock()
17170            .unwrap_or_else(|e| e.into_inner());
17171        let pool = crate::pool::Pool::from_env();
17172        let reps: usize = std::env::var("CMF_BENCH_REPS")
17173            .ok()
17174            .and_then(|v| v.parse().ok())
17175            .unwrap_or(30);
17176        for &(rows, cols) in &[(2048usize, 4096usize), (4096, 2048), (4096, 4096)] {
17177            let total = cortiq_core::quant::expected_nbytes(
17178                cortiq_core::TensorDtype::Q4TiledP,
17179                &[rows, cols],
17180            )
17181            .unwrap();
17182            let (params_off, _, _) = cortiq_core::quant::q4tp_sections(rows, cols);
17183            let mut bytes: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
17184            let lo = cortiq_core::quant::f32_to_f16(-4.0);
17185            let step = cortiq_core::quant::f32_to_f16(0.1);
17186            for r in 0..rows {
17187                let o = params_off + r * 4;
17188                bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
17189                bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
17190            }
17191            for &b in &[2usize, 4, 5, 8] {
17192                let xs: Vec<f32> = (0..b * cols)
17193                    .map(|i| ((i % 97) as f32 - 48.0) / 48.0)
17194                    .collect();
17195                let mut out = vec![0f32; b * rows];
17196                let mut best = [f64::MAX; 3];
17197                for _ in 0..reps {
17198                    for (k, best_k) in best.iter_mut().enumerate() {
17199                        let t = std::time::Instant::now();
17200                        match k {
17201                            0 | 1 => super::q4tp_matmat_with(
17202                                &bytes,
17203                                &xs,
17204                                b,
17205                                rows,
17206                                cols,
17207                                &mut out,
17208                                pool.as_deref(),
17209                                k == 1,
17210                            ),
17211                            _ => {
17212                                for (bi, o) in out.chunks_mut(rows).enumerate() {
17213                                    super::q4tp_matvec(
17214                                        &bytes,
17215                                        &xs[bi * cols..(bi + 1) * cols],
17216                                        rows,
17217                                        cols,
17218                                        o,
17219                                        pool.as_deref(),
17220                                    );
17221                                }
17222                            }
17223                        }
17224                        *best_k = best_k.min(t.elapsed().as_secs_f64());
17225                    }
17226                }
17227                println!(
17228                    "q4tp {rows}x{cols} b={b}: fast {:.3} ms, row-exact {:.3} ms, \
17229                     {b} matvecs {:.3} ms",
17230                    best[0] * 1e3,
17231                    best[1] * 1e3,
17232                    best[2] * 1e3
17233                );
17234            }
17235        }
17236    }
17237
17238    /// The q4t twin of the throughput bench, same shape and rules, so the
17239    /// two quantisations' batch kernels can be read against each other.
17240    /// `cargo test -p cortiq-engine --release q4t_matmat_throughput -- --ignored --nocapture`
17241    #[test]
17242    #[ignore]
17243    fn q4t_matmat_throughput() {
17244        let (rows, cols, b) = (9216usize, 2304usize, 296usize);
17245        let total =
17246            cortiq_core::quant::expected_nbytes(cortiq_core::TensorDtype::Q4Tiled, &[rows, cols])
17247                .unwrap();
17248        // q4t carries a per-group f16 scale in the tile's first two bytes;
17249        // random bytes there decode to inf and the bench would time NaNs.
17250        let mut bytes: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
17251        let sc = cortiq_core::quant::f32_to_f16(0.02);
17252        for t in bytes.chunks_mut(super::Q4_TILE) {
17253            t[..2].copy_from_slice(&sc.to_le_bytes());
17254        }
17255        let xs: Vec<f32> = (0..b * cols)
17256            .map(|i| ((i % 97) as f32 - 48.0) / 48.0)
17257            .collect();
17258        let mut out = vec![0f32; b * rows];
17259        let pool = crate::pool::Pool::from_env();
17260        super::q4t_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
17261        let reps: usize = std::env::var("CMF_BENCH_REPS")
17262            .ok()
17263            .and_then(|v| v.parse().ok())
17264            .unwrap_or(10);
17265        let mut best = f64::MAX;
17266        for _ in 0..reps {
17267            let t = std::time::Instant::now();
17268            super::q4t_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
17269            best = best.min(t.elapsed().as_secs_f64());
17270        }
17271        let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
17272        println!(
17273            "q4t matmat {rows}x{cols} b={b}: {:.1} ms  {:.1} GFLOP/s  (checksum {:.3})",
17274            best * 1e3,
17275            flops / best / 1e9,
17276            out.iter().take(64).sum::<f32>()
17277        );
17278    }
17279}
17280
17281#[cfg(all(test, target_arch = "x86_64"))]
17282mod x86_gemm_tests {
17283    use super::x86_gemm_rm;
17284
17285    /// The AVX2 two-row FMA q4tp batch kernel against the row-exact
17286    /// vertical cell (the matvec's bits): odd row counts (the last row
17287    /// takes the vertical cell), columns past a group of four, a wide
17288    /// row of 2560 columns (80 groups), outliers in the activations, and
17289    /// the pool. Measured against the output's scale, as the blocked
17290    /// kernel's test does.
17291    #[test]
17292    fn q4tp_fma_batch_matches_the_row_exact_cell() {
17293        if !super::q4tp_fma_x86() || !super::a8w8_enabled() {
17294            return;
17295        }
17296        let pool = crate::pool::Pool::from_env();
17297        for &(rows, cols, b) in &[(33usize, 96usize, 7usize), (64, 256, 13), (257, 2560, 37)] {
17298            let total = cortiq_core::quant::expected_nbytes(
17299                cortiq_core::TensorDtype::Q4TiledP,
17300                &[rows, cols],
17301            )
17302            .unwrap();
17303            let (params_off, _, _) = cortiq_core::quant::q4tp_sections(rows, cols);
17304            let mut bytes: Vec<u8> = (0..total).map(|i| (i * 61 % 251) as u8).collect();
17305            let lo = cortiq_core::quant::f32_to_f16(-4.0);
17306            let step = cortiq_core::quant::f32_to_f16(0.1);
17307            for r in 0..rows {
17308                let o = params_off + r * 4;
17309                bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
17310                bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
17311            }
17312            let mut xs: Vec<f32> = (0..b * cols)
17313                .map(|i| ((i * 37 % 89) as f32 - 44.0) / 44.0)
17314                .collect();
17315            for bi in (0..b).step_by(3) {
17316                xs[bi * cols + (bi * 7) % cols] = 40.0; // an outlier column
17317            }
17318            let mut got = vec![0f32; b * rows];
17319            let mut want = vec![0f32; b * rows];
17320            super::q4tp_matmat_with(&bytes, &xs, b, rows, cols, &mut got, pool.as_deref(), false);
17321            super::q4tp_matmat_with(&bytes, &xs, b, rows, cols, &mut want, pool.as_deref(), true);
17322            let scale = want.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
17323            let worst = got.iter().zip(&want).map(|(g, w)| (g - w).abs()).fold(0f32, f32::max);
17324            assert!(
17325                worst <= 1e-5 * scale,
17326                "{rows}x{cols} b={b}: FMA batch and row-exact cell disagree by {worst:.3e} (scale {scale:.3e})"
17327            );
17328        }
17329    }
17330
17331    #[allow(clippy::too_many_arguments)]
17332    fn naive(
17333        m: usize,
17334        n: usize,
17335        k: usize,
17336        alpha: f32,
17337        a: &[f32],
17338        lda: usize,
17339        b: &[f32],
17340        ldb: usize,
17341        b_rows_are_n: bool,
17342    ) -> Vec<f64> {
17343        let mut c = vec![0f64; m * n];
17344        for i in 0..m {
17345            for j in 0..n {
17346                let mut acc = 0f64;
17347                for p in 0..k {
17348                    let bv = if b_rows_are_n { b[j * ldb + p] } else { b[p * ldb + j] };
17349                    acc += a[i * lda + p] as f64 * bv as f64;
17350                }
17351                c[i * n + j] = acc * alpha as f64;
17352            }
17353        }
17354        c
17355    }
17356
17357    /// The AVX2 micro-GEMM against a double-precision triple loop: both
17358    /// B layouts, padded leading dimensions, shapes that exercise the
17359    /// 4×16 kernel, the row and column tails, and the k blocking (k >
17360    /// 256 resumes the accumulators from C).
17361    #[test]
17362    fn x86_gemm_matches_naive_on_odd_shapes() {
17363        for &(m, n, k) in &[
17364            (1usize, 1usize, 1usize),
17365            (4, 16, 256),
17366            (7, 37, 19),
17367            (9, 300, 257),
17368            (64, 1000, 256),
17369            (5, 256, 1001),
17370            (13, 129, 512),
17371            (8, 16, 0),
17372        ] {
17373            for &bt in &[false, true] {
17374                let (lda, ldb, ldc) = (k + 3, if bt { k + 2 } else { n + 5 }, n + 1);
17375                let a: Vec<f32> = (0..m * lda)
17376                    .map(|i| ((i * 7919) % 1000) as f32 / 500.0 - 1.0)
17377                    .collect();
17378                let brows = if bt { n } else { k.max(1) };
17379                let b: Vec<f32> = (0..brows * ldb)
17380                    .map(|i| ((i * 104_729) % 997) as f32 / 498.5 - 1.0)
17381                    .collect();
17382                let mut c = vec![f32::NAN; m * ldc];
17383                x86_gemm_rm(m, n, k, 0.37, &a, lda, &b, ldb, bt, &mut c, ldc);
17384                let r = naive(m, n, k, 0.37, &a, lda, &b, ldb, bt);
17385                for i in 0..m {
17386                    for j in 0..n {
17387                        let (x, y) = (c[i * ldc + j] as f64, r[i * n + j]);
17388                        assert!(
17389                            (x - y).abs() <= 1e-3 * (1.0 + y.abs()),
17390                            "m{m} n{n} k{k} bt{bt} [{i},{j}]: {x} vs {y}"
17391                        );
17392                    }
17393                    assert!(c[i * ldc + n].is_nan(), "m{m} n{n} k{k}: padding column written");
17394                }
17395            }
17396        }
17397    }
17398}