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, so a value differs from
4450    /// theirs by rounding only; the work is one pass over the weights at
4451    /// the speed of the memory bus (see `qtensor/moe_fast.rs`). False =
4452    /// not covered (no AVX-512 VNNI, not q4tp, float activations, a shape
4453    /// outside 64-wide steps): the caller takes the two row kernels.
4454    pub fn moe_expert_rows_fast(
4455        experts: &[(&QTensor, &QTensor, &QTensor)],
4456        groups: &[Vec<usize>],
4457        xs: &[f32],
4458        out: &mut [f32],
4459        pool: Option<&Pool>,
4460    ) -> bool {
4461        #[cfg(target_arch = "x86_64")]
4462        {
4463            moe_fast::run(experts, groups, xs, out, pool)
4464        }
4465        #[cfg(not(target_arch = "x86_64"))]
4466        {
4467            let _ = (experts, groups, xs, out, pool);
4468            false
4469        }
4470    }
4471
4472    /// Whether this host has the cores for [`QTensor::moe_expert_rows_fast`] (AVX-512 VNNI with the int8
4473    /// activation path on).
4474    pub fn moe_expert_rows_fast_available() -> bool {
4475        #[cfg(target_arch = "x86_64")]
4476        {
4477            avx512vnni_enabled() && a8w8_enabled()
4478        }
4479        #[cfg(not(target_arch = "x86_64"))]
4480        {
4481            false
4482        }
4483    }
4484}
4485
4486/// Batched q8 kernel: same math as qmatvec, the row makes a single
4487/// pass from memory for the whole batch.
4488/// Accelerate CBLAS — the Apple AMX matrix units, the same engine
4489/// llama.cpp's `-ngl 0` prefill rides via ggml-blas.
4490#[cfg(target_os = "macos")]
4491mod accel_blas {
4492    #[link(name = "Accelerate", kind = "framework")]
4493    unsafe extern "C" {
4494        pub fn cblas_sgemm(
4495            order: i32,
4496            trans_a: i32,
4497            trans_b: i32,
4498            m: i32,
4499            n: i32,
4500            k: i32,
4501            alpha: f32,
4502            a: *const f32,
4503            lda: i32,
4504            b: *const f32,
4505            ldb: i32,
4506            beta: f32,
4507            c: *mut f32,
4508            ldc: i32,
4509        );
4510    }
4511}
4512
4513#[cfg(target_os = "macos")]
4514pub(crate) fn accel_gemm_enabled() -> bool {
4515    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4516    *ON.get_or_init(|| std::env::var("CMF_ACCEL").map(|v| v != "0").unwrap_or(true))
4517}
4518
4519/// Off macOS the "accel" GEMM is the portable NEON micro-kernel below —
4520/// same entry point, so the batched-attention path opens on mobile.
4521#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
4522pub(crate) fn accel_gemm_enabled() -> bool {
4523    true
4524}
4525
4526/// On x86-64 the "accel" GEMM is the AVX2 micro-kernel below, so the
4527/// batched attention opens there too; `CMF_CPU_ATTEND=0` keeps the
4528/// per-position attend, and a CPU without AVX2/FMA never leaves it.
4529#[cfg(all(target_arch = "x86_64", not(target_os = "macos")))]
4530pub(crate) fn accel_gemm_enabled() -> bool {
4531    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4532    *ON.get_or_init(|| {
4533        std::is_x86_feature_detected!("avx2")
4534            && std::is_x86_feature_detected!("fma")
4535            && std::env::var("CMF_CPU_ATTEND").map(|v| v != "0").unwrap_or(true)
4536    })
4537}
4538
4539/// Any other architecture has no CPU micro-GEMM: the per-position attend.
4540#[cfg(not(any(target_os = "macos", target_arch = "aarch64", target_arch = "x86_64")))]
4541#[allow(dead_code)]
4542pub(crate) fn accel_gemm_enabled() -> bool {
4543    false
4544}
4545
4546/// Portable NEON f32 GEMM (row-major, optional Bᵀ): a 4×8 fmla
4547/// micro-kernel with A broadcast against B panels — the mobile stand-in
4548/// for Accelerate in the batched causal attention (QKᵀ and P·V). Not a
4549/// BLAS: shapes here are the attention panels (m ≤ heads·chunk,
4550/// k = head_dim or context), and the goal is removing the per-position
4551/// quadratic wall, not peak GEMM.
4552#[cfg(target_arch = "aarch64")]
4553#[allow(clippy::too_many_arguments)]
4554pub(crate) fn neon_gemm_rm(
4555    m: usize,
4556    n: usize,
4557    k: usize,
4558    alpha: f32,
4559    a: &[f32],
4560    lda: usize,
4561    b_mat: &[f32],
4562    ldb: usize,
4563    b_rows_are_n: bool,
4564    c: &mut [f32],
4565    ldc: usize,
4566) {
4567    debug_assert!(a.len() >= (m - 1) * lda + k);
4568    debug_assert!(c.len() >= (m - 1) * ldc + n);
4569    // SAFETY: bounds asserted above; NEON is baseline on aarch64.
4570    unsafe {
4571        use core::arch::aarch64::*;
4572        let mut i = 0usize;
4573        while i < m {
4574            let mi = (m - i).min(4);
4575            let mut j = 0usize;
4576            while j < n {
4577                let nj = (n - j).min(8);
4578                if mi == 4 && nj == 8 {
4579                    let (mut c0a, mut c0b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
4580                    let (mut c1a, mut c1b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
4581                    let (mut c2a, mut c2b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
4582                    let (mut c3a, mut c3b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
4583                    for p in 0..k {
4584                        let (b0, b1) = if b_rows_are_n {
4585                            // B is [n, k]: column p of Bᵀ = element p of
4586                            // eight consecutive B rows — gathered.
4587                            let base = b_mat.as_ptr().add(j * ldb + p);
4588                            let g = |o: usize| *base.add(o * ldb);
4589                            ([g(0), g(1), g(2), g(3)], [g(4), g(5), g(6), g(7)])
4590                        } else {
4591                            let base = b_mat.as_ptr().add(p * ldb + j);
4592                            (
4593                                [*base, *base.add(1), *base.add(2), *base.add(3)],
4594                                [*base.add(4), *base.add(5), *base.add(6), *base.add(7)],
4595                            )
4596                        };
4597                        let bv0 = vld1q_f32(b0.as_ptr());
4598                        let bv1 = vld1q_f32(b1.as_ptr());
4599                        let a0 = vdupq_n_f32(*a.as_ptr().add(i * lda + p));
4600                        let a1 = vdupq_n_f32(*a.as_ptr().add((i + 1) * lda + p));
4601                        let a2 = vdupq_n_f32(*a.as_ptr().add((i + 2) * lda + p));
4602                        let a3 = vdupq_n_f32(*a.as_ptr().add((i + 3) * lda + p));
4603                        c0a = vfmaq_f32(c0a, a0, bv0);
4604                        c0b = vfmaq_f32(c0b, a0, bv1);
4605                        c1a = vfmaq_f32(c1a, a1, bv0);
4606                        c1b = vfmaq_f32(c1b, a1, bv1);
4607                        c2a = vfmaq_f32(c2a, a2, bv0);
4608                        c2b = vfmaq_f32(c2b, a2, bv1);
4609                        c3a = vfmaq_f32(c3a, a3, bv0);
4610                        c3b = vfmaq_f32(c3b, a3, bv1);
4611                    }
4612                    let al = vdupq_n_f32(alpha);
4613                    for (r, (ca, cb)) in [(c0a, c0b), (c1a, c1b), (c2a, c2b), (c3a, c3b)]
4614                        .iter()
4615                        .enumerate()
4616                    {
4617                        let dst = c.as_mut_ptr().add((i + r) * ldc + j);
4618                        vst1q_f32(dst, vmulq_f32(*ca, al));
4619                        vst1q_f32(dst.add(4), vmulq_f32(*cb, al));
4620                    }
4621                } else {
4622                    for r in 0..mi {
4623                        for q in 0..nj {
4624                            let mut acc = 0f32;
4625                            for p in 0..k {
4626                                let bv = if b_rows_are_n {
4627                                    b_mat[(j + q) * ldb + p]
4628                                } else {
4629                                    b_mat[p * ldb + j + q]
4630                                };
4631                                acc += a[(i + r) * lda + p] * bv;
4632                            }
4633                            c[(i + r) * ldc + j + q] = acc * alpha;
4634                        }
4635                    }
4636                }
4637                j += nj;
4638            }
4639            i += mi;
4640        }
4641    }
4642}
4643
4644/// Off-macOS aarch64: the batched attention rides the NEON micro-GEMM.
4645#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
4646#[allow(clippy::too_many_arguments)]
4647pub(crate) fn sgemm_rm(
4648    m: usize,
4649    n: usize,
4650    k: usize,
4651    alpha: f32,
4652    a: &[f32],
4653    lda: usize,
4654    b_mat: &[f32],
4655    ldb: usize,
4656    b_rows_are_n: bool,
4657    c: &mut [f32],
4658    ldc: usize,
4659) {
4660    neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
4661}
4662
4663/// The portable micro-GEMM of the batched causal attention: NEON on
4664/// aarch64, AVX2 on x86-64. Same contract as `neon_gemm_rm`.
4665#[cfg(target_arch = "aarch64")]
4666#[allow(clippy::too_many_arguments)]
4667pub(crate) fn micro_gemm_rm(
4668    m: usize,
4669    n: usize,
4670    k: usize,
4671    alpha: f32,
4672    a: &[f32],
4673    lda: usize,
4674    b_mat: &[f32],
4675    ldb: usize,
4676    b_rows_are_n: bool,
4677    c: &mut [f32],
4678    ldc: usize,
4679) {
4680    neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
4681}
4682
4683/// The portable micro-GEMM of the batched causal attention (x86-64 arm).
4684#[cfg(target_arch = "x86_64")]
4685#[allow(clippy::too_many_arguments)]
4686pub(crate) fn micro_gemm_rm(
4687    m: usize,
4688    n: usize,
4689    k: usize,
4690    alpha: f32,
4691    a: &[f32],
4692    lda: usize,
4693    b_mat: &[f32],
4694    ldb: usize,
4695    b_rows_are_n: bool,
4696    c: &mut [f32],
4697    ldc: usize,
4698) {
4699    x86_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
4700}
4701
4702/// Off-macOS x86-64: the batched attention rides the AVX2 micro-GEMM.
4703#[cfg(all(target_arch = "x86_64", not(target_os = "macos")))]
4704#[allow(clippy::too_many_arguments)]
4705pub(crate) fn sgemm_rm(
4706    m: usize,
4707    n: usize,
4708    k: usize,
4709    alpha: f32,
4710    a: &[f32],
4711    lda: usize,
4712    b_mat: &[f32],
4713    ldb: usize,
4714    b_rows_are_n: bool,
4715    c: &mut [f32],
4716    ldc: usize,
4717) {
4718    x86_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
4719}
4720
4721/// Portable x86-64 f32 GEMM (row-major, optional Bᵀ): the x86 twin of
4722/// `neon_gemm_rm` for the batched causal attention (Q·Kᵀ and P·V). An
4723/// AVX2+FMA 4×16 micro-kernel with A broadcast against B panels, blocked
4724/// so a 128-column × 256-row slab of B stays in L2 while the rows of A
4725/// stream past it (the P·V multiply has k = context, so without the k
4726/// blocking V would be re-read once per four rows). Same contract:
4727/// C[m,n] = alpha·A[m,k] × B, B stored [k,n] — or [n,k] with
4728/// `b_rows_are_n`, which packs Bᵀ first (the attention callers never
4729/// take that arm: they pack Kᵀ themselves). Not a BLAS: the shapes are
4730/// attention panels, and the goal is the end of the per-position
4731/// quadratic wall on x86, not peak GEMM. Scalar on a CPU without
4732/// AVX2/FMA.
4733#[cfg(target_arch = "x86_64")]
4734#[allow(clippy::too_many_arguments)]
4735pub(crate) fn x86_gemm_rm(
4736    m: usize,
4737    n: usize,
4738    k: usize,
4739    alpha: f32,
4740    a: &[f32],
4741    lda: usize,
4742    b_mat: &[f32],
4743    ldb: usize,
4744    b_rows_are_n: bool,
4745    c: &mut [f32],
4746    ldc: usize,
4747) {
4748    if m == 0 || n == 0 {
4749        return;
4750    }
4751    debug_assert!(a.len() >= (m - 1) * lda + k);
4752    debug_assert!(c.len() >= (m - 1) * ldc + n);
4753    if b_rows_are_n {
4754        debug_assert!(k == 0 || b_mat.len() >= (n - 1) * ldb + k);
4755        let mut bt = vec![0f32; k * n];
4756        for j in 0..n {
4757            for (p, &v) in b_mat[j * ldb..j * ldb + k].iter().enumerate() {
4758                bt[p * n + j] = v;
4759            }
4760        }
4761        return x86_gemm_rm(m, n, k, alpha, a, lda, &bt, n, false, c, ldc);
4762    }
4763    if k == 0 {
4764        for r in 0..m {
4765            c[r * ldc..r * ldc + n].fill(0.0);
4766        }
4767        return;
4768    }
4769    debug_assert!(b_mat.len() >= (k - 1) * ldb + n);
4770    if std::is_x86_feature_detected!("avx2") && std::is_x86_feature_detected!("fma") {
4771        // SAFETY: the bounds hold (asserted above), and both features were
4772        // just detected on this CPU.
4773        unsafe { x86_gemm_avx2(m, n, k, alpha, a, lda, b_mat, ldb, c, ldc) }
4774    } else {
4775        for r in 0..m {
4776            for j in 0..n {
4777                let mut acc = 0f32;
4778                for p in 0..k {
4779                    acc += a[r * lda + p] * b_mat[p * ldb + j];
4780                }
4781                c[r * ldc + j] = acc * alpha;
4782            }
4783        }
4784    }
4785}
4786
4787/// The AVX2+FMA body of `x86_gemm_rm` (B stored [k,n]). Loop nest:
4788/// column block (128) → k block (256) → four rows of A → 16-column
4789/// panel; the first k block starts the accumulators at zero, later ones
4790/// resume from C, the last one applies `alpha`.
4791///
4792/// # Safety
4793/// The caller has checked `avx2` and `fma` and the slice bounds
4794/// (`a` ≥ (m−1)·lda + k, `b` ≥ (k−1)·ldb + n, `c` ≥ (m−1)·ldc + n).
4795#[cfg(target_arch = "x86_64")]
4796#[target_feature(enable = "avx2,fma")]
4797#[allow(clippy::too_many_arguments)]
4798unsafe fn x86_gemm_avx2(
4799    m: usize,
4800    n: usize,
4801    k: usize,
4802    alpha: f32,
4803    a: &[f32],
4804    lda: usize,
4805    b: &[f32],
4806    ldb: usize,
4807    c: &mut [f32],
4808    ldc: usize,
4809) {
4810    use core::arch::x86_64::*;
4811    const NC: usize = 128;
4812    const KC: usize = 256;
4813    unsafe {
4814        let ap = a.as_ptr();
4815        let bp = b.as_ptr();
4816        let cp = c.as_mut_ptr();
4817        let al = _mm256_set1_ps(alpha);
4818        let mut jc = 0usize;
4819        while jc < n {
4820            let jend = (jc + NC).min(n);
4821            let mut pc = 0usize;
4822            while pc < k {
4823                let pend = (pc + KC).min(k);
4824                let (first, last) = (pc == 0, pend == k);
4825                let mut i = 0usize;
4826                while i < m {
4827                    let mi = (m - i).min(4);
4828                    let mut j = jc;
4829                    while j < jend {
4830                        let nj = (jend - j).min(16);
4831                        if mi == 4 && nj == 16 {
4832                            let a0 = ap.add(i * lda);
4833                            let a1 = ap.add((i + 1) * lda);
4834                            let a2 = ap.add((i + 2) * lda);
4835                            let a3 = ap.add((i + 3) * lda);
4836                            let c0 = cp.add(i * ldc + j);
4837                            let c1 = cp.add((i + 1) * ldc + j);
4838                            let c2 = cp.add((i + 2) * ldc + j);
4839                            let c3 = cp.add((i + 3) * ldc + j);
4840                            let (mut c00, mut c01, mut c10, mut c11, mut c20, mut c21, mut c30, mut c31) =
4841                                if first {
4842                                    let z = _mm256_setzero_ps();
4843                                    (z, z, z, z, z, z, z, z)
4844                                } else {
4845                                    (
4846                                        _mm256_loadu_ps(c0),
4847                                        _mm256_loadu_ps(c0.add(8)),
4848                                        _mm256_loadu_ps(c1),
4849                                        _mm256_loadu_ps(c1.add(8)),
4850                                        _mm256_loadu_ps(c2),
4851                                        _mm256_loadu_ps(c2.add(8)),
4852                                        _mm256_loadu_ps(c3),
4853                                        _mm256_loadu_ps(c3.add(8)),
4854                                    )
4855                                };
4856                            for p in pc..pend {
4857                                let bb = bp.add(p * ldb + j);
4858                                let b0 = _mm256_loadu_ps(bb);
4859                                let b1 = _mm256_loadu_ps(bb.add(8));
4860                                let x0 = _mm256_broadcast_ss(&*a0.add(p));
4861                                c00 = _mm256_fmadd_ps(x0, b0, c00);
4862                                c01 = _mm256_fmadd_ps(x0, b1, c01);
4863                                let x1 = _mm256_broadcast_ss(&*a1.add(p));
4864                                c10 = _mm256_fmadd_ps(x1, b0, c10);
4865                                c11 = _mm256_fmadd_ps(x1, b1, c11);
4866                                let x2 = _mm256_broadcast_ss(&*a2.add(p));
4867                                c20 = _mm256_fmadd_ps(x2, b0, c20);
4868                                c21 = _mm256_fmadd_ps(x2, b1, c21);
4869                                let x3 = _mm256_broadcast_ss(&*a3.add(p));
4870                                c30 = _mm256_fmadd_ps(x3, b0, c30);
4871                                c31 = _mm256_fmadd_ps(x3, b1, c31);
4872                            }
4873                            if last {
4874                                c00 = _mm256_mul_ps(c00, al);
4875                                c01 = _mm256_mul_ps(c01, al);
4876                                c10 = _mm256_mul_ps(c10, al);
4877                                c11 = _mm256_mul_ps(c11, al);
4878                                c20 = _mm256_mul_ps(c20, al);
4879                                c21 = _mm256_mul_ps(c21, al);
4880                                c30 = _mm256_mul_ps(c30, al);
4881                                c31 = _mm256_mul_ps(c31, al);
4882                            }
4883                            _mm256_storeu_ps(c0, c00);
4884                            _mm256_storeu_ps(c0.add(8), c01);
4885                            _mm256_storeu_ps(c1, c10);
4886                            _mm256_storeu_ps(c1.add(8), c11);
4887                            _mm256_storeu_ps(c2, c20);
4888                            _mm256_storeu_ps(c2.add(8), c21);
4889                            _mm256_storeu_ps(c3, c30);
4890                            _mm256_storeu_ps(c3.add(8), c31);
4891                        } else {
4892                            // Row-at-a-time tail: fewer than four rows left,
4893                            // or a panel narrower than 16 columns.
4894                            for r in 0..mi {
4895                                let ar = ap.add((i + r) * lda);
4896                                let cr = cp.add((i + r) * ldc + j);
4897                                let mut q = 0usize;
4898                                while q + 8 <= nj {
4899                                    let mut acc = if first {
4900                                        _mm256_setzero_ps()
4901                                    } else {
4902                                        _mm256_loadu_ps(cr.add(q))
4903                                    };
4904                                    for p in pc..pend {
4905                                        acc = _mm256_fmadd_ps(
4906                                            _mm256_broadcast_ss(&*ar.add(p)),
4907                                            _mm256_loadu_ps(bp.add(p * ldb + j + q)),
4908                                            acc,
4909                                        );
4910                                    }
4911                                    if last {
4912                                        acc = _mm256_mul_ps(acc, al);
4913                                    }
4914                                    _mm256_storeu_ps(cr.add(q), acc);
4915                                    q += 8;
4916                                }
4917                                while q < nj {
4918                                    let mut acc = if first { 0.0 } else { *cr.add(q) };
4919                                    for p in pc..pend {
4920                                        acc += *ar.add(p) * *bp.add(p * ldb + j + q);
4921                                    }
4922                                    if last {
4923                                        acc *= alpha;
4924                                    }
4925                                    *cr.add(q) = acc;
4926                                    q += 1;
4927                                }
4928                            }
4929                        }
4930                        j += nj;
4931                    }
4932                    i += mi;
4933                }
4934                pc = pend;
4935            }
4936            jc = jend;
4937        }
4938    }
4939}
4940
4941/// Row-major f32 GEMM, exposed for offline tools (the AWNP pass builds a
4942/// per-layer projection and applies it to every expert; a naive triple loop
4943/// would turn a two-minute job into half an hour).
4944#[allow(clippy::too_many_arguments)]
4945pub fn sgemm_public(
4946    m: usize,
4947    n: usize,
4948    k: usize,
4949    alpha: f32,
4950    a: &[f32],
4951    lda: usize,
4952    b_mat: &[f32],
4953    ldb: usize,
4954    b_rows_are_n: bool,
4955    c: &mut [f32],
4956    ldc: usize,
4957) {
4958    #[cfg(any(target_os = "macos", target_arch = "aarch64"))]
4959    {
4960        sgemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
4961    }
4962    // x86 without Accelerate has no sgemm_rm (the specialized paths there are
4963    // quantized kernels): the packed AVX/FMA/AVX-512 kernels of
4964    // matrixmultiply, single-threaded — offline callers split rows across
4965    // their own threads.
4966    #[cfg(not(any(target_os = "macos", target_arch = "aarch64")))]
4967    {
4968        if m == 0 || n == 0 {
4969            return;
4970        }
4971        assert!(a.len() >= (m - 1) * lda + k && c.len() >= (m - 1) * ldc + n);
4972        let (rsb, csb) = if b_rows_are_n { (1, ldb as isize) } else { (ldb as isize, 1) };
4973        assert!(k == 0 || b_mat.len() > if b_rows_are_n { (n - 1) * ldb + k - 1 } else { (k - 1) * ldb + n - 1 });
4974        // SAFETY: the asserts above bound every element the strides reach.
4975        unsafe {
4976            matrixmultiply::sgemm(
4977                m,
4978                k,
4979                n,
4980                alpha,
4981                a.as_ptr(),
4982                lda as isize,
4983                1,
4984                b_mat.as_ptr(),
4985                rsb,
4986                csb,
4987                0.0,
4988                c.as_mut_ptr(),
4989                ldc as isize,
4990                1,
4991            );
4992        }
4993    }
4994}
4995
4996/// Row-major f32 GEMM on Accelerate: C[m,n] = alpha·A[m,k] × B(ᵀ).
4997/// `b_rows_are_n` = true multiplies by Bᵀ where B is stored [n, k].
4998#[cfg(target_os = "macos")]
4999#[allow(clippy::too_many_arguments)]
5000pub(crate) fn sgemm_rm(
5001    m: usize,
5002    n: usize,
5003    k: usize,
5004    alpha: f32,
5005    a: &[f32],
5006    lda: usize,
5007    b_mat: &[f32],
5008    ldb: usize,
5009    b_rows_are_n: bool,
5010    c: &mut [f32],
5011    ldc: usize,
5012) {
5013    debug_assert!(a.len() >= (m - 1) * lda + k);
5014    debug_assert!(c.len() >= (m - 1) * ldc + n);
5015    // Test hook: route the attention GEMMs through the portable NEON
5016    // micro-kernel ON APPLE SILICON — how the mobile batched attend is
5017    // measured without a phone in the loop. (Intel macOS has no NEON —
5018    // the hook is a no-op there, Accelerate continues below.)
5019    #[cfg(target_arch = "aarch64")]
5020    if std::env::var("CMF_FORCE_NEON_GEMM")
5021        .map(|v| v == "1")
5022        .unwrap_or(false)
5023    {
5024        return neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
5025    }
5026    unsafe {
5027        accel_blas::cblas_sgemm(
5028            101, // RowMajor
5029            111, // NoTrans A
5030            if b_rows_are_n { 112 } else { 111 },
5031            m as i32,
5032            n as i32,
5033            k as i32,
5034            alpha,
5035            a.as_ptr(),
5036            lda as i32,
5037            b_mat.as_ptr(),
5038            ldb as i32,
5039            0.0,
5040            c.as_mut_ptr(),
5041            ldc as i32,
5042        );
5043    }
5044}
5045
5046/// Prefill GEMM through Accelerate (macOS): dequantize q8 rows into
5047/// f32 tiles (scale folded in, pool-parallel) and multiply each tile
5048/// on the AMX with one row-major sgemm. Tiles live in cache, weights
5049/// stream once. Numerics are f32-GEMM (not the int8 dot): prefill
5050/// logits shift within f32 rounding — tolerance-class, like every
5051/// reduction-order change; decode (M=1) never takes this path.
5052#[cfg(target_os = "macos")]
5053fn qmatmat_accel(
5054    q: &[u8],
5055    row_scale: &[f32],
5056    pre: &[std::borrow::Cow<'_, [f32]>],
5057    rows: usize,
5058    cols: usize,
5059    out: &mut [f32],
5060    pool: Option<&Pool>,
5061) {
5062    // NOTE: double-buffering the dequant against the sgemm (a scoped
5063    // thread driving the pool on tile k+1 while the caller multiplies
5064    // tile k) was tried and LOST ~6%: Accelerate's sgemm is itself
5065    // multithreaded, and the dequant workers just steal its cores.
5066    const TR: usize = 2048;
5067    let b = pre.len();
5068    thread_local! {
5069        static XPANEL: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
5070        static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
5071    }
5072    XPANEL.with(|xp| {
5073        WTILE.with(|wt| {
5074            let mut xpanel = xp.borrow_mut();
5075            xpanel.clear();
5076            for x in pre {
5077                xpanel.extend_from_slice(x);
5078            }
5079            let mut wtile = wt.borrow_mut();
5080            wtile.resize(TR * cols, 0.0);
5081            let mut r0 = 0usize;
5082            while r0 < rows {
5083                let tr = TR.min(rows - r0);
5084                // Dequant the tile (scale folded) — pool-parallel.
5085                let wt_addr = SendMut(wtile.as_mut_ptr());
5086                let run = |start: usize, end: usize| {
5087                    for r in start..end {
5088                        let row = &q[(r0 + r) * cols..(r0 + r + 1) * cols];
5089                        let s = row_scale[r0 + r];
5090                        // SAFETY: workers cover disjoint r ranges.
5091                        let dst =
5092                            unsafe { std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols) };
5093                        for (d, &v) in dst.iter_mut().zip(row) {
5094                            *d = (v as i8) as f32 * s;
5095                        }
5096                    }
5097                };
5098                dispatch_rows(pool, tr, &run);
5099                // C[b, tr] (at column r0 of out[b, rows]) = X · Wtileᵀ
5100                unsafe {
5101                    accel_blas::cblas_sgemm(
5102                        101, // RowMajor
5103                        111, // NoTrans A
5104                        112, // Trans B
5105                        b as i32,
5106                        tr as i32,
5107                        cols as i32,
5108                        1.0,
5109                        xpanel.as_ptr(),
5110                        cols as i32,
5111                        wtile.as_ptr(),
5112                        cols as i32,
5113                        0.0,
5114                        out.as_mut_ptr().add(r0),
5115                        rows as i32,
5116                    );
5117                }
5118                r0 += tr;
5119            }
5120        })
5121    });
5122}
5123
5124fn qmatmat(
5125    q: &[u8],
5126    row_scale: &[f32],
5127    pre: &[std::borrow::Cow<'_, [f32]>],
5128    rows: usize,
5129    cols: usize,
5130    out: &mut [f32],
5131    pool: Option<&Pool>,
5132) {
5133    let b = pre.len();
5134    debug_assert_eq!(out.len(), b * rows);
5135    let _p = crate::cpuprof::spans::span("mm.qmatmat");
5136    // Big prefill batches ride the AMX (roadmap PR3): the row×batch
5137    // SDOT loop below peaks near the CPU's dot throughput, an order
5138    // below the matrix units. Small tensors and tiny test models stay
5139    // on the exact integer path.
5140    #[cfg(target_os = "macos")]
5141    if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
5142        qmatmat_accel(q, row_scale, pre, rows, cols, out, pool);
5143        return;
5144    }
5145    #[cfg(target_arch = "aarch64")]
5146    if sdot_enabled() {
5147        let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
5148        let out_addr = SendMut(out.as_mut_ptr());
5149        // Blocked 2×4 (mobile prefill: no AMX to fall back on — this
5150        // path IS the ARM prefill GEMM off Apple silicon).
5151        let blocked_ok = blocked_enabled();
5152        let use_i8mm = i8mm_enabled();
5153        if blocked_ok {
5154            let run = |start: usize, end: usize| {
5155                let mut o = start;
5156                while o < end {
5157                    if o + 2 <= end {
5158                        let r0 = &q[o * cols..(o + 1) * cols];
5159                        let r1 = &q[(o + 1) * cols..(o + 2) * cols];
5160                        let mut bi = 0usize;
5161                        while bi + 4 <= acts.len() {
5162                            let xs = [
5163                                acts[bi].xq.as_slice(),
5164                                acts[bi + 1].xq.as_slice(),
5165                                acts[bi + 2].xq.as_slice(),
5166                                acts[bi + 3].xq.as_slice(),
5167                            ];
5168                            let d = if use_i8mm {
5169                                unsafe { dot_i8_smmla_2x4(r0, r1, xs) }
5170                            } else {
5171                                unsafe { dot_i8_sdot_2x4(r0, r1, xs) }
5172                            };
5173                            for (r, row) in [r0, r1].into_iter().enumerate() {
5174                                for k in 0..4 {
5175                                    let act = &acts[bi + k];
5176                                    let mut v = d[r][k] as f32 * act.sx;
5177                                    for &(j, xv) in &act.outliers {
5178                                        v += (row[j] as i8) as f32 * xv;
5179                                    }
5180                                    unsafe {
5181                                        *out_addr.at((bi + k) * rows + o + r) = v * row_scale[o + r]
5182                                    };
5183                                }
5184                            }
5185                            bi += 4;
5186                        }
5187                        while bi < acts.len() {
5188                            for (r, row) in [r0, r1].into_iter().enumerate() {
5189                                let v = row_dot_sdot(row, &acts[bi]) * row_scale[o + r];
5190                                unsafe { *out_addr.at(bi * rows + o + r) = v };
5191                            }
5192                            bi += 1;
5193                        }
5194                        o += 2;
5195                    } else {
5196                        let row = &q[o * cols..(o + 1) * cols];
5197                        for (bi, act) in acts.iter().enumerate() {
5198                            let v = row_dot_sdot(row, act) * row_scale[o];
5199                            unsafe { *out_addr.at(bi * rows + o) = v };
5200                        }
5201                        o += 1;
5202                    }
5203                }
5204            };
5205            dispatch_rows(pool, rows, &run);
5206            return;
5207        }
5208        let run = |start: usize, end: usize| {
5209            for o in start..end {
5210                let row = &q[o * cols..(o + 1) * cols];
5211                for (bi, act) in acts.iter().enumerate() {
5212                    let v = row_dot_sdot(row, act) * row_scale[o];
5213                    unsafe { *out_addr.at(bi * rows + o) = v };
5214                }
5215            }
5216        };
5217        dispatch_rows(pool, rows, &run);
5218        return;
5219    }
5220    // x86 A8W8 batch. Non-VNNI parts take the BLOCKED 2×4 kernel
5221    // (roadmap P0: two weight rows' abs() stay in registers across four
5222    // activation streams); VNNI machines keep the per-row bias-trick
5223    // dot, which is already throughput-bound there.
5224    #[cfg(target_arch = "x86_64")]
5225    if avx2_a8w8_enabled() {
5226        let acts = {
5227            let _psa = crate::cpuprof::spans::span("mm.split_act");
5228            let a: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
5229            drop(_psa);
5230            a
5231        };
5232        let out_addr = SendMut(out.as_mut_ptr());
5233        // CMF_X86_BLOCKED=0 forces the per-row path (paired in-process
5234        // A/B on noisy shared-vCPU hosts).
5235        let blocked_ok = blocked_enabled();
5236        if !avx512vnni_enabled() && blocked_ok && !row_exact() {
5237            let run = |start: usize, end: usize| {
5238                let mut o = start;
5239                while o < end {
5240                    if o + 2 <= end {
5241                        let r0 = &q[o * cols..(o + 1) * cols];
5242                        let r1 = &q[(o + 1) * cols..(o + 2) * cols];
5243                        let mut bi = 0usize;
5244                        while bi + 4 <= acts.len() {
5245                            let xs = [
5246                                acts[bi].xq.as_slice(),
5247                                acts[bi + 1].xq.as_slice(),
5248                                acts[bi + 2].xq.as_slice(),
5249                                acts[bi + 3].xq.as_slice(),
5250                            ];
5251                            let d = unsafe { dot_i8_i8_avx2_2x4(r0, r1, xs) };
5252                            for (r, row) in [r0, r1].into_iter().enumerate() {
5253                                for k in 0..4 {
5254                                    let act = &acts[bi + k];
5255                                    let mut v = d[r][k] as f32 * act.sx;
5256                                    for &(j, xv) in &act.outliers {
5257                                        v += (row[j] as i8) as f32 * xv;
5258                                    }
5259                                    unsafe {
5260                                        *out_addr.at((bi + k) * rows + o + r) = v * row_scale[o + r]
5261                                    };
5262                                }
5263                            }
5264                            bi += 4;
5265                        }
5266                        while bi < acts.len() {
5267                            for (r, row) in [r0, r1].into_iter().enumerate() {
5268                                let v = row_dot_avx2(row, &acts[bi]) * row_scale[o + r];
5269                                unsafe { *out_addr.at(bi * rows + o + r) = v };
5270                            }
5271                            bi += 1;
5272                        }
5273                        o += 2;
5274                    } else {
5275                        let row = &q[o * cols..(o + 1) * cols];
5276                        for (bi, act) in acts.iter().enumerate() {
5277                            let v = row_dot_avx2(row, act) * row_scale[o];
5278                            unsafe { *out_addr.at(bi * rows + o) = v };
5279                        }
5280                        o += 1;
5281                    }
5282                }
5283            };
5284            dispatch_rows(pool, rows, &run);
5285            return;
5286        }
5287        let run = |start: usize, end: usize| {
5288            for o in start..end {
5289                let row = &q[o * cols..(o + 1) * cols];
5290                for (bi, act) in acts.iter().enumerate() {
5291                    let v = row_dot_avx2(row, act) * row_scale[o];
5292                    unsafe { *out_addr.at(bi * rows + o) = v };
5293                }
5294            }
5295        };
5296        dispatch_rows(pool, rows, &run);
5297        return;
5298    }
5299    let out_addr = SendMut(out.as_mut_ptr());
5300    let run = |start: usize, end: usize| {
5301        for o in start..end {
5302            let row = &q[o * cols..(o + 1) * cols];
5303            for (bi, x) in pre.iter().enumerate() {
5304                let mut acc = 0f32;
5305                for j in 0..cols {
5306                    acc += (row[j] as i8) as f32 * x[j];
5307                }
5308                unsafe { *out_addr.at(bi * rows + o) = acc * row_scale[o] };
5309            }
5310        }
5311    };
5312    dispatch_rows(pool, rows, &run);
5313}
5314
5315/// Split rows across pool workers (shared qmatvec pattern). Self-balancing
5316/// — see `Pool::run_rows` for why a static 1/n split is wrong here.
5317fn dispatch_rows(pool: Option<&Pool>, rows: usize, run: &(dyn Fn(usize, usize) + Sync)) {
5318    match pool {
5319        Some(pool) if rows >= 256 => pool.run_rows(rows, run),
5320        _ => run(0, rows),
5321    }
5322}
5323
5324/// Split a q4_block blob into (packed nibbles, f16 group scales).
5325fn q4_split(bytes: &[u8], rows: usize, cols: usize) -> (&[u8], &[u8]) {
5326    let groups = rows * cols / GROUP_SIZE;
5327    bytes.split_at(groups * 16)
5328}
5329
5330/// SIMD unpack for the dominant vbit width B=4 (94% of rows on the
5331/// log2-shape calibration): 16 packed bytes -> 32 centered i8 values.
5332/// vbit packs MSB-first, so the HIGH nibble is the even element
5333/// (opposite of q4_block's lo-first interleave). Centering is u-7.
5334#[inline]
5335fn vbit_fill4(data: &[u8], buf: &mut [u8]) {
5336    #[cfg(target_arch = "aarch64")]
5337    unsafe {
5338        return vbit_fill4_neon(data, buf);
5339    }
5340    #[cfg(target_arch = "x86_64")]
5341    if avx2_enabled() {
5342        return unsafe { vbit_fill4_avx2(data, buf) };
5343    }
5344    #[allow(unreachable_code)]
5345    for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
5346        let u = unpack8::<4>(&data[blk * 4..]);
5347        for k in 0..8 {
5348            chunk[k] = (u[k] - 7) as i8 as u8;
5349        }
5350    }
5351}
5352
5353#[cfg(target_arch = "aarch64")]
5354#[target_feature(enable = "neon")]
5355unsafe fn vbit_fill4_neon(data: &[u8], buf: &mut [u8]) {
5356    // SAFETY: buf.len() is a multiple of GROUP_SIZE=32; data holds
5357    // buf.len()/2 packed bytes (validated at load).
5358    unsafe {
5359        use core::arch::aarch64::*;
5360        let n = buf.len();
5361        let mask = vdupq_n_u8(0x0F);
5362        let seven = vdupq_n_s8(7);
5363        let mut g = 0usize;
5364        while g * 32 + 32 <= n {
5365            let b = vld1q_u8(data.as_ptr().add(g * 16));
5366            let hi = vshrq_n_u8::<4>(b);
5367            let lo = vandq_u8(b, mask);
5368            let z0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(hi, lo)), seven);
5369            let z1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(hi, lo)), seven);
5370            vst1q_u8(buf.as_mut_ptr().add(g * 32), vreinterpretq_u8_s8(z0));
5371            vst1q_u8(buf.as_mut_ptr().add(g * 32 + 16), vreinterpretq_u8_s8(z1));
5372            g += 1;
5373        }
5374    }
5375}
5376
5377#[cfg(target_arch = "x86_64")]
5378#[target_feature(enable = "avx2")]
5379unsafe fn vbit_fill4_avx2(data: &[u8], buf: &mut [u8]) {
5380    // SAFETY: see vbit_fill4_neon.
5381    unsafe {
5382        use core::arch::x86_64::*;
5383        let n = buf.len();
5384        let mask = _mm_set1_epi8(0x0F);
5385        let seven = _mm256_set1_epi8(7);
5386        let mut g = 0usize;
5387        while g * 32 + 32 <= n {
5388            let b = _mm_loadu_si128(data.as_ptr().add(g * 16) as *const __m128i);
5389            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), mask);
5390            let lo = _mm_and_si128(b, mask);
5391            let z = _mm256_sub_epi8(
5392                _mm256_set_m128i(_mm_unpackhi_epi8(hi, lo), _mm_unpacklo_epi8(hi, lo)),
5393                seven,
5394            );
5395            _mm256_storeu_si256(buf.as_mut_ptr().add(g * 32) as *mut __m256i, z);
5396            g += 1;
5397        }
5398    }
5399}
5400
5401/// Unpack 8 MSB-first B-bit values from exactly B bytes (fixed shifts —
5402/// no serial bit-buffer, auto-vectorizable). Every 32-value group starts
5403/// byte-aligned (32·B/8 is integral for B∈3..8), so groups decompose
5404/// into 4 such blocks.
5405#[inline(always)]
5406fn unpack8<const B: usize>(data: &[u8]) -> [i32; 8] {
5407    let mut acc = 0u64;
5408    for i in 0..B {
5409        acc = (acc << 8) | data[i] as u64;
5410    }
5411    let mask = (1u64 << B) - 1;
5412    let mut out = [0i32; 8];
5413    for (k, o) in out.iter_mut().enumerate() {
5414        *o = ((acc >> ((7 - k) * B)) & mask) as i32;
5415    }
5416    out
5417}
5418
5419/// Fused vbit matvec straight from the mapped bytes (spec §3, P13
5420/// FIG.3): [u8 bits: rows][f16 scales: rows·cols/32][bit-packed rows,
5421/// MSB-first, byte-padded]. Row data offsets are precomputed at load
5422/// (`vbit_row_offsets`) — the per-call prefix scan was O(rows) pure
5423/// overhead on every matvec.
5424#[allow(clippy::too_many_arguments)]
5425fn vbitmatvec(
5426    bytes: &[u8],
5427    offsets: &[usize],
5428    x: &[f32],
5429    rows: usize,
5430    cols: usize,
5431    out: &mut [f32],
5432    pool: Option<&Pool>,
5433) {
5434    debug_assert_eq!(out.len(), rows);
5435    debug_assert_eq!(offsets.len(), rows + 1);
5436
5437    // SDOT path: unpack the row to centered i8 once, then per-group
5438    // int8 dot against the quantized activations — same A8W8 contract
5439    // as q8 (bounded noise; CMF_SDOT=0 keeps the exact scalar path).
5440    if a8w8_enabled() {
5441        let act = split_act(x);
5442        let out_addr = SendMut(out.as_mut_ptr());
5443        let run = move |start: usize, end: usize| {
5444            vbit_range_a8w8(bytes, offsets, x, &act, rows, cols, out_addr, start, end)
5445        };
5446        dispatch_rows(pool, rows, &run);
5447        return;
5448    }
5449
5450    let out_addr = SendMut(out.as_mut_ptr());
5451    let run = move |start: usize, end: usize| {
5452        vbit_range_f32(bytes, offsets, x, rows, cols, out_addr, start, end)
5453    };
5454    dispatch_rows(pool, rows, &run);
5455}
5456
5457/// One vbit row range via the A8W8 int8 path — kernel body of
5458/// `vbitmatvec`, extracted so multi-matrix jobs can drive it for
5459/// several tensors in one dispatch (b=8 rows go exact f32).
5460#[allow(clippy::too_many_arguments)]
5461fn vbit_range_a8w8(
5462    bytes: &[u8],
5463    offsets: &[usize],
5464    x: &[f32],
5465    act: &SplitAct,
5466    rows: usize,
5467    cols: usize,
5468    out: SendMut,
5469    start: usize,
5470    end: usize,
5471) {
5472    let ng = cols / GROUP_SIZE;
5473    let bits = &bytes[..rows];
5474    let sc_off = rows;
5475    let row_dot = |r: usize| -> f32 {
5476        let b = bits[r] as usize;
5477        let l = (1i32 << (b - 1)) - 1;
5478        let mask = (1u64 << b) - 1;
5479        let data = &bytes[offsets[r]..offsets[r + 1]];
5480        if b == 8 {
5481            // u−L reaches 128 → does not fit i8; exact f32 path.
5482            let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
5483            let mut dot = 0f32;
5484            for g in 0..ng {
5485                let so = (r * ng + g) * 2;
5486                let sgf = f16_to_f32(u16::from_le_bytes([
5487                    bytes[sc_off + so],
5488                    bytes[sc_off + so + 1],
5489                ]));
5490                let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
5491                let mut gd = 0f32;
5492                for &xv in xg.iter() {
5493                    if nbits < 8 {
5494                        acc = (acc << 8) | data[idx] as u64;
5495                        idx += 1;
5496                        nbits += 8;
5497                    }
5498                    let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
5499                    nbits -= 8;
5500                    gd += (u - l) as f32 * xv;
5501                }
5502                dot += gd * sgf;
5503            }
5504            return dot;
5505        }
5506        // Per-worker scratch: this closure runs for every row of the
5507        // tensor (lm_head ≈ 150k rows/token) — a heap allocation per
5508        // row was measurable pure overhead.
5509        thread_local! {
5510            static VBIT_SCRATCH: std::cell::RefCell<Vec<u8>> =
5511                const { std::cell::RefCell::new(Vec::new()) };
5512        }
5513        #[inline(always)]
5514        fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
5515            for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
5516                let u = unpack8::<B>(&data[blk * B..]);
5517                for k in 0..8 {
5518                    chunk[k] = (u[k] - l) as i8 as u8;
5519                }
5520            }
5521        }
5522        let _ = mask;
5523        VBIT_SCRATCH.with(|scratch| {
5524            let mut buf = scratch.borrow_mut();
5525            buf.resize(cols, 0);
5526            match b {
5527                3 => fill::<3>(data, l, &mut buf),
5528                4 => vbit_fill4(data, &mut buf),
5529                5 => fill::<5>(data, l, &mut buf),
5530                6 => fill::<6>(data, l, &mut buf),
5531                _ => unreachable!(),
5532            }
5533            let mut dot = 0f32;
5534            for g in 0..ng {
5535                let so = (r * ng + g) * 2;
5536                let s = f16_to_f32(u16::from_le_bytes([
5537                    bytes[sc_off + so],
5538                    bytes[sc_off + so + 1],
5539                ]));
5540                let d = dot_i8_i8(
5541                    &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
5542                    &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
5543                ) as f32
5544                    * act.sx;
5545                dot += d * s;
5546            }
5547            for &(j, xv) in &act.outliers {
5548                let so = (r * ng + j / GROUP_SIZE) * 2;
5549                let s = f16_to_f32(u16::from_le_bytes([
5550                    bytes[sc_off + so],
5551                    bytes[sc_off + so + 1],
5552                ]));
5553                // xq is zeroed at outlier slots — add the exact term.
5554                dot += (buf[j] as i8) as f32 * s * xv;
5555            }
5556            dot
5557        })
5558    };
5559    for r in start..end {
5560        // SAFETY: disjoint row ranges per worker.
5561        unsafe { *out.at(r) = row_dot(r) };
5562    }
5563}
5564
5565/// Exact scalar vbit row range (same extraction, non-SDOT path).
5566#[allow(clippy::too_many_arguments)]
5567fn vbit_range_f32(
5568    bytes: &[u8],
5569    offsets: &[usize],
5570    x: &[f32],
5571    rows: usize,
5572    cols: usize,
5573    out: SendMut,
5574    start: usize,
5575    end: usize,
5576) {
5577    let ng = cols / GROUP_SIZE;
5578    let bits = &bytes[..rows];
5579    let sc_off = rows;
5580    // Per-bit-width specialized inner loops: the compiler unrolls the
5581    // constant shifts (the generic bit-buffer loop was branch-bound —
5582    // 5.6 vs 13.2 tok/s q4 on the 0.8B).
5583    #[inline(always)]
5584    fn dot_row<const B: usize>(
5585        data: &[u8],
5586        bytes: &[u8],
5587        sc_off: usize,
5588        r: usize,
5589        ng: usize,
5590        x: &[f32],
5591    ) -> f32 {
5592        let l = ((1i32 << (B - 1)) - 1) as f32;
5593        let gbytes = GROUP_SIZE * B / 8;
5594        let mut dot = 0f32;
5595        for g in 0..ng {
5596            let so = (r * ng + g) * 2;
5597            let s = f16_to_f32(u16::from_le_bytes([
5598                bytes[sc_off + so],
5599                bytes[sc_off + so + 1],
5600            ]));
5601            let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
5602            let gd0 = &data[g * gbytes..(g + 1) * gbytes];
5603            let mut gd = 0f32;
5604            for blk in 0..GROUP_SIZE / 8 {
5605                let u = unpack8::<B>(&gd0[blk * B..]);
5606                let xb = &xg[blk * 8..blk * 8 + 8];
5607                for k in 0..8 {
5608                    gd += (u[k] as f32 - l) * xb[k];
5609                }
5610            }
5611            dot += gd * s;
5612        }
5613        dot
5614    }
5615    for r in start..end {
5616        let data = &bytes[offsets[r]..offsets[r + 1]];
5617        let v = match bits[r] {
5618            3 => dot_row::<3>(data, bytes, sc_off, r, ng, x),
5619            4 => dot_row::<4>(data, bytes, sc_off, r, ng, x),
5620            5 => dot_row::<5>(data, bytes, sc_off, r, ng, x),
5621            6 => dot_row::<6>(data, bytes, sc_off, r, ng, x),
5622            8 => dot_row::<8>(data, bytes, sc_off, r, ng, x),
5623            b => unreachable!("vbit bit-width {b} (validated at load)"),
5624        };
5625        // SAFETY: disjoint row ranges per worker.
5626        unsafe { *out.at(r) = v };
5627    }
5628}
5629
5630/// Fused two-input vbit matvec: each row is unpacked from the mmap ONCE
5631/// and dotted against BOTH activations (MTP verify / pair prefill used
5632/// to run two full matvecs — double weight traffic and double unpack).
5633/// Per-input math is identical to `vbitmatvec` → same accuracy contract.
5634#[allow(clippy::too_many_arguments)]
5635fn vbitmatvec2(
5636    bytes: &[u8],
5637    offsets: &[usize],
5638    x1: &[f32],
5639    x2: &[f32],
5640    rows: usize,
5641    cols: usize,
5642    o1: &mut [f32],
5643    o2: &mut [f32],
5644    pool: Option<&Pool>,
5645) {
5646    debug_assert_eq!(o1.len(), rows);
5647    debug_assert_eq!(o2.len(), rows);
5648
5649    if a8w8_enabled() {
5650        let a1 = split_act(x1);
5651        let a2 = split_act(x2);
5652        let p1 = SendMut(o1.as_mut_ptr());
5653        let p2 = SendMut(o2.as_mut_ptr());
5654        let run = move |start: usize, end: usize| {
5655            vbit_range2_a8w8(
5656                bytes, offsets, x1, x2, &a1, &a2, rows, cols, p1, p2, start, end,
5657            )
5658        };
5659        dispatch_rows(pool, rows, &run);
5660        return;
5661    }
5662
5663    let p1 = SendMut(o1.as_mut_ptr());
5664    let p2 = SendMut(o2.as_mut_ptr());
5665    let run = move |start: usize, end: usize| {
5666        vbit_range2_f32(bytes, offsets, x1, x2, rows, cols, p1, p2, start, end)
5667    };
5668    dispatch_rows(pool, rows, &run);
5669}
5670
5671/// Two-input vbit row range via the A8W8 int8 path — kernel body of
5672/// `vbitmatvec2`, extracted for pair multi-matrix jobs (b=8 rows go
5673/// exact f32 for both lanes, bits streamed once).
5674#[allow(clippy::too_many_arguments)]
5675fn vbit_range2_a8w8(
5676    bytes: &[u8],
5677    offsets: &[usize],
5678    x1: &[f32],
5679    x2: &[f32],
5680    a1: &SplitAct,
5681    a2: &SplitAct,
5682    rows: usize,
5683    cols: usize,
5684    p1: SendMut,
5685    p2: SendMut,
5686    start: usize,
5687    end: usize,
5688) {
5689    let ng = cols / GROUP_SIZE;
5690    let bits = &bytes[..rows];
5691    let sc_off = rows;
5692    let row_dots = |r: usize| -> (f32, f32) {
5693        let b = bits[r] as usize;
5694        let l = (1i32 << (b - 1)) - 1;
5695        let data = &bytes[offsets[r]..offsets[r + 1]];
5696        if b == 8 {
5697            // u−L reaches 128 → does not fit i8; exact f32 path,
5698            // bits still streamed once for both lanes.
5699            let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
5700            let (mut d1, mut d2) = (0f32, 0f32);
5701            for g in 0..ng {
5702                let so = (r * ng + g) * 2;
5703                let sgf = f16_to_f32(u16::from_le_bytes([
5704                    bytes[sc_off + so],
5705                    bytes[sc_off + so + 1],
5706                ]));
5707                let (mut g1, mut g2) = (0f32, 0f32);
5708                for k in 0..GROUP_SIZE {
5709                    if nbits < 8 {
5710                        acc = (acc << 8) | data[idx] as u64;
5711                        idx += 1;
5712                        nbits += 8;
5713                    }
5714                    let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
5715                    nbits -= 8;
5716                    let w = (u - l) as f32;
5717                    g1 += w * x1[g * GROUP_SIZE + k];
5718                    g2 += w * x2[g * GROUP_SIZE + k];
5719                }
5720                d1 += g1 * sgf;
5721                d2 += g2 * sgf;
5722            }
5723            return (d1, d2);
5724        }
5725        thread_local! {
5726            static VBIT_SCRATCH2: std::cell::RefCell<Vec<u8>> =
5727                const { std::cell::RefCell::new(Vec::new()) };
5728        }
5729        #[inline(always)]
5730        fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
5731            for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
5732                let u = unpack8::<B>(&data[blk * B..]);
5733                for k in 0..8 {
5734                    chunk[k] = (u[k] - l) as i8 as u8;
5735                }
5736            }
5737        }
5738        VBIT_SCRATCH2.with(|scratch| {
5739            let mut buf = scratch.borrow_mut();
5740            buf.resize(cols, 0);
5741            match b {
5742                3 => fill::<3>(data, l, &mut buf),
5743                4 => vbit_fill4(data, &mut buf),
5744                5 => fill::<5>(data, l, &mut buf),
5745                6 => fill::<6>(data, l, &mut buf),
5746                _ => unreachable!(),
5747            }
5748            let (mut d1, mut d2) = (0f32, 0f32);
5749            for g in 0..ng {
5750                let so = (r * ng + g) * 2;
5751                let s = f16_to_f32(u16::from_le_bytes([
5752                    bytes[sc_off + so],
5753                    bytes[sc_off + so + 1],
5754                ]));
5755                let wg = &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
5756                let v1 = dot_i8_i8(wg, &a1.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a1.sx;
5757                let v2 = dot_i8_i8(wg, &a2.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a2.sx;
5758                d1 += v1 * s;
5759                d2 += v2 * s;
5760            }
5761            for &(j, xv) in &a1.outliers {
5762                let so = (r * ng + j / GROUP_SIZE) * 2;
5763                let s = f16_to_f32(u16::from_le_bytes([
5764                    bytes[sc_off + so],
5765                    bytes[sc_off + so + 1],
5766                ]));
5767                d1 += (buf[j] as i8) as f32 * s * xv;
5768            }
5769            for &(j, xv) in &a2.outliers {
5770                let so = (r * ng + j / GROUP_SIZE) * 2;
5771                let s = f16_to_f32(u16::from_le_bytes([
5772                    bytes[sc_off + so],
5773                    bytes[sc_off + so + 1],
5774                ]));
5775                d2 += (buf[j] as i8) as f32 * s * xv;
5776            }
5777            (d1, d2)
5778        })
5779    };
5780    for r in start..end {
5781        let (v1, v2) = row_dots(r);
5782        // SAFETY: disjoint row ranges per worker.
5783        unsafe {
5784            *p1.at(r) = v1;
5785            *p2.at(r) = v2;
5786        }
5787    }
5788}
5789
5790/// Two-input exact scalar vbit row range (same extraction) —
5791/// per-bit-width specialized, two accumulators per row; per-lane
5792/// accumulation order matches `vbitmatvec` exactly.
5793#[allow(clippy::too_many_arguments)]
5794fn vbit_range2_f32(
5795    bytes: &[u8],
5796    offsets: &[usize],
5797    x1: &[f32],
5798    x2: &[f32],
5799    rows: usize,
5800    cols: usize,
5801    p1: SendMut,
5802    p2: SendMut,
5803    start: usize,
5804    end: usize,
5805) {
5806    let ng = cols / GROUP_SIZE;
5807    let bits = &bytes[..rows];
5808    let sc_off = rows;
5809    #[inline(always)]
5810    #[allow(clippy::too_many_arguments)]
5811    fn dot_row2<const B: usize>(
5812        data: &[u8],
5813        bytes: &[u8],
5814        sc_off: usize,
5815        r: usize,
5816        ng: usize,
5817        x1: &[f32],
5818        x2: &[f32],
5819    ) -> (f32, f32) {
5820        let l = ((1i32 << (B - 1)) - 1) as f32;
5821        let gbytes = GROUP_SIZE * B / 8;
5822        let (mut d1, mut d2) = (0f32, 0f32);
5823        for g in 0..ng {
5824            let so = (r * ng + g) * 2;
5825            let s = f16_to_f32(u16::from_le_bytes([
5826                bytes[sc_off + so],
5827                bytes[sc_off + so + 1],
5828            ]));
5829            let x1g = &x1[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
5830            let x2g = &x2[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
5831            let gd0 = &data[g * gbytes..(g + 1) * gbytes];
5832            let (mut g1, mut g2) = (0f32, 0f32);
5833            for blk in 0..GROUP_SIZE / 8 {
5834                let u = unpack8::<B>(&gd0[blk * B..]);
5835                for k in 0..8 {
5836                    let w = u[k] as f32 - l;
5837                    g1 += w * x1g[blk * 8 + k];
5838                    g2 += w * x2g[blk * 8 + k];
5839                }
5840            }
5841            d1 += g1 * s;
5842            d2 += g2 * s;
5843        }
5844        (d1, d2)
5845    }
5846    for r in start..end {
5847        let data = &bytes[offsets[r]..offsets[r + 1]];
5848        let (v1, v2) = match bits[r] {
5849            3 => dot_row2::<3>(data, bytes, sc_off, r, ng, x1, x2),
5850            4 => dot_row2::<4>(data, bytes, sc_off, r, ng, x1, x2),
5851            5 => dot_row2::<5>(data, bytes, sc_off, r, ng, x1, x2),
5852            6 => dot_row2::<6>(data, bytes, sc_off, r, ng, x1, x2),
5853            8 => dot_row2::<8>(data, bytes, sc_off, r, ng, x1, x2),
5854            b => unreachable!("vbit bit-width {b} (validated at load)"),
5855        };
5856        // SAFETY: disjoint row ranges per worker.
5857        unsafe {
5858            *p1.at(r) = v1;
5859            *p2.at(r) = v2;
5860        }
5861    }
5862}
5863
5864// ───────────────────── q4_tiled kernels (§4.3) ─────────────────────
5865
5866/// One q4_tiled row dot on the A8W8 int8 path: per 32-group the tile
5867/// is ONE sequential read — [f16 scale][16B nibbles] — versus the two
5868/// distant streams of the split layout. Values/order identical to the
5869/// split kernels.
5870#[inline]
5871#[allow(unreachable_code)]
5872fn dot_q4t_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
5873    #[cfg(target_arch = "aarch64")]
5874    unsafe {
5875        return dot_q4t_row_sdot(bytes, r, gpr, xq);
5876    }
5877    #[cfg(target_arch = "x86_64")]
5878    unsafe {
5879        if vnni_tiles_enabled() {
5880            return dot_q4t_row_vnni(bytes, r, gpr, xq);
5881        }
5882        return dot_q4t_row_avx2(bytes, r, gpr, xq);
5883    }
5884    let mut acc = 0f32;
5885    for gi in 0..gpr {
5886        let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
5887        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
5888        let mut d = 0i32;
5889        for (k, &b) in tile[2..].iter().enumerate() {
5890            d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
5891                + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
5892        }
5893        acc += d as f32 * s;
5894    }
5895    acc
5896}
5897
5898#[cfg(target_arch = "aarch64")]
5899#[target_feature(enable = "neon,dotprod")]
5900unsafe fn dot_q4t_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
5901    // SAFETY: callers uphold slice-length contracts (18B tile per group,
5902    // xq.len() == gpr·GROUP_SIZE).
5903    unsafe {
5904        use core::arch::aarch64::*;
5905        use core::arch::asm;
5906        let lomask = vdupq_n_u8(0x0F);
5907        let eight = vdupq_n_s8(8);
5908        let mut acc = 0f32;
5909        for gi in 0..gpr {
5910            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
5911            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5912            let b = vld1q_u8(t.add(2));
5913            let lo = vandq_u8(b, lomask);
5914            let hi = vshrq_n_u8::<4>(b);
5915            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
5916            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
5917            let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
5918            let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
5919            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5920            asm!(
5921                "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
5922                "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
5923                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5924                e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
5925                options(pure, nomem, nostack),
5926            );
5927            acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
5928        }
5929        acc
5930    }
5931}
5932
5933#[cfg(target_arch = "x86_64")]
5934#[target_feature(enable = "avx2")]
5935unsafe fn dot_q4t_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
5936    // SAFETY: see dot_q4t_row_sdot.
5937    unsafe {
5938        use core::arch::x86_64::*;
5939        let lomask = _mm_set1_epi8(0x0F);
5940        let eight = _mm256_set1_epi8(8);
5941        let ones = _mm256_set1_epi16(1);
5942        let mut acc = 0f32;
5943        for gi in 0..gpr {
5944            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
5945            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5946            let b = _mm_loadu_si128(t.add(2) as *const __m128i);
5947            let lo = _mm_and_si128(b, lomask);
5948            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
5949            let w = _mm256_sub_epi8(
5950                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5951                eight,
5952            );
5953            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5954            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
5955            let d = _mm256_madd_epi16(p16, ones);
5956            let hi128 = _mm256_extracti128_si256::<1>(d);
5957            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
5958            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
5959            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
5960            acc += _mm_cvtsi128_si32(s32) as f32 * s;
5961        }
5962        acc
5963    }
5964}
5965
5966/// VNNI twin of `dot_q4t_row_avx2`: same unpack, `vpdpbusd` replaces
5967/// the maddubs+madd pair (see `dpbusd_hsum` — sums are bit-identical).
5968/// 256-bit VL encoding, so the VEX `vpsignb` stays usable.
5969#[cfg(target_arch = "x86_64")]
5970#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
5971unsafe fn dot_q4t_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
5972    // SAFETY: see dot_q4t_row_sdot.
5973    unsafe {
5974        use core::arch::x86_64::*;
5975        let lomask = _mm_set1_epi8(0x0F);
5976        let eight = _mm256_set1_epi8(8);
5977        let mut acc = 0f32;
5978        for gi in 0..gpr {
5979            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
5980            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5981            let b = _mm_loadu_si128(t.add(2) as *const __m128i);
5982            let lo = _mm_and_si128(b, lomask);
5983            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
5984            let w = _mm256_sub_epi8(
5985                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5986                eight,
5987            );
5988            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5989            let d = dpbusd_hsum(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
5990            acc += d as f32 * s;
5991        }
5992        acc
5993    }
5994}
5995
5996/// One q4_tiled row against FOUR activation streams: the nibble unpack
5997/// and abs() happen once per group instead of once per (group,
5998/// activation) — the unpack is the dominant per-element cost of the
5999/// tiled format (roadmap P0 portable blocking, q4t leg).
6000#[cfg(target_arch = "x86_64")]
6001// `fma` is NOT implied by `avx2`: without it LLVM lowers _mm256_fmadd_ps
6002// to a libm call per lane — measured 2x slower than the reduction this
6003// kernel replaces. The runtime gate (`avx2_enabled`) already requires
6004// both features, so declaring it here is safe.
6005#[target_feature(enable = "avx2,fma")]
6006unsafe fn dot_q4t_row_1x4_avx2(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
6007    // SAFETY: callers uphold the 18B-tile and xq-length contracts.
6008    unsafe {
6009        use core::arch::x86_64::*;
6010        let lomask = _mm_set1_epi8(0x0F);
6011        let eight = _mm256_set1_epi8(8);
6012        let ones = _mm256_set1_epi16(1);
6013        // One f32 accumulator VECTOR per activation, reduced once at the
6014        // end. Folding each group's i32 lanes to a scalar inside the loop
6015        // costs an extracti128 + three shift/add + a movd — a cross-lane
6016        // dependency chain per (group, activation), 288 of them per row at
6017        // cols=2304. The per-group scale is what forces a float
6018        // accumulator; it does not force a horizontal sum.
6019        //
6020        // The four accumulators are NAMED, not an array: as `[__m256; 4]`
6021        // indexed by a loop variable LLVM keeps them in memory and every
6022        // group pays four 32-byte loads and stores. That alone made this
6023        // kernel 2x SLOWER than the per-group reduction it replaces
6024        // (measured on the EPYC box: 150 s vs 71 s for two 256² steps).
6025        let mut f0 = _mm256_setzero_ps();
6026        let mut f1 = _mm256_setzero_ps();
6027        let mut f2 = _mm256_setzero_ps();
6028        let mut f3 = _mm256_setzero_ps();
6029        for gi in 0..gpr {
6030            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
6031            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
6032            let sv = _mm256_set1_ps(s);
6033            let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
6034            let lo = _mm_and_si128(bb, lomask);
6035            let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
6036            let w = _mm256_sub_epi8(
6037                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
6038                eight,
6039            );
6040            let aw = _mm256_abs_epi8(w);
6041            let off = gi * GROUP_SIZE;
6042            let dot = |xq: &[i8]| {
6043                let x = _mm256_loadu_si256(xq.as_ptr().add(off) as *const __m256i);
6044                let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
6045                _mm256_cvtepi32_ps(_mm256_madd_epi16(p16, ones))
6046            };
6047            f0 = _mm256_fmadd_ps(dot(xs[0]), sv, f0);
6048            f1 = _mm256_fmadd_ps(dot(xs[1]), sv, f1);
6049            f2 = _mm256_fmadd_ps(dot(xs[2]), sv, f2);
6050            f3 = _mm256_fmadd_ps(dot(xs[3]), sv, f3);
6051        }
6052        [
6053            hsum256_ps(f0),
6054            hsum256_ps(f1),
6055            hsum256_ps(f2),
6056            hsum256_ps(f3),
6057        ]
6058    }
6059}
6060
6061/// Horizontal sum of eight f32 lanes — the one cross-lane reduction the
6062/// blocked kernels pay, once per row instead of once per group.
6063#[cfg(target_arch = "x86_64")]
6064#[target_feature(enable = "avx2")]
6065#[inline]
6066unsafe fn hsum256_ps(v: core::arch::x86_64::__m256) -> f32 {
6067    // SAFETY: pure register arithmetic on the caller's vector.
6068    unsafe {
6069        use core::arch::x86_64::*;
6070        let hi = _mm256_extractf128_ps::<1>(v);
6071        let s = _mm_add_ps(_mm256_castps256_ps128(v), hi);
6072        let s = _mm_add_ps(s, _mm_movehl_ps(s, s));
6073        let s = _mm_add_ss(s, _mm_shuffle_ps::<0x55>(s, s));
6074        _mm_cvtss_f32(s)
6075    }
6076}
6077
6078/// VNNI twin of `dot_q4t_row_1x4_avx2` (see `dpbusd_hsum`).
6079#[cfg(target_arch = "x86_64")]
6080#[target_feature(enable = "avx2,fma,avx512f,avx512bw,avx512vl,avx512vnni")]
6081unsafe fn dot_q4t_row_1x4_vnni(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
6082    // SAFETY: callers uphold the 18B-tile and xq-length contracts.
6083    unsafe {
6084        use core::arch::x86_64::*;
6085        let lomask = _mm_set1_epi8(0x0F);
6086        let eight = _mm256_set1_epi8(8);
6087        // Same shape as the AVX2 twin: accumulate in f32 vectors and pay
6088        // one cross-lane reduction per row, not per (group, activation).
6089        let mut f0 = _mm256_setzero_ps();
6090        let mut f1 = _mm256_setzero_ps();
6091        let mut f2 = _mm256_setzero_ps();
6092        let mut f3 = _mm256_setzero_ps();
6093        for gi in 0..gpr {
6094            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
6095            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
6096            let sv = _mm256_set1_ps(s);
6097            let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
6098            let lo = _mm_and_si128(bb, lomask);
6099            let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
6100            let w = _mm256_sub_epi8(
6101                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
6102                eight,
6103            );
6104            let aw = _mm256_abs_epi8(w);
6105            let off = gi * GROUP_SIZE;
6106            let dot = |xq: &[i8]| {
6107                let x = _mm256_loadu_si256(xq.as_ptr().add(off) as *const __m256i);
6108                _mm256_cvtepi32_ps(_mm256_dpbusd_epi32(
6109                    _mm256_setzero_si256(),
6110                    aw,
6111                    _mm256_sign_epi8(x, w),
6112                ))
6113            };
6114            f0 = _mm256_fmadd_ps(dot(xs[0]), sv, f0);
6115            f1 = _mm256_fmadd_ps(dot(xs[1]), sv, f1);
6116            f2 = _mm256_fmadd_ps(dot(xs[2]), sv, f2);
6117            f3 = _mm256_fmadd_ps(dot(xs[3]), sv, f3);
6118        }
6119        let acc = [
6120            hsum256_ps(f0),
6121            hsum256_ps(f1),
6122            hsum256_ps(f2),
6123            hsum256_ps(f3),
6124        ];
6125        acc
6126    }
6127}
6128
6129/// ARM twin of `dot_q4t_row_1x4_avx2`: one nibble unpack per group
6130/// serves FOUR activation streams. Per stream the group order and f32
6131/// accumulation match `dot_q4t_row_sdot` exactly — batch == matvec
6132/// bit-for-bit.
6133#[cfg(target_arch = "aarch64")]
6134#[target_feature(enable = "neon,dotprod")]
6135unsafe fn dot_q4t_row_1x4_sdot(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
6136    // SAFETY: callers uphold the 18B-tile and xq-length contracts.
6137    unsafe {
6138        use core::arch::aarch64::*;
6139        use core::arch::asm;
6140        let lomask = vdupq_n_u8(0x0F);
6141        let eight = vdupq_n_s8(8);
6142        let mut acc = [0f32; 4];
6143        for gi in 0..gpr {
6144            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
6145            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
6146            let b = vld1q_u8(t.add(2));
6147            let lo = vandq_u8(b, lomask);
6148            let hi = vshrq_n_u8::<4>(b);
6149            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
6150            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
6151            for (k, xq) in xs.iter().enumerate() {
6152                let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
6153                let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
6154                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6155                asm!(
6156                    "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
6157                    "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
6158                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6159                    e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
6160                    options(pure, nomem, nostack),
6161                );
6162                acc[k] += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
6163            }
6164        }
6165        acc
6166    }
6167}
6168
6169/// Exact-term correction for A8W8 outliers on a tiled row.
6170#[inline]
6171fn q4t_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
6172    let gi = j / GROUP_SIZE;
6173    let k = j % GROUP_SIZE;
6174    let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
6175    let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
6176    let byte = tile[2 + k / 2];
6177    let nib = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
6178    ((nib as i32 - 8) as f32, s)
6179}
6180
6181/// Exact scalar q4_tiled row (CMF_SDOT=0 contract) — same pairwise
6182/// accumulation shape as `q4_range_f32`.
6183#[inline]
6184fn q4t_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
6185    let mut acc = 0f32;
6186    for gi in 0..gpr {
6187        let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
6188        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
6189        let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
6190        let mut ga = 0f32;
6191        for (k, &b) in tile[2..].iter().enumerate() {
6192            ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
6193                + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
6194        }
6195        acc += ga * s;
6196    }
6197    acc
6198}
6199
6200/// Split view of a `q4tp` payload. The three planes are resolved once per
6201/// matvec instead of per row — `q4tp_sections` is cheap, but doing it inside
6202/// the row loop would put a division on the hot path for nothing.
6203/// Where row `r` of a coded matrix lives inside its bytes, as up to three
6204/// (offset, length) ranges (a zero length marks an unused one). A q4t row is
6205/// one run of tiles; a q4tp row is its nibbles, its two ladder parameters
6206/// and its scale codes, in three sections of the tensor.
6207fn row_span_offsets(
6208    dtype: TensorDtype,
6209    rows: usize,
6210    cols: usize,
6211    r: usize,
6212) -> Option<[(usize, usize); 3]> {
6213    if r >= rows || cols % GROUP_SIZE != 0 {
6214        return None;
6215    }
6216    let gpr = cols / GROUP_SIZE;
6217    match dtype {
6218        TensorDtype::Q4Tiled => Some([(r * gpr * Q4_TILE, gpr * Q4_TILE), (0, 0), (0, 0)]),
6219        TensorDtype::Q4TiledP => {
6220            let (params, codes, stride) = q4tp_sections(rows, cols);
6221            Some([
6222                (r * gpr * Q4TP_NIB, gpr * Q4TP_NIB),
6223                (params + r * 4, 4),
6224                (codes + r * stride, stride),
6225            ])
6226        }
6227        _ => None,
6228    }
6229}
6230
6231struct Q4tpView<'a> {
6232    nib: &'a [u8],
6233    params: &'a [u8],
6234    codes: &'a [u8],
6235    stride: usize,
6236    /// q2tp reads the ladder with rung 0 = exact zero.
6237    zero_rung: bool,
6238}
6239
6240impl<'a> Q4tpView<'a> {
6241    fn new(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
6242        let (params_off, codes_off, stride) = q4tp_sections(rows, cols);
6243        Self {
6244            nib: &bytes[..params_off],
6245            params: &bytes[params_off..codes_off],
6246            codes: &bytes[codes_off..],
6247            stride,
6248            zero_rung: false,
6249        }
6250    }
6251
6252    /// The q2tp view: identical params/codes planes, 8 B weight chunks.
6253    fn new_q2(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
6254        let (params_off, codes_off, stride) = q2tp_sections(rows, cols);
6255        Self {
6256            nib: &bytes[..params_off],
6257            params: &bytes[params_off..codes_off],
6258            codes: &bytes[codes_off..],
6259            stride,
6260            zero_rung: true,
6261        }
6262    }
6263
6264    /// Expand row `r`'s per-tile scales into `out` (length `gpr`).
6265    ///
6266    /// Doing this once per row — rather than decoding a 5-bit code inside the
6267    /// tile loop — is what makes the format free at runtime. Random access to
6268    /// a packed 5-bit field costs a division, two bounds checks and a branch;
6269    /// the tile's actual work is two `sdot`s, so per-tile decoding dominated
6270    /// the kernel and cost 5x (measured: 1.4 vs 6.9 tok/s on Nanbeige-3B).
6271    /// Walking the plane sequentially with a bit accumulator is ~3 ops.
6272    /// Eight 5-bit codes are exactly five bytes, so a whole group of
6273    /// eight decodes from one little-endian word at fixed shifts. The
6274    /// bit-accumulator this replaces carried a data-dependent `while
6275    /// have < 5` refill whose branch sat in the innermost loop of every
6276    /// q4tp row; a decode profile put this function above the dot
6277    /// products it feeds. Same bitstream, same codes — just no branch
6278    /// and eight independent extractions.
6279    /// `avx2_enabled()` and `ladder_4` are compiled only for x86-64, so a
6280    /// plain `cfg!(target_arch = ...)` guard is NOT enough: `cfg!` is a
6281    /// runtime `false` on other targets, yet both arms still get compiled
6282    /// and the x86-only symbols do not exist there at all (aarch64 build:
6283    /// E0425). This associated fn IS a compile-time `cfg` — the whole
6284    /// branch, symbol references included, is dropped on non-x86 targets,
6285    /// and the non-x86 arm falls straight through to the scalar ladder.
6286    #[cfg(target_arch = "x86_64")]
6287    #[inline]
6288    fn ladder_via_avx2(&self) -> bool {
6289        avx2_enabled()
6290    }
6291    #[cfg(not(target_arch = "x86_64"))]
6292    #[inline]
6293    fn ladder_via_avx2(&self) -> bool {
6294        false
6295    }
6296
6297    #[inline]
6298    fn scales_into(&self, r: usize, gpr: usize, out: &mut [f32]) {
6299        let tab = if self.zero_rung {
6300            q2tp_ladder(self.params, r)
6301        } else if self.ladder_via_avx2() {
6302            // One lane of the 4-wide vector ladder (the other three are
6303            // duplicates) — every q4tp scale in the process then comes
6304            // from the SAME ladder bits, whichever call site decoded it.
6305            // `ladder_via_avx2()` is cfg-gated, so on non-x86 targets
6306            // this arm is never taken and never even compiled with the
6307            // x86-only symbols in scope.
6308            #[cfg(target_arch = "x86_64")]
6309            let tab = {
6310                let p = &self.params[r * 4..r * 4 + 4];
6311                let lo = f16_to_f32(u16::from_le_bytes([p[0], p[1]]));
6312                let st = f16_to_f32(u16::from_le_bytes([p[2], p[3]]));
6313                let mut tabc = [0f32; 128];
6314                unsafe { ladder_4(&[lo; 4], &[st; 4], &mut tabc) };
6315                let mut t = [0f32; 32];
6316                for (g, o) in t.iter_mut().enumerate() {
6317                    *o = tabc[g * 4];
6318                }
6319                t
6320            };
6321            // On a non-x86 target the arm is unreachable (the `cfg!`
6322            // above is false), so it decodes through the scalar ladder —
6323            // the very path this branch would have fallen back to.
6324            #[cfg(not(target_arch = "x86_64"))]
6325            let tab = q4tp_ladder(self.params, r);
6326            tab
6327        } else {
6328            q4tp_ladder(self.params, r)
6329        };
6330        let codes = &self.codes[r * self.stride..(r + 1) * self.stride];
6331        let out = &mut out[..gpr];
6332        let mut chunks = out.chunks_exact_mut(8);
6333        let mut ci = 0usize;
6334        for c in &mut chunks {
6335            let w = u64::from(codes[ci])
6336                | u64::from(codes[ci + 1]) << 8
6337                | u64::from(codes[ci + 2]) << 16
6338                | u64::from(codes[ci + 3]) << 24
6339                | u64::from(codes[ci + 4]) << 32;
6340            for (k, o) in c.iter_mut().enumerate() {
6341                *o = tab[((w >> (5 * k)) & 31) as usize];
6342            }
6343            ci += 5;
6344        }
6345        // Fewer than eight codes left: the shared total accessor, which
6346        // tolerates a 5-bit field whose spill byte is past the stride.
6347        let tail = &codes[ci..];
6348        for (k, o) in chunks.into_remainder().iter_mut().enumerate() {
6349            *o = tab[q4tp_code(tail, k)];
6350        }
6351    }
6352}
6353
6354#[inline]
6355fn dot_q4tp_row_i8(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
6356    #[cfg(target_arch = "aarch64")]
6357    unsafe {
6358        return dot_q4tp_row_sdot(nib, r, gpr, xq, scales);
6359    }
6360    #[cfg(target_arch = "x86_64")]
6361    unsafe {
6362        if vnni_tiles_enabled() {
6363            return dot_q4tp_row_vnni(nib, r, gpr, xq, scales);
6364        }
6365        return dot_q4tp_row_avx2(nib, r, gpr, xq, scales);
6366    }
6367    #[allow(unreachable_code)]
6368    {
6369        let mut acc = 0f32;
6370        for gi in 0..gpr {
6371            let tile = &nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
6372            let s = scales[gi];
6373            let mut d = 0i32;
6374            for (k, &b) in tile.iter().enumerate() {
6375                d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
6376                    + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
6377            }
6378            acc += d as f32 * s;
6379        }
6380        acc
6381    }
6382}
6383
6384/// q4tp twin of `dot_q4t_row_sdot`: identical nibble math, but the tile
6385/// stride is 16 B (no inline scale) and the scale is a ladder lookup.
6386#[cfg(target_arch = "aarch64")]
6387#[target_feature(enable = "neon,dotprod")]
6388unsafe fn dot_q4tp_row_sdot(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
6389    // SAFETY: callers uphold slice-length contracts (16B tile per group,
6390    // xq.len() == gpr·GROUP_SIZE, codes covering gpr 5-bit fields).
6391    unsafe {
6392        use core::arch::aarch64::*;
6393        use core::arch::asm;
6394        let lomask = vdupq_n_u8(0x0F);
6395        let eight = vdupq_n_s8(8);
6396        let mut acc = 0f32;
6397        for gi in 0..gpr {
6398            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
6399            let s = *scales.get_unchecked(gi);
6400            let b = vld1q_u8(t);
6401            let lo = vandq_u8(b, lomask);
6402            let hi = vshrq_n_u8::<4>(b);
6403            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
6404            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
6405            let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
6406            let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
6407            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6408            asm!(
6409                "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
6410                "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
6411                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6412                e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
6413                options(pure, nomem, nostack),
6414            );
6415            acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
6416        }
6417        acc
6418    }
6419}
6420
6421#[cfg(target_arch = "x86_64")]
6422#[target_feature(enable = "avx2")]
6423unsafe fn dot_q4tp_row_avx2(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
6424    // SAFETY: see dot_q4tp_row_sdot.
6425    unsafe {
6426        use core::arch::x86_64::*;
6427        let lomask = _mm_set1_epi8(0x0F);
6428        let eight = _mm256_set1_epi8(8);
6429        let ones = _mm256_set1_epi16(1);
6430        let mut acc = 0f32;
6431        for gi in 0..gpr {
6432            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
6433            let s = *scales.get_unchecked(gi);
6434            let b = _mm_loadu_si128(t as *const __m128i);
6435            let lo = _mm_and_si128(b, lomask);
6436            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
6437            let w = _mm256_sub_epi8(
6438                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
6439                eight,
6440            );
6441            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6442            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
6443            let d = _mm256_madd_epi16(p16, ones);
6444            let hi128 = _mm256_extracti128_si256::<1>(d);
6445            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
6446            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
6447            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
6448            acc += _mm_cvtsi128_si32(s32) as f32 * s;
6449        }
6450        acc
6451    }
6452}
6453
6454/// VNNI twin of `dot_q4tp_row_avx2` (see `dot_q4t_row_vnni` for why the
6455/// 256-bit VL encoding is the one to use here).
6456#[cfg(target_arch = "x86_64")]
6457#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
6458unsafe fn dot_q4tp_row_vnni(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
6459    // SAFETY: see dot_q4tp_row_sdot.
6460    unsafe {
6461        use core::arch::x86_64::*;
6462        let lomask = _mm_set1_epi8(0x0F);
6463        let eight = _mm256_set1_epi8(8);
6464        let mut acc = 0f32;
6465        for gi in 0..gpr {
6466            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
6467            let s = *scales.get_unchecked(gi);
6468            let b = _mm_loadu_si128(t as *const __m128i);
6469            let lo = _mm_and_si128(b, lomask);
6470            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
6471            let w = _mm256_sub_epi8(
6472                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
6473                eight,
6474            );
6475            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6476            acc += dpbusd_hsum(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w)) as f32 * s;
6477        }
6478        acc
6479    }
6480}
6481
6482/// `2^x` for a 4-lane vector (range-reduced degree-6 polynomial).
6483/// ≤1 ulp vs libm `exp2f` — f32's own precision rounds the polynomial
6484/// home. Used only by [`q4tp_scales_lanes`]; see the ladder note there.
6485///
6486/// # Safety
6487/// AVX2.
6488#[cfg(target_arch = "x86_64")]
6489#[target_feature(enable = "avx2")]
6490#[inline]
6491unsafe fn exp2_128(x: core::arch::x86_64::__m128) -> core::arch::x86_64::__m128 {
6492    unsafe {
6493        use core::arch::x86_64::*;
6494        // floor via trunc with a below-integer correction.
6495        let tr = _mm_cvttps_epi32(x);
6496        // Compare against cvt(tr), NOT a bit-cast: the bit pattern of a
6497        // negative i32 is a NaN and cmplt against it is always false.
6498        let below = _mm_and_si128(
6499            _mm_castps_si128(_mm_cmplt_ps(x, _mm_cvtepi32_ps(tr))),
6500            _mm_set1_epi32(1),
6501        );
6502        let n = _mm_sub_epi32(tr, below); // floor(x)
6503        let r = _mm_sub_ps(x, _mm_cvtepi32_ps(n)); // x - floor(x) ∈ [0,1)
6504        // 2^r = exp(r·ln2) ≈ Σ aₖrᵏ, aₖ = ln2ᵏ/k!, k = 1..9 (Horner).
6505        // At r ∈ [0,1) the a₁₀ term is 7.1e-9 — below f32's own ulp.
6506        // The coefficients are carried at full decimal precision on
6507        // purpose: truncating them to an f32 spelling (what
6508        // `excessive_precision` suggests) would round the constants before
6509        // they reach the AVX registers. `approx_constant` fires on a₁ =
6510        // LN_2, but the table must stay internally consistent, so both
6511        // lints are allowed for this table only.
6512        #[allow(clippy::excessive_precision, clippy::approx_constant)]
6513        let a = [
6514            6.931_471_805_599_453e-1,
6515            2.402_265_069_591_007e-1,
6516            5.550_410_866_482_158e-2,
6517            9.618_129_107_628_477e-3,
6518            1.333_355_814_642_844e-3,
6519            1.540_353_039_338_161e-4,
6520            1.525_273_380_405_984e-5,
6521            1.321_548_679_014_431e-6,
6522            1.017_808_600_923_969_6e-7,
6523        ];
6524        let mut p = _mm_set1_ps(a[8]);
6525        for &c in a[..8].iter().rev() {
6526            p = _mm_add_ps(_mm_mul_ps(p, r), _mm_set1_ps(c));
6527        }
6528        p = _mm_add_ps(_mm_mul_ps(p, r), _mm_set1_ps(1.0));
6529        // Combine: p ∈ [1,2) carries the mantissa; n the exponent.
6530        // Underflow (n < −126) → 0; our args are f16-derived, so this
6531        // only matters for pathological rows.
6532        // bits(p) already carries exponent 127; adding n<<23 (mod 2^32,
6533        // correct for negative n too) retargets it to 127+n.
6534        let ncl = _mm_min_epi32(
6535            _mm_max_epi32(n, _mm_set1_epi32(-126)),
6536            _mm_set1_epi32(127),
6537        );
6538        let mut bits = _mm_add_epi32(_mm_castps_si128(p), _mm_slli_epi32(ncl, 23));
6539        let over = _mm_cmpgt_epi32(n, _mm_set1_epi32(127));
6540        bits = _mm_or_si128(
6541            _mm_andnot_si128(over, bits),
6542            _mm_and_si128(over, _mm_set1_epi32(0x007F_8000)),
6543        );
6544        let zero = _mm_cmplt_ps(x, _mm_set1_ps(-149.0));
6545        _mm_andnot_ps(zero, _mm_castsi128_ps(bits))
6546    }
6547}
6548
6549/// Horizontal helper: the four lanes of `v` as `[f32; 4]`.
6550#[cfg(target_arch = "x86_64")]
6551#[inline]
6552unsafe fn lanes4(v: core::arch::x86_64::__m128) -> [f32; 4] {
6553    use core::arch::x86_64::*;
6554    unsafe {
6555        [
6556            _mm_cvtss_f32(v),
6557            _mm_cvtss_f32(_mm_shuffle_ps(v, v, 0x55)),
6558            _mm_cvtss_f32(_mm_shuffle_ps(v, v, 0xAA)),
6559            _mm_cvtss_f32(_mm_shuffle_ps(v, v, 0xFF)),
6560        ]
6561    }
6562}
6563
6564/// FOUR rows' 32-run q4tp ladders at once into `[group * 4 + lane]`.
6565/// The chain is `t = exp2(lo); t *= exp2(st)` — per lane this is the
6566/// scalar `q4tp_ladder` recurrence run as `vmulps`, bit-identical; only
6567/// `exp2` is replaced (≤1 ulp, see [`exp2_128`]).
6568///
6569/// # Safety
6570/// AVX2.
6571#[cfg(target_arch = "x86_64")]
6572#[target_feature(enable = "avx2")]
6573#[inline]
6574unsafe fn ladder_4(lo4: &[f32; 4], st4: &[f32; 4], tabc: &mut [f32; 128]) {
6575    unsafe {
6576        use core::arch::x86_64::*;
6577        let ratio = exp2_128(_mm_loadu_ps(st4.as_ptr()));
6578        let mut t = exp2_128(_mm_loadu_ps(lo4.as_ptr()));
6579        tabc[0..4].copy_from_slice(&lanes4(t));
6580        for c in 1..32 {
6581            t = _mm_mul_ps(t, ratio);
6582            tabc[c * 4..c * 4 + 4].copy_from_slice(&lanes4(t));
6583        }
6584    }
6585}
6586
6587/// Decode FOUR rows' q4tp group scales at once. The scalar
6588/// [`Q4tpView::scales_into`] costs ~364 cycles a row — two libm `exp2`
6589/// calls plus the 31-step serial `t *= ratio` ladder chain — and the
6590/// ladder is 17 % of every decode token. Here the chains run as one
6591/// 4-lane vector: `t[c] = t[c-1] * ratio` is `vmulps`, bit-identical to
6592/// the scalar f32 chain (and to `q4tp_ladder`'s definition), so the only
6593/// bit change is `exp2` (≤1 ulp, see [`exp2_128`]); the dequant reference
6594/// shares this function. Output layout is `[row][group]`.
6595///
6596/// Only for `Q4TiledP` views (`!zero_rung`); rows need not be
6597/// contiguous. Falls back per incomplete tail chunk.
6598///
6599/// # Safety
6600/// AVX2; `out` must hold `rows.len() * gpr` floats.
6601#[cfg(target_arch = "x86_64")]
6602#[target_feature(enable = "avx2")]
6603unsafe fn q4tp_scales_lanes(
6604    vs: &[&Q4tpView<'_>],
6605    rows: &[usize],
6606    gpr: usize,
6607    out: &mut [f32],
6608) {
6609    use core::arch::x86_64::*;
6610    debug_assert_eq!(vs.len(), rows.len());
6611    debug_assert!(out.len() >= rows.len() * gpr);
6612    let mut done = 0usize;
6613    for (ci, chunk) in rows.chunks(4).enumerate() {
6614        let vchunk = &vs[ci * 4..ci * 4 + chunk.len()];
6615        if chunk.len() < 4 || vchunk.iter().any(|v| v.zero_rung) {
6616            for (l, &r) in chunk.iter().enumerate() {
6617                vchunk[l].scales_into(r, gpr, &mut out[(done + l) * gpr..(done + l + 1) * gpr]);
6618            }
6619            done += chunk.len();
6620            continue;
6621        }
6622        let (r0, r1, r2, r3) = (chunk[0], chunk[1], chunk[2], chunk[3]);
6623        let v0 = vchunk[0];
6624        let v1 = vchunk[1];
6625        let v2 = vchunk[2];
6626        let v3 = vchunk[3];
6627        // Gather the four (lo, st) f16 pairs — 4 B per row.
6628        let mut lo = [0f32; 4];
6629        let mut st = [0f32; 4];
6630        for (l, r) in [r0, r1, r2, r3].iter().enumerate() {
6631            let o = *r * 4;
6632            let vv = vchunk[l];
6633            lo[l] = cortiq_core::quant::f16_to_f32(u16::from_le_bytes([
6634                vv.params[o],
6635                vv.params[o + 1],
6636            ]));
6637            st[l] = cortiq_core::quant::f16_to_f32(u16::from_le_bytes([
6638                vv.params[o + 2],
6639                vv.params[o + 3],
6640            ]));
6641        }
6642        let lo_v = _mm_loadu_ps(lo.as_ptr());
6643        let st_v = _mm_loadu_ps(st.as_ptr());
6644        let ratio = exp2_128(st_v);
6645        let mut t = exp2_128(lo_v);
6646        // tab layout [group * 4 + lane] — contiguous stores per step.
6647        let mut tabc = [0f32; 32 * 4];
6648        tabc[0..4].copy_from_slice(&[
6649            _mm_cvtss_f32(t),
6650            _mm_cvtss_f32(_mm_shuffle_ps(t, t, 0x55)),
6651            _mm_cvtss_f32(_mm_shuffle_ps(t, t, 0xAA)),
6652            _mm_cvtss_f32(_mm_shuffle_ps(t, t, 0xFF)),
6653        ]);
6654        for c in 1..32 {
6655            t = _mm_mul_ps(t, ratio);
6656            tabc[c * 4..c * 4 + 4].copy_from_slice(&[
6657                _mm_cvtss_f32(t),
6658                _mm_cvtss_f32(_mm_shuffle_ps(t, t, 0x55)),
6659                _mm_cvtss_f32(_mm_shuffle_ps(t, t, 0xAA)),
6660                _mm_cvtss_f32(_mm_shuffle_ps(t, t, 0xFF)),
6661            ]);
6662        }
6663        // Scalar code walks (per row), reading tabc[group * 4 + lane].
6664        let vs4 = [v0, v1, v2, v3];
6665        for (l, &r) in [r0, r1, r2, r3].iter().enumerate() {
6666            let v = vs4[l];
6667            let codes = &v.codes[r * v.stride..(r + 1) * v.stride];
6668            let dst = &mut out[(done + l) * gpr..(done + l + 1) * gpr];
6669            let mut chunks = dst.chunks_exact_mut(8);
6670            let mut ci = 0usize;
6671            for c in &mut chunks {
6672                let w = u64::from(codes[ci])
6673                    | u64::from(codes[ci + 1]) << 8
6674                    | u64::from(codes[ci + 2]) << 16
6675                    | u64::from(codes[ci + 3]) << 24
6676                    | u64::from(codes[ci + 4]) << 32;
6677                for (k, o) in c.iter_mut().enumerate() {
6678                    *o = tabc[(((w >> (5 * k)) & 31) as usize) * 4 + l];
6679                }
6680                ci += 5;
6681            }
6682            let tail = &codes[ci..];
6683            for (k, o) in chunks.into_remainder().iter_mut().enumerate() {
6684                *o = tabc[q4tp_code(tail, k) * 4 + l];
6685            }
6686        }
6687        done += 4;
6688    }
6689}
6690
6691/// Decode-lane count for the q4tp row loops: how many rows one worker
6692/// steps through with independent `acc += d * scale` chains in ONE body.
6693/// The per-group AVX2 dot is a ~17-cycle dependent chain per 32 weights;
6694/// a single row serializes on it, overlapping rows reaches the uop-issue
6695/// bound instead. Each row keeps its own group order, so every output is
6696/// bit-identical to the plain row loop. `CMF_Q4TP_LANES=1` restores it;
6697/// 2/4/8 are the supported lane counts (default 4).
6698fn q4tp_lanes() -> usize {
6699    #[cfg(not(target_arch = "x86_64"))]
6700    {
6701        1
6702    }
6703    #[cfg(target_arch = "x86_64")]
6704    {
6705        static ON: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
6706        *ON.get_or_init(|| {
6707            let n: usize = std::env::var("CMF_Q4TP_LANES")
6708                .ok()
6709                .and_then(|v| v.parse().ok())
6710                .unwrap_or(4);
6711            match n {
6712                1 => 1,
6713                2 => 2,
6714                8 => 8,
6715                _ => 4,
6716            }
6717        })
6718    }
6719}
6720
6721/// One q4tp group tile unpacked to interleaved weights in 0..=16 — the
6722/// raw nibbles, WITHOUT the −8 shift. The correction travels in
6723/// `SplitAct::xg8` instead (each vertical lane subtracts `Σx`, and the
6724/// 8-lane reduction makes that `8·Σx` overall), which lets `vpmaddubsw`
6725/// consume the unsigned codes directly and deletes the per-byte `vpaddb`
6726/// + `vpabsb`/`vpsignb` — two of them port-5 ops in the old row
6727/// kernel's innermost loop.
6728///
6729/// # Safety
6730/// AVX2; `tile` must point at 16 valid weight bytes.
6731#[cfg(target_arch = "x86_64")]
6732#[target_feature(enable = "avx2")]
6733#[inline]
6734unsafe fn q4tp_unpack_w(tile: *const u8) -> core::arch::x86_64::__m256i {
6735    unsafe {
6736        use core::arch::x86_64::*;
6737        let lomask = _mm_set1_epi8(0x0F);
6738        let b = _mm_loadu_si128(tile as *const __m128i);
6739        let lo = _mm_and_si128(b, lomask);
6740        let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
6741        _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi))
6742    }
6743}
6744
6745/// Σ over 32 columns of `w·x` as i32x8 lane partials: `vpmaddubsw`
6746/// (|w| ≤ 15, |x| ≤ 127 — pair sums ≤ 3810, no saturation) then `vpmaddwd`.
6747/// Exact in any association.
6748///
6749/// # Safety
6750/// AVX2; `xp` must point at 32 valid i8 activation bytes.
6751#[cfg(target_arch = "x86_64")]
6752#[target_feature(enable = "avx2")]
6753#[inline]
6754unsafe fn q4tp_madd_wx(
6755    w: core::arch::x86_64::__m256i,
6756    xp: *const i8,
6757) -> core::arch::x86_64::__m256i {
6758    unsafe {
6759        use core::arch::x86_64::*;
6760        let ones = _mm256_set1_epi16(1);
6761        let x = _mm256_loadu_si256(xp as *const __m256i);
6762        _mm256_madd_epi16(_mm256_maddubs_epi16(w, x), ones)
6763    }
6764}
6765
6766/// THE shared end-of-row reduction of the vertical f32 accumulator.
6767/// Every a8w8 q4tp path — the decode lane drivers and the row-exact
6768/// matmat cells — must reduce through this one function, or the
6769/// matvec↔matmat bitwise parity test breaks.
6770#[cfg(target_arch = "x86_64")]
6771#[target_feature(enable = "avx2")]
6772#[inline]
6773unsafe fn q4tp_hsum256(acc: core::arch::x86_64::__m256) -> f32 {
6774    unsafe {
6775        use core::arch::x86_64::*;
6776        let a = _mm_add_ps(_mm256_castps256_ps128(acc), _mm256_extractf128_ps::<1>(acc));
6777        let b = _mm_add_ps(a, _mm_movehl_ps(a, a));
6778        let c = _mm_add_ss(b, _mm_shuffle_ps(b, b, 0x55));
6779        _mm_cvtss_f32(c)
6780    }
6781}
6782
6783/// The `− 8·Σx` term of a q4tp a8w8 row: `8 · Σ_g Σx_g · scale_g`, an f32
6784/// dot over the groups. Every vertical-cell path subtracts it after its
6785/// horizontal sum through this one function, so their bits agree.
6786#[cfg(target_arch = "x86_64")]
6787#[inline]
6788fn q4tp_xg8_corr(xg8: &[f32], sc: &[f32]) -> f32 {
6789    let n = sc.len().min(xg8.len());
6790    8.0 * crate::attention::dot_f32(&xg8[..n], &sc[..n])
6791}
6792
6793/// One (row, activation) cell computed the vertical way: per group,
6794/// `Σ(w+8)·x as f32 · scale` fused into lane-wise f32x8 accumulators, one
6795/// horizontal reduction, then the `8·Σx` term (`q4tp_xg8_corr`). Used by the lane drivers' tails
6796/// and by the row-exact matmat column loop so both land on identical bits.
6797///
6798/// # Safety
6799/// AVX2; `gpr` groups must be covered by `xq`/`sc`/`xg8`.
6800#[cfg(target_arch = "x86_64")]
6801#[target_feature(enable = "avx2,fma")]
6802#[inline]
6803unsafe fn q4tp_cell_dot_v(
6804    nib: *const u8,
6805    r: usize,
6806    gpr: usize,
6807    xp: *const i8,
6808    sc: *const f32,
6809    xg8: *const f32,
6810) -> f32 {
6811    unsafe {
6812        use core::arch::x86_64::*;
6813        let mut acc = _mm256_setzero_ps();
6814        for gi in 0..gpr {
6815            let w = q4tp_unpack_w(nib.add((r * gpr + gi) * Q4TP_NIB));
6816            let d = q4tp_madd_wx(w, xp.add(gi * GROUP_SIZE));
6817            acc = _mm256_fmadd_ps(_mm256_cvtepi32_ps(d), _mm256_set1_ps(*sc.add(gi)), acc);
6818        }
6819        q4tp_hsum256(acc)
6820            - q4tp_xg8_corr(
6821                std::slice::from_raw_parts(xg8, gpr),
6822                std::slice::from_raw_parts(sc, gpr),
6823            )
6824    }
6825}
6826
6827/// Row `r` of a row-exact q4tp matmat against one activation panel (up to
6828/// four columns at a time): each weight group is unpacked ONCE and fed to
6829/// four independent column accumulators — the cell math of
6830/// [`q4tp_cell_dot_v`], gi-ascending, so every cell equals the decode
6831/// matvec's bits while the columns overlap their latency chains and the
6832/// unpack is shared. This is the AVX2 twin of the avx512-only blocked
6833/// kernel: without it, Broadwell prefill re-decodes every row once per
6834/// column (7 tok/s measured against q8_2f's 24).
6835///
6836/// # Safety
6837/// AVX2 (caller gates on `a8w8_enabled()` + `xg8.len() == gpr`);
6838/// `out` cells for this row are this worker's alone.
6839#[cfg(target_arch = "x86_64")]
6840#[target_feature(enable = "avx2,fma")]
6841unsafe fn q4tp_matmat_row_cols_v(
6842    v: &Q4tpView<'_>,
6843    r: usize,
6844    gpr: usize,
6845    acts: &[SplitAct],
6846    sc: &[f32],
6847    abase: usize,
6848    rows: usize,
6849    out: SendMut,
6850) {
6851    use core::arch::x86_64::*;
6852    debug_assert_eq!(sc.len(), gpr);
6853    let n = acts.len();
6854    let mut bi = 0usize;
6855    while bi + 4 <= n {
6856        let a: [&SplitAct; 4] = [&acts[bi], &acts[bi + 1], &acts[bi + 2], &acts[bi + 3]];
6857        let mut accs = [_mm256_setzero_ps(); 4];
6858        for gi in 0..gpr {
6859            let w = q4tp_unpack_w(v.nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB));
6860            let sb = _mm256_set1_ps(*sc.get_unchecked(gi));
6861            for (k, act) in a.iter().enumerate() {
6862                let xp = act.xq.as_ptr().add(gi * GROUP_SIZE);
6863                let d = q4tp_madd_wx(w, xp);
6864                accs[k] = _mm256_fmadd_ps(_mm256_cvtepi32_ps(d), sb, accs[k]);
6865            }
6866        }
6867        for k in 0..4 {
6868            let act = a[k];
6869            let mut acc = (q4tp_hsum256(accs[k]) - q4tp_xg8_corr(&act.xg8, sc)) * act.sx;
6870            for &(j, xv) in &act.outliers {
6871                let (w, s) = q4tp_outlier(v.nib, r, gpr, j, sc);
6872                acc += w * s * xv;
6873            }
6874            // SAFETY: disjoint (bi, r) cells per worker.
6875            unsafe { *out.at((abase + bi + k) * rows + r) = acc };
6876        }
6877        bi += 4;
6878    }
6879    // Remainder columns (1..=3): one vertical cell each.
6880    while bi < n {
6881        let act = &acts[bi];
6882        let mut acc = q4tp_cell_dot_v(
6883            v.nib.as_ptr(),
6884            r,
6885            gpr,
6886            act.xq.as_ptr(),
6887            sc.as_ptr(),
6888            act.xg8.as_ptr(),
6889        ) * act.sx;
6890        for &(j, xv) in &act.outliers {
6891            let (w, s) = q4tp_outlier(v.nib, r, gpr, j, sc);
6892            acc += w * s * xv;
6893        }
6894        // SAFETY: disjoint (bi, r) cells per worker.
6895        unsafe { *out.at((abase + bi) * rows + r) = acc };
6896        bi += 1;
6897    }
6898}
6899
6900/// Columns per activation panel of the q4tp batch: `CMF_Q4TP_PANEL`, else
6901/// 256 — and on x86-64 as many as fill half the core's L2 (`cols` int8
6902/// bytes a column), a multiple of four from 8 to 256. A panel of 256
6903/// columns of a 2560-wide row is 655 KB: on a 512 KB L2 (Zen 1/2) every
6904/// weight row streamed the whole panel from L3 again.
6905fn q4tp_panel_cols(cols: usize) -> usize {
6906    static ENV: std::sync::OnceLock<Option<usize>> = std::sync::OnceLock::new();
6907    if let Some(v) = *ENV.get_or_init(|| {
6908        std::env::var("CMF_Q4TP_PANEL")
6909            .ok()
6910            .and_then(|v| v.parse().ok())
6911            .filter(|v: &usize| *v > 0)
6912    }) {
6913        return v;
6914    }
6915    #[cfg(target_arch = "x86_64")]
6916    {
6917        ((l2_bytes() / 2) / cols.max(1)).clamp(8, 256) & !3
6918    }
6919    #[cfg(not(target_arch = "x86_64"))]
6920    {
6921        let _ = cols;
6922        256
6923    }
6924}
6925
6926/// The per-core L2 size (Linux sysfs; 1 MB when unknown).
6927#[cfg(target_arch = "x86_64")]
6928fn l2_bytes() -> usize {
6929    static L2: std::sync::OnceLock<usize> = std::sync::OnceLock::new();
6930    *L2.get_or_init(|| {
6931        let s = std::fs::read_to_string("/sys/devices/system/cpu/cpu0/cache/index2/size").unwrap_or_default();
6932        let s = s.trim();
6933        let (num, mul) = match s.strip_suffix('K') {
6934            Some(n) => (n, 1024),
6935            None => match s.strip_suffix('M') {
6936                Some(n) => (n, 1024 * 1024),
6937                None => (s, 1),
6938            },
6939        };
6940        num.parse::<usize>().ok().map(|n| n * mul).filter(|&b| b >= 64 * 1024).unwrap_or(1 << 20)
6941    })
6942}
6943
6944/// `CMF_Q4TP_FMA=0` keeps the row-exact vertical cell
6945/// (`q4tp_matmat_row_cols_v`) for a batch on an AVX2 CPU (A/B only).
6946#[cfg(target_arch = "x86_64")]
6947fn q4tp_fma_x86() -> bool {
6948    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
6949    *ON.get_or_init(|| {
6950        std::is_x86_feature_detected!("avx2")
6951            && std::is_x86_feature_detected!("fma")
6952            && std::env::var("CMF_Q4TP_FMA").map(|v| v != "0").unwrap_or(true)
6953    })
6954}
6955
6956/// Two weight rows (`r`, `r + 1`) against a panel of activation columns,
6957/// the tuned AVX2 shape for a batch that need not be row-exact: four
6958/// columns at a time, so each group's tile is unpacked once for four
6959/// columns and each column's 32 bytes are loaded once for two rows (eight
6960/// lane accumulators, the register file's worth). Per group and cell the
6961/// loop is `vpmaddubsw`, `vpmaddwd`, a convert and one FMA by the group
6962/// scale — the vertical cell's `− 8·Σx` correction leaves the loop and is
6963/// taken once per cell as `8 · Σ_g Σx_g · scale_g` (an f32 dot over the
6964/// groups). The same products as the vertical cell, summed in another
6965/// order: tolerance-class against the matvec, like the ARM tuned shape.
6966/// Columns past the last group of four take the vertical cell, row by row.
6967///
6968/// # Safety
6969/// AVX2 and FMA (`q4tp_fma_x86`); every column carries `xg8` of `gpr`
6970/// groups; `out` cells (`(abase + k)·rows + r`, `r + 1`) disjoint per worker.
6971#[cfg(target_arch = "x86_64")]
6972#[target_feature(enable = "avx2,fma")]
6973#[allow(clippy::too_many_arguments)]
6974unsafe fn q4tp_matmat_2rows_fma(
6975    v: &Q4tpView<'_>,
6976    r: usize,
6977    gpr: usize,
6978    acts: &[SplitAct],
6979    sc0: &[f32],
6980    sc1: &[f32],
6981    abase: usize,
6982    rows: usize,
6983    out: SendMut,
6984) {
6985    use core::arch::x86_64::*;
6986    debug_assert!(sc0.len() == gpr && sc1.len() == gpr);
6987    let n = acts.len();
6988    let nib = v.nib.as_ptr();
6989    let ones = _mm256_set1_epi16(1);
6990    let mut bi = 0usize;
6991    unsafe {
6992        while bi + 4 <= n {
6993            let a: [&SplitAct; 4] = [&acts[bi], &acts[bi + 1], &acts[bi + 2], &acts[bi + 3]];
6994            let xp: [*const i8; 4] = [
6995                a[0].xq.as_ptr(),
6996                a[1].xq.as_ptr(),
6997                a[2].xq.as_ptr(),
6998                a[3].xq.as_ptr(),
6999            ];
7000            let mut c0 = [_mm256_setzero_ps(); 4];
7001            let mut c1 = [_mm256_setzero_ps(); 4];
7002            for gi in 0..gpr {
7003                let w0 = q4tp_unpack_w(nib.add((r * gpr + gi) * Q4TP_NIB));
7004                let w1 = q4tp_unpack_w(nib.add(((r + 1) * gpr + gi) * Q4TP_NIB));
7005                let s0 = _mm256_set1_ps(*sc0.get_unchecked(gi));
7006                let s1 = _mm256_set1_ps(*sc1.get_unchecked(gi));
7007                for k in 0..4 {
7008                    let x = _mm256_loadu_si256(xp[k].add(gi * GROUP_SIZE) as *const __m256i);
7009                    let d0 = _mm256_madd_epi16(_mm256_maddubs_epi16(w0, x), ones);
7010                    let d1 = _mm256_madd_epi16(_mm256_maddubs_epi16(w1, x), ones);
7011                    c0[k] = _mm256_fmadd_ps(_mm256_cvtepi32_ps(d0), s0, c0[k]);
7012                    c1[k] = _mm256_fmadd_ps(_mm256_cvtepi32_ps(d1), s1, c1[k]);
7013                }
7014            }
7015            for (row, cr, scr) in [(r, &c0, sc0), (r + 1, &c1, sc1)] {
7016                for k in 0..4 {
7017                    let act = a[k];
7018                    let mut acc = (q4tp_hsum256(cr[k]) - q4tp_xg8_corr(&act.xg8, scr)) * act.sx;
7019                    for &(j, xv) in &act.outliers {
7020                        let (w, s) = q4tp_outlier(v.nib, row, gpr, j, scr);
7021                        acc += w * s * xv;
7022                    }
7023                    // SAFETY: disjoint (bi, r) cells per worker.
7024                    *out.at((abase + bi + k) * rows + row) = acc;
7025                }
7026            }
7027            bi += 4;
7028        }
7029        if bi < n {
7030            q4tp_matmat_row_cols_v(v, r, gpr, &acts[bi..], sc0, abase + bi, rows, out);
7031            q4tp_matmat_row_cols_v(v, r + 1, gpr, &acts[bi..], sc1, abase + bi, rows, out);
7032        }
7033    }
7034}
7035
7036/// Rows `start..end` of a single-output q4tp a8w8 matvec, `L` independent
7037/// accumulator chains wide. Per row: scales decoded once, groups stepped
7038/// in ascending `gi` with `acc += d as f32 * scale` (the exact plain-row
7039/// order), then the `* sx` scale and the outlier corrections in list
7040/// order — the plain row loop's arithmetic, line for line.
7041///
7042/// # Safety
7043/// AVX2 (callers gate on `a8w8_enabled()`); `out` disjoint per row.
7044#[cfg(target_arch = "x86_64")]
7045#[target_feature(enable = "avx2,fma")]
7046unsafe fn q4tp_matvec_lanes<const L: usize>(
7047    v: &Q4tpView<'_>,
7048    start: usize,
7049    end: usize,
7050    gpr: usize,
7051    xq: &[i8],
7052    xg8: &[f32],
7053    sx: f32,
7054    outliers: &[(usize, f32)],
7055    sc: &mut [f32],
7056    out: SendMut,
7057) {
7058    use core::arch::x86_64::*;
7059    debug_assert!(L >= 2 && sc.len() >= L * gpr && xg8.len() == gpr);
7060    crate::cpuprof::count_lane_rows((end - start) as u64);
7061    let mut base = start;
7062    while base + L <= end {
7063        let _t0 = crate::cpuprof::time(crate::cpuprof::Slot::GuScales);
7064        let mut rows = [0usize; L];
7065        for l in 0..L {
7066            rows[l] = base + l;
7067        }
7068        let v4: [&Q4tpView; L] = [v; L];
7069        q4tp_scales_lanes(&v4, &rows, gpr, sc);
7070        drop(_t0);
7071        let _t1 = crate::cpuprof::time(crate::cpuprof::Slot::GuDots);
7072        let mut accs = [_mm256_setzero_ps(); L];
7073        for gi in 0..gpr {
7074            let xp = xq.as_ptr().add(gi * GROUP_SIZE);
7075            for l in 0..L {
7076                let w = q4tp_unpack_w(v.nib.as_ptr().add((rows[l] * gpr + gi) * Q4TP_NIB));
7077                let d = q4tp_madd_wx(w, xp);
7078                let sb = _mm256_set1_ps(*sc.get_unchecked(l * gpr + gi));
7079                accs[l] = _mm256_fmadd_ps(_mm256_cvtepi32_ps(d), sb, accs[l]);
7080            }
7081        }
7082        drop(_t1);
7083        let _t2 = crate::cpuprof::time(crate::cpuprof::Slot::GuPost);
7084        for l in 0..L {
7085            let r = rows[l];
7086            let scl = &sc[l * gpr..(l + 1) * gpr];
7087            let mut acc = (q4tp_hsum256(accs[l]) - q4tp_xg8_corr(xg8, scl)) * sx;
7088            for &(j, xv) in outliers {
7089                let (w, s) = q4tp_outlier(v.nib, r, gpr, j, scl);
7090                acc += w * s * xv;
7091            }
7092            // SAFETY: one lane owns each row.
7093            unsafe { *out.at(r) = acc };
7094        }
7095        drop(_t2);
7096        base += L;
7097    }
7098    // Partial lane batch (and any tail): one vertical cell per row —
7099    // the same reduction shape as the lanes above and as the row-exact
7100    // matmat, so bits agree across every path.
7101    for r in base..end {
7102        let scl = &mut sc[..gpr];
7103        v.scales_into(r, gpr, scl);
7104        let scl: &[f32] = scl;
7105        let mut acc = q4tp_cell_dot_v(v.nib.as_ptr(), r, gpr, xq.as_ptr(), scl.as_ptr(), xg8.as_ptr())
7106            * sx;
7107        for &(j, xv) in outliers {
7108            let (w, s) = q4tp_outlier(v.nib, r, gpr, j, scl);
7109            acc += w * s * xv;
7110        }
7111        // SAFETY: disjoint row ranges per worker.
7112        unsafe { *out.at(r) = acc };
7113    }
7114}
7115
7116/// Rows `start..end` of the fused SwiGLU gate+up q4tp arm, `L` chains
7117/// wide: both projections' scales decoded per lane, both group loops
7118/// stepped in the plain arm's per-row order, then `silu_mul_limited`
7119/// over the outlier-corrected pair — bit-identical to the plain arm.
7120///
7121/// # Safety
7122/// AVX2 (callers gate on `a8w8_enabled()`); `out` disjoint per row.
7123#[cfg(target_arch = "x86_64")]
7124#[target_feature(enable = "avx2,fma")]
7125unsafe fn q4tp_gateup_lanes<const L: usize>(
7126    gv: &Q4tpView<'_>,
7127    uv: &Q4tpView<'_>,
7128    start: usize,
7129    end: usize,
7130    gpr: usize,
7131    xq: &[i8],
7132    xg8: &[f32],
7133    sx: f32,
7134    limit: f32,
7135    outliers: &[(usize, f32)],
7136    gsc: &mut [f32],
7137    usc: &mut [f32],
7138    out: SendMut,
7139) {
7140    use core::arch::x86_64::*;
7141    debug_assert!(
7142        L >= 2 && gsc.len() >= L * gpr && usc.len() >= L * gpr && xg8.len() == gpr
7143    );
7144    crate::cpuprof::count_lane_rows((end - start) as u64);
7145    let mut base = start;
7146    while base + L <= end {
7147        let _t0 = crate::cpuprof::time(crate::cpuprof::Slot::GuScales);
7148        let mut rows = [0usize; L];
7149        for l in 0..L {
7150            rows[l] = base + l;
7151        }
7152        let gv4: [&Q4tpView; L] = [gv; L];
7153        let uv4: [&Q4tpView; L] = [uv; L];
7154        q4tp_scales_lanes(&gv4, &rows, gpr, gsc);
7155        q4tp_scales_lanes(&uv4, &rows, gpr, usc);
7156        drop(_t0);
7157        let _t1 = crate::cpuprof::time(crate::cpuprof::Slot::GuDots);
7158        let mut ag = [_mm256_setzero_ps(); L];
7159        let mut au = [_mm256_setzero_ps(); L];
7160        for gi in 0..gpr {
7161            let xp = xq.as_ptr().add(gi * GROUP_SIZE);
7162            for l in 0..L {
7163                let w = q4tp_unpack_w(gv.nib.as_ptr().add((rows[l] * gpr + gi) * Q4TP_NIB));
7164                let d = q4tp_madd_wx(w, xp);
7165                let sb = _mm256_set1_ps(*gsc.get_unchecked(l * gpr + gi));
7166                ag[l] = _mm256_fmadd_ps(_mm256_cvtepi32_ps(d), sb, ag[l]);
7167            }
7168            for l in 0..L {
7169                let w = q4tp_unpack_w(uv.nib.as_ptr().add((rows[l] * gpr + gi) * Q4TP_NIB));
7170                let d = q4tp_madd_wx(w, xp);
7171                let sb = _mm256_set1_ps(*usc.get_unchecked(l * gpr + gi));
7172                au[l] = _mm256_fmadd_ps(_mm256_cvtepi32_ps(d), sb, au[l]);
7173            }
7174        }
7175        drop(_t1);
7176        let _t2 = crate::cpuprof::time(crate::cpuprof::Slot::GuPost);
7177        for l in 0..L {
7178            let r = rows[l];
7179            let gs = &gsc[l * gpr..(l + 1) * gpr];
7180            let us = &usc[l * gpr..(l + 1) * gpr];
7181            let mut g = (q4tp_hsum256(ag[l]) - q4tp_xg8_corr(xg8, gs)) * sx;
7182            let mut u = (q4tp_hsum256(au[l]) - q4tp_xg8_corr(xg8, us)) * sx;
7183            for &(j, xv) in outliers {
7184                let (w, s) = q4tp_outlier(gv.nib, r, gpr, j, gs);
7185                g += w * s * xv;
7186                let (w, s) = q4tp_outlier(uv.nib, r, gpr, j, us);
7187                u += w * s * xv;
7188            }
7189            // SAFETY: one lane owns each row.
7190            unsafe { *out.at(r) = silu_mul_limited(g, u, limit) };
7191        }
7192        drop(_t2);
7193        base += L;
7194    }
7195    // Partial lane batch (and any tail): one vertical cell per stream.
7196    for r in base..end {
7197        let g = &mut gsc[..gpr];
7198        let u = &mut usc[..gpr];
7199        gv.scales_into(r, gpr, g);
7200        uv.scales_into(r, gpr, u);
7201        let (g, u): (&[f32], &[f32]) = (g, u);
7202        let mut gv_ =
7203            q4tp_cell_dot_v(gv.nib.as_ptr(), r, gpr, xq.as_ptr(), g.as_ptr(), xg8.as_ptr()) * sx;
7204        let mut uv_ =
7205            q4tp_cell_dot_v(uv.nib.as_ptr(), r, gpr, xq.as_ptr(), u.as_ptr(), xg8.as_ptr()) * sx;
7206        for &(j, xv) in outliers {
7207            let (w, s) = q4tp_outlier(gv.nib, r, gpr, j, g);
7208            gv_ += w * s * xv;
7209            let (w, s) = q4tp_outlier(uv.nib, r, gpr, j, u);
7210            uv_ += w * s * xv;
7211        }
7212        // SAFETY: disjoint row ranges per worker.
7213        unsafe { *out.at(r) = silu_mul_limited(gv_, uv_, limit) };
7214    }
7215}
7216
7217/// Flat-space sibling of [`q4tp_matvec_lanes`] for `matvec_many`: the
7218/// virtual row space spans several tensors, so each lane resolves its
7219/// (tensor, row) pair through the caller's `locate`, then runs the plain
7220/// per-row math of the fused q4tp arm.
7221///
7222/// # Safety
7223/// AVX2 (callers gate on `a8w8_enabled()`); `outs` disjoint per row.
7224#[cfg(target_arch = "x86_64")]
7225#[target_feature(enable = "avx2,fma")]
7226unsafe fn q4tp_many_lanes<const L: usize>(
7227    views: &[Q4tpView<'_>],
7228    locate: &(dyn Fn(usize) -> (usize, usize)),
7229    start: usize,
7230    end: usize,
7231    gpr: usize,
7232    xq: &[i8],
7233    xg8: &[f32],
7234    sx: f32,
7235    outliers: &[(usize, f32)],
7236    sc: &mut [f32],
7237    outs: &[SendMut],
7238) {
7239    use core::arch::x86_64::*;
7240    debug_assert!(L >= 2 && sc.len() >= L * gpr && xg8.len() == gpr);
7241    crate::cpuprof::count_lane_rows((end - start) as u64);
7242    let mut base = start;
7243    while base + L <= end {
7244        let _t0 = crate::cpuprof::time(crate::cpuprof::Slot::GuScales);
7245        let mut tk = [0usize; L];
7246        let mut rows = [0usize; L];
7247        for l in 0..L {
7248            let (t, r) = locate(base + l);
7249            tk[l] = t;
7250            rows[l] = r;
7251        }
7252        let v4: [&Q4tpView; L] = std::array::from_fn(|l| &views[tk[l]]);
7253        q4tp_scales_lanes(&v4, &rows, gpr, sc);
7254        drop(_t0);
7255        let _t1 = crate::cpuprof::time(crate::cpuprof::Slot::GuDots);
7256        let mut accs = [_mm256_setzero_ps(); L];
7257        for gi in 0..gpr {
7258            let xp = xq.as_ptr().add(gi * GROUP_SIZE);
7259            for l in 0..L {
7260                let w = q4tp_unpack_w(
7261                    views[tk[l]].nib.as_ptr().add((rows[l] * gpr + gi) * Q4TP_NIB),
7262                );
7263                let d = q4tp_madd_wx(w, xp);
7264                let sb = _mm256_set1_ps(*sc.get_unchecked(l * gpr + gi));
7265                accs[l] = _mm256_fmadd_ps(_mm256_cvtepi32_ps(d), sb, accs[l]);
7266            }
7267        }
7268        drop(_t1);
7269        let _t2 = crate::cpuprof::time(crate::cpuprof::Slot::GuPost);
7270        for l in 0..L {
7271            let (t, r) = (tk[l], rows[l]);
7272            let scl = &sc[l * gpr..(l + 1) * gpr];
7273            let mut acc = (q4tp_hsum256(accs[l]) - q4tp_xg8_corr(xg8, scl)) * sx;
7274            for &(j, xv) in outliers {
7275                let (w, s) = q4tp_outlier(views[t].nib, r, gpr, j, scl);
7276                acc += w * s * xv;
7277            }
7278            // SAFETY: one lane owns each (tensor, row) pair.
7279            unsafe { *outs[t].at(r) = acc };
7280        }
7281        drop(_t2);
7282        base += L;
7283    }
7284    // Partial lane batch (and any tail): one vertical cell per row.
7285    for flat in base..end {
7286        let (t, r) = locate(flat);
7287        let v = &views[t];
7288        let scl = &mut sc[..gpr];
7289        v.scales_into(r, gpr, scl);
7290        let scl: &[f32] = scl;
7291        let mut acc =
7292            q4tp_cell_dot_v(v.nib.as_ptr(), r, gpr, xq.as_ptr(), scl.as_ptr(), xg8.as_ptr()) * sx;
7293        for &(j, xv) in outliers {
7294            let (w, s) = q4tp_outlier(v.nib, r, gpr, j, scl);
7295            acc += w * s * xv;
7296        }
7297        // SAFETY: one worker owns each (tensor, row) pair.
7298        unsafe { *outs[t].at(r) = acc };
7299    }
7300}
7301
7302/// Exact scalar q4tp row — the `CMF_SDOT=0` contract, same pairwise
7303/// accumulation shape as `q4t_row_exact`.
7304#[inline]
7305fn q4tp_row_exact(nib: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
7306    #[cfg(target_arch = "x86_64")]
7307    if avx2_enabled() {
7308        // Keep the scalar pair/group reduction order, not a vector sum.
7309        return unsafe { q4tp_row_float_avx2(nib, r, gpr, x, scales) };
7310    }
7311    q4tp_row_float_scalar(nib, r, gpr, x, scales)
7312}
7313
7314#[inline]
7315fn q4tp_row_float_scalar(nib: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
7316    let mut acc = 0f32;
7317    for gi in 0..gpr {
7318        let tile = &nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
7319        let s = scales[gi];
7320        let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
7321        let mut ga = 0f32;
7322        for (k, &b) in tile.iter().enumerate() {
7323            ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
7324                + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
7325        }
7326        acc += ga * s;
7327    }
7328    acc
7329}
7330
7331/// Vectorize unpack, conversion and multiplication, but preserve every
7332/// pair addition and the scalar accumulation order. No activation rounding
7333/// or FMA: bit-identical to the float scalar row, including its group scale.
7334#[cfg(target_arch = "x86_64")]
7335#[target_feature(enable = "avx2")]
7336unsafe fn q4tp_row_float_avx2(nib: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
7337    // SAFETY: caller checks AVX2 and provides the same complete 32-element
7338    // groups as the scalar row. Loads/stores are explicitly unaligned.
7339    unsafe {
7340        use core::arch::x86_64::*;
7341        let mask = _mm_set1_epi8(15);
7342        let eight = _mm_set1_epi8(8);
7343        let order = _mm256_setr_epi32(0, 1, 4, 5, 2, 3, 6, 7);
7344        let mut acc = 0.0f32;
7345        for gi in 0..gpr {
7346            let packed = _mm_loadu_si128(nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB).cast());
7347            let lo = _mm_and_si128(packed, mask);
7348            let hi = _mm_and_si128(_mm_srli_epi16::<4>(packed), mask);
7349            let w0 = _mm_sub_epi8(_mm_unpacklo_epi8(lo, hi), eight);
7350            let w1 = _mm_sub_epi8(_mm_unpackhi_epi8(lo, hi), eight);
7351            let xp = x.as_ptr().add(gi * GROUP_SIZE);
7352            let a = _mm256_mul_ps(
7353                _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(w0)),
7354                _mm256_loadu_ps(xp),
7355            );
7356            let b = _mm256_mul_ps(
7357                _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(w0))),
7358                _mm256_loadu_ps(xp.add(8)),
7359            );
7360            let c = _mm256_mul_ps(
7361                _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(w1)),
7362                _mm256_loadu_ps(xp.add(16)),
7363            );
7364            let d = _mm256_mul_ps(
7365                _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(w1))),
7366                _mm256_loadu_ps(xp.add(24)),
7367            );
7368            let mut pairs = [0.0f32; 16];
7369            _mm256_storeu_ps(
7370                pairs.as_mut_ptr(),
7371                _mm256_permutevar8x32_ps(_mm256_hadd_ps(a, b), order),
7372            );
7373            _mm256_storeu_ps(
7374                pairs.as_mut_ptr().add(8),
7375                _mm256_permutevar8x32_ps(_mm256_hadd_ps(c, d), order),
7376            );
7377            let mut ga = 0.0f32;
7378            for v in pairs {
7379                ga += v;
7380            }
7381            acc += ga * scales[gi];
7382        }
7383        acc
7384    }
7385}
7386
7387/// Single weight of a q4tp tensor — the a8w8 outlier path, which restores
7388/// activation outliers at full precision after the int8 pass.
7389#[inline]
7390fn q4tp_outlier(nib: &[u8], r: usize, gpr: usize, j: usize, scales: &[f32]) -> (f32, f32) {
7391    let (gi, k) = (j / GROUP_SIZE, j % GROUP_SIZE);
7392    let byte = nib[(r * gpr + gi) * Q4TP_NIB + k / 2];
7393    let n = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
7394    ((n as i32 - 8) as f32, scales[gi])
7395}
7396
7397/// Fused q4tp matvec (dispatch mirrors `q4t_matvec`).
7398fn q4tp_matvec(
7399    bytes: &[u8],
7400    x: &[f32],
7401    rows: usize,
7402    cols: usize,
7403    out: &mut [f32],
7404    pool: Option<&Pool>,
7405) {
7406    debug_assert_eq!(out.len(), rows);
7407    let gpr = cols / GROUP_SIZE;
7408    let v = Q4tpView::new(bytes, rows, cols);
7409    let out_addr = SendMut(out.as_mut_ptr());
7410    if a8w8_enabled() {
7411        let act = split_act(x);
7412        #[cfg(target_arch = "x86_64")]
7413        let lanes = q4tp_lanes();
7414        let run = |start: usize, end: usize| {
7415            #[cfg(target_arch = "x86_64")]
7416            if lanes > 1 && act.xg8.len() == gpr {
7417                // L independent accumulator chains per worker — the plain
7418                // row body's math, lane-interleaved.
7419                with_krow(gpr * lanes, |sc| {
7420                    // SAFETY: a8w8 implies AVX2; disjoint row ranges.
7421                    unsafe {
7422                        match lanes {
7423                            2 => q4tp_matvec_lanes::<2>(
7424                                &v, start, end, gpr, &act.xq, &act.xg8, act.sx,
7425                                &act.outliers, sc, out_addr,
7426                            ),
7427                            8 => q4tp_matvec_lanes::<8>(
7428                                &v, start, end, gpr, &act.xq, &act.xg8, act.sx,
7429                                &act.outliers, sc, out_addr,
7430                            ),
7431                            _ => q4tp_matvec_lanes::<4>(
7432                                &v, start, end, gpr, &act.xq, &act.xg8, act.sx,
7433                                &act.outliers, sc, out_addr,
7434                            ),
7435                        }
7436                    }
7437                });
7438                return;
7439            }
7440            crate::cpuprof::count_plain_rows((end - start) as u64);
7441            // One scratch row of scales per worker — borrowed, not minted.
7442            with_krow(gpr, |sc| {
7443                for r in start..end {
7444                    v.scales_into(r, gpr, sc);
7445                    let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, sc) * act.sx;
7446                    for &(j, xv) in &act.outliers {
7447                        let (w, s) = q4tp_outlier(v.nib, r, gpr, j, sc);
7448                        acc += w * s * xv;
7449                    }
7450                    // SAFETY: disjoint row ranges per worker.
7451                    unsafe { *out_addr.at(r) = acc }
7452                }
7453            })
7454        };
7455        dispatch_rows(pool, rows, &run);
7456        return;
7457    }
7458    let run = |start: usize, end: usize| {
7459        with_krow(gpr, |sc| {
7460            for r in start..end {
7461                v.scales_into(r, gpr, sc);
7462                // SAFETY: disjoint row ranges per worker.
7463                unsafe { *out_addr.at(r) = q4tp_row_exact(v.nib, r, gpr, x, sc) };
7464            }
7465        })
7466    };
7467    dispatch_rows(pool, rows, &run);
7468}
7469
7470/// Fused two-input q4tp matvec — the SwiGLU gate/up pair. Weights and the
7471/// row ladder are read once and spent on both activation streams.
7472#[allow(clippy::too_many_arguments)]
7473fn q4tp_matvec2(
7474    bytes: &[u8],
7475    x1: &[f32],
7476    x2: &[f32],
7477    rows: usize,
7478    cols: usize,
7479    o1: &mut [f32],
7480    o2: &mut [f32],
7481    pool: Option<&Pool>,
7482) {
7483    let gpr = cols / GROUP_SIZE;
7484    let v = Q4tpView::new(bytes, rows, cols);
7485    let (p1, p2) = (SendMut(o1.as_mut_ptr()), SendMut(o2.as_mut_ptr()));
7486    let run = |start: usize, end: usize| {
7487        let mut sc = vec![0f32; gpr];
7488        for r in start..end {
7489            v.scales_into(r, gpr, &mut sc);
7490            // SAFETY: disjoint row ranges per worker.
7491            unsafe {
7492                *p1.at(r) = q4tp_row_exact(v.nib, r, gpr, x1, &sc);
7493                *p2.at(r) = q4tp_row_exact(v.nib, r, gpr, x2, &sc);
7494            }
7495        }
7496    };
7497    dispatch_rows(pool, rows, &run);
7498}
7499
7500/// One q2tp outlier weight at column `j` of row `r`: the 2-bit code and
7501/// its group scale, mirrored on `q4tp_outlier`.
7502#[inline]
7503fn q2tp_outlier(chunks: &[u8], r: usize, gpr: usize, j: usize, scales: &[f32]) -> (f32, f32) {
7504    let (gi, k) = (j / GROUP_SIZE, j % GROUP_SIZE);
7505    let byte = chunks[(r * gpr + gi) * Q2TP_CHUNK + k / 4];
7506    let c = (byte >> (2 * (k % 4))) & 3;
7507    (c as f32 - 1.5, scales[gi])
7508}
7509
7510#[cfg(target_arch = "x86_64")]
7511const Q2TP_DECODE_U32: [u32; 256] = {
7512    let mut tab = [0u32; 256];
7513    let mut b = 0usize;
7514    while b < 256 {
7515        tab[b] = ((b as u32) & 3)
7516            | ((((b as u32) >> 2) & 3) << 8)
7517            | ((((b as u32) >> 4) & 3) << 16)
7518            | ((((b as u32) >> 6) & 3) << 24);
7519        b += 1;
7520    }
7521    tab
7522};
7523
7524/// Eight packed q2tp bytes against 32 signed activation bytes. `maddubs`
7525/// exactly computes unsigned 2-bit code × signed i8; its pair sums cannot
7526/// saturate (2 × 3 × 127 < i16::MAX), and the second madd widens to i32.
7527#[cfg(target_arch = "x86_64")]
7528#[target_feature(enable = "avx2")]
7529unsafe fn q2tp_code_dot_avx2(ch: &[u8], x: &[i8]) -> i32 {
7530    use core::arch::x86_64::*;
7531    debug_assert!(ch.len() >= Q2TP_CHUNK && x.len() >= GROUP_SIZE);
7532    let codes = _mm256_setr_epi32(
7533        Q2TP_DECODE_U32[ch[0] as usize] as i32,
7534        Q2TP_DECODE_U32[ch[1] as usize] as i32,
7535        Q2TP_DECODE_U32[ch[2] as usize] as i32,
7536        Q2TP_DECODE_U32[ch[3] as usize] as i32,
7537        Q2TP_DECODE_U32[ch[4] as usize] as i32,
7538        Q2TP_DECODE_U32[ch[5] as usize] as i32,
7539        Q2TP_DECODE_U32[ch[6] as usize] as i32,
7540        Q2TP_DECODE_U32[ch[7] as usize] as i32,
7541    );
7542    let xv = unsafe { _mm256_loadu_si256(x.as_ptr().cast()) };
7543    let pair = _mm256_maddubs_epi16(codes, xv);
7544    let quad = _mm256_madd_epi16(pair, _mm256_set1_epi16(1));
7545    let sum128 = _mm_add_epi32(
7546        _mm256_castsi256_si128(quad),
7547        _mm256_extracti128_si256(quad, 1),
7548    );
7549    let sum64 = _mm_hadd_epi32(sum128, sum128);
7550    _mm_cvtsi128_si32(_mm_hadd_epi32(sum64, sum64))
7551}
7552
7553#[cfg(target_arch = "x86_64")]
7554#[inline]
7555fn q2tp_avx2_enabled() -> bool {
7556    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
7557    *ON.get_or_init(|| std::arch::is_x86_feature_detected!("avx2"))
7558}
7559
7560/// Integer dot of one q2tp row against pre-quantized activations:
7561/// Σ_g s_g · (Σ c·xq − 1.5·Σ xq). The half-integer grid (c − 1.5)
7562/// becomes exact integer math through the group sums — the same trick
7563/// every a8w8 kernel in this file rides. The codes decode into a
7564/// 32-byte scratch in natural order and the dot itself is the shared
7565/// SDOT primitive; elsewhere a scalar integer loop.
7566#[inline]
7567fn dot_q2tp_row_i8(
7568    chunks: &[u8],
7569    r: usize,
7570    gpr: usize,
7571    xq: &[i8],
7572    gsum: &[i32],
7573    scales: &[f32],
7574) -> f32 {
7575    let mut acc = 0f32;
7576    let base = r * gpr * Q2TP_CHUNK;
7577    #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
7578    let mut codes = [0i8; GROUP_SIZE];
7579    #[cfg(target_arch = "x86_64")]
7580    // Cache the process-wide AVX2 decision; probing it for every 32-weight
7581    // group adds a branch to the hottest Q2TP row loop while retaining the
7582    // established table-load AVX2 arithmetic (and without coupling this path
7583    // to the separate FMA-gated A8W8 switch).
7584    let avx2 = q2tp_avx2_enabled();
7585    for gi in 0..gpr {
7586        let ch = &chunks[base + gi * Q2TP_CHUNK..base + (gi + 1) * Q2TP_CHUNK];
7587        let xg = &xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
7588        #[cfg(target_arch = "aarch64")]
7589        // NEON: the byte's four 2-bit fields land in four lane vectors
7590        // (shift+mask), vld4 de-interleaves xq to match (xj[k] =
7591        // xq[4k+j]), widening MACs accumulate exactly in i32. A scalar
7592        // decode here cost as much as the dot it fed — the profile put
7593        // it at the top of the whole W2 decode.
7594        let dot = unsafe {
7595            use core::arch::aarch64::*;
7596            let b = vld1_u8(ch.as_ptr());
7597            let three = vdup_n_u8(3);
7598            let c0 = vreinterpret_s8_u8(vand_u8(b, three));
7599            let c1 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 2), three));
7600            let c2 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 4), three));
7601            let c3 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 6), three));
7602            let x4 = vld4_s8(xg.as_ptr());
7603            let mut acc4 = vdupq_n_s32(0);
7604            acc4 = vpadalq_s16(acc4, vmull_s8(c0, x4.0));
7605            acc4 = vpadalq_s16(acc4, vmull_s8(c1, x4.1));
7606            acc4 = vpadalq_s16(acc4, vmull_s8(c2, x4.2));
7607            acc4 = vpadalq_s16(acc4, vmull_s8(c3, x4.3));
7608            vaddvq_s32(acc4)
7609        };
7610        #[cfg(target_arch = "x86_64")]
7611        let dot: i32 = if avx2 {
7612            // SAFETY: the runtime feature check gates the target-feature body;
7613            // the group slices above are exactly 8 and 32 bytes long.
7614            unsafe { q2tp_code_dot_avx2(ch, xg) }
7615        } else {
7616            ch.iter()
7617                .enumerate()
7618                .map(|(k, &b)| {
7619                    ((b & 3) as i32) * xg[k * 4] as i32
7620                        + (((b >> 2) & 3) as i32) * xg[k * 4 + 1] as i32
7621                        + (((b >> 4) & 3) as i32) * xg[k * 4 + 2] as i32
7622                        + (((b >> 6) & 3) as i32) * xg[k * 4 + 3] as i32
7623                })
7624                .sum()
7625        };
7626        #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
7627        let dot: i32 = {
7628            for (k, &b) in ch.iter().enumerate() {
7629                codes[k * 4] = (b & 3) as i8;
7630                codes[k * 4 + 1] = ((b >> 2) & 3) as i8;
7631                codes[k * 4 + 2] = ((b >> 4) & 3) as i8;
7632                codes[k * 4 + 3] = ((b >> 6) & 3) as i8;
7633            }
7634            codes
7635                .iter()
7636                .zip(xg)
7637                .map(|(&c, &x)| c as i32 * x as i32)
7638                .sum()
7639        };
7640        acc += scales[gi] * (dot as f32 - 1.5 * gsum[gi] as f32);
7641    }
7642    acc
7643}
7644
7645/// Exact f32 dot of one q2tp row: 2-bit fields LSB-first, (c − 1.5)·s.
7646/// Scalar on purpose — the 2-bit class targets the GPU graph; the CPU
7647/// path exists for parity gates and small-machine fallback.
7648fn q2tp_row_exact(chunks: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
7649    q2tp_row_exact_center(chunks, r, gpr, x, scales, 1.5)
7650}
7651
7652/// Fused Prism affine row: the derived correction is applied inside the
7653/// decoded code, avoiding a second accumulated dot and avoiding cancellation
7654/// between `B=(c-1.5)s` and `+.5s` for long 5120/17408 rows.
7655#[inline]
7656fn q2tp_affine_row_exact(chunks: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
7657    q2tp_row_exact_center(chunks, r, gpr, x, scales, 1.0)
7658}
7659
7660#[inline]
7661fn q2tp_row_exact_center(
7662    chunks: &[u8],
7663    r: usize,
7664    gpr: usize,
7665    x: &[f32],
7666    scales: &[f32],
7667    center: f32,
7668) -> f32 {
7669    let mut acc = 0f32;
7670    for gi in 0..gpr {
7671        let ch = &chunks[(r * gpr + gi) * Q2TP_CHUNK..(r * gpr + gi + 1) * Q2TP_CHUNK];
7672        let s = scales[gi];
7673        let xb = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
7674        let mut g = 0f32;
7675        for (k, &b) in ch.iter().enumerate() {
7676            g += ((b & 3) as f32 - center) * xb[k * 4]
7677                + (((b >> 2) & 3) as f32 - center) * xb[k * 4 + 1]
7678                + (((b >> 4) & 3) as f32 - center) * xb[k * 4 + 2]
7679                + (((b >> 6) & 3) as f32 - center) * xb[k * 4 + 3];
7680        }
7681        acc += s * g;
7682    }
7683    acc
7684}
7685
7686fn q2tp_matvec(
7687    bytes: &[u8],
7688    x: &[f32],
7689    rows: usize,
7690    cols: usize,
7691    out: &mut [f32],
7692    pool: Option<&Pool>,
7693) {
7694    q2tp_matvec_mode(bytes, x, rows, cols, out, pool, false);
7695}
7696
7697fn q2tp_affine_matvec(
7698    bytes: &[u8],
7699    x: &[f32],
7700    rows: usize,
7701    cols: usize,
7702    out: &mut [f32],
7703    pool: Option<&Pool>,
7704) {
7705    q2tp_matvec_mode(bytes, x, rows, cols, out, pool, true);
7706}
7707
7708fn q2tp_matvec_mode(
7709    bytes: &[u8],
7710    x: &[f32],
7711    rows: usize,
7712    cols: usize,
7713    out: &mut [f32],
7714    pool: Option<&Pool>,
7715    affine: bool,
7716) {
7717    debug_assert_eq!(out.len(), rows);
7718    let gpr = cols / GROUP_SIZE;
7719    let v = Q4tpView::new_q2(bytes, rows, cols);
7720    let out_addr = SendMut(out.as_mut_ptr());
7721    // a8w8 fast path (CMF_SDOT=0 keeps the exact scalar walk): integer
7722    // code dots + group sums, exact outlier correction — the same
7723    // contract as every sibling kernel; measured 2-bit rows were the
7724    // only scalar holdout in the family.
7725    if !affine && a8w8_enabled() {
7726        let act = split_act(x);
7727        let gsum = q1_group_sums(&act.xq, gpr);
7728        let (act, gsum) = (&act, &gsum);
7729        let run = move |start: usize, end: usize| {
7730            with_krow(gpr, |sc| {
7731                for r in start..end {
7732                    v.scales_into(r, gpr, sc);
7733                    let mut acc = dot_q2tp_row_i8(v.nib, r, gpr, &act.xq, gsum, sc) * act.sx;
7734                    for &(j, xv) in &act.outliers {
7735                        let (w, s) = q2tp_outlier(v.nib, r, gpr, j, sc);
7736                        acc += w * s * xv;
7737                    }
7738                    // SAFETY: disjoint row ranges per worker.
7739                    unsafe { *out_addr.at(r) = acc };
7740                }
7741            })
7742        };
7743        dispatch_rows(pool, rows, &run);
7744        return;
7745    }
7746    let run = |start: usize, end: usize| {
7747        with_krow(gpr, |sc| {
7748            for r in start..end {
7749                v.scales_into(r, gpr, sc);
7750                // SAFETY: disjoint row ranges per worker.
7751                unsafe {
7752                    *out_addr.at(r) = if affine {
7753                        q2tp_affine_row_exact(v.nib, r, gpr, x, sc)
7754                    } else {
7755                        q2tp_row_exact(v.nib, r, gpr, x, sc)
7756                    }
7757                };
7758            }
7759        })
7760    };
7761    dispatch_rows(pool, rows, &run);
7762}
7763
7764/// Fused two-input q2tp matvec — the SwiGLU gate/up pair.
7765#[allow(clippy::too_many_arguments)]
7766fn q2tp_matvec2(
7767    bytes: &[u8],
7768    x1: &[f32],
7769    x2: &[f32],
7770    rows: usize,
7771    cols: usize,
7772    o1: &mut [f32],
7773    o2: &mut [f32],
7774    pool: Option<&Pool>,
7775) {
7776    q2tp_matvec2_mode(bytes, x1, x2, rows, cols, o1, o2, pool, false);
7777}
7778
7779#[allow(clippy::too_many_arguments)]
7780fn q2tp_affine_matvec2(
7781    bytes: &[u8],
7782    x1: &[f32],
7783    x2: &[f32],
7784    rows: usize,
7785    cols: usize,
7786    o1: &mut [f32],
7787    o2: &mut [f32],
7788    pool: Option<&Pool>,
7789) {
7790    q2tp_matvec2_mode(bytes, x1, x2, rows, cols, o1, o2, pool, true);
7791}
7792
7793#[allow(clippy::too_many_arguments)]
7794fn q2tp_matvec2_mode(
7795    bytes: &[u8],
7796    x1: &[f32],
7797    x2: &[f32],
7798    rows: usize,
7799    cols: usize,
7800    o1: &mut [f32],
7801    o2: &mut [f32],
7802    pool: Option<&Pool>,
7803    affine: bool,
7804) {
7805    let gpr = cols / GROUP_SIZE;
7806    let v = Q4tpView::new_q2(bytes, rows, cols);
7807    let (p1, p2) = (SendMut(o1.as_mut_ptr()), SendMut(o2.as_mut_ptr()));
7808    let run = |start: usize, end: usize| {
7809        let mut sc = vec![0f32; gpr];
7810        for r in start..end {
7811            v.scales_into(r, gpr, &mut sc);
7812            // SAFETY: disjoint row ranges per worker.
7813            unsafe {
7814                *p1.at(r) = if affine {
7815                    q2tp_affine_row_exact(v.nib, r, gpr, x1, &sc)
7816                } else {
7817                    q2tp_row_exact(v.nib, r, gpr, x1, &sc)
7818                };
7819                *p2.at(r) = if affine {
7820                    q2tp_affine_row_exact(v.nib, r, gpr, x2, &sc)
7821                } else {
7822                    q2tp_row_exact(v.nib, r, gpr, x2, &sc)
7823                };
7824            }
7825        }
7826    };
7827    dispatch_rows(pool, rows, &run);
7828}
7829
7830/// Batched q2tp matmat: scalar row kernel over every batch column. CPU
7831/// prefill only — decode rides the graph, so plain and correct beats
7832/// clever here.
7833/// Test doors into the host 2-bit kernels: the stand's heap corruption
7834/// pointed at down-shaped tensors, and the private fns need a way to be
7835/// held to a reference without a model file around them.
7836pub fn q2tp_matvec_for_test(bytes: &[u8], x: &[f32], rows: usize, cols: usize, out: &mut [f32]) {
7837    // The facade IS the reference: encoder oracles hold requant output
7838    // to the exact scalar walk. The production dispatch may take the i8
7839    // fast path, whose error scale is the ACTIVATIONS' — a different
7840    // claim than the encoder correctness these tests pin.
7841    let gpr = cols / GROUP_SIZE;
7842    let v = Q4tpView::new_q2(bytes, rows, cols);
7843    with_krow(gpr, |sc| {
7844        for r in 0..rows {
7845            v.scales_into(r, gpr, sc);
7846            out[r] = q2tp_row_exact(v.nib, r, gpr, x, sc);
7847        }
7848    });
7849}
7850
7851/// Test door for the descriptor-specific fused affine decode.  Production
7852/// callers select this through a validated Prism header, never by dtype alone.
7853pub fn q2tp_affine_matvec_for_test(
7854    bytes: &[u8],
7855    x: &[f32],
7856    rows: usize,
7857    cols: usize,
7858    out: &mut [f32],
7859) {
7860    q2tp_affine_matvec(bytes, x, rows, cols, out, None);
7861}
7862
7863pub fn q2tp_matmat_for_test(
7864    bytes: &[u8],
7865    xs_all: &[f32],
7866    b: usize,
7867    rows: usize,
7868    cols: usize,
7869    out: &mut [f32],
7870) {
7871    q2tp_matmat(bytes, xs_all, b, rows, cols, out, None);
7872}
7873
7874fn q2tp_matmat(
7875    bytes: &[u8],
7876    xs_all: &[f32],
7877    b: usize,
7878    rows: usize,
7879    cols: usize,
7880    out: &mut [f32],
7881    pool: Option<&Pool>,
7882) {
7883    q2tp_matmat_mode(bytes, xs_all, b, rows, cols, out, pool, false);
7884}
7885
7886fn q2tp_affine_matmat(
7887    bytes: &[u8],
7888    xs_all: &[f32],
7889    b: usize,
7890    rows: usize,
7891    cols: usize,
7892    out: &mut [f32],
7893    pool: Option<&Pool>,
7894) {
7895    q2tp_matmat_mode(bytes, xs_all, b, rows, cols, out, pool, true);
7896}
7897
7898fn q2tp_matmat_mode(
7899    bytes: &[u8],
7900    xs_all: &[f32],
7901    b: usize,
7902    rows: usize,
7903    cols: usize,
7904    out: &mut [f32],
7905    pool: Option<&Pool>,
7906    affine: bool,
7907) {
7908    debug_assert_eq!(out.len(), b * rows);
7909    let gpr = cols / GROUP_SIZE;
7910    let v = Q4tpView::new_q2(bytes, rows, cols);
7911    let out_addr = SendMut(out.as_mut_ptr());
7912    let run = |start: usize, end: usize| {
7913        let mut sc = vec![0f32; gpr];
7914        for r in start..end {
7915            v.scales_into(r, gpr, &mut sc);
7916            for bi in 0..b {
7917                let x = &xs_all[bi * cols..(bi + 1) * cols];
7918                // SAFETY: disjoint row ranges per worker.
7919                unsafe {
7920                    *out_addr.at(bi * rows + r) = if affine {
7921                        q2tp_affine_row_exact(v.nib, r, gpr, x, &sc)
7922                    } else {
7923                        q2tp_row_exact(v.nib, r, gpr, x, &sc)
7924                    }
7925                };
7926            }
7927        }
7928    };
7929    dispatch_rows(pool, rows, &run);
7930}
7931
7932/// The pre-vectorised shape, kept for A/B (`CMF_Q4TP_V1=1`): the
7933/// horizontal add lands once per group per column instead of once per
7934/// row. Same weights, same activations — only the reduction differs.
7935/// It is also the row-exact batch kernel: per group and column it forms
7936/// `int dot as f32 * scale` and adds it to a scalar running sum, which is
7937/// `dot_q4tp_row_sdot` step for step (Rust never contracts to an fma),
7938/// so each column is bit-identical to that column's matvec.
7939#[cfg(target_arch = "aarch64")]
7940#[target_feature(enable = "neon,dotprod")]
7941unsafe fn dot_q4tp_row_1x4_sdot_v1(
7942    nib: &[u8],
7943    r: usize,
7944    gpr: usize,
7945    xs: [&[i8]; 4],
7946    scales: &[f32],
7947) -> [f32; 4] {
7948    unsafe {
7949        use core::arch::aarch64::*;
7950        use core::arch::asm;
7951        let lomask = vdupq_n_u8(0x0F);
7952        let eight = vdupq_n_s8(8);
7953        let (mut f0, mut f1, mut f2, mut f3) = (0f32, 0f32, 0f32, 0f32);
7954        for gi in 0..gpr {
7955            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
7956            let s = *scales.get_unchecked(gi);
7957            let bb = vld1q_u8(t);
7958            let lo = vandq_u8(bb, lomask);
7959            let hi = vshrq_n_u8::<4>(bb);
7960            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
7961            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
7962            let mut d = [0f32; 4];
7963            for (k, dk) in d.iter_mut().enumerate() {
7964                let x0 = vld1q_s8(xs[k].as_ptr().add(gi * GROUP_SIZE));
7965                let x1 = vld1q_s8(xs[k].as_ptr().add(gi * GROUP_SIZE + 16));
7966                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
7967                asm!(
7968                    "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
7969                    "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
7970                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
7971                    e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
7972                    options(pure, nomem, nostack),
7973                );
7974                *dk = vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
7975            }
7976            f0 += d[0];
7977            f1 += d[1];
7978            f2 += d[2];
7979            f3 += d[3];
7980        }
7981        [f0, f1, f2, f3]
7982    }
7983}
7984
7985/// Which q4tp batch kernel to run: 1 = the previous one, 2 = the tuned
7986/// one, 0 = decide from the CPU. An atomic rather than a `OnceLock` so a
7987/// benchmark can alternate the two inside one process, where the machine's
7988/// mood — a shared box drifts ±25% between runs — is the same for both.
7989/// What the two mean is per-architecture: on x86 the blocked AVX-512 path
7990/// against the per-column one, on ARM the two reduction shapes.
7991#[allow(dead_code)]
7992static Q4TP_ALT: std::sync::atomic::AtomicU8 = std::sync::atomic::AtomicU8::new(0);
7993
7994/// Tests that store `Q4TP_ALT` hold this, so one test's kernel pick does
7995/// not leak into another's bit-exact comparison running in parallel.
7996#[cfg(test)]
7997static Q4TP_ALT_TEST_LOCK: std::sync::Mutex<()> = std::sync::Mutex::new(());
7998
7999/// Blocking pays on x86 only with 512-bit VNNI. With AVX2 alone, four
8000/// columns sharing an unpack still measured slower than the per-column
8001/// path (23.2 ms against 19.4 on a 48-thread EPYC), because that path
8002/// already dequantizes the row once — so the blocked kernel bought a
8003/// second unpack-free pass at the price of half the vector width.
8004#[cfg(target_arch = "x86_64")]
8005fn q4tp_blocked_x86() -> bool {
8006    match Q4TP_ALT.load(std::sync::atomic::Ordering::Relaxed) {
8007        1 => false,
8008        // A forced ON still asks the CPU. The switch exists so a bench can
8009        // pick a kernel, not so it can promise instructions the machine
8010        // does not have — CI caught that as a SIGILL on a runner without
8011        // AVX-512, where the parity test had turned the path on by hand.
8012        2 => avx512vnni_enabled(),
8013        // Deliberately not cached back into the switch: both gates below
8014        // hold their own `OnceLock`, and latching their answer here would
8015        // make a test's override outlive the test that set it.
8016        _ => blocked_enabled() && avx512vnni_enabled(),
8017    }
8018}
8019
8020/// `CMF_Q4TP_V1=1` picks the old reduction shape (A/B only).
8021#[cfg(target_arch = "aarch64")]
8022#[allow(dead_code)]
8023fn q4tp_v1() -> bool {
8024    match Q4TP_ALT.load(std::sync::atomic::Ordering::Relaxed) {
8025        1 => true,
8026        2 => false,
8027        _ => {
8028            static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8029            *ON.get_or_init(|| std::env::var("CMF_Q4TP_V1").is_ok_and(|v| v != "0"))
8030        }
8031    }
8032}
8033
8034/// Two weight rows against eight columns. The activation load is the
8035/// same for both rows, so it is paid once for twice the arithmetic, and
8036/// sixteen accumulator chains run where eight did — which is what a kernel
8037/// retiring 0.29 instructions a cycle is short of. Register pressure is
8038/// the limit: sixteen `zmm` accumulators, two weight tiles, one
8039/// activation, of thirty-two.
8040///
8041/// Four rows by four columns spends the same sixteen accumulators the
8042/// other way and measured worse — 1488 GFLOP/s against 1644 — so the
8043/// unpack, which four rows pay twice as often, costs more than the extra
8044/// sharing of one activation load buys.
8045#[cfg(target_arch = "x86_64")]
8046#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
8047unsafe fn dot_q4tp_2x8_avx512(
8048    nib: &[u8],
8049    r0: usize,
8050    gpr: usize,
8051    xs: [&[i8]; 8],
8052    sc0: &[f32],
8053    sc1: &[f32],
8054) -> [[f32; 8]; 2] {
8055    // SAFETY: as dot_q4tp_row_1x8_avx512, two adjacent rows at once; the
8056    // caller guarantees r0 + 1 < rows and the ISA.
8057    unsafe {
8058        use core::arch::x86_64::*;
8059        let lomask = _mm256_set1_epi8(0x0F);
8060        let eight = _mm256_set1_epi8(8);
8061        let zero = _mm512_setzero_si512();
8062        let mut v0 = [_mm512_setzero_ps(); 8];
8063        let mut v1 = [_mm512_setzero_ps(); 8];
8064        let pairs = gpr / 2;
8065        let unpack = |r: usize, gi: usize| -> (__m512i, __mmask64) {
8066            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
8067            let bb = _mm256_loadu_si256(t as *const __m256i);
8068            let lo = _mm256_and_si256(bb, lomask);
8069            let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
8070            let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
8071            let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
8072            let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
8073            let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
8074            (_mm512_abs_epi8(w), _mm512_movepi8_mask(w))
8075        };
8076        for gp in 0..pairs {
8077            let gi = gp * 2;
8078            let (wa0, neg0) = unpack(r0, gi);
8079            let (wa1, neg1) = unpack(r0 + 1, gi);
8080            let off = gi * GROUP_SIZE;
8081            let sv = |sc: &[f32]| {
8082                _mm512_insertf32x8::<1>(
8083                    _mm512_castps256_ps512(_mm256_set1_ps(*sc.get_unchecked(gi))),
8084                    _mm256_set1_ps(*sc.get_unchecked(gi + 1)),
8085                )
8086            };
8087            let s0 = sv(sc0);
8088            let s1 = sv(sc1);
8089            for k in 0..8 {
8090                let xv = _mm512_loadu_si512(xs[k].as_ptr().add(off) as *const __m512i);
8091                let d0 = _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(
8092                    zero,
8093                    wa0,
8094                    _mm512_mask_sub_epi8(xv, neg0, zero, xv),
8095                ));
8096                let d1 = _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(
8097                    zero,
8098                    wa1,
8099                    _mm512_mask_sub_epi8(xv, neg1, zero, xv),
8100                ));
8101                v0[k] = _mm512_fmadd_ps(d0, s0, v0[k]);
8102                v1[k] = _mm512_fmadd_ps(d1, s1, v1[k]);
8103            }
8104        }
8105        let mut acc = [[0f32; 8]; 2];
8106        for k in 0..8 {
8107            acc[0][k] = _mm512_reduce_add_ps(v0[k]);
8108            acc[1][k] = _mm512_reduce_add_ps(v1[k]);
8109        }
8110        if gpr % 2 == 1 {
8111            let off = (gpr - 1) * GROUP_SIZE;
8112            for j in off..off + GROUP_SIZE {
8113                let (w0, sa) = q4tp_outlier(nib, r0, gpr, j, sc0);
8114                let (w1, sb) = q4tp_outlier(nib, r0 + 1, gpr, j, sc1);
8115                for k in 0..8 {
8116                    let x = *xs[k].get_unchecked(j) as f32;
8117                    acc[0][k] += w0 * sa * x;
8118                    acc[1][k] += w1 * sb * x;
8119                }
8120            }
8121        }
8122        acc
8123    }
8124}
8125
8126/// The same, eight columns at a time. One unpack then feeds twice as many
8127/// activation streams, so a wide batch reads the weight tile half as
8128/// often; the price is eight accumulators live at once. Measured 9.0 ->
8129/// 8.3 ms at 9216x2304, b=296 on a 48-thread EPYC 9B45.
8130#[cfg(target_arch = "x86_64")]
8131#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
8132unsafe fn dot_q4tp_row_1x8_avx512(
8133    nib: &[u8],
8134    r: usize,
8135    gpr: usize,
8136    xs: [&[i8]; 8],
8137    scales: &[f32],
8138) -> [f32; 8] {
8139    // SAFETY: as dot_q4tp_row_1x4_avx2; caller guarantees the ISA.
8140    unsafe {
8141        use core::arch::x86_64::*;
8142        let lomask = _mm256_set1_epi8(0x0F);
8143        let eight = _mm256_set1_epi8(8);
8144        let zero = _mm512_setzero_si512();
8145        let (mut v0, mut v1, mut v2, mut v3) = (
8146            _mm512_setzero_ps(),
8147            _mm512_setzero_ps(),
8148            _mm512_setzero_ps(),
8149            _mm512_setzero_ps(),
8150        );
8151        let (mut v4, mut v5, mut v6, mut v7) = (
8152            _mm512_setzero_ps(),
8153            _mm512_setzero_ps(),
8154            _mm512_setzero_ps(),
8155            _mm512_setzero_ps(),
8156        );
8157        let pairs = gpr / 2;
8158        for gp in 0..pairs {
8159            let gi = gp * 2;
8160            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
8161            let bb = _mm256_loadu_si256(t as *const __m256i);
8162            let lo = _mm256_and_si256(bb, lomask);
8163            let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
8164            // `unpack` works per 128-bit lane, so the halves come out as
8165            // [A.lo, B.lo] and [A.hi, B.hi]; the shuffle reorders the four
8166            // 128-bit lanes into the weights' natural order, which is what
8167            // the straight activation load expects.
8168            let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
8169            let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
8170            let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
8171            let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
8172            let wabs = _mm512_abs_epi8(w);
8173            let neg = _mm512_movepi8_mask(w);
8174            let off = gi * GROUP_SIZE;
8175            let sv = _mm512_insertf32x8::<1>(
8176                _mm512_castps256_ps512(_mm256_set1_ps(*scales.get_unchecked(gi))),
8177                _mm256_set1_ps(*scales.get_unchecked(gi + 1)),
8178            );
8179            let dot = |x: &[i8]| -> __m512 {
8180                let xv = _mm512_loadu_si512(x.as_ptr().add(off) as *const __m512i);
8181                let sx = _mm512_mask_sub_epi8(xv, neg, zero, xv);
8182                _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(zero, wabs, sx))
8183            };
8184            v0 = _mm512_fmadd_ps(dot(xs[0]), sv, v0);
8185            v1 = _mm512_fmadd_ps(dot(xs[1]), sv, v1);
8186            v2 = _mm512_fmadd_ps(dot(xs[2]), sv, v2);
8187            v3 = _mm512_fmadd_ps(dot(xs[3]), sv, v3);
8188            v4 = _mm512_fmadd_ps(dot(xs[4]), sv, v4);
8189            v5 = _mm512_fmadd_ps(dot(xs[5]), sv, v5);
8190            v6 = _mm512_fmadd_ps(dot(xs[6]), sv, v6);
8191            v7 = _mm512_fmadd_ps(dot(xs[7]), sv, v7);
8192        }
8193        let mut acc = [
8194            _mm512_reduce_add_ps(v0),
8195            _mm512_reduce_add_ps(v1),
8196            _mm512_reduce_add_ps(v2),
8197            _mm512_reduce_add_ps(v3),
8198            _mm512_reduce_add_ps(v4),
8199            _mm512_reduce_add_ps(v5),
8200            _mm512_reduce_add_ps(v6),
8201            _mm512_reduce_add_ps(v7),
8202        ];
8203        // An odd group count leaves one group over; the narrow kernel
8204        // finishes it rather than the tail being a special case here.
8205        if gpr % 2 == 1 {
8206            let off = (gpr - 1) * GROUP_SIZE;
8207            for j in off..off + GROUP_SIZE {
8208                let (w, s) = q4tp_outlier(nib, r, gpr, j, scales);
8209                let ws = w * s;
8210                for k in 0..8 {
8211                    acc[k] += ws * *xs[k].get_unchecked(j) as f32;
8212                }
8213            }
8214        }
8215        acc
8216    }
8217}
8218
8219/// The same four columns, 512 bits wide. Two groups (64 weights) ride one
8220/// unpack and one `vpdpbusd`, where AVX2 needs two unpacks and four
8221/// `maddubs`/`madd` pairs — about 2.3x fewer instructions for the same
8222/// arithmetic. The two groups carry different scales, so the fma takes a
8223/// vector whose halves hold each group's scale rather than a broadcast.
8224///
8225/// There is no 512-bit `vpsignb`, so the activation's sign is applied by
8226/// negating under a mask taken from the weight's sign bits. That mask is
8227/// per-tile, so it is hoisted out of the column loop and the per-column
8228/// cost stays exactly one instruction, as with `sign_epi8`. Weights of
8229/// zero are not zeroed by the mask trick and do not need to be: their
8230/// magnitude is zero, so the product is.
8231#[cfg(target_arch = "x86_64")]
8232#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
8233unsafe fn dot_q4tp_row_1x4_avx512(
8234    nib: &[u8],
8235    r: usize,
8236    gpr: usize,
8237    xs: [&[i8]; 4],
8238    scales: &[f32],
8239) -> [f32; 4] {
8240    // SAFETY: as dot_q4tp_row_1x4_avx2; caller guarantees the ISA.
8241    unsafe {
8242        use core::arch::x86_64::*;
8243        let lomask = _mm256_set1_epi8(0x0F);
8244        let eight = _mm256_set1_epi8(8);
8245        let zero = _mm512_setzero_si512();
8246        let (mut v0, mut v1, mut v2, mut v3) = (
8247            _mm512_setzero_ps(),
8248            _mm512_setzero_ps(),
8249            _mm512_setzero_ps(),
8250            _mm512_setzero_ps(),
8251        );
8252        let pairs = gpr / 2;
8253        for gp in 0..pairs {
8254            let gi = gp * 2;
8255            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
8256            let bb = _mm256_loadu_si256(t as *const __m256i);
8257            let lo = _mm256_and_si256(bb, lomask);
8258            let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
8259            // `unpack` works per 128-bit lane, so the halves come out as
8260            // [A.lo, B.lo] and [A.hi, B.hi]; the shuffle reorders the four
8261            // 128-bit lanes into the weights' natural order, which is what
8262            // the straight activation load expects.
8263            let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
8264            let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
8265            let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
8266            let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
8267            let wabs = _mm512_abs_epi8(w);
8268            let neg = _mm512_movepi8_mask(w);
8269            let off = gi * GROUP_SIZE;
8270            let sv = _mm512_insertf32x8::<1>(
8271                _mm512_castps256_ps512(_mm256_set1_ps(*scales.get_unchecked(gi))),
8272                _mm256_set1_ps(*scales.get_unchecked(gi + 1)),
8273            );
8274            let dot = |x: &[i8]| -> __m512 {
8275                let xv = _mm512_loadu_si512(x.as_ptr().add(off) as *const __m512i);
8276                let sx = _mm512_mask_sub_epi8(xv, neg, zero, xv);
8277                _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(zero, wabs, sx))
8278            };
8279            v0 = _mm512_fmadd_ps(dot(xs[0]), sv, v0);
8280            v1 = _mm512_fmadd_ps(dot(xs[1]), sv, v1);
8281            v2 = _mm512_fmadd_ps(dot(xs[2]), sv, v2);
8282            v3 = _mm512_fmadd_ps(dot(xs[3]), sv, v3);
8283        }
8284        let mut acc = [
8285            _mm512_reduce_add_ps(v0),
8286            _mm512_reduce_add_ps(v1),
8287            _mm512_reduce_add_ps(v2),
8288            _mm512_reduce_add_ps(v3),
8289        ];
8290        // An odd group count leaves one group over; the narrow kernel
8291        // finishes it rather than the tail being a special case here.
8292        if gpr % 2 == 1 {
8293            let off = (gpr - 1) * GROUP_SIZE;
8294            for j in off..off + GROUP_SIZE {
8295                let (w, s) = q4tp_outlier(nib, r, gpr, j, scales);
8296                let ws = w * s;
8297                for k in 0..4 {
8298                    acc[k] += ws * *xs[k].get_unchecked(j) as f32;
8299                }
8300            }
8301        }
8302        acc
8303    }
8304}
8305
8306/// Four batch columns against one q4tp row: the tile is unpacked ONCE and
8307/// spent on four activation streams, which is where a prefill batch stops
8308/// being weight-bandwidth-bound. Twin of `dot_q4t_row_1x4_sdot`.
8309#[cfg(target_arch = "aarch64")]
8310#[target_feature(enable = "neon,dotprod")]
8311unsafe fn dot_q4tp_row_1x4_sdot(
8312    nib: &[u8],
8313    r: usize,
8314    gpr: usize,
8315    xs: [&[i8]; 4],
8316    scales: &[f32],
8317) -> [f32; 4] {
8318    // SAFETY: see dot_q4tp_row_sdot; every xs[k] is gpr·GROUP_SIZE long.
8319    unsafe {
8320        use core::arch::aarch64::*;
8321        use core::arch::asm;
8322        let lomask = vdupq_n_u8(0x0F);
8323        let eight = vdupq_n_s8(8);
8324        // Named accumulators, NOT an array indexed by a loop variable: the
8325        // latter does not stay in registers (the same defect cost 2x in the
8326        // AVX2 q4t kernel and again in WGSL).
8327        //
8328        // They are VECTORS, and the horizontal add happens once at the end
8329        // instead of once per group per column. `vaddvq` is a cross-lane
8330        // reduction — with 72 groups and four columns the old shape paid
8331        // 288 of them per row, each one a dependency stall the pipeline
8332        // cannot hide, to save four float adds. The group's scale now
8333        // rides an fma into the lane accumulators, so the arithmetic per
8334        // group is one convert and one fma. Summation order changes (the
8335        // lanes carry independent partial sums), which is the same
8336        // round-off class the SDOT path already lives in — the strict
8337        // kernel (`CMF_SDOT=0`, what `cortiq ppl` runs) is unchanged and
8338        // stays the reference.
8339        let (mut v0, mut v1, mut v2, mut v3) = (
8340            vdupq_n_f32(0.0),
8341            vdupq_n_f32(0.0),
8342            vdupq_n_f32(0.0),
8343            vdupq_n_f32(0.0),
8344        );
8345        for gi in 0..gpr {
8346            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
8347            let s = *scales.get_unchecked(gi);
8348            let bb = vld1q_u8(t);
8349            let lo = vandq_u8(bb, lomask);
8350            let hi = vshrq_n_u8::<4>(bb);
8351            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
8352            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
8353            let off = gi * GROUP_SIZE;
8354            let dot4 = |x: &[i8]| -> int32x4_t {
8355                let x0 = vld1q_s8(x.as_ptr().add(off));
8356                let x1 = vld1q_s8(x.as_ptr().add(off + 16));
8357                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
8358                asm!(
8359                    "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
8360                    "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
8361                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
8362                    e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
8363                    options(pure, nomem, nostack),
8364                );
8365                vaddq_s32(a0, a1)
8366            };
8367            v0 = vfmaq_n_f32(v0, vcvtq_f32_s32(dot4(xs[0])), s);
8368            v1 = vfmaq_n_f32(v1, vcvtq_f32_s32(dot4(xs[1])), s);
8369            v2 = vfmaq_n_f32(v2, vcvtq_f32_s32(dot4(xs[2])), s);
8370            v3 = vfmaq_n_f32(v3, vcvtq_f32_s32(dot4(xs[3])), s);
8371        }
8372        [
8373            vaddvq_f32(v0),
8374            vaddvq_f32(v1),
8375            vaddvq_f32(v2),
8376            vaddvq_f32(v3),
8377        ]
8378    }
8379}
8380
8381/// Fused q4tp matmat — the same three arms `q4t_matmat` has. Shipping only
8382/// the scalar one made Nanbeige-3B decode at 1.2 tok/s against q4t's 5.9:
8383/// the format was fine, the missing arms were the whole regression.
8384///
8385/// Under `row_exact()` every cell is summed in `q4tp_matvec`'s order on
8386/// every architecture; the mode is read once, so one call never mixes
8387/// the two contracts when a concurrent scope opens or closes mid-call.
8388fn q4tp_matmat(
8389    bytes: &[u8],
8390    xs_all: &[f32],
8391    b: usize,
8392    rows: usize,
8393    cols: usize,
8394    out: &mut [f32],
8395    pool: Option<&Pool>,
8396) {
8397    q4tp_matmat_with(bytes, xs_all, b, rows, cols, out, pool, row_exact())
8398}
8399
8400/// `q4tp_matmat` with the row-exact mode passed in rather than read from
8401/// the shared counter: `exact` = each cell equals its token's matvec bit
8402/// for bit, otherwise the fast blocked / AMX arms are free to reorder.
8403#[allow(clippy::too_many_arguments)]
8404fn q4tp_matmat_with(
8405    bytes: &[u8],
8406    xs_all: &[f32],
8407    b: usize,
8408    rows: usize,
8409    cols: usize,
8410    out: &mut [f32],
8411    pool: Option<&Pool>,
8412    exact: bool,
8413) {
8414    debug_assert_eq!(out.len(), b * rows);
8415    let gpr = cols / GROUP_SIZE;
8416    let v = Q4tpView::new(bytes, rows, cols);
8417
8418    // Wide batches ride the AMX through a dequant-tile sgemm, as in q4t.
8419    // An f32 GEMM over dequantized weights is not the matvec's int8 sum,
8420    // so a row-exact batch never takes it.
8421    #[cfg(target_os = "macos")]
8422    if !exact && b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
8423        dequant_matmat_accel(
8424            &|r, dst| {
8425                let mut sc = [0f32; 32];
8426                let mut scv;
8427                let s: &[f32] = if gpr <= 32 {
8428                    v.scales_into(r, gpr, &mut sc);
8429                    &sc[..gpr]
8430                } else {
8431                    scv = vec![0f32; gpr];
8432                    v.scales_into(r, gpr, &mut scv);
8433                    &scv
8434                };
8435                for gi in 0..gpr {
8436                    let tile = &v.nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
8437                    for (k, &bb) in tile.iter().enumerate() {
8438                        dst[gi * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s[gi];
8439                        dst[gi * GROUP_SIZE + k * 2 + 1] =
8440                            (((bb >> 4) & 0x0F) as f32 - 8.0) * s[gi];
8441                    }
8442                }
8443            },
8444            xs_all,
8445            b,
8446            rows,
8447            cols,
8448            out,
8449            pool,
8450        );
8451        return;
8452    }
8453
8454    let out_addr = SendMut(out.as_mut_ptr());
8455    if a8w8_enabled() {
8456        let acts = split_acts_pooled(xs_all, b, cols, pool);
8457        let acts = &acts;
8458        // ARM stays blocked under `exact` too: the 1x4 kernel then runs in
8459        // its v1 shape, whose per-group `int dot as f32 * scale` and scalar
8460        // running sum are exactly `dot_q4tp_row_sdot`'s, so the tile is
8461        // still unpacked once for four columns and each column equals its
8462        // matvec. The tuned shape (fma into lane partials, one horizontal
8463        // add per row) is 1-16 ulp off the matvec and is kept for `!exact`.
8464        #[cfg(target_arch = "aarch64")]
8465        let blocked_ok = sdot_enabled() && blocked_enabled();
8466        // x86 gets the same blocking: one tile unpack spent on four
8467        // columns. Without it every column re-decoded the row, which is
8468        // why a 48-core EPYC measured a sixth of an M4's per-core rate.
8469        // The gate is `avx2_enabled`, as in q4t — `sdot_enabled` answers
8470        // for ARM's dotprod and is hard-wired false everywhere else, so
8471        // asking it here left the whole blocked path unreachable on x86.
8472        #[cfg(target_arch = "x86_64")]
8473        let blocked_ok = q4tp_blocked_x86() && !exact;
8474        #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
8475        let blocked_ok = {
8476            let _ = exact;
8477            false
8478        };
8479        // Columns are swept in panels that fit L2. Without this a
8480        // row-pair walks every activation in the batch — 4.8 MB at
8481        // 512x512 — and does it again for the next pair, so the whole
8482        // batch streams out of the shared cache once per row. Measured
8483        // 800 GB/s of it, flat across batch sizes, which is the signature
8484        // of a loop bound by traffic rather than by arithmetic. A panel of
8485        // 256 columns is 590 KB beside 221 KB of this worker's weights:
8486        // both stay resident and the batch crosses L3 once instead of
8487        // once per row.
8488        let panel_cols = q4tp_panel_cols(cols);
8489        let run = |start: usize, end: usize| {
8490            for abase in (0..acts.len()).step_by(panel_cols) {
8491                let alen = (acts.len() - abase).min(panel_cols);
8492                let mut sc = vec![0f32; gpr];
8493                #[cfg(target_arch = "x86_64")]
8494                let mut r_lo = start;
8495                #[cfg(target_arch = "x86_64")]
8496                if blocked_ok && alen >= 8 {
8497                    let mut sc1 = vec![0f32; gpr];
8498                    while r_lo + 2 <= end {
8499                        v.scales_into(r_lo, gpr, &mut sc);
8500                        v.scales_into(r_lo + 1, gpr, &mut sc1);
8501                        let mut bi = 0usize;
8502                        while bi + 8 <= alen {
8503                            let xs = [
8504                                acts[abase + bi].xq.as_slice(),
8505                                acts[abase + bi + 1].xq.as_slice(),
8506                                acts[abase + bi + 2].xq.as_slice(),
8507                                acts[abase + bi + 3].xq.as_slice(),
8508                                acts[abase + bi + 4].xq.as_slice(),
8509                                acts[abase + bi + 5].xq.as_slice(),
8510                                acts[abase + bi + 6].xq.as_slice(),
8511                                acts[abase + bi + 7].xq.as_slice(),
8512                            ];
8513                            let d = unsafe { dot_q4tp_2x8_avx512(v.nib, r_lo, gpr, xs, &sc, &sc1) };
8514                            for (row, dr, scr) in [(r_lo, &d[0], &sc), (r_lo + 1, &d[1], &sc1)] {
8515                                for k in 0..8 {
8516                                    let act = &acts[abase + bi + k];
8517                                    let mut acc = dr[k] * act.sx;
8518                                    for &(j, xv) in &act.outliers {
8519                                        let (w, s) = q4tp_outlier(v.nib, row, gpr, j, scr);
8520                                        acc += w * s * xv;
8521                                    }
8522                                    // SAFETY: disjoint (bi, r) cells per worker.
8523                                    unsafe { *out_addr.at((abase + bi + k) * rows + row) = acc };
8524                                }
8525                            }
8526                            bi += 8;
8527                        }
8528                        // Columns past the last group of eight, both rows —
8529                        // the same single-row kernel the tail below uses.
8530                        for row in [r_lo, r_lo + 1] {
8531                            let scr: &[f32] = if row == r_lo { &sc } else { &sc1 };
8532                            for b2 in bi..alen {
8533                                let act = &acts[abase + b2];
8534                                let xs4 = [
8535                                    act.xq.as_slice(),
8536                                    act.xq.as_slice(),
8537                                    act.xq.as_slice(),
8538                                    act.xq.as_slice(),
8539                                ];
8540                                let d =
8541                                    unsafe { dot_q4tp_row_1x4_avx512(v.nib, row, gpr, xs4, scr) };
8542                                let mut acc = d[0] * act.sx;
8543                                for &(j, xv) in &act.outliers {
8544                                    let (w, s) = q4tp_outlier(v.nib, row, gpr, j, scr);
8545                                    acc += w * s * xv;
8546                                }
8547                                // SAFETY: disjoint (bi, r) cells per worker.
8548                                unsafe { *out_addr.at((abase + b2) * rows + row) = acc };
8549                            }
8550                        }
8551                        r_lo += 2;
8552                    }
8553                }
8554                // AVX2 without VNNI, a batch that need not be row-exact:
8555                // two rows against four columns with FMA
8556                // (`q4tp_matmat_2rows_fma`); an odd last row falls to the
8557                // vertical cell below.
8558                #[cfg(target_arch = "x86_64")]
8559                if !blocked_ok
8560                    && !exact
8561                    && alen >= 4
8562                    && q4tp_fma_x86()
8563                    && acts[abase..abase + alen].iter().all(|a| a.xg8.len() == gpr)
8564                {
8565                    let mut sc1 = vec![0f32; gpr];
8566                    while r_lo + 2 <= end {
8567                        v.scales_into(r_lo, gpr, &mut sc);
8568                        v.scales_into(r_lo + 1, gpr, &mut sc1);
8569                        // SAFETY: AVX2+FMA checked; cells disjoint per
8570                        // (worker, bi, r).
8571                        unsafe {
8572                            q4tp_matmat_2rows_fma(
8573                                &v,
8574                                r_lo,
8575                                gpr,
8576                                &acts[abase..abase + alen],
8577                                &sc,
8578                                &sc1,
8579                                abase,
8580                                rows,
8581                                out_addr,
8582                            )
8583                        };
8584                        r_lo += 2;
8585                    }
8586                }
8587                #[cfg(target_arch = "x86_64")]
8588                let row_start = r_lo;
8589                #[cfg(not(target_arch = "x86_64"))]
8590                let row_start = start;
8591                for r in row_start..end {
8592                    v.scales_into(r, gpr, &mut sc);
8593                    let mut bi = 0usize;
8594                    #[cfg(target_arch = "x86_64")]
8595                    if blocked_ok {
8596                        while bi + 8 <= alen {
8597                            let xs = [
8598                                acts[abase + bi].xq.as_slice(),
8599                                acts[abase + bi + 1].xq.as_slice(),
8600                                acts[abase + bi + 2].xq.as_slice(),
8601                                acts[abase + bi + 3].xq.as_slice(),
8602                                acts[abase + bi + 4].xq.as_slice(),
8603                                acts[abase + bi + 5].xq.as_slice(),
8604                                acts[abase + bi + 6].xq.as_slice(),
8605                                acts[abase + bi + 7].xq.as_slice(),
8606                            ];
8607                            let d = unsafe { dot_q4tp_row_1x8_avx512(v.nib, r, gpr, xs, &sc) };
8608                            for k in 0..8 {
8609                                let act = &acts[abase + bi + k];
8610                                let mut acc = d[k] * act.sx;
8611                                for &(j, xv) in &act.outliers {
8612                                    let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
8613                                    acc += w * s * xv;
8614                                }
8615                                // SAFETY: disjoint (bi, r) cells per worker.
8616                                unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
8617                            }
8618                            bi += 8;
8619                        }
8620                        while bi + 4 <= alen {
8621                            let xs = [
8622                                acts[abase + bi].xq.as_slice(),
8623                                acts[abase + bi + 1].xq.as_slice(),
8624                                acts[abase + bi + 2].xq.as_slice(),
8625                                acts[abase + bi + 3].xq.as_slice(),
8626                            ];
8627                            let d = unsafe { dot_q4tp_row_1x4_avx512(v.nib, r, gpr, xs, &sc) };
8628                            for k in 0..4 {
8629                                let act = &acts[abase + bi + k];
8630                                let mut acc = d[k] * act.sx;
8631                                for &(j, xv) in &act.outliers {
8632                                    let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
8633                                    acc += w * s * xv;
8634                                }
8635                                // SAFETY: disjoint (bi, r) cells per worker.
8636                                unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
8637                            }
8638                            bi += 4;
8639                        }
8640                    }
8641                    #[cfg(target_arch = "aarch64")]
8642                    if blocked_ok {
8643                        while bi + 4 <= alen {
8644                            let xs = [
8645                                acts[abase + bi].xq.as_slice(),
8646                                acts[abase + bi + 1].xq.as_slice(),
8647                                acts[abase + bi + 2].xq.as_slice(),
8648                                acts[abase + bi + 3].xq.as_slice(),
8649                            ];
8650                            let d = unsafe {
8651                                if exact || q4tp_v1() {
8652                                    dot_q4tp_row_1x4_sdot_v1(v.nib, r, gpr, xs, &sc)
8653                                } else {
8654                                    dot_q4tp_row_1x4_sdot(v.nib, r, gpr, xs, &sc)
8655                                }
8656                            };
8657                            for k in 0..4 {
8658                                let act = &acts[abase + bi + k];
8659                                let mut acc = d[k] * act.sx;
8660                                for &(j, xv) in &act.outliers {
8661                                    let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
8662                                    acc += w * s * xv;
8663                                }
8664                                // SAFETY: disjoint (bi, r) cells per worker.
8665                                unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
8666                            }
8667                            bi += 4;
8668                        }
8669                    }
8670                    let _ = blocked_ok;
8671                    // AVX2: the vertical column panel shares each group's
8672                    // unpack across four columns and keeps the row-exact
8673                    // cell bits (see `q4tp_matmat_row_cols_v`). Falls
8674                    // through to the single-column loop only when a column
8675                    // lacks its per-group correction.
8676                    #[cfg(target_arch = "x86_64")]
8677                    if acts[abase..abase + alen]
8678                        .iter()
8679                        .all(|a| a.xg8.len() == gpr)
8680                    {
8681                        // SAFETY: a8w8 implies AVX2; cells stay disjoint
8682                        // per (worker, bi, r).
8683                        unsafe {
8684                            q4tp_matmat_row_cols_v(
8685                                &v,
8686                                r,
8687                                gpr,
8688                                &acts[abase + bi..abase + alen],
8689                                &sc,
8690                                abase + bi,
8691                                rows,
8692                                out_addr,
8693                            )
8694                        };
8695                        bi = alen;
8696                    }
8697                    while bi < alen {
8698                        let act = &acts[abase + bi];
8699                        let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, &sc) * act.sx;
8700                        for &(j, xv) in &act.outliers {
8701                            let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
8702                            acc += w * s * xv;
8703                        }
8704                        // SAFETY: disjoint (bi, r) cells per worker range.
8705                        unsafe { *out_addr.at((abase + bi) * rows + r) = acc };
8706                        bi += 1;
8707                    }
8708                }
8709            }
8710        };
8711        dispatch_rows(pool, rows, &run);
8712        return;
8713    }
8714
8715    let run = |start: usize, end: usize| {
8716        let mut sc = vec![0f32; gpr];
8717        for r in start..end {
8718            v.scales_into(r, gpr, &mut sc);
8719            for bi in 0..b {
8720                let x = &xs_all[bi * cols..(bi + 1) * cols];
8721                // SAFETY: disjoint (bi, r) cells per worker range.
8722                unsafe { *out_addr.at(bi * rows + r) = q4tp_row_exact(v.nib, r, gpr, x, &sc) };
8723            }
8724        }
8725    };
8726    dispatch_rows(pool, rows, &run);
8727}
8728
8729/// Fused q4_tiled matvec (dispatch mirrors `q4matvec`).
8730fn q4t_matvec(
8731    bytes: &[u8],
8732    x: &[f32],
8733    rows: usize,
8734    cols: usize,
8735    out: &mut [f32],
8736    pool: Option<&Pool>,
8737) {
8738    debug_assert_eq!(out.len(), rows);
8739    let gpr = cols / GROUP_SIZE;
8740    let out_addr = SendMut(out.as_mut_ptr());
8741    if a8w8_enabled() {
8742        let act = split_act(x);
8743        let run = move |start: usize, end: usize| {
8744            for r in start..end {
8745                let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
8746                for &(j, xv) in &act.outliers {
8747                    let (w, s) = q4t_outlier(bytes, r, gpr, j);
8748                    acc += w * s * xv;
8749                }
8750                // SAFETY: disjoint row ranges per worker.
8751                unsafe { *out_addr.at(r) = acc };
8752            }
8753        };
8754        dispatch_rows(pool, rows, &run);
8755        return;
8756    }
8757    let run = move |start: usize, end: usize| {
8758        for r in start..end {
8759            // SAFETY: disjoint row ranges per worker.
8760            unsafe { *out_addr.at(r) = q4t_row_exact(bytes, r, gpr, x) };
8761        }
8762    };
8763    dispatch_rows(pool, rows, &run);
8764}
8765
8766/// Fused two-input q4_tiled matvec (weights read once per pair).
8767#[allow(clippy::too_many_arguments)]
8768fn q4t_matvec2(
8769    bytes: &[u8],
8770    x1: &[f32],
8771    x2: &[f32],
8772    rows: usize,
8773    cols: usize,
8774    o1: &mut [f32],
8775    o2: &mut [f32],
8776    pool: Option<&Pool>,
8777) {
8778    let gpr = cols / GROUP_SIZE;
8779    let p1 = SendMut(o1.as_mut_ptr());
8780    let p2 = SendMut(o2.as_mut_ptr());
8781    if a8w8_enabled() {
8782        let a1 = split_act(x1);
8783        let a2 = split_act(x2);
8784        let run = move |start: usize, end: usize| {
8785            for r in start..end {
8786                let mut v1 = dot_q4t_row_i8(bytes, r, gpr, &a1.xq) * a1.sx;
8787                let mut v2 = dot_q4t_row_i8(bytes, r, gpr, &a2.xq) * a2.sx;
8788                for &(j, xv) in &a1.outliers {
8789                    let (w, s) = q4t_outlier(bytes, r, gpr, j);
8790                    v1 += w * s * xv;
8791                }
8792                for &(j, xv) in &a2.outliers {
8793                    let (w, s) = q4t_outlier(bytes, r, gpr, j);
8794                    v2 += w * s * xv;
8795                }
8796                // SAFETY: disjoint row ranges per worker.
8797                unsafe {
8798                    *p1.at(r) = v1;
8799                    *p2.at(r) = v2;
8800                }
8801            }
8802        };
8803        dispatch_rows(pool, rows, &run);
8804        return;
8805    }
8806    let run = move |start: usize, end: usize| {
8807        for r in start..end {
8808            // SAFETY: disjoint row ranges per worker.
8809            unsafe {
8810                *p1.at(r) = q4t_row_exact(bytes, r, gpr, x1);
8811                *p2.at(r) = q4t_row_exact(bytes, r, gpr, x2);
8812            }
8813        }
8814    };
8815    dispatch_rows(pool, rows, &run);
8816}
8817
8818/// Batched q4_tiled matmat: each row's tiles stream once per microbatch.
8819#[allow(clippy::too_many_arguments)]
8820/// Prefill GEMM through Accelerate for group-quantized codecs: a
8821/// caller-supplied row dequantizer fills f32 tiles (pool-parallel) and
8822/// each tile rides the AMX with one sgemm — the generic sibling of
8823/// `qmatmat_accel` (q8). Numerics are f32-GEMM (tolerance class);
8824/// decode (b=1) never takes this path.
8825#[cfg(target_os = "macos")]
8826fn dequant_matmat_accel(
8827    dequant_row: &(dyn Fn(usize, &mut [f32]) + Sync),
8828    xs_all: &[f32],
8829    b: usize,
8830    rows: usize,
8831    cols: usize,
8832    out: &mut [f32],
8833    pool: Option<&Pool>,
8834) {
8835    const TR: usize = 2048;
8836    thread_local! {
8837        static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
8838    }
8839    WTILE.with(|wt| {
8840        let mut wtile = wt.borrow_mut();
8841        wtile.resize(TR * cols, 0.0);
8842        let mut r0 = 0usize;
8843        while r0 < rows {
8844            let tr = TR.min(rows - r0);
8845            let wt_addr = SendMut(wtile.as_mut_ptr());
8846            let run = |start: usize, end: usize| {
8847                for r in start..end {
8848                    // SAFETY: workers cover disjoint r ranges.
8849                    let dst = unsafe { std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols) };
8850                    dequant_row(r0 + r, dst);
8851                }
8852            };
8853            dispatch_rows(pool, tr, &run);
8854            unsafe {
8855                accel_blas::cblas_sgemm(
8856                    101, // RowMajor
8857                    111, // NoTrans A
8858                    112, // Trans B
8859                    b as i32,
8860                    tr as i32,
8861                    cols as i32,
8862                    1.0,
8863                    xs_all.as_ptr(),
8864                    cols as i32,
8865                    wtile.as_ptr(),
8866                    cols as i32,
8867                    0.0,
8868                    out.as_mut_ptr().add(r0),
8869                    rows as i32,
8870                );
8871            }
8872            r0 += tr;
8873        }
8874    });
8875}
8876
8877fn q4t_matmat(
8878    bytes: &[u8],
8879    xs_all: &[f32],
8880    b: usize,
8881    rows: usize,
8882    cols: usize,
8883    out: &mut [f32],
8884    pool: Option<&Pool>,
8885) {
8886    debug_assert_eq!(out.len(), b * rows);
8887    let gpr = cols / GROUP_SIZE;
8888    // Wide batches ride the AMX like q8's qmatmat: on Apple silicon
8889    // the dequant-tile sgemm is an order above the SDOT row loop for
8890    // prefill shapes (imagegen DiT forwards are exactly this).
8891    #[cfg(target_os = "macos")]
8892    if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
8893        dequant_matmat_accel(
8894            &|r, dst| {
8895                for gi in 0..gpr {
8896                    let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
8897                    let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
8898                    for (k, &bb) in tile[2..].iter().enumerate() {
8899                        dst[gi * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s;
8900                        dst[gi * GROUP_SIZE + k * 2 + 1] = (((bb >> 4) & 0x0F) as f32 - 8.0) * s;
8901                    }
8902                }
8903            },
8904            xs_all,
8905            b,
8906            rows,
8907            cols,
8908            out,
8909            pool,
8910        );
8911        return;
8912    }
8913    let out_addr = SendMut(out.as_mut_ptr());
8914    if a8w8_enabled() {
8915        let acts: Vec<SplitAct> = (0..b)
8916            .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
8917            .collect();
8918        let acts = &acts;
8919        #[cfg(target_arch = "x86_64")]
8920        let blocked_ok = avx2_enabled() && blocked_enabled();
8921        #[cfg(target_arch = "aarch64")]
8922        let blocked_ok = sdot_enabled() && blocked_enabled();
8923        #[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))]
8924        let blocked_ok = false;
8925        let run = move |start: usize, end: usize| {
8926            for r in start..end {
8927                let mut bi = 0usize;
8928                #[cfg(target_arch = "aarch64")]
8929                if blocked_ok {
8930                    while bi + 4 <= acts.len() {
8931                        let xs = [
8932                            acts[bi].xq.as_slice(),
8933                            acts[bi + 1].xq.as_slice(),
8934                            acts[bi + 2].xq.as_slice(),
8935                            acts[bi + 3].xq.as_slice(),
8936                        ];
8937                        let d = unsafe { dot_q4t_row_1x4_sdot(bytes, r, gpr, xs) };
8938                        for k in 0..4 {
8939                            let act = &acts[bi + k];
8940                            let mut acc = d[k] * act.sx;
8941                            for &(j, xv) in &act.outliers {
8942                                let (w, sc) = q4t_outlier(bytes, r, gpr, j);
8943                                acc += w * sc * xv;
8944                            }
8945                            // SAFETY: disjoint (bi, r) cells per worker.
8946                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
8947                        }
8948                        bi += 4;
8949                    }
8950                }
8951                #[cfg(target_arch = "x86_64")]
8952                if blocked_ok {
8953                    while bi + 4 <= acts.len() {
8954                        let xs = [
8955                            acts[bi].xq.as_slice(),
8956                            acts[bi + 1].xq.as_slice(),
8957                            acts[bi + 2].xq.as_slice(),
8958                            acts[bi + 3].xq.as_slice(),
8959                        ];
8960                        let d = unsafe {
8961                            if vnni_tiles_enabled() {
8962                                dot_q4t_row_1x4_vnni(bytes, r, gpr, xs)
8963                            } else {
8964                                dot_q4t_row_1x4_avx2(bytes, r, gpr, xs)
8965                            }
8966                        };
8967                        for k in 0..4 {
8968                            let act = &acts[bi + k];
8969                            let mut acc = d[k] * act.sx;
8970                            for &(j, xv) in &act.outliers {
8971                                let (w, sc) = q4t_outlier(bytes, r, gpr, j);
8972                                acc += w * sc * xv;
8973                            }
8974                            // SAFETY: disjoint (bi, r) cells per worker.
8975                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
8976                        }
8977                        bi += 4;
8978                    }
8979                }
8980                let _ = blocked_ok;
8981                while bi < acts.len() {
8982                    let act = &acts[bi];
8983                    let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
8984                    for &(j, xv) in &act.outliers {
8985                        let (w, s) = q4t_outlier(bytes, r, gpr, j);
8986                        acc += w * s * xv;
8987                    }
8988                    // SAFETY: disjoint (bi, r) cells per worker range.
8989                    unsafe { *out_addr.at(bi * rows + r) = acc };
8990                    bi += 1;
8991                }
8992            }
8993        };
8994        dispatch_rows(pool, rows, &run);
8995        return;
8996    }
8997    let run = move |start: usize, end: usize| {
8998        for r in start..end {
8999            for bi in 0..b {
9000                let x = &xs_all[bi * cols..(bi + 1) * cols];
9001                // SAFETY: disjoint (bi, r) cells per worker range.
9002                unsafe { *out_addr.at(bi * rows + r) = q4t_row_exact(bytes, r, gpr, x) };
9003            }
9004        }
9005    };
9006    dispatch_rows(pool, rows, &run);
9007}
9008
9009// ── q1 (dtype 12): binary weights, [f16 scale][4B sign bits] per
9010// 32-group tile. The kernel family mirrors q4_tiled: one sequential
9011// stream of 6-byte tiles, per-tile integer dot × scale, exact outlier
9012// correction (A8W8 contract), exact scalar path under CMF_SDOT=0. ──
9013
9014/// Per-32-group sums of the quantized activation — the ±1 identity's
9015/// shared half: `dot = −2·sdot(mask, x) − gsum[g]`, computed ONCE per
9016/// matvec and reused by every row.
9017fn q1_group_sums(xq: &[i8], gpr: usize) -> Vec<i32> {
9018    (0..gpr)
9019        .map(|gi| {
9020            xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE]
9021                .iter()
9022                .map(|&v| v as i32)
9023                .sum()
9024        })
9025        .collect()
9026}
9027
9028/// One q1 row via the A8W8 int8 path — mask-SDOT on ARM (no ±1
9029/// expansion at all), scalar bit loop elsewhere (AVX2 queued with the
9030/// x86 pass).
9031#[inline]
9032#[allow(unreachable_code)]
9033/// AVX2 q1 row via the same ±1 identity as the ARM sdot kernel: the
9034/// sign bits expand to a {0, −1} byte mask through shuffle+cmpeq, the
9035/// masked activation sums through maddubs(1, x&mask), and
9036/// `dot = −(2·masked_sum + Σx_group)` — bit-identical integer math.
9037#[cfg(target_arch = "x86_64")]
9038#[target_feature(enable = "avx2")]
9039unsafe fn dot_q1_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
9040    // SAFETY: callers uphold the 6B-tile and xq/gsum length contracts.
9041    unsafe {
9042        use core::arch::x86_64::*;
9043        // Byte j of the mask must replicate bits-byte j/8.
9044        let expand = _mm256_setr_epi8(
9045            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,
9046            3, 3, 3,
9047        );
9048        let bitsel = _mm256_setr_epi8(
9049            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
9050            -128, 1, 2, 4, 8, 16, 32, 64, -128,
9051        );
9052        let ones8 = _mm256_set1_epi8(1);
9053        let ones16 = _mm256_set1_epi16(1);
9054        let mut acc = 0f32;
9055        for gi in 0..gpr {
9056            let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
9057            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
9058            let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
9059            let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
9060            let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
9061            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9062            let sel = _mm256_and_si256(x, mask);
9063            // Σ of selected i8 lanes: maddubs(1u8, sel_i8) pairs → madd.
9064            let p16 = _mm256_maddubs_epi16(ones8, sel);
9065            let d32 = _mm256_madd_epi16(p16, ones16);
9066            let hi128 = _mm256_extracti128_si256::<1>(d32);
9067            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
9068            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
9069            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
9070            let msum = _mm_cvtsi128_si32(s32);
9071            // The and-select keeps x UN-negated (unlike ARM's −1-mask
9072            // sdot): d = Σ_set − Σ_unset = 2·Σ_set − Σ_all.
9073            let d = 2 * msum - gsum[gi];
9074            acc += d as f32 * s;
9075        }
9076        acc
9077    }
9078}
9079
9080/// VNNI twin of `dot_q1_row_avx2`: the masked-select sum goes through
9081/// one `vpdpbusd(1u8, sel)` (see `dpbusd_hsum` — bit-identical).
9082#[cfg(target_arch = "x86_64")]
9083#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
9084unsafe fn dot_q1_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
9085    // SAFETY: callers uphold the 6B-tile and xq/gsum length contracts.
9086    unsafe {
9087        use core::arch::x86_64::*;
9088        let expand = _mm256_setr_epi8(
9089            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,
9090            3, 3, 3,
9091        );
9092        let bitsel = _mm256_setr_epi8(
9093            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
9094            -128, 1, 2, 4, 8, 16, 32, 64, -128,
9095        );
9096        let ones8 = _mm256_set1_epi8(1);
9097        let mut acc = 0f32;
9098        for gi in 0..gpr {
9099            let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
9100            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
9101            let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
9102            let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
9103            let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
9104            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9105            let msum = dpbusd_hsum(ones8, _mm256_and_si256(x, mask));
9106            let d = 2 * msum - gsum[gi];
9107            acc += d as f32 * s;
9108        }
9109        acc
9110    }
9111}
9112
9113/// VNNI twin of `dot_q1_row_1x4_avx2` (see `dpbusd_hsum`).
9114#[cfg(target_arch = "x86_64")]
9115#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
9116unsafe fn dot_q1_row_1x4_vnni(
9117    bytes: &[u8],
9118    r: usize,
9119    gpr: usize,
9120    xs: [&[i8]; 4],
9121    gsums: [&[i32]; 4],
9122) -> [f32; 4] {
9123    // SAFETY: callers uphold the 6B-tile and xq/gsum length contracts.
9124    unsafe {
9125        use core::arch::x86_64::*;
9126        let expand = _mm256_setr_epi8(
9127            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,
9128            3, 3, 3,
9129        );
9130        let bitsel = _mm256_setr_epi8(
9131            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
9132            -128, 1, 2, 4, 8, 16, 32, 64, -128,
9133        );
9134        let ones8 = _mm256_set1_epi8(1);
9135        let mut acc = [0f32; 4];
9136        for gi in 0..gpr {
9137            let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
9138            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
9139            let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
9140            let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
9141            let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
9142            for (k, xq) in xs.iter().enumerate() {
9143                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9144                let msum = dpbusd_hsum(ones8, _mm256_and_si256(x, mask));
9145                let d = 2 * msum - gsums[k][gi];
9146                acc[k] += d as f32 * s;
9147            }
9148        }
9149        acc
9150    }
9151}
9152
9153/// The blocked 1×4 flavor: the expanded bit mask serves four activation
9154/// streams per group (mask build once, four select+reduce chains).
9155#[cfg(target_arch = "x86_64")]
9156#[target_feature(enable = "avx2")]
9157unsafe fn dot_q1_row_1x4_avx2(
9158    bytes: &[u8],
9159    r: usize,
9160    gpr: usize,
9161    xs: [&[i8]; 4],
9162    gsums: [&[i32]; 4],
9163) -> [f32; 4] {
9164    // SAFETY: callers uphold the 6B-tile and xq/gsum length contracts.
9165    unsafe {
9166        use core::arch::x86_64::*;
9167        let expand = _mm256_setr_epi8(
9168            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,
9169            3, 3, 3,
9170        );
9171        let bitsel = _mm256_setr_epi8(
9172            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
9173            -128, 1, 2, 4, 8, 16, 32, 64, -128,
9174        );
9175        let ones8 = _mm256_set1_epi8(1);
9176        let ones16 = _mm256_set1_epi16(1);
9177        let mut acc = [0f32; 4];
9178        for gi in 0..gpr {
9179            let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
9180            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
9181            let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
9182            let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
9183            let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
9184            for (k, xq) in xs.iter().enumerate() {
9185                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9186                let sel = _mm256_and_si256(x, mask);
9187                let p16 = _mm256_maddubs_epi16(ones8, sel);
9188                let d32 = _mm256_madd_epi16(p16, ones16);
9189                let hi128 = _mm256_extracti128_si256::<1>(d32);
9190                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
9191                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
9192                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
9193                let msum = _mm_cvtsi128_si32(s32);
9194                let d = 2 * msum - gsums[k][gi];
9195                acc[k] += d as f32 * s;
9196            }
9197        }
9198        acc
9199    }
9200}
9201
9202#[allow(unreachable_code)]
9203fn dot_q1_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
9204    #[cfg(target_arch = "aarch64")]
9205    unsafe {
9206        return dot_q1_row_sdot(bytes, r, gpr, xq, gsum);
9207    }
9208    #[cfg(target_arch = "x86_64")]
9209    if avx2_enabled() {
9210        unsafe {
9211            if vnni_tiles_enabled() {
9212                return dot_q1_row_vnni(bytes, r, gpr, xq, gsum);
9213            }
9214            return dot_q1_row_avx2(bytes, r, gpr, xq, gsum);
9215        }
9216    }
9217    let _ = gsum;
9218    let mut acc = 0f32;
9219    for gi in 0..gpr {
9220        let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
9221        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
9222        let mut d = 0i32;
9223        for (j, &b) in tile[2..].iter().enumerate() {
9224            for k in 0..8 {
9225                let w = ((b >> k) & 1) as i32 * 2 - 1;
9226                d += w * xq[gi * GROUP_SIZE + j * 8 + k] as i32;
9227            }
9228        }
9229        acc += d as f32 * s;
9230    }
9231    acc
9232}
9233
9234/// SDOT q1 row via the ±1 identity: the vtst mask (0xFF where the bit
9235/// is set, i.e. −1 as i8) feeds `sdot` DIRECTLY — no expansion to ±1
9236/// lanes at all — and `dot = −(2·sdot(mask, x) + Σx_group)`, with the
9237/// per-group activation sums shared across every row of the matvec.
9238/// Four tiles (128 weights) per iteration: integer dots reduce through
9239/// a vpaddq tree into ONE i32x4 that meets its four scales in a single
9240/// fused f32 multiply-add. Integer math throughout — bit-identical to
9241/// the scalar ±1 reference.
9242#[cfg(target_arch = "aarch64")]
9243#[target_feature(enable = "neon,dotprod")]
9244unsafe fn dot_q1_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
9245    // SAFETY: callers uphold slice-length contracts (6B tile per group,
9246    // xq.len() == gpr·GROUP_SIZE, gsum.len() == gpr).
9247    unsafe {
9248        use core::arch::aarch64::*;
9249        use core::arch::asm;
9250        const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
9251        let m = vld1q_u8(MASKS.as_ptr());
9252        // One tile's −Σ_set(x) as an UNREDUCED i32x4 (two mask-sdots).
9253        macro_rules! tile_dot {
9254            ($t:expr, $x:expr) => {{
9255                let v0 = vcombine_u8(vdup_n_u8(*$t.add(2)), vdup_n_u8(*$t.add(3)));
9256                let v1 = vcombine_u8(vdup_n_u8(*$t.add(4)), vdup_n_u8(*$t.add(5)));
9257                let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
9258                let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
9259                let x0 = vld1q_s8($x);
9260                let x1 = vld1q_s8($x.add(16));
9261                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
9262                asm!(
9263                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
9264                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
9265                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
9266                    w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
9267                    options(pure, nomem, nostack),
9268                );
9269                vaddq_s32(a0, a1)
9270            }};
9271        }
9272        // TBL unpack over PAIR loads: one vld1q covers two 6B tiles
9273        // ([s s b b b b][s s b b b b] + 4B slack), TBL replicates each
9274        // bit-byte across 8 lanes for vtst, and the four scales gather
9275        // through tbl2 into one fcvtl — the 16 ld1r broadcast loads and
9276        // 4 branchy software f16 conversions per 128 weights (the
9277        // measured load-port wall of this kernel) become 2 vector
9278        // loads + 9 table lookups. Integer math order is unchanged —
9279        // bit-identical results (FCVTL is exact on every f16).
9280        const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
9281        const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
9282        const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
9283        const IW11: [u8; 16] = [
9284            10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
9285        ];
9286        const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
9287        let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
9288        let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
9289        let isc = vld1_u8(ISC.as_ptr());
9290        // One tile's −Σ_set(x) from a TBL-unpacked pair load.
9291        macro_rules! tile_dot_tbl {
9292            ($ld:expr, $i0:expr, $i1:expr, $x:expr) => {{
9293                let w0 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i0), m));
9294                let w1 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i1), m));
9295                let x0 = vld1q_s8($x);
9296                let x1 = vld1q_s8($x.add(16));
9297                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
9298                asm!(
9299                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
9300                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
9301                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
9302                    w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
9303                    options(pure, nomem, nostack),
9304                );
9305                vaddq_s32(a0, a1)
9306            }};
9307        }
9308        let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
9309        let row_base = r * gpr * Q1_TILE;
9310        let abs_end = bytes.len();
9311        let xp = xq.as_ptr();
9312        let gp = gsum.as_ptr();
9313        let mut accv = vdupq_n_f32(0.0);
9314        let mut gi = 0;
9315        // The second pair load reads 4B past tile gi+3 — stay inside
9316        // the payload slice (only the file's final tiles fall back).
9317        while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
9318            let t0 = base.add(gi * Q1_TILE);
9319            let ld_a = vld1q_u8(t0);
9320            let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
9321            let d0 = tile_dot_tbl!(ld_a, iw00, iw01, xp.add(gi * GROUP_SIZE));
9322            let d1 = tile_dot_tbl!(ld_a, iw10, iw11, xp.add((gi + 1) * GROUP_SIZE));
9323            let d2 = tile_dot_tbl!(ld_b, iw00, iw01, xp.add((gi + 2) * GROUP_SIZE));
9324            let d3 = tile_dot_tbl!(ld_b, iw10, iw11, xp.add((gi + 3) * GROUP_SIZE));
9325            // [−Σ0, −Σ1, −Σ2, −Σ3] → dots = −(2·Σset_neg + gsum)
9326            let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
9327            let g = vld1q_s32(gp.add(gi));
9328            let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
9329            let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
9330            let scf: float32x4_t;
9331            asm!(
9332                "fcvtl {o:v}.4s, {i:v}.4h",
9333                o = out(vreg) scf, i = in(vreg) sc16,
9334                options(pure, nomem, nostack),
9335            );
9336            accv = vfmaq_f32(accv, vcvtq_f32_s32(dots), scf);
9337            gi += 4;
9338        }
9339        let mut acc = vaddvq_f32(accv);
9340        while gi < gpr {
9341            let t = base.add(gi * Q1_TILE);
9342            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
9343            let d = vaddvq_s32(tile_dot!(t, xp.add(gi * GROUP_SIZE)));
9344            acc += (-(2 * d + *gp.add(gi))) as f32 * s;
9345            gi += 1;
9346        }
9347        acc
9348    }
9349}
9350
9351/// Blocked q1 1×4: one TBL unpack of the tile pair serves FOUR
9352/// activation streams (prefill amortization — the same idea as the
9353/// AVX2 twin; per stream the group order, fma order and tail match the
9354/// single-row kernel exactly, so batch == matvec bit-for-bit).
9355#[cfg(target_arch = "aarch64")]
9356#[target_feature(enable = "neon,dotprod")]
9357unsafe fn dot_q1_row_1x4_sdot(
9358    bytes: &[u8],
9359    r: usize,
9360    gpr: usize,
9361    xs: [&[i8]; 4],
9362    gs: [&[i32]; 4],
9363) -> [f32; 4] {
9364    // SAFETY: same slice-length contracts as `dot_q1_row_sdot`, ×4.
9365    unsafe {
9366        use core::arch::aarch64::*;
9367        use core::arch::asm;
9368        const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
9369        const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
9370        const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
9371        const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
9372        const IW11: [u8; 16] = [
9373            10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
9374        ];
9375        const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
9376        let m = vld1q_u8(MASKS.as_ptr());
9377        let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
9378        let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
9379        let isc = vld1_u8(ISC.as_ptr());
9380        macro_rules! sdot2 {
9381            ($w0:expr, $w1:expr, $x:expr) => {{
9382                let x0 = vld1q_s8($x);
9383                let x1 = vld1q_s8($x.add(16));
9384                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
9385                asm!(
9386                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
9387                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
9388                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
9389                    w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
9390                    options(pure, nomem, nostack),
9391                );
9392                vaddq_s32(a0, a1)
9393            }};
9394        }
9395        let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
9396        let row_base = r * gpr * Q1_TILE;
9397        let abs_end = bytes.len();
9398        let mut accv = [vdupq_n_f32(0.0); 4];
9399        let mut gi = 0;
9400        while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
9401            let t0 = base.add(gi * Q1_TILE);
9402            let ld_a = vld1q_u8(t0);
9403            let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
9404            // Unpack ONCE — eight ±mask vectors serve all four streams.
9405            let w00 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw00), m));
9406            let w01 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw01), m));
9407            let w10 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw10), m));
9408            let w11 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw11), m));
9409            let w20 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw00), m));
9410            let w21 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw01), m));
9411            let w30 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw10), m));
9412            let w31 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw11), m));
9413            let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
9414            let scf: float32x4_t;
9415            asm!(
9416                "fcvtl {o:v}.4s, {i:v}.4h",
9417                o = out(vreg) scf, i = in(vreg) sc16,
9418                options(pure, nomem, nostack),
9419            );
9420            for k in 0..4 {
9421                let xp = xs[k].as_ptr();
9422                let d0 = sdot2!(w00, w01, xp.add(gi * GROUP_SIZE));
9423                let d1 = sdot2!(w10, w11, xp.add((gi + 1) * GROUP_SIZE));
9424                let d2 = sdot2!(w20, w21, xp.add((gi + 2) * GROUP_SIZE));
9425                let d3 = sdot2!(w30, w31, xp.add((gi + 3) * GROUP_SIZE));
9426                let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
9427                let g = vld1q_s32(gs[k].as_ptr().add(gi));
9428                let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
9429                accv[k] = vfmaq_f32(accv[k], vcvtq_f32_s32(dots), scf);
9430            }
9431            gi += 4;
9432        }
9433        let mut acc = [
9434            vaddvq_f32(accv[0]),
9435            vaddvq_f32(accv[1]),
9436            vaddvq_f32(accv[2]),
9437            vaddvq_f32(accv[3]),
9438        ];
9439        while gi < gpr {
9440            let t = base.add(gi * Q1_TILE);
9441            let sc = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
9442            let v0 = vcombine_u8(vdup_n_u8(*t.add(2)), vdup_n_u8(*t.add(3)));
9443            let v1 = vcombine_u8(vdup_n_u8(*t.add(4)), vdup_n_u8(*t.add(5)));
9444            let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
9445            let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
9446            for k in 0..4 {
9447                let d = vaddvq_s32(sdot2!(w0, w1, xs[k].as_ptr().add(gi * GROUP_SIZE)));
9448                acc[k] += (-(2 * d + *gs[k].as_ptr().add(gi))) as f32 * sc;
9449            }
9450            gi += 1;
9451        }
9452        acc
9453    }
9454}
9455
9456/// (weight ±1, scale) of one q1 element — the exact outlier term.
9457#[inline]
9458fn q1_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
9459    let gi = j / GROUP_SIZE;
9460    let k = j % GROUP_SIZE;
9461    let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
9462    let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
9463    let bit = (tile[2 + k / 8] >> (k % 8)) & 1;
9464    ((bit as i32 * 2 - 1) as f32, s)
9465}
9466
9467/// Exact scalar q1 row (CMF_SDOT=0 contract).
9468#[inline]
9469fn q1_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
9470    let mut acc = 0f32;
9471    for gi in 0..gpr {
9472        let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
9473        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
9474        let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
9475        let mut ga = 0f32;
9476        for (j, &b) in tile[2..].iter().enumerate() {
9477            for k in 0..8 {
9478                ga += (((b >> k) & 1) as f32 * 2.0 - 1.0) * xg[j * 8 + k];
9479            }
9480        }
9481        acc += ga * s;
9482    }
9483    acc
9484}
9485
9486/// One binary-weight row range via A8W8 (the body of `bin_matvec`'s hot
9487/// loop, extracted so multi-matrix jobs drive the same kernel).
9488#[allow(clippy::too_many_arguments)]
9489fn bin_range_a8w8(
9490    w: BinW,
9491    gpr: usize,
9492    act: &SplitAct,
9493    gsum: &[i32],
9494    out: SendMut,
9495    start: usize,
9496    end: usize,
9497) {
9498    for r in start..end {
9499        let mut acc = w.dot_i8(r, gpr, &act.xq, gsum) * act.sx;
9500        for &(j, xv) in &act.outliers {
9501            let (wv, s) = w.outlier(r, gpr, j);
9502            acc += wv * s * xv;
9503        }
9504        // SAFETY: disjoint row ranges per worker.
9505        unsafe { *out.at(r) = acc };
9506    }
9507}
9508
9509/// Exact-scalar binary-weight row range (CMF_SDOT=0 contract).
9510fn bin_range_f32(w: BinW, gpr: usize, x: &[f32], out: SendMut, start: usize, end: usize) {
9511    for r in start..end {
9512        // SAFETY: disjoint row ranges per worker.
9513        unsafe { *out.at(r) = w.row_exact(r, gpr, x) };
9514    }
9515}
9516
9517/// A binary (±s) weight in either of its layouts: `q1` 6-byte tiles, or
9518/// `q1g` planes with one scale per 128-group. The row drivers — matvec,
9519/// matvec2, matmat, the multi-projection job and the fused gate/up pair —
9520/// are written once against this; only the per-row kernels differ, and
9521/// the match costs one predictable branch per row. `gpr` is always the
9522/// row's count of 32-tiles (cols / 32): the activation group sums and the
9523/// tile walk are shared by both layouts.
9524#[derive(Clone, Copy)]
9525pub(crate) enum BinW<'a> {
9526    Q1(&'a [u8]),
9527    Q1g {
9528        bytes: &'a [u8],
9529        /// Byte offset of the f16 scale plane (= rows · cols/8).
9530        sc_off: usize,
9531        /// Row count: the q1 kernels treat the payload's final row
9532        /// specially (see `dot_q1g_row_sdot`), and q1g mirrors that.
9533        rows: usize,
9534    },
9535}
9536
9537impl<'a> BinW<'a> {
9538    fn q1g(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
9539        Self::Q1g {
9540            bytes,
9541            sc_off: cortiq_core::quant::q1g_scale_offset(rows, cols),
9542            rows,
9543        }
9544    }
9545
9546    #[inline]
9547    fn dot_i8(self, r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
9548        match self {
9549            Self::Q1(b) => dot_q1_row_i8(b, r, gpr, xq, gsum),
9550            Self::Q1g {
9551                bytes,
9552                sc_off,
9553                rows,
9554            } => dot_q1g_row_i8(bytes, sc_off, rows, r, gpr, xq, gsum),
9555        }
9556    }
9557
9558    #[inline]
9559    fn row_exact(self, r: usize, gpr: usize, x: &[f32]) -> f32 {
9560        match self {
9561            Self::Q1(b) => q1_row_exact(b, r, gpr, x),
9562            Self::Q1g { bytes, sc_off, .. } => q1g_row_exact(bytes, sc_off, r, gpr, x),
9563        }
9564    }
9565
9566    #[inline]
9567    fn outlier(self, r: usize, gpr: usize, j: usize) -> (f32, f32) {
9568        match self {
9569            Self::Q1(b) => q1_outlier(b, r, gpr, j),
9570            Self::Q1g { bytes, sc_off, .. } => q1g_outlier(bytes, sc_off, r, gpr, j),
9571        }
9572    }
9573
9574    /// Blocked 1×4 row (one unpack, four activation streams).
9575    ///
9576    /// # Safety
9577    /// The SDOT/dotprod target feature must be present (`sdot_enabled`).
9578    #[cfg(target_arch = "aarch64")]
9579    #[inline]
9580    unsafe fn dot_1x4_sdot(
9581        self,
9582        r: usize,
9583        gpr: usize,
9584        xs: [&[i8]; 4],
9585        gs: [&[i32]; 4],
9586    ) -> [f32; 4] {
9587        // SAFETY: forwarded caller contract.
9588        unsafe {
9589            match self {
9590                Self::Q1(b) => dot_q1_row_1x4_sdot(b, r, gpr, xs, gs),
9591                Self::Q1g {
9592                    bytes,
9593                    sc_off,
9594                    rows,
9595                } => dot_q1g_row_1x4_sdot(bytes, sc_off, rows, r, gpr, xs, gs),
9596            }
9597        }
9598    }
9599
9600    /// Blocked 1×4 row on x86 (VNNI when the tile switch is on).
9601    ///
9602    /// # Safety
9603    /// AVX2 must be present (`avx2_enabled`).
9604    #[cfg(target_arch = "x86_64")]
9605    #[inline]
9606    unsafe fn dot_1x4_x86(self, r: usize, gpr: usize, xs: [&[i8]; 4], gs: [&[i32]; 4]) -> [f32; 4] {
9607        // SAFETY: forwarded caller contract.
9608        unsafe {
9609            match (self, vnni_tiles_enabled()) {
9610                (Self::Q1(b), true) => dot_q1_row_1x4_vnni(b, r, gpr, xs, gs),
9611                (Self::Q1(b), false) => dot_q1_row_1x4_avx2(b, r, gpr, xs, gs),
9612                (Self::Q1g { bytes, sc_off, .. }, true) => {
9613                    dot_q1g_row_1x4_vnni(bytes, sc_off, r, gpr, xs, gs)
9614                }
9615                (Self::Q1g { bytes, sc_off, .. }, false) => {
9616                    dot_q1g_row_1x4_avx2(bytes, sc_off, r, gpr, xs, gs)
9617                }
9618            }
9619        }
9620    }
9621}
9622
9623// ── q1g (dtype 20): binary weights, ONE f16 scale per 128-group, in sign
9624// and scale PLANES (`[bits: rows·cols/8][f16: rows·cols/128]`). Same
9625// kernel family as q1 and — deliberately — the same float accumulation
9626// order (one `acc += d·s` per 32-tile, the group's scale reused four
9627// times), so a q1g repack of a q1 tensor whose tile scales agree gives
9628// the same f32 outputs bit for bit on every path. The win is the stream:
9629// 18 bytes per 128 weights instead of 24, the bits 16-byte aligned. ──
9630
9631/// Sign bits of q1g tile `t` (flat 32-tile index `r·gpr + gi`).
9632#[inline]
9633fn q1g_bits(bytes: &[u8], t: usize) -> u32 {
9634    u32::from_le_bytes([
9635        bytes[t * 4],
9636        bytes[t * 4 + 1],
9637        bytes[t * 4 + 2],
9638        bytes[t * 4 + 3],
9639    ])
9640}
9641
9642/// Raw f16 bits of the scale of the 128-group holding q1g tile `t`.
9643#[inline]
9644fn q1g_scale_bits(bytes: &[u8], sc_off: usize, t: usize) -> u16 {
9645    let o = sc_off + (t / 4) * 2;
9646    u16::from_le_bytes([bytes[o], bytes[o + 1]])
9647}
9648
9649/// One q1g row via the A8W8 int8 path (dispatch twin of `dot_q1_row_i8`).
9650#[allow(unreachable_code)]
9651#[allow(clippy::too_many_arguments)]
9652fn dot_q1g_row_i8(
9653    bytes: &[u8],
9654    sc_off: usize,
9655    rows: usize,
9656    r: usize,
9657    gpr: usize,
9658    xq: &[i8],
9659    gsum: &[i32],
9660) -> f32 {
9661    #[cfg(target_arch = "aarch64")]
9662    unsafe {
9663        return dot_q1g_row_sdot(bytes, sc_off, rows, r, gpr, xq, gsum);
9664    }
9665    #[cfg(target_arch = "x86_64")]
9666    if avx2_enabled() {
9667        unsafe {
9668            if vnni_tiles_enabled() {
9669                return dot_q1g_row_vnni(bytes, sc_off, r, gpr, xq, gsum);
9670            }
9671            return dot_q1g_row_avx2(bytes, sc_off, r, gpr, xq, gsum);
9672        }
9673    }
9674    let _ = (gsum, rows);
9675    let mut acc = 0f32;
9676    for gi in 0..gpr {
9677        let t = r * gpr + gi;
9678        let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, t));
9679        let bits = q1g_bits(bytes, t);
9680        let mut d = 0i32;
9681        for k in 0..GROUP_SIZE {
9682            let w = ((bits >> k) & 1) as i32 * 2 - 1;
9683            d += w * xq[gi * GROUP_SIZE + k] as i32;
9684        }
9685        acc += d as f32 * s;
9686    }
9687    acc
9688}
9689
9690/// Four tiles' unreduced `i32×8` partial sums → `[Σ0, Σ1, Σ2, Σ3]` in one
9691/// hadd tree: 3 hadds + 1 add for four tiles where the per-tile reduce of
9692/// the q1 kernels costs 7 ops each. Integer math, so any order gives the
9693/// q1 kernel's per-tile sums exactly.
9694#[cfg(target_arch = "x86_64")]
9695#[target_feature(enable = "avx2")]
9696#[inline]
9697unsafe fn q1g_hsum4(
9698    d0: core::arch::x86_64::__m256i,
9699    d1: core::arch::x86_64::__m256i,
9700    d2: core::arch::x86_64::__m256i,
9701    d3: core::arch::x86_64::__m256i,
9702) -> core::arch::x86_64::__m128i {
9703    // SAFETY: pure register math.
9704    unsafe {
9705        use core::arch::x86_64::*;
9706        let h = _mm256_hadd_epi32(_mm256_hadd_epi32(d0, d1), _mm256_hadd_epi32(d2, d3));
9707        _mm_add_epi32(_mm256_castsi256_si128(h), _mm256_extracti128_si256::<1>(h))
9708    }
9709}
9710
9711/// One 128-group's epilogue, in q1's float order: `d = 2·Σset − Σx` per
9712/// tile (exact i32 → exact f32), then `acc += d·s` tile by tile — four
9713/// separately rounded multiply/add pairs, the q1 kernels' sequence.
9714#[cfg(target_arch = "x86_64")]
9715#[target_feature(enable = "avx2")]
9716#[inline]
9717unsafe fn q1g_group_acc(
9718    acc: f32,
9719    msum4: core::arch::x86_64::__m128i,
9720    gsum4: &[i32],
9721    s: f32,
9722) -> f32 {
9723    // SAFETY: gsum4 holds the group's four tile sums (callers slice it).
9724    unsafe {
9725        use core::arch::x86_64::*;
9726        let g = _mm_loadu_si128(gsum4.as_ptr() as *const __m128i);
9727        let d = _mm_cvtepi32_ps(_mm_sub_epi32(_mm_slli_epi32::<1>(msum4), g));
9728        let mut f = [0f32; 4];
9729        _mm_storeu_ps(f.as_mut_ptr(), d);
9730        let mut acc = acc;
9731        for v in f {
9732            acc += v * s;
9733        }
9734        acc
9735    }
9736}
9737
9738/// AVX2 q1g row: `dot_q1_row_avx2`'s mask/select/maddubs per tile, but the
9739/// four tiles of a 128-group reduce together (`q1g_hsum4`) and the scale
9740/// converts once per group — bit-identical results (integer sums, then
9741/// q1's per-tile float sequence), fewer horizontal ops and f16 decodes.
9742#[cfg(target_arch = "x86_64")]
9743#[target_feature(enable = "avx2")]
9744unsafe fn dot_q1g_row_avx2(
9745    bytes: &[u8],
9746    sc_off: usize,
9747    r: usize,
9748    gpr: usize,
9749    xq: &[i8],
9750    gsum: &[i32],
9751) -> f32 {
9752    // SAFETY: callers uphold the plane-length and xq/gsum contracts
9753    // (gpr % 4 == 0, xq.len() == gpr·32, gsum.len() == gpr).
9754    unsafe {
9755        use core::arch::x86_64::*;
9756        let expand = _mm256_setr_epi8(
9757            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,
9758            3, 3, 3,
9759        );
9760        let bitsel = _mm256_setr_epi8(
9761            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
9762            -128, 1, 2, 4, 8, 16, 32, 64, -128,
9763        );
9764        let ones8 = _mm256_set1_epi8(1);
9765        let ones16 = _mm256_set1_epi16(1);
9766        let bp = bytes.as_ptr().add(r * gpr * 4);
9767        let mut acc = 0f32;
9768        let mut gi = 0;
9769        while gi < gpr {
9770            let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
9771            let mut d = [_mm256_setzero_si256(); 4];
9772            for (k, dk) in d.iter_mut().enumerate() {
9773                let t = gi + k;
9774                let bits = (bp.add(t * 4) as *const u32).read_unaligned();
9775                let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
9776                let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
9777                let x = _mm256_loadu_si256(xq.as_ptr().add(t * GROUP_SIZE) as *const __m256i);
9778                let sel = _mm256_and_si256(x, mask);
9779                *dk = _mm256_madd_epi16(_mm256_maddubs_epi16(ones8, sel), ones16);
9780            }
9781            let m4 = q1g_hsum4(d[0], d[1], d[2], d[3]);
9782            acc = q1g_group_acc(acc, m4, &gsum[gi..gi + 4], s);
9783            gi += 4;
9784        }
9785        acc
9786    }
9787}
9788
9789/// VNNI twin of `dot_q1g_row_avx2`: `vpdpbusd(1u8, sel)` per tile (the
9790/// same i32 lanes as maddubs+madd — bounded products), same group reduce.
9791#[cfg(target_arch = "x86_64")]
9792#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
9793unsafe fn dot_q1g_row_vnni(
9794    bytes: &[u8],
9795    sc_off: usize,
9796    r: usize,
9797    gpr: usize,
9798    xq: &[i8],
9799    gsum: &[i32],
9800) -> f32 {
9801    // SAFETY: callers uphold the plane-length and xq/gsum contracts.
9802    unsafe {
9803        use core::arch::x86_64::*;
9804        let expand = _mm256_setr_epi8(
9805            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,
9806            3, 3, 3,
9807        );
9808        let bitsel = _mm256_setr_epi8(
9809            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
9810            -128, 1, 2, 4, 8, 16, 32, 64, -128,
9811        );
9812        let ones8 = _mm256_set1_epi8(1);
9813        let bp = bytes.as_ptr().add(r * gpr * 4);
9814        let mut acc = 0f32;
9815        let mut gi = 0;
9816        while gi < gpr {
9817            let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
9818            let mut d = [_mm256_setzero_si256(); 4];
9819            for (k, dk) in d.iter_mut().enumerate() {
9820                let t = gi + k;
9821                let bits = (bp.add(t * 4) as *const u32).read_unaligned();
9822                let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
9823                let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
9824                let x = _mm256_loadu_si256(xq.as_ptr().add(t * GROUP_SIZE) as *const __m256i);
9825                *dk = _mm256_dpbusd_epi32(_mm256_setzero_si256(), ones8, _mm256_and_si256(x, mask));
9826            }
9827            let m4 = q1g_hsum4(d[0], d[1], d[2], d[3]);
9828            acc = q1g_group_acc(acc, m4, &gsum[gi..gi + 4], s);
9829            gi += 4;
9830        }
9831        acc
9832    }
9833}
9834
9835/// Blocked 1×4 AVX2 q1g row: each tile's mask is built once and serves
9836/// four activation streams; per stream, the group reduce and float order
9837/// of `dot_q1g_row_avx2` (so batch == matvec bit for bit).
9838#[cfg(target_arch = "x86_64")]
9839#[target_feature(enable = "avx2")]
9840unsafe fn dot_q1g_row_1x4_avx2(
9841    bytes: &[u8],
9842    sc_off: usize,
9843    r: usize,
9844    gpr: usize,
9845    xs: [&[i8]; 4],
9846    gsums: [&[i32]; 4],
9847) -> [f32; 4] {
9848    // SAFETY: callers uphold the plane-length and xq/gsum contracts.
9849    unsafe {
9850        use core::arch::x86_64::*;
9851        let expand = _mm256_setr_epi8(
9852            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,
9853            3, 3, 3,
9854        );
9855        let bitsel = _mm256_setr_epi8(
9856            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
9857            -128, 1, 2, 4, 8, 16, 32, 64, -128,
9858        );
9859        let ones8 = _mm256_set1_epi8(1);
9860        let ones16 = _mm256_set1_epi16(1);
9861        let bp = bytes.as_ptr().add(r * gpr * 4);
9862        let mut acc = [0f32; 4];
9863        let mut gi = 0;
9864        while gi < gpr {
9865            let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
9866            // d[stream][tile]
9867            let mut d = [[_mm256_setzero_si256(); 4]; 4];
9868            for t in 0..4 {
9869                let bits = (bp.add((gi + t) * 4) as *const u32).read_unaligned();
9870                let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
9871                let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
9872                for (k, xq) in xs.iter().enumerate() {
9873                    let x = _mm256_loadu_si256(
9874                        xq.as_ptr().add((gi + t) * GROUP_SIZE) as *const __m256i
9875                    );
9876                    let sel = _mm256_and_si256(x, mask);
9877                    d[k][t] = _mm256_madd_epi16(_mm256_maddubs_epi16(ones8, sel), ones16);
9878                }
9879            }
9880            for k in 0..4 {
9881                let m4 = q1g_hsum4(d[k][0], d[k][1], d[k][2], d[k][3]);
9882                acc[k] = q1g_group_acc(acc[k], m4, &gsums[k][gi..gi + 4], s);
9883            }
9884            gi += 4;
9885        }
9886        acc
9887    }
9888}
9889
9890/// VNNI twin of `dot_q1g_row_1x4_avx2`.
9891#[cfg(target_arch = "x86_64")]
9892#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
9893unsafe fn dot_q1g_row_1x4_vnni(
9894    bytes: &[u8],
9895    sc_off: usize,
9896    r: usize,
9897    gpr: usize,
9898    xs: [&[i8]; 4],
9899    gsums: [&[i32]; 4],
9900) -> [f32; 4] {
9901    // SAFETY: callers uphold the plane-length and xq/gsum contracts.
9902    unsafe {
9903        use core::arch::x86_64::*;
9904        let expand = _mm256_setr_epi8(
9905            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,
9906            3, 3, 3,
9907        );
9908        let bitsel = _mm256_setr_epi8(
9909            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
9910            -128, 1, 2, 4, 8, 16, 32, 64, -128,
9911        );
9912        let ones8 = _mm256_set1_epi8(1);
9913        let bp = bytes.as_ptr().add(r * gpr * 4);
9914        let mut acc = [0f32; 4];
9915        let mut gi = 0;
9916        while gi < gpr {
9917            let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
9918            let mut d = [[_mm256_setzero_si256(); 4]; 4];
9919            for t in 0..4 {
9920                let bits = (bp.add((gi + t) * 4) as *const u32).read_unaligned();
9921                let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
9922                let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
9923                for (k, xq) in xs.iter().enumerate() {
9924                    let x = _mm256_loadu_si256(
9925                        xq.as_ptr().add((gi + t) * GROUP_SIZE) as *const __m256i
9926                    );
9927                    d[k][t] = _mm256_dpbusd_epi32(
9928                        _mm256_setzero_si256(),
9929                        ones8,
9930                        _mm256_and_si256(x, mask),
9931                    );
9932                }
9933            }
9934            for k in 0..4 {
9935                let m4 = q1g_hsum4(d[k][0], d[k][1], d[k][2], d[k][3]);
9936                acc[k] = q1g_group_acc(acc[k], m4, &gsums[k][gi..gi + 4], s);
9937            }
9938            gi += 4;
9939        }
9940        acc
9941    }
9942}
9943
9944/// SDOT q1g row: the `dot_q1_row_sdot` math on the plane layout. One
9945/// aligned 16-byte load carries a whole 128-group's signs (q1 needs two
9946/// pair loads and a TBL gather of four scales for the same weights), TBL
9947/// replicates each sign byte across 8 lanes for vtst, and the group's ONE
9948/// scale converts once (FCVTL, exact) and broadcasts. Lane order, the
9949/// vpaddq tree, the fma into the four lane accumulators and the tail are
9950/// the q1 kernel's — including its quirk: q1's pair load reads 4 bytes
9951/// past the 4th tile, so the payload's final four tiles (last row) take
9952/// the scalar tail there. q1g has no overread, but takes the same tail on
9953/// the last row to stay bit-identical to q1.
9954#[cfg(target_arch = "aarch64")]
9955#[target_feature(enable = "neon,dotprod")]
9956#[allow(clippy::too_many_arguments)]
9957unsafe fn dot_q1g_row_sdot(
9958    bytes: &[u8],
9959    sc_off: usize,
9960    rows: usize,
9961    r: usize,
9962    gpr: usize,
9963    xq: &[i8],
9964    gsum: &[i32],
9965) -> f32 {
9966    // SAFETY: callers uphold the plane-length and xq/gsum contracts
9967    // (gpr % 4 == 0, xq.len() == gpr·32, gsum.len() == gpr).
9968    unsafe {
9969        use core::arch::aarch64::*;
9970        use core::arch::asm;
9971        const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
9972        // Tile t of the group = sign bytes 4t..4t+4 of the 16-byte load.
9973        const IW: [[u8; 16]; 8] = [
9974            [0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1],
9975            [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3],
9976            [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5],
9977            [6, 6, 6, 6, 6, 6, 6, 6, 7, 7, 7, 7, 7, 7, 7, 7],
9978            [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9],
9979            [
9980                10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
9981            ],
9982            [
9983                12, 12, 12, 12, 12, 12, 12, 12, 13, 13, 13, 13, 13, 13, 13, 13,
9984            ],
9985            [
9986                14, 14, 14, 14, 14, 14, 14, 14, 15, 15, 15, 15, 15, 15, 15, 15,
9987            ],
9988        ];
9989        let m = vld1q_u8(MASKS.as_ptr());
9990        let iw: [uint8x16_t; 8] = core::array::from_fn(|i| vld1q_u8(IW[i].as_ptr()));
9991        macro_rules! tile_dot {
9992            ($w0:expr, $w1:expr, $x:expr) => {{
9993                let x0 = vld1q_s8($x);
9994                let x1 = vld1q_s8($x.add(16));
9995                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
9996                asm!(
9997                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
9998                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
9999                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
10000                    w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
10001                    options(pure, nomem, nostack),
10002                );
10003                vaddq_s32(a0, a1)
10004            }};
10005        }
10006        macro_rules! tbl_mask {
10007            ($ld:expr, $i:expr) => {
10008                vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, iw[$i]), m))
10009            };
10010        }
10011        let bp = bytes.as_ptr().add(r * gpr * 4);
10012        let xp = xq.as_ptr();
10013        let gp = gsum.as_ptr();
10014        let vec_end = if r + 1 == rows { gpr - 4 } else { gpr };
10015        let mut accv = vdupq_n_f32(0.0);
10016        let mut gi = 0;
10017        while gi + 4 <= vec_end {
10018            let ld = vld1q_u8(bp.add(gi * 4));
10019            let d0 = tile_dot!(tbl_mask!(ld, 0), tbl_mask!(ld, 1), xp.add(gi * GROUP_SIZE));
10020            let d1 = tile_dot!(
10021                tbl_mask!(ld, 2),
10022                tbl_mask!(ld, 3),
10023                xp.add((gi + 1) * GROUP_SIZE)
10024            );
10025            let d2 = tile_dot!(
10026                tbl_mask!(ld, 4),
10027                tbl_mask!(ld, 5),
10028                xp.add((gi + 2) * GROUP_SIZE)
10029            );
10030            let d3 = tile_dot!(
10031                tbl_mask!(ld, 6),
10032                tbl_mask!(ld, 7),
10033                xp.add((gi + 3) * GROUP_SIZE)
10034            );
10035            let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
10036            let g = vld1q_s32(gp.add(gi));
10037            let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
10038            let h = vdup_n_u16(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
10039            let scf: float32x4_t;
10040            asm!(
10041                "fcvtl {o:v}.4s, {i:v}.4h",
10042                o = out(vreg) scf, i = in(vreg) h,
10043                options(pure, nomem, nostack),
10044            );
10045            accv = vfmaq_f32(accv, vcvtq_f32_s32(dots), scf);
10046            gi += 4;
10047        }
10048        let mut acc = vaddvq_f32(accv);
10049        while gi < gpr {
10050            let t = bp.add(gi * 4);
10051            let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
10052            let v0 = vcombine_u8(vdup_n_u8(*t), vdup_n_u8(*t.add(1)));
10053            let v1 = vcombine_u8(vdup_n_u8(*t.add(2)), vdup_n_u8(*t.add(3)));
10054            let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
10055            let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
10056            let d = vaddvq_s32(tile_dot!(w0, w1, xp.add(gi * GROUP_SIZE)));
10057            acc += (-(2 * d + *gp.add(gi))) as f32 * s;
10058            gi += 1;
10059        }
10060        acc
10061    }
10062}
10063
10064/// Blocked q1g 1×4: one unpack of the group serves four activation
10065/// streams; per stream bit-identical to `dot_q1g_row_sdot` (and so to the
10066/// q1 1×4 kernel, same last-row tail).
10067#[cfg(target_arch = "aarch64")]
10068#[target_feature(enable = "neon,dotprod")]
10069#[allow(clippy::too_many_arguments)]
10070unsafe fn dot_q1g_row_1x4_sdot(
10071    bytes: &[u8],
10072    sc_off: usize,
10073    rows: usize,
10074    r: usize,
10075    gpr: usize,
10076    xs: [&[i8]; 4],
10077    gs: [&[i32]; 4],
10078) -> [f32; 4] {
10079    // SAFETY: same contracts as `dot_q1g_row_sdot`, ×4.
10080    unsafe {
10081        use core::arch::aarch64::*;
10082        use core::arch::asm;
10083        const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
10084        const IW: [[u8; 16]; 8] = [
10085            [0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1],
10086            [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3],
10087            [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5],
10088            [6, 6, 6, 6, 6, 6, 6, 6, 7, 7, 7, 7, 7, 7, 7, 7],
10089            [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9],
10090            [
10091                10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
10092            ],
10093            [
10094                12, 12, 12, 12, 12, 12, 12, 12, 13, 13, 13, 13, 13, 13, 13, 13,
10095            ],
10096            [
10097                14, 14, 14, 14, 14, 14, 14, 14, 15, 15, 15, 15, 15, 15, 15, 15,
10098            ],
10099        ];
10100        let m = vld1q_u8(MASKS.as_ptr());
10101        let iw: [uint8x16_t; 8] = core::array::from_fn(|i| vld1q_u8(IW[i].as_ptr()));
10102        macro_rules! sdot2 {
10103            ($w0:expr, $w1:expr, $x:expr) => {{
10104                let x0 = vld1q_s8($x);
10105                let x1 = vld1q_s8($x.add(16));
10106                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
10107                asm!(
10108                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
10109                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
10110                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
10111                    w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
10112                    options(pure, nomem, nostack),
10113                );
10114                vaddq_s32(a0, a1)
10115            }};
10116        }
10117        let bp = bytes.as_ptr().add(r * gpr * 4);
10118        let vec_end = if r + 1 == rows { gpr - 4 } else { gpr };
10119        let mut accv = [vdupq_n_f32(0.0); 4];
10120        let mut gi = 0;
10121        while gi + 4 <= vec_end {
10122            let ld = vld1q_u8(bp.add(gi * 4));
10123            // Unpack ONCE — eight ±mask vectors serve all four streams.
10124            let w: [int8x16_t; 8] =
10125                core::array::from_fn(|i| vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld, iw[i]), m)));
10126            let h = vdup_n_u16(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
10127            let scf: float32x4_t;
10128            asm!(
10129                "fcvtl {o:v}.4s, {i:v}.4h",
10130                o = out(vreg) scf, i = in(vreg) h,
10131                options(pure, nomem, nostack),
10132            );
10133            for k in 0..4 {
10134                let xp = xs[k].as_ptr();
10135                let d0 = sdot2!(w[0], w[1], xp.add(gi * GROUP_SIZE));
10136                let d1 = sdot2!(w[2], w[3], xp.add((gi + 1) * GROUP_SIZE));
10137                let d2 = sdot2!(w[4], w[5], xp.add((gi + 2) * GROUP_SIZE));
10138                let d3 = sdot2!(w[6], w[7], xp.add((gi + 3) * GROUP_SIZE));
10139                let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
10140                let g = vld1q_s32(gs[k].as_ptr().add(gi));
10141                let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
10142                accv[k] = vfmaq_f32(accv[k], vcvtq_f32_s32(dots), scf);
10143            }
10144            gi += 4;
10145        }
10146        let mut acc = [
10147            vaddvq_f32(accv[0]),
10148            vaddvq_f32(accv[1]),
10149            vaddvq_f32(accv[2]),
10150            vaddvq_f32(accv[3]),
10151        ];
10152        while gi < gpr {
10153            let t = bp.add(gi * 4);
10154            let sc = f16_to_f32(q1g_scale_bits(bytes, sc_off, r * gpr + gi));
10155            let v0 = vcombine_u8(vdup_n_u8(*t), vdup_n_u8(*t.add(1)));
10156            let v1 = vcombine_u8(vdup_n_u8(*t.add(2)), vdup_n_u8(*t.add(3)));
10157            let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
10158            let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
10159            for k in 0..4 {
10160                let d = vaddvq_s32(sdot2!(w0, w1, xs[k].as_ptr().add(gi * GROUP_SIZE)));
10161                acc[k] += (-(2 * d + *gs[k].as_ptr().add(gi))) as f32 * sc;
10162            }
10163            gi += 1;
10164        }
10165        acc
10166    }
10167}
10168
10169/// (weight ±1, scale) of one q1g element — the exact outlier term.
10170#[inline]
10171fn q1g_outlier(bytes: &[u8], sc_off: usize, r: usize, gpr: usize, j: usize) -> (f32, f32) {
10172    let t = r * gpr + j / GROUP_SIZE;
10173    let k = j % GROUP_SIZE;
10174    let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, t));
10175    let bit = (bytes[t * 4 + k / 8] >> (k % 8)) & 1;
10176    ((bit as i32 * 2 - 1) as f32, s)
10177}
10178
10179/// Exact scalar q1g row (CMF_SDOT=0 contract) — `q1_row_exact`'s order.
10180#[inline]
10181fn q1g_row_exact(bytes: &[u8], sc_off: usize, r: usize, gpr: usize, x: &[f32]) -> f32 {
10182    let mut acc = 0f32;
10183    for gi in 0..gpr {
10184        let t = r * gpr + gi;
10185        let s = f16_to_f32(q1g_scale_bits(bytes, sc_off, t));
10186        let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
10187        let mut ga = 0f32;
10188        for (j, &b) in bytes[t * 4..t * 4 + 4].iter().enumerate() {
10189            for k in 0..8 {
10190                ga += (((b >> k) & 1) as f32 * 2.0 - 1.0) * xg[j * 8 + k];
10191            }
10192        }
10193        acc += ga * s;
10194    }
10195    acc
10196}
10197
10198/// q1t per-row overlay locator. After the base (`base_len`) come
10199/// `[u32 row_ptr[rows+1]]` then `[(u16 col, f16 val)]` grouped by row (row
10200/// `r`'s entries are `[row_ptr[r], row_ptr[r+1])`). Returns
10201/// `(row_ptr offset, entries offset, present)`.
10202fn q1t_overlay(bytes: &[u8], base_len: usize, rows: usize) -> (usize, usize, bool) {
10203    let entries = base_len + (rows + 1) * 4;
10204    (base_len, entries, entries <= bytes.len())
10205}
10206
10207/// Read `row_ptr[r]` from the overlay's prefix-sum table.
10208#[inline]
10209fn q1t_rowptr(bytes: &[u8], rp_off: usize, r: usize) -> usize {
10210    let o = rp_off + r * 4;
10211    u32::from_le_bytes([bytes[o], bytes[o + 1], bytes[o + 2], bytes[o + 3]]) as usize
10212}
10213
10214/// Byte → the 5 ternary signs it packs `{−1,0,+1}` as f32, precomputed so
10215/// decoding a q1t code is a table load, not the base-3 divide/modulo per
10216/// weight (division is ~20–40× the cost of a load). Built at compile time.
10217const SIGN5: [[f32; 5]; 256] = {
10218    let mut lut = [[0.0f32; 5]; 256];
10219    let pow3 = [1u16, 3, 9, 27, 81];
10220    let mut byte = 0usize;
10221    while byte < 256 {
10222        let mut i = 0usize;
10223        while i < 5 {
10224            let code = (byte as u16 / pow3[i]) % 3;
10225            lut[byte][i] = if code == 1 {
10226                1.0
10227            } else if code == 2 {
10228                -1.0
10229            } else {
10230                0.0
10231            };
10232            i += 1;
10233        }
10234        byte += 1;
10235    }
10236    lut
10237};
10238
10239/// Same table, as i8 signs — the operand for the int8 SDOT base kernel.
10240const SIGN5_I8: [[i8; 5]; 256] = {
10241    let mut lut = [[0i8; 5]; 256];
10242    let pow3 = [1u16, 3, 9, 27, 81];
10243    let mut byte = 0usize;
10244    while byte < 256 {
10245        let mut i = 0usize;
10246        while i < 5 {
10247            let code = (byte as u16 / pow3[i]) % 3;
10248            lut[byte][i] = if code == 1 {
10249                1
10250            } else if code == 2 {
10251                -1
10252            } else {
10253                0
10254            };
10255            i += 1;
10256        }
10257        byte += 1;
10258    }
10259    lut
10260};
10261
10262/// The same 5 i8 signs packed into a u64 (`[s0 s1 s2 s3 s4 0 0 0]`, LE) so the
10263/// group unpack is 7 unaligned u64 stores at offsets 0,5,10,…,30 instead of
10264/// six 5-byte copies + LUT indexing — each store's trailing zeros are fixed by
10265/// the next store, and the last one runs 6 B past the 32nd weight (the unpack
10266/// buffer is padded to 40). This is the decode/prefill hot inner op.
10267const SIGN5_U64: [u64; 256] = {
10268    let mut lut = [0u64; 256];
10269    let pow3 = [1u16, 3, 9, 27, 81];
10270    let mut byte = 0usize;
10271    while byte < 256 {
10272        let mut v = 0u64;
10273        let mut i = 0usize;
10274        while i < 5 {
10275            let code = (byte as u16 / pow3[i]) % 3;
10276            let s: u8 = if code == 1 {
10277                1
10278            } else if code == 2 {
10279                0xFF
10280            } else {
10281                0
10282            };
10283            v |= (s as u64) << (i * 8);
10284            i += 1;
10285        }
10286        lut[byte] = v;
10287        byte += 1;
10288    }
10289    lut
10290};
10291
10292/// Ternary base weight at `(row r, col j)` = `sign(code)·s_group`. Used to add
10293/// back activation-outlier columns, whose `x` was zeroed for the int8 bulk dot
10294/// (`split_act`). At a weight-outlier position the code is 0, so this is 0 and
10295/// the overlay correction owns that column — no double counting.
10296#[inline]
10297fn q1t_base_weight(bytes: &[u8], r: usize, gpr: usize, j: usize) -> f32 {
10298    const TILE: usize = cortiq_core::quant::Q1T_TILE;
10299    let off = (r * gpr + j / GROUP_SIZE) * TILE;
10300    let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
10301    let within = j % GROUP_SIZE;
10302    SIGN5[bytes[off + 2 + within / 5] as usize][within % 5] * s
10303}
10304
10305/// One 32-group int8 dot via two SDOTs. Bit-exact vs the scalar i8 sum
10306/// (integer accumulation is order-independent).
10307#[cfg(target_arch = "aarch64")]
10308#[target_feature(enable = "neon,dotprod")]
10309#[inline]
10310unsafe fn sdot32_i8(w: *const i8, x: *const i8) -> i32 {
10311    // SAFETY: caller guarantees 32 readable i8 at each pointer.
10312    unsafe {
10313        use core::arch::aarch64::*;
10314        use core::arch::asm;
10315        let w0 = vld1q_s8(w);
10316        let w1 = vld1q_s8(w.add(16));
10317        let x0 = vld1q_s8(x);
10318        let x1 = vld1q_s8(x.add(16));
10319        let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
10320        asm!(
10321            "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
10322            "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
10323            a0 = inout(vreg) a0, a1 = inout(vreg) a1,
10324            w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
10325            options(pure, nomem, nostack),
10326        );
10327        vaddvq_s32(vaddq_s32(a0, a1))
10328    }
10329}
10330
10331/// One 32-group int8 dot via AVX2: signed·signed as `maddubs(|w|, sign(x,w))`
10332/// then `madd` and a horizontal reduce (the same idiom as `dot_q4t_row_avx2`).
10333#[cfg(target_arch = "x86_64")]
10334#[target_feature(enable = "avx2")]
10335#[inline]
10336unsafe fn i8dot32_avx2(w: *const i8, x: *const i8) -> i32 {
10337    // SAFETY: caller guarantees 32 readable i8 at each pointer.
10338    unsafe {
10339        use core::arch::x86_64::*;
10340        let wv = _mm256_loadu_si256(w as *const __m256i);
10341        let xv = _mm256_loadu_si256(x as *const __m256i);
10342        let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
10343        let d = _mm256_madd_epi16(p16, _mm256_set1_epi16(1));
10344        let hi128 = _mm256_extracti128_si256::<1>(d);
10345        let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
10346        let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
10347        let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
10348        _mm_cvtsi128_si32(s32)
10349    }
10350}
10351
10352/// Unpack one q1t group's base-3 codes into 32 i8 signs via 7 unaligned u64
10353/// stores (see `SIGN5_U64`). `dst` MUST have ≥ 40 bytes: the 7th store writes
10354/// `dst[30..38]`. Stores go in order so each one's trailing zeros are
10355/// overwritten by the next; the final 6 padding bytes are unused by the dot.
10356#[inline]
10357fn q1t_unpack_group_i8(codes: *const u8, dst: &mut [i8]) {
10358    debug_assert!(dst.len() >= 40);
10359    // SAFETY: codes points at 7 readable bytes; dst has ≥ 40 bytes so every
10360    // 8-byte store at offset bi*5 (bi ≤ 6 → ≤ 30) stays in bounds.
10361    unsafe {
10362        let p = dst.as_mut_ptr();
10363        for bi in 0..7 {
10364            core::ptr::write_unaligned(
10365                p.add(bi * 5) as *mut u64,
10366                SIGN5_U64[*codes.add(bi) as usize],
10367            );
10368        }
10369    }
10370}
10371
10372/// One 32-group int8 dot, arch-dispatched (the matmat inner loop, where the
10373/// row's signs are unpacked once and dotted against every batch input).
10374/// Callers are gated by `a8w8_enabled()`, so the target-feature arms are
10375/// reachable; the scalar arm is a non-SIMD-arch fallback.
10376#[inline]
10377fn q1t_i8dot32(w: *const i8, x: *const i8) -> i32 {
10378    #[cfg(target_arch = "aarch64")]
10379    unsafe {
10380        return sdot32_i8(w, x);
10381    }
10382    #[cfg(target_arch = "x86_64")]
10383    unsafe {
10384        return i8dot32_avx2(w, x);
10385    }
10386    #[allow(unreachable_code)]
10387    unsafe {
10388        let mut s = 0i32;
10389        for k in 0..GROUP_SIZE {
10390            s += *w.add(k) as i32 * *x.add(k) as i32;
10391        }
10392        s
10393    }
10394}
10395
10396#[inline]
10397unsafe fn q1t_unpack_reg_u64s(codes: *const u8) -> (u64, u64, u64, u64) {
10398    let (s0, s1, s2, s3, s4, s5, s6) = unsafe {
10399        (
10400            SIGN5_U64[*codes as usize],
10401            SIGN5_U64[*codes.add(1) as usize],
10402            SIGN5_U64[*codes.add(2) as usize],
10403            SIGN5_U64[*codes.add(3) as usize],
10404            SIGN5_U64[*codes.add(4) as usize],
10405            SIGN5_U64[*codes.add(5) as usize],
10406            SIGN5_U64[*codes.add(6) as usize],
10407        )
10408    };
10409
10410    let u0 = s0 | (s1 << 40);
10411    let u1 = (s1 >> 24) | (s2 << 16) | (s3 << 56);
10412    let u2 = (s3 >> 8) | (s4 << 32);
10413    let u3 = (s4 >> 32) | (s5 << 8) | (s6 << 48);
10414
10415    (u0, u1, u2, u3)
10416}
10417
10418/// One q1t row's int8 base dot: `Σ_group s·dot(signs, xq)` (before the shared
10419/// `sx`). Direct register unpacking (zero stack stores/loads, no STLF stalls).
10420/// ARM SDOT.
10421#[cfg(target_arch = "aarch64")]
10422#[target_feature(enable = "neon,dotprod")]
10423unsafe fn q1t_dot_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
10424    use core::arch::aarch64::*;
10425    use core::arch::asm;
10426    unsafe {
10427        const TILE: usize = cortiq_core::quant::Q1T_TILE;
10428        let mut acc = 0f32;
10429        let bytes_ptr = bytes.as_ptr();
10430        let xq_ptr = xq.as_ptr();
10431        let row_off = r * gpr * TILE;
10432
10433        let gpr2 = gpr & !1;
10434        let mut gi = 0;
10435        while gi < gpr2 {
10436            let off0 = row_off + gi * TILE;
10437            let off1 = off0 + TILE;
10438            let s0 = f16_to_f32(u16::from_le_bytes([
10439                *bytes_ptr.add(off0),
10440                *bytes_ptr.add(off0 + 1),
10441            ]));
10442            let s1 = f16_to_f32(u16::from_le_bytes([
10443                *bytes_ptr.add(off1),
10444                *bytes_ptr.add(off1 + 1),
10445            ]));
10446
10447            let (u0_0, u1_0, u2_0, u3_0) = q1t_unpack_reg_u64s(bytes_ptr.add(off0 + 2));
10448            let (u0_1, u1_1, u2_1, u3_1) = q1t_unpack_reg_u64s(bytes_ptr.add(off1 + 2));
10449
10450            let w0_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_0), vcreate_u64(u1_0)));
10451            let w1_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_0), vcreate_u64(u3_0)));
10452            let w0_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_1), vcreate_u64(u1_1)));
10453            let w1_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_1), vcreate_u64(u3_1)));
10454
10455            let x0_0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE));
10456            let x1_0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE + 16));
10457            let x0_1 = vld1q_s8(xq_ptr.add((gi + 1) * GROUP_SIZE));
10458            let x1_1 = vld1q_s8(xq_ptr.add((gi + 1) * GROUP_SIZE + 16));
10459
10460            let (mut a0_0, mut a1_0) = (vdupq_n_s32(0), vdupq_n_s32(0));
10461            let (mut a0_1, mut a1_1) = (vdupq_n_s32(0), vdupq_n_s32(0));
10462            asm!(
10463                "sdot {a0_0:v}.4s, {w0_0:v}.16b, {x0_0:v}.16b",
10464                "sdot {a1_0:v}.4s, {w1_0:v}.16b, {x1_0:v}.16b",
10465                "sdot {a0_1:v}.4s, {w0_1:v}.16b, {x0_1:v}.16b",
10466                "sdot {a1_1:v}.4s, {w1_1:v}.16b, {x1_1:v}.16b",
10467                a0_0 = inout(vreg) a0_0, a1_0 = inout(vreg) a1_0,
10468                a0_1 = inout(vreg) a0_1, a1_1 = inout(vreg) a1_1,
10469                w0_0 = in(vreg) w0_0, x0_0 = in(vreg) x0_0, w1_0 = in(vreg) w1_0, x1_0 = in(vreg) x1_0,
10470                w0_1 = in(vreg) w0_1, x0_1 = in(vreg) x0_1, w1_1 = in(vreg) w1_1, x1_1 = in(vreg) x1_1,
10471                options(pure, nomem, nostack),
10472            );
10473            let d0 = vaddvq_s32(vaddq_s32(a0_0, a1_0));
10474            let d1 = vaddvq_s32(vaddq_s32(a0_1, a1_1));
10475            acc += d0 as f32 * s0 + d1 as f32 * s1;
10476            gi += 2;
10477        }
10478
10479        if gi < gpr {
10480            let off = row_off + gi * TILE;
10481            let s = f16_to_f32(u16::from_le_bytes([
10482                *bytes_ptr.add(off),
10483                *bytes_ptr.add(off + 1),
10484            ]));
10485            let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
10486            let w0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0), vcreate_u64(u1)));
10487            let w1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2), vcreate_u64(u3)));
10488            let x0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE));
10489            let x1 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE + 16));
10490            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
10491            asm!(
10492                "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
10493                "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
10494                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
10495                w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
10496                options(pure, nomem, nostack),
10497            );
10498            let d = vaddvq_s32(vaddq_s32(a0, a1));
10499            acc += d as f32 * s;
10500        }
10501        acc
10502    }
10503}
10504
10505/// x86 AVX2 mirror of `q1t_dot_row_sdot` (maddubs int8 dot per group).
10506#[cfg(target_arch = "x86_64")]
10507#[target_feature(enable = "avx2")]
10508unsafe fn q1t_dot_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
10509    use core::arch::x86_64::*;
10510    unsafe {
10511        const TILE: usize = cortiq_core::quant::Q1T_TILE;
10512        let mut acc = 0f32;
10513        let bytes_ptr = bytes.as_ptr();
10514        let xq_ptr = xq.as_ptr();
10515        let row_off = r * gpr * TILE;
10516
10517        let ones = _mm256_set1_epi16(1);
10518        for gi in 0..gpr {
10519            let off = row_off + gi * TILE;
10520            let s = f16_to_f32(u16::from_le_bytes([
10521                *bytes_ptr.add(off),
10522                *bytes_ptr.add(off + 1),
10523            ]));
10524            let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
10525            let wv = _mm256_set_epi64x(u3 as i64, u2 as i64, u1 as i64, u0 as i64);
10526            let xv = _mm256_loadu_si256(xq_ptr.add(gi * GROUP_SIZE) as *const __m256i);
10527            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
10528            let d256 = _mm256_madd_epi16(p16, ones);
10529            let d128 = _mm_add_epi32(
10530                _mm256_castsi256_si128(d256),
10531                _mm256_extracti128_si256(d256, 1),
10532            );
10533            let d64 = _mm_add_epi32(d128, _mm_shuffle_epi32(d128, 0xee));
10534            let d32 = _mm_cvtsi128_si32(_mm_add_epi32(d64, _mm_shuffle_epi32(d64, 0x55)));
10535            acc += d32 as f32 * s;
10536        }
10537        acc
10538    }
10539}
10540
10541/// VNNI twin of `q1t_dot_row_avx2` (see `dpbusd_hsum`).
10542#[cfg(target_arch = "x86_64")]
10543#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
10544unsafe fn q1t_dot_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
10545    use core::arch::x86_64::*;
10546    // SAFETY: same tile/xq contracts as `q1t_dot_row_avx2`.
10547    unsafe {
10548        const TILE: usize = cortiq_core::quant::Q1T_TILE;
10549        let mut acc = 0f32;
10550        let bytes_ptr = bytes.as_ptr();
10551        let xq_ptr = xq.as_ptr();
10552        let row_off = r * gpr * TILE;
10553        for gi in 0..gpr {
10554            let off = row_off + gi * TILE;
10555            let s = f16_to_f32(u16::from_le_bytes([
10556                *bytes_ptr.add(off),
10557                *bytes_ptr.add(off + 1),
10558            ]));
10559            let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
10560            let wv = _mm256_set_epi64x(u3 as i64, u2 as i64, u1 as i64, u0 as i64);
10561            let xv = _mm256_loadu_si256(xq_ptr.add(gi * GROUP_SIZE) as *const __m256i);
10562            let d = dpbusd_hsum(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
10563            acc += d as f32 * s;
10564        }
10565        acc
10566    }
10567}
10568
10569/// Per-row int8 base dot, dispatched once per row (matvec decode hot path).
10570/// Callers are gated by `a8w8_enabled()`, so the target-feature kernels are
10571/// reachable.
10572#[inline]
10573fn q1t_dot_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
10574    #[cfg(target_arch = "aarch64")]
10575    unsafe {
10576        return q1t_dot_row_sdot(bytes, r, gpr, xq);
10577    }
10578    #[cfg(target_arch = "x86_64")]
10579    unsafe {
10580        if vnni_tiles_enabled() {
10581            return q1t_dot_row_vnni(bytes, r, gpr, xq);
10582        }
10583        return q1t_dot_row_avx2(bytes, r, gpr, xq);
10584    }
10585    #[allow(unreachable_code)]
10586    {
10587        const TILE: usize = cortiq_core::quant::Q1T_TILE;
10588        let mut acc = 0f32;
10589        let mut sg = [0i8; GROUP_SIZE + 8]; // +8 slack for the u64-store unpack
10590        for gi in 0..gpr {
10591            let off = (r * gpr + gi) * TILE;
10592            let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
10593            q1t_unpack_group_i8(bytes.as_ptr().wrapping_add(off + 2), &mut sg);
10594            let mut d = 0i32;
10595            for k in 0..GROUP_SIZE {
10596                d += sg[k] as i32 * xq[gi * GROUP_SIZE + k] as i32;
10597            }
10598            acc += d as f32 * s;
10599        }
10600        acc
10601    }
10602}
10603
10604/// Σ over a row's outliers of `value·x[col]` — the correction that adds the
10605/// overlay's exact weights on top of the base dot. INVARIANT: the encoder
10606/// writes ternary code 0 at every outlier position (`quantize_q1t`), so the
10607/// base contributes nothing there and this is a plain `value·x`, not
10608/// `(value − base)·x` — no scattered per-outlier scale read. Row `r`'s entries
10609/// are the contiguous slice `[row_ptr[r], row_ptr[r+1])`, so no binary search.
10610fn q1t_row_outlier_correction(
10611    bytes: &[u8],
10612    r: usize,
10613    rp_off: usize,
10614    entries_off: usize,
10615    has_ov: bool,
10616    x: &[f32],
10617) -> f32 {
10618    if !has_ov {
10619        return 0.0;
10620    }
10621    let (c0, c1) = (
10622        q1t_rowptr(bytes, rp_off, r),
10623        q1t_rowptr(bytes, rp_off, r + 1),
10624    );
10625    let mut corr = 0f32;
10626    for p in c0..c1 {
10627        let e = entries_off + p * 4;
10628        let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
10629        let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
10630        corr += val * x[col];
10631    }
10632    corr
10633}
10634
10635/// Dequantize one q1t row into `buf[..cols]` via the sign LUT (no division),
10636/// then apply the row's outliers (its `[row_ptr[r], row_ptr[r+1])` slice).
10637/// Used by the batched (prefill) path where the decode amortizes over the batch.
10638fn q1t_dequant_row(
10639    bytes: &[u8],
10640    r: usize,
10641    gpr: usize,
10642    rp_off: usize,
10643    entries_off: usize,
10644    has_ov: bool,
10645    buf: &mut [f32],
10646) {
10647    const TILE: usize = cortiq_core::quant::Q1T_TILE;
10648    for g in 0..gpr {
10649        let off = (r * gpr + g) * TILE;
10650        let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
10651        let codes = &bytes[off + 2..off + TILE];
10652        let bc = g * GROUP_SIZE;
10653        // 6 full bytes (30 codes) + a 7th byte holding the last 2.
10654        for bi in 0..6 {
10655            let lut = &SIGN5[codes[bi] as usize];
10656            let d = &mut buf[bc + bi * 5..bc + bi * 5 + 5];
10657            for i in 0..5 {
10658                d[i] = lut[i] * s;
10659            }
10660        }
10661        let lut = &SIGN5[codes[6] as usize];
10662        buf[bc + 30] = lut[0] * s;
10663        buf[bc + 31] = lut[1] * s;
10664    }
10665    if !has_ov {
10666        return;
10667    }
10668    let (c0, c1) = (
10669        q1t_rowptr(bytes, rp_off, r),
10670        q1t_rowptr(bytes, rp_off, r + 1),
10671    );
10672    for p in c0..c1 {
10673        let e = entries_off + p * 4;
10674        let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
10675        buf[col] = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
10676    }
10677}
10678
10679/// Add the sparse outlier overlay onto a base dot already in `out` (the GPU
10680/// computes the ternary base; the overlay stays on the CPU — its entries are
10681/// few and its per-row gather doesn't vectorize on the GPU). Row-parallel.
10682fn q1t_add_overlay(
10683    bytes: &[u8],
10684    x: &[f32],
10685    rows: usize,
10686    cols: usize,
10687    out: &mut [f32],
10688    pool: Option<&Pool>,
10689) {
10690    const TILE: usize = cortiq_core::quant::Q1T_TILE;
10691    let gpr = cols / GROUP_SIZE;
10692    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
10693    if !has_ov {
10694        return;
10695    }
10696    let out_addr = SendMut(out.as_mut_ptr());
10697    let run = move |start: usize, end: usize| {
10698        for r in start..end {
10699            let corr = q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
10700            // SAFETY: disjoint rows; add onto the base the GPU already wrote.
10701            unsafe { *out_addr.at(r) += corr };
10702        }
10703    };
10704    dispatch_rows(pool, rows, &run);
10705}
10706
10707/// Q1T row range via the A8W8 int8 path — shared activation split,
10708/// per-row: base SDOT dot + outlier correction + overlay.
10709#[allow(clippy::too_many_arguments)]
10710fn q1t_range_a8w8(
10711    bytes: &[u8],
10712    gpr: usize,
10713    rp_off: usize,
10714    ent_off: usize,
10715    has_ov: bool,
10716    act: &SplitAct,
10717    x: &[f32],
10718    out: SendMut,
10719    start: usize,
10720    end: usize,
10721) {
10722    for r in start..end {
10723        let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
10724        for &(j, xv) in &act.outliers {
10725            acc += q1t_base_weight(bytes, r, gpr, j) * xv;
10726        }
10727        acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
10728        // SAFETY: disjoint row ranges per worker.
10729        unsafe { *out.at(r) = acc };
10730    }
10731}
10732
10733/// Q1T row range via the f32 path (no SDOT) — for matvec_many batched
10734/// dispatch when a8w8 is unavailable.
10735#[allow(clippy::too_many_arguments)]
10736fn q1t_range_f32_batch(
10737    bytes: &[u8],
10738    gpr: usize,
10739    rp_off: usize,
10740    ent_off: usize,
10741    has_ov: bool,
10742    x: &[f32],
10743    out: SendMut,
10744    start: usize,
10745    end: usize,
10746) {
10747    const TILE: usize = cortiq_core::quant::Q1T_TILE;
10748    let mut sg = [0f32; GROUP_SIZE];
10749    for r in start..end {
10750        let mut acc = 0f32;
10751        for g in 0..gpr {
10752            let off = (r * gpr + g) * TILE;
10753            let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
10754            let codes = &bytes[off + 2..off + TILE];
10755            let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
10756            for bi in 0..6 {
10757                sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
10758            }
10759            let lut = &SIGN5[codes[6] as usize];
10760            sg[30] = lut[0];
10761            sg[31] = lut[1];
10762            let mut gsum = 0f32;
10763            for k in 0..GROUP_SIZE {
10764                gsum += sg[k] * xg[k];
10765            }
10766            acc += s * gsum;
10767        }
10768        acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
10769        // SAFETY: disjoint row ranges per worker.
10770        unsafe { *out.at(r) = acc };
10771    }
10772}
10773
10774/// Ternary (q1t) matvec — decode+dot straight from mmap, one group at a time:
10775/// no per-ROW buffer, no division (the sign LUT), and a tiny per-group sign
10776/// buffer so the 32-wide dot vectorizes. This is the decode hot path.
10777fn q1t_matvec(
10778    bytes: &[u8],
10779    x: &[f32],
10780    rows: usize,
10781    cols: usize,
10782    out: &mut [f32],
10783    pool: Option<&Pool>,
10784) {
10785    debug_assert_eq!(out.len(), rows);
10786    const TILE: usize = cortiq_core::quant::Q1T_TILE;
10787    let gpr = cols / GROUP_SIZE;
10788    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
10789    let out_addr = SendMut(out.as_mut_ptr());
10790    // int8 SDOT base dot (ARM dotprod): ~4× the f32 arithmetic. x → i8 once
10791    // (`split_act`), activation outliers added back exactly in f32, weight
10792    // overlay on top. ARM SDOT / x86 AVX2; CMF_SDOT=0 keeps the exact f32 path.
10793    if a8w8_enabled() {
10794        let act = split_act(x);
10795        let act = &act;
10796        let run = move |start: usize, end: usize| {
10797            for r in start..end {
10798                let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
10799                for &(j, xv) in &act.outliers {
10800                    acc += q1t_base_weight(bytes, r, gpr, j) * xv;
10801                }
10802                acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
10803                // SAFETY: disjoint row ranges per worker.
10804                unsafe { *out_addr.at(r) = acc };
10805            }
10806        };
10807        dispatch_rows(pool, rows, &run);
10808        return;
10809    }
10810    let run = move |start: usize, end: usize| {
10811        // Per-group signs, unpacked contiguously so the dot below is a clean
10812        // 32-wide reduction the autovectorizer turns into f32x4 FMAs — the
10813        // 5-values-per-byte base-3 layout won't SIMD in place.
10814        let mut sg = [0f32; GROUP_SIZE];
10815        for r in start..end {
10816            let mut acc = 0f32;
10817            for g in 0..gpr {
10818                let off = (r * gpr + g) * TILE;
10819                let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
10820                let codes = &bytes[off + 2..off + TILE];
10821                let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
10822                for bi in 0..6 {
10823                    sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
10824                }
10825                let lut = &SIGN5[codes[6] as usize];
10826                sg[30] = lut[0];
10827                sg[31] = lut[1];
10828                let mut gsum = 0f32;
10829                for k in 0..GROUP_SIZE {
10830                    gsum += sg[k] * xg[k];
10831                }
10832                acc += s * gsum;
10833            }
10834            acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
10835            unsafe { *out_addr.at(r) = acc };
10836        }
10837    };
10838    dispatch_rows(pool, rows, &run);
10839}
10840
10841/// Fused-pair twin of `q1t_dot_row_sdot`: ONE register unpack of the
10842/// ternary codes serves BOTH activation streams (the unpack chain is
10843/// the dominant per-row cost — MTP verify pairs paid it twice). Per
10844/// stream the group order and f32 accumulation match the single-row
10845/// kernel exactly, so pair == 2×matvec bit-for-bit.
10846#[cfg(target_arch = "aarch64")]
10847#[target_feature(enable = "neon,dotprod")]
10848unsafe fn q1t_dot_row_sdot2(bytes: &[u8], r: usize, gpr: usize, xa: &[i8], xb: &[i8]) -> [f32; 2] {
10849    use core::arch::aarch64::*;
10850    use core::arch::asm;
10851    // SAFETY: same slice-length contracts as `q1t_dot_row_sdot`, ×2.
10852    unsafe {
10853        const TILE: usize = cortiq_core::quant::Q1T_TILE;
10854        let bytes_ptr = bytes.as_ptr();
10855        let row_off = r * gpr * TILE;
10856        let xp = [xa.as_ptr(), xb.as_ptr()];
10857        let mut acc = [0f32; 2];
10858        macro_rules! sdot2 {
10859            ($w0:expr, $w1:expr, $x:expr) => {{
10860                let x0 = vld1q_s8($x);
10861                let x1 = vld1q_s8($x.add(16));
10862                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
10863                asm!(
10864                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
10865                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
10866                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
10867                    w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
10868                    options(pure, nomem, nostack),
10869                );
10870                vaddvq_s32(vaddq_s32(a0, a1))
10871            }};
10872        }
10873        let gpr2 = gpr & !1;
10874        let mut gi = 0;
10875        while gi < gpr2 {
10876            let off0 = row_off + gi * TILE;
10877            let off1 = off0 + TILE;
10878            let s0 = f16_to_f32(u16::from_le_bytes([
10879                *bytes_ptr.add(off0),
10880                *bytes_ptr.add(off0 + 1),
10881            ]));
10882            let s1 = f16_to_f32(u16::from_le_bytes([
10883                *bytes_ptr.add(off1),
10884                *bytes_ptr.add(off1 + 1),
10885            ]));
10886            let (u0_0, u1_0, u2_0, u3_0) = q1t_unpack_reg_u64s(bytes_ptr.add(off0 + 2));
10887            let (u0_1, u1_1, u2_1, u3_1) = q1t_unpack_reg_u64s(bytes_ptr.add(off1 + 2));
10888            let w0_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_0), vcreate_u64(u1_0)));
10889            let w1_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_0), vcreate_u64(u3_0)));
10890            let w0_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_1), vcreate_u64(u1_1)));
10891            let w1_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_1), vcreate_u64(u3_1)));
10892            for k in 0..2 {
10893                let d0 = sdot2!(w0_0, w1_0, xp[k].add(gi * GROUP_SIZE));
10894                let d1 = sdot2!(w0_1, w1_1, xp[k].add((gi + 1) * GROUP_SIZE));
10895                acc[k] += d0 as f32 * s0 + d1 as f32 * s1;
10896            }
10897            gi += 2;
10898        }
10899        if gi < gpr {
10900            let off = row_off + gi * TILE;
10901            let s = f16_to_f32(u16::from_le_bytes([
10902                *bytes_ptr.add(off),
10903                *bytes_ptr.add(off + 1),
10904            ]));
10905            let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
10906            let w0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0), vcreate_u64(u1)));
10907            let w1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2), vcreate_u64(u3)));
10908            for k in 0..2 {
10909                let d = sdot2!(w0, w1, xp[k].add(gi * GROUP_SIZE));
10910                acc[k] += d as f32 * s;
10911            }
10912        }
10913        acc
10914    }
10915}
10916
10917/// Fused Q1T pair matvec: ONE pass over the rows serves both
10918/// activation streams — on ARM the ternary register unpack happens
10919/// once per tile pair (`q1t_dot_row_sdot2`); elsewhere the second dot
10920/// rides the row's L1-warm tile bytes. Per stream the math matches
10921/// `q1t_matvec` exactly.
10922fn q1t_matvec2(
10923    bytes: &[u8],
10924    x1: &[f32],
10925    x2: &[f32],
10926    rows: usize,
10927    cols: usize,
10928    o1: &mut [f32],
10929    o2: &mut [f32],
10930    pool: Option<&Pool>,
10931) {
10932    debug_assert_eq!(o1.len(), rows);
10933    debug_assert_eq!(o2.len(), rows);
10934    const TILE: usize = cortiq_core::quant::Q1T_TILE;
10935    let gpr = cols / GROUP_SIZE;
10936    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
10937    let out1 = SendMut(o1.as_mut_ptr());
10938    let out2 = SendMut(o2.as_mut_ptr());
10939    if a8w8_enabled() {
10940        let a1 = split_act(x1);
10941        let a2 = split_act(x2);
10942        let (a1, a2) = (&a1, &a2);
10943        let run = move |start: usize, end: usize| {
10944            for r in start..end {
10945                #[cfg(target_arch = "aarch64")]
10946                // a8w8 on aarch64 ⇔ sdot_enabled(), so the kernel's
10947                // target features are present.
10948                let ds = unsafe { q1t_dot_row_sdot2(bytes, r, gpr, &a1.xq, &a2.xq) };
10949                #[cfg(not(target_arch = "aarch64"))]
10950                let ds = [
10951                    q1t_dot_row_i8(bytes, r, gpr, &a1.xq),
10952                    q1t_dot_row_i8(bytes, r, gpr, &a2.xq),
10953                ];
10954                let mut acc1 = ds[0] * a1.sx;
10955                for &(j, xv) in &a1.outliers {
10956                    acc1 += q1t_base_weight(bytes, r, gpr, j) * xv;
10957                }
10958                acc1 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x1);
10959                let mut acc2 = ds[1] * a2.sx;
10960                for &(j, xv) in &a2.outliers {
10961                    acc2 += q1t_base_weight(bytes, r, gpr, j) * xv;
10962                }
10963                acc2 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x2);
10964                // SAFETY: disjoint row ranges per worker.
10965                unsafe {
10966                    *out1.at(r) = acc1;
10967                    *out2.at(r) = acc2;
10968                }
10969            }
10970        };
10971        dispatch_rows(pool, rows, &run);
10972        return;
10973    }
10974    let run = move |start: usize, end: usize| {
10975        // Exact path (CMF_SDOT=0): unpack the sign LUT once per group,
10976        // dot both streams — same op order per stream as `q1t_matvec`.
10977        let mut sg = [0f32; GROUP_SIZE];
10978        for r in start..end {
10979            let mut acc1 = 0f32;
10980            let mut acc2 = 0f32;
10981            for g in 0..gpr {
10982                let off = (r * gpr + g) * TILE;
10983                let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
10984                let codes = &bytes[off + 2..off + TILE];
10985                for bi in 0..6 {
10986                    sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
10987                }
10988                let lut = &SIGN5[codes[6] as usize];
10989                sg[30] = lut[0];
10990                sg[31] = lut[1];
10991                let xg1 = &x1[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
10992                let xg2 = &x2[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
10993                let mut gsum1 = 0f32;
10994                for k in 0..GROUP_SIZE {
10995                    gsum1 += sg[k] * xg1[k];
10996                }
10997                acc1 += s * gsum1;
10998                let mut gsum2 = 0f32;
10999                for k in 0..GROUP_SIZE {
11000                    gsum2 += sg[k] * xg2[k];
11001                }
11002                acc2 += s * gsum2;
11003            }
11004            acc1 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x1);
11005            acc2 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x2);
11006            // SAFETY: disjoint row ranges per worker.
11007            unsafe {
11008                *out1.at(r) = acc1;
11009                *out2.at(r) = acc2;
11010            }
11011        }
11012    };
11013    dispatch_rows(pool, rows, &run);
11014}
11015
11016/// Ternary (q1t) matmat (prefill) — dequant each row once, dot the whole
11017/// batch against it (amortizes the per-row decode).
11018fn q1t_matmat(
11019    bytes: &[u8],
11020    xs: &[f32],
11021    b: usize,
11022    rows: usize,
11023    cols: usize,
11024    out: &mut [f32],
11025    pool: Option<&Pool>,
11026) {
11027    debug_assert_eq!(out.len(), b * rows);
11028    const TILE: usize = cortiq_core::quant::Q1T_TILE;
11029    let gpr = cols / GROUP_SIZE;
11030    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
11031    let out_addr = SendMut(out.as_mut_ptr());
11032    // int8 prefill (ARM SDOT / x86 AVX2): quantize the B inputs once, unpack
11033    // each weight row's signs to i8 ONCE, then int8-dot against every input —
11034    // the row sign-decode amortizes over the whole batch. CMF_SDOT=0 → f32.
11035    if a8w8_enabled() {
11036        let acts: Vec<SplitAct> = (0..b)
11037            .map(|bi| split_act(&xs[bi * cols..(bi + 1) * cols]))
11038            .collect();
11039        let acts = &acts;
11040        let run = move |start: usize, end: usize| {
11041            let mut sg = vec![0i8; cols + 8]; // row signs, i8 (+8 unpack slack)
11042            let mut sc = vec![0f32; gpr]; // per-group scales
11043            let mut accs = vec![0f32; b]; // per-batch accumulators, reused per row
11044            for r in start..end {
11045                for g in 0..gpr {
11046                    let off = (r * gpr + g) * TILE;
11047                    sc[g] = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
11048                    q1t_unpack_group_i8(
11049                        bytes.as_ptr().wrapping_add(off + 2),
11050                        &mut sg[g * GROUP_SIZE..],
11051                    );
11052                }
11053                for bi in 0..b {
11054                    let act = &acts[bi];
11055                    let mut isum = 0f32;
11056                    for g in 0..gpr {
11057                        let d = q1t_i8dot32(
11058                            sg.as_ptr().wrapping_add(g * GROUP_SIZE),
11059                            act.xq.as_ptr().wrapping_add(g * GROUP_SIZE),
11060                        );
11061                        isum += d as f32 * sc[g];
11062                    }
11063                    let mut acc = isum * act.sx;
11064                    for &(j, xv) in &act.outliers {
11065                        acc += q1t_base_weight(bytes, r, gpr, j) * xv;
11066                    }
11067                    accs[bi] = acc;
11068                }
11069                // Overlay ONCE per row for the whole batch: read each (col, val)
11070                // from mmap a single time (was b× — the re-read dominated prefill)
11071                // and fan it out over the batch via the cached inputs.
11072                if has_ov {
11073                    let (c0, c1) = (
11074                        q1t_rowptr(bytes, rp_off, r),
11075                        q1t_rowptr(bytes, rp_off, r + 1),
11076                    );
11077                    for p in c0..c1 {
11078                        let e = ent_off + p * 4;
11079                        let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
11080                        let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
11081                        for bi in 0..b {
11082                            accs[bi] += val * xs[bi * cols + col];
11083                        }
11084                    }
11085                }
11086                for bi in 0..b {
11087                    unsafe { *out_addr.at(bi * rows + r) = accs[bi] };
11088                }
11089            }
11090        };
11091        dispatch_rows(pool, rows, &run);
11092        return;
11093    }
11094    let run = move |start: usize, end: usize| {
11095        let mut buf = vec![0f32; cols];
11096        for r in start..end {
11097            q1t_dequant_row(bytes, r, gpr, rp_off, ent_off, has_ov, &mut buf);
11098            for bi in 0..b {
11099                let xr = &xs[bi * cols..(bi + 1) * cols];
11100                let mut acc = 0f32;
11101                for j in 0..cols {
11102                    acc += buf[j] * xr[j];
11103                }
11104                unsafe { *out_addr.at(bi * rows + r) = acc };
11105            }
11106        }
11107    };
11108    dispatch_rows(pool, rows, &run);
11109}
11110
11111/// q1-tile entry points of the binary drivers (the tests' handles; the
11112/// runtime goes through `QTensor::bin_w`).
11113#[cfg(test)]
11114fn q1_matvec(
11115    bytes: &[u8],
11116    x: &[f32],
11117    rows: usize,
11118    cols: usize,
11119    out: &mut [f32],
11120    pool: Option<&Pool>,
11121) {
11122    bin_matvec(BinW::Q1(bytes), x, rows, cols, out, pool)
11123}
11124
11125/// Binary-weight matvec (q1 or q1g): A8W8 rows with the shared group
11126/// sums, or the exact scalar walk under CMF_SDOT=0.
11127fn bin_matvec(w: BinW, x: &[f32], rows: usize, cols: usize, out: &mut [f32], pool: Option<&Pool>) {
11128    debug_assert_eq!(out.len(), rows);
11129    let gpr = cols / GROUP_SIZE;
11130    let out_addr = SendMut(out.as_mut_ptr());
11131    if a8w8_enabled() {
11132        let act = split_act(x);
11133        let gsum = q1_group_sums(&act.xq, gpr);
11134        let (act, gsum) = (&act, &gsum);
11135        let run =
11136            move |start: usize, end: usize| bin_range_a8w8(w, gpr, act, gsum, out_addr, start, end);
11137        dispatch_rows(pool, rows, &run);
11138        return;
11139    }
11140    let run = move |start: usize, end: usize| bin_range_f32(w, gpr, x, out_addr, start, end);
11141    dispatch_rows(pool, rows, &run);
11142}
11143
11144/// Fused two-input q1 matvec (weights read once per pair).
11145#[cfg(test)]
11146#[allow(clippy::too_many_arguments)]
11147fn q1_matvec2(
11148    bytes: &[u8],
11149    x1: &[f32],
11150    x2: &[f32],
11151    rows: usize,
11152    cols: usize,
11153    o1: &mut [f32],
11154    o2: &mut [f32],
11155    pool: Option<&Pool>,
11156) {
11157    bin_matvec2(BinW::Q1(bytes), x1, x2, rows, cols, o1, o2, pool)
11158}
11159
11160/// Fused two-input binary-weight matvec (q1 or q1g).
11161#[allow(clippy::too_many_arguments)]
11162fn bin_matvec2(
11163    w: BinW,
11164    x1: &[f32],
11165    x2: &[f32],
11166    rows: usize,
11167    cols: usize,
11168    o1: &mut [f32],
11169    o2: &mut [f32],
11170    pool: Option<&Pool>,
11171) {
11172    let gpr = cols / GROUP_SIZE;
11173    let p1 = SendMut(o1.as_mut_ptr());
11174    let p2 = SendMut(o2.as_mut_ptr());
11175    if a8w8_enabled() {
11176        let a1 = split_act(x1);
11177        let a2 = split_act(x2);
11178        let g1 = q1_group_sums(&a1.xq, gpr);
11179        let g2 = q1_group_sums(&a2.xq, gpr);
11180        let (a1, a2, g1, g2) = (&a1, &a2, &g1, &g2);
11181        let run = move |start: usize, end: usize| {
11182            for r in start..end {
11183                let mut v1 = w.dot_i8(r, gpr, &a1.xq, g1) * a1.sx;
11184                let mut v2 = w.dot_i8(r, gpr, &a2.xq, g2) * a2.sx;
11185                for &(j, xv) in &a1.outliers {
11186                    let (wv, s) = w.outlier(r, gpr, j);
11187                    v1 += wv * s * xv;
11188                }
11189                for &(j, xv) in &a2.outliers {
11190                    let (wv, s) = w.outlier(r, gpr, j);
11191                    v2 += wv * s * xv;
11192                }
11193                // SAFETY: disjoint row ranges per worker.
11194                unsafe {
11195                    *p1.at(r) = v1;
11196                    *p2.at(r) = v2;
11197                }
11198            }
11199        };
11200        dispatch_rows(pool, rows, &run);
11201        return;
11202    }
11203    let run = move |start: usize, end: usize| {
11204        for r in start..end {
11205            // SAFETY: disjoint row ranges per worker.
11206            unsafe {
11207                *p1.at(r) = w.row_exact(r, gpr, x1);
11208                *p2.at(r) = w.row_exact(r, gpr, x2);
11209            }
11210        }
11211    };
11212    dispatch_rows(pool, rows, &run);
11213}
11214
11215/// Batched q1 matmat: each row's tiles stream once per microbatch.
11216#[cfg(test)]
11217#[allow(clippy::too_many_arguments)]
11218fn q1_matmat(
11219    bytes: &[u8],
11220    xs_all: &[f32],
11221    b: usize,
11222    rows: usize,
11223    cols: usize,
11224    out: &mut [f32],
11225    pool: Option<&Pool>,
11226) {
11227    bin_matmat(BinW::Q1(bytes), xs_all, b, rows, cols, out, pool)
11228}
11229
11230/// Batched binary-weight matmat (q1 or q1g).
11231#[allow(clippy::too_many_arguments)]
11232fn bin_matmat(
11233    w: BinW,
11234    xs_all: &[f32],
11235    b: usize,
11236    rows: usize,
11237    cols: usize,
11238    out: &mut [f32],
11239    pool: Option<&Pool>,
11240) {
11241    debug_assert_eq!(out.len(), b * rows);
11242    let gpr = cols / GROUP_SIZE;
11243    let out_addr = SendMut(out.as_mut_ptr());
11244    if a8w8_enabled() {
11245        let acts: Vec<(SplitAct, Vec<i32>)> = (0..b)
11246            .map(|bi| {
11247                let act = split_act(&xs_all[bi * cols..(bi + 1) * cols]);
11248                let gsum = q1_group_sums(&act.xq, gpr);
11249                (act, gsum)
11250            })
11251            .collect();
11252        let acts = &acts;
11253        #[cfg(target_arch = "x86_64")]
11254        let blocked_ok = avx2_enabled() && blocked_enabled();
11255        #[cfg(target_arch = "aarch64")]
11256        let blocked_ok = sdot_enabled() && blocked_enabled();
11257        let run = move |start: usize, end: usize| {
11258            for r in start..end {
11259                let mut bi = 0usize;
11260                // Blocked 1×4: the unpacked bit mask serves four
11261                // activation streams per group.
11262                #[cfg(target_arch = "aarch64")]
11263                if blocked_ok {
11264                    while bi + 4 <= acts.len() {
11265                        let xs = [
11266                            acts[bi].0.xq.as_slice(),
11267                            acts[bi + 1].0.xq.as_slice(),
11268                            acts[bi + 2].0.xq.as_slice(),
11269                            acts[bi + 3].0.xq.as_slice(),
11270                        ];
11271                        let gs = [
11272                            acts[bi].1.as_slice(),
11273                            acts[bi + 1].1.as_slice(),
11274                            acts[bi + 2].1.as_slice(),
11275                            acts[bi + 3].1.as_slice(),
11276                        ];
11277                        // SAFETY: blocked_ok checked SDOT.
11278                        let d = unsafe { w.dot_1x4_sdot(r, gpr, xs, gs) };
11279                        for k in 0..4 {
11280                            let (act, _) = &acts[bi + k];
11281                            let mut acc = d[k] * act.sx;
11282                            for &(j, xv) in &act.outliers {
11283                                let (wv, sc) = w.outlier(r, gpr, j);
11284                                acc += wv * sc * xv;
11285                            }
11286                            // SAFETY: disjoint (bi, r) cells per worker.
11287                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
11288                        }
11289                        bi += 4;
11290                    }
11291                }
11292                #[cfg(target_arch = "x86_64")]
11293                if blocked_ok {
11294                    while bi + 4 <= acts.len() {
11295                        let xs = [
11296                            acts[bi].0.xq.as_slice(),
11297                            acts[bi + 1].0.xq.as_slice(),
11298                            acts[bi + 2].0.xq.as_slice(),
11299                            acts[bi + 3].0.xq.as_slice(),
11300                        ];
11301                        let gs = [
11302                            acts[bi].1.as_slice(),
11303                            acts[bi + 1].1.as_slice(),
11304                            acts[bi + 2].1.as_slice(),
11305                            acts[bi + 3].1.as_slice(),
11306                        ];
11307                        // SAFETY: blocked_ok checked AVX2.
11308                        let d = unsafe { w.dot_1x4_x86(r, gpr, xs, gs) };
11309                        for k in 0..4 {
11310                            let (act, _) = &acts[bi + k];
11311                            let mut acc = d[k] * act.sx;
11312                            for &(j, xv) in &act.outliers {
11313                                let (wv, sc) = w.outlier(r, gpr, j);
11314                                acc += wv * sc * xv;
11315                            }
11316                            // SAFETY: disjoint (bi, r) cells per worker.
11317                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
11318                        }
11319                        bi += 4;
11320                    }
11321                }
11322                while bi < acts.len() {
11323                    let (act, gsum) = &acts[bi];
11324                    let mut acc = w.dot_i8(r, gpr, &act.xq, gsum) * act.sx;
11325                    for &(j, xv) in &act.outliers {
11326                        let (wv, s) = w.outlier(r, gpr, j);
11327                        acc += wv * s * xv;
11328                    }
11329                    // SAFETY: disjoint (bi, r) cells per worker range.
11330                    unsafe { *out_addr.at(bi * rows + r) = acc };
11331                    bi += 1;
11332                }
11333            }
11334        };
11335        dispatch_rows(pool, rows, &run);
11336        return;
11337    }
11338    let run = move |start: usize, end: usize| {
11339        for r in start..end {
11340            for bi in 0..b {
11341                let x = &xs_all[bi * cols..(bi + 1) * cols];
11342                // SAFETY: disjoint (bi, r) cells per worker range.
11343                unsafe { *out_addr.at(bi * rows + r) = w.row_exact(r, gpr, x) };
11344            }
11345        }
11346    };
11347    dispatch_rows(pool, rows, &run);
11348}
11349
11350/// Fused q4_block matvec straight from the mapped bytes. SDOT path when
11351/// dotprod is available (port of vmfcore `dot_q4_block_sdot`, measured
11352/// +23% on q4 decode): nibbles → centered i8, int8×int8 `sdot` per
11353/// 32-group, exact outlier correction — the same A8W8 contract as q8.
11354/// `CMF_SDOT=0` keeps the exact scalar path.
11355fn q4matvec(
11356    bytes: &[u8],
11357    x: &[f32],
11358    rows: usize,
11359    cols: usize,
11360    out: &mut [f32],
11361    pool: Option<&Pool>,
11362) {
11363    debug_assert_eq!(out.len(), rows);
11364    let (packed, scales) = q4_split(bytes, rows, cols);
11365    let gpr = cols / GROUP_SIZE;
11366    let out_addr = SendMut(out.as_mut_ptr());
11367
11368    if a8w8_enabled() {
11369        let act = split_act(x);
11370        let run = move |start: usize, end: usize| {
11371            q4_range_a8w8(packed, scales, gpr, cols, &act, out_addr, start, end)
11372        };
11373        dispatch_rows(pool, rows, &run);
11374        return;
11375    }
11376
11377    let run =
11378        move |start: usize, end: usize| q4_range_f32(packed, scales, gpr, x, out_addr, start, end);
11379    dispatch_rows(pool, rows, &run);
11380}
11381
11382/// One q4 row via the A8W8 int8 path — SDOT on ARM, AVX2 maddubs on
11383/// x86 (scalar fallback is unreachable: callers gate on a8w8_enabled).
11384#[inline]
11385#[allow(unreachable_code)]
11386/// One UNPACKED q4 row (centered i8 in `buf`) against four activation
11387/// streams: the 32-byte weight chunk and its abs() load once per group,
11388/// the per-group f16 scale decodes once — four maddubs+reduce chains
11389/// instead of four full (load, abs, dot) rounds.
11390#[cfg(target_arch = "x86_64")]
11391#[target_feature(enable = "avx2")]
11392unsafe fn dot_q4b_row_1x4_avx2(
11393    buf: &[u8],
11394    scales: &[u8],
11395    g0: usize,
11396    gpr: usize,
11397    xs: [&[i8]; 4],
11398) -> [f32; 4] {
11399    // SAFETY: callers uphold buffer contracts (buf.len() == gpr·32).
11400    unsafe {
11401        use core::arch::x86_64::*;
11402        let ones = _mm256_set1_epi16(1);
11403        let mut acc = [0f32; 4];
11404        for gi in 0..gpr {
11405            let s = f16_to_f32(u16::from_le_bytes([
11406                scales[(g0 + gi) * 2],
11407                scales[(g0 + gi) * 2 + 1],
11408            ]));
11409            let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
11410            let aw = _mm256_abs_epi8(w);
11411            for (k, xq) in xs.iter().enumerate() {
11412                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
11413                let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
11414                let d = _mm256_madd_epi16(p16, ones);
11415                let hi128 = _mm256_extracti128_si256::<1>(d);
11416                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
11417                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
11418                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
11419                acc[k] += _mm_cvtsi128_si32(s32) as f32 * s;
11420            }
11421        }
11422        acc
11423    }
11424}
11425
11426/// VNNI twin of `dot_q4b_row_1x4_avx2` (see `dpbusd_hsum`).
11427#[cfg(target_arch = "x86_64")]
11428#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
11429unsafe fn dot_q4b_row_1x4_vnni(
11430    buf: &[u8],
11431    scales: &[u8],
11432    g0: usize,
11433    gpr: usize,
11434    xs: [&[i8]; 4],
11435) -> [f32; 4] {
11436    // SAFETY: callers uphold buffer contracts (buf.len() == gpr·32).
11437    unsafe {
11438        use core::arch::x86_64::*;
11439        let mut acc = [0f32; 4];
11440        for gi in 0..gpr {
11441            let s = f16_to_f32(u16::from_le_bytes([
11442                scales[(g0 + gi) * 2],
11443                scales[(g0 + gi) * 2 + 1],
11444            ]));
11445            let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
11446            let aw = _mm256_abs_epi8(w);
11447            for (k, xq) in xs.iter().enumerate() {
11448                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
11449                let d = dpbusd_hsum(aw, _mm256_sign_epi8(x, w));
11450                acc[k] += d as f32 * s;
11451            }
11452        }
11453        acc
11454    }
11455}
11456
11457/// The vbit flavor of the blocked 1×4: the per-activation A8W8 scale
11458/// folds in PER GROUP as `(d·sx)·s` — bit-matching the single-matvec
11459/// accumulation order (the q4_block flavor applies sx once at the end,
11460/// matching ITS single path; the two conventions are historical and
11461/// each blocked leg must mirror its own).
11462#[cfg(target_arch = "x86_64")]
11463#[target_feature(enable = "avx2")]
11464unsafe fn dot_q4b_row_1x4_sx_avx2(
11465    buf: &[u8],
11466    scales: &[u8],
11467    g0: usize,
11468    gpr: usize,
11469    xs: [&[i8]; 4],
11470    sxs: [f32; 4],
11471) -> [f32; 4] {
11472    // SAFETY: callers uphold buffer contracts (buf.len() == gpr·32).
11473    unsafe {
11474        use core::arch::x86_64::*;
11475        let ones = _mm256_set1_epi16(1);
11476        let mut acc = [0f32; 4];
11477        for gi in 0..gpr {
11478            let s = f16_to_f32(u16::from_le_bytes([
11479                scales[(g0 + gi) * 2],
11480                scales[(g0 + gi) * 2 + 1],
11481            ]));
11482            let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
11483            let aw = _mm256_abs_epi8(w);
11484            for (k, xq) in xs.iter().enumerate() {
11485                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
11486                let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
11487                let d = _mm256_madd_epi16(p16, ones);
11488                let hi128 = _mm256_extracti128_si256::<1>(d);
11489                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
11490                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
11491                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
11492                acc[k] += (_mm_cvtsi128_si32(s32) as f32 * sxs[k]) * s;
11493            }
11494        }
11495        acc
11496    }
11497}
11498
11499/// VNNI twin of `dot_q4b_row_1x4_sx_avx2` (see `dpbusd_hsum`; the
11500/// per-group `(d·sx)·s` fold mirrors the vbit single path).
11501#[cfg(target_arch = "x86_64")]
11502#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
11503unsafe fn dot_q4b_row_1x4_sx_vnni(
11504    buf: &[u8],
11505    scales: &[u8],
11506    g0: usize,
11507    gpr: usize,
11508    xs: [&[i8]; 4],
11509    sxs: [f32; 4],
11510) -> [f32; 4] {
11511    // SAFETY: callers uphold buffer contracts (buf.len() == gpr·32).
11512    unsafe {
11513        use core::arch::x86_64::*;
11514        let mut acc = [0f32; 4];
11515        for gi in 0..gpr {
11516            let s = f16_to_f32(u16::from_le_bytes([
11517                scales[(g0 + gi) * 2],
11518                scales[(g0 + gi) * 2 + 1],
11519            ]));
11520            let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
11521            let aw = _mm256_abs_epi8(w);
11522            for (k, xq) in xs.iter().enumerate() {
11523                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
11524                let d = dpbusd_hsum(aw, _mm256_sign_epi8(x, w));
11525                acc[k] += (d as f32 * sxs[k]) * s;
11526            }
11527        }
11528        acc
11529    }
11530}
11531
11532#[allow(unreachable_code)]
11533fn dot_q4_row_i8(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
11534    #[cfg(target_arch = "aarch64")]
11535    unsafe {
11536        return dot_q4_row_sdot(packed, scales, g0, gpr, xq);
11537    }
11538    #[cfg(target_arch = "x86_64")]
11539    unsafe {
11540        return dot_q4_row_avx2(packed, scales, g0, gpr, xq);
11541    }
11542    let mut acc = 0f32;
11543    for gi in 0..gpr {
11544        let g = g0 + gi;
11545        let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
11546        let mut d = 0i32;
11547        for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
11548            d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
11549                + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
11550        }
11551        acc += d as f32 * s;
11552    }
11553    acc
11554}
11555
11556/// Two-activation q4 row via the A8W8 int8 path (see `dot_q4_row_i8`).
11557#[inline]
11558#[allow(unreachable_code)]
11559fn dot_q4_row_i8_2(
11560    packed: &[u8],
11561    scales: &[u8],
11562    g0: usize,
11563    gpr: usize,
11564    xq1: &[i8],
11565    xq2: &[i8],
11566) -> (f32, f32) {
11567    #[cfg(target_arch = "aarch64")]
11568    unsafe {
11569        return dot_q4_row_sdot2(packed, scales, g0, gpr, xq1, xq2);
11570    }
11571    #[cfg(target_arch = "x86_64")]
11572    unsafe {
11573        return dot_q4_row_avx2_2(packed, scales, g0, gpr, xq1, xq2);
11574    }
11575    (
11576        dot_q4_row_i8(packed, scales, g0, gpr, xq1),
11577        dot_q4_row_i8(packed, scales, g0, gpr, xq2),
11578    )
11579}
11580
11581/// One q4 row range via SDOT (kernel body of `q4matvec`, extracted so
11582/// multi-matrix jobs can drive it for several tensors in one dispatch).
11583#[allow(clippy::too_many_arguments)]
11584fn q4_range_a8w8(
11585    packed: &[u8],
11586    scales: &[u8],
11587    gpr: usize,
11588    cols: usize,
11589    act: &SplitAct,
11590    out: SendMut,
11591    start: usize,
11592    end: usize,
11593) {
11594    for r in start..end {
11595        let mut acc = dot_q4_row_i8(packed, scales, r * gpr, gpr, &act.xq) * act.sx;
11596        // xq is zeroed at outlier slots — add the exact terms.
11597        for &(j, xv) in &act.outliers {
11598            let flat = r * cols + j;
11599            let byte = packed[flat / 2];
11600            let nib = if flat & 1 == 0 {
11601                byte & 0x0F
11602            } else {
11603                byte >> 4
11604            };
11605            let s = f16_to_f32(u16::from_le_bytes([
11606                scales[(flat / GROUP_SIZE) * 2],
11607                scales[(flat / GROUP_SIZE) * 2 + 1],
11608            ]));
11609            acc += ((nib as i32 - 8) as f32) * s * xv;
11610        }
11611        // SAFETY: disjoint row ranges per worker.
11612        unsafe { *out.at(r) = acc };
11613    }
11614}
11615
11616/// Two-input q4 row range via the A8W8 int8 path — kernel body of
11617/// `q4matvec2`, extracted for pair multi-matrix jobs.
11618#[allow(clippy::too_many_arguments)]
11619fn q4_range2_a8w8(
11620    packed: &[u8],
11621    scales: &[u8],
11622    gpr: usize,
11623    cols: usize,
11624    a1: &SplitAct,
11625    a2: &SplitAct,
11626    p1: SendMut,
11627    p2: SendMut,
11628    start: usize,
11629    end: usize,
11630) {
11631    for r in start..end {
11632        let (s1, s2) = dot_q4_row_i8_2(packed, scales, r * gpr, gpr, &a1.xq, &a2.xq);
11633        let mut acc1 = s1 * a1.sx;
11634        let mut acc2 = s2 * a2.sx;
11635        // xq is zeroed at outlier slots — add the exact terms.
11636        let fix = |outliers: &[(usize, f32)], acc: &mut f32| {
11637            for &(j, xv) in outliers {
11638                let flat = r * cols + j;
11639                let byte = packed[flat / 2];
11640                let nib = if flat & 1 == 0 {
11641                    byte & 0x0F
11642                } else {
11643                    byte >> 4
11644                };
11645                let s = f16_to_f32(u16::from_le_bytes([
11646                    scales[(flat / GROUP_SIZE) * 2],
11647                    scales[(flat / GROUP_SIZE) * 2 + 1],
11648                ]));
11649                *acc += ((nib as i32 - 8) as f32) * s * xv;
11650            }
11651        };
11652        fix(&a1.outliers, &mut acc1);
11653        fix(&a2.outliers, &mut acc2);
11654        // SAFETY: disjoint row ranges per worker.
11655        unsafe {
11656            *p1.at(r) = acc1;
11657            *p2.at(r) = acc2;
11658        }
11659    }
11660}
11661
11662/// Exact scalar q4 row range (same extraction, non-SDOT path).
11663fn q4_range_f32(
11664    packed: &[u8],
11665    scales: &[u8],
11666    gpr: usize,
11667    x: &[f32],
11668    out: SendMut,
11669    start: usize,
11670    end: usize,
11671) {
11672    for r in start..end {
11673        let mut acc = 0f32;
11674        for gi in 0..gpr {
11675            let g = r * gpr + gi;
11676            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
11677            let pk = &packed[g * 16..(g + 1) * 16];
11678            let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
11679            let mut ga = 0f32;
11680            for (k, &b) in pk.iter().enumerate() {
11681                ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
11682                    + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
11683            }
11684            acc += ga * s;
11685        }
11686        // SAFETY: disjoint row ranges per worker.
11687        unsafe { *out.at(r) = acc };
11688    }
11689}
11690
11691/// Fused two-input q4 matvec: nibbles are unpacked ONCE per group and
11692/// dotted against both activations (was: two full matvecs — double
11693/// weight traffic). Per-lane math matches `q4matvec` exactly.
11694#[allow(clippy::too_many_arguments)]
11695fn q4matvec2(
11696    bytes: &[u8],
11697    x1: &[f32],
11698    x2: &[f32],
11699    rows: usize,
11700    cols: usize,
11701    o1: &mut [f32],
11702    o2: &mut [f32],
11703    pool: Option<&Pool>,
11704) {
11705    debug_assert_eq!(o1.len(), rows);
11706    debug_assert_eq!(o2.len(), rows);
11707    let (packed, scales) = q4_split(bytes, rows, cols);
11708    let gpr = cols / GROUP_SIZE;
11709
11710    if a8w8_enabled() {
11711        let a1 = split_act(x1);
11712        let a2 = split_act(x2);
11713        let p1 = SendMut(o1.as_mut_ptr());
11714        let p2 = SendMut(o2.as_mut_ptr());
11715        let run = move |start: usize, end: usize| {
11716            q4_range2_a8w8(packed, scales, gpr, cols, &a1, &a2, p1, p2, start, end)
11717        };
11718        dispatch_rows(pool, rows, &run);
11719        return;
11720    }
11721
11722    let p1 = SendMut(o1.as_mut_ptr());
11723    let p2 = SendMut(o2.as_mut_ptr());
11724    let run = move |start: usize, end: usize| {
11725        q4_range2_f32(packed, scales, gpr, x1, x2, p1, p2, start, end)
11726    };
11727    dispatch_rows(pool, rows, &run);
11728}
11729
11730/// Two-input exact scalar q4 row range (same extraction).
11731#[allow(clippy::too_many_arguments)]
11732fn q4_range2_f32(
11733    packed: &[u8],
11734    scales: &[u8],
11735    gpr: usize,
11736    x1: &[f32],
11737    x2: &[f32],
11738    p1: SendMut,
11739    p2: SendMut,
11740    start: usize,
11741    end: usize,
11742) {
11743    for r in start..end {
11744        let (mut acc1, mut acc2) = (0f32, 0f32);
11745        for gi in 0..gpr {
11746            let g = r * gpr + gi;
11747            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
11748            let pk = &packed[g * 16..(g + 1) * 16];
11749            let x1g = &x1[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
11750            let x2g = &x2[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
11751            let (mut g1, mut g2) = (0f32, 0f32);
11752            for (k, &b) in pk.iter().enumerate() {
11753                let wl = (b & 0x0F) as f32 - 8.0;
11754                let wh = ((b >> 4) & 0x0F) as f32 - 8.0;
11755                g1 += wl * x1g[k * 2] + wh * x1g[k * 2 + 1];
11756                g2 += wl * x2g[k * 2] + wh * x2g[k * 2 + 1];
11757            }
11758            acc1 += g1 * s;
11759            acc2 += g2 * s;
11760        }
11761        // SAFETY: disjoint row ranges per worker.
11762        unsafe {
11763            *p1.at(r) = acc1;
11764            *p2.at(r) = acc2;
11765        }
11766    }
11767}
11768
11769thread_local! {
11770    /// Per-worker decoded-row scratch for the batched q4/vbit kernels
11771    /// (centered i8 for SDOT, f32 for the exact/scalar paths).
11772    static ROW_I8: std::cell::RefCell<Vec<u8>> = const { std::cell::RefCell::new(Vec::new()) };
11773    static ROW_F32: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
11774}
11775
11776/// Batched q4 matmat: each weight row is unpacked from the mmap ONCE
11777/// and dotted against ALL b activations (prefill used to fall back to b
11778/// full matvecs — b× weight traffic and b× nibble decode). Per-position
11779/// math matches `q4matvec` exactly: same group order, same accumulation.
11780/// `out` is row-major [b, rows] like `qmatmat`.
11781#[allow(clippy::too_many_arguments)]
11782fn q4matmat(
11783    bytes: &[u8],
11784    xs_all: &[f32],
11785    b: usize,
11786    rows: usize,
11787    cols: usize,
11788    out: &mut [f32],
11789    pool: Option<&Pool>,
11790) {
11791    debug_assert_eq!(xs_all.len(), b * cols);
11792    debug_assert_eq!(out.len(), b * rows);
11793    let (packed, scales) = q4_split(bytes, rows, cols);
11794    let gpr = cols / GROUP_SIZE;
11795    let gscale = |g: usize| f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
11796
11797    if a8w8_enabled() {
11798        let acts: Vec<SplitAct> = (0..b)
11799            .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
11800            .collect();
11801        let acts = &acts;
11802        let out_addr = SendMut(out.as_mut_ptr());
11803        let run = move |start: usize, end: usize| {
11804            ROW_I8.with(|rb| {
11805                let mut buf = rb.borrow_mut();
11806                buf.resize(cols, 0);
11807                for r in start..end {
11808                    // Unpack the row's nibbles to centered i8 once
11809                    // (element 2k = low nibble, 2k+1 = high — flat order,
11810                    // same as dot_q4_row_sdot's zip).
11811                    for gi in 0..gpr {
11812                        let g = r * gpr + gi;
11813                        for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
11814                            buf[gi * GROUP_SIZE + k * 2] = ((bt & 0x0F) as i32 - 8) as i8 as u8;
11815                            buf[gi * GROUP_SIZE + k * 2 + 1] =
11816                                (((bt >> 4) & 0x0F) as i32 - 8) as i8 as u8;
11817                        }
11818                    }
11819                    let mut bi = 0usize;
11820                    #[cfg(target_arch = "x86_64")]
11821                    if avx2_enabled() && blocked_enabled() {
11822                        while bi + 4 <= acts.len() {
11823                            let xs = [
11824                                acts[bi].xq.as_slice(),
11825                                acts[bi + 1].xq.as_slice(),
11826                                acts[bi + 2].xq.as_slice(),
11827                                acts[bi + 3].xq.as_slice(),
11828                            ];
11829                            let d = unsafe {
11830                                if vnni_tiles_enabled() {
11831                                    dot_q4b_row_1x4_vnni(&buf, scales, r * gpr, gpr, xs)
11832                                } else {
11833                                    dot_q4b_row_1x4_avx2(&buf, scales, r * gpr, gpr, xs)
11834                                }
11835                            };
11836                            for k in 0..4 {
11837                                let act = &acts[bi + k];
11838                                let mut acc = d[k] * act.sx;
11839                                for &(j, xv) in &act.outliers {
11840                                    acc += (buf[j] as i8) as f32
11841                                        * gscale((r * cols + j) / GROUP_SIZE)
11842                                        * xv;
11843                                }
11844                                // SAFETY: disjoint (bi, r) cells per worker.
11845                                unsafe { *out_addr.at((bi + k) * rows + r) = acc };
11846                            }
11847                            bi += 4;
11848                        }
11849                    }
11850                    while bi < acts.len() {
11851                        let act = &acts[bi];
11852                        let mut acc = 0f32;
11853                        for gi in 0..gpr {
11854                            let d = dot_i8_i8(
11855                                &buf[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
11856                                &act.xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
11857                            );
11858                            acc += d as f32 * gscale(r * gpr + gi);
11859                        }
11860                        acc *= act.sx;
11861                        // xq is zeroed at outlier slots — exact terms.
11862                        for &(j, xv) in &act.outliers {
11863                            acc += (buf[j] as i8) as f32 * gscale((r * cols + j) / GROUP_SIZE) * xv;
11864                        }
11865                        // SAFETY: disjoint (bi, r) cells per worker row range.
11866                        unsafe { *out_addr.at(bi * rows + r) = acc };
11867                        bi += 1;
11868                    }
11869                }
11870            })
11871        };
11872        dispatch_rows(pool, rows, &run);
11873        return;
11874    }
11875
11876    let out_addr = SendMut(out.as_mut_ptr());
11877    let run = move |start: usize, end: usize| {
11878        ROW_F32.with(|rb| {
11879            let mut buf = rb.borrow_mut();
11880            buf.resize(cols, 0.0);
11881            for r in start..end {
11882                // Decode raw (nib − 8) values once; scales stay per-group
11883                // so the accumulation order matches q4matvec bit-for-bit.
11884                for gi in 0..gpr {
11885                    let g = r * gpr + gi;
11886                    for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
11887                        buf[gi * GROUP_SIZE + k * 2] = (bt & 0x0F) as f32 - 8.0;
11888                        buf[gi * GROUP_SIZE + k * 2 + 1] = ((bt >> 4) & 0x0F) as f32 - 8.0;
11889                    }
11890                }
11891                for bi in 0..b {
11892                    let x = &xs_all[bi * cols..(bi + 1) * cols];
11893                    let mut acc = 0f32;
11894                    for gi in 0..gpr {
11895                        let mut ga = 0f32;
11896                        // Pairwise (lo + hi) addition, matching
11897                        // q4matvec's `ga += lo·x + hi·x` shape exactly —
11898                        // a flat one-per-element loop rounds differently
11899                        // and broke bit-parity on the scalar (x86) path.
11900                        for k in 0..GROUP_SIZE / 2 {
11901                            let e = gi * GROUP_SIZE + k * 2;
11902                            ga += buf[e] * x[e] + buf[e + 1] * x[e + 1];
11903                        }
11904                        acc += ga * gscale(r * gpr + gi);
11905                    }
11906                    // SAFETY: disjoint (bi, r) cells per worker row range.
11907                    unsafe { *out_addr.at(bi * rows + r) = acc };
11908                }
11909            }
11910        })
11911    };
11912    dispatch_rows(pool, rows, &run);
11913}
11914
11915/// Batched vbit matmat: each variable-bit row is decoded from the mmap
11916/// ONCE for the whole microbatch. Same per-position math as
11917/// `vbitmatvec` (SDOT A8W8 with exact outliers / exact f32 for b=8 rows
11918/// and the scalar path).
11919#[allow(clippy::too_many_arguments)]
11920fn vbitmatmat(
11921    bytes: &[u8],
11922    offsets: &[usize],
11923    xs_all: &[f32],
11924    b: usize,
11925    rows: usize,
11926    cols: usize,
11927    out: &mut [f32],
11928    pool: Option<&Pool>,
11929) {
11930    debug_assert_eq!(xs_all.len(), b * cols);
11931    debug_assert_eq!(out.len(), b * rows);
11932    debug_assert_eq!(offsets.len(), rows + 1);
11933    let ng = cols / GROUP_SIZE;
11934    let bits = &bytes[..rows];
11935    let sc_off = rows;
11936    let gscale = |r: usize, g: usize| {
11937        let so = (r * ng + g) * 2;
11938        f16_to_f32(u16::from_le_bytes([
11939            bytes[sc_off + so],
11940            bytes[sc_off + so + 1],
11941        ]))
11942    };
11943
11944    // Decode row r's raw (u − L) values into `dst` (f32, unscaled).
11945    let decode_f32 = |r: usize, dst: &mut [f32]| {
11946        let bw = bits[r] as usize;
11947        let l = ((1i32 << (bw - 1)) - 1) as f32;
11948        let data = &bytes[offsets[r]..offsets[r + 1]];
11949        let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
11950        for d in dst.iter_mut() {
11951            while nbits < bw {
11952                acc = (acc << 8) | data[idx] as u64;
11953                idx += 1;
11954                nbits += 8;
11955            }
11956            let u = ((acc >> (nbits - bw)) & ((1u64 << bw) - 1)) as f32;
11957            nbits -= bw;
11958            *d = u - l;
11959        }
11960    };
11961
11962    if a8w8_enabled() {
11963        let acts: Vec<SplitAct> = (0..b)
11964            .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
11965            .collect();
11966        let acts = &acts;
11967        let out_addr = SendMut(out.as_mut_ptr());
11968        let run = move |start: usize, end: usize| {
11969            for r in start..end {
11970                let bw = bits[r] as usize;
11971                if bw == 8 {
11972                    // u−L reaches 128 → no i8 path; decode once, exact
11973                    // f32 dots for every position (same as vbitmatvec).
11974                    ROW_F32.with(|rb| {
11975                        let mut buf = rb.borrow_mut();
11976                        buf.resize(cols, 0.0);
11977                        decode_f32(r, &mut buf);
11978                        for bi in 0..b {
11979                            let x = &xs_all[bi * cols..(bi + 1) * cols];
11980                            let mut dot = 0f32;
11981                            for g in 0..ng {
11982                                let mut gd = 0f32;
11983                                for k in 0..GROUP_SIZE {
11984                                    gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
11985                                }
11986                                dot += gd * gscale(r, g);
11987                            }
11988                            // SAFETY: disjoint (bi, r) cells per worker range.
11989                            unsafe { *out_addr.at(bi * rows + r) = dot };
11990                        }
11991                    });
11992                    continue;
11993                }
11994                let l = (1i32 << (bw - 1)) - 1;
11995                let data = &bytes[offsets[r]..offsets[r + 1]];
11996                ROW_I8.with(|rb| {
11997                    let mut buf = rb.borrow_mut();
11998                    buf.resize(cols, 0);
11999                    #[inline(always)]
12000                    fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
12001                        for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
12002                            let u = unpack8::<B>(&data[blk * B..]);
12003                            for k in 0..8 {
12004                                chunk[k] = (u[k] - l) as i8 as u8;
12005                            }
12006                        }
12007                    }
12008                    match bw {
12009                        3 => fill::<3>(data, l, &mut buf),
12010                        4 => vbit_fill4(data, &mut buf),
12011                        5 => fill::<5>(data, l, &mut buf),
12012                        6 => fill::<6>(data, l, &mut buf),
12013                        _ => unreachable!("vbit bit-width {bw} (validated at load)"),
12014                    }
12015                    let mut bi = 0usize;
12016                    // The vbit scale table shares q4_block's layout
12017                    // (contiguous f16 per (row·ng + g)), so the same
12018                    // blocked 1×4 kernel serves the decoded row.
12019                    #[cfg(target_arch = "x86_64")]
12020                    if avx2_enabled() && blocked_enabled() {
12021                        while bi + 4 <= acts.len() {
12022                            let xs = [
12023                                acts[bi].xq.as_slice(),
12024                                acts[bi + 1].xq.as_slice(),
12025                                acts[bi + 2].xq.as_slice(),
12026                                acts[bi + 3].xq.as_slice(),
12027                            ];
12028                            let sxs = [
12029                                acts[bi].sx,
12030                                acts[bi + 1].sx,
12031                                acts[bi + 2].sx,
12032                                acts[bi + 3].sx,
12033                            ];
12034                            let d = unsafe {
12035                                if vnni_tiles_enabled() {
12036                                    dot_q4b_row_1x4_sx_vnni(
12037                                        &buf,
12038                                        &bytes[sc_off..],
12039                                        r * ng,
12040                                        ng,
12041                                        xs,
12042                                        sxs,
12043                                    )
12044                                } else {
12045                                    dot_q4b_row_1x4_sx_avx2(
12046                                        &buf,
12047                                        &bytes[sc_off..],
12048                                        r * ng,
12049                                        ng,
12050                                        xs,
12051                                        sxs,
12052                                    )
12053                                }
12054                            };
12055                            for k in 0..4 {
12056                                let act = &acts[bi + k];
12057                                let mut dot = d[k];
12058                                for &(j, xv) in &act.outliers {
12059                                    dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
12060                                }
12061                                // SAFETY: disjoint (bi, r) cells per worker.
12062                                unsafe { *out_addr.at((bi + k) * rows + r) = dot };
12063                            }
12064                            bi += 4;
12065                        }
12066                    }
12067                    while bi < acts.len() {
12068                        let act = &acts[bi];
12069                        let mut dot = 0f32;
12070                        for g in 0..ng {
12071                            let d = dot_i8_i8(
12072                                &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
12073                                &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
12074                            ) as f32
12075                                * act.sx;
12076                            dot += d * gscale(r, g);
12077                        }
12078                        for &(j, xv) in &act.outliers {
12079                            dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
12080                        }
12081                        // SAFETY: disjoint (bi, r) cells per worker range.
12082                        unsafe { *out_addr.at(bi * rows + r) = dot };
12083                        bi += 1;
12084                    }
12085                });
12086            }
12087        };
12088        dispatch_rows(pool, rows, &run);
12089        return;
12090    }
12091
12092    let out_addr = SendMut(out.as_mut_ptr());
12093    let run = move |start: usize, end: usize| {
12094        ROW_F32.with(|rb| {
12095            let mut buf = rb.borrow_mut();
12096            buf.resize(cols, 0.0);
12097            for r in start..end {
12098                decode_f32(r, &mut buf);
12099                for bi in 0..b {
12100                    let x = &xs_all[bi * cols..(bi + 1) * cols];
12101                    let mut dot = 0f32;
12102                    for g in 0..ng {
12103                        let mut gd = 0f32;
12104                        for k in 0..GROUP_SIZE {
12105                            gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
12106                        }
12107                        dot += gd * gscale(r, g);
12108                    }
12109                    // SAFETY: disjoint (bi, r) cells per worker range.
12110                    unsafe { *out_addr.at(bi * rows + r) = dot };
12111                }
12112            }
12113        })
12114    };
12115    dispatch_rows(pool, rows, &run);
12116}
12117
12118/// Build a GPU batch job for a q8-family mapped tensor (primary
12119/// shard): prescaled input + directory coordinates. None → not
12120/// GPU-eligible, caller stays on the CPU.
12121pub(crate) fn gpu_batch_job<'a>(
12122    t: &'a QTensor,
12123    x: &[f32],
12124) -> Option<(std::sync::Arc<CmfModel>, crate::gpu::BatchJob<'a>)> {
12125    match t {
12126        QTensor::Mapped {
12127            model,
12128            idx,
12129            dtype: dt @ (TensorDtype::Q8Row | TensorDtype::Q8_2f),
12130            rows,
12131            cols,
12132            row_scale,
12133            col_field,
12134            ..
12135        } => Some((
12136            model.clone(),
12137            crate::gpu::BatchJob {
12138                idx: *idx,
12139                rows: *rows,
12140                cols: *cols,
12141                row_scale,
12142                xs: prescale(x, col_field, *dt).into_owned(),
12143                layout: crate::gpu::BatchLayout::Q8,
12144            },
12145        )),
12146        // q1: raw f32 activations, tile-embedded scales.
12147        QTensor::Mapped {
12148            model,
12149            idx,
12150            dtype: TensorDtype::Q1,
12151            rows,
12152            cols,
12153            ..
12154        } => Some((
12155            model.clone(),
12156            crate::gpu::BatchJob {
12157                idx: *idx,
12158                rows: *rows,
12159                cols: *cols,
12160                row_scale: &[],
12161                xs: x.to_vec(),
12162                layout: crate::gpu::BatchLayout::Q1,
12163            },
12164        )),
12165        // q1g: same raw-f32 contract, sign + scale planes.
12166        QTensor::Mapped {
12167            model,
12168            idx,
12169            dtype: TensorDtype::Q1G,
12170            rows,
12171            cols,
12172            ..
12173        } => Some((
12174            model.clone(),
12175            crate::gpu::BatchJob {
12176                idx: *idx,
12177                rows: *rows,
12178                cols: *cols,
12179                row_scale: &[],
12180                xs: x.to_vec(),
12181                layout: crate::gpu::BatchLayout::Q1g,
12182            },
12183        )),
12184        // q4_tiled / q4tp: raw f32 activations; the scales live in the
12185        // payload (inline tiles / row ladder), so row_scale stays empty.
12186        // The GDN projection batch already runs these layouts on Metal —
12187        // this arm lets the attention QKV batch reach the same kernels.
12188        QTensor::Mapped {
12189            model,
12190            idx,
12191            dtype: dt @ (TensorDtype::Q4Tiled | TensorDtype::Q4TiledP),
12192            rows,
12193            cols,
12194            ..
12195        } => Some((
12196            model.clone(),
12197            crate::gpu::BatchJob {
12198                idx: *idx,
12199                rows: *rows,
12200                cols: *cols,
12201                row_scale: &[],
12202                xs: x.to_vec(),
12203                layout: if *dt == TensorDtype::Q4Tiled {
12204                    crate::gpu::BatchLayout::Q4t
12205                } else {
12206                    crate::gpu::BatchLayout::Q4tp
12207                },
12208            },
12209        )),
12210        _ => None,
12211    }
12212}
12213
12214thread_local! {
12215    static PRESCALE_BUF1: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
12216    static PRESCALE_BUF2: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
12217}
12218
12219pub(crate) fn prescale<'a>(
12220    x: &'a [f32],
12221    col_field: &[f32],
12222    dtype: TensorDtype,
12223) -> std::borrow::Cow<'a, [f32]> {
12224    if dtype == TensorDtype::Q8_2f {
12225        x.iter().zip(col_field).map(|(a, c)| a * c).collect()
12226    } else {
12227        std::borrow::Cow::Borrowed(x)
12228    }
12229}
12230
12231/// θ col-field fold for q8_2f activations. Borrowed pass-through for
12232/// every other dtype, using thread-local buffers to eliminate per-matvec allocations.
12233pub(crate) fn prescale_with<R, F: FnOnce(&[f32]) -> R>(
12234    x: &[f32],
12235    col_field: &[f32],
12236    dtype: TensorDtype,
12237    buf_id: u8,
12238    f: F,
12239) -> R {
12240    if dtype == TensorDtype::Q8_2f {
12241        if buf_id == 1 {
12242            PRESCALE_BUF1.with(|b| {
12243                let mut buf = b.borrow_mut();
12244                buf.clear();
12245                buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
12246                f(&buf)
12247            })
12248        } else {
12249            PRESCALE_BUF2.with(|b| {
12250                let mut buf = b.borrow_mut();
12251                buf.clear();
12252                buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
12253                f(&buf)
12254            })
12255        }
12256    } else {
12257        f(x)
12258    }
12259}
12260
12261// ───────────────────── x86-64 AVX2 kernels (roadmap этап 2) ─────────────────────
12262
12263/// AVX2+FMA available? Default ON when the CPU supports both;
12264/// `CMF_AVX2=0` disables (falls back to the autovectorized loops).
12265#[cfg(target_arch = "x86_64")]
12266pub(crate) fn avx2_enabled() -> bool {
12267    use std::sync::OnceLock;
12268    static ON: OnceLock<bool> = OnceLock::new();
12269    *ON.get_or_init(|| {
12270        std::env::var("CMF_AVX2").map(|v| v != "0").unwrap_or(true)
12271            && std::arch::is_x86_feature_detected!("avx2")
12272            && std::arch::is_x86_feature_detected!("fma")
12273    })
12274}
12275
12276/// AVX2 A8W8 allowed? The quantized-activation contract is switched by
12277/// the SAME env as the ARM SDOT path: `CMF_SDOT=0` keeps exact kernels
12278/// (the golden-parity exact gate relies on it) — AVX2 f32 kernels stay
12279/// active either way, they are exact (regrouped sums only).
12280#[cfg(target_arch = "x86_64")]
12281fn avx2_a8w8_enabled() -> bool {
12282    if FLOAT_ACTIVATIONS.get() {
12283        return false;
12284    }
12285    use std::sync::OnceLock;
12286    static ON: OnceLock<bool> = OnceLock::new();
12287    *ON.get_or_init(|| {
12288        avx2_enabled() && std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true)
12289    })
12290}
12291
12292thread_local! {
12293    static FULL_GPU_Q8: std::cell::Cell<bool> = const { std::cell::Cell::new(false) };
12294}
12295
12296/// Match the graph's full-device q8 projection precision on MiMo's host
12297/// tail. Does not enable the GPU or bypass a CPU-only/device-refusal gate.
12298pub(crate) fn enter_full_gpu_q8_scope() -> impl Drop {
12299    struct Restore(bool, std::marker::PhantomData<std::rc::Rc<()>>);
12300    impl Drop for Restore {
12301        fn drop(&mut self) {
12302            FULL_GPU_Q8.set(self.0);
12303        }
12304    }
12305    Restore(FULL_GPU_Q8.replace(true), std::marker::PhantomData)
12306}
12307
12308// Dynamic MiMo experts must not change activation precision when a cache
12309// fill moves them from CPU to GPU. Thread-local: only the cold-expert
12310// dispatch selects float kernels; concurrent pipelines keep their policy.
12311thread_local! {
12312    static FLOAT_ACTIVATIONS: std::cell::Cell<bool> = const { std::cell::Cell::new(false) };
12313}
12314
12315/// Flint's released Q2/Q4 profiles are unusually sensitive to activation
12316/// quantization in the host-side attention path. With the normal A8W8
12317/// default, a short greedy decode can diverge after the first token even
12318/// though the selected-expert Metal path itself is unchanged. An Embryo file
12319/// carries a router fitted on the exact φ, which A8W8 moves by 4e-3, and its
12320/// CPU decode costs the same without it. Keep the reference float-activation
12321/// contract for both unless the operator has explicitly selected an
12322/// `CMF_SDOT` policy.
12323fn exact_cpu_default_enabled() -> bool {
12324    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
12325    *ON.get_or_init(|| std::env::var_os("CMF_SDOT").is_none())
12326}
12327
12328impl QTensor {
12329    #[inline]
12330    fn exact_cpu_default(&self) -> bool {
12331        !FLOAT_ACTIVATIONS.get()
12332            && exact_cpu_default_enabled()
12333            && matches!(
12334                self,
12335                Self::Mapped { model, dtype, .. } if matches!(model.header.arch.arch_name.as_str(), "flint" | "cortiq_embryo")
12336                    // Dory on q8 runs as fast on float activations and exact; on
12337                    // q4tp int8 is 4.7× on an M4 for 0.1 % of perplexity.
12338                    || (model.header.arch.arch_name == "dory" && *dtype == TensorDtype::Q8_2f)
12339            )
12340    }
12341
12342    #[inline]
12343    fn any_exact_cpu_default(ts: &[&QTensor]) -> bool {
12344        ts.iter().any(|t| t.exact_cpu_default())
12345    }
12346
12347    #[inline]
12348    fn any_exact_cpu_default_pairs(pairs: &[(&QTensor, &QTensor)]) -> bool {
12349        pairs
12350            .iter()
12351            .any(|(gate, up)| gate.exact_cpu_default() || up.exact_cpu_default())
12352    }
12353}
12354
12355/// Run `f` with the exact float-activation contract on this thread: the
12356/// 4-bit and 8-bit kernels skip the int8 activation split (`CMF_SDOT=0`
12357/// semantics, scoped). For parity tests and diagnostics — e.g. an `mxfp4`
12358/// model then computes bit for bit what its dequantized f32 twin computes.
12359pub fn with_float_activations<R>(f: impl FnOnce() -> R) -> R {
12360    float_activations_scope(f)
12361}
12362
12363/// Active on the current thread. `Pool` snapshots this with its job
12364/// descriptor and installs it on persistent workers, so a scoped quantized
12365/// matvec cannot accidentally choose a different kernel on a worker.
12366#[inline]
12367pub(crate) fn float_activations_active() -> bool {
12368    FLOAT_ACTIVATIONS.get()
12369}
12370
12371/// Restores the current thread's float-activation state on drop. The `Rc`
12372/// marker deliberately keeps the guard on the thread that created it: this
12373/// policy is thread-local, not a portable capability.
12374pub(crate) struct FloatActivationsScope {
12375    previous: bool,
12376    _not_send: std::marker::PhantomData<std::rc::Rc<()>>,
12377}
12378
12379impl Drop for FloatActivationsScope {
12380    fn drop(&mut self) {
12381        FLOAT_ACTIVATIONS.set(self.previous);
12382    }
12383}
12384
12385#[inline]
12386pub(crate) fn enter_float_activations_scope() -> FloatActivationsScope {
12387    FloatActivationsScope {
12388        previous: FLOAT_ACTIVATIONS.replace(true),
12389        _not_send: std::marker::PhantomData,
12390    }
12391}
12392
12393/// Run `f` with exact float activations (no A8W8 int8 activation
12394/// quantization in the quantized matvecs). The persistent row pool carries
12395/// this scope into a dispatched job and restores each worker afterwards.
12396/// Public for parity tests that compare the default kernels with the strict
12397/// ones in one process; `CMF_SDOT=0` is the process-wide equivalent.
12398#[doc(hidden)]
12399pub fn float_activations_scope<R>(f: impl FnOnce() -> R) -> R {
12400    let _scope = enter_float_activations_scope();
12401    f()
12402}
12403
12404/// Row-exact batching: while set, the x86 batched kernels (`qmatmat`,
12405/// `q4tp_matmat`) compute every (weight row, token) cell with the
12406/// single-token kernel instead of the blocked 2×4 / 1×4 / 1×8 tiles, so a
12407/// token's result does not depend on the batch it rides in and equals its
12408/// matvec. On ARM `q4tp_matmat` keeps its 1×4 tile but in the matvec's
12409/// reduction order, and no q4tp batch takes the AMX or device GEMM. The
12410/// MiMo speculative verify holds it (`row_exact_scope`) — its accepted
12411/// rows must be the rows plain decode would have produced.
12412// Shared with pool workers, so overlapping requests must keep the mode
12413// enabled until the LAST scope leaves. Saving/restoring a global bool is
12414// incorrect when two threads enter and leave in a non-LIFO order.
12415static ROW_EXACT: std::sync::atomic::AtomicUsize = std::sync::atomic::AtomicUsize::new(0);
12416
12417pub(crate) fn row_exact() -> bool {
12418    ROW_EXACT.load(std::sync::atomic::Ordering::Acquire) != 0
12419}
12420
12421fn counted_row_exact_scope<R>(active: &std::sync::atomic::AtomicUsize, f: impl FnOnce() -> R) -> R {
12422    struct Restore<'a>(&'a std::sync::atomic::AtomicUsize);
12423    impl Drop for Restore<'_> {
12424        fn drop(&mut self) {
12425            self.0.fetch_sub(1, std::sync::atomic::Ordering::AcqRel);
12426        }
12427    }
12428    active.fetch_add(1, std::sync::atomic::Ordering::AcqRel);
12429    let _restore = Restore(active);
12430    f()
12431}
12432
12433/// Run `f` with row-exact batching on (also released on unwind).
12434pub(crate) fn row_exact_scope<R>(f: impl FnOnce() -> R) -> R {
12435    counted_row_exact_scope(&ROW_EXACT, f)
12436}
12437
12438/// A8W8 quantized-activation path available on THIS machine? One
12439/// switch across architectures: ARM dotprod (CMF_SDOT) or x86 AVX2
12440/// (CMF_AVX2 + the same CMF_SDOT exact-contract override).
12441#[inline]
12442pub(crate) fn a8w8_enabled() -> bool {
12443    #[cfg(target_arch = "aarch64")]
12444    {
12445        sdot_enabled()
12446    }
12447    #[cfg(target_arch = "x86_64")]
12448    {
12449        avx2_a8w8_enabled()
12450    }
12451    #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
12452    {
12453        false
12454    }
12455}
12456
12457/// int8·int8 dot dispatch: SDOT on ARM; AVX-512 VNNI (vpdpbusd) or AVX2
12458/// maddubs on x86. Callers are gated by `a8w8_enabled()`.
12459#[inline]
12460#[allow(unreachable_code)]
12461fn dot_i8_i8(w: &[u8], xq: &[i8]) -> i32 {
12462    #[cfg(target_arch = "aarch64")]
12463    unsafe {
12464        return dot_i8_sdot(w, xq);
12465    }
12466    #[cfg(target_arch = "x86_64")]
12467    unsafe {
12468        if avx512vnni_enabled() {
12469            return dot_i8_i8_vnni(w, xq);
12470        }
12471        return dot_i8_i8_avx2(w, xq);
12472    }
12473    w.iter()
12474        .zip(xq)
12475        .map(|(&a, &b)| (a as i8) as i32 * b as i32)
12476        .sum()
12477}
12478
12479/// AVX-512 VNNI available? (F+BW+VL+VNNI; `CMF_AVX512=0` falls back to
12480/// AVX2.) VL matters: short 32-byte groups (q4/vbit) ride the 256-bit
12481/// `vpdpbusd` encoding.
12482#[cfg(target_arch = "x86_64")]
12483fn avx512vnni_enabled() -> bool {
12484    use std::sync::OnceLock;
12485    static ON: OnceLock<bool> = OnceLock::new();
12486    *ON.get_or_init(|| {
12487        std::env::var("CMF_AVX512")
12488            .map(|v| v != "0")
12489            .unwrap_or(true)
12490            && std::arch::is_x86_feature_detected!("avx512f")
12491            && std::arch::is_x86_feature_detected!("avx512bw")
12492            && std::arch::is_x86_feature_detected!("avx512vl")
12493            && std::arch::is_x86_feature_detected!("avx512vnni")
12494    })
12495}
12496
12497/// Grouped-codec VNNI arms (the q4t/q4b/q1/q1t tile kernels): default
12498/// ON where AVX-512 VNNI exists (`CMF_VNNI_TILES=0` opt-out). Measured
12499/// on Ryzen 7950X (Zen4, 3 alternating process pairs, blocked GEMM
12500/// 4864×896 b=256): q4t 63→68 GF/s (+8%), q1 53→56 (+6%), q4b 72→75
12501/// (+4%) — consistent, no leg regressed. The tile kernels keep a
12502/// horizontal reduce per 32-weight group, so the `vpdpbusd` saving is
12503/// smaller than the long-dot q8 win (+13%), but it is real and free.
12504#[cfg(target_arch = "x86_64")]
12505fn vnni_tiles_enabled() -> bool {
12506    use std::sync::OnceLock;
12507    static ON: OnceLock<bool> = OnceLock::new();
12508    *ON.get_or_init(|| {
12509        std::env::var("CMF_VNNI_TILES")
12510            .map(|v| v != "0")
12511            .unwrap_or(true)
12512            && avx512vnni_enabled()
12513    })
12514}
12515
12516/// One 256-bit u8×i8 dot → i32 via `vpdpbusd` into a fresh accumulator
12517/// plus the same horizontal reduce the AVX2 kernels use. Products are
12518/// bounded (|w| ≤ 8 or ≤ 1), so maddubs never saturated — the i32 sum
12519/// is bit-identical to the maddubs+madd pair it replaces.
12520#[cfg(target_arch = "x86_64")]
12521#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
12522#[inline]
12523unsafe fn dpbusd_hsum(aw: core::arch::x86_64::__m256i, xs: core::arch::x86_64::__m256i) -> i32 {
12524    // SAFETY: pure register math.
12525    unsafe {
12526        use core::arch::x86_64::*;
12527        let d = _mm256_dpbusd_epi32(_mm256_setzero_si256(), aw, xs);
12528        let hi128 = _mm256_extracti128_si256::<1>(d);
12529        let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
12530        let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
12531        let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
12532        _mm_cvtsi128_si32(s32)
12533    }
12534}
12535
12536/// int8·int8 via AVX-512 VNNI: `vpdpbusd` fuses the maddubs+madd+add
12537/// triple into one u8×i8 dot-accumulate. AVX-512 has no vpsignb, so the
12538/// |w|·sign(x,w) trick becomes |w| × (x negated where w<0) via a mask
12539/// subtract — w==0 lanes contribute 0 through |w|=0 either way.
12540#[cfg(target_arch = "x86_64")]
12541#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
12542unsafe fn dot_i8_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
12543    // SAFETY: callers uphold slice-length contracts (see call sites).
12544    unsafe {
12545        use core::arch::x86_64::*;
12546        let n = w.len();
12547        let mut j = 0usize;
12548        let mut total: i32;
12549        // 4 independent accumulators: vpdpbusd is its own loop-carried
12550        // dependency (~5-cycle latency) — a single-acc loop runs
12551        // latency-bound and LOSES to the AVX2 maddubs kernel, measured
12552        // on Granite Rapids.
12553        {
12554            #[inline(always)]
12555            unsafe fn step(
12556                w: *const u8,
12557                x: *const i8,
12558                acc: core::arch::x86_64::__m512i,
12559            ) -> core::arch::x86_64::__m512i {
12560                unsafe {
12561                    use core::arch::x86_64::*;
12562                    let wv = _mm512_loadu_si512(w as *const _);
12563                    let xv = _mm512_loadu_si512(x as *const _);
12564                    let aw = _mm512_abs_epi8(wv);
12565                    let neg = _mm512_movepi8_mask(wv);
12566                    let sx = _mm512_mask_sub_epi8(xv, neg, _mm512_setzero_si512(), xv);
12567                    _mm512_dpbusd_epi32(acc, aw, sx)
12568                }
12569            }
12570            let (mut a0, mut a1, mut a2, mut a3) = (
12571                _mm512_setzero_si512(),
12572                _mm512_setzero_si512(),
12573                _mm512_setzero_si512(),
12574                _mm512_setzero_si512(),
12575            );
12576            while j + 256 <= n {
12577                a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
12578                a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), a1);
12579                a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), a2);
12580                a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), a3);
12581                j += 256;
12582            }
12583            while j + 64 <= n {
12584                a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
12585                j += 64;
12586            }
12587            let s01 = _mm512_add_epi32(a0, a1);
12588            let s23 = _mm512_add_epi32(a2, a3);
12589            total = _mm512_reduce_add_epi32(_mm512_add_epi32(s01, s23));
12590        }
12591        // 32-wide (q4/vbit groups are exactly 32 bytes).
12592        if j + 32 <= n {
12593            let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
12594            let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
12595            let d = _mm256_dpbusd_epi32(
12596                _mm256_setzero_si256(),
12597                _mm256_abs_epi8(wv),
12598                _mm256_sign_epi8(xv, wv),
12599            );
12600            let hi128 = _mm256_extracti128_si256::<1>(d);
12601            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
12602            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
12603            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
12604            total += _mm_cvtsi128_si32(s32);
12605            j += 32;
12606        }
12607        while j < n {
12608            total += (w[j] as i8) as i32 * xq[j] as i32;
12609            j += 1;
12610        }
12611        total
12612    }
12613}
12614
12615/// i8 row · f32 x via AVX2/FMA (x86 mirror of `dot_i8_f32_neon`).
12616#[cfg(target_arch = "x86_64")]
12617#[target_feature(enable = "avx2,fma")]
12618unsafe fn dot_i8_f32_avx2(w: &[u8], x: &[f32]) -> f32 {
12619    // SAFETY: callers uphold slice-length contracts (see call sites).
12620    unsafe {
12621        use core::arch::x86_64::*;
12622        let n = x.len();
12623        let wp = w.as_ptr();
12624        let xp = x.as_ptr();
12625        let (mut a0, mut a1) = (_mm256_setzero_ps(), _mm256_setzero_ps());
12626        let mut j = 0usize;
12627        while j + 16 <= n {
12628            let wb = _mm_loadu_si128(wp.add(j) as *const __m128i);
12629            let lo = _mm256_cvtepi8_epi32(wb);
12630            let hi = _mm256_cvtepi8_epi32(_mm_srli_si128::<8>(wb));
12631            a0 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(lo), _mm256_loadu_ps(xp.add(j)), a0);
12632            a1 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(hi), _mm256_loadu_ps(xp.add(j + 8)), a1);
12633            j += 16;
12634        }
12635        let acc = _mm256_add_ps(a0, a1);
12636        let hi128 = _mm256_extractf128_ps::<1>(acc);
12637        let s128 = _mm_add_ps(_mm256_castps256_ps128(acc), hi128);
12638        let s64 = _mm_add_ps(s128, _mm_movehl_ps(s128, s128));
12639        let s32 = _mm_add_ss(s64, _mm_shuffle_ps::<1>(s64, s64));
12640        let mut sum = _mm_cvtss_f32(s32);
12641        while j < n {
12642            sum += (*wp.add(j) as i8) as f32 * *xp.add(j);
12643            j += 1;
12644        }
12645        sum
12646    }
12647}
12648
12649/// int8(weight)·int8(activation) → i32 via AVX2 maddubs — the x86
12650/// analogue of the SDOT A8W8 path. `maddubs` takes u8×i8, so the
12651/// standard sign trick applies: |w| × sign(x, w) ≡ w × x per lane.
12652/// Pair saturation is safe: |w|≤128, |x|≤127 → 2·128·127 < 32767.
12653#[cfg(target_arch = "x86_64")]
12654#[target_feature(enable = "avx2")]
12655unsafe fn dot_i8_i8_avx2(w: &[u8], xq: &[i8]) -> i32 {
12656    // SAFETY: callers uphold slice-length contracts (see call sites).
12657    unsafe {
12658        use core::arch::x86_64::*;
12659        let n = w.len();
12660        let ones = _mm256_set1_epi16(1);
12661        let mut acc = _mm256_setzero_si256();
12662        let mut j = 0usize;
12663        while j + 32 <= n {
12664            let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
12665            let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
12666            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
12667            acc = _mm256_add_epi32(acc, _mm256_madd_epi16(p16, ones));
12668            j += 32;
12669        }
12670        let hi128 = _mm256_extracti128_si256::<1>(acc);
12671        let s128 = _mm_add_epi32(_mm256_castsi256_si128(acc), hi128);
12672        let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
12673        let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
12674        let mut s = _mm_cvtsi128_si32(s32);
12675        while j < n {
12676            s += (w[j] as i8) as i32 * xq[j] as i32;
12677            j += 1;
12678        }
12679        s
12680    }
12681}
12682
12683/// smmla 2×4: one instruction covers a 2-row × 2-activation × 8-deep
12684/// tile (32 MACs vs sdot's 16) — the weight pair loads once per 8-k
12685/// slice as a combined 2×8 register and meets two activation pairs.
12686#[cfg(target_arch = "aarch64")]
12687#[target_feature(enable = "neon,i8mm")]
12688unsafe fn dot_i8_smmla_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
12689    // SAFETY: callers uphold slice-length contracts.
12690    unsafe {
12691        use core::arch::aarch64::*;
12692        use core::arch::asm;
12693        let n = w0.len();
12694        let w0p = w0.as_ptr() as *const i8;
12695        let w1p = w1.as_ptr() as *const i8;
12696        // acc01 holds [c(r0,x0) c(r0,x1) c(r1,x0) c(r1,x1)]; acc23 the
12697        // same for x2/x3.
12698        let mut acc01 = vdupq_n_s32(0);
12699        let mut acc23 = vdupq_n_s32(0);
12700        let mut i = 0usize;
12701        while i + 8 <= n {
12702            let wa = vcombine_s8(vld1_s8(w0p.add(i)), vld1_s8(w1p.add(i)));
12703            let xb01 = vcombine_s8(
12704                vld1_s8(xs[0].as_ptr().add(i)),
12705                vld1_s8(xs[1].as_ptr().add(i)),
12706            );
12707            let xb23 = vcombine_s8(
12708                vld1_s8(xs[2].as_ptr().add(i)),
12709                vld1_s8(xs[3].as_ptr().add(i)),
12710            );
12711            asm!(
12712                "smmla {a01:v}.4s, {w:v}.16b, {x01:v}.16b",
12713                "smmla {a23:v}.4s, {w:v}.16b, {x23:v}.16b",
12714                a01 = inout(vreg) acc01, a23 = inout(vreg) acc23,
12715                w = in(vreg) wa, x01 = in(vreg) xb01, x23 = in(vreg) xb23,
12716                options(pure, nomem, nostack),
12717            );
12718            i += 8;
12719        }
12720        let mut out = [[0i32; 4]; 2];
12721        let a01: [i32; 4] = core::mem::transmute(acc01);
12722        let a23: [i32; 4] = core::mem::transmute(acc23);
12723        out[0][0] = a01[0];
12724        out[0][1] = a01[1];
12725        out[1][0] = a01[2];
12726        out[1][1] = a01[3];
12727        out[0][2] = a23[0];
12728        out[0][3] = a23[1];
12729        out[1][2] = a23[2];
12730        out[1][3] = a23[3];
12731        if i < n {
12732            for (k, x) in xs.iter().enumerate() {
12733                for j in i..n {
12734                    out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
12735                    out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
12736                }
12737            }
12738        }
12739        out
12740    }
12741}
12742
12743/// ARM twin of the x86 blocked prefill GEMM: two weight rows stay in
12744/// registers across four activation streams, eight sdot accumulators.
12745/// (The per-row form re-read each W row once per activation.)
12746#[cfg(target_arch = "aarch64")]
12747#[target_feature(enable = "neon,dotprod")]
12748unsafe fn dot_i8_sdot_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
12749    // SAFETY: callers uphold slice-length contracts.
12750    unsafe {
12751        use core::arch::aarch64::*;
12752        use core::arch::asm;
12753        let n = w0.len();
12754        let w0p = w0.as_ptr() as *const i8;
12755        let w1p = w1.as_ptr() as *const i8;
12756        let mut acc = [[vdupq_n_s32(0); 4]; 2];
12757        let mut i = 0usize;
12758        while i + 16 <= n {
12759            let wv0 = vld1q_s8(w0p.add(i));
12760            let wv1 = vld1q_s8(w1p.add(i));
12761            for (k, x) in xs.iter().enumerate() {
12762                let xv = vld1q_s8(x.as_ptr().add(i));
12763                let (mut a0, mut a1) = (acc[0][k], acc[1][k]);
12764                asm!(
12765                    "sdot {a0:v}.4s, {w0:v}.16b, {x:v}.16b",
12766                    "sdot {a1:v}.4s, {w1:v}.16b, {x:v}.16b",
12767                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
12768                    w0 = in(vreg) wv0, w1 = in(vreg) wv1, x = in(vreg) xv,
12769                    options(pure, nomem, nostack),
12770                );
12771                acc[0][k] = a0;
12772                acc[1][k] = a1;
12773            }
12774            i += 16;
12775        }
12776        let mut out = [[0i32; 4]; 2];
12777        for r in 0..2 {
12778            for k in 0..4 {
12779                out[r][k] = vaddvq_s32(acc[r][k]);
12780            }
12781        }
12782        if i < n {
12783            for (k, x) in xs.iter().enumerate() {
12784                for j in i..n {
12785                    out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
12786                    out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
12787                }
12788            }
12789        }
12790        out
12791    }
12792}
12793
12794/// Blocked 2 weight rows × 4 activations for the prefill GEMM
12795/// (roadmap P0: packed panels + multi-row accumulators). The two rows'
12796/// abs() live in registers across all four activation streams; the
12797/// sign-fixup is recomputed per pair (the price of the maddubs trick).
12798/// Returns raw i8·i8 dots; the caller applies scales and outliers.
12799#[cfg(target_arch = "x86_64")]
12800#[target_feature(enable = "avx2")]
12801unsafe fn dot_i8_i8_avx2_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
12802    // SAFETY: callers uphold slice-length contracts.
12803    unsafe {
12804        use core::arch::x86_64::*;
12805        let n = w0.len();
12806        let ones = _mm256_set1_epi16(1);
12807        let mut acc = [[_mm256_setzero_si256(); 4]; 2];
12808        let mut j = 0usize;
12809        while j + 32 <= n {
12810            let wv0 = _mm256_loadu_si256(w0.as_ptr().add(j) as *const __m256i);
12811            let wv1 = _mm256_loadu_si256(w1.as_ptr().add(j) as *const __m256i);
12812            let aw0 = _mm256_abs_epi8(wv0);
12813            let aw1 = _mm256_abs_epi8(wv1);
12814            for (k, x) in xs.iter().enumerate() {
12815                let xv = _mm256_loadu_si256(x.as_ptr().add(j) as *const __m256i);
12816                let p0 = _mm256_maddubs_epi16(aw0, _mm256_sign_epi8(xv, wv0));
12817                acc[0][k] = _mm256_add_epi32(acc[0][k], _mm256_madd_epi16(p0, ones));
12818                let p1 = _mm256_maddubs_epi16(aw1, _mm256_sign_epi8(xv, wv1));
12819                acc[1][k] = _mm256_add_epi32(acc[1][k], _mm256_madd_epi16(p1, ones));
12820            }
12821            j += 32;
12822        }
12823        let mut out = [[0i32; 4]; 2];
12824        for r in 0..2 {
12825            for k in 0..4 {
12826                let a = acc[r][k];
12827                let hi128 = _mm256_extracti128_si256::<1>(a);
12828                let s128 = _mm_add_epi32(_mm256_castsi256_si128(a), hi128);
12829                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
12830                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
12831                out[r][k] = _mm_cvtsi128_si32(s32);
12832            }
12833        }
12834        if j < n {
12835            for (k, x) in xs.iter().enumerate() {
12836                for i in j..n {
12837                    out[0][k] += (w0[i] as i8) as i32 * x[i] as i32;
12838                    out[1][k] += (w1[i] as i8) as i32 * x[i] as i32;
12839                }
12840            }
12841        }
12842        out
12843    }
12844}
12845
12846/// AVX2/VNNI q8 row dot with exact outlier correction (x86 mirror of
12847/// `row_dot_sdot` — same A8W8 contract). With AVX-512 VNNI the row goes
12848/// through the bias trick: Σ(w+128)·x via pure `vpdpbusd` (no per-lane
12849/// sign fixups), corrected by −128·Σx with Σx precomputed per split.
12850#[cfg(target_arch = "x86_64")]
12851#[inline]
12852fn row_dot_avx2(row: &[u8], act: &SplitAct) -> f32 {
12853    let dot = if avx512vnni_enabled() && row.len() >= 64 {
12854        (unsafe { dot_u8p128_i8_vnni(row, &act.xq) }) - 128 * act.xsum
12855    } else {
12856        unsafe { dot_i8_i8_avx2(row, &act.xq) }
12857    };
12858    let mut acc = dot as f32 * act.sx;
12859    for &(j, xv) in &act.outliers {
12860        acc += (row[j] as i8) as f32 * xv;
12861    }
12862    acc
12863}
12864
12865/// Σ (w[i]+128)·x[i] via pure `vpdpbusd` — the caller subtracts
12866/// 128·Σx. Four independent accumulators (dpbusd is ~5-cycle latency;
12867/// a single-acc loop runs latency-bound, measured on Granite Rapids).
12868#[cfg(target_arch = "x86_64")]
12869#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
12870unsafe fn dot_u8p128_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
12871    // SAFETY: callers uphold slice-length contracts (see call sites).
12872    unsafe {
12873        use core::arch::x86_64::*;
12874        let n = w.len();
12875        let flip = _mm512_set1_epi8(-128); // XOR 0x80: i8 w → u8 (w+128)
12876        #[inline(always)]
12877        unsafe fn step(
12878            w: *const u8,
12879            x: *const i8,
12880            flip: core::arch::x86_64::__m512i,
12881            acc: core::arch::x86_64::__m512i,
12882        ) -> core::arch::x86_64::__m512i {
12883            unsafe {
12884                use core::arch::x86_64::*;
12885                let wv = _mm512_xor_si512(_mm512_loadu_si512(w as *const _), flip);
12886                _mm512_dpbusd_epi32(acc, wv, _mm512_loadu_si512(x as *const _))
12887            }
12888        }
12889        let (mut a0, mut a1, mut a2, mut a3) = (
12890            _mm512_setzero_si512(),
12891            _mm512_setzero_si512(),
12892            _mm512_setzero_si512(),
12893            _mm512_setzero_si512(),
12894        );
12895        let mut j = 0usize;
12896        while j + 256 <= n {
12897            a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
12898            a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), flip, a1);
12899            a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), flip, a2);
12900            a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), flip, a3);
12901            j += 256;
12902        }
12903        while j + 64 <= n {
12904            a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
12905            j += 64;
12906        }
12907        let mut total = _mm512_reduce_add_epi32(_mm512_add_epi32(
12908            _mm512_add_epi32(a0, a1),
12909            _mm512_add_epi32(a2, a3),
12910        ));
12911        // Scalar tail: (w as i8) + 128 ≡ (w as u8) ^ 0x80.
12912        while j < n {
12913            total += ((w[j] ^ 0x80) as i32) * xq[j] as i32;
12914            j += 1;
12915        }
12916        total
12917    }
12918}
12919
12920/// One q4 row via AVX2: nibbles → centered i8 (unpacklo/hi restores the
12921/// writer's flat order, same as the NEON vzip pair), maddubs against
12922/// the pre-quantized activation group, × the group's f16 scale. Pair
12923/// saturation safe: |w|≤8, |x|≤127 → 2·8·127 ≪ 32767. Mirror of
12924/// `dot_q4_row_sdot`.
12925#[cfg(target_arch = "x86_64")]
12926#[target_feature(enable = "avx2")]
12927unsafe fn dot_q4_row_avx2(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
12928    // SAFETY: callers uphold slice-length contracts (16 packed bytes and
12929    // 2 scale bytes per group; xq.len() == gpr·GROUP_SIZE).
12930    unsafe {
12931        use core::arch::x86_64::*;
12932        let lomask = _mm_set1_epi8(0x0F);
12933        let eight = _mm256_set1_epi8(8);
12934        let ones = _mm256_set1_epi16(1);
12935        let mut acc = 0f32;
12936        for gi in 0..gpr {
12937            let g = g0 + gi;
12938            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
12939            let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
12940            let lo = _mm_and_si128(b, lomask);
12941            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
12942            let w = _mm256_sub_epi8(
12943                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
12944                eight,
12945            );
12946            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
12947            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
12948            let d = _mm256_madd_epi16(p16, ones);
12949            let hi128 = _mm256_extracti128_si256::<1>(d);
12950            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
12951            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
12952            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
12953            acc += _mm_cvtsi128_si32(s32) as f32 * s;
12954        }
12955        acc
12956    }
12957}
12958
12959/// Two-activation q4 row via AVX2: nibbles unpacked ONCE per group,
12960/// both activations dotted against the same centered i8 register.
12961#[cfg(target_arch = "x86_64")]
12962#[target_feature(enable = "avx2")]
12963unsafe fn dot_q4_row_avx2_2(
12964    packed: &[u8],
12965    scales: &[u8],
12966    g0: usize,
12967    gpr: usize,
12968    xq1: &[i8],
12969    xq2: &[i8],
12970) -> (f32, f32) {
12971    // SAFETY: callers uphold slice-length contracts (see dot_q4_row_avx2).
12972    unsafe {
12973        use core::arch::x86_64::*;
12974        let lomask = _mm_set1_epi8(0x0F);
12975        let eight = _mm256_set1_epi8(8);
12976        let ones = _mm256_set1_epi16(1);
12977        let (mut acc1, mut acc2) = (0f32, 0f32);
12978        #[inline(always)]
12979        unsafe fn hsum(d: core::arch::x86_64::__m256i) -> i32 {
12980            unsafe {
12981                use core::arch::x86_64::*;
12982                let hi128 = _mm256_extracti128_si256::<1>(d);
12983                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
12984                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
12985                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
12986                _mm_cvtsi128_si32(s32)
12987            }
12988        }
12989        for gi in 0..gpr {
12990            let g = g0 + gi;
12991            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
12992            let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
12993            let lo = _mm_and_si128(b, lomask);
12994            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
12995            let w = _mm256_sub_epi8(
12996                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
12997                eight,
12998            );
12999            let aw = _mm256_abs_epi8(w);
13000            let x1 = _mm256_loadu_si256(xq1.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
13001            let x2 = _mm256_loadu_si256(xq2.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
13002            let d1 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x1, w)), ones);
13003            let d2 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x2, w)), ones);
13004            acc1 += hsum(d1) as f32 * s;
13005            acc2 += hsum(d2) as f32 * s;
13006        }
13007        (acc1, acc2)
13008    }
13009}
13010
13011/// One q8 row range via AVX2 (x86 mirror of `q8_range_sdot`).
13012#[cfg(target_arch = "x86_64")]
13013fn q8_range_avx2(
13014    q: &[u8],
13015    row_scale: &[f32],
13016    act: &SplitAct,
13017    cols: usize,
13018    out_addr: SendMut,
13019    start: usize,
13020    end: usize,
13021) {
13022    for o in start..end {
13023        let v = row_dot_avx2(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
13024        // SAFETY: disjoint row ranges per worker.
13025        unsafe { *out_addr.at(o) = v };
13026    }
13027}
13028
13029/// Two-input q8 row range via AVX2 (x86 mirror of `q8_range2_sdot`).
13030#[cfg(target_arch = "x86_64")]
13031#[allow(clippy::too_many_arguments)]
13032fn q8_range2_avx2(
13033    q: &[u8],
13034    row_scale: &[f32],
13035    a1: &SplitAct,
13036    a2: &SplitAct,
13037    cols: usize,
13038    p1: SendMut,
13039    p2: SendMut,
13040    start: usize,
13041    end: usize,
13042) {
13043    for o in start..end {
13044        let row = &q[o * cols..(o + 1) * cols];
13045        // SAFETY: disjoint row ranges per worker.
13046        unsafe {
13047            *p1.at(o) = row_dot_avx2(row, a1) * row_scale[o];
13048            *p2.at(o) = row_dot_avx2(row, a2) * row_scale[o];
13049        }
13050    }
13051}
13052
13053// ───────────────────── A8W8 SDOT path (port of vmfcore, ×1.78 decode) ─────────────────────
13054
13055/// ARMv8.6 i8mm (smmla): 32 int8 MACs per instruction vs sdot's 16 —
13056/// yet MEASURED 2.4× SLOWER than the blocked sdot on Apple silicon
13057/// (108 vs 264 GF/s): the on-the-fly vcombine packing and the two-
13058/// accumulator dependency chain swamp the MAC advantage, and Apple's
13059/// four SIMD pipes already keep sdot fed. OPT-IN (CMF_I8MM=1) for
13060/// field trials on Cortex-A710/X-class parts with two pipes, where the
13061/// balance may differ; a pre-interleaved weight layout (repack infra)
13062/// is the known path if it ever earns its keep.
13063#[cfg(target_arch = "aarch64")]
13064fn i8mm_enabled() -> bool {
13065    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
13066    *ON.get_or_init(|| {
13067        std::env::var("CMF_I8MM").map(|v| v == "1").unwrap_or(false)
13068            && std::arch::is_aarch64_feature_detected!("i8mm")
13069    })
13070}
13071
13072/// SDOT enabled? Default ON when the CPU has ARMv8.2 dotprod;
13073/// `CMF_SDOT=0` disables (falls back to i8×f32 NEON).
13074/// (On non-ARM release builds only the test tolerance switch calls it.)
13075#[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
13076fn sdot_enabled() -> bool {
13077    if FLOAT_ACTIVATIONS.get() {
13078        return false;
13079    }
13080    use std::sync::OnceLock;
13081    static ON: OnceLock<bool> = OnceLock::new();
13082    *ON.get_or_init(|| {
13083        let want = std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true);
13084        if !want {
13085            return false;
13086        }
13087
13088        #[cfg(target_arch = "aarch64")]
13089        {
13090            if std::arch::is_aarch64_feature_detected!("dotprod") {
13091                return true;
13092            }
13093            #[cfg(target_os = "android")]
13094            {
13095                if let Ok(cpuinfo) = std::fs::read_to_string("/proc/cpuinfo") {
13096                    if cpuinfo.lines().any(|l| {
13097                        (l.starts_with("Features") || l.starts_with("features"))
13098                            && l.contains("asimddp")
13099                    }) {
13100                        return true;
13101                    }
13102                }
13103            }
13104            false
13105        }
13106        #[cfg(not(target_arch = "aarch64"))]
13107        {
13108            false
13109        }
13110    })
13111}
13112
13113/// Two-field activation split (≡ vmfcore `q8_split_prep`): outlier
13114/// channels (>8·rms) are computed exactly in f32; the bulk (outliers
13115/// zeroed → clean absmax) goes through int8 SDOT. Computed ONCE per
13116/// matvec, shared by all rows/workers.
13117struct SplitAct {
13118    xq: Vec<i8>,
13119    sx: f32,
13120    outliers: Vec<(usize, f32)>,
13121    /// Σ xq — the VNNI bias-trick correction (`(w+128)·x` sums need
13122    /// `−128·Σx`); one i32 per split, computed once per matvec.
13123    #[cfg_attr(not(target_arch = "x86_64"), allow(dead_code))]
13124    xsum: i32,
13125    /// `Σ xq` per 32-element group, already f32 — the q4tp `(w+8)·x − 8·Σx`
13126    /// correction for the VERTICAL accumulator: the f32x8 partials reduce
13127    /// through an 8-lane horizontal sum, and each lane subtracts this once,
13128    /// so the row total carries exactly `8·Σx`. Precomputing it lets the
13129    /// AVX2 kernels drop the per-byte `vpaddb`/`vpabsb`/`vpsignb` (two of
13130    /// the three sat on port 5) and feed the raw unsigned nibbles straight
13131    /// to `vpmaddubsw`. Every value is an integer ≤ 32·127, so both the
13132    /// f32 conversion and each lane's subtraction are exact and the cell
13133    /// equals the old `(nib − 8)` dot bit for bit.
13134    #[cfg_attr(not(target_arch = "x86_64"), allow(dead_code))]
13135    xg8: Vec<f32>,
13136}
13137
13138thread_local! {
13139    /// Recycled xq buffers: split_act runs for every matvec (~200/token)
13140    /// and its hidden-size allocation was steady-state heap churn.
13141    static XQ_FREE: std::cell::RefCell<Vec<Vec<i8>>> =
13142        const { std::cell::RefCell::new(Vec::new()) };
13143    /// Recycled per-group correction buffers (see `SplitAct::xg8`).
13144    static XG8_FREE: std::cell::RefCell<Vec<Vec<f32>>> =
13145        const { std::cell::RefCell::new(Vec::new()) };
13146}
13147
13148/// Pop a recycled `xg8` buffer, or mint one; filled by [`xg8_of`].
13149fn xg8_take() -> Vec<f32> {
13150    let mut v = XG8_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
13151    v.clear();
13152    v
13153}
13154
13155impl Drop for SplitAct {
13156    fn drop(&mut self) {
13157        let buf = std::mem::take(&mut self.xq);
13158        if buf.capacity() > 0 {
13159            XQ_FREE.with(|f| {
13160                let mut f = f.borrow_mut();
13161                if f.len() < 16 {
13162                    f.push(buf);
13163                }
13164            });
13165        }
13166        let buf = std::mem::take(&mut self.xg8);
13167        if buf.capacity() > 0 {
13168            XG8_FREE.with(|f| {
13169                let mut f = f.borrow_mut();
13170                if f.len() < 16 {
13171                    f.push(buf);
13172                }
13173            });
13174        }
13175    }
13176}
13177
13178/// Fill `out` with `Σ xq` per group — see [`SplitAct::xg8`]. A tail
13179/// shorter than one group is dropped (`chunks_exact`); kernels guard on
13180/// `xg8.len() == gpr`.
13181fn xg8_fill(xq: &[i8], out: &mut Vec<f32>) {
13182    out.clear();
13183    out.extend(xq.chunks_exact(GROUP_SIZE).map(|c| {
13184        let mut s = 0i32;
13185        for &v in c {
13186            s += v as i32;
13187        }
13188        s as f32
13189    }));
13190}
13191
13192thread_local! {
13193    /// One scratch row per WORKER, kept for the life of the thread.
13194    ///
13195    /// The kernels take a row of group scales per dispatch, and a fresh
13196    /// `vec![0f32; gpr]` inside the closure is one allocation per worker per
13197    /// dispatch — on the release checkpoint about six thousand a token, a
13198    /// quarter of everything the benchmark counts.
13199    static KROW: UnsafeCell<Vec<f32>> = const { UnsafeCell::new(Vec::new()) };
13200    /// Two scratch rows for a gate/up pair.  Q2TP and Q4TP fused SwiGLU
13201    /// decode each projection's ladder independently, but both ladders have
13202    /// the same lifetime; retaining them per worker removes two heap trips
13203    /// from every dispatched expert pair.
13204    static KROWS: UnsafeCell<[Vec<f32>; 2]> =
13205        const { UnsafeCell::new([Vec::new(), Vec::new()]) };
13206}
13207
13208/// Borrow `n` floats of the calling worker's scratch. Nothing inside a
13209/// kernel body borrows it again, which is what keeps the RefCell honest.
13210#[inline]
13211fn with_krow<R>(n: usize, f: impl FnOnce(&mut [f32]) -> R) -> R {
13212    KROW.with(|s| {
13213        // SAFETY: `KROW` is thread-local and this function does not recurse;
13214        // each worker has exclusive access to its own scratch row.
13215        let b = unsafe { &mut *s.get() };
13216        if b.len() < n {
13217            b.resize(n, 0.0);
13218        }
13219        f(&mut b[..n])
13220    })
13221}
13222
13223/// `t.round().clamp(-127.0, 127.0) as i8`, bit for bit, without the libm
13224/// call. On baseline x86-64 (no SSE4.1 `roundps`) `f32::round` is a
13225/// function call per element, and split_act runs it over every hidden
13226/// state before every matvec: measured 27 us a call on a 2048-wide
13227/// activation on an EPYC 7763 — 5.4 ms of a 55 ms decode token, all of
13228/// it on the caller's thread while thirty workers wait. Clamping first is
13229/// equivalent (round is monotonic and ±127 are integers), and after the
13230/// clamp `t - trunc(t)` is exact, so the half-away-from-zero decision is
13231/// the one `round` makes. NaN clamps to NaN and converts to 0, as before.
13232/// The loop vectorizes (cvttps2dq + compare/select).
13233#[inline(always)]
13234fn q8_round(t: f32) -> i8 {
13235    let t = t.clamp(-127.0, 127.0);
13236    let i = t as i32;
13237    let f = t - i as f32;
13238    let r = if f >= 0.5 {
13239        i + 1
13240    } else if f <= -0.5 {
13241        i - 1
13242    } else {
13243        i
13244    };
13245    r as i8
13246}
13247
13248#[inline]
13249fn with_krows<R>(n: usize, f: impl FnOnce(&mut [f32], &mut [f32]) -> R) -> R {
13250    KROWS.with(|s| {
13251        // SAFETY: KROWS is thread-local and no kernel body recursively calls
13252        // this helper; each participant owns its two vectors exclusively.
13253        let b = unsafe { &mut *s.get() };
13254        for row in b.iter_mut() {
13255            if row.len() < n {
13256                row.resize(n, 0.0);
13257            }
13258        }
13259        let (a, b) = b.split_at_mut(1);
13260        f(&mut a[0][..n], &mut b[0][..n])
13261    })
13262}
13263
13264#[inline]
13265fn silu_mul_limited(mut gate: f32, mut up: f32, limit: f32) -> f32 {
13266    if limit > 0.0 {
13267        up = up.clamp(-limit, limit);
13268        gate = gate.min(limit);
13269    }
13270    gate / (1.0 + (-gate).exp()) * up
13271}
13272
13273/// `split_act` of every column of a batch, the columns spread over the
13274/// pool from 16 on. Each column's split is its own (the same bits as the
13275/// serial loop); serially it was 25 µs a column, 6 ms of a Spark-4B
13276/// prefill token on an 8-thread EPYC while the workers idled.
13277fn split_acts_pooled(xs_all: &[f32], b: usize, cols: usize, pool: Option<&Pool>) -> Vec<SplitAct> {
13278    match pool {
13279        Some(p) if b >= 16 => {
13280            let slots: Vec<std::sync::Mutex<Option<SplitAct>>> =
13281                (0..b).map(|_| std::sync::Mutex::new(None)).collect();
13282            p.run_rows(b, &|start, end| {
13283                for bi in start..end {
13284                    let a = split_act(&xs_all[bi * cols..(bi + 1) * cols]);
13285                    *slots[bi].lock().unwrap_or_else(|e| e.into_inner()) = Some(a);
13286                }
13287            });
13288            slots
13289                .into_iter()
13290                .map(|m| {
13291                    m.into_inner()
13292                        .unwrap_or_else(|e| e.into_inner())
13293                        .expect("every column split")
13294                })
13295                .collect()
13296        }
13297        _ => (0..b)
13298            .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
13299            .collect(),
13300    }
13301}
13302
13303fn split_act(x: &[f32]) -> SplitAct {
13304    let _prof = crate::cpuprof::time(crate::cpuprof::Slot::SplitAct);
13305    let n = x.len();
13306    let rms = (x.iter().map(|&v| (v * v) as f64).sum::<f64>() / n.max(1) as f64).sqrt() as f32;
13307    let thr = 8.0 * rms;
13308    // One pass: collect outliers and the bulk absmax (outliers excluded —
13309    // identical to the old zero-then-fold over a copied buffer, minus the
13310    // full-vector copy).
13311    let mut outliers: Vec<(usize, f32)> = Vec::new();
13312    let mut amax = 0f32;
13313    for (j, &v) in x.iter().enumerate() {
13314        let a = v.abs();
13315        if a > thr {
13316            outliers.push((j, v));
13317        } else if a > amax {
13318            amax = a;
13319        }
13320    }
13321    let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
13322    let inv = 1.0 / sx;
13323    crate::cpuprof::count_outliers(outliers.len() as u64);
13324    let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
13325    xq.clear();
13326    xq.reserve(n);
13327    if outliers.is_empty() {
13328        xq.extend(
13329            x.iter()
13330                .map(|&v| q8_round(v * inv)),
13331        );
13332    } else {
13333        // Outlier slots quantize to 0 (their exact term is added later).
13334        xq.extend(x.iter().map(|&v| {
13335            if v.abs() > thr {
13336                0
13337            } else {
13338                q8_round(v * inv)
13339            }
13340        }));
13341    }
13342    let xsum = xq.iter().map(|&v| v as i32).sum();
13343    let mut xg8 = xg8_take();
13344    xg8_fill(&xq, &mut xg8);
13345    SplitAct {
13346        xq,
13347        sx,
13348        outliers,
13349        xsum,
13350        xg8,
13351    }
13352}
13353
13354fn split_act_q8_2f(x: &[f32], col: &[f32]) -> SplitAct {
13355    let _prof = crate::cpuprof::time(crate::cpuprof::Slot::SplitAct);
13356    let n = x.len();
13357    let rms = (x
13358        .iter()
13359        .zip(col)
13360        .map(|(&a, &c)| {
13361            let v = a * c;
13362            (v * v) as f64
13363        })
13364        .sum::<f64>()
13365        / n.max(1) as f64)
13366        .sqrt() as f32;
13367    let thr = 8.0 * rms;
13368
13369    let mut outliers = Vec::new();
13370    let mut amax = 0f32;
13371    for (j, (&a, &c)) in x.iter().zip(col).enumerate() {
13372        let v = a * c;
13373        let s = v.abs();
13374        if s > thr {
13375            outliers.push((j, v));
13376        } else if s > amax {
13377            amax = s;
13378        }
13379    }
13380
13381    let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
13382    let inv = 1.0 / sx;
13383    let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
13384    xq.clear();
13385    xq.reserve(n);
13386    if outliers.is_empty() {
13387        xq.extend(
13388            x.iter()
13389                .zip(col)
13390                .map(|(&a, &c)| q8_round((a * c) * inv)),
13391        );
13392    } else {
13393        xq.extend(x.iter().zip(col).map(|(&a, &c)| {
13394            let v = a * c;
13395            if v.abs() > thr {
13396                0
13397            } else {
13398                q8_round(v * inv)
13399            }
13400        }));
13401    }
13402    let xsum = xq.iter().map(|&v| v as i32).sum();
13403    let mut xg8 = xg8_take();
13404    xg8_fill(&xq, &mut xg8);
13405    SplitAct {
13406        xq,
13407        sx,
13408        outliers,
13409        xsum,
13410        xg8,
13411    }
13412}
13413
13414/// int8(weight)·int8(activation) → i32 via `sdot` (inline asm — the
13415/// vdotq intrinsic is unstable; port of vmfcore `dot_i8_sdot`).
13416#[cfg(target_arch = "aarch64")]
13417#[target_feature(enable = "neon,dotprod")]
13418unsafe fn dot_i8_sdot(w: &[u8], xq: &[i8]) -> i32 {
13419    // SAFETY: callers uphold slice-length contracts (see call sites).
13420    unsafe {
13421        use core::arch::aarch64::*;
13422        use core::arch::asm;
13423        let wp = w.as_ptr() as *const i8;
13424        let n = w.len();
13425        let (mut a0, mut a1, mut a2, mut a3) = (
13426            vdupq_n_s32(0),
13427            vdupq_n_s32(0),
13428            vdupq_n_s32(0),
13429            vdupq_n_s32(0),
13430        );
13431        let mut i = 0;
13432        while i + 64 <= n {
13433            let (w0, x0) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
13434            let (w1, x1) = (vld1q_s8(wp.add(i + 16)), vld1q_s8(xq.as_ptr().add(i + 16)));
13435            let (w2, x2) = (vld1q_s8(wp.add(i + 32)), vld1q_s8(xq.as_ptr().add(i + 32)));
13436            let (w3, x3) = (vld1q_s8(wp.add(i + 48)), vld1q_s8(xq.as_ptr().add(i + 48)));
13437            asm!(
13438                "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
13439                "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
13440                "sdot {a2:v}.4s, {w2:v}.16b, {x2:v}.16b",
13441                "sdot {a3:v}.4s, {w3:v}.16b, {x3:v}.16b",
13442                a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
13443                w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
13444                w2 = in(vreg) w2, x2 = in(vreg) x2, w3 = in(vreg) w3, x3 = in(vreg) x3,
13445                options(pure, nomem, nostack),
13446            );
13447            i += 64;
13448        }
13449        while i + 16 <= n {
13450            let (wv, xv) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
13451            asm!("sdot {a:v}.4s, {w:v}.16b, {x:v}.16b",
13452                 a = inout(vreg) a0, w = in(vreg) wv, x = in(vreg) xv, options(pure, nomem, nostack));
13453            i += 16;
13454        }
13455        let mut s = vaddvq_s32(vaddq_s32(vaddq_s32(a0, a1), vaddq_s32(a2, a3)));
13456        while i < n {
13457            s += (*wp.add(i)) as i32 * xq[i] as i32;
13458            i += 1;
13459        }
13460        s
13461    }
13462}
13463
13464/// Row-blocked SDOT: 4 output rows per pass — the activation chunk is
13465/// loaded once and reused, 4 independent accumulators hide sdot latency
13466/// (port of vmfcore `dot_i8_sdot_4rows`).
13467#[cfg(target_arch = "aarch64")]
13468#[target_feature(enable = "neon,dotprod")]
13469unsafe fn dot_i8_sdot_4rows(w0: &[u8], w1: &[u8], w2: &[u8], w3: &[u8], xq: &[i8]) -> [i32; 4] {
13470    // SAFETY: callers uphold slice-length contracts (see call sites).
13471    unsafe {
13472        use core::arch::aarch64::*;
13473        use core::arch::asm;
13474        let n = xq.len();
13475        let px = xq.as_ptr();
13476        let (p0, p1, p2, p3) = (
13477            w0.as_ptr() as *const i8,
13478            w1.as_ptr() as *const i8,
13479            w2.as_ptr() as *const i8,
13480            w3.as_ptr() as *const i8,
13481        );
13482        let (mut a0, mut a1, mut a2, mut a3) = (
13483            vdupq_n_s32(0),
13484            vdupq_n_s32(0),
13485            vdupq_n_s32(0),
13486            vdupq_n_s32(0),
13487        );
13488        let mut i = 0;
13489        while i + 16 <= n {
13490            let x = vld1q_s8(px.add(i));
13491            let v0 = vld1q_s8(p0.add(i));
13492            let v1 = vld1q_s8(p1.add(i));
13493            let v2 = vld1q_s8(p2.add(i));
13494            let v3 = vld1q_s8(p3.add(i));
13495            asm!(
13496                "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
13497                "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
13498                "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
13499                "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
13500                a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
13501                v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
13502                options(pure, nomem, nostack),
13503            );
13504            i += 16;
13505        }
13506        let mut r = [
13507            vaddvq_s32(a0),
13508            vaddvq_s32(a1),
13509            vaddvq_s32(a2),
13510            vaddvq_s32(a3),
13511        ];
13512        while i < n {
13513            let xi = *px.add(i) as i32;
13514            r[0] += (*p0.add(i)) as i32 * xi;
13515            r[1] += (*p1.add(i)) as i32 * xi;
13516            r[2] += (*p2.add(i)) as i32 * xi;
13517            r[3] += (*p3.add(i)) as i32 * xi;
13518            i += 1;
13519        }
13520        r
13521    }
13522}
13523
13524/// 4 interleaved rows in one pass: the repacked group is [r0[c], r1[c],
13525/// r2[c], r3[c]] per 16-byte chunk, so each iteration reads ONE 64-byte
13526/// line plus the shared activation chunk — a single sequential weight
13527/// stream per worker. Per-row accumulation is the same one-accumulator
13528/// scheme as `dot_i8_sdot_4rows`; integer sums are exact, so outputs
13529/// are bit-identical to the mmap-layout kernel.
13530#[cfg(target_arch = "aarch64")]
13531#[target_feature(enable = "neon,dotprod")]
13532unsafe fn dot_i8_sdot_4rows_il(g: &[u8], xq: &[i8]) -> [i32; 4] {
13533    // SAFETY: callers uphold slice-length contracts (g.len() == 4·n,
13534    // n % 16 == 0 — guaranteed by the repack gate).
13535    unsafe {
13536        use core::arch::aarch64::*;
13537        use core::arch::asm;
13538        let n = xq.len();
13539        let px = xq.as_ptr();
13540        let pg = g.as_ptr() as *const i8;
13541        let (mut a0, mut a1, mut a2, mut a3) = (
13542            vdupq_n_s32(0),
13543            vdupq_n_s32(0),
13544            vdupq_n_s32(0),
13545            vdupq_n_s32(0),
13546        );
13547        let mut i = 0;
13548        while i + 16 <= n {
13549            let x = vld1q_s8(px.add(i));
13550            let base = pg.add(4 * i);
13551            let v0 = vld1q_s8(base);
13552            let v1 = vld1q_s8(base.add(16));
13553            let v2 = vld1q_s8(base.add(32));
13554            let v3 = vld1q_s8(base.add(48));
13555            asm!(
13556                "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
13557                "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
13558                "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
13559                "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
13560                a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
13561                v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
13562                options(pure, nomem, nostack),
13563            );
13564            i += 16;
13565        }
13566        [
13567            vaddvq_s32(a0),
13568            vaddvq_s32(a1),
13569            vaddvq_s32(a2),
13570            vaddvq_s32(a3),
13571        ]
13572    }
13573}
13574
13575/// One q8 row range via SDOT (4-row blocks + tail) — the body of
13576/// `qmatvec`'s hot loop, extracted so multi-matrix jobs can drive the
13577/// SAME kernel for several tensors under one pool dispatch. `rep` — the
13578/// load-time interleaved repack (empty = mmap layout only); rows outside
13579/// full 4-row groups always come from the mmap layout.
13580#[cfg(target_arch = "aarch64")]
13581fn q8_range_sdot(
13582    q: &[u8],
13583    rep: &[u8],
13584    row_scale: &[f32],
13585    act: &SplitAct,
13586    cols: usize,
13587    out_addr: SendMut,
13588    start: usize,
13589    end: usize,
13590) {
13591    let mut o = start;
13592    // Leading rows to the group boundary (repack path only): the pool
13593    // splits row ranges arbitrarily, groups are absolute.
13594    if !rep.is_empty() {
13595        while o < end && o % 4 != 0 {
13596            let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
13597            unsafe { *out_addr.at(o) = v };
13598            o += 1;
13599        }
13600    }
13601    while o + 4 <= end {
13602        let r = if rep.is_empty() {
13603            unsafe {
13604                dot_i8_sdot_4rows(
13605                    &q[o * cols..(o + 1) * cols],
13606                    &q[(o + 1) * cols..(o + 2) * cols],
13607                    &q[(o + 2) * cols..(o + 3) * cols],
13608                    &q[(o + 3) * cols..(o + 4) * cols],
13609                    &act.xq,
13610                )
13611            }
13612        } else {
13613            unsafe { dot_i8_sdot_4rows_il(&rep[o * cols..(o + 4) * cols], &act.xq) }
13614        };
13615        for k in 0..4 {
13616            let mut acc = r[k] as f32 * act.sx;
13617            for &(j, xv) in &act.outliers {
13618                acc += (q[(o + k) * cols + j] as i8) as f32 * xv;
13619            }
13620            // SAFETY: disjoint row ranges per worker.
13621            unsafe { *out_addr.at(o + k) = acc * row_scale[o + k] };
13622        }
13623        o += 4;
13624    }
13625    while o < end {
13626        let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
13627        unsafe { *out_addr.at(o) = v };
13628        o += 1;
13629    }
13630}
13631
13632/// Two-input q8 row range via SDOT — `qmatvec2`'s hot loop, extracted
13633/// for the fused pair multi-matrix job (`matvec2_many`).
13634#[cfg(target_arch = "aarch64")]
13635#[allow(clippy::too_many_arguments)]
13636fn q8_range2_sdot(
13637    q: &[u8],
13638    row_scale: &[f32],
13639    a1: &SplitAct,
13640    a2: &SplitAct,
13641    cols: usize,
13642    p1: SendMut,
13643    p2: SendMut,
13644    start: usize,
13645    end: usize,
13646) {
13647    for o in start..end {
13648        let row = &q[o * cols..(o + 1) * cols];
13649        // SAFETY: disjoint row ranges per worker.
13650        unsafe {
13651            *p1.at(o) = row_dot_sdot(row, a1) * row_scale[o];
13652            *p2.at(o) = row_dot_sdot(row, a2) * row_scale[o];
13653        }
13654    }
13655}
13656
13657/// Two-input q8 row range, f32 kernel (non-SDOT) — same extraction.
13658#[allow(clippy::too_many_arguments)]
13659fn q8_range2_f32(
13660    q: &[u8],
13661    row_scale: &[f32],
13662    x1: &[f32],
13663    x2: &[f32],
13664    cols: usize,
13665    p1: SendMut,
13666    p2: SendMut,
13667    start: usize,
13668    end: usize,
13669) {
13670    for o in start..end {
13671        let row = &q[o * cols..(o + 1) * cols];
13672        // SAFETY: disjoint row ranges per worker.
13673        unsafe {
13674            *p1.at(o) = dot_i8_f32(row, x1) * row_scale[o];
13675            *p2.at(o) = dot_i8_f32(row, x2) * row_scale[o];
13676        }
13677    }
13678}
13679
13680/// Scalar/NEON-f32 q8 row range (non-SDOT platforms) — same extraction.
13681fn q8_range_f32(
13682    q: &[u8],
13683    row_scale: &[f32],
13684    xs: &[f32],
13685    cols: usize,
13686    out_addr: SendMut,
13687    start: usize,
13688    end: usize,
13689) {
13690    for o in start..end {
13691        let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
13692        // SAFETY: disjoint row ranges per worker.
13693        unsafe { *out_addr.at(o) = v };
13694    }
13695}
13696
13697/// One q8 row against a split activation, portable: the per-arch fast
13698/// dots where they exist, the exact scalar loop elsewhere. The scalar
13699/// arm is also the test oracle for both fast arms.
13700#[inline]
13701fn q8_row_dot(row: &[u8], act: &SplitAct) -> f32 {
13702    #[cfg(target_arch = "aarch64")]
13703    return row_dot_sdot(row, act);
13704    #[cfg(target_arch = "x86_64")]
13705    return row_dot_avx2(row, act);
13706    #[allow(unreachable_code)]
13707    q8_row_dot_scalar(row, act)
13708}
13709
13710#[allow(dead_code)]
13711fn q8_row_dot_scalar(row: &[u8], act: &SplitAct) -> f32 {
13712    let mut acc = 0i32;
13713    for (k, &b) in row.iter().enumerate() {
13714        acc += (b as i8) as i32 * act.xq[k] as i32;
13715    }
13716    let mut acc = acc as f32 * act.sx;
13717    for &(j, xv) in &act.outliers {
13718        acc += (row[j] as i8) as f32 * xv;
13719    }
13720    acc
13721}
13722
13723/// SDOT row dot with exact outlier correction:
13724/// `dot = sdot(w, xq)·sx + Σ_outl w[j]·x[j]` (then × row_scale by caller).
13725#[cfg(target_arch = "aarch64")]
13726#[inline]
13727fn row_dot_sdot(row: &[u8], act: &SplitAct) -> f32 {
13728    let mut acc = unsafe { dot_i8_sdot(row, &act.xq) } as f32 * act.sx;
13729    for &(j, xv) in &act.outliers {
13730        acc += (row[j] as i8) as f32 * xv;
13731    }
13732    acc
13733}
13734
13735/// One q4 row via SDOT: each 32-group's nibbles unpack to centered i8
13736/// (nib−8 ∈ [−8,7]), int8×int8 `sdot` against the pre-quantized
13737/// activation group, × the group's f16 scale. Returns Σ_g dot_g·s_g;
13738/// the caller multiplies by the activation scale and adds the exact
13739/// outlier terms (port of vmfcore `dot_q4_block_sdot`, +23% measured).
13740/// Nibble order matches the writer: element 2k = low nibble, 2k+1 = high
13741/// → zip(lo,hi) restores flat order.
13742#[cfg(target_arch = "aarch64")]
13743#[target_feature(enable = "neon,dotprod")]
13744unsafe fn dot_q4_row_sdot(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
13745    // SAFETY: callers uphold slice-length contracts (16 packed bytes and
13746    // 2 scale bytes per group; xq.len() == gpr·GROUP_SIZE).
13747    unsafe {
13748        use core::arch::aarch64::*;
13749        use core::arch::asm;
13750        let lomask = vdupq_n_u8(0x0F);
13751        let eight = vdupq_n_s8(8);
13752        let mut acc = 0f32;
13753        for gi in 0..gpr {
13754            let g = g0 + gi;
13755            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
13756            let b = vld1q_u8(packed.as_ptr().add(g * 16));
13757            let lo = vandq_u8(b, lomask);
13758            let hi = vshrq_n_u8::<4>(b);
13759            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
13760            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
13761            let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
13762            let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
13763            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
13764            asm!(
13765                "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
13766                "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
13767                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
13768                e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
13769                options(pure, nomem, nostack),
13770            );
13771            acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
13772        }
13773        acc
13774    }
13775}
13776
13777/// Two-activation q4 row via SDOT: the nibble unpack (the expensive
13778/// part) happens ONCE per group; both pre-quantized activations are
13779/// dotted against the same centered i8 registers. Per-lane math matches
13780/// `dot_q4_row_sdot` exactly.
13781#[cfg(target_arch = "aarch64")]
13782#[target_feature(enable = "neon,dotprod")]
13783unsafe fn dot_q4_row_sdot2(
13784    packed: &[u8],
13785    scales: &[u8],
13786    g0: usize,
13787    gpr: usize,
13788    xq1: &[i8],
13789    xq2: &[i8],
13790) -> (f32, f32) {
13791    // SAFETY: callers uphold slice-length contracts (16 packed bytes and
13792    // 2 scale bytes per group; xq*.len() == gpr·GROUP_SIZE).
13793    unsafe {
13794        use core::arch::aarch64::*;
13795        use core::arch::asm;
13796        let lomask = vdupq_n_u8(0x0F);
13797        let eight = vdupq_n_s8(8);
13798        let (mut acc1, mut acc2) = (0f32, 0f32);
13799        for gi in 0..gpr {
13800            let g = g0 + gi;
13801            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
13802            let b = vld1q_u8(packed.as_ptr().add(g * 16));
13803            let lo = vandq_u8(b, lomask);
13804            let hi = vshrq_n_u8::<4>(b);
13805            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
13806            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
13807            let x10 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE));
13808            let x11 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE + 16));
13809            let x20 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE));
13810            let x21 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE + 16));
13811            let (mut a0, mut a1, mut b0, mut b1) = (
13812                vdupq_n_s32(0),
13813                vdupq_n_s32(0),
13814                vdupq_n_s32(0),
13815                vdupq_n_s32(0),
13816            );
13817            asm!(
13818                "sdot {a0:v}.4s, {e0:v}.16b, {x10:v}.16b",
13819                "sdot {a1:v}.4s, {e1:v}.16b, {x11:v}.16b",
13820                "sdot {b0:v}.4s, {e0:v}.16b, {x20:v}.16b",
13821                "sdot {b1:v}.4s, {e1:v}.16b, {x21:v}.16b",
13822                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
13823                b0 = inout(vreg) b0, b1 = inout(vreg) b1,
13824                e0 = in(vreg) e0, e1 = in(vreg) e1,
13825                x10 = in(vreg) x10, x11 = in(vreg) x11,
13826                x20 = in(vreg) x20, x21 = in(vreg) x21,
13827                options(pure, nomem, nostack),
13828            );
13829            acc1 += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
13830            acc2 += vaddvq_s32(vaddq_s32(b0, b1)) as f32 * s;
13831        }
13832        (acc1, acc2)
13833    }
13834}
13835
13836// ───────────────────── fused int8 kernels ─────────────────────
13837
13838/// `acc += w · row` where the row is centered i8 — NEON widen+fma on
13839/// aarch64, scalar elsewhere. The KV-cache q8 value path rides on this.
13840#[inline]
13841pub(crate) fn axpy_i8_f32(acc: &mut [f32], row: &[i8], w: f32) {
13842    #[cfg(target_arch = "aarch64")]
13843    unsafe {
13844        return axpy_i8_f32_neon(acc, row, w);
13845    }
13846    #[cfg(target_arch = "x86_64")]
13847    if avx2_enabled() {
13848        return unsafe { axpy_i8_f32_avx2(acc, row, w) };
13849    }
13850    #[allow(unreachable_code)]
13851    {
13852        for (a, &b) in acc.iter_mut().zip(row) {
13853            *a += w * b as f32;
13854        }
13855    }
13856}
13857
13858/// i8→f32 axpy via AVX2/FMA (x86 mirror of `axpy_i8_f32_neon`).
13859#[cfg(target_arch = "x86_64")]
13860#[target_feature(enable = "avx2,fma")]
13861unsafe fn axpy_i8_f32_avx2(acc: &mut [f32], row: &[i8], w: f32) {
13862    // SAFETY: callers uphold slice-length contracts (see call sites).
13863    unsafe {
13864        use core::arch::x86_64::*;
13865        let n = acc.len().min(row.len());
13866        let ap = acc.as_mut_ptr();
13867        let rp = row.as_ptr();
13868        let wv = _mm256_set1_ps(w);
13869        let mut j = 0usize;
13870        while j + 16 <= n {
13871            let rb = _mm_loadu_si128(rp.add(j) as *const __m128i);
13872            let lo = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(rb));
13873            let hi = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(rb)));
13874            let v0 = _mm256_fmadd_ps(wv, lo, _mm256_loadu_ps(ap.add(j)));
13875            let v1 = _mm256_fmadd_ps(wv, hi, _mm256_loadu_ps(ap.add(j + 8)));
13876            _mm256_storeu_ps(ap.add(j), v0);
13877            _mm256_storeu_ps(ap.add(j + 8), v1);
13878            j += 16;
13879        }
13880        while j < n {
13881            *ap.add(j) += w * (*rp.add(j)) as f32;
13882            j += 1;
13883        }
13884    }
13885}
13886
13887#[cfg(target_arch = "aarch64")]
13888#[target_feature(enable = "neon")]
13889unsafe fn axpy_i8_f32_neon(acc: &mut [f32], row: &[i8], w: f32) {
13890    // SAFETY: callers uphold slice-length contracts (see call sites).
13891    unsafe {
13892        use core::arch::aarch64::*;
13893        let n = acc.len().min(row.len());
13894        let ap = acc.as_mut_ptr();
13895        let rp = row.as_ptr();
13896        let wv = vdupq_n_f32(w);
13897        let mut j = 0usize;
13898        while j + 16 <= n {
13899            let rb = vld1q_s8(rp.add(j));
13900            let lo = vmovl_s8(vget_low_s8(rb));
13901            let hi = vmovl_s8(vget_high_s8(rb));
13902            for (off, half) in [(0, lo), (8, hi)] {
13903                let f0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(half)));
13904                let f1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(half)));
13905                let o = j + off;
13906                vst1q_f32(ap.add(o), vfmaq_f32(vld1q_f32(ap.add(o)), wv, f0));
13907                vst1q_f32(ap.add(o + 4), vfmaq_f32(vld1q_f32(ap.add(o + 4)), wv, f1));
13908            }
13909            j += 16;
13910        }
13911        while j < n {
13912            *ap.add(j) += w * (*rp.add(j)) as f32;
13913            j += 1;
13914        }
13915    }
13916}
13917
13918/// i8 row · f32 x. NEON on aarch64 (ported from vmfcore `dot_i8_f32_neon`,
13919/// ≈9× scalar), scalar elsewhere.
13920#[inline]
13921pub(crate) fn dot_i8_f32(w: &[u8], x: &[f32]) -> f32 {
13922    #[cfg(target_arch = "aarch64")]
13923    unsafe {
13924        return dot_i8_f32_neon(w, x);
13925    }
13926    #[cfg(target_arch = "x86_64")]
13927    if avx2_enabled() {
13928        return unsafe { dot_i8_f32_avx2(w, x) };
13929    }
13930    #[allow(unreachable_code)]
13931    {
13932        let mut sum = 0.0f32;
13933        for (j, &b) in w.iter().enumerate() {
13934            sum += (b as i8) as f32 * x[j];
13935        }
13936        sum
13937    }
13938}
13939
13940/// i8 row · (x ⊙ col_field) — the q8_2f row dot with the θ col-field
13941/// folded into the product (no prescaled copy of x). NEON on aarch64,
13942/// scalar elsewhere. Used by the active-neuron path `row_dot`.
13943#[inline]
13944fn dot_i8_col_f32(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
13945    #[cfg(target_arch = "aarch64")]
13946    unsafe {
13947        return dot_i8_col_f32_neon(w, x, col);
13948    }
13949    #[allow(unreachable_code)]
13950    {
13951        let mut sum = 0.0f32;
13952        for (j, &b) in w.iter().enumerate() {
13953            sum += (b as i8) as f32 * x[j] * col[j];
13954        }
13955        sum
13956    }
13957}
13958
13959#[cfg(target_arch = "aarch64")]
13960#[target_feature(enable = "neon")]
13961unsafe fn dot_i8_col_f32_neon(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
13962    // SAFETY: callers uphold slice-length contracts (see call sites).
13963    unsafe {
13964        use core::arch::aarch64::*;
13965        let n = x.len();
13966        let wp = w.as_ptr() as *const i8;
13967        let xp = x.as_ptr();
13968        let cp = col.as_ptr();
13969        let (mut a0, mut a1, mut a2, mut a3) = (
13970            vdupq_n_f32(0.0),
13971            vdupq_n_f32(0.0),
13972            vdupq_n_f32(0.0),
13973            vdupq_n_f32(0.0),
13974        );
13975        let mut j = 0usize;
13976        while j + 16 <= n {
13977            let wb = vld1q_s8(wp.add(j));
13978            let lo = vmovl_s8(vget_low_s8(wb));
13979            let hi = vmovl_s8(vget_high_s8(wb));
13980            let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
13981            let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
13982            let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
13983            let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
13984            a0 = vfmaq_f32(
13985                a0,
13986                w0,
13987                vmulq_f32(vld1q_f32(xp.add(j)), vld1q_f32(cp.add(j))),
13988            );
13989            a1 = vfmaq_f32(
13990                a1,
13991                w1,
13992                vmulq_f32(vld1q_f32(xp.add(j + 4)), vld1q_f32(cp.add(j + 4))),
13993            );
13994            a2 = vfmaq_f32(
13995                a2,
13996                w2,
13997                vmulq_f32(vld1q_f32(xp.add(j + 8)), vld1q_f32(cp.add(j + 8))),
13998            );
13999            a3 = vfmaq_f32(
14000                a3,
14001                w3,
14002                vmulq_f32(vld1q_f32(xp.add(j + 12)), vld1q_f32(cp.add(j + 12))),
14003            );
14004            j += 16;
14005        }
14006        let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
14007        while j < n {
14008            sum += (*wp.add(j)) as f32 * *xp.add(j) * *cp.add(j);
14009            j += 1;
14010        }
14011        sum
14012    }
14013}
14014
14015#[cfg(target_arch = "aarch64")]
14016#[target_feature(enable = "neon")]
14017unsafe fn dot_i8_f32_neon(w: &[u8], x: &[f32]) -> f32 {
14018    // SAFETY: callers uphold slice-length contracts (see call sites).
14019    unsafe {
14020        use core::arch::aarch64::*;
14021        let n = x.len();
14022        let wp = w.as_ptr() as *const i8;
14023        let xp = x.as_ptr();
14024        let (mut a0, mut a1, mut a2, mut a3) = (
14025            vdupq_n_f32(0.0),
14026            vdupq_n_f32(0.0),
14027            vdupq_n_f32(0.0),
14028            vdupq_n_f32(0.0),
14029        );
14030        let mut j = 0usize;
14031        while j + 16 <= n {
14032            let wb = vld1q_s8(wp.add(j));
14033            let lo = vmovl_s8(vget_low_s8(wb));
14034            let hi = vmovl_s8(vget_high_s8(wb));
14035            let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
14036            let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
14037            let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
14038            let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
14039            a0 = vfmaq_f32(a0, w0, vld1q_f32(xp.add(j)));
14040            a1 = vfmaq_f32(a1, w1, vld1q_f32(xp.add(j + 4)));
14041            a2 = vfmaq_f32(a2, w2, vld1q_f32(xp.add(j + 8)));
14042            a3 = vfmaq_f32(a3, w3, vld1q_f32(xp.add(j + 12)));
14043            j += 16;
14044        }
14045        let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
14046        while j < n {
14047            sum += (*wp.add(j)) as f32 * *xp.add(j);
14048            j += 1;
14049        }
14050        sum
14051    }
14052}
14053
14054#[allow(clippy::too_many_arguments)]
14055fn qmatvec(
14056    q: &[u8],
14057    rep: &[u8],
14058    row_scale: &[f32],
14059    x: &[f32],
14060    col_field: &[f32],
14061    dtype: TensorDtype,
14062    rows: usize,
14063    cols: usize,
14064    out: &mut [f32],
14065    pool: Option<&Pool>,
14066) {
14067    debug_assert_eq!(out.len(), rows);
14068    #[cfg(not(target_arch = "aarch64"))]
14069    let _ = rep;
14070
14071    #[cfg(target_arch = "aarch64")]
14072    if sdot_enabled() {
14073        let act = if dtype == TensorDtype::Q8_2f {
14074            split_act_q8_2f(x, col_field)
14075        } else {
14076            split_act(x)
14077        };
14078        let out_addr = SendMut(out.as_mut_ptr());
14079        let run_range = |start: usize, end: usize| {
14080            q8_range_sdot(q, rep, row_scale, &act, cols, out_addr, start, end)
14081        };
14082        match pool {
14083            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
14084            _ => run_range(0, rows),
14085        }
14086        return;
14087    }
14088    // x86 A8W8 via AVX2 maddubs — same quantized-activation contract as
14089    // the SDOT path (CMF_AVX2=0 keeps the exact i8×f32 loop).
14090    #[cfg(target_arch = "x86_64")]
14091    if avx2_a8w8_enabled() {
14092        let act = if dtype == TensorDtype::Q8_2f {
14093            split_act_q8_2f(x, col_field)
14094        } else {
14095            split_act(x)
14096        };
14097        let out_addr = SendMut(out.as_mut_ptr());
14098        let run_range = |start: usize, end: usize| {
14099            q8_range_avx2(q, row_scale, &act, cols, out_addr, start, end)
14100        };
14101        match pool {
14102            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
14103            _ => run_range(0, rows),
14104        }
14105        return;
14106    }
14107
14108    prescale_with(x, col_field, dtype, 1, |xs| {
14109        let out_addr = SendMut(out.as_mut_ptr());
14110        let run_range = move |start: usize, end: usize| {
14111            for o in start..end {
14112                let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
14113                // SAFETY: disjoint row ranges per worker.
14114                unsafe { *out_addr.at(o) = v };
14115            }
14116        };
14117        match pool {
14118            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
14119            _ => run_range(0, rows),
14120        }
14121    });
14122}
14123
14124#[allow(clippy::too_many_arguments)]
14125fn qmatvec2(
14126    q: &[u8],
14127    row_scale: &[f32],
14128    x1: &[f32],
14129    x2: &[f32],
14130    col_field: &[f32],
14131    dtype: TensorDtype,
14132    rows: usize,
14133    cols: usize,
14134    o1: &mut [f32],
14135    o2: &mut [f32],
14136    pool: Option<&Pool>,
14137) {
14138    #[cfg(target_arch = "aarch64")]
14139    if sdot_enabled() {
14140        let a1s = if dtype == TensorDtype::Q8_2f {
14141            split_act_q8_2f(x1, col_field)
14142        } else {
14143            split_act(x1)
14144        };
14145        let a2s = if dtype == TensorDtype::Q8_2f {
14146            split_act_q8_2f(x2, col_field)
14147        } else {
14148            split_act(x2)
14149        };
14150        let p1 = SendMut(o1.as_mut_ptr());
14151        let p2 = SendMut(o2.as_mut_ptr());
14152        let run_range = |start: usize, end: usize| {
14153            q8_range2_sdot(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
14154        };
14155        match pool {
14156            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
14157            _ => run_range(0, rows),
14158        }
14159        return;
14160    }
14161    #[cfg(target_arch = "x86_64")]
14162    if avx2_a8w8_enabled() {
14163        let a1s = if dtype == TensorDtype::Q8_2f {
14164            split_act_q8_2f(x1, col_field)
14165        } else {
14166            split_act(x1)
14167        };
14168        let a2s = if dtype == TensorDtype::Q8_2f {
14169            split_act_q8_2f(x2, col_field)
14170        } else {
14171            split_act(x2)
14172        };
14173        let p1 = SendMut(o1.as_mut_ptr());
14174        let p2 = SendMut(o2.as_mut_ptr());
14175        let run_range = |start: usize, end: usize| {
14176            q8_range2_avx2(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
14177        };
14178        match pool {
14179            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
14180            _ => run_range(0, rows),
14181        }
14182        return;
14183    }
14184
14185    prescale_with(x1, col_field, dtype, 1, |x1s| {
14186        prescale_with(x2, col_field, dtype, 2, |x2s| {
14187            let p1 = SendMut(o1.as_mut_ptr());
14188            let p2 = SendMut(o2.as_mut_ptr());
14189            let run_range = move |start: usize, end: usize| {
14190                for o in start..end {
14191                    let row = &q[o * cols..(o + 1) * cols];
14192                    let s1 = dot_i8_f32(row, x1s) * row_scale[o];
14193                    let s2 = dot_i8_f32(row, x2s) * row_scale[o];
14194                    // SAFETY: disjoint row ranges per worker.
14195                    unsafe {
14196                        *p1.at(o) = s1;
14197                        *p2.at(o) = s2;
14198                    }
14199                }
14200            };
14201            match pool {
14202                Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
14203                _ => run_range(0, rows),
14204            }
14205        });
14206    });
14207}
14208
14209#[derive(Clone, Copy)]
14210struct SendMut(*mut f32);
14211unsafe impl Send for SendMut {}
14212unsafe impl Sync for SendMut {}
14213
14214impl SendMut {
14215    #[inline]
14216    fn at(self, i: usize) -> *mut f32 {
14217        unsafe { self.0.add(i) }
14218    }
14219}
14220
14221#[cfg(test)]
14222mod tests {
14223
14224    /// `q8_round` must be `round().clamp(±127) as i8` bit for bit: every
14225    /// half-integer, their neighbours one ulp either side, the clamp
14226    /// boundary, huge values, infinities and NaN, plus a dense sweep.
14227    #[test]
14228    fn q8_round_is_round_clamp() {
14229        let reference = |t: f32| t.round().clamp(-127.0, 127.0) as i8;
14230        let mut probe = vec![
14231            0.0f32,
14232            -0.0,
14233            f32::NAN,
14234            f32::INFINITY,
14235            f32::NEG_INFINITY,
14236            f32::MAX,
14237            f32::MIN,
14238            1e30,
14239            -1e30,
14240            f32::MIN_POSITIVE,
14241            -f32::MIN_POSITIVE,
14242        ];
14243        for k in -300i32..=300 {
14244            let h = k as f32 * 0.5;
14245            let up = f32::from_bits(h.to_bits() + 1);
14246            let down = f32::from_bits(h.to_bits().wrapping_sub(1));
14247            for t in [h, up, down] {
14248                probe.push(t);
14249                probe.push(-t);
14250            }
14251        }
14252        let mut t = -140.0f32;
14253        while t < 140.0 {
14254            probe.push(t);
14255            t += 0.000_731;
14256        }
14257        for t in probe {
14258            assert_eq!(q8_round(t), reference(t), "t = {t:e} ({:#x})", t.to_bits());
14259        }
14260    }
14261
14262    use super::*;
14263
14264    /// The f32 `matmat` walks the batch four positions a row at a time; every
14265    /// output must still be the scalar-order sum `matvec` gives, bit for bit
14266    /// (a 7-row batch covers the four-wide block and the remainder).
14267    #[test]
14268    fn f32_matmat_equals_matvec_bitwise() {
14269        let (rows, cols, b) = (5usize, 37usize, 7usize);
14270        let w: Vec<f32> = (0..rows * cols).map(|i| ((i * 7919 % 113) as f32 - 56.0) / 37.0).collect();
14271        let xs: Vec<f32> = (0..b * cols).map(|i| ((i * 104_729 % 97) as f32 - 48.0) / 29.0).collect();
14272        let t = QTensor::from_f32(w, rows, cols);
14273        let mut out = vec![0f32; b * rows];
14274        t.matmat(&xs, b, &mut out, None);
14275        for bi in 0..b {
14276            let mut one = vec![0f32; rows];
14277            t.matvec(&xs[bi * cols..(bi + 1) * cols], &mut one, None);
14278            for o in 0..rows {
14279                assert_eq!(out[bi * rows + o].to_bits(), one[o].to_bits(), "bi {bi} row {o}");
14280            }
14281        }
14282    }
14283
14284    /// A short batch on a pool splits the f32 GEMM over its output rows
14285    /// (it used to run on one core); the sums must not change by a bit.
14286    #[test]
14287    fn f32_matmat_row_split_on_pool_is_bitwise() {
14288        let (rows, cols, b) = (300usize, 37usize, 7usize);
14289        let w: Vec<f32> = (0..rows * cols).map(|i| ((i * 7919 % 113) as f32 - 56.0) / 37.0).collect();
14290        let xs: Vec<f32> = (0..b * cols).map(|i| ((i * 104_729 % 97) as f32 - 48.0) / 29.0).collect();
14291        let t = QTensor::from_f32(w, rows, cols);
14292        let mut alone = vec![0f32; b * rows];
14293        t.matmat(&xs, b, &mut alone, None);
14294        let pool = Pool::new(3);
14295        let mut split = vec![0f32; b * rows];
14296        t.matmat(&xs, b, &mut split, Some(&pool));
14297        let bits = |v: &[f32]| v.iter().map(|x| x.to_bits()).collect::<Vec<_>>();
14298        assert_eq!(bits(&alone), bits(&split));
14299    }
14300
14301    #[test]
14302    fn q2tp_i8_dot_matches_exact_on_grid() {
14303        // On-grid activations (±1 → sx=1/127, xq=±127 dequantizes
14304        // exactly, no outliers) must make the integer path agree with
14305        // the exact scalar walk to f32 rounding.
14306        let (rows, cols) = (5, 64);
14307        let gpr = cols / GROUP_SIZE;
14308        // Synthetic codes plane + a flat ladder: scales_into is not under
14309        // test here, so drive dot_q2tp_row_i8 / q2tp_row_exact directly
14310        // with hand-made scales.
14311        let chunks: Vec<u8> = (0..rows * gpr * Q2TP_CHUNK)
14312            .map(|i| (i as u32).wrapping_mul(2654435761) as u8)
14313            .collect();
14314        let scales: Vec<f32> = (0..gpr).map(|g| 0.5 + g as f32 * 0.25).collect();
14315        let x: Vec<f32> = (0..cols)
14316            .map(|i| if i % 3 == 0 { -1.0 } else { 1.0 })
14317            .collect();
14318        let act = split_act(&x);
14319        assert!(
14320            act.outliers.is_empty(),
14321            "on-grid input must have no outliers"
14322        );
14323        let gsum = q1_group_sums(&act.xq, gpr);
14324        for r in 0..rows {
14325            let exact = q2tp_row_exact(&chunks, r, gpr, &x, &scales);
14326            let fast = dot_q2tp_row_i8(&chunks, r, gpr, &act.xq, &gsum, &scales) * act.sx;
14327            assert!(
14328                (exact - fast).abs() <= exact.abs() * 1e-5 + 1e-5,
14329                "row {r}: exact {exact} vs i8 {fast}"
14330            );
14331        }
14332    }
14333
14334    #[cfg(target_arch = "x86_64")]
14335    #[test]
14336    fn q2tp_avx2_dot_matches_scalar_for_random_patterns() {
14337        // Compare the release AVX2 integer dot against the scalar oracle over
14338        // arbitrary packed bytes/activation signs.  This guards the exact
14339        // table-load path used after rejecting a faster-looking decoder whose
14340        // full-checkpoint greedy output drifted.
14341        if !std::arch::is_x86_feature_detected!("avx2") {
14342            return;
14343        }
14344        let mut seed = 0x9e3779b9u32;
14345        let mut next = || {
14346            seed = seed.wrapping_mul(1664525).wrapping_add(1013904223);
14347            seed
14348        };
14349        for _ in 0..20_000 {
14350            let mut ch = [0u8; Q2TP_CHUNK];
14351            let mut x = [0i8; GROUP_SIZE];
14352            for b in &mut ch {
14353                *b = next() as u8;
14354            }
14355            for v in &mut x {
14356                *v = (next() >> 24) as i8;
14357            }
14358            let mut reference = 0i32;
14359            for (k, &b) in ch.iter().enumerate() {
14360                reference += (b & 3) as i32 * x[k * 4] as i32;
14361                reference += ((b >> 2) & 3) as i32 * x[k * 4 + 1] as i32;
14362                reference += ((b >> 4) & 3) as i32 * x[k * 4 + 2] as i32;
14363                reference += ((b >> 6) & 3) as i32 * x[k * 4 + 3] as i32;
14364            }
14365            // SAFETY: guarded by the runtime AVX2 feature check and fixed
14366            // 8-byte/32-byte slice lengths above.
14367            let got = unsafe { q2tp_code_dot_avx2(&ch, &x) };
14368            assert_eq!(got, reference, "packed q2 lane mismatch");
14369        }
14370    }
14371
14372    #[test]
14373    fn q2tp_affine_fuses_half_scale_correction_without_changing_raw_decode() {
14374        let (rows, cols) = (1usize, GROUP_SIZE);
14375        let mut bytes = vec![0u8; Q2TP_CHUNK + 4 + 1];
14376        // Repeating symbols 0,1,2,0 at unit scale.  q2tp's raw B is
14377        // (c-1.5), while the affine Prism operator is (c-1.0).
14378        bytes[..Q2TP_CHUNK].fill(0x24); // codes 0,1,2,0 in LSB-first order
14379        bytes[Q2TP_CHUNK..Q2TP_CHUNK + 2].copy_from_slice(&0u16.to_le_bytes());
14380        bytes[Q2TP_CHUNK + 2..Q2TP_CHUNK + 4].copy_from_slice(&0u16.to_le_bytes());
14381        bytes[Q2TP_CHUNK + 4] = 1; // dtype16 rung 1 = 1.0
14382        let x = vec![1.0f32; cols];
14383        let mut raw = vec![0.0f32; rows];
14384        let mut affine = vec![0.0f32; rows];
14385        q2tp_matvec_for_test(&bytes, &x, rows, cols, &mut raw);
14386        q2tp_affine_matvec_for_test(&bytes, &x, rows, cols, &mut affine);
14387        assert_eq!(raw, vec![-24.0]);
14388        assert_eq!(affine, vec![-8.0]);
14389        assert!((affine[0] - (raw[0] + 0.5 * cols as f32)).abs() < 1e-6);
14390    }
14391
14392    #[test]
14393    fn q8_row_dot_fast_matches_scalar() {
14394        // The per-arch fast dot must agree with the exact scalar oracle
14395        // (same contract the fused q8 FFN arm rides on).
14396        let cols = 96;
14397        let row: Vec<u8> = (0..cols)
14398            .map(|i| ((i * 37 % 251) - 125) as i8 as u8)
14399            .collect();
14400        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.13).sin()).collect();
14401        let act = split_act(&x);
14402        let fast = q8_row_dot(&row, &act);
14403        let scalar = q8_row_dot_scalar(&row, &act);
14404        assert!(
14405            (fast - scalar).abs() <= scalar.abs() * 1e-5 + 1e-5,
14406            "fast {fast} vs scalar {scalar}"
14407        );
14408    }
14409
14410    #[test]
14411    fn f32_matvec_matches_matvec_rows_bitexact() {
14412        let (rows, cols) = (300, 40);
14413        let w: Vec<f32> = (0..rows * cols).map(|i| (i as f32 * 0.017).sin()).collect();
14414        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.05).cos()).collect();
14415        let qt = QTensor::from_f32(w.clone(), rows, cols);
14416
14417        let mut a = vec![0.0f32; rows];
14418        matvec_rows(None, &w, &x, &mut a);
14419        let mut b = vec![0.0f32; rows];
14420        qt.matvec(&x, &mut b, None);
14421        assert_eq!(a, b);
14422    }
14423
14424    #[test]
14425    fn sdot_kernel_exact_on_grid() {
14426        // Activations already on the i8 grid (±1 with amax=1 → sx=1/127,
14427        // xq=±127 dequantizes EXACTLY) → the SDOT path must match the
14428        // exact f32 dot to float rounding. This isolates kernel
14429        // correctness from quantization noise.
14430        eprintln!("sdot_enabled = {}", sdot_enabled());
14431        let (rows, cols) = (9, 80); // odd rows → exercises 4-row + tail
14432        let w: Vec<u8> = (0..rows * cols)
14433            .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
14434            .collect();
14435        let scales: Vec<f32> = (0..rows).map(|o| 0.005 + o as f32 * 0.001).collect();
14436        let x: Vec<f32> = (0..cols)
14437            .map(|i| match i % 3 {
14438                0 => 1.0,
14439                1 => -1.0,
14440                _ => 0.0,
14441            })
14442            .collect();
14443        let mut a = vec![0.0f32; rows];
14444        qmatvec(
14445            &w,
14446            &[],
14447            &scales,
14448            &x,
14449            &[],
14450            TensorDtype::Q8Row,
14451            rows,
14452            cols,
14453            &mut a,
14454            None,
14455        );
14456        for o in 0..rows {
14457            let mut acc = 0.0f32;
14458            for j in 0..cols {
14459                acc += (w[o * cols + j] as i8) as f32 * x[j];
14460            }
14461            let expect = acc * scales[o];
14462            assert!(
14463                (a[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
14464                "row {o}: {} vs {expect}",
14465                a[o]
14466            );
14467        }
14468    }
14469
14470    #[test]
14471    fn q1_tbl_fast_path_matches_reference() {
14472        // gpr = 8 exercises the TBL pair-load fast loop, and the LAST
14473        // row's final 4-tile window trips the 4B-overread guard (the
14474        // payload ends exactly at the last tile) — both paths must
14475        // agree with the dequant reference.
14476        let (rows, cols) = (5, 256);
14477        let gpr = cols / GROUP_SIZE;
14478        let mut bytes = Vec::new();
14479        for t in 0..rows * gpr {
14480            let s = 0.007 + (t % 11) as f32 * 0.004;
14481            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
14482            for j in 0..4 {
14483                bytes.push(((t * 53 + j * 89 + 7) % 249) as u8);
14484            }
14485        }
14486        let x: Vec<f32> = (0..cols)
14487            .map(|i| if (i * 5) % 7 < 3 { 1.0 } else { -1.0 })
14488            .collect();
14489        let mut w = vec![0.0f32; rows * cols];
14490        cortiq_core::quant::dequant_q1(&bytes, &mut w);
14491        let mut got = vec![0.0f32; rows];
14492        q1_matvec(&bytes, &x, rows, cols, &mut got, None);
14493        for o in 0..rows {
14494            let expect: f32 = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
14495            assert!(
14496                (got[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
14497                "row {o}: {} vs {expect}",
14498                got[o]
14499            );
14500        }
14501        // Blocked 1×4 batch (b=5: one quad + remainder) must equal the
14502        // single-matvec path bit-for-bit.
14503        let b = 5usize;
14504        let mut xs_all = Vec::new();
14505        for bi in 0..b {
14506            xs_all.extend(x.iter().map(|v| if bi % 2 == 0 { *v } else { -*v }));
14507        }
14508        let mut mm = vec![0.0f32; b * rows];
14509        q1_matmat(&bytes, &xs_all, b, rows, cols, &mut mm, None);
14510        for bi in 0..b {
14511            let mut single = vec![0.0f32; rows];
14512            q1_matvec(
14513                &bytes,
14514                &xs_all[bi * cols..(bi + 1) * cols],
14515                rows,
14516                cols,
14517                &mut single,
14518                None,
14519            );
14520            assert_eq!(&mm[bi * rows..(bi + 1) * rows], &single[..], "stream {bi}");
14521        }
14522    }
14523
14524    #[test]
14525    fn q1_kernels_match_exact_reference() {
14526        // Synthetic q1 payload: 6-byte tiles [f16 scale][4B bits].
14527        let (rows, cols) = (7, 96);
14528        let gpr = cols / GROUP_SIZE;
14529        let mut bytes = Vec::new();
14530        for t in 0..rows * gpr {
14531            let s = 0.01 + (t % 13) as f32 * 0.003;
14532            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
14533            for j in 0..4 {
14534                bytes.push(((t * 31 + j * 97) % 251) as u8);
14535            }
14536        }
14537        // On-grid activations (±1, amax 1) → the SDOT path is exact.
14538        let x: Vec<f32> = (0..cols)
14539            .map(|i| if i % 3 == 0 { 1.0 } else { -1.0 })
14540            .collect();
14541        // Reference through the core dequant.
14542        let mut w = vec![0.0f32; rows * cols];
14543        cortiq_core::quant::dequant_q1(&bytes, &mut w);
14544        let mut expect = vec![0.0f32; rows];
14545        for o in 0..rows {
14546            expect[o] = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
14547        }
14548        let mut got = vec![0.0f32; rows];
14549        q1_matvec(&bytes, &x, rows, cols, &mut got, None);
14550        for o in 0..rows {
14551            assert!(
14552                (got[o] - expect[o]).abs() < 1e-3 * expect[o].abs().max(1e-3),
14553                "row {o}: {} vs {}",
14554                got[o],
14555                expect[o]
14556            );
14557        }
14558        // Pair and batch paths agree with the single path.
14559        let x2: Vec<f32> = x.iter().map(|v| -v).collect();
14560        let (mut a1, mut a2) = (vec![0.0f32; rows], vec![0.0f32; rows]);
14561        q1_matvec2(&bytes, &x, &x2, rows, cols, &mut a1, &mut a2, None);
14562        assert_eq!(a1, got);
14563        let mut xs = x.clone();
14564        xs.extend_from_slice(&x2);
14565        let mut mm = vec![0.0f32; 2 * rows];
14566        q1_matmat(&bytes, &xs, 2, rows, cols, &mut mm, None);
14567        assert_eq!(&mm[..rows], got.as_slice());
14568        assert_eq!(&mm[rows..], a2.as_slice());
14569    }
14570
14571    /// A Bonsai-shaped q1 payload (one f16 scale per 128-group, stored in
14572    /// all four tiles; pseudo-random signs) and its q1g repack.
14573    fn q1_q1g_pair(rows: usize, cols: usize, salt: usize) -> (Vec<u8>, Vec<u8>) {
14574        let tiles = rows * cols / GROUP_SIZE;
14575        let mut q1 = Vec::with_capacity(tiles * Q1_TILE);
14576        for t in 0..tiles {
14577            let s = 0.004 + ((t / 4 * 37 + salt) % 29) as f32 * 0.0011;
14578            q1.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
14579            let bits = ((t * 2_654_435_761 + salt * 40_503) as u32).rotate_left((t % 31) as u32);
14580            q1.extend_from_slice(&bits.to_le_bytes());
14581        }
14582        let rp = cortiq_core::quant::repack_q1_to_q1g(&q1, rows, cols, 0.0);
14583        assert_eq!(rp.differing, 0, "uniform scales repack exactly");
14584        (q1, rp.data)
14585    }
14586
14587    /// Activations off the int8 grid with two outlier channels, so the
14588    /// A8W8 path quantizes AND takes the exact outlier correction.
14589    fn q1g_test_x(cols: usize, salt: usize) -> Vec<f32> {
14590        let mut x: Vec<f32> = (0..cols)
14591            .map(|i| ((i * 13 + salt * 7) as f32 * 0.071).sin() * 0.9)
14592            .collect();
14593        x[3 % cols] = 41.0;
14594        x[(cols / 2 + 5) % cols] = -37.5;
14595        x
14596    }
14597
14598    /// Every CPU row kernel of q1g returns the q1 kernel's f32 BIT FOR BIT
14599    /// on a uniform-scale repack: A8W8 rows (SIMD of this machine, and the
14600    /// x86 AVX2/VNNI twins directly), the exact scalar rows (CMF_SDOT=0
14601    /// contract), the outlier term, the drivers (matvec, matvec2, matmat
14602    /// with the blocked 1×4 quad + remainder) — including gpr = 4 and the
14603    /// last row, where q1's pair load falls back to its scalar tail.
14604    #[test]
14605    fn q1g_cpu_kernels_equal_q1_bit_for_bit() {
14606        for &(rows, cols) in &[(3usize, 128usize), (5, 256), (9, 2048), (2, 6144)] {
14607            let gpr = cols / GROUP_SIZE;
14608            let (q1, q1g) = q1_q1g_pair(rows, cols, rows + cols);
14609            let (wa, wb) = (BinW::Q1(&q1), BinW::q1g(&q1g, rows, cols));
14610            let x = q1g_test_x(cols, rows);
14611            let act = split_act(&x);
14612            assert!(!act.outliers.is_empty(), "test wants the outlier path");
14613            let gsum = q1_group_sums(&act.xq, gpr);
14614            for r in 0..rows {
14615                let (a, b) = (wa.row_exact(r, gpr, &x), wb.row_exact(r, gpr, &x));
14616                assert_eq!(a.to_bits(), b.to_bits(), "{rows}x{cols} exact row {r}");
14617                let (a, b) = (
14618                    wa.dot_i8(r, gpr, &act.xq, &gsum),
14619                    wb.dot_i8(r, gpr, &act.xq, &gsum),
14620                );
14621                assert_eq!(a.to_bits(), b.to_bits(), "{rows}x{cols} i8 row {r}");
14622                for j in [0, 31, 32, 127, cols - 1] {
14623                    let (a, b) = (wa.outlier(r, gpr, j), wb.outlier(r, gpr, j));
14624                    assert_eq!(
14625                        (a.0.to_bits(), a.1.to_bits()),
14626                        (b.0.to_bits(), b.1.to_bits())
14627                    );
14628                }
14629                #[cfg(target_arch = "x86_64")]
14630                if avx2_enabled() {
14631                    let sc = cortiq_core::quant::q1g_scale_offset(rows, cols);
14632                    // SAFETY: AVX2 checked; payloads sized for (rows, cols).
14633                    unsafe {
14634                        let a = dot_q1_row_avx2(&q1, r, gpr, &act.xq, &gsum);
14635                        let b = dot_q1g_row_avx2(&q1g, sc, r, gpr, &act.xq, &gsum);
14636                        assert_eq!(a.to_bits(), b.to_bits(), "avx2 row {r}");
14637                    }
14638                    if std::is_x86_feature_detected!("avx512vnni")
14639                        && std::is_x86_feature_detected!("avx512bw")
14640                        && std::is_x86_feature_detected!("avx512vl")
14641                    {
14642                        // SAFETY: VNNI family detected just above.
14643                        unsafe {
14644                            let a = dot_q1_row_vnni(&q1, r, gpr, &act.xq, &gsum);
14645                            let b = dot_q1g_row_vnni(&q1g, sc, r, gpr, &act.xq, &gsum);
14646                            assert_eq!(a.to_bits(), b.to_bits(), "vnni row {r}");
14647                        }
14648                    }
14649                }
14650            }
14651            // Drivers: single, pair, batch (b = 5: one blocked quad + one).
14652            let (mut ga, mut gb) = (vec![0f32; rows], vec![0f32; rows]);
14653            bin_matvec(wa, &x, rows, cols, &mut ga, None);
14654            bin_matvec(wb, &x, rows, cols, &mut gb, None);
14655            assert_eq!(ga, gb, "{rows}x{cols} matvec");
14656            let x2 = q1g_test_x(cols, rows + 1);
14657            let (mut a1, mut a2, mut b1, mut b2) = (
14658                vec![0f32; rows],
14659                vec![0f32; rows],
14660                vec![0f32; rows],
14661                vec![0f32; rows],
14662            );
14663            bin_matvec2(wa, &x, &x2, rows, cols, &mut a1, &mut a2, None);
14664            bin_matvec2(wb, &x, &x2, rows, cols, &mut b1, &mut b2, None);
14665            assert_eq!((&a1, &a2), (&b1, &b2), "{rows}x{cols} matvec2");
14666            let b = 5usize;
14667            let xs: Vec<f32> = (0..b).flat_map(|k| q1g_test_x(cols, k * 3 + 2)).collect();
14668            let (mut ma, mut mb) = (vec![0f32; b * rows], vec![0f32; b * rows]);
14669            bin_matmat(wa, &xs, b, rows, cols, &mut ma, None);
14670            bin_matmat(wb, &xs, b, rows, cols, &mut mb, None);
14671            assert_eq!(ma, mb, "{rows}x{cols} matmat");
14672            // The blocked 1×4 kernels directly (the driver only takes them
14673            // under CMF_X86_BLOCKED / SDOT, a process-wide switch).
14674            let acts: Vec<(SplitAct, Vec<i32>)> = (0..4)
14675                .map(|k| {
14676                    let a = split_act(&xs[k * cols..(k + 1) * cols]);
14677                    let g = q1_group_sums(&a.xq, gpr);
14678                    (a, g)
14679                })
14680                .collect();
14681            let xq4: [&[i8]; 4] = std::array::from_fn(|k| acts[k].0.xq.as_slice());
14682            let gs4: [&[i32]; 4] = std::array::from_fn(|k| acts[k].1.as_slice());
14683            for r in 0..rows {
14684                #[cfg(target_arch = "x86_64")]
14685                if avx2_enabled() {
14686                    // SAFETY: AVX2 checked; payloads sized for (rows, cols).
14687                    let (a, b) = unsafe {
14688                        (
14689                            wa.dot_1x4_x86(r, gpr, xq4, gs4),
14690                            wb.dot_1x4_x86(r, gpr, xq4, gs4),
14691                        )
14692                    };
14693                    assert_eq!(a.map(f32::to_bits), b.map(f32::to_bits), "1x4 row {r}");
14694                }
14695                #[cfg(target_arch = "aarch64")]
14696                if sdot_enabled() {
14697                    // SAFETY: SDOT checked; payloads sized for (rows, cols).
14698                    let (a, b) = unsafe {
14699                        (
14700                            wa.dot_1x4_sdot(r, gpr, xq4, gs4),
14701                            wb.dot_1x4_sdot(r, gpr, xq4, gs4),
14702                        )
14703                    };
14704                    assert_eq!(a.map(f32::to_bits), b.map(f32::to_bits), "1x4 row {r}");
14705                }
14706                let _ = (r, &xq4, &gs4);
14707            }
14708        }
14709    }
14710
14711    /// CPU kernel A/B, q1 vs q1g matvec (A8W8 path of this machine, the
14712    /// exact path too), Bonsai-1.7B decode shapes, arms alternated every
14713    /// round, medians printed. `CMF_THREADS` sets the pool (default 8).
14714    /// `cargo test -p cortiq-engine --release --lib q1g_cpu_matvec_speed --
14715    /// --ignored --nocapture`
14716    #[test]
14717    #[ignore]
14718    fn q1g_cpu_matvec_speed() {
14719        use crate::pool::Pool;
14720        let threads = std::env::var("CMF_THREADS")
14721            .ok()
14722            .and_then(|v| v.parse().ok())
14723            .unwrap_or(8usize);
14724        let pool = Pool::new(threads);
14725        for &(name, rows, cols, reps) in &[
14726            ("gate/up_proj", 6144usize, 2048usize, 41usize),
14727            ("down_proj", 2048, 6144, 41),
14728            ("lm_head", 151_669, 2048, 9),
14729        ] {
14730            let (q1, q1g) = q1_q1g_pair(rows, cols, 3);
14731            let (wa, wb) = (BinW::Q1(&q1), BinW::q1g(&q1g, rows, cols));
14732            let x = q1g_test_x(cols, 1);
14733            let mut out = vec![0f32; rows];
14734            let mut time = |w: BinW| {
14735                let t0 = std::time::Instant::now();
14736                bin_matvec(w, &x, rows, cols, &mut out, Some(&pool));
14737                t0.elapsed().as_secs_f64()
14738            };
14739            let _ = (time(wa), time(wb));
14740            let (mut ta, mut tb) = (Vec::new(), Vec::new());
14741            for r in 0..reps {
14742                if r % 2 == 0 {
14743                    ta.push(time(wa));
14744                    tb.push(time(wb));
14745                } else {
14746                    tb.push(time(wb));
14747                    ta.push(time(wa));
14748                }
14749            }
14750            ta.sort_by(|a, b| a.partial_cmp(b).unwrap());
14751            tb.sort_by(|a, b| a.partial_cmp(b).unwrap());
14752            let (ma, mb) = (ta[reps / 2], tb[reps / 2]);
14753            eprintln!(
14754                "cpu x{threads} a8w8={} {name:>12} {rows}x{cols}: q1 {:8.1} us | q1g {:8.1} us | q1g/q1 time {:.3}",
14755                a8w8_enabled(),
14756                ma * 1e6,
14757                mb * 1e6,
14758                mb / ma
14759            );
14760        }
14761    }
14762
14763    /// The mapped QTensor paths a real Bonsai q1g file takes — matvec, the
14764    /// multi-projection job (all-q1g and q1/q1g mixed), the fused gate/up
14765    /// with SiLU, matmat, matvec2 and the embedding row — equal the q1
14766    /// tensor they were repacked from, bit for bit.
14767    #[test]
14768    fn q1g_mapped_tensor_paths_equal_q1() {
14769        use crate::pool::Pool;
14770        use cortiq_core::{CmfHeader, CmfModel, TensorSpec};
14771        let (rows, cols) = (300usize, 256usize);
14772        let (gq1, gq1g) = q1_q1g_pair(rows, cols, 1);
14773        let (uq1, uq1g) = q1_q1g_pair(rows, cols, 2);
14774        let header: CmfHeader = serde_json::from_value(serde_json::json!({
14775            "format": "cmf", "version": cortiq_core::CMF_VERSION,
14776            "arch": {"arch_name": "tiny-q1g", "hidden_size": cols,
14777                "intermediate_size": rows, "num_layers": 1, "num_attention_heads": 2,
14778                "num_kv_heads": 1, "head_dim": 128, "vocab_size": rows,
14779                "layer_types": ["FullAttention"], "rms_norm_eps": 1e-6,
14780                "max_position_embeddings": 8},
14781            "quant_type": "VBIT"
14782        }))
14783        .unwrap();
14784        let spec = |name: &str, dtype, data: &Vec<u8>| TensorSpec {
14785            name: name.into(),
14786            dtype,
14787            shape: vec![rows, cols],
14788            data: data.clone(),
14789        };
14790        let specs = [
14791            spec("g1", TensorDtype::Q1, &gq1),
14792            spec("u1", TensorDtype::Q1, &uq1),
14793            spec("gg", TensorDtype::Q1G, &gq1g),
14794            spec("ug", TensorDtype::Q1G, &uq1g),
14795        ];
14796        let dir = std::env::temp_dir().join(format!("cmf-q1g-mapped-{}", std::process::id()));
14797        std::fs::create_dir_all(&dir).unwrap();
14798        let path = dir.join("m.cmf");
14799        CmfModel::write(&path, &header, &specs, None, None).unwrap();
14800        let model = Arc::new(CmfModel::open(&path).unwrap());
14801        assert!(model.verify().is_empty(), "{:?}", model.verify());
14802        let t = |n: &str| QTensor::from_model(&model, n).unwrap();
14803        let (g1, u1, gg, ug) = (t("g1"), t("u1"), t("gg"), t("ug"));
14804        assert_eq!(gg.model_dtype(), Some(TensorDtype::Q1G));
14805        assert!(gg.is_q1() && gg.graph_weight().map(|w| w.2) == Some(10));
14806        let pool = Pool::new(3);
14807        let x = q1g_test_x(cols, 4);
14808        let run = |a: &QTensor| {
14809            let mut o = vec![0f32; rows];
14810            a.matvec(&x, &mut o, Some(&pool));
14811            o
14812        };
14813        assert_eq!(run(&g1), run(&gg), "matvec");
14814        let many = |a: &QTensor, b: &QTensor| {
14815            let (mut oa, mut ob) = (vec![0f32; rows], vec![0f32; rows]);
14816            QTensor::matvec_many([a, b], &x, [&mut oa, &mut ob], Some(&pool));
14817            (oa, ob)
14818        };
14819        let want = many(&g1, &u1);
14820        assert_eq!(want, many(&gg, &ug), "matvec_many q1g");
14821        assert_eq!(want, many(&g1, &ug), "matvec_many mixed");
14822        assert_eq!(want, (run(&g1), run(&u1)), "matvec_many == singles");
14823        let fused = |g: &QTensor, u: &QTensor| {
14824            let mut o = vec![0f32; rows];
14825            let ok = QTensor::matvec_silu_mul(g, u, &x, &mut o, Some(&pool));
14826            (ok, o)
14827        };
14828        let (ok_a, fa) = fused(&g1, &u1);
14829        let (ok_b, fb) = fused(&gg, &ug);
14830        assert_eq!(ok_a, ok_b, "q1g takes the fused arm whenever q1 does");
14831        assert_eq!(fa, fb, "fused gate/up");
14832        let b = 6usize;
14833        let xs: Vec<f32> = (0..b).flat_map(|k| q1g_test_x(cols, k + 9)).collect();
14834        let mm = |a: &QTensor| {
14835            let mut o = vec![0f32; b * rows];
14836            a.matmat(&xs, b, &mut o, Some(&pool));
14837            o
14838        };
14839        assert_eq!(mm(&g1), mm(&gg), "matmat");
14840        let x2 = q1g_test_x(cols, 77);
14841        let pair = |a: &QTensor| {
14842            let (mut o1, mut o2) = (vec![0f32; rows], vec![0f32; rows]);
14843            a.matvec2(&x, &x2, &mut o1, &mut o2, Some(&pool));
14844            (o1, o2)
14845        };
14846        assert_eq!(pair(&g1), pair(&gg), "matvec2");
14847        for r in [0, 1, rows / 2, rows - 1] {
14848            let (mut ra, mut rb) = (vec![0f32; cols], vec![0f32; cols]);
14849            g1.row_f32(r, &mut ra);
14850            gg.row_f32(r, &mut rb);
14851            assert_eq!(ra, rb, "embedding row {r}");
14852        }
14853        drop((g1, u1, gg, ug, model));
14854        let _ = std::fs::remove_dir_all(&dir);
14855    }
14856
14857    /// A `blr` tensor read from a real CMF file: it maps to the factored
14858    /// arm (not a dense copy), declines every device-graph accessor, and
14859    /// every QTensor entry point a layer may take — matvec, matvec2,
14860    /// matmat (table loop and GEMM sizes), the multi-tensor job mixed with
14861    /// a q8_row sibling, the fused gate/up refusal, rows and columns — equals
14862    /// the dense reference built by the core dequant.
14863    #[test]
14864    fn blr_mapped_tensor_paths_equal_dense() {
14865        use crate::pool::Pool;
14866        use cortiq_core::{CmfHeader, CmfModel, TensorSpec};
14867        let (rows, cols, r) = (96usize, 128usize, 16usize);
14868        let mut rng = crate::blr::tests::Rng(41);
14869        let gate = crate::blr::tests::make(&mut rng, rows, cols, r, 2);
14870        let up = crate::blr::tests::make(&mut rng, rows, cols, r, 1);
14871        // q8_row sibling: [int8 rows·cols][f16 row scale · rows].
14872        let mut q8: Vec<u8> = (0..rows * cols)
14873            .map(|i| (((i * 89) % 253) as i32 - 126) as i8 as u8)
14874            .collect();
14875        for o in 0..rows {
14876            q8.extend_from_slice(&cortiq_core::quant::f32_to_f16(0.002 + o as f32 * 1e-4).to_le_bytes());
14877        }
14878        let header: CmfHeader = serde_json::from_value(serde_json::json!({
14879            "format": "cmf", "version": cortiq_core::CMF_VERSION,
14880            "arch": {"arch_name": "tiny-blr", "hidden_size": cols,
14881                "intermediate_size": rows, "num_layers": 1, "num_attention_heads": 2,
14882                "num_kv_heads": 1, "head_dim": 64, "vocab_size": rows,
14883                "layer_types": ["FullAttention"], "rms_norm_eps": 1e-6,
14884                "max_position_embeddings": 8},
14885            "quant_type": "Q8_2F"
14886        }))
14887        .unwrap();
14888        let specs = [
14889            TensorSpec {
14890                name: "g".into(),
14891                dtype: TensorDtype::BinaryLowRank,
14892                shape: vec![rows, cols],
14893                data: gate.clone(),
14894            },
14895            TensorSpec {
14896                name: "u".into(),
14897                dtype: TensorDtype::BinaryLowRank,
14898                shape: vec![rows, cols],
14899                data: up.clone(),
14900            },
14901            TensorSpec {
14902                name: "q".into(),
14903                dtype: TensorDtype::Q8Row,
14904                shape: vec![rows, cols],
14905                data: q8,
14906            },
14907        ];
14908        let dir = std::env::temp_dir().join(format!("cmf-blr-mapped-{}", std::process::id()));
14909        std::fs::create_dir_all(&dir).unwrap();
14910        let path = dir.join("m.cmf");
14911        CmfModel::write(&path, &header, &specs, None, None).unwrap();
14912        let model = Arc::new(CmfModel::open(&path).unwrap());
14913        assert!(model.verify().is_empty(), "{:?}", model.verify());
14914        let t = |n: &str| QTensor::from_model(&model, n).unwrap();
14915        let (g, u, q) = (t("g"), t("u"), t("q"));
14916        assert!(g.is_blr() && u.is_blr() && !q.is_blr());
14917        assert!(g.graph_weight().is_none() && g.graph_weight_descriptor().is_none());
14918        assert!(g.as_f32().is_none() && g.model_idx().is_none());
14919        assert!(g.mapped_device_gemm().is_none() && g.q1_parts().is_none());
14920        assert!(!g.sparse_col_ok());
14921        assert_eq!((g.rows(), g.cols()), (rows, cols));
14922        let deq = |bytes: &[u8]| {
14923            let mut w = vec![0f32; rows * cols];
14924            cortiq_core::blr::dequant_blr(bytes, rows, cols, &mut w).unwrap();
14925            w
14926        };
14927        let (wg, wu) = (deq(&gate), deq(&up));
14928        let mv = |w: &[f32], x: &[f32]| -> Vec<f32> {
14929            (0..rows)
14930                .map(|i| (0..cols).map(|j| w[i * cols + j] * x[j]).sum())
14931                .collect()
14932        };
14933        let close = |a: &[f32], b: &[f32], what: &str| {
14934            let s = b.iter().map(|v| v * v).sum::<f32>().sqrt().max(1e-30);
14935            let e = a.iter().zip(b).map(|(x, y)| (x - y) * (x - y)).sum::<f32>().sqrt();
14936            assert!(e / s < 1e-5, "{what}: rel {}", e / s);
14937        };
14938        let pool = Pool::new(3);
14939        let x: Vec<f32> = (0..cols).map(|_| rng.f()).collect();
14940        let x2: Vec<f32> = (0..cols).map(|_| rng.f()).collect();
14941        let mut o = vec![0f32; rows];
14942        g.matvec(&x, &mut o, Some(&pool));
14943        close(&o, &mv(&wg, &x), "matvec");
14944        let (mut o1, mut o2) = (vec![0f32; rows], vec![0f32; rows]);
14945        g.matvec2(&x, &x2, &mut o1, &mut o2, Some(&pool));
14946        close(&o1, &mv(&wg, &x), "matvec2 a");
14947        close(&o2, &mv(&wg, &x2), "matvec2 b");
14948        for b in [2usize, 5] {
14949            let xs: Vec<f32> = (0..b * cols).map(|_| rng.f()).collect();
14950            let mut om = vec![0f32; b * rows];
14951            g.matmat(&xs, b, &mut om, Some(&pool));
14952            for t in 0..b {
14953                close(
14954                    &om[t * rows..(t + 1) * rows],
14955                    &mv(&wg, &xs[t * cols..(t + 1) * cols]),
14956                    "matmat",
14957                );
14958            }
14959        }
14960        // Multi-tensor job mixed with a q8_row sibling: per-tensor fallback,
14961        // the q8_row output identical to its own matvec.
14962        let (mut a, mut bq, mut c) = (vec![0f32; rows], vec![0f32; rows], vec![0f32; rows]);
14963        QTensor::matvec_many([&g, &q, &u], &x, [&mut a, &mut bq, &mut c], Some(&pool));
14964        close(&a, &mv(&wg, &x), "matvec_many blr gate");
14965        close(&c, &mv(&wu, &x), "matvec_many blr up");
14966        let mut qo = vec![0f32; rows];
14967        q.matvec(&x, &mut qo, Some(&pool));
14968        assert_eq!(bq, qo, "matvec_many q8_row sibling");
14969        let mut f = vec![0f32; rows];
14970        assert!(!QTensor::matvec_silu_mul(&g, &u, &x, &mut f, Some(&pool)));
14971        let mut row = vec![0f32; cols];
14972        g.row_f32(rows / 3, &mut row);
14973        close(&row, &wg[(rows / 3) * cols..(rows / 3 + 1) * cols], "row_f32");
14974        let mut scratch = vec![0f32; cols];
14975        let d = g.row_dot(7, &x, &mut scratch);
14976        let want: f32 = (0..cols).map(|j| wg[7 * cols + j] * x[j]).sum();
14977        assert!((d - want).abs() < 1e-4 * (1.0 + want.abs()), "row_dot");
14978        let mut col = vec![0f32; rows];
14979        g.add_col_scaled(cols / 2, 1.5, &mut col);
14980        let want: Vec<f32> = (0..rows).map(|i| 1.5 * wg[i * cols + cols / 2]).collect();
14981        close(&col, &want, "add_col_scaled");
14982        drop((g, u, q, model));
14983        let _ = std::fs::remove_dir_all(&dir);
14984    }
14985
14986    #[test]
14987    fn repack_is_bit_identical() {
14988        // The interleaved-repack kernel must produce EXACTLY the same
14989        // bits as the mmap-layout kernel: integer accumulation is order-
14990        // exact, the f32 epilogue is identical. Odd rows exercise the
14991        // tail; direct range calls exercise unaligned pool splits.
14992        let (rows, cols) = (267, 96); // 66 groups + 3 tail rows, cols % 16 == 0
14993        let w: Vec<u8> = (0..rows * cols)
14994            .map(|i| (((i * 89) % 253) as i32 - 126) as i8 as u8)
14995            .collect();
14996        let scales: Vec<f32> = (0..rows).map(|o| 0.003 + o as f32 * 0.0007).collect();
14997        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.37).sin() * 2.0).collect();
14998        let rep = q8_repack_layout(&w, rows, cols);
14999        // Group interleave round-trips.
15000        for g in 0..rows / 4 {
15001            for c in 0..cols / 16 {
15002                for lane in 0..4 {
15003                    assert_eq!(
15004                        &rep[g * 4 * cols + c * 64 + lane * 16
15005                            ..g * 4 * cols + c * 64 + lane * 16 + 16],
15006                        &w[(g * 4 + lane) * cols + c * 16..(g * 4 + lane) * cols + c * 16 + 16],
15007                    );
15008                }
15009            }
15010        }
15011        let mut a = vec![0.0f32; rows];
15012        qmatvec(
15013            &w,
15014            &[],
15015            &scales,
15016            &x,
15017            &[],
15018            TensorDtype::Q8Row,
15019            rows,
15020            cols,
15021            &mut a,
15022            None,
15023        );
15024        let mut b = vec![0.0f32; rows];
15025        qmatvec(
15026            &w,
15027            &rep,
15028            &scales,
15029            &x,
15030            &[],
15031            TensorDtype::Q8Row,
15032            rows,
15033            cols,
15034            &mut b,
15035            None,
15036        );
15037        assert_eq!(a, b, "full-range repack output diverged");
15038
15039        #[cfg(target_arch = "aarch64")]
15040        if sdot_enabled() {
15041            // Unaligned range split (pool workers get arbitrary bounds).
15042            let act = split_act(&x);
15043            let mut c1 = vec![0.0f32; rows];
15044            let mut c2 = vec![0.0f32; rows];
15045            q8_range_sdot(
15046                &w,
15047                &[],
15048                &scales,
15049                &act,
15050                cols,
15051                SendMut(c1.as_mut_ptr()),
15052                3,
15053                rows - 2,
15054            );
15055            q8_range_sdot(
15056                &w,
15057                &rep,
15058                &scales,
15059                &act,
15060                cols,
15061                SendMut(c2.as_mut_ptr()),
15062                3,
15063                rows - 2,
15064            );
15065            assert_eq!(c1, c2, "unaligned-range repack output diverged");
15066        }
15067    }
15068
15069    #[test]
15070    fn sdot_a8w8_noise_is_bounded() {
15071        // Off-grid activations: A8 quantization noise must stay small in
15072        // relative L2 over the whole output (realistic accuracy contract;
15073        // vmfcore measured argmax-identical decode on real models).
15074        let (rows, cols) = (16, 512);
15075        let w: Vec<u8> = (0..rows * cols)
15076            .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
15077            .collect();
15078        let scales = vec![0.01f32; rows];
15079        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.21).sin()).collect();
15080        let mut a = vec![0.0f32; rows];
15081        qmatvec(
15082            &w,
15083            &[],
15084            &scales,
15085            &x,
15086            &[],
15087            TensorDtype::Q8Row,
15088            rows,
15089            cols,
15090            &mut a,
15091            None,
15092        );
15093        let (mut num, mut den) = (0f64, 0f64);
15094        for o in 0..rows {
15095            let mut acc = 0.0f32;
15096            for j in 0..cols {
15097                acc += (w[o * cols + j] as i8) as f32 * x[j];
15098            }
15099            let expect = acc * scales[o];
15100            num += ((a[o] - expect) as f64).powi(2);
15101            den += (expect as f64).powi(2);
15102        }
15103        let rel = (num / den.max(1e-12)).sqrt();
15104        assert!(rel < 0.05, "A8W8 relative L2 error too high: {rel}");
15105    }
15106
15107    #[test]
15108    fn i8_dot_neon_matches_scalar() {
15109        let n = 100;
15110        let w: Vec<u8> = (0..n).map(|i| ((i * 37 + 11) % 251) as u8).collect();
15111        let x: Vec<f32> = (0..n).map(|i| (i as f32 * 0.13).sin()).collect();
15112        let mut scalar = 0.0f32;
15113        for j in 0..n {
15114            scalar += (w[j] as i8) as f32 * x[j];
15115        }
15116        let fast = dot_i8_f32(&w, &x);
15117        assert!((scalar - fast).abs() < 1e-3 * scalar.abs().max(1.0));
15118    }
15119
15120    /// Fused vbit matvec must match full dequant_vbit + dense matvec.
15121    #[test]
15122    fn vbitmatvec_matches_full_dequant() {
15123        let (rows, cols) = (6, 64);
15124        let ng = cols / GROUP_SIZE;
15125        // Hand-craft: bits per row, f16 scales, packed rows.
15126        let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
15127        let mut bytes = bits.clone();
15128        for g in 0..rows * ng {
15129            let s = 0.02 + 0.001 * g as f32;
15130            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
15131        }
15132        for r in 0..rows {
15133            let b = bits[r] as usize;
15134            let (mut acc, mut nb) = (0u64, 0usize);
15135            let mut rowbytes = Vec::new();
15136            for i in 0..cols {
15137                let v = ((i * 7 + r * 13) % (1 << b)) as u64;
15138                acc = (acc << b) | v;
15139                nb += b;
15140                while nb >= 8 {
15141                    nb -= 8;
15142                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
15143                }
15144            }
15145            if nb > 0 {
15146                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
15147            }
15148            bytes.extend_from_slice(&rowbytes);
15149        }
15150        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
15151
15152        let mut reference = vec![0f32; rows * cols];
15153        cortiq_core::quant::dequant_vbit(&bytes, rows, cols, &mut reference).unwrap();
15154        let mut expect = vec![0f32; rows];
15155        for r in 0..rows {
15156            expect[r] = reference[r * cols..(r + 1) * cols]
15157                .iter()
15158                .zip(&x)
15159                .map(|(w, xv)| w * xv)
15160                .sum();
15161        }
15162        let mut got = vec![0f32; rows];
15163        let offsets = vbit_row_offsets(&bytes, rows, cols);
15164        vbitmatvec(&bytes, &offsets, &x, rows, cols, &mut got, None);
15165        // SDOT path quantizes activations to i8 (A8W8): bounded noise,
15166        // same contract as q8 (exact path is pinned by CMF_SDOT=0 in
15167        // the golden-parity gate).
15168        let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
15169        let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
15170        for r in 0..rows {
15171            assert!(
15172                (got[r] - expect[r]).abs() < tol * scale,
15173                "row {r}: {} vs {}",
15174                got[r],
15175                expect[r]
15176            );
15177        }
15178    }
15179
15180    /// Fused q4 matvec must match the reference full-dequant + dense
15181    /// matvec bit-for-bit in structure (same f32 math, group order).
15182    /// vbit matmat: the blocked 1×4 leg must match the per-row path
15183    /// (paired env toggle; larger shape so both code paths engage).
15184    #[test]
15185    #[cfg(target_arch = "x86_64")]
15186    fn vbit_matmat_blocked_matches_per_row() {
15187        let (rows, cols, b) = (64usize, 128usize, 9usize);
15188        let ng = cols / GROUP_SIZE;
15189        let bits: Vec<u8> = (0..rows).map(|r| [3u8, 4, 5, 6][r % 4]).collect();
15190        let mut bytes = bits.clone();
15191        for g in 0..rows * ng {
15192            let sc = 0.02 + 0.0005 * g as f32;
15193            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
15194        }
15195        for r in 0..rows {
15196            let bw = bits[r] as usize;
15197            let (mut acc, mut nb) = (0u64, 0usize);
15198            let mut rowbytes = Vec::new();
15199            for i in 0..cols {
15200                let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
15201                acc = (acc << bw) | v;
15202                nb += bw;
15203                while nb >= 8 {
15204                    nb -= 8;
15205                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
15206                }
15207            }
15208            if nb > 0 {
15209                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
15210            }
15211            bytes.extend_from_slice(&rowbytes);
15212        }
15213        let x: Vec<f32> = (0..b * cols)
15214            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
15215            .collect();
15216        let offsets = vbit_row_offsets(&bytes, rows, cols);
15217        let mut y_a = vec![0f32; b * rows];
15218        let mut y_b = vec![0f32; b * rows];
15219        unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
15220        vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_a, None);
15221        unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
15222        vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_b, None);
15223        unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
15224        let max_d = y_a
15225            .iter()
15226            .zip(&y_b)
15227            .map(|(p, q)| (p - q).abs())
15228            .fold(0.0f32, f32::max);
15229        assert!(max_d < 1e-4, "vbit blocked ≠ per-row: max|Δ| = {max_d}");
15230    }
15231
15232    /// q4t blocked 1×4 (SDOT on ARM, AVX2 on x86) must equal the
15233    /// per-row path exactly: same nibble unpack, same group order,
15234    /// same f32 accumulation — batch == matvec bit-for-bit. b=9 covers
15235    /// two full 1×4 blocks plus a remainder through the single-row
15236    /// kernel. (Both paths produce identical output, so the shared
15237    /// CMF_X86_BLOCKED env var racing with other tests cannot flip
15238    /// the verdict — worst case both sides take the same path.)
15239    #[test]
15240    fn q4t_matmat_blocked_matches_per_row() {
15241        let (rows, cols, b) = (16usize, 64usize, 9usize);
15242        let gpr = cols / GROUP_SIZE;
15243        let mut bytes = vec![0u8; rows * gpr * Q4_TILE];
15244        for r in 0..rows {
15245            for g in 0..gpr {
15246                let t = (r * gpr + g) * Q4_TILE;
15247                let sc = 0.02 + 0.001 * (r * gpr + g) as f32;
15248                bytes[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
15249                for k in 0..16 {
15250                    bytes[t + 2 + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
15251                }
15252            }
15253        }
15254        let x: Vec<f32> = (0..b * cols)
15255            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
15256            .collect();
15257        let mut y_blk = vec![0f32; b * rows];
15258        let mut y_row = vec![0f32; b * rows];
15259        unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
15260        q4t_matmat(&bytes, &x, b, rows, cols, &mut y_blk, None);
15261        unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
15262        q4t_matmat(&bytes, &x, b, rows, cols, &mut y_row, None);
15263        unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
15264        assert_eq!(y_blk, y_row, "q4t blocked 1x4 ≠ per-row");
15265    }
15266
15267    /// The wide-batch Accelerate arm of q4t_matmat vs a brute-force
15268    /// f32 dequant matmul: both are f32 GEMMs, so only reduction
15269    /// order differs — tight tolerance.
15270    /// A synthetic q4tp payload: random nibbles plus a per-row ladder whose
15271    /// span varies row to row, so the codes actually exercise the full 0..31
15272    /// range rather than clustering on one rung.
15273    fn synth_q4tp(rows: usize, cols: usize) -> Vec<u8> {
15274        use cortiq_core::quant::{f32_to_f16, q4tp_code_stride, q4tp_put_code};
15275        let gpr = cols / GROUP_SIZE;
15276        let stride = q4tp_code_stride(gpr);
15277        let (params_off, codes_off, _) = q4tp_sections(rows, cols);
15278        let mut b = vec![0u8; codes_off + rows * stride];
15279        for r in 0..rows {
15280            for g in 0..gpr {
15281                let t = (r * gpr + g) * Q4TP_NIB;
15282                for k in 0..16 {
15283                    b[t + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
15284                }
15285            }
15286            let lo = -6.0 - 0.03 * (r % 17) as f32;
15287            let step = 0.01 + 0.004 * (r % 11) as f32;
15288            let p = params_off + r * 4;
15289            b[p..p + 2].copy_from_slice(&f32_to_f16(lo).to_le_bytes());
15290            b[p + 2..p + 4].copy_from_slice(&f32_to_f16(step).to_le_bytes());
15291            let crow = &mut b[codes_off + r * stride..codes_off + (r + 1) * stride];
15292            for g in 0..gpr {
15293                q4tp_put_code(crow, g, (r * 5 + g * 3) % 32);
15294            }
15295        }
15296        b
15297    }
15298
15299    /// The same weights re-expressed as q4_tiled, so the proven kernel can
15300    /// be the reference: each tile stores the ladder scale its code selects.
15301    /// Only the f16 rounding of that scale separates the two payloads.
15302    fn q4tp_as_q4t(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
15303        let gpr = cols / GROUP_SIZE;
15304        let v = Q4tpView::new(bytes, rows, cols);
15305        let mut out = vec![0u8; rows * gpr * Q4_TILE];
15306        let mut sc = vec![0f32; gpr];
15307        for r in 0..rows {
15308            v.scales_into(r, gpr, &mut sc);
15309            for g in 0..gpr {
15310                let t = (r * gpr + g) * Q4_TILE;
15311                let s = sc[g];
15312                out[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
15313                let src = (r * gpr + g) * Q4TP_NIB;
15314                out[t + 2..t + Q4_TILE].copy_from_slice(&v.nib[src..src + Q4TP_NIB]);
15315            }
15316        }
15317        out
15318    }
15319
15320    /// The exact (`CMF_SDOT=0`) path must reproduce `dequant_q4tp` to f32
15321    /// rounding — that scalar routine is the format's definition, and the
15322    /// kernels re-derive the scale from the ladder independently. Call the
15323    /// row kernel directly: `matmat` picks the int8 arm when a8w8 is on,
15324    /// so routing through it would test the other path by accident.
15325    #[test]
15326    fn q4tp_exact_path_matches_dequant_reference() {
15327        let (rows, cols) = (256usize, 512usize);
15328        let gpr = cols / GROUP_SIZE;
15329        let bytes = synth_q4tp(rows, cols);
15330        let mut w = vec![0f32; rows * cols];
15331        cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
15332
15333        let x: Vec<f32> = (0..cols)
15334            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
15335            .collect();
15336        let v = Q4tpView::new(&bytes, rows, cols);
15337        let mut sc = vec![0f32; gpr];
15338        for r in 0..rows {
15339            v.scales_into(r, gpr, &mut sc);
15340            let got = q4tp_row_exact(v.nib, r, gpr, &x, &sc);
15341            let want: f32 = (0..cols).map(|c| w[r * cols + c] * x[c]).sum();
15342            // These dot products cancel down to ~1e-3 from terms of ~5e-2, so
15343            // the meaningful yardstick is the summed magnitude, not the result:
15344            // against the result any reordering of a 512-term f32 sum "fails".
15345            let mag: f32 = (0..cols).map(|c| (w[r * cols + c] * x[c]).abs()).sum();
15346            assert!(
15347                (got - want).abs() <= 1e-5 * mag,
15348                "row {r}: kernel {got} vs dequant {want}"
15349            );
15350        }
15351    }
15352
15353    /// The int8 (a8w8) path can't be checked against an f32 reference — the
15354    /// activation quantization dominates. Check it against the q4t kernel it
15355    /// was ported from instead, on payloads holding the same weights: that
15356    /// isolates exactly what the port could break (16 B stride, ladder
15357    /// lookup, nibble unpack) from what it deliberately shares.
15358    #[test]
15359    fn q4tp_matvec_matches_the_q4t_kernel_it_was_ported_from() {
15360        let (rows, cols) = (256usize, 512usize);
15361        let bytes = synth_q4tp(rows, cols);
15362        let twin = q4tp_as_q4t(&bytes, rows, cols);
15363        let x: Vec<f32> = (0..cols)
15364            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
15365            .collect();
15366
15367        let mut got = vec![0f32; rows];
15368        q4tp_matvec(&bytes, &x, rows, cols, &mut got, None);
15369        let mut want = vec![0f32; rows];
15370        q4t_matvec(&twin, &x, rows, cols, &mut want, None);
15371
15372        // Scale is f16 in the twin and f32 here, so allow that rounding on
15373        // top of the summed magnitude (same cancellation argument as above).
15374        let mut w = vec![0f32; rows * cols];
15375        cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
15376        for r in 0..rows {
15377            let mag: f32 = (0..cols).map(|c| (w[r * cols + c] * x[c]).abs()).sum();
15378            assert!(
15379                (got[r] - want[r]).abs() <= 1e-3 * mag,
15380                "row {r}: q4tp {} vs q4t {}",
15381                got[r],
15382                want[r]
15383            );
15384        }
15385    }
15386
15387    /// `matmat` carries three arms (Accelerate, blocked int8 1x4, scalar).
15388    /// Batch 5 crosses the blocked kernel's stride, so this exercises the
15389    /// 1x4 path AND its scalar tail in one run — the blocked kernel is new
15390    /// code and its four accumulators are exactly what tends to go wrong.
15391    #[test]
15392    fn q4tp_matmat_matches_the_q4t_kernel_it_was_ported_from() {
15393        let (rows, cols, b) = (256usize, 512usize, 5usize);
15394        let bytes = synth_q4tp(rows, cols);
15395        let twin = q4tp_as_q4t(&bytes, rows, cols);
15396        let xs: Vec<f32> = (0..b * cols)
15397            .map(|i| ((i * 29 + 11) % 89) as f32 / 89.0 - 0.5)
15398            .collect();
15399
15400        let mut got = vec![0f32; b * rows];
15401        q4tp_matmat(&bytes, &xs, b, rows, cols, &mut got, None);
15402        let mut want = vec![0f32; b * rows];
15403        q4t_matmat(&twin, &xs, b, rows, cols, &mut want, None);
15404
15405        let mut w = vec![0f32; rows * cols];
15406        cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
15407        for t in 0..b {
15408            for r in 0..rows {
15409                let mag: f32 = (0..cols)
15410                    .map(|c| (w[r * cols + c] * xs[t * cols + c]).abs())
15411                    .sum();
15412                let (g, wa) = (got[t * rows + r], want[t * rows + r]);
15413                assert!(
15414                    (g - wa).abs() <= 1e-3 * mag,
15415                    "batch {t} row {r}: q4tp {g} vs q4t {wa}"
15416                );
15417            }
15418        }
15419    }
15420
15421    #[test]
15422    fn q4tp_matvec2_matches_the_single_stream_kernel() {
15423        let (rows, cols) = (128usize, 256usize);
15424        let gpr = cols / GROUP_SIZE;
15425        let bytes = synth_q4tp(rows, cols);
15426        let xs: Vec<f32> = (0..2 * cols)
15427            .map(|i| ((i * 29 + 11) % 89) as f32 / 89.0 - 0.5)
15428            .collect();
15429
15430        let (mut o1, mut o2) = (vec![0f32; rows], vec![0f32; rows]);
15431        q4tp_matvec2(
15432            &bytes,
15433            &xs[..cols],
15434            &xs[cols..],
15435            rows,
15436            cols,
15437            &mut o1,
15438            &mut o2,
15439            None,
15440        );
15441
15442        // matvec2 takes the exact path for both streams, so the single-row
15443        // kernel is an exact reference — no tolerance for path differences.
15444        let v = Q4tpView::new(&bytes, rows, cols);
15445        let mut sc = vec![0f32; gpr];
15446        for r in 0..rows {
15447            v.scales_into(r, gpr, &mut sc);
15448            assert_eq!(o1[r], q4tp_row_exact(v.nib, r, gpr, &xs[..cols], &sc));
15449            assert_eq!(o2[r], q4tp_row_exact(v.nib, r, gpr, &xs[cols..], &sc));
15450        }
15451    }
15452
15453    /// q4tp must not COST speed — it exists to save bytes, and a format that
15454    /// trades 7% of a file for a slower model is a bad trade. This guard is
15455    /// here because correctness tests happily passed while `q4tp_matmat` was
15456    /// missing its int8 and Accelerate arms and the model ran 5x slower.
15457    /// Measured on M-series: 0.97-1.04x, i.e. parity (16 B tiles are better
15458    /// aligned than q4t's 18 B, which pays for the scale indirection).
15459    #[test]
15460    fn q4tp_matvec_keeps_pace_with_q4t() {
15461        let (rows, cols) = (4096usize, 3072usize);
15462        let bytes = synth_q4tp(rows, cols);
15463        let twin = q4tp_as_q4t(&bytes, rows, cols);
15464        let x: Vec<f32> = (0..cols).map(|i| (i % 97) as f32 / 97.0 - 0.5).collect();
15465        let mut o = vec![0f32; rows];
15466        let n = 12;
15467        let mut best = (f64::MAX, f64::MAX);
15468        // Interleaved A/B, minimum statistic: this machine throttles, and a
15469        // mean over a thermal ramp reliably indicts whichever ran second.
15470        for _ in 0..3 {
15471            let t0 = std::time::Instant::now();
15472            for _ in 0..n {
15473                q4t_matvec(&twin, &x, rows, cols, &mut o, None);
15474            }
15475            best.0 = best.0.min(t0.elapsed().as_secs_f64());
15476            let t0 = std::time::Instant::now();
15477            for _ in 0..n {
15478                q4tp_matvec(&bytes, &x, rows, cols, &mut o, None);
15479            }
15480            best.1 = best.1.min(t0.elapsed().as_secs_f64());
15481        }
15482        let ratio = best.1 / best.0;
15483        println!(
15484            "q4t {:.3} ms | q4tp {:.3} ms | {ratio:.2}x",
15485            best.0 * 1e3 / n as f64,
15486            best.1 * 1e3 / n as f64
15487        );
15488        assert!(ratio < 2.0, "q4tp matvec {ratio:.2}x slower than q4t");
15489    }
15490
15491    #[cfg(target_os = "macos")]
15492    #[test]
15493    fn q4t_matmat_accel_matches_dequant_reference() {
15494        if !accel_gemm_enabled() {
15495            return; // CMF_ACCEL=0
15496        }
15497        let (rows, cols, b) = (512usize, 1024usize, 8usize); // ≥500K → accel arm
15498        let gpr = cols / GROUP_SIZE;
15499        let mut bytes = vec![0u8; rows * gpr * Q4_TILE];
15500        for r in 0..rows {
15501            for g in 0..gpr {
15502                let t = (r * gpr + g) * Q4_TILE;
15503                let sc = 0.02 + 0.0005 * ((r * gpr + g) % 64) as f32;
15504                bytes[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
15505                for k in 0..16 {
15506                    bytes[t + 2 + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
15507                }
15508            }
15509        }
15510        let x: Vec<f32> = (0..b * cols)
15511            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
15512            .collect();
15513        let mut got = vec![0f32; b * rows];
15514        q4t_matmat(&bytes, &x, b, rows, cols, &mut got, None);
15515        // Brute-force reference off the same tiles.
15516        let mut w = vec![0f32; rows * cols];
15517        for r in 0..rows {
15518            for g in 0..gpr {
15519                let t = (r * gpr + g) * Q4_TILE;
15520                let s = f16_to_f32(u16::from_le_bytes([bytes[t], bytes[t + 1]]));
15521                for (k, &bb) in bytes[t + 2..t + Q4_TILE].iter().enumerate() {
15522                    w[r * cols + g * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s;
15523                    w[r * cols + g * GROUP_SIZE + k * 2 + 1] =
15524                        (((bb >> 4) & 0x0F) as f32 - 8.0) * s;
15525                }
15526            }
15527        }
15528        for bi in 0..b {
15529            for r in 0..rows {
15530                let want: f32 = (0..cols).map(|j| x[bi * cols + j] * w[r * cols + j]).sum();
15531                let d = (got[bi * rows + r] - want).abs();
15532                assert!(
15533                    d <= want.abs().max(1.0) * 1e-4,
15534                    "accel q4t GEMM diverged at ({bi},{r}): {} vs {want}",
15535                    got[bi * rows + r]
15536                );
15537            }
15538        }
15539    }
15540
15541    #[test]
15542    fn q4matvec_matches_full_dequant() {
15543        let (rows, cols) = (8, 64);
15544        let groups = rows * cols / GROUP_SIZE;
15545        // Hand-craft a q4_block blob: nibbles then f16 scales.
15546        let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
15547        for i in 0..groups * 16 {
15548            bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
15549        }
15550        for g in 0..groups {
15551            let s = 0.01 + 0.003 * g as f32;
15552            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
15553        }
15554        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
15555
15556        let mut reference = vec![0.0f32; rows * cols];
15557        cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
15558        let mut expect = vec![0.0f32; rows];
15559        for r in 0..rows {
15560            expect[r] = reference[r * cols..(r + 1) * cols]
15561                .iter()
15562                .zip(&x)
15563                .map(|(w, xv)| w * xv)
15564                .sum();
15565        }
15566
15567        let mut got = vec![0.0f32; rows];
15568        q4matvec(&bytes, &x, rows, cols, &mut got, None);
15569        // SDOT path quantizes activations to i8 (A8W8): bounded noise,
15570        // same contract as q8/vbit (exact path is pinned by CMF_SDOT=0
15571        // in the golden-parity gate).
15572        let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
15573        let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
15574        for r in 0..rows {
15575            assert!(
15576                (got[r] - expect[r]).abs() < tol * scale,
15577                "row {r}: {} vs {}",
15578                got[r],
15579                expect[r]
15580            );
15581        }
15582    }
15583
15584    /// Fused two-input vbit matvec must equal two single matvecs exactly
15585    /// (same per-lane accumulation order on both scalar and SDOT paths).
15586    #[test]
15587    fn vbitmatvec2_equals_two_singles() {
15588        let (rows, cols) = (6, 64);
15589        let ng = cols / GROUP_SIZE;
15590        let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
15591        let mut bytes = bits.clone();
15592        for g in 0..rows * ng {
15593            let s = 0.02 + 0.001 * g as f32;
15594            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
15595        }
15596        for r in 0..rows {
15597            let b = bits[r] as usize;
15598            let (mut acc, mut nb) = (0u64, 0usize);
15599            let mut rowbytes = Vec::new();
15600            for i in 0..cols {
15601                let v = ((i * 7 + r * 13) % (1 << b)) as u64;
15602                acc = (acc << b) | v;
15603                nb += b;
15604                while nb >= 8 {
15605                    nb -= 8;
15606                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
15607                }
15608            }
15609            if nb > 0 {
15610                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
15611            }
15612            bytes.extend_from_slice(&rowbytes);
15613        }
15614        let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
15615        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).cos()).collect();
15616        let offsets = vbit_row_offsets(&bytes, rows, cols);
15617
15618        let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
15619        vbitmatvec(&bytes, &offsets, &x1, rows, cols, &mut a1, None);
15620        vbitmatvec(&bytes, &offsets, &x2, rows, cols, &mut a2, None);
15621        let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
15622        vbitmatvec2(
15623            &bytes, &offsets, &x1, &x2, rows, cols, &mut b1, &mut b2, None,
15624        );
15625        assert_eq!(a1, b1, "fused vbit lane 1 must be bit-identical");
15626        assert_eq!(a2, b2, "fused vbit lane 2 must be bit-identical");
15627    }
15628
15629    /// Fused two-input q4 matvec must equal two single matvecs exactly.
15630    #[test]
15631    fn q4matvec2_equals_two_singles() {
15632        let (rows, cols) = (8, 128);
15633        let groups = rows * cols / GROUP_SIZE;
15634        let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
15635        for i in 0..groups * 16 {
15636            bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
15637        }
15638        for g in 0..groups {
15639            let s = 0.01 + 0.003 * g as f32;
15640            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
15641        }
15642        // Include an outlier channel so the SDOT correction path is
15643        // exercised in the pair kernel too.
15644        let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
15645        x1[9] = 250.0;
15646        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
15647
15648        let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
15649        q4matvec(&bytes, &x1, rows, cols, &mut a1, None);
15650        q4matvec(&bytes, &x2, rows, cols, &mut a2, None);
15651        let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
15652        q4matvec2(&bytes, &x1, &x2, rows, cols, &mut b1, &mut b2, None);
15653        assert_eq!(a1, b1, "fused q4 lane 1 must be bit-identical");
15654        assert_eq!(a2, b2, "fused q4 lane 2 must be bit-identical");
15655    }
15656
15657    /// Multi-matrix job must equal separate matvecs exactly — same
15658    /// kernels, only the dispatch is fused.
15659    #[test]
15660    fn matvec_many_equals_separate_matvecs() {
15661        use crate::pool::Pool;
15662        let (r1, r2, cols) = (300, 200, 64);
15663        let mk = |salt: usize, rows: usize| {
15664            QTensor::from_f32(
15665                (0..rows * cols)
15666                    .map(|i| ((i * 7 + salt) % 97) as f32 / 97.0 - 0.5)
15667                    .collect(),
15668                rows,
15669                cols,
15670            )
15671        };
15672        let (a, b) = (mk(1, r1), mk(5, r2));
15673        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).sin()).collect();
15674        let pool = Pool::new(3);
15675
15676        let (mut ea, mut eb) = (vec![0f32; r1], vec![0f32; r2]);
15677        a.matvec(&x, &mut ea, Some(&pool));
15678        b.matvec(&x, &mut eb, Some(&pool));
15679        let (mut ga, mut gb) = (vec![0f32; r1], vec![0f32; r2]);
15680        QTensor::matvec_many([&a, &b], &x, [&mut ga, &mut gb], Some(&pool));
15681        assert_eq!(ea, ga, "fused multi-matrix lane 1 must be bit-identical");
15682        assert_eq!(eb, gb, "fused multi-matrix lane 2 must be bit-identical");
15683    }
15684
15685    /// The public Q4TP operator must take the real mapped matvec_many arm,
15686    /// rather than the F32 fallback above.  Build a tiny valid CMF so both
15687    /// handles retain their mmap payloads, then compare the fused dispatch
15688    /// with two ordinary mapped matvec calls bit-for-bit.
15689    #[test]
15690    fn q4tp_matvec_many_equals_separate_matvecs() {
15691        use crate::pool::Pool;
15692        use cortiq_core::{CMF_VERSION, CmfHeader, CmfModel, QuantType, TensorSpec};
15693
15694        let (r1, r2, cols) = (300usize, 200usize, 64usize);
15695        let arch: cortiq_core::ModelArch = serde_json::from_value(serde_json::json!({
15696            "arch_name": "tiny-q4tp",
15697            "hidden_size": cols,
15698            "intermediate_size": cols * 2,
15699            "num_layers": 1,
15700            "num_attention_heads": 2,
15701            "num_kv_heads": 1,
15702            "head_dim": 32,
15703            "vocab_size": r1,
15704            "layer_types": ["FullAttention"],
15705            "rms_norm_eps": 1e-6,
15706            "max_position_embeddings": 8,
15707            "linear_conv_kernel_dim": 0,
15708            "linear_num_key_heads": 0,
15709            "linear_num_value_heads": 0
15710        }))
15711        .unwrap();
15712        let header = CmfHeader {
15713            format: "cmf".into(),
15714            version: CMF_VERSION,
15715            arch,
15716            quant_type: QuantType::Q4Block,
15717            provenance: None,
15718            tokenizer_config: None,
15719            section_hashes: None,
15720            skills: Vec::new(),
15721            shard: None,
15722            calibration: None,
15723            routing: None,
15724            genome: None,
15725            lineage: Vec::new(),
15726            router: None,
15727            segments: Vec::new(),
15728        };
15729        let specs = [
15730            TensorSpec {
15731                name: "q".into(),
15732                dtype: TensorDtype::Q4TiledP,
15733                shape: vec![r1, cols],
15734                data: synth_q4tp(r1, cols),
15735            },
15736            TensorSpec {
15737                name: "kv".into(),
15738                dtype: TensorDtype::Q4TiledP,
15739                shape: vec![r2, cols],
15740                data: synth_q4tp(r2, cols),
15741            },
15742        ];
15743        let dir = std::env::temp_dir().join(format!("cmf-q4tp-many-{}", std::process::id()));
15744        std::fs::create_dir_all(&dir).unwrap();
15745        let path = dir.join("m.cmf");
15746        CmfModel::write(&path, &header, &specs, None, None).unwrap();
15747        let model = Arc::new(CmfModel::open(&path).unwrap());
15748        let (a, b) = (
15749            QTensor::from_model(&model, "q").unwrap(),
15750            QTensor::from_model(&model, "kv").unwrap(),
15751        );
15752        assert_eq!(a.model_dtype(), Some(TensorDtype::Q4TiledP));
15753        assert_eq!(b.model_dtype(), Some(TensorDtype::Q4TiledP));
15754        let x: Vec<f32> = (0..cols)
15755            .map(|i| ((i * 17 + 3) % 97) as f32 / 97.0 - 0.5)
15756            .collect();
15757        let pool = Pool::new(3);
15758        let (mut ea, mut eb) = (vec![0.0f32; r1], vec![0.0f32; r2]);
15759        a.matvec(&x, &mut ea, Some(&pool));
15760        b.matvec(&x, &mut eb, Some(&pool));
15761        let (mut ga, mut gb) = (vec![0.0f32; r1], vec![0.0f32; r2]);
15762        QTensor::matvec_many([&a, &b], &x, [&mut ga, &mut gb], Some(&pool));
15763        assert_eq!(ea, ga, "Q4TP fused lane 1 must be bit-identical");
15764        assert_eq!(eb, gb, "Q4TP fused lane 2 must be bit-identical");
15765        let _ = std::fs::remove_dir_all(&dir);
15766    }
15767
15768    /// The bus-speed expert kernel (`moe_fast`): its scales are bit for bit
15769    /// those of `scales_into`, whole experts equal `moe_gate_up_rows` +
15770    /// `moe_down_rows` up to the rounding of a sixteen-lane sum, with
15771    /// activation outliers on both stages, empty groups, groups past four
15772    /// tokens, and the same bits with and without a pool.
15773    #[cfg(target_arch = "x86_64")]
15774    #[test]
15775    fn fast_expert_rows_match_the_row_kernels() {
15776        use crate::pool::Pool;
15777        use cortiq_core::{CMF_VERSION, CmfHeader, CmfModel, QuantType, TensorSpec};
15778
15779        // The scales, every row width the code decode has a path for.
15780        for (rows, cols) in [(32usize, 2560usize), (16, 640), (48, 192), (16, 64)] {
15781            let bytes = synth_q4tp(rows, cols);
15782            let v = Q4tpView::new(&bytes, rows, cols);
15783            let gpr = cols / GROUP_SIZE;
15784            for r0 in (0..rows).step_by(16) {
15785                let Some(got) = moe_fast::block_scales(&v, r0, gpr) else {
15786                    eprintln!("fast_expert_rows: no AVX-512 VNNI here, skipped");
15787                    return;
15788                };
15789                for i in 0..16 {
15790                    let mut want = vec![0f32; gpr];
15791                    v.scales_into(r0 + i, gpr, &mut want);
15792                    assert_eq!(
15793                        want.iter().map(|x| x.to_bits()).collect::<Vec<_>>(),
15794                        got[i * gpr..(i + 1) * gpr].iter().map(|x| x.to_bits()).collect::<Vec<_>>(),
15795                        "scales of row {} ({rows}x{cols})",
15796                        r0 + i
15797                    );
15798                }
15799            }
15800        }
15801
15802        let (h, inter, ne) = (192usize, 640usize, 5usize);
15803        let arch: cortiq_core::ModelArch = serde_json::from_value(serde_json::json!({
15804            "arch_name": "tiny-q4tp-moe",
15805            "hidden_size": h,
15806            "intermediate_size": inter,
15807            "num_layers": 1,
15808            "num_attention_heads": h / 32,
15809            "num_kv_heads": 1,
15810            "head_dim": 32,
15811            "vocab_size": 8,
15812            "layer_types": ["FullAttention"],
15813            "rms_norm_eps": 1e-6,
15814            "max_position_embeddings": 8,
15815            "linear_conv_kernel_dim": 0,
15816            "linear_num_key_heads": 0,
15817            "linear_num_value_heads": 0
15818        }))
15819        .unwrap();
15820        let header = CmfHeader {
15821            format: "cmf".into(),
15822            version: CMF_VERSION,
15823            arch,
15824            quant_type: QuantType::Q4Block,
15825            provenance: None,
15826            tokenizer_config: None,
15827            section_hashes: None,
15828            skills: Vec::new(),
15829            shard: None,
15830            calibration: None,
15831            routing: None,
15832            genome: None,
15833            lineage: Vec::new(),
15834            router: None,
15835            segments: Vec::new(),
15836        };
15837        let mut specs = Vec::new();
15838        for e in 0..ne {
15839            for (k, (n, r, c)) in [("g", inter, h), ("u", inter, h), ("d", h, inter)]
15840                .into_iter()
15841                .enumerate()
15842            {
15843                let mut data = synth_q4tp(r, c);
15844                for (i, byte) in data[..r * (c / GROUP_SIZE) * Q4TP_NIB]
15845                    .iter_mut()
15846                    .enumerate()
15847                {
15848                    *byte ^= ((i * (e * 3 + k + 1)) % 251) as u8;
15849                }
15850                specs.push(TensorSpec {
15851                    name: format!("{n}{e}"),
15852                    dtype: TensorDtype::Q4TiledP,
15853                    shape: vec![r, c],
15854                    data,
15855                });
15856            }
15857        }
15858        let dir = std::env::temp_dir().join(format!("cmf-moe-fast-{}", std::process::id()));
15859        std::fs::create_dir_all(&dir).unwrap();
15860        let path = dir.join("m.cmf");
15861        CmfModel::write(&path, &header, &specs, None, None).unwrap();
15862        let model = Arc::new(CmfModel::open(&path).unwrap());
15863        let t = |n: String| QTensor::from_model(&model, &n).unwrap();
15864        let g: Vec<QTensor> = (0..ne).map(|e| t(format!("g{e}"))).collect();
15865        let u: Vec<QTensor> = (0..ne).map(|e| t(format!("u{e}"))).collect();
15866        let d: Vec<QTensor> = (0..ne).map(|e| t(format!("d{e}"))).collect();
15867        let b = 6usize;
15868        let mut xs: Vec<f32> = (0..b * h)
15869            .map(|i| ((i * 31 + 7) % 89) as f32 / 89.0 - 0.5)
15870            .collect();
15871        xs[5] = 9.0; // an activation outlier on token 0
15872        // Expert -> its tokens: five (a full chunk and one more), one, none, three, two.
15873        let groups: Vec<Vec<usize>> = vec![vec![0, 1, 2, 3, 4], vec![5], vec![], vec![2, 0, 5], vec![1, 3]];
15874        let lens: Vec<usize> = groups.iter().map(|v| v.len()).collect();
15875        let n_pairs: usize = lens.iter().sum();
15876        let pool = Pool::new(3);
15877        let close = |want: &[f32], got: &[f32], rel: f32, what: &str| {
15878            assert_eq!(want.len(), got.len(), "{what}: length");
15879            let top = want.iter().fold(0f32, |m, v| m.max(v.abs()));
15880            assert!(top > 0.0, "{what}: an all-zero reference proves nothing");
15881            for (i, (w, v)) in want.iter().zip(got).enumerate() {
15882                assert!((w - v).abs() <= rel * top, "{what}: [{i}] {w} vs {v} (scale {top})");
15883            }
15884        };
15885
15886        // The reference: the two row kernels.
15887        let pairs: Vec<(&QTensor, &QTensor)> = (0..ne).map(|e| (&g[e], &u[e])).collect();
15888        let mut gs: Vec<Vec<f32>> = (0..n_pairs).map(|_| vec![0f32; inter]).collect();
15889        assert!(QTensor::moe_gate_up_rows(&pairs, &groups, &xs, &mut gs, Some(&pool)));
15890        let downs: Vec<&QTensor> = d.iter().collect();
15891        let mut ds: Vec<Vec<f32>> = (0..n_pairs).map(|_| vec![0f32; h]).collect();
15892        assert!(QTensor::moe_down_rows(&downs, &lens, &gs, &mut ds, Some(&pool)));
15893
15894        let triples: Vec<(&QTensor, &QTensor, &QTensor)> = (0..ne).map(|e| (&g[e], &u[e], &d[e])).collect();
15895        let mut out = vec![0f32; n_pairs * h];
15896        let mut gs_fast = Vec::new();
15897        if !moe_fast::run_stages(&triples, &groups, &xs, &mut out, Some(&pool), None, Some(&mut gs_fast)) {
15898            eprintln!("fast_expert_rows: the int8 activation kernels are off here, skipped");
15899            let _ = std::fs::remove_dir_all(&dir);
15900            return;
15901        }
15902        close(&gs.concat(), &gs_fast, 1e-4, "gate and up");
15903        close(&ds.concat(), &out, 1e-2, "whole experts");
15904
15905        // The down stage alone on the reference's gate rows, one of them
15906        // with an outlier (the gate rows of a real expert have them).
15907        let mut gs_out = gs.clone();
15908        let top = gs.iter().flatten().fold(0f32, |m, v| m.max(v.abs()));
15909        gs_out[0][3] = 40.0 * top;
15910        let mut ds_out: Vec<Vec<f32>> = (0..n_pairs).map(|_| vec![0f32; h]).collect();
15911        assert!(QTensor::moe_down_rows(&downs, &lens, &gs_out, &mut ds_out, Some(&pool)));
15912        let mut out2 = vec![0f32; n_pairs * h];
15913        assert!(moe_fast::run_stages(&triples, &groups, &xs, &mut out2, Some(&pool), Some(&gs_out.concat()), None));
15914        close(&ds_out.concat(), &out2, 1e-4, "down");
15915
15916        // The same bits on the calling thread alone, and a refusal the caller can fall back from.
15917        let mut out3 = vec![0f32; n_pairs * h];
15918        assert!(QTensor::moe_expert_rows_fast(&triples, &groups, &xs, &mut out3, None));
15919        assert_eq!(
15920            out.iter().map(|x| x.to_bits()).collect::<Vec<_>>(),
15921            out3.iter().map(|x| x.to_bits()).collect::<Vec<_>>(),
15922            "pooled and serial runs must agree bit for bit"
15923        );
15924        let mut short = vec![0f32; n_pairs * h - 1];
15925        assert!(!QTensor::moe_expert_rows_fast(&triples, &groups, &xs, &mut short, Some(&pool)));
15926        let _ = std::fs::remove_dir_all(&dir);
15927    }
15928
15929    /// The MiMo speculative verify's kernels: several tokens' MoE through
15930    /// `moe_gate_up_rows` / `moe_down_rows` (+ the caller's route-order sum)
15931    /// is bit-identical to each token alone through `moe_gate_up_many` /
15932    /// `moe_down_many` (decode), and a row-exact `q4tp_matmat` of five
15933    /// tokens (wide enough for the blocked tiles) equals five matvecs.
15934    #[test]
15935    fn multi_token_moe_rows_equal_single_token_decode() {
15936        use crate::pool::Pool;
15937        use cortiq_core::{CMF_VERSION, CmfHeader, CmfModel, QuantType, TensorSpec};
15938
15939        let (h, inter, ne) = (64usize, 128usize, 3usize);
15940        let arch: cortiq_core::ModelArch = serde_json::from_value(serde_json::json!({
15941            "arch_name": "tiny-q4tp-moe",
15942            "hidden_size": h,
15943            "intermediate_size": inter,
15944            "num_layers": 1,
15945            "num_attention_heads": 2,
15946            "num_kv_heads": 1,
15947            "head_dim": 32,
15948            "vocab_size": 8,
15949            "layer_types": ["FullAttention"],
15950            "rms_norm_eps": 1e-6,
15951            "max_position_embeddings": 8,
15952            "linear_conv_kernel_dim": 0,
15953            "linear_num_key_heads": 0,
15954            "linear_num_value_heads": 0
15955        }))
15956        .unwrap();
15957        let header = CmfHeader {
15958            format: "cmf".into(),
15959            version: CMF_VERSION,
15960            arch,
15961            quant_type: QuantType::Q4Block,
15962            provenance: None,
15963            tokenizer_config: None,
15964            section_hashes: None,
15965            skills: Vec::new(),
15966            shard: None,
15967            calibration: None,
15968            routing: None,
15969            genome: None,
15970            lineage: Vec::new(),
15971            router: None,
15972            segments: Vec::new(),
15973        };
15974        let mut specs = Vec::new();
15975        for e in 0..ne {
15976            for (k, (n, r, c)) in [("g", inter, h), ("u", inter, h), ("d", h, inter)]
15977                .into_iter()
15978                .enumerate()
15979            {
15980                // Distinct experts: perturb only the nibble plane (any byte
15981                // is a valid pair of codes; the ladder stays intact).
15982                let mut data = synth_q4tp(r, c);
15983                for (i, byte) in data[..r * (c / GROUP_SIZE) * Q4TP_NIB]
15984                    .iter_mut()
15985                    .enumerate()
15986                {
15987                    *byte ^= ((i * (e * 3 + k + 1)) % 251) as u8;
15988                }
15989                specs.push(TensorSpec {
15990                    name: format!("{n}{e}"),
15991                    dtype: TensorDtype::Q4TiledP,
15992                    shape: vec![r, c],
15993                    data,
15994                });
15995            }
15996        }
15997        let dir = std::env::temp_dir().join(format!(
15998            "cmf-moe-rows-{}-{}",
15999            std::process::id(),
16000            FLOAT_ACTIVATIONS.get()
16001        ));
16002        std::fs::create_dir_all(&dir).unwrap();
16003        let path = dir.join("m.cmf");
16004        CmfModel::write(&path, &header, &specs, None, None).unwrap();
16005        let model = Arc::new(CmfModel::open(&path).unwrap());
16006        let t = |n: String| QTensor::from_model(&model, &n).unwrap();
16007        let g: Vec<QTensor> = (0..ne).map(|e| t(format!("g{e}"))).collect();
16008        let u: Vec<QTensor> = (0..ne).map(|e| t(format!("u{e}"))).collect();
16009        let d: Vec<QTensor> = (0..ne).map(|e| t(format!("d{e}"))).collect();
16010        let b = 4usize;
16011        let mut xs: Vec<f32> = (0..b * h)
16012            .map(|i| ((i * 31 + 7) % 89) as f32 / 89.0 - 0.5)
16013            .collect();
16014        xs[5] = 9.0; // an activation outlier on token 0
16015        // Token -> (experts in route order, weights).
16016        let routes: Vec<(Vec<usize>, Vec<f32>)> = vec![
16017            (vec![2, 0], vec![0.6, 0.4]),
16018            (vec![0, 1, 2], vec![0.2, 0.5, 0.3]),
16019            (vec![1], vec![1.0]),
16020            (vec![2, 1, 0], vec![0.25, 0.25, 0.5]),
16021        ];
16022        let pool = Pool::new(3);
16023        // Decode reference, token by token.
16024        let mut want = vec![0f32; b * h];
16025        for (tk, (idx, w)) in routes.iter().enumerate() {
16026            let x = &xs[tk * h..(tk + 1) * h];
16027            let pairs: Vec<(&QTensor, &QTensor)> = idx.iter().map(|&e| (&g[e], &u[e])).collect();
16028            let mut gs: Vec<Vec<f32>> = idx.iter().map(|_| vec![0f32; inter]).collect();
16029            assert!(QTensor::moe_gate_up_many(&pairs, x, &mut gs, Some(&pool)));
16030            if FLOAT_ACTIVATIONS.get() {
16031                for (slot, &e) in idx.iter().enumerate() {
16032                    let (mut gate, mut up) = (vec![0.0; inter], vec![0.0; inter]);
16033                    g[e].matvec(x, &mut gate, Some(&pool));
16034                    u[e].matvec(x, &mut up, Some(&pool));
16035                    for (v, u) in gate.iter_mut().zip(up) {
16036                        *v = (*v / (1.0 + (-*v).exp())) * u;
16037                    }
16038                    assert_eq!(gs[slot], gate, "float gate/up must equal ordinary matvecs");
16039                }
16040            }
16041            let downs: Vec<&QTensor> = idx.iter().map(|&e| &d[e]).collect();
16042            assert!(QTensor::moe_down_many(
16043                &downs,
16044                &gs,
16045                w,
16046                &mut want[tk * h..(tk + 1) * h],
16047                Some(&pool)
16048            ));
16049        }
16050        if FLOAT_ACTIVATIONS.get() {
16051            for (tk, (idx, w)) in routes.iter().enumerate() {
16052                let mut scalar = vec![0.0; h];
16053                for (&e, &weight) in idx.iter().zip(w) {
16054                    let (mut gate, mut up, mut down) =
16055                        (vec![0.0; inter], vec![0.0; inter], vec![0.0; h]);
16056                    g[e].matvec(&xs[tk * h..(tk + 1) * h], &mut gate, Some(&pool));
16057                    u[e].matvec(&xs[tk * h..(tk + 1) * h], &mut up, Some(&pool));
16058                    for (v, u) in gate.iter_mut().zip(up) {
16059                        *v = (*v / (1.0 + (-*v).exp())) * u;
16060                    }
16061                    d[e].matvec(&gate, &mut down, Some(&pool));
16062                    for (v, d) in scalar.iter_mut().zip(down) {
16063                        *v += weight * d;
16064                    }
16065                }
16066                assert_eq!(
16067                    &want[tk * h..(tk + 1) * h],
16068                    scalar,
16069                    "float many equals scalar experts"
16070                );
16071            }
16072        }
16073        // All four tokens at once, grouped by expert.
16074        let mut experts: Vec<usize> = Vec::new();
16075        let mut groups: Vec<Vec<usize>> = Vec::new();
16076        for (tk, (idx, _)) in routes.iter().enumerate() {
16077            for &e in idx {
16078                match experts.iter().position(|&x| x == e) {
16079                    Some(k) => groups[k].push(tk),
16080                    None => {
16081                        experts.push(e);
16082                        groups.push(vec![tk]);
16083                    }
16084                }
16085            }
16086        }
16087        let n_pairs: usize = groups.iter().map(|g| g.len()).sum();
16088        let pairs: Vec<(&QTensor, &QTensor)> = experts.iter().map(|&e| (&g[e], &u[e])).collect();
16089        let mut gs: Vec<Vec<f32>> = (0..n_pairs).map(|_| vec![0f32; inter]).collect();
16090        assert!(QTensor::moe_gate_up_rows(
16091            &pairs,
16092            &groups,
16093            &xs,
16094            &mut gs,
16095            Some(&pool)
16096        ));
16097        let downs: Vec<&QTensor> = experts.iter().map(|&e| &d[e]).collect();
16098        let lens: Vec<usize> = groups.iter().map(|g| g.len()).collect();
16099        let mut ds: Vec<Vec<f32>> = (0..n_pairs).map(|_| vec![0f32; h]).collect();
16100        assert!(QTensor::moe_down_rows(
16101            &downs,
16102            &lens,
16103            &gs,
16104            &mut ds,
16105            Some(&pool)
16106        ));
16107        let slot = |tk: usize, e: usize| {
16108            let k = experts.iter().position(|&x| x == e).unwrap();
16109            groups[..k].iter().map(|g| g.len()).sum::<usize>()
16110                + groups[k].iter().position(|&x| x == tk).unwrap()
16111        };
16112        let mut got = vec![0f32; b * h];
16113        for (tk, (idx, w)) in routes.iter().enumerate() {
16114            for i in 0..h {
16115                let mut acc = 0f32;
16116                for (&e, &we) in idx.iter().zip(w) {
16117                    acc += we * ds[slot(tk, e)][i];
16118                }
16119                got[tk * h + i] = acc;
16120            }
16121        }
16122        assert!(want.iter().any(|v| *v != 0.0));
16123        assert_eq!(
16124            want.iter().map(|v| v.to_bits()).collect::<Vec<_>>(),
16125            got.iter().map(|v| v.to_bits()).collect::<Vec<_>>(),
16126            "multi-token MoE must equal decode bit for bit"
16127        );
16128
16129        // Row-exact q4tp_matmat: five tokens (a blocked 1x4 tile + tail
16130        // otherwise) equal five matvecs.
16131        let b5 = 5usize;
16132        let x5: Vec<f32> = (0..b5 * h)
16133            .map(|i| ((i * 13 + 5) % 71) as f32 / 71.0 - 0.5)
16134            .collect();
16135        let mut mm = vec![0f32; b5 * inter];
16136        row_exact_scope(|| g[1].matmat(&x5, b5, &mut mm, Some(&pool)));
16137        for tk in 0..b5 {
16138            let mut mv = vec![0f32; inter];
16139            g[1].matvec(&x5[tk * h..(tk + 1) * h], &mut mv, Some(&pool));
16140            assert_eq!(
16141                mv.iter().map(|v| v.to_bits()).collect::<Vec<_>>(),
16142                mm[tk * inter..(tk + 1) * inter]
16143                    .iter()
16144                    .map(|v| v.to_bits())
16145                    .collect::<Vec<_>>(),
16146                "row-exact matmat token {tk}"
16147            );
16148        }
16149        // Other concurrent tests/requests may still hold the shared mode.
16150        // Nested, overlapping and unwind restoration is checked separately.
16151        let _ = std::fs::remove_dir_all(&dir);
16152    }
16153
16154    /// The row-exact fix must not touch the fast path. Outside the scope
16155    /// `q4tp_matmat` has to produce exactly what it did before, and on ARM
16156    /// "before" is spelled out below: the tuned 1x4 SDOT tile for every
16157    /// four columns and the single-row kernel for the tail. Inside the
16158    /// scope every column equals its token's matvec. `q4tp_matmat_with`
16159    /// takes the mode as an argument, so a concurrent test holding the
16160    /// shared scope cannot flip it under this one.
16161    #[test]
16162    fn q4tp_matmat_fast_path_unchanged_outside_row_exact() {
16163        use crate::pool::Pool;
16164        use std::sync::atomic::Ordering::Relaxed;
16165        let _alt = Q4TP_ALT_TEST_LOCK.lock().unwrap_or_else(|e| e.into_inner());
16166        // The tuned ARM shape, which is also what an unset switch picks
16167        // unless CMF_Q4TP_V1 is exported.
16168        Q4TP_ALT.store(2, Relaxed);
16169        let pool = Pool::new(3);
16170        // Under 500k cells, so macOS keeps the matmat off the AMX; the
16171        // second shape runs across pool workers (rows >= 256) with 32
16172        // groups of accumulation and a tail after two 1x4 tiles.
16173        for &(rows, cols, b) in &[(64usize, 256usize, 7usize), (320, 1024, 9)] {
16174            let bytes = synth_q4tp(rows, cols);
16175            let mut xs: Vec<f32> = (0..b * cols)
16176                .map(|i| ((i * 29 + 11) % 83) as f32 / 83.0 - 0.5)
16177                .collect();
16178            xs[3] = 7.5; // an activation outlier on token 0
16179            let run = |exact: bool| {
16180                let mut out = vec![0f32; b * rows];
16181                q4tp_matmat_with(&bytes, &xs, b, rows, cols, &mut out, Some(&pool), exact);
16182                out
16183            };
16184            let (fast, exact) = (run(false), run(true));
16185            #[cfg(not(target_arch = "aarch64"))]
16186            let _ = fast;
16187            let bits = |v: &[f32]| v.iter().map(|x| x.to_bits()).collect::<Vec<_>>();
16188            let mut matvecs = vec![0f32; b * rows];
16189            for (bi, o) in matvecs.chunks_mut(rows).enumerate() {
16190                q4tp_matvec(
16191                    &bytes,
16192                    &xs[bi * cols..(bi + 1) * cols],
16193                    rows,
16194                    cols,
16195                    o,
16196                    Some(&pool),
16197                );
16198            }
16199            assert!(matvecs.iter().any(|v| *v != 0.0));
16200            assert_eq!(
16201                bits(&exact),
16202                bits(&matvecs),
16203                "{rows}x{cols} b={b}: row-exact matmat must equal per-token matvecs"
16204            );
16205            #[cfg(target_arch = "aarch64")]
16206            {
16207                // The pre-fix ARM loop, cell for cell.
16208                let gpr = cols / GROUP_SIZE;
16209                let v = Q4tpView::new(&bytes, rows, cols);
16210                let mut old = vec![0f32; b * rows];
16211                let mut sc = vec![0f32; gpr];
16212                let a8w8 = a8w8_enabled();
16213                let blocked = sdot_enabled() && blocked_enabled();
16214                let acts: Vec<SplitAct> = (0..b)
16215                    .map(|bi| split_act(&xs[bi * cols..(bi + 1) * cols]))
16216                    .collect();
16217                for r in 0..rows {
16218                    v.scales_into(r, gpr, &mut sc);
16219                    if !a8w8 {
16220                        for bi in 0..b {
16221                            let x = &xs[bi * cols..(bi + 1) * cols];
16222                            old[bi * rows + r] = q4tp_row_exact(v.nib, r, gpr, x, &sc);
16223                        }
16224                        continue;
16225                    }
16226                    let finish = |d: f32, act: &SplitAct| {
16227                        let mut acc = d * act.sx;
16228                        for &(j, xv) in &act.outliers {
16229                            let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
16230                            acc += w * s * xv;
16231                        }
16232                        acc
16233                    };
16234                    let mut bi = 0usize;
16235                    while blocked && bi + 4 <= b {
16236                        let xs4 = [
16237                            acts[bi].xq.as_slice(),
16238                            acts[bi + 1].xq.as_slice(),
16239                            acts[bi + 2].xq.as_slice(),
16240                            acts[bi + 3].xq.as_slice(),
16241                        ];
16242                        let d = unsafe { dot_q4tp_row_1x4_sdot(v.nib, r, gpr, xs4, &sc) };
16243                        for k in 0..4 {
16244                            old[(bi + k) * rows + r] = finish(d[k], &acts[bi + k]);
16245                        }
16246                        bi += 4;
16247                    }
16248                    for (bi, act) in acts.iter().enumerate().skip(bi) {
16249                        let d = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, &sc);
16250                        old[bi * rows + r] = finish(d, act);
16251                    }
16252                }
16253                assert_eq!(
16254                    bits(&fast),
16255                    bits(&old),
16256                    "{rows}x{cols} b={b}: the fast path changed outside row_exact"
16257                );
16258                // And it is still the fast tile that runs: its lane-parallel
16259                // fma sum rounds differently from the matvec somewhere.
16260                if blocked {
16261                    assert_ne!(
16262                        bits(&fast),
16263                        bits(&matvecs),
16264                        "{rows}x{cols} b={b}: the tuned tile no longer runs outside row_exact"
16265                    );
16266                }
16267            }
16268        }
16269        Q4TP_ALT.store(0, Relaxed);
16270    }
16271
16272    #[test]
16273    fn row_exact_scopes_survive_overlap_nesting_and_unwind() {
16274        use std::sync::{Barrier, atomic::{AtomicUsize, Ordering}};
16275        // A private counter makes this restoration test independent of
16276        // numerical tests concurrently using the production counter.
16277        let active = AtomicUsize::new(0);
16278        counted_row_exact_scope(&active, || {
16279            assert_eq!(active.load(Ordering::Acquire), 1);
16280            counted_row_exact_scope(&active, || {
16281                assert_eq!(active.load(Ordering::Acquire), 2);
16282            });
16283            assert_eq!(active.load(Ordering::Acquire), 1);
16284        });
16285        assert_eq!(active.load(Ordering::Acquire), 0);
16286
16287        let both_entered = Barrier::new(2);
16288        let release_last = Barrier::new(2);
16289        std::thread::scope(|s| {
16290            let first = s.spawn(|| counted_row_exact_scope(&active, || {
16291                both_entered.wait();
16292            }));
16293            let last = s.spawn(|| counted_row_exact_scope(&active, || {
16294                both_entered.wait();
16295                release_last.wait();
16296            }));
16297            first.join().unwrap();
16298            let after_first = active.load(Ordering::Acquire);
16299            release_last.wait();
16300            last.join().unwrap();
16301            assert_eq!(after_first, 1, "second request must remain exact");
16302        });
16303        assert_eq!(active.load(Ordering::Acquire), 0);
16304        let panic = std::panic::catch_unwind(|| {
16305            counted_row_exact_scope(&active, || panic!("scope unwind"));
16306        });
16307        assert!(panic.is_err());
16308        assert_eq!(active.load(Ordering::Acquire), 0);
16309    }
16310
16311    #[test]
16312    #[cfg(target_arch = "x86_64")]
16313    fn q4tp_float_avx2_is_bitwise_scalar() {
16314        if !avx2_enabled() {
16315            return;
16316        }
16317        for cols in [32, 64, 96, 2048, 4096] {
16318            let rows = 9;
16319            let bytes = synth_q4tp(rows, cols);
16320            let v = Q4tpView::new(&bytes, rows, cols);
16321            let gpr = cols / GROUP_SIZE;
16322            let mut sc = vec![0.0; gpr];
16323            for seed in 1..=5 {
16324                let xs: Vec<f32> = (0..cols)
16325                    .map(|i| (((i * 104729 + seed * 8191) % 100003) as f32 - 50001.0) / 7919.0)
16326                    .collect();
16327                for r in 0..rows {
16328                    v.scales_into(r, gpr, &mut sc);
16329                    let scalar = q4tp_row_float_scalar(v.nib, r, gpr, &xs, &sc);
16330                    let vector = unsafe { q4tp_row_float_avx2(v.nib, r, gpr, &xs, &sc) };
16331                    assert_eq!(
16332                        scalar.to_bits(),
16333                        vector.to_bits(),
16334                        "cols={cols} row={r} seed={seed}"
16335                    );
16336                }
16337            }
16338        }
16339    }
16340
16341    #[test]
16342    fn multi_token_moe_rows_float_equal_single_token_decode() {
16343        float_activations_scope(multi_token_moe_rows_equal_single_token_decode);
16344    }
16345
16346    #[test]
16347    fn full_gpu_q8_scope_is_nested_and_thread_local() {
16348        assert!(!FULL_GPU_Q8.get());
16349        let before = gpu_split_frac();
16350        {
16351            let _guard = enter_full_gpu_q8_scope();
16352            assert_eq!(gpu_split_frac(), 1.0);
16353            {
16354                let _nested = enter_full_gpu_q8_scope();
16355            }
16356            assert_eq!(gpu_split_frac(), 1.0);
16357            std::thread::spawn(|| assert!(!FULL_GPU_Q8.get())).join().unwrap();
16358        }
16359        assert!(!FULL_GPU_Q8.get());
16360        assert_eq!(gpu_split_frac(), before);
16361    }
16362
16363    #[test]
16364    fn float_activation_scope_is_nested_thread_local_and_unwind_safe() {
16365        assert!(!FLOAT_ACTIVATIONS.get());
16366        let before = a8w8_enabled();
16367        float_activations_scope(|| {
16368            assert!(!a8w8_enabled());
16369            float_activations_scope(|| assert!(!a8w8_enabled()));
16370            assert!(FLOAT_ACTIVATIONS.get());
16371            std::thread::spawn(|| assert!(!FLOAT_ACTIVATIONS.get()))
16372                .join()
16373                .unwrap();
16374        });
16375        assert!(!FLOAT_ACTIVATIONS.get());
16376        assert_eq!(a8w8_enabled(), before);
16377        let _ = std::panic::catch_unwind(|| float_activations_scope(|| panic!("test unwind")));
16378        assert!(!FLOAT_ACTIVATIONS.get());
16379    }
16380
16381    /// Batched q4/vbit matmat must equal per-position matvec calls
16382    /// exactly (the fallback it replaced) — same kernels, same order.
16383    #[test]
16384    fn batched_matmat_equals_per_position_matvec() {
16385        let (rows, cols, b) = (8, 64, 5);
16386        // q4 blob.
16387        let groups = rows * cols / GROUP_SIZE;
16388        let mut q4 = Vec::new();
16389        for i in 0..groups * 16 {
16390            q4.push((((i * 7 + 3) % 256) & 0xFF) as u8);
16391        }
16392        for g in 0..groups {
16393            q4.extend_from_slice(
16394                &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
16395            );
16396        }
16397        // vbit blob (mixed widths incl. 8).
16398        let ng = cols / GROUP_SIZE;
16399        let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4, 5, 3];
16400        let mut vb = bits.clone();
16401        for g in 0..rows * ng {
16402            vb.extend_from_slice(
16403                &cortiq_core::quant::f32_to_f16(0.02 + 0.001 * g as f32).to_le_bytes(),
16404            );
16405        }
16406        for r in 0..rows {
16407            let bw = bits[r] as usize;
16408            let (mut acc, mut nb) = (0u64, 0usize);
16409            let mut rowbytes = Vec::new();
16410            for i in 0..cols {
16411                let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
16412                acc = (acc << bw) | v;
16413                nb += bw;
16414                while nb >= 8 {
16415                    nb -= 8;
16416                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
16417                }
16418            }
16419            if nb > 0 {
16420                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
16421            }
16422            vb.extend_from_slice(&rowbytes);
16423        }
16424        let offsets = vbit_row_offsets(&vb, rows, cols);
16425
16426        let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
16427
16428        // q4: batch vs singles.
16429        let mut got = vec![0f32; b * rows];
16430        q4matmat(&q4, &xs, b, rows, cols, &mut got, None);
16431        for bi in 0..b {
16432            let mut expect = vec![0f32; rows];
16433            q4matvec(
16434                &q4,
16435                &xs[bi * cols..(bi + 1) * cols],
16436                rows,
16437                cols,
16438                &mut expect,
16439                None,
16440            );
16441            assert_eq!(
16442                &got[bi * rows..(bi + 1) * rows],
16443                &expect[..],
16444                "q4 batch pos {bi}"
16445            );
16446        }
16447
16448        // vbit: batch vs singles.
16449        let mut got = vec![0f32; b * rows];
16450        vbitmatmat(&vb, &offsets, &xs, b, rows, cols, &mut got, None);
16451        for bi in 0..b {
16452            let mut expect = vec![0f32; rows];
16453            vbitmatvec(
16454                &vb,
16455                &offsets,
16456                &xs[bi * cols..(bi + 1) * cols],
16457                rows,
16458                cols,
16459                &mut expect,
16460                None,
16461            );
16462            assert_eq!(
16463                &got[bi * rows..(bi + 1) * rows],
16464                &expect[..],
16465                "vbit batch pos {bi}"
16466            );
16467        }
16468    }
16469
16470    /// q4_tiled kernels must produce BIT-identical outputs to the q4
16471    /// split kernels on the same values (same ints, same order — only
16472    /// the byte placement differs).
16473    #[test]
16474    fn q4_tiled_matches_q4_block_bitexact() {
16475        let (rows, cols, b) = (8usize, 128usize, 3usize);
16476        let groups = rows * cols / GROUP_SIZE;
16477        let mut split = Vec::with_capacity(groups * 18);
16478        for i in 0..groups * 16 {
16479            split.push((((i * 7 + 3) % 256) & 0xFF) as u8);
16480        }
16481        for g in 0..groups {
16482            split.extend_from_slice(
16483                &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
16484            );
16485        }
16486        // Re-tile: [scale][nibbles] per group.
16487        let (packed, scales) = split.split_at(groups * 16);
16488        let mut tiled = Vec::with_capacity(groups * Q4_TILE);
16489        for g in 0..groups {
16490            tiled.extend_from_slice(&scales[g * 2..g * 2 + 2]);
16491            tiled.extend_from_slice(&packed[g * 16..(g + 1) * 16]);
16492        }
16493
16494        let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
16495        x1[9] = 250.0; // exercise the outlier path
16496        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
16497
16498        let (mut a, mut t) = (vec![0f32; rows], vec![0f32; rows]);
16499        q4matvec(&split, &x1, rows, cols, &mut a, None);
16500        q4t_matvec(&tiled, &x1, rows, cols, &mut t, None);
16501        assert_eq!(a, t, "q4t matvec must match q4 bit-for-bit");
16502
16503        let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
16504        let (mut t1, mut t2) = (vec![0f32; rows], vec![0f32; rows]);
16505        q4matvec2(&split, &x1, &x2, rows, cols, &mut a1, &mut a2, None);
16506        q4t_matvec2(&tiled, &x1, &x2, rows, cols, &mut t1, &mut t2, None);
16507        assert_eq!(a1, t1);
16508        assert_eq!(a2, t2);
16509
16510        let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
16511        let (mut am, mut tm) = (vec![0f32; b * rows], vec![0f32; b * rows]);
16512        q4matmat(&split, &xs, b, rows, cols, &mut am, None);
16513        q4t_matmat(&tiled, &xs, b, rows, cols, &mut tm, None);
16514        assert_eq!(am, tm, "q4t matmat must match q4 bit-for-bit");
16515    }
16516
16517    /// q4 SDOT outlier correction: a single huge activation channel
16518    /// (>8·rms → outlier, zeroed in xq) must still contribute its EXACT
16519    /// term. On-grid bulk (±1/0 → xq dequantizes exactly) isolates the
16520    /// correction from A8W8 noise. cols must exceed 64: at n=64 the
16521    /// 8·rms threshold equals sqrt(v²+rest) ≥ v, so a single outlier
16522    /// can never qualify (8² = n).
16523    #[test]
16524    fn q4matvec_sdot_outlier_exact() {
16525        let (rows, cols) = (4, 128);
16526        let groups = rows * cols / GROUP_SIZE;
16527        let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
16528        for i in 0..groups * 16 {
16529            bytes.push(((i * 11 + 5) % 256) as u8);
16530        }
16531        for g in 0..groups {
16532            let s = 0.02 + 0.002 * g as f32;
16533            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
16534        }
16535        let mut x: Vec<f32> = (0..cols)
16536            .map(|i| match i % 3 {
16537                0 => 1.0,
16538                1 => -1.0,
16539                _ => 0.0,
16540            })
16541            .collect();
16542        x[17] = 300.0; // ≫ 8·rms → outlier channel
16543
16544        let mut reference = vec![0.0f32; rows * cols];
16545        cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
16546        let mut expect = vec![0.0f32; rows];
16547        for r in 0..rows {
16548            expect[r] = reference[r * cols..(r + 1) * cols]
16549                .iter()
16550                .zip(&x)
16551                .map(|(w, xv)| w * xv)
16552                .sum();
16553        }
16554        let mut got = vec![0.0f32; rows];
16555        q4matvec(&bytes, &x, rows, cols, &mut got, None);
16556        let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
16557        for r in 0..rows {
16558            assert!(
16559                (got[r] - expect[r]).abs() < 2e-3 * scale,
16560                "row {r}: {} vs {} (outlier term must be exact)",
16561                got[r],
16562                expect[r]
16563            );
16564        }
16565    }
16566
16567    /// The fused q1t matvec must equal the reference (dequant_q1t → dot),
16568    /// including the ternary zero level and the binary-searched outlier
16569    /// overlay. Guards the mmap kernel that makes a 12B q1t runnable.
16570    #[test]
16571    fn q1t_matvec_matches_reference() {
16572        use cortiq_core::quant::{dequant_q1t, f32_to_f16};
16573        let (rows, cols) = (3usize, 64usize); // gpr = 2
16574        let gpr = cols / GROUP_SIZE;
16575        let scales = [0.5f32, 0.3, 0.7, 0.2, 0.6, 0.15];
16576        // Overlay (must be sorted by flat index): a few spikes across rows.
16577        let outliers: [(u32, f32); 3] = [(5, 9.0), (70, -4.5), (150, 3.25)];
16578        let is_out = |flat: usize| outliers.iter().any(|&(i, _)| i as usize == flat);
16579        let mut bytes = Vec::new();
16580        for r in 0..rows {
16581            for g in 0..gpr {
16582                bytes.extend_from_slice(&f32_to_f16(scales[r * gpr + g]).to_le_bytes());
16583                let mut c = [0u8; 7];
16584                for k in 0..GROUP_SIZE {
16585                    // Encoder invariant: code 0 at outlier positions.
16586                    let code = if is_out(r * cols + g * GROUP_SIZE + k) {
16587                        0
16588                    } else {
16589                        ((k + r * 3 + g) % 3) as u8 // 0,1,2
16590                    };
16591                    cortiq_core::quant::q1t_pack(&mut c, k, code);
16592                }
16593                bytes.extend_from_slice(&c);
16594            }
16595        }
16596        // Per-row overlay: [u32 row_ptr[rows+1]] then [(u16 col, f16 val)] by
16597        // row (outliers are sorted by flat index → already grouped by row).
16598        let mut row_ptr = vec![0u32; rows + 1];
16599        for &(idx, _) in &outliers {
16600            row_ptr[idx as usize / cols + 1] += 1;
16601        }
16602        for r in 0..rows {
16603            row_ptr[r + 1] += row_ptr[r];
16604        }
16605        for &p in &row_ptr {
16606            bytes.extend_from_slice(&p.to_le_bytes());
16607        }
16608        for &(idx, v) in &outliers {
16609            bytes.extend_from_slice(&((idx as usize % cols) as u16).to_le_bytes());
16610            bytes.extend_from_slice(&f32_to_f16(v).to_le_bytes());
16611        }
16612
16613        let mut refw = vec![0f32; rows * cols];
16614        dequant_q1t(&bytes, rows, cols, &mut refw);
16615        // On-grid activations (±1, amax 1) so the int8 SDOT path reconstructs
16616        // x exactly and matches the f32 reference (same trick as the q1 test).
16617        let x: Vec<f32> = (0..cols)
16618            .map(|j| if j % 3 == 0 { 1.0 } else { -1.0 })
16619            .collect();
16620        let mut expect = vec![0f32; rows];
16621        for r in 0..rows {
16622            let mut a = 0.0f32;
16623            for j in 0..cols {
16624                a += refw[r * cols + j] * x[j];
16625            }
16626            expect[r] = a;
16627        }
16628        let tol = |e: f32| 1e-3 * e.abs().max(1e-3);
16629        let mut got = vec![0f32; rows];
16630        q1t_matvec(&bytes, &x, rows, cols, &mut got, None);
16631        for r in 0..rows {
16632            assert!(
16633                (got[r] - expect[r]).abs() < tol(expect[r]),
16634                "row {r}: {} vs {}",
16635                got[r],
16636                expect[r]
16637            );
16638        }
16639        // matmat (b=2, f32 decode path) must agree too.
16640        let x2: Vec<f32> = x.iter().chain(x.iter()).copied().collect();
16641        let mut gm = vec![0f32; 2 * rows];
16642        q1t_matmat(&bytes, &x2, 2, rows, cols, &mut gm, None);
16643        for r in 0..rows {
16644            assert!((gm[r] - expect[r]).abs() < tol(expect[r]));
16645            assert!((gm[rows + r] - expect[r]).abs() < tol(expect[r]));
16646        }
16647        // Fused pair (q1t_matvec2) must equal two single matvecs
16648        // bit-for-bit: same unpack, same group order, same f32
16649        // accumulation per stream. Distinct x2 exercises both lanes.
16650        let xb: Vec<f32> = (0..cols)
16651            .map(|j| if j % 5 == 0 { -1.0 } else { 1.0 })
16652            .collect();
16653        let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
16654        q1t_matvec(&bytes, &x, rows, cols, &mut s1, None);
16655        q1t_matvec(&bytes, &xb, rows, cols, &mut s2, None);
16656        let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
16657        q1t_matvec2(&bytes, &x, &xb, rows, cols, &mut p1, &mut p2, None);
16658        assert_eq!(p1, s1, "q1t pair lane 1 ≠ single matvec");
16659        assert_eq!(p2, s2, "q1t pair lane 2 ≠ single matvec");
16660    }
16661
16662    /// Pair == 2×matvec with an ODD group count (the kernel's tail
16663    /// group) and no overlay section.
16664    #[test]
16665    fn q1t_matvec2_odd_gpr_matches_singles() {
16666        use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_pack};
16667        let (rows, cols) = (5usize, 96usize); // gpr = 3 → paired + tail
16668        let gpr = cols / GROUP_SIZE;
16669        let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE);
16670        for r in 0..rows {
16671            for g in 0..gpr {
16672                bytes.extend_from_slice(&f32_to_f16(0.1 + 0.05 * (r + g) as f32).to_le_bytes());
16673                let mut c = [0u8; 7];
16674                for k in 0..GROUP_SIZE {
16675                    q1t_pack(&mut c, k, ((k * 7 + r * 5 + g * 3) % 3) as u8);
16676                }
16677                bytes.extend_from_slice(&c);
16678            }
16679        }
16680        let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.31).sin()).collect();
16681        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).cos()).collect();
16682        let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
16683        q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
16684        q1t_matvec(&bytes, &x2, rows, cols, &mut s2, None);
16685        let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
16686        q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
16687        assert_eq!(p1, s1, "odd-gpr pair lane 1 ≠ single");
16688        assert_eq!(p2, s2, "odd-gpr pair lane 2 ≠ single");
16689    }
16690
16691    // Speed A/B: fused pair (one unpack, two streams) vs two single
16692    // matvecs. Single-threaded, FFN-sized, min-of paired in-process.
16693    //   cargo test -p cortiq-engine --release q1t_matvec2_speed -- --ignored --nocapture
16694    #[test]
16695    #[ignore]
16696    fn q1t_matvec2_speed() {
16697        use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_pack};
16698        use std::time::Instant;
16699        let (rows, cols) = (8192usize, 4096usize);
16700        let gpr = cols / GROUP_SIZE;
16701        let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE);
16702        for r in 0..rows {
16703            for g in 0..gpr {
16704                let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
16705                bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
16706                let mut c = [0u8; 7];
16707                for k in 0..GROUP_SIZE {
16708                    q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
16709                }
16710                bytes.extend_from_slice(&c);
16711            }
16712        }
16713        let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.31).sin()).collect();
16714        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).cos()).collect();
16715        let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
16716        let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
16717        // Warm both paths once.
16718        q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
16719        q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
16720        let (mut t_pair, mut t_two) = (f64::MAX, f64::MAX);
16721        for _ in 0..8 {
16722            let t0 = Instant::now();
16723            q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
16724            t_pair = t_pair.min(t0.elapsed().as_secs_f64() * 1000.0);
16725            let t1 = Instant::now();
16726            q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
16727            q1t_matvec(&bytes, &x2, rows, cols, &mut s2, None);
16728            t_two = t_two.min(t1.elapsed().as_secs_f64() * 1000.0);
16729        }
16730        assert_eq!(p1, s1);
16731        assert_eq!(p2, s2);
16732        println!("q1t pair {rows}x{cols}: fused {t_pair:.2} ms | two singles {t_two:.2} ms");
16733    }
16734
16735    // Speed A/B: the base-3-division decode (what the packing commit left in
16736    // place) vs the fused sign-LUT matvec. Both single-threaded, same bytes.
16737    //   cargo test -p cortiq-engine q1t_matvec_speed -- --ignored --nocapture
16738    #[test]
16739    #[ignore]
16740    fn q1t_matvec_speed() {
16741        use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_code, q1t_pack};
16742        use std::time::Instant;
16743        let (rows, cols) = (8192usize, 4096usize); // FFN-sized
16744        let gpr = cols / GROUP_SIZE;
16745        let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE + 16);
16746        for r in 0..rows {
16747            for g in 0..gpr {
16748                let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
16749                bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
16750                let mut c = [0u8; 7];
16751                for k in 0..GROUP_SIZE {
16752                    q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
16753                }
16754                bytes.extend_from_slice(&c);
16755            }
16756        }
16757        let (n, stride) = (rows * cols, 40usize); // ~2.5% outliers, per-row overlay
16758        let mut row_ptr = vec![0u32; rows + 1];
16759        let mut idx = 0usize;
16760        while idx < n {
16761            row_ptr[idx / cols + 1] += 1;
16762            idx += stride;
16763        }
16764        for r in 0..rows {
16765            row_ptr[r + 1] += row_ptr[r];
16766        }
16767        for &p in &row_ptr {
16768            bytes.extend_from_slice(&p.to_le_bytes());
16769        }
16770        let mut idx = 0usize;
16771        while idx < n {
16772            bytes.extend_from_slice(&((idx % cols) as u16).to_le_bytes());
16773            bytes.extend_from_slice(&f32_to_f16((idx % 13) as f32 * 0.1 - 0.6).to_le_bytes());
16774            idx += stride;
16775        }
16776        // On-grid ±1 so the fast path's int8 SDOT is exact vs the f32 "slow"
16777        // reference (the A/B is a timing check; values must still agree).
16778        let x: Vec<f32> = (0..cols)
16779            .map(|j| if j % 3 == 0 { 1.0 } else { -1.0 })
16780            .collect();
16781        let (rp_off, ent_off, has_ov) = q1t_overlay(&bytes, rows * gpr * Q1T_TILE, rows);
16782
16783        // "before": base-3 division decode into a buffer, then dot.
16784        let slow = |out: &mut [f32]| {
16785            let mut buf = vec![0f32; cols];
16786            for r in 0..rows {
16787                for g in 0..gpr {
16788                    let off = (r * gpr + g) * Q1T_TILE;
16789                    let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
16790                    let codes = &bytes[off + 2..off + Q1T_TILE];
16791                    for k in 0..GROUP_SIZE {
16792                        buf[g * GROUP_SIZE + k] = match q1t_code(codes, k) {
16793                            1 => s,
16794                            2 => -s,
16795                            _ => 0.0,
16796                        };
16797                    }
16798                }
16799                out[r] = q1t_row_outlier_correction(&bytes, r, rp_off, ent_off, has_ov, &x)
16800                    + (0..cols).map(|j| buf[j] * x[j]).sum::<f32>();
16801            }
16802        };
16803        let iters = 5;
16804        let mut a = vec![0f32; rows];
16805        slow(&mut a); // warm
16806        let t = Instant::now();
16807        for _ in 0..iters {
16808            slow(&mut a);
16809        }
16810        let slow_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
16811
16812        let mut b = vec![0f32; rows];
16813        q1t_matvec(&bytes, &x, rows, cols, &mut b, None); // warm
16814        let t = Instant::now();
16815        for _ in 0..iters {
16816            q1t_matvec(&bytes, &x, rows, cols, &mut b, None);
16817        }
16818        let fast_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
16819
16820        for r in 0..rows {
16821            assert!((a[r] - b[r]).abs() < 1e-2, "mismatch row {r}");
16822        }
16823        println!(
16824            "q1t matvec {rows}x{cols} (1 thread): div-decode {slow_ms:.2} ms  fused-LUT {fast_ms:.2} ms  => {:.2}x",
16825            slow_ms / fast_ms
16826        );
16827    }
16828}
16829
16830#[cfg(test)]
16831mod row_span_tests {
16832    use cortiq_core::TensorDtype;
16833
16834    /// The three ranges `row_span_offsets` names are exactly what decoding a
16835    /// q4tp row reads: with every other byte of the tensor overwritten the
16836    /// row's nibbles and scales come back unchanged.
16837    #[test]
16838    fn q4tp_row_spans_cover_the_row() {
16839        let (rows, cols) = (37usize, 160usize);
16840        let gpr = cols / 32;
16841        let total = cortiq_core::quant::expected_nbytes(TensorDtype::Q4TiledP, &[rows, cols]).unwrap();
16842        let (params_off, _, _) = cortiq_core::quant::q4tp_sections(rows, cols);
16843        let mut bytes: Vec<u8> = (0..total).map(|i| (i * 61 % 251) as u8).collect();
16844        let lo = cortiq_core::quant::f32_to_f16(-4.0);
16845        let step = cortiq_core::quant::f32_to_f16(0.1);
16846        for r in 0..rows {
16847            let o = params_off + r * 4;
16848            bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
16849            bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
16850        }
16851        for r in [0usize, 1, 17, rows - 1] {
16852            let spans = super::row_span_offsets(TensorDtype::Q4TiledP, rows, cols, r).unwrap();
16853            assert_eq!(spans.iter().map(|s| s.1).sum::<usize>(), total / rows);
16854            let mut keep = vec![false; total];
16855            for &(o, n) in &spans {
16856                keep[o..o + n].fill(true);
16857            }
16858            let other: Vec<u8> =
16859                bytes.iter().zip(&keep).map(|(&b, &k)| if k { b } else { 0x5A }).collect();
16860            let (a, b) = (super::Q4tpView::new(&bytes, rows, cols), super::Q4tpView::new(&other, rows, cols));
16861            let (mut sa, mut sb) = (vec![0f32; gpr], vec![0f32; gpr]);
16862            a.scales_into(r, gpr, &mut sa);
16863            b.scales_into(r, gpr, &mut sb);
16864            assert_eq!(sa, sb, "row {r} scales");
16865            let nib = r * gpr * super::Q4TP_NIB..(r + 1) * gpr * super::Q4TP_NIB;
16866            assert_eq!(a.nib[nib.clone()], b.nib[nib], "row {r} nibbles");
16867        }
16868        assert!(super::row_span_offsets(TensorDtype::Q4TiledP, rows, cols, rows).is_none());
16869        assert!(super::row_span_offsets(TensorDtype::Q8_2f, rows, cols, 0).is_none());
16870        let t = super::row_span_offsets(TensorDtype::Q4Tiled, rows, cols, 3).unwrap();
16871        assert_eq!(t[0], (3 * gpr * super::Q4_TILE, gpr * super::Q4_TILE));
16872    }
16873}
16874
16875#[cfg(test)]
16876mod gemm_bench {
16877    /// `cargo test -p cortiq-engine --release q4tp_matmat_throughput -- --ignored --nocapture`
16878    /// Times the batched q4tp GEMM at the shapes the image DiT runs
16879    /// (b=296 tokens, 2304 -> 9216), on synthetic bytes: no model, no
16880    /// mmap, no thermal drift over minutes — a kernel change shows up
16881    /// here in seconds where a full render hides it in noise.
16882    ///
16883    /// On macOS add `CMF_ACCEL=0`: this shape is over the 500k-cell mark
16884    /// where the matmat hands off to Accelerate's dequant sgemm, and
16885    /// without the opt-out both rows below measure the AMX, not the
16886    /// kernel under test.
16887    #[test]
16888    #[ignore]
16889    fn q4tp_matmat_throughput() {
16890        let _alt = super::Q4TP_ALT_TEST_LOCK
16891            .lock()
16892            .unwrap_or_else(|e| e.into_inner());
16893        // 296 is a prompt-encode batch; the image DiT runs 2085 at
16894        // 512x512, where the activation panel stops fitting L2 and the
16895        // loop's shape starts to matter more than its instructions.
16896        let b: usize = std::env::var("CMF_BENCH_B")
16897            .ok()
16898            .and_then(|v| v.parse().ok())
16899            .unwrap_or(296);
16900        let (rows, cols) = (9216usize, 2304usize);
16901        let (_, _, _) = (rows, cols, b);
16902        let total =
16903            cortiq_core::quant::expected_nbytes(cortiq_core::TensorDtype::Q4TiledP, &[rows, cols])
16904                .unwrap();
16905        // Random nibbles are fine, but the row params are f16 (lo, step)
16906        // of a geometric ladder: garbage there gives exp2 of a huge
16907        // exponent, the scales come back inf, and the whole bench times
16908        // NaN arithmetic instead of the kernel.
16909        let (params_off, codes_off, _) = cortiq_core::quant::q4tp_sections(rows, cols);
16910        let mut bytes: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
16911        let lo = cortiq_core::quant::f32_to_f16(-4.0);
16912        let step = cortiq_core::quant::f32_to_f16(0.1);
16913        for r in 0..rows {
16914            let o = params_off + r * 4;
16915            bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
16916            bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
16917        }
16918        let _ = codes_off;
16919        let xs: Vec<f32> = (0..b * cols)
16920            .map(|i| ((i % 97) as f32 - 48.0) / 48.0)
16921            .collect();
16922        let mut out = vec![0f32; b * rows];
16923        let pool = crate::pool::Pool::from_env();
16924        // A shared 48-core stand drifts ±25% run to run, which is wider
16925        // than any kernel change worth making. So: alternate the two
16926        // kernels inside one process and keep the BEST time for
16927        // each. Interleaving makes both see the same interference, and a
16928        // minimum is the one statistic another tenant cannot inflate.
16929        super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
16930        let reps: usize = std::env::var("CMF_BENCH_REPS")
16931            .ok()
16932            .and_then(|v| v.parse().ok())
16933            .unwrap_or(10);
16934        let mut best = [f64::MAX; 2];
16935        let mut sums = [0f32; 2];
16936        for _ in 0..reps {
16937            for (k, w) in [(0usize, 1u8), (1usize, 2u8)] {
16938                super::Q4TP_ALT.store(w, std::sync::atomic::Ordering::Relaxed);
16939                let t = std::time::Instant::now();
16940                super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
16941                best[k] = best[k].min(t.elapsed().as_secs_f64());
16942                sums[k] = out.iter().take(64).sum::<f32>();
16943            }
16944        }
16945        let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
16946        for (k, name) in ["previous", "tuned   "].iter().enumerate() {
16947            println!(
16948                "q4tp matmat {rows}x{cols} b={b} {name}: {:.1} ms  {:.1} GFLOP/s  (checksum {:.3})",
16949                best[k] * 1e3,
16950                flops / best[k] / 1e9,
16951                sums[k]
16952            );
16953        }
16954        assert!(
16955            (sums[0] - sums[1]).abs() < 1e-2,
16956            "the tuned kernel changed the result: {} vs {}",
16957            sums[0],
16958            sums[1]
16959        );
16960    }
16961
16962    /// The blocked kernel must agree with the per-column path exactly —
16963    /// same weights, same activation split, only a different instruction
16964    /// mix. Shapes are chosen to hit the awkward cases: a column count
16965    /// that leaves an odd group (the 512-bit kernel does two at a time),
16966    /// and a batch that does not divide by four.
16967    #[test]
16968    fn q4tp_matmat_blocked_matches_scalar() {
16969        use std::sync::atomic::Ordering::Relaxed;
16970        let _alt = super::Q4TP_ALT_TEST_LOCK
16971            .lock()
16972            .unwrap_or_else(|e| e.into_inner());
16973        // The last shape carries the image DiT's column count — 2304, so
16974        // 72 groups of accumulation, which is where a reordered sum can
16975        // actually drift — and runs through the thread pool, since the
16976        // blocked path splits rows across workers. Its row count stays
16977        // under 500k cells on purpose: above that, macOS diverts the whole
16978        // matmat to the Accelerate/AMX dequant sgemm and neither kernel
16979        // here would run.
16980        for &(rows, cols, b) in &[
16981            (64usize, 128usize, 7usize),
16982            (33, 96, 4),
16983            (16, 256, 9),
16984            (192, 2304, 37),
16985        ] {
16986            let total = cortiq_core::quant::expected_nbytes(
16987                cortiq_core::TensorDtype::Q4TiledP,
16988                &[rows, cols],
16989            )
16990            .unwrap();
16991            let (params_off, _, _) = cortiq_core::quant::q4tp_sections(rows, cols);
16992            let mut bytes: Vec<u8> = (0..total).map(|i| (i * 61 % 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 xs: Vec<f32> = (0..b * cols)
17001                .map(|i| ((i % 89) as f32 - 44.0) / 44.0)
17002                .collect();
17003            let mut got = vec![0f32; b * rows];
17004            let mut want = vec![0f32; b * rows];
17005            let gpr = cols / 32;
17006            let view = super::Q4tpView::new(&bytes, rows, cols);
17007            let pool = crate::pool::Pool::from_env();
17008            super::Q4TP_ALT.store(2, Relaxed);
17009            super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut got, pool.as_deref());
17010            super::Q4TP_ALT.store(1, Relaxed);
17011            super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut want, pool.as_deref());
17012            super::Q4TP_ALT.store(0, Relaxed);
17013            // Measured against the output's scale, not cell by cell: a
17014            // dot product of 2304 terms lands near zero wherever the row
17015            // and the activation nearly cancel, and there a per-cell
17016            // ratio reports 1e-3 for an absolute error of 5e-6 — f32's
17017            // own rounding, reordered. What must stay small is the error
17018            // relative to what the layer actually outputs.
17019            let scale = want.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
17020            let (mut worst, mut at) = (0f32, 0usize);
17021            for (i, (g, w)) in got.iter().zip(&want).enumerate() {
17022                if (g - w).abs() > worst {
17023                    worst = (g - w).abs();
17024                    at = i;
17025                }
17026            }
17027            assert!(
17028                worst <= 1e-4 * scale,
17029                "{rows}x{cols} b={b}: blocked and scalar disagree by {worst:.3e} \
17030                 (scale {scale:.3e}) at cell {at}: {} vs {}",
17031                got[at],
17032                want[at]
17033            );
17034
17035            // "Same speed, no quality loss" is a claim about which answer
17036            // is RIGHT, not about which two agree. Both paths sum the same
17037            // 2304 products in different orders, so f64 decides: the
17038            // blocked kernel keeps sixteen partial sums and folds them at
17039            // the end, which is a shallower addition tree than the
17040            // per-column path's running scalar, and it must not be worse.
17041            let (mut e_blocked, mut e_scalar) = (0f64, 0f64);
17042            for bi in 0..b {
17043                let act = super::split_act(&xs[bi * cols..(bi + 1) * cols]);
17044                for r in 0..rows {
17045                    let mut sc = vec![0f32; gpr];
17046                    view.scales_into(r, gpr, &mut sc);
17047                    let mut exact = 0f64;
17048                    for j in 0..cols {
17049                        let (w, sq) = super::q4tp_outlier(view.nib, r, gpr, j, &sc);
17050                        exact += w as f64 * sq as f64 * act.xq[j] as f64;
17051                    }
17052                    exact *= act.sx as f64;
17053                    for &(j, xv) in &act.outliers {
17054                        let (w, sq) = super::q4tp_outlier(view.nib, r, gpr, j, &sc);
17055                        exact += w as f64 * sq as f64 * xv as f64;
17056                    }
17057                    let i = bi * rows + r;
17058                    e_blocked = e_blocked.max((got[i] as f64 - exact).abs());
17059                    e_scalar = e_scalar.max((want[i] as f64 - exact).abs());
17060                }
17061            }
17062            println!(
17063                "{rows}x{cols} b={b}: worst error vs f64 — blocked {e_blocked:.3e}, \
17064                 per-column {e_scalar:.3e}"
17065            );
17066            // An absolute bar, not a race between the two: at these
17067            // magnitudes both sit in f32's last bits, and on a small shape
17068            // whichever one happens to round the unluckiest cell "wins" by
17069            // a factor the next seed reverses.
17070            assert!(
17071                e_blocked <= 1e-5 * scale as f64 && e_scalar <= 1e-5 * scale as f64,
17072                "{rows}x{cols} b={b}: error against f64 too large — blocked \
17073                 {e_blocked:.3e}, per-column {e_scalar:.3e}, scale {scale:.3e}"
17074            );
17075        }
17076    }
17077
17078    /// What the row-exact contract costs a speculative-verify panel on the
17079    /// host: the same batch through the fast arms, through the row-exact
17080    /// arms, and as one matvec per token (the other way to be exact).
17081    /// Arms alternate inside one process and keep their best time.
17082    /// `CMF_GPU=0 cargo test -p cortiq-engine --release q4tp_matmat_row_exact_cost -- --ignored --nocapture`
17083    #[test]
17084    #[ignore]
17085    fn q4tp_matmat_row_exact_cost() {
17086        let _alt = super::Q4TP_ALT_TEST_LOCK
17087            .lock()
17088            .unwrap_or_else(|e| e.into_inner());
17089        let pool = crate::pool::Pool::from_env();
17090        let reps: usize = std::env::var("CMF_BENCH_REPS")
17091            .ok()
17092            .and_then(|v| v.parse().ok())
17093            .unwrap_or(30);
17094        for &(rows, cols) in &[(2048usize, 4096usize), (4096, 2048), (4096, 4096)] {
17095            let total = cortiq_core::quant::expected_nbytes(
17096                cortiq_core::TensorDtype::Q4TiledP,
17097                &[rows, cols],
17098            )
17099            .unwrap();
17100            let (params_off, _, _) = cortiq_core::quant::q4tp_sections(rows, cols);
17101            let mut bytes: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
17102            let lo = cortiq_core::quant::f32_to_f16(-4.0);
17103            let step = cortiq_core::quant::f32_to_f16(0.1);
17104            for r in 0..rows {
17105                let o = params_off + r * 4;
17106                bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
17107                bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
17108            }
17109            for &b in &[2usize, 4, 5, 8] {
17110                let xs: Vec<f32> = (0..b * cols)
17111                    .map(|i| ((i % 97) as f32 - 48.0) / 48.0)
17112                    .collect();
17113                let mut out = vec![0f32; b * rows];
17114                let mut best = [f64::MAX; 3];
17115                for _ in 0..reps {
17116                    for (k, best_k) in best.iter_mut().enumerate() {
17117                        let t = std::time::Instant::now();
17118                        match k {
17119                            0 | 1 => super::q4tp_matmat_with(
17120                                &bytes,
17121                                &xs,
17122                                b,
17123                                rows,
17124                                cols,
17125                                &mut out,
17126                                pool.as_deref(),
17127                                k == 1,
17128                            ),
17129                            _ => {
17130                                for (bi, o) in out.chunks_mut(rows).enumerate() {
17131                                    super::q4tp_matvec(
17132                                        &bytes,
17133                                        &xs[bi * cols..(bi + 1) * cols],
17134                                        rows,
17135                                        cols,
17136                                        o,
17137                                        pool.as_deref(),
17138                                    );
17139                                }
17140                            }
17141                        }
17142                        *best_k = best_k.min(t.elapsed().as_secs_f64());
17143                    }
17144                }
17145                println!(
17146                    "q4tp {rows}x{cols} b={b}: fast {:.3} ms, row-exact {:.3} ms, \
17147                     {b} matvecs {:.3} ms",
17148                    best[0] * 1e3,
17149                    best[1] * 1e3,
17150                    best[2] * 1e3
17151                );
17152            }
17153        }
17154    }
17155
17156    /// The q4t twin of the throughput bench, same shape and rules, so the
17157    /// two quantisations' batch kernels can be read against each other.
17158    /// `cargo test -p cortiq-engine --release q4t_matmat_throughput -- --ignored --nocapture`
17159    #[test]
17160    #[ignore]
17161    fn q4t_matmat_throughput() {
17162        let (rows, cols, b) = (9216usize, 2304usize, 296usize);
17163        let total =
17164            cortiq_core::quant::expected_nbytes(cortiq_core::TensorDtype::Q4Tiled, &[rows, cols])
17165                .unwrap();
17166        // q4t carries a per-group f16 scale in the tile's first two bytes;
17167        // random bytes there decode to inf and the bench would time NaNs.
17168        let mut bytes: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
17169        let sc = cortiq_core::quant::f32_to_f16(0.02);
17170        for t in bytes.chunks_mut(super::Q4_TILE) {
17171            t[..2].copy_from_slice(&sc.to_le_bytes());
17172        }
17173        let xs: Vec<f32> = (0..b * cols)
17174            .map(|i| ((i % 97) as f32 - 48.0) / 48.0)
17175            .collect();
17176        let mut out = vec![0f32; b * rows];
17177        let pool = crate::pool::Pool::from_env();
17178        super::q4t_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
17179        let reps: usize = std::env::var("CMF_BENCH_REPS")
17180            .ok()
17181            .and_then(|v| v.parse().ok())
17182            .unwrap_or(10);
17183        let mut best = f64::MAX;
17184        for _ in 0..reps {
17185            let t = std::time::Instant::now();
17186            super::q4t_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
17187            best = best.min(t.elapsed().as_secs_f64());
17188        }
17189        let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
17190        println!(
17191            "q4t matmat {rows}x{cols} b={b}: {:.1} ms  {:.1} GFLOP/s  (checksum {:.3})",
17192            best * 1e3,
17193            flops / best / 1e9,
17194            out.iter().take(64).sum::<f32>()
17195        );
17196    }
17197}
17198
17199#[cfg(all(test, target_arch = "x86_64"))]
17200mod x86_gemm_tests {
17201    use super::x86_gemm_rm;
17202
17203    /// The AVX2 two-row FMA q4tp batch kernel against the row-exact
17204    /// vertical cell (the matvec's bits): odd row counts (the last row
17205    /// takes the vertical cell), columns past a group of four, a wide
17206    /// row of 2560 columns (80 groups), outliers in the activations, and
17207    /// the pool. Measured against the output's scale, as the blocked
17208    /// kernel's test does.
17209    #[test]
17210    fn q4tp_fma_batch_matches_the_row_exact_cell() {
17211        if !super::q4tp_fma_x86() || !super::a8w8_enabled() {
17212            return;
17213        }
17214        let pool = crate::pool::Pool::from_env();
17215        for &(rows, cols, b) in &[(33usize, 96usize, 7usize), (64, 256, 13), (257, 2560, 37)] {
17216            let total = cortiq_core::quant::expected_nbytes(
17217                cortiq_core::TensorDtype::Q4TiledP,
17218                &[rows, cols],
17219            )
17220            .unwrap();
17221            let (params_off, _, _) = cortiq_core::quant::q4tp_sections(rows, cols);
17222            let mut bytes: Vec<u8> = (0..total).map(|i| (i * 61 % 251) as u8).collect();
17223            let lo = cortiq_core::quant::f32_to_f16(-4.0);
17224            let step = cortiq_core::quant::f32_to_f16(0.1);
17225            for r in 0..rows {
17226                let o = params_off + r * 4;
17227                bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
17228                bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
17229            }
17230            let mut xs: Vec<f32> = (0..b * cols)
17231                .map(|i| ((i * 37 % 89) as f32 - 44.0) / 44.0)
17232                .collect();
17233            for bi in (0..b).step_by(3) {
17234                xs[bi * cols + (bi * 7) % cols] = 40.0; // an outlier column
17235            }
17236            let mut got = vec![0f32; b * rows];
17237            let mut want = vec![0f32; b * rows];
17238            super::q4tp_matmat_with(&bytes, &xs, b, rows, cols, &mut got, pool.as_deref(), false);
17239            super::q4tp_matmat_with(&bytes, &xs, b, rows, cols, &mut want, pool.as_deref(), true);
17240            let scale = want.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
17241            let worst = got.iter().zip(&want).map(|(g, w)| (g - w).abs()).fold(0f32, f32::max);
17242            assert!(
17243                worst <= 1e-5 * scale,
17244                "{rows}x{cols} b={b}: FMA batch and row-exact cell disagree by {worst:.3e} (scale {scale:.3e})"
17245            );
17246        }
17247    }
17248
17249    #[allow(clippy::too_many_arguments)]
17250    fn naive(
17251        m: usize,
17252        n: usize,
17253        k: usize,
17254        alpha: f32,
17255        a: &[f32],
17256        lda: usize,
17257        b: &[f32],
17258        ldb: usize,
17259        b_rows_are_n: bool,
17260    ) -> Vec<f64> {
17261        let mut c = vec![0f64; m * n];
17262        for i in 0..m {
17263            for j in 0..n {
17264                let mut acc = 0f64;
17265                for p in 0..k {
17266                    let bv = if b_rows_are_n { b[j * ldb + p] } else { b[p * ldb + j] };
17267                    acc += a[i * lda + p] as f64 * bv as f64;
17268                }
17269                c[i * n + j] = acc * alpha as f64;
17270            }
17271        }
17272        c
17273    }
17274
17275    /// The AVX2 micro-GEMM against a double-precision triple loop: both
17276    /// B layouts, padded leading dimensions, shapes that exercise the
17277    /// 4×16 kernel, the row and column tails, and the k blocking (k >
17278    /// 256 resumes the accumulators from C).
17279    #[test]
17280    fn x86_gemm_matches_naive_on_odd_shapes() {
17281        for &(m, n, k) in &[
17282            (1usize, 1usize, 1usize),
17283            (4, 16, 256),
17284            (7, 37, 19),
17285            (9, 300, 257),
17286            (64, 1000, 256),
17287            (5, 256, 1001),
17288            (13, 129, 512),
17289            (8, 16, 0),
17290        ] {
17291            for &bt in &[false, true] {
17292                let (lda, ldb, ldc) = (k + 3, if bt { k + 2 } else { n + 5 }, n + 1);
17293                let a: Vec<f32> = (0..m * lda)
17294                    .map(|i| ((i * 7919) % 1000) as f32 / 500.0 - 1.0)
17295                    .collect();
17296                let brows = if bt { n } else { k.max(1) };
17297                let b: Vec<f32> = (0..brows * ldb)
17298                    .map(|i| ((i * 104_729) % 997) as f32 / 498.5 - 1.0)
17299                    .collect();
17300                let mut c = vec![f32::NAN; m * ldc];
17301                x86_gemm_rm(m, n, k, 0.37, &a, lda, &b, ldb, bt, &mut c, ldc);
17302                let r = naive(m, n, k, 0.37, &a, lda, &b, ldb, bt);
17303                for i in 0..m {
17304                    for j in 0..n {
17305                        let (x, y) = (c[i * ldc + j] as f64, r[i * n + j]);
17306                        assert!(
17307                            (x - y).abs() <= 1e-3 * (1.0 + y.abs()),
17308                            "m{m} n{n} k{k} bt{bt} [{i},{j}]: {x} vs {y}"
17309                        );
17310                    }
17311                    assert!(c[i * ldc + n].is_nan(), "m{m} n{n} k{k}: padding column written");
17312                }
17313            }
17314        }
17315    }
17316}