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::sync::Arc;
21
22pub enum QTensor {
23    F32 {
24        data: Vec<f32>,
25        rows: usize,
26        cols: usize,
27    },
28    Mapped {
29        model: Arc<CmfModel>,
30        /// Index into the model's tensor directory.
31        idx: usize,
32        dtype: TensorDtype,
33        rows: usize,
34        cols: usize,
35        /// Per-row scales, dequantized to f32 up front (tiny).
36        row_scale: Vec<f32>,
37        /// q8_2f column field (θ), dequantized up front; empty for q8_row.
38        col_field: Vec<f32>,
39        /// Vbit only: byte offset of each row's packed data within the
40        /// tensor blob (`[rows + 1]`, computed once at load — the per-
41        /// matvec prefix scan over row bit-widths was O(rows) each call).
42        vbit_offsets: Vec<usize>,
43        /// q8-family decode repack (load-time, optional): rows in groups
44        /// of 4, interleaved in 16-byte units — one 64-byte line per
45        /// iteration feeds all 4 sdot lanes, ONE sequential weight
46        /// stream per worker instead of four (this is where llama.cpp's
47        /// repacked Q8 kernels get their bandwidth). Empty = off
48        /// (CMF_REPACK=0, non-SDOT arch, or an ineligible shape). Trades
49        /// an anonymous copy of the quants for mmap pages that go cold.
50        repack: Vec<u8>,
51    },
52}
53
54/// Load-time q8 repack gate (see `Mapped::repack`). OPT-IN
55/// (`CMF_REPACK=1`): the single-stream hypothesis LOST on Apple Silicon
56/// (M4, interleaved A/B: decode 101 vs 94 tok/s — four adjacent row
57/// streams per worker feed the prefetcher MORE memory-level parallelism
58/// than one); kept as an experiment flag for x86, where the tradeoff
59/// may land differently.
60fn repack_enabled() -> bool {
61    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
62    *ON.get_or_init(|| {
63        std::env::var("CMF_REPACK")
64            .map(|v| v == "1")
65            .unwrap_or(cfg!(target_os = "android"))
66    })
67}
68
69/// Interleave q8 rows for the decode kernel: group g holds rows
70/// 4g..4g+4 as [r0[c], r1[c], r2[c], r3[c]] per 16-byte chunk c. Only
71/// full groups are packed — tail rows keep reading the mmap layout.
72fn q8_repack(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
73    #[cfg(target_arch = "aarch64")]
74    let arch_ok = sdot_enabled();
75    #[cfg(not(target_arch = "aarch64"))]
76    let arch_ok = false;
77    if !arch_ok || !repack_enabled() || rows < 256 || cols % 16 != 0 {
78        return Vec::new();
79    }
80    q8_repack_layout(bytes, rows, cols)
81}
82
83/// The pure layout transform behind `q8_repack` (tested directly —
84/// the gate depends on arch and env).
85fn q8_repack_layout(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
86    let groups = rows / 4;
87    let mut rep = vec![0u8; groups * 4 * cols];
88    for g in 0..groups {
89        let dst = &mut rep[g * 4 * cols..(g + 1) * 4 * cols];
90        for c in 0..cols / 16 {
91            for lane in 0..4 {
92                let src = (g * 4 + lane) * cols + c * 16;
93                dst[c * 64 + lane * 16..c * 64 + lane * 16 + 16]
94                    .copy_from_slice(&bytes[src..src + 16]);
95            }
96        }
97    }
98    rep
99}
100
101/// Prefix-sum of vbit row payload offsets (absolute within the tensor
102/// bytes). `offsets[r]..offsets[r+1]` is row r's packed data.
103fn vbit_row_offsets(bytes: &[u8], rows: usize, cols: usize) -> Vec<usize> {
104    let ng = cols / GROUP_SIZE;
105    let bits = &bytes[..rows];
106    let mut offsets = Vec::with_capacity(rows + 1);
107    let mut off = rows + rows * ng * 2;
108    for r in 0..rows {
109        offsets.push(off);
110        off += (cols * bits[r] as usize).div_ceil(8);
111    }
112    offsets.push(off);
113    offsets
114}
115
116/// `CMF_X86_BLOCKED` / `CMF_GPU_LMHEAD` / `CMF_GPU_SPLIT`, read once. They
117/// used to be read from the environment on every large matvec and on every
118/// matmat in six places — microseconds each, but also a knob that could
119/// change under a running process, which is not a thing a kernel choice
120/// should be able to do mid-sequence.
121fn blocked_enabled() -> bool {
122    use std::sync::atomic::Ordering::Relaxed;
123    match BLOCKED_OVERRIDE.load(Relaxed) {
124        1 => false,
125        2 => true,
126        _ => {
127            static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
128            *ON.get_or_init(|| {
129                std::env::var("CMF_X86_BLOCKED")
130                    .map(|v| v != "0")
131                    .unwrap_or(true)
132            })
133        }
134    }
135}
136
137static BLOCKED_OVERRIDE: std::sync::atomic::AtomicU8 = std::sync::atomic::AtomicU8::new(0);
138
139/// Force the blocked GEMM on or off, ignoring the environment; `None`
140/// restores it. For tests that need to run BOTH paths and compare them:
141/// `blocked_enabled` caches its answer for the life of the process, which
142/// is right when the environment is the only input, but leaves a test that
143/// flips `CMF_X86_BLOCKED` between two calls comparing a path against
144/// itself — or against whatever a test running in parallel latched first.
145pub fn set_blocked_override(on: Option<bool>) {
146    let v = match on {
147        None => 0,
148        Some(false) => 1,
149        Some(true) => 2,
150    };
151    BLOCKED_OVERRIDE.store(v, std::sync::atomic::Ordering::Relaxed);
152}
153
154fn gpu_lmhead_enabled() -> bool {
155    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
156    *ON.get_or_init(|| {
157        std::env::var("CMF_GPU_LMHEAD")
158            .map(|v| v != "0")
159            .unwrap_or(true)
160    })
161}
162
163fn gpu_split_frac() -> f32 {
164    static F: std::sync::OnceLock<f32> = std::sync::OnceLock::new();
165    *F.get_or_init(|| {
166        std::env::var("CMF_GPU_SPLIT")
167            .ok()
168            .and_then(|v| v.parse::<f32>().ok())
169            .unwrap_or(0.5)
170            .clamp(0.0, 1.0)
171    })
172}
173
174impl QTensor {
175    pub fn from_f32(data: Vec<f32>, rows: usize, cols: usize) -> Self {
176        debug_assert_eq!(data.len(), rows * cols);
177        Self::F32 { data, rows, cols }
178    }
179
180    /// Wrap a directory tensor without dequantizing the payload.
181    /// Falls back to dequantized f32 for dtypes without a fused kernel.
182    pub fn from_model(model: &Arc<CmfModel>, name: &str) -> Result<Self, String> {
183        // Indexed lookup: the linear directory scan made pipeline build
184        // O(N²) on MoE/skills files with thousands of tensors.
185        let idx = model
186            .tensor_index(name)
187            .ok_or_else(|| format!("tensor '{name}' not found in CMF directory"))?;
188        let entry = &model.tensors[idx];
189        if entry.shape.len() != 2 {
190            return Err(format!("QTensor::from_model needs 2-D, got '{name}'"));
191        }
192        let (rows, cols) = (entry.shape[0], entry.shape[1]);
193        let bytes = model.entry_bytes(entry);
194
195        match entry.dtype {
196            TensorDtype::Q8Row | TensorDtype::Q8_2f => {
197                let n = rows * cols;
198                let scales_off = n;
199                let row_scale: Vec<f32> = (0..rows)
200                    .map(|o| {
201                        f16_to_f32(u16::from_le_bytes([
202                            bytes[scales_off + o * 2],
203                            bytes[scales_off + o * 2 + 1],
204                        ]))
205                    })
206                    .collect();
207                let col_field: Vec<f32> = if entry.dtype == TensorDtype::Q8_2f {
208                    let col_off = n + rows * 2;
209                    (0..cols)
210                        .map(|i| {
211                            f16_to_f32(u16::from_le_bytes([
212                                bytes[col_off + i * 2],
213                                bytes[col_off + i * 2 + 1],
214                            ]))
215                        })
216                        .collect()
217                } else {
218                    Vec::new()
219                };
220                Ok(Self::Mapped {
221                    model: model.clone(),
222                    idx,
223                    dtype: entry.dtype,
224                    rows,
225                    cols,
226                    row_scale,
227                    col_field,
228                    vbit_offsets: Vec::new(),
229                    repack: q8_repack(bytes, rows, cols),
230                })
231            }
232            // vbit: fused kernel unpacks variable-bit rows from mmap.
233            TensorDtype::Vbit if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
234                model: model.clone(),
235                idx,
236                dtype: entry.dtype,
237                rows,
238                cols,
239                row_scale: Vec::new(),
240                col_field: Vec::new(),
241                vbit_offsets: vbit_row_offsets(bytes, rows, cols),
242                repack: Vec::new(),
243            }),
244            // vbit_ro (§4.2): the offset table comes straight from the
245            // file — no load-time prefix scan; kernels are shared with
246            // legacy vbit (they consume absolute offsets either way).
247            TensorDtype::VbitRo if cols % GROUP_SIZE == 0 => {
248                let (_, off_off, packed_off) = cortiq_core::quant::vbit_ro_sections(rows, cols);
249                let offsets: Vec<usize> = (0..=rows)
250                    .map(|r| packed_off + cortiq_core::quant::vbit_ro_offset(bytes, off_off, r))
251                    .collect();
252                Ok(Self::Mapped {
253                    model: model.clone(),
254                    idx,
255                    dtype: entry.dtype,
256                    rows,
257                    cols,
258                    row_scale: Vec::new(),
259                    col_field: Vec::new(),
260                    vbit_offsets: offsets,
261                    repack: Vec::new(),
262                })
263            }
264            // q4_block: fused kernel reads nibbles straight from mmap —
265            // a 14B q4 file no longer explodes into ×8 f32 RAM.
266            // q4_tiled (§4.3): interleaved [scale][nibbles] tiles — one
267            // sequential memory stream (measured ×1.66 ARM / ×1.13 AVX2
268            // at kernel level over the split layout).
269            TensorDtype::Q4Tiled if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
270                model: model.clone(),
271                idx,
272                dtype: entry.dtype,
273                rows,
274                cols,
275                row_scale: Vec::new(),
276                col_field: Vec::new(),
277                vbit_offsets: Vec::new(),
278                repack: Vec::new(),
279            }),
280            // q4tp (§4.10): nibbles from mmap, scale from the row ladder —
281            // 7.3% less file than q4t at the same 4-bit grid.
282            TensorDtype::Q4TiledP if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
283                model: model.clone(),
284                idx,
285                dtype: entry.dtype,
286                rows,
287                cols,
288                row_scale: Vec::new(),
289                col_field: Vec::new(),
290                vbit_offsets: Vec::new(),
291                repack: Vec::new(),
292            }),
293            // q2tp: 2-bit chunks from mmap, scale from the same row ladder.
294            TensorDtype::Q2TiledP if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
295                model: model.clone(),
296                idx,
297                dtype: entry.dtype,
298                rows,
299                cols,
300                row_scale: Vec::new(),
301                col_field: Vec::new(),
302                vbit_offsets: Vec::new(),
303                repack: Vec::new(),
304            }),
305            TensorDtype::Q4Block if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
306                model: model.clone(),
307                idx,
308                dtype: entry.dtype,
309                rows,
310                cols,
311                row_scale: Vec::new(),
312                col_field: Vec::new(),
313                vbit_offsets: Vec::new(),
314                repack: Vec::new(),
315            }),
316            // q1: binary sign-bit tiles from mmap (1-bit-trained models).
317            TensorDtype::Q1 if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
318                model: model.clone(),
319                idx,
320                dtype: entry.dtype,
321                rows,
322                cols,
323                row_scale: Vec::new(),
324                col_field: Vec::new(),
325                vbit_offsets: Vec::new(),
326                repack: Vec::new(),
327            }),
328            // q1t (ternary + outlier overlay): fused per-row dequant kernel
329            // reads straight from mmap — a 12B q1t stays ~its file size in
330            // RAM instead of dequantizing to ~48 GB of f32.
331            TensorDtype::Q1T 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            // No fused kernel yet → dequantize once (correct, more RAM).
343            _ => {
344                let mut data = vec![0.0f32; rows * cols];
345                cortiq_core::quant::dequant_tensor(entry, bytes, &mut data)?;
346                Ok(Self::from_f32(data, rows, cols))
347            }
348        }
349    }
350
351    /// q1-mapped tensor? (GPU gates: the q1 CPU kernel is
352    /// compute-bound, so offload pays at much smaller shapes than q8.)
353    pub(crate) fn is_q1(&self) -> bool {
354        matches!(
355            self,
356            Self::Mapped {
357                dtype: TensorDtype::Q1,
358                ..
359            }
360        )
361    }
362
363    /// Owned-f32 view (data, rows, cols) — the GDN a/b gate projections
364    /// arrive dequantized (force-f16 in the converter → F32 in RAM).
365    pub(crate) fn f32_parts(&self) -> Option<(&[f32], usize, usize)> {
366        match self {
367            Self::F32 { data, rows, cols } => Some((data, *rows, *cols)),
368            _ => None,
369        }
370    }
371
372    /// (directory idx, rows, cols) of a q1-mapped tensor — the
373    /// whole-block GPU path resolves offsets itself.
374    /// (idx, rows, cols) of a mapped tensor the whole-token GPU graph can drive
375    /// — Q1, Q1T or Q4-block (it resolves the offset and picks the kernel by
376    /// dtype). Q4-block lets a precise down_proj/lm_head stay on-device.
377    /// Named `q1_parts` for historical reasons.
378    pub(crate) fn q1_parts(&self) -> Option<(usize, usize, usize)> {
379        if self.has_prism_contract() {
380            return None;
381        }
382        match self {
383            #[cfg(target_os = "macos")]
384            Self::Mapped {
385                dtype: TensorDtype::Q1T,
386                ..
387            } if !crate::gpu::metal_q1t_enabled() => None,
388            Self::Mapped {
389                idx,
390                dtype:
391                    TensorDtype::Q1
392                    | TensorDtype::Q1T
393                    | TensorDtype::Q4Block
394                    | TensorDtype::Q4Tiled
395                    // Q2TiledP deliberately absent: the Metal graph has no
396                    // q2tp kernel, and advertising it here made the block
397                    // plan truncate mid-run at the first q2tp layer.
398                    | TensorDtype::Q4TiledP
399                    | TensorDtype::Q8Row
400                    | TensorDtype::Q8_2f,
401                rows,
402                cols,
403                ..
404            } => Some((*idx, *rows, *cols)),
405            _ => None,
406        }
407    }
408
409    /// `(directory idx, rows, cols)` for the native Metal token graph.  The
410    /// historical q1 graph gate intentionally refuses every Prism tensor so
411    /// an untransformed q2 payload cannot slip into the resident path.  The
412    /// Metal2 graph is descriptor-aware and admits only the production
413    /// q2tp-affine forward targets; ordinary q1/q4 callers retain the old
414    /// `q1_parts` behaviour.
415    #[cfg(target_os = "macos")]
416    pub(crate) fn metal_graph_parts(&self) -> Option<(usize, usize, usize)> {
417        if let Some((model, idx, kind, _)) = self.graph_weight_descriptor() {
418            let name = &model.tensors[idx].name;
419            let forward = kind == 9 && crate::prism::is_forward_weight(model, name);
420            let affine = kind == 9 && crate::prism::is_affine_target(model, name);
421            if forward && affine {
422                let e = model.tensors.get(idx)?;
423                return Some((idx, *e.shape.first()?, *e.shape.get(1)?));
424            }
425        }
426        self.q1_parts()
427    }
428
429    /// (directory idx, rows, cols) of a q4_tiled mapped tensor. The
430    /// chunk-prefill graph takes it in the same 4-tuple slot as
431    /// `q8_row_parts` with an EMPTY row_scale — q4t carries its scales
432    /// inside the 18-byte tiles, and the empty slice is what tells the
433    /// encoder to reach for the q4t kernels.
434    pub(crate) fn q4t_parts(&self) -> Option<(usize, usize, usize)> {
435        if self.has_prism_contract() {
436            return None;
437        }
438        match self {
439            Self::Mapped {
440                idx,
441                dtype: TensorDtype::Q4Tiled,
442                rows,
443                cols,
444                ..
445            } => Some((*idx, *rows, *cols)),
446            _ => None,
447        }
448    }
449
450    /// (directory idx, rows, cols) of a q4tp mapped tensor. Same empty-scale
451    /// slot as `q4t_parts` in the chunk graph — the encoder tells the two
452    /// apart by the tensor's dtype, not by the slot.
453    pub(crate) fn q4tp_parts(&self) -> Option<(usize, usize, usize)> {
454        if self.has_prism_contract() {
455            return None;
456        }
457        match self {
458            Self::Mapped {
459                idx,
460                dtype: TensorDtype::Q4TiledP,
461                rows,
462                cols,
463                ..
464            } => Some((*idx, *rows, *cols)),
465            _ => None,
466        }
467    }
468
469    /// (directory idx, rows, cols, row_scale) of a plain q8_row mapped
470    /// tensor — the chunk-prefill GPU graph resolves offsets itself.
471    /// q8_2f is excluded on purpose: its column field would need a
472    /// prescale stage on the device.
473    pub(crate) fn q8_row_parts(&self) -> Option<(usize, usize, usize, &[f32])> {
474        if self.has_prism_contract() {
475            return None;
476        }
477        match self {
478            Self::Mapped {
479                idx,
480                dtype: TensorDtype::Q8Row,
481                rows,
482                cols,
483                row_scale,
484                col_field,
485                ..
486            } if col_field.is_empty() => Some((*idx, *rows, *cols, row_scale)),
487            _ => None,
488        }
489    }
490
491    /// The layout this tensor is stored in, when it is mapped from a model.
492    /// The frames branch on it — a q2tp gate against a q4tp down is a real
493    /// combination in the 2-bit profile and needs a different kernel.
494    pub fn model_dtype(&self) -> Option<cortiq_core::TensorDtype> {
495        match self {
496            Self::Mapped { dtype, .. } => Some(*dtype),
497            _ => None,
498        }
499    }
500
501    /// The tensor's index in the model directory, when it is mapped from one.
502    /// The GPU frames bind by index rather than by name — a name lookup per
503    /// layer per token is not free, and the index is what the device cache is
504    /// keyed on anyway.
505    pub fn model_idx(&self) -> Option<usize> {
506        match self {
507            Self::Mapped { idx, .. } => Some(*idx),
508            _ => None,
509        }
510    }
511
512    /// The model this tensor is mapped from, when it is mapped at all. The
513    /// GPU frames need the container to reach the bytes; a QTensor already
514    /// holds it, and threading a second handle down every call site to say
515    /// the same thing invites the two to disagree.
516    pub fn model_arc(&self) -> Option<std::sync::Arc<cortiq_core::CmfModel>> {
517        match self {
518            Self::Mapped { model, .. } => Some(model.clone()),
519            _ => None,
520        }
521    }
522
523    /// Whether this mapped tensor belongs to the Prism/Bonsai transform
524    /// contract.  Device graphs do not carry the descriptor, so callers use
525    /// this conservative predicate to stay on the descriptor-aware CPU path
526    /// instead of silently executing an unrotated matrix.
527    pub(crate) fn has_prism_contract(&self) -> bool {
528        matches!(self, Self::Mapped { model, .. } if crate::prism::has_contract(model))
529    }
530
531    pub fn rows(&self) -> usize {
532        match self {
533            Self::F32 { rows, .. } | Self::Mapped { rows, .. } => *rows,
534        }
535    }
536
537    /// Mapped q4t handle (model + directory index) — the fused GPU FFN
538    /// needs the raw file coordinates of its three projections.
539    pub(crate) fn mapped_q4t(&self) -> Option<(&Arc<CmfModel>, usize)> {
540        if self.has_prism_contract() {
541            return None;
542        }
543        match self {
544            Self::Mapped {
545                model,
546                idx,
547                dtype: TensorDtype::Q4Tiled,
548                ..
549            } => Some((model, *idx)),
550            _ => None,
551        }
552    }
553
554    /// Same slot as `mapped_q4t` for a q4tp tensor — the fused DiT FFN picks
555    /// its kernels by which of the two answers.
556    pub fn mapped_q4tp(&self) -> Option<(&Arc<CmfModel>, usize)> {
557        if self.has_prism_contract() {
558            return None;
559        }
560        match self {
561            Self::Mapped {
562                model,
563                idx,
564                dtype: TensorDtype::Q4TiledP,
565                ..
566            } => Some((model, *idx)),
567            _ => None,
568        }
569    }
570
571    /// (model, tensor idx) for a mapped weight in ANY codec the fused device
572    /// paths can run — four-bit tiled or either int8 layout.
573    ///
574    /// The fused DiT chains asked for `mapped_q4tp` by name, so an eight-bit
575    /// container never reached them and rendered through per-op GEMMs even
576    /// after those kernels learned its codec. The gate is what the codec has
577    /// a device GEMM for, not which codec it is.
578    pub fn mapped_device_gemm(&self) -> Option<(&Arc<CmfModel>, usize)> {
579        if self.has_prism_contract() {
580            return None;
581        }
582        match self {
583            Self::Mapped {
584                model,
585                idx,
586                dtype: TensorDtype::Q4TiledP | TensorDtype::Q8Row | TensorDtype::Q8_2f,
587                ..
588            } => Some((model, *idx)),
589            _ => None,
590        }
591    }
592
593    /// (model, tensor idx) for a q2tp mapped weight — the 2-bit twin of
594    /// `mapped_q4tp`, used by the mixed MoE profile.
595    pub fn mapped_q2tp(&self) -> Option<(&Arc<CmfModel>, usize)> {
596        if self.has_prism_contract() {
597            return None;
598        }
599        match self {
600            Self::Mapped {
601                model,
602                idx,
603                dtype: TensorDtype::Q2TiledP,
604                ..
605            } => Some((model, *idx)),
606            _ => None,
607        }
608    }
609
610    pub fn cols(&self) -> usize {
611        match self {
612            Self::F32 { cols, .. } | Self::Mapped { cols, .. } => *cols,
613        }
614    }
615
616    /// (model, tensor idx) for a q1 mapped weight — the wgpu token graph
617    /// keys its resident VRAM cache by idx. None for any other dtype/kind.
618    pub fn mapped_q1(&self) -> Option<(&std::sync::Arc<CmfModel>, usize)> {
619        if self.has_prism_contract() {
620            return None;
621        }
622        match self {
623            Self::Mapped {
624                model,
625                idx,
626                dtype: TensorDtype::Q1,
627                ..
628            } => Some((model, *idx)),
629            _ => None,
630        }
631    }
632
633    /// (model, idx, kind, row_scale) for a graph-capable mapped weight.
634    /// kind: 0=q8_row (per-row scales), 1=q1, 2=q4_block, 3=q1t
635    /// (tile-embedded, no rs), 5=q4_tiled, 6=q4tp, 7=q8_2f (both scale
636    /// planes live inside the tensor). None only for `vbit`.
637    ///
638    /// The old comment here claimed q4_block was unhandled while the arm
639    /// right below mapped it, and it named q8_2f as unhandled after that
640    /// stopped being true — a stale comment on this function is how a
641    /// model silently loses the graph, so it is worth keeping honest.
642    pub fn graph_weight(&self) -> Option<(&std::sync::Arc<CmfModel>, usize, u8, &[f32])> {
643        if self.has_prism_contract() {
644            return None;
645        }
646        self.graph_weight_descriptor()
647    }
648
649    /// Descriptor-aware graph handle used only by the Prism token graph.
650    /// Ordinary graph callers continue to use [`graph_weight`] and therefore
651    /// remain fail-closed until they provide the same explicit transform
652    /// contract.
653    pub(crate) fn graph_weight_descriptor(
654        &self,
655    ) -> Option<(&std::sync::Arc<CmfModel>, usize, u8, &[f32])> {
656        match self {
657            Self::Mapped {
658                model,
659                idx,
660                dtype: TensorDtype::Q8Row,
661                row_scale,
662                ..
663            } => Some((model, *idx, 0, row_scale.as_slice())),
664            Self::Mapped {
665                model,
666                idx,
667                dtype: TensorDtype::Q1,
668                ..
669            } => Some((model, *idx, 1, &[])),
670            // Q4Tiled is kind 5, NOT 2: both carried 2 historically, and
671            // the wgpu token graph fed 18B interleaved tiles to the
672            // split-layout q4b kernel — garbage output on q4t models
673            // (caught by an end-to-end answer check on real Vulkan).
674            Self::Mapped {
675                model,
676                idx,
677                dtype: TensorDtype::Q4Tiled,
678                ..
679            } => Some((model, *idx, 5, &[])),
680            // Kind 6, not 5: q4tp's nibble stride and scale planes differ,
681            // and feeding them to the q4t kernel is exactly the mistake that
682            // produced garbage when Q4Tiled shared kind 2 with Q4Block.
683            Self::Mapped {
684                model,
685                idx,
686                dtype: TensorDtype::Q4TiledP,
687                ..
688            } => Some((model, *idx, 6, &[])),
689            Self::Mapped {
690                model,
691                idx,
692                dtype: TensorDtype::Q4Block,
693                ..
694            } => Some((model, *idx, 2, &[])),
695            // q8_2f carries BOTH scale planes after the int8 body (rows
696            // f16, then cols f16), so the graph takes the whole tensor
697            // and the kernel reads them where they lie — no host-side
698            // prescale, which is what the per-op path does instead.
699            Self::Mapped {
700                model,
701                idx,
702                dtype: TensorDtype::Q8_2f,
703                ..
704            } => Some((model, *idx, 7, &[])),
705            Self::Mapped {
706                model,
707                idx,
708                dtype: TensorDtype::Q1T,
709                ..
710            } => Some((model, *idx, 3, &[])),
711            // Kind 9: the 2-bit plane on the q4tp ladder (dense FFN gate/up
712            // of the q2tp profile). Its own kernel — 8 bytes a group where
713            // q4tp has 16, and rung 0 is the exact zero.
714            Self::Mapped {
715                model,
716                idx,
717                dtype: TensorDtype::Q2TiledP,
718                ..
719            } => Some((model, *idx, 9, &[])),
720            _ => None,
721        }
722    }
723
724    /// Dense f32 view — only for owned tensors. Masked/sparse execution
725    /// paths require it; quantized weights don't support masks yet.
726    pub fn as_f32(&self) -> Option<&[f32]> {
727        match self {
728            Self::F32 { data, .. } => Some(data),
729            Self::Mapped { .. } => None,
730        }
731    }
732
733    fn quant_bytes(&self) -> &[u8] {
734        match self {
735            Self::Mapped { model, idx, .. } => model.entry_bytes(&model.tensors[*idx]),
736            Self::F32 { .. } => unreachable!("quant_bytes on F32"),
737        }
738    }
739
740    /// Dequantize one row into `dst` (embedding lookup).
741    pub fn row_f32(&self, r: usize, dst: &mut [f32]) {
742        let cols = self.cols();
743        debug_assert_eq!(dst.len(), cols);
744        match self {
745            Self::F32 { data, .. } => dst.copy_from_slice(&data[r * cols..(r + 1) * cols]),
746            Self::Mapped {
747                model,
748                idx,
749                dtype,
750                row_scale,
751                col_field,
752                vbit_offsets,
753                ..
754            } => {
755                if *dtype == TensorDtype::Q4Tiled {
756                    let bytes = self.quant_bytes();
757                    let gpr = cols / GROUP_SIZE;
758                    for gi in 0..gpr {
759                        let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
760                        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
761                        for (k, &b) in tile[2..].iter().enumerate() {
762                            dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
763                            dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
764                        }
765                    }
766                    if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
767                        crate::prism::inverse_embedding(model, dst);
768                    }
769                    return;
770                }
771                if *dtype == TensorDtype::Q4TiledP {
772                    let bytes = self.quant_bytes();
773                    let gpr = cols / GROUP_SIZE;
774                    let v = Q4tpView::new(bytes, self.rows(), cols);
775                    let mut sc = vec![0f32; gpr];
776                    v.scales_into(r, gpr, &mut sc);
777                    for gi in 0..gpr {
778                        let tile = &v.nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
779                        let s = sc[gi];
780                        for (k, &b) in tile.iter().enumerate() {
781                            dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
782                            dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
783                        }
784                    }
785                    if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
786                        crate::prism::inverse_embedding(model, dst);
787                    }
788                    return;
789                }
790                if *dtype == TensorDtype::Q2TiledP {
791                    let bytes = self.quant_bytes();
792                    let gpr = cols / GROUP_SIZE;
793                    let v = Q4tpView::new_q2(bytes, self.rows(), cols);
794                    let mut sc = vec![0f32; gpr];
795                    v.scales_into(r, gpr, &mut sc);
796                    for gi in 0..gpr {
797                        let ch =
798                            &v.nib[(r * gpr + gi) * Q2TP_CHUNK..(r * gpr + gi + 1) * Q2TP_CHUNK];
799                        let s = sc[gi];
800                        for (k, &b) in ch.iter().enumerate() {
801                            for j in 0..4 {
802                                let center = if crate::prism::is_affine_target(
803                                    model,
804                                    &model.tensors[*idx].name,
805                                ) {
806                                    1.0
807                                } else {
808                                    1.5
809                                };
810                                dst[gi * GROUP_SIZE + k * 4 + j] =
811                                    (((b >> (2 * j)) & 3) as f32 - center) * s;
812                            }
813                        }
814                    }
815                    if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
816                        crate::prism::inverse_embedding(model, dst);
817                    }
818                    return;
819                }
820                if *dtype == TensorDtype::Q4Block {
821                    let (packed, scales) = q4_split(self.quant_bytes(), self.rows(), cols);
822                    let gpr = cols / GROUP_SIZE;
823                    for gi in 0..gpr {
824                        let g = r * gpr + gi;
825                        let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
826                        for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
827                            dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
828                            dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
829                        }
830                    }
831                    if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
832                        crate::prism::inverse_embedding(model, dst);
833                    }
834                    return;
835                }
836                if *dtype == TensorDtype::Q1 {
837                    let bytes = self.quant_bytes();
838                    let gpr = cols / GROUP_SIZE;
839                    for gi in 0..gpr {
840                        let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
841                        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
842                        for (j, &b) in tile[2..].iter().enumerate() {
843                            for k in 0..8 {
844                                dst[gi * GROUP_SIZE + j * 8 + k] =
845                                    (((b >> k) & 1) as f32 * 2.0 - 1.0) * s;
846                            }
847                        }
848                    }
849                    if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
850                        crate::prism::inverse_embedding(model, dst);
851                    }
852                    return;
853                }
854                if *dtype == TensorDtype::Q1T {
855                    let bytes = self.quant_bytes();
856                    let gpr = cols / GROUP_SIZE;
857                    let base_len = self.rows() * gpr * cortiq_core::quant::Q1T_TILE;
858                    for gi in 0..gpr {
859                        let off = (r * gpr + gi) * cortiq_core::quant::Q1T_TILE;
860                        let s = cortiq_core::quant::f16_to_f32(u16::from_le_bytes([
861                            bytes[off],
862                            bytes[off + 1],
863                        ]));
864                        let codes = &bytes[off + 2..off + cortiq_core::quant::Q1T_TILE];
865                        for k in 0..GROUP_SIZE {
866                            dst[gi * GROUP_SIZE + k] = match cortiq_core::quant::q1t_code(codes, k)
867                            {
868                                1 => s,
869                                2 => -s,
870                                _ => 0.0,
871                            };
872                        }
873                    }
874                    // Overlay
875                    let rows = self.rows();
876                    let entries = base_len + (rows + 1) * 4;
877                    if entries <= bytes.len() {
878                        let ptrs = &bytes[base_len..base_len + (rows + 1) * 4];
879                        let r0 = u32::from_le_bytes([
880                            ptrs[r * 4],
881                            ptrs[r * 4 + 1],
882                            ptrs[r * 4 + 2],
883                            ptrs[r * 4 + 3],
884                        ]) as usize;
885                        let r1 = u32::from_le_bytes([
886                            ptrs[(r + 1) * 4],
887                            ptrs[(r + 1) * 4 + 1],
888                            ptrs[(r + 1) * 4 + 2],
889                            ptrs[(r + 1) * 4 + 3],
890                        ]) as usize;
891                        let off = entries + r0 * 4;
892                        for i in 0..r1 - r0 {
893                            let item = &bytes[off + i * 4..off + i * 4 + 4];
894                            let c = u16::from_le_bytes([item[0], item[1]]) as usize;
895                            let v = cortiq_core::quant::f16_to_f32(u16::from_le_bytes([
896                                item[2], item[3],
897                            ]));
898                            if c < cols {
899                                dst[c] = v;
900                            }
901                        }
902                    }
903                    if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
904                        crate::prism::inverse_embedding(model, dst);
905                    }
906                    return;
907                }
908                if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
909                    let bytes = self.quant_bytes();
910                    let rows = self.rows();
911                    let ng = cols / GROUP_SIZE;
912                    let bits = &bytes[..rows];
913                    let sc_off = rows;
914                    // Precomputed at load — embedding lookup used to scan
915                    // the bit-widths of every preceding row (O(token_id)).
916                    let off = vbit_offsets[r];
917                    let b = bits[r] as usize;
918                    let l = ((1usize << (b - 1)) - 1) as f32;
919                    let data = &bytes[off..];
920                    let (mut acc, mut nbits, mut byte_idx) = (0u64, 0usize, 0usize);
921                    for (i, d) in dst.iter_mut().enumerate() {
922                        while nbits < b {
923                            acc = (acc << 8) | data[byte_idx] as u64;
924                            byte_idx += 1;
925                            nbits += 8;
926                        }
927                        let u = ((acc >> (nbits - b)) & ((1u64 << b) - 1)) as f32;
928                        nbits -= b;
929                        let so = (r * ng + i / GROUP_SIZE) * 2;
930                        let sv = f16_to_f32(u16::from_le_bytes([
931                            bytes[sc_off + so],
932                            bytes[sc_off + so + 1],
933                        ]));
934                        *d = (u - l) * sv;
935                    }
936                    if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
937                        crate::prism::inverse_embedding(model, dst);
938                    }
939                    return;
940                }
941                let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
942                let s = row_scale[r];
943                match dtype {
944                    TensorDtype::Q8Row => {
945                        for (d, &b) in dst.iter_mut().zip(q) {
946                            *d = (b as i8) as f32 * s;
947                        }
948                    }
949                    TensorDtype::Q8_2f => {
950                        for (i, (d, &b)) in dst.iter_mut().zip(q).enumerate() {
951                            *d = (b as i8) as f32 * s * col_field[i];
952                        }
953                    }
954                    _ => unreachable!(),
955                }
956                if crate::prism::is_inverse_embedding(model, &model.tensors[*idx].name) {
957                    crate::prism::inverse_embedding(model, dst);
958                }
959            }
960        }
961    }
962
963    /// Can this tensor's columns be read cheaply (for sparse down_proj)?
964    /// True for F32/Q8Row/Q8_2f (per-row scale, direct strided access);
965    /// false for group-packed q4/vbit (column access would unpack whole
966    /// groups — sparse execution falls back to f32 for those).
967    pub fn sparse_col_ok(&self) -> bool {
968        match self {
969            Self::F32 { .. } => true,
970            Self::Mapped { dtype, .. } => {
971                matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
972            }
973        }
974    }
975
976    /// down_proj [hidden, inter]: accumulate `w · col(c)` into `out`
977    /// [hidden] — reads ONLY column `c` (one neuron) from the mmap,
978    /// no full-matrix dequant. `out[k] += w · down[k, c]`.
979    pub fn add_col_scaled(&self, c: usize, w: f32, out: &mut [f32]) {
980        let inter = self.cols();
981        let hidden = self.rows();
982        debug_assert_eq!(out.len(), hidden);
983        match self {
984            Self::F32 { data, .. } => {
985                for (k, o) in out.iter_mut().enumerate() {
986                    *o += w * data[k * inter + c];
987                }
988            }
989            Self::Mapped {
990                dtype,
991                row_scale,
992                col_field,
993                ..
994            } => {
995                let q = self.quant_bytes();
996                let colf = if *dtype == TensorDtype::Q8_2f {
997                    col_field[c]
998                } else {
999                    1.0
1000                };
1001                let wc = w * colf;
1002                for (k, o) in out.iter_mut().enumerate() {
1003                    let b = q[k * inter + c] as i8 as f32;
1004                    *o += wc * b * row_scale[k];
1005                }
1006            }
1007        }
1008    }
1009
1010    /// Touch the head of row `r` so the DRAM latency of the next
1011    /// neuron's weights overlaps the current one's arithmetic.
1012    ///
1013    /// Scattered rows are what per-token sparsity reads, and a 2 KB
1014    /// stride is past what the hardware prefetcher follows: without this
1015    /// every row starts with a cold miss that nothing hides. One touch
1016    /// per 512 bytes is enough — the rest of the row is a sequential run
1017    /// the prefetcher does pick up.
1018    #[inline]
1019    pub fn prefetch_row(&self, r: usize) {
1020        let Self::Mapped { dtype, .. } = self else {
1021            return;
1022        };
1023        if !matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f) {
1024            return;
1025        }
1026        let cols = self.cols();
1027        let q = self.quant_bytes();
1028        let (a, b) = (r * cols, (r + 1) * cols);
1029        if b > q.len() {
1030            return;
1031        }
1032        let mut j = a;
1033        while j < b {
1034            unsafe { std::ptr::read_volatile(q.as_ptr().add(j)) };
1035            j += 512;
1036        }
1037    }
1038
1039    /// `out += w · row(r)` — the transposed twin of `add_col_scaled`.
1040    ///
1041    /// A neuron's `down` weights are a COLUMN of `[hidden, inter]`, and a
1042    /// column is strided: reading one costs a cache line per element, so
1043    /// per-neuron dynamic sparsity saves arithmetic and no bytes. Stored
1044    /// transposed (`down_proj.t.weight`, `[inter, hidden]`) the same
1045    /// weights are a contiguous ROW, and this accumulate reads exactly
1046    /// the neurons the token asked for.
1047    pub fn add_row_scaled(&self, r: usize, w: f32, out: &mut [f32], scratch: &mut [f32]) {
1048        let cols = self.cols();
1049        debug_assert_eq!(out.len(), cols);
1050        match self {
1051            Self::F32 { data, .. } => {
1052                let row = &data[r * cols..(r + 1) * cols];
1053                for (o, v) in out.iter_mut().zip(row) {
1054                    *o += w * v;
1055                }
1056            }
1057            Self::Mapped {
1058                dtype,
1059                row_scale,
1060                col_field,
1061                ..
1062            } => match dtype {
1063                TensorDtype::Q8Row => {
1064                    let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
1065                    let ws = w * row_scale[r];
1066                    let row: &[i8] =
1067                        unsafe { std::slice::from_raw_parts(q.as_ptr() as *const i8, q.len()) };
1068                    axpy_i8_f32(out, row, ws);
1069                }
1070                TensorDtype::Q8_2f => {
1071                    let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
1072                    let ws = w * row_scale[r];
1073                    for ((o, b), c) in out.iter_mut().zip(q).zip(col_field) {
1074                        *o += ws * c * (*b as i8 as f32);
1075                    }
1076                }
1077                _ => {
1078                    self.row_f32(r, scratch);
1079                    for (o, v) in out.iter_mut().zip(scratch.iter()) {
1080                        *o += w * v;
1081                    }
1082                }
1083            },
1084        }
1085    }
1086
1087    /// Dot of row `r` with `x` (gate/up active-neuron path). Reads only
1088    /// row `r` from the mmap — no full dequant. q4/vbit dequant the row
1089    /// into `scratch` first (rare for active-FFN weights).
1090    pub fn row_dot(&self, r: usize, x: &[f32], scratch: &mut [f32]) -> f32 {
1091        let cols = self.cols();
1092        match self {
1093            Self::F32 { data, .. } => {
1094                let row = &data[r * cols..(r + 1) * cols];
1095                row.iter().zip(x).map(|(w, v)| w * v).sum()
1096            }
1097            Self::Mapped {
1098                model,
1099                idx,
1100                dtype,
1101                row_scale,
1102                col_field,
1103                ..
1104            } => {
1105                let prism_forward =
1106                    crate::prism::is_forward_weight(model, &model.tensors[*idx].name);
1107                if prism_forward {
1108                    let transformed = crate::prism::forward(model, &x[..cols]);
1109                    let gpr = cols / GROUP_SIZE;
1110                    match dtype {
1111                        TensorDtype::Q2TiledP => {
1112                            let v = Q4tpView::new_q2(self.quant_bytes(), self.rows(), cols);
1113                            let mut sc = vec![0f32; gpr];
1114                            v.scales_into(r, gpr, &mut sc);
1115                            if crate::prism::is_affine_target(model, &model.tensors[*idx].name) {
1116                                return q2tp_affine_row_exact(v.nib, r, gpr, &transformed, &sc);
1117                            }
1118                            return q2tp_row_exact(v.nib, r, gpr, &transformed, &sc);
1119                        }
1120                        _ => {
1121                            self.row_f32(r, scratch);
1122                            return scratch.iter().zip(&transformed).map(|(w, v)| w * v).sum();
1123                        }
1124                    }
1125                }
1126                match dtype {
1127                    TensorDtype::Q8Row => {
1128                        let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
1129                        dot_i8_f32(q, x) * row_scale[r]
1130                    }
1131                    TensorDtype::Q8_2f => {
1132                        let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
1133                        dot_i8_col_f32(q, x, col_field) * row_scale[r]
1134                    }
1135                    _ => {
1136                        self.row_f32(r, scratch);
1137                        scratch.iter().zip(x).map(|(w, v)| w * v).sum()
1138                    }
1139                }
1140            }
1141        }
1142    }
1143
1144    /// `out = W · x` (row-major). F32 delegates to the historical
1145    /// bit-exact path; Mapped runs the fused int8 kernel.
1146    pub fn matvec(&self, x: &[f32], out: &mut [f32], pool: Option<&Pool>) {
1147        match self {
1148            // NOTE: `out.len()` DRIVES this arm — it computes that many rows,
1149            // and `x.len()` is the stride. A short `out` is legitimate here,
1150            // which is why the check below lives in the Mapped arm only.
1151            Self::F32 { data, .. } => matvec_rows(pool, data, x, out),
1152            Self::Mapped {
1153                model,
1154                idx,
1155                dtype,
1156                rows,
1157                cols,
1158                row_scale,
1159                col_field,
1160                vbit_offsets,
1161                repack,
1162            } => {
1163                let _ = (model, idx);
1164                // Every kernel below writes `rows` entries through a raw
1165                // pointer, so a short `out` is an out-of-bounds WRITE, not a
1166                // wrong answer: it scribbles on the allocator's metadata and
1167                // the process aborts much later, somewhere innocent
1168                // (`double free or corruption`, `corrupted double-linked
1169                // list`). The debug_assert two of the kernels carried is
1170                // compiled out of the release — exactly the build where it
1171                // matters. Fail here instead, while the caller is still on
1172                // the stack to be named.
1173                assert!(
1174                    out.len() >= *rows && x.len() >= *cols,
1175                    "matvec {rows}x{cols}: out {} (need {rows}), x {} (need {cols})",
1176                    out.len(),
1177                    x.len(),
1178                );
1179                let prism_forward =
1180                    crate::prism::is_forward_weight(model, &model.tensors[*idx].name);
1181                if *dtype == TensorDtype::Q2TiledP
1182                    && std::env::var("CMF_Q2TP_TRACE").as_deref() == Ok("1")
1183                {
1184                    use std::sync::atomic::{AtomicUsize, Ordering};
1185                    static N: AtomicUsize = AtomicUsize::new(0);
1186                    let n = N.fetch_add(1, Ordering::Relaxed);
1187                    if n < 128 {
1188                        eprintln!(
1189                            "q2tp-dispatch #{n} name={} prism={} rows={} cols={} gpu={} optin={} layer={}",
1190                            model.tensors[*idx].name,
1191                            prism_forward,
1192                            rows,
1193                            cols,
1194                            crate::gpu::enabled_here(),
1195                            crate::gpu::q2tp_gpu_opt_in(),
1196                            crate::gpu::cur_layer(),
1197                        );
1198                    }
1199                }
1200                // Prism stores every manifest-listed forward matrix in the
1201                // signed-Hadamard basis.  The q2tp WGSL path receives that
1202                // transformed vector and an explicit affine bit; codecs
1203                // without a descriptor-aware kernel remain on CPU below.
1204                if prism_forward {
1205                    let transformed = crate::prism::forward(model, &x[..*cols]);
1206                    match dtype {
1207                        TensorDtype::Q4Block => {
1208                            q4matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1209                        }
1210                        TensorDtype::Q4Tiled => {
1211                            q4t_matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1212                        }
1213                        TensorDtype::Q4TiledP => {
1214                            q4tp_matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1215                        }
1216                        TensorDtype::Q2TiledP => {
1217                            let affine =
1218                                crate::prism::is_affine_target(model, &model.tensors[*idx].name);
1219                            if *rows * *cols >= 8_388_608
1220                                && crate::gpu::enabled_here()
1221                                && crate::gpu::q2tp_gpu_opt_in()
1222                            {
1223                                let gpu_ok = if affine {
1224                                    crate::gpu::q2tp_affine_matvec(
1225                                        model,
1226                                        *idx,
1227                                        &transformed,
1228                                        *rows,
1229                                        *cols,
1230                                        out,
1231                                    )
1232                                } else {
1233                                    crate::gpu::q2tp_matvec(
1234                                        model,
1235                                        *idx,
1236                                        &transformed,
1237                                        *rows,
1238                                        *cols,
1239                                        out,
1240                                    )
1241                                };
1242                                if gpu_ok {
1243                                    return;
1244                                }
1245                            }
1246                            if affine {
1247                                q2tp_affine_matvec(
1248                                    self.quant_bytes(),
1249                                    &transformed,
1250                                    *rows,
1251                                    *cols,
1252                                    out,
1253                                    pool,
1254                                )
1255                            } else {
1256                                q2tp_matvec(
1257                                    self.quant_bytes(),
1258                                    &transformed,
1259                                    *rows,
1260                                    *cols,
1261                                    out,
1262                                    pool,
1263                                )
1264                            }
1265                        }
1266                        TensorDtype::Q1 => {
1267                            q1_matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1268                        }
1269                        TensorDtype::Q1T => {
1270                            q1t_matvec(self.quant_bytes(), &transformed, *rows, *cols, out, pool)
1271                        }
1272                        TensorDtype::Vbit | TensorDtype::VbitRo => vbitmatvec(
1273                            self.quant_bytes(),
1274                            vbit_offsets,
1275                            &transformed,
1276                            *rows,
1277                            *cols,
1278                            out,
1279                            pool,
1280                        ),
1281                        TensorDtype::Q8Row | TensorDtype::Q8_2f => qmatvec(
1282                            self.quant_bytes(),
1283                            repack,
1284                            row_scale,
1285                            &transformed,
1286                            col_field,
1287                            *dtype,
1288                            *rows,
1289                            *cols,
1290                            out,
1291                            pool,
1292                        ),
1293                        _ => unreachable!("unsupported mapped Prism dtype {dtype:?}"),
1294                    }
1295                    return;
1296                }
1297                if *dtype == TensorDtype::Q4Block {
1298                    // GPU route (wgpu q4b kernel) for large q4_block matvecs —
1299                    // gives NVIDIA/AMD/Intel q4 models a GPU path. Probe keeps
1300                    // the winner; Metal returns false → the CPU kernel below.
1301                    if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1302                        let t0 = std::time::Instant::now();
1303                        match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
1304                            crate::gpu::ProbeArm::Gpu => {
1305                                if crate::gpu::q4b_matvec(model, *idx, x, *rows, *cols, out) {
1306                                    crate::gpu::probe_record(
1307                                        crate::gpu::OpClass::Matvec,
1308                                        true,
1309                                        t0.elapsed(),
1310                                    );
1311                                    return;
1312                                }
1313                            }
1314                            crate::gpu::ProbeArm::CpuTimed => {
1315                                q4matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1316                                crate::gpu::probe_record(
1317                                    crate::gpu::OpClass::Matvec,
1318                                    false,
1319                                    t0.elapsed(),
1320                                );
1321                                return;
1322                            }
1323                            crate::gpu::ProbeArm::Cpu => {}
1324                        }
1325                    }
1326                    q4matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1327                    return;
1328                }
1329                if *dtype == TensorDtype::Q4Tiled {
1330                    // GPU route for large q4t matvecs — the lm_head class,
1331                    // same shape as the q4tp arm below. The probe keeps the
1332                    // winner; a backend without the kernel refuses and the
1333                    // CPU path stays.
1334                    if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1335                        let t0 = std::time::Instant::now();
1336                        let cls = crate::gpu::matvec_class(*rows, *cols);
1337                        match crate::gpu::probe_arm(cls) {
1338                            crate::gpu::ProbeArm::Gpu => {
1339                                if crate::gpu::q4t_matvec(model, *idx, x, *rows, *cols, out) {
1340                                    crate::gpu::probe_record(cls, true, t0.elapsed());
1341                                    return;
1342                                }
1343                            }
1344                            crate::gpu::ProbeArm::CpuTimed => {
1345                                q4t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1346                                crate::gpu::probe_record(cls, false, t0.elapsed());
1347                                return;
1348                            }
1349                            crate::gpu::ProbeArm::Cpu => {}
1350                        }
1351                    }
1352                    q4t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1353                    return;
1354                }
1355                if *dtype == TensorDtype::Q4TiledP {
1356                    // GPU route for large q4tp matvecs — the lm_head class.
1357                    // On a q4tp checkpoint the head is the biggest single
1358                    // host matvec left in the decode step, and the batched
1359                    // kernel at b=1 already exists on both backends. Probe
1360                    // keeps the winner, same as q4_block above.
1361                    if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1362                        let t0 = std::time::Instant::now();
1363                        let cls = crate::gpu::matvec_class(*rows, *cols);
1364                        match crate::gpu::probe_arm(cls) {
1365                            crate::gpu::ProbeArm::Gpu => {
1366                                if crate::gpu::q4tp_matvec(model, *idx, x, *rows, *cols, out) {
1367                                    crate::gpu::probe_record(cls, true, t0.elapsed());
1368                                    return;
1369                                }
1370                            }
1371                            crate::gpu::ProbeArm::CpuTimed => {
1372                                q4tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1373                                crate::gpu::probe_record(cls, false, t0.elapsed());
1374                                return;
1375                            }
1376                            crate::gpu::ProbeArm::Cpu => {}
1377                        }
1378                    }
1379                    q4tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1380                    return;
1381                }
1382                if *dtype == TensorDtype::Q2TiledP {
1383                    q2tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1384                    return;
1385                }
1386                if *dtype == TensorDtype::Q1 {
1387                    // GPU route for large q1 matvecs (out_proj / lm_head
1388                    // class): the CPU q1 kernel is load-port-bound at
1389                    // ~4 GB/s/core, the GPU one is bandwidth-bound — the
1390                    // probe measures both arms and keeps the winner.
1391                    if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1392                        let t0 = std::time::Instant::now();
1393                        let arm = if crate::gpu::q1_force() {
1394                            crate::gpu::ProbeArm::Gpu
1395                        } else {
1396                            crate::gpu::probe_arm(crate::gpu::OpClass::Matvec)
1397                        };
1398                        match arm {
1399                            crate::gpu::ProbeArm::Gpu => {
1400                                if crate::gpu::q1_matvec(model, *idx, x, *rows, *cols, out) {
1401                                    crate::gpu::probe_record(
1402                                        crate::gpu::OpClass::Matvec,
1403                                        true,
1404                                        t0.elapsed(),
1405                                    );
1406                                    return;
1407                                }
1408                            }
1409                            crate::gpu::ProbeArm::CpuTimed => {
1410                                q1_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1411                                crate::gpu::probe_record(
1412                                    crate::gpu::OpClass::Matvec,
1413                                    false,
1414                                    t0.elapsed(),
1415                                );
1416                                return;
1417                            }
1418                            crate::gpu::ProbeArm::Cpu => {}
1419                        }
1420                    }
1421                    q1_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1422                    return;
1423                }
1424                if *dtype == TensorDtype::Q1T {
1425                    // GPU route for large q1t matvecs: the ternary BASE dot runs
1426                    // on the GPU (load-port-bound on CPU, like q1), then the
1427                    // sparse overlay is added on the CPU. Probe keeps the winner.
1428                    if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1429                        let t0 = std::time::Instant::now();
1430                        match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
1431                            crate::gpu::ProbeArm::Gpu => {
1432                                if crate::gpu::q1t_matvec(model, *idx, x, *rows, *cols, out) {
1433                                    q1t_add_overlay(self.quant_bytes(), x, *rows, *cols, out, pool);
1434                                    crate::gpu::probe_record(
1435                                        crate::gpu::OpClass::Matvec,
1436                                        true,
1437                                        t0.elapsed(),
1438                                    );
1439                                    return;
1440                                }
1441                            }
1442                            crate::gpu::ProbeArm::CpuTimed => {
1443                                q1t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1444                                crate::gpu::probe_record(
1445                                    crate::gpu::OpClass::Matvec,
1446                                    false,
1447                                    t0.elapsed(),
1448                                );
1449                                return;
1450                            }
1451                            crate::gpu::ProbeArm::Cpu => {}
1452                        }
1453                    }
1454                    q1t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1455                    return;
1456                }
1457                if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1458                    vbitmatvec(self.quant_bytes(), vbit_offsets, x, *rows, *cols, out, pool);
1459                    return;
1460                }
1461                let xs = prescale(x, col_field, *dtype);
1462                // D5: large q8 matrices (lm_head-class) — hybrid
1463                // CPU∥GPU: split the rows, both sides compute
1464                // SIMULTANEOUSLY (same math, shared prescale).
1465                // GPU share: CMF_GPU_SPLIT (0..1, default 0.5).
1466                if *rows >= crate::gpu::min_rows()
1467                    && matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
1468                    && gpu_lmhead_enabled()
1469                    && crate::gpu::enabled_here()
1470                {
1471                    // Runtime probe: alternate the hybrid against the
1472                    // pure-CPU matvec, keep whichever is faster HERE.
1473                    let t0 = std::time::Instant::now();
1474                    match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
1475                        crate::gpu::ProbeArm::Gpu => {}
1476                        crate::gpu::ProbeArm::CpuTimed => {
1477                            qmatvec(
1478                                self.quant_bytes(),
1479                                repack,
1480                                row_scale,
1481                                x,
1482                                col_field,
1483                                *dtype,
1484                                *rows,
1485                                *cols,
1486                                out,
1487                                pool,
1488                            );
1489                            crate::gpu::probe_record(
1490                                crate::gpu::OpClass::Matvec,
1491                                false,
1492                                t0.elapsed(),
1493                            );
1494                            return;
1495                        }
1496                        crate::gpu::ProbeArm::Cpu => {
1497                            qmatvec(
1498                                self.quant_bytes(),
1499                                repack,
1500                                row_scale,
1501                                x,
1502                                col_field,
1503                                *dtype,
1504                                *rows,
1505                                *cols,
1506                                out,
1507                                pool,
1508                            );
1509                            return;
1510                        }
1511                    }
1512                    let frac = gpu_split_frac();
1513                    let cpu_rows = ((*rows as f32) * (1.0 - frac)) as usize;
1514                    let (out_cpu, out_gpu) = out.split_at_mut(cpu_rows);
1515                    let bytes = self.quant_bytes();
1516                    let ok = std::thread::scope(|sc| {
1517                        let g = sc.spawn(|| {
1518                            crate::gpu::q8_matvec_range(
1519                                model,
1520                                *idx,
1521                                cpu_rows,
1522                                &row_scale[cpu_rows..],
1523                                &xs,
1524                                *rows - cpu_rows,
1525                                *cols,
1526                                out_gpu,
1527                            )
1528                        });
1529                        if cpu_rows > 0 {
1530                            // Repack prefix covers the full groups of the
1531                            // CPU half (the split starts at row 0).
1532                            let rep_cpu = if repack.is_empty() {
1533                                &[][..]
1534                            } else {
1535                                &repack[..(cpu_rows / 4) * 4 * *cols]
1536                            };
1537                            qmatvec(
1538                                &bytes[..cpu_rows * *cols],
1539                                rep_cpu,
1540                                &row_scale[..cpu_rows],
1541                                x,
1542                                col_field,
1543                                *dtype,
1544                                cpu_rows,
1545                                *cols,
1546                                out_cpu,
1547                                pool,
1548                            );
1549                        }
1550                        g.join().unwrap_or(false)
1551                    });
1552                    if ok {
1553                        crate::gpu::probe_record(crate::gpu::OpClass::Matvec, true, t0.elapsed());
1554                        return;
1555                    }
1556                    // GPU failed — CPU finishes its half (rows rebased —
1557                    // group offsets don't line up, mmap layout only).
1558                    qmatvec(
1559                        &bytes[cpu_rows * *cols..(*rows) * *cols],
1560                        &[],
1561                        &row_scale[cpu_rows..],
1562                        x,
1563                        col_field,
1564                        *dtype,
1565                        *rows - cpu_rows,
1566                        *cols,
1567                        out_gpu,
1568                        pool,
1569                    );
1570                    return;
1571                }
1572                qmatvec(
1573                    self.quant_bytes(),
1574                    repack,
1575                    row_scale,
1576                    x,
1577                    col_field,
1578                    *dtype,
1579                    *rows,
1580                    *cols,
1581                    out,
1582                    pool,
1583                );
1584            }
1585        }
1586    }
1587
1588    /// Fused two-input matvec (MTP verify pair): weights streamed once.
1589    pub fn matvec2(
1590        &self,
1591        x1: &[f32],
1592        x2: &[f32],
1593        o1: &mut [f32],
1594        o2: &mut [f32],
1595        pool: Option<&Pool>,
1596    ) {
1597        match self {
1598            Self::F32 { data, .. } => matvec_rows2(pool, data, x1, x2, o1, o2),
1599            Self::Mapped {
1600                model,
1601                idx,
1602                dtype,
1603                rows,
1604                cols,
1605                row_scale,
1606                col_field,
1607                vbit_offsets,
1608                ..
1609            } => {
1610                if crate::prism::is_forward_weight(model, &model.tensors[*idx].name) {
1611                    let tx1 = crate::prism::forward(model, &x1[..*cols]);
1612                    let tx2 = crate::prism::forward(model, &x2[..*cols]);
1613                    match dtype {
1614                        TensorDtype::Q4Block => {
1615                            q4matvec2(self.quant_bytes(), &tx1, &tx2, *rows, *cols, o1, o2, pool)
1616                        }
1617                        TensorDtype::Q4Tiled => {
1618                            q4t_matvec2(self.quant_bytes(), &tx1, &tx2, *rows, *cols, o1, o2, pool)
1619                        }
1620                        TensorDtype::Q4TiledP => {
1621                            q4tp_matvec2(self.quant_bytes(), &tx1, &tx2, *rows, *cols, o1, o2, pool)
1622                        }
1623                        TensorDtype::Q2TiledP => {
1624                            if crate::prism::is_affine_target(model, &model.tensors[*idx].name) {
1625                                q2tp_affine_matvec2(
1626                                    self.quant_bytes(),
1627                                    &tx1,
1628                                    &tx2,
1629                                    *rows,
1630                                    *cols,
1631                                    o1,
1632                                    o2,
1633                                    pool,
1634                                )
1635                            } else {
1636                                q2tp_matvec2(
1637                                    self.quant_bytes(),
1638                                    &tx1,
1639                                    &tx2,
1640                                    *rows,
1641                                    *cols,
1642                                    o1,
1643                                    o2,
1644                                    pool,
1645                                )
1646                            }
1647                        }
1648                        TensorDtype::Q1 => {
1649                            q1_matvec2(self.quant_bytes(), &tx1, &tx2, *rows, *cols, o1, o2, pool)
1650                        }
1651                        TensorDtype::Q1T => {
1652                            q1t_matvec2(self.quant_bytes(), &tx1, &tx2, *rows, *cols, o1, o2, pool)
1653                        }
1654                        TensorDtype::Vbit | TensorDtype::VbitRo => vbitmatvec2(
1655                            self.quant_bytes(),
1656                            vbit_offsets,
1657                            &tx1,
1658                            &tx2,
1659                            *rows,
1660                            *cols,
1661                            o1,
1662                            o2,
1663                            pool,
1664                        ),
1665                        TensorDtype::Q8Row | TensorDtype::Q8_2f => qmatvec2(
1666                            self.quant_bytes(),
1667                            row_scale,
1668                            &tx1,
1669                            &tx2,
1670                            col_field,
1671                            *dtype,
1672                            *rows,
1673                            *cols,
1674                            o1,
1675                            o2,
1676                            pool,
1677                        ),
1678                        _ => unreachable!("unsupported mapped Prism dtype {dtype:?}"),
1679                    }
1680                    return;
1681                }
1682                if *dtype == TensorDtype::Q4Block {
1683                    q4matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1684                    return;
1685                }
1686                if *dtype == TensorDtype::Q4Tiled {
1687                    q4t_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1688                    return;
1689                }
1690                if *dtype == TensorDtype::Q4TiledP {
1691                    q4tp_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1692                    return;
1693                }
1694                if *dtype == TensorDtype::Q2TiledP {
1695                    q2tp_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1696                    return;
1697                }
1698                if *dtype == TensorDtype::Q1 {
1699                    q1_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1700                    return;
1701                }
1702                if *dtype == TensorDtype::Q1T {
1703                    // Fused ternary pair: one row pass, the register
1704                    // unpack shared across both streams on ARM. (Q1T
1705                    // lacks a row_scale array — scales live inline in
1706                    // the tiles — so it must not fall through to the
1707                    // q8 qmatvec2 below.)
1708                    q1t_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1709                    return;
1710                }
1711                if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1712                    vbitmatvec2(
1713                        self.quant_bytes(),
1714                        vbit_offsets,
1715                        x1,
1716                        x2,
1717                        *rows,
1718                        *cols,
1719                        o1,
1720                        o2,
1721                        pool,
1722                    );
1723                    return;
1724                }
1725                qmatvec2(
1726                    self.quant_bytes(),
1727                    row_scale,
1728                    x1,
1729                    x2,
1730                    col_field,
1731                    *dtype,
1732                    *rows,
1733                    *cols,
1734                    o1,
1735                    o2,
1736                    pool,
1737                );
1738            }
1739        }
1740    }
1741}
1742
1743impl QTensor {
1744    /// Batched matvec (prefill-GEMM): xs — row-major [b, cols],
1745    /// out — row-major [b, rows]. Element-wise semantics are IDENTICAL
1746    /// to b matvec calls (same dot kernels in the same order); the win —
1747    /// the weight row streams from DRAM once per batch, not b times.
1748    /// `(model, index)` when this is a memory-mapped q4tp tensor — the
1749    /// identity a device-resident chain needs to hand `tp_matmat` the
1750    /// weight without going through this struct's own dispatch.
1751    pub fn q4tp_mapped(&self) -> Option<(&std::sync::Arc<CmfModel>, usize)> {
1752        if self.has_prism_contract() {
1753            return None;
1754        }
1755        match self {
1756            Self::Mapped {
1757                model, idx, dtype, ..
1758            } if *dtype == TensorDtype::Q4TiledP => Some((model, *idx)),
1759            _ => None,
1760        }
1761    }
1762
1763    pub fn matmat(&self, xs_all: &[f32], b: usize, out: &mut [f32], pool: Option<&Pool>) {
1764        let cols = self.cols();
1765        let rows = self.rows();
1766        debug_assert_eq!(xs_all.len(), b * cols);
1767        debug_assert_eq!(out.len(), b * rows);
1768        // GPTQ calibration: fold this layer's inputs into its Hessian. Only
1769        // Mapped tensors carry a directory name; the check is a relaxed
1770        // atomic load, free when not calibrating.
1771        if crate::gptq_capture::capturing() {
1772            if let Self::Mapped { model, idx, .. } = self {
1773                crate::gptq_capture::accumulate(&model.tensors[*idx].name, xs_all, b, cols);
1774            }
1775        }
1776        match self {
1777            Self::F32 { data, .. } => {
1778                let out_addr = SendMut(out.as_mut_ptr());
1779                let run = |start: usize, end: usize| {
1780                    for o in start..end {
1781                        let row = &data[o * cols..(o + 1) * cols];
1782                        for bi in 0..b {
1783                            let x = &xs_all[bi * cols..(bi + 1) * cols];
1784                            let mut acc = 0f32;
1785                            for j in 0..cols {
1786                                acc += row[j] * x[j];
1787                            }
1788                            unsafe { *out_addr.at(bi * rows + o) = acc };
1789                        }
1790                    }
1791                };
1792                dispatch_rows(pool, rows, &run);
1793            }
1794            Self::Mapped {
1795                model,
1796                idx,
1797                dtype,
1798                row_scale,
1799                col_field,
1800                vbit_offsets,
1801                ..
1802            } => {
1803                if crate::prism::is_forward_weight(model, &model.tensors[*idx].name) {
1804                    let mut transformed = Vec::with_capacity(xs_all.len());
1805                    for bi in 0..b {
1806                        transformed.extend_from_slice(&crate::prism::forward(
1807                            model,
1808                            &xs_all[bi * cols..(bi + 1) * cols],
1809                        ));
1810                    }
1811                    match dtype {
1812                        TensorDtype::Q4Block => {
1813                            q4matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
1814                        }
1815                        TensorDtype::Q4Tiled => {
1816                            q4t_matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
1817                        }
1818                        TensorDtype::Q4TiledP => {
1819                            q4tp_matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
1820                        }
1821                        TensorDtype::Q2TiledP => {
1822                            let affine =
1823                                crate::prism::is_affine_target(model, &model.tensors[*idx].name);
1824                            // Affine Prism Q2TP has a descriptor-aware GPU
1825                            // kernel for short/tail batches too.  Unlike the
1826                            // ordinary Q2TP path, don't force b<32 back to a
1827                            // scalar CPU matmat: prefill chunks and the final
1828                            // tail both need to stay on the tested GPU arm.
1829                            let gpu_batch_ok = if affine {
1830                                b >= 2
1831                            } else {
1832                                b >= 32 && b * rows * cols >= 128_000_000
1833                            };
1834                            if gpu_batch_ok
1835                                && cols % 32 == 0
1836                                && crate::gpu::enabled_here()
1837                                && crate::gpu::q2tp_gpu_opt_in()
1838                            {
1839                                let gpu_ok = if affine {
1840                                    crate::gpu::q2tp_affine_matmat(
1841                                        model,
1842                                        *idx,
1843                                        &transformed,
1844                                        b,
1845                                        rows,
1846                                        cols,
1847                                        out,
1848                                    )
1849                                } else {
1850                                    crate::gpu::q2tp_matmat(
1851                                        model,
1852                                        *idx,
1853                                        &transformed,
1854                                        b,
1855                                        rows,
1856                                        cols,
1857                                        out,
1858                                    )
1859                                };
1860                                if gpu_ok {
1861                                    return;
1862                                }
1863                            }
1864                            // A one-token Prism decode is the other short
1865                            // case.  Use the descriptor-aware matvec kernel
1866                            // before falling back to the exact CPU path.
1867                            if affine
1868                                && b == 1
1869                                && cols % 32 == 0
1870                                && crate::gpu::enabled_here()
1871                                && crate::gpu::q2tp_gpu_opt_in()
1872                                && crate::gpu::q2tp_affine_matvec(
1873                                    model,
1874                                    *idx,
1875                                    &transformed[..cols],
1876                                    rows,
1877                                    cols,
1878                                    &mut out[..rows],
1879                                )
1880                            {
1881                                return;
1882                            }
1883                            if affine {
1884                                q2tp_affine_matmat(
1885                                    self.quant_bytes(),
1886                                    &transformed,
1887                                    b,
1888                                    rows,
1889                                    cols,
1890                                    out,
1891                                    pool,
1892                                )
1893                            } else {
1894                                q2tp_matmat(
1895                                    self.quant_bytes(),
1896                                    &transformed,
1897                                    b,
1898                                    rows,
1899                                    cols,
1900                                    out,
1901                                    pool,
1902                                )
1903                            }
1904                        }
1905                        TensorDtype::Q1 => {
1906                            q1_matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
1907                        }
1908                        TensorDtype::Q1T => {
1909                            q1t_matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
1910                        }
1911                        TensorDtype::Vbit | TensorDtype::VbitRo => vbitmatmat(
1912                            self.quant_bytes(),
1913                            vbit_offsets,
1914                            &transformed,
1915                            b,
1916                            rows,
1917                            cols,
1918                            out,
1919                            pool,
1920                        ),
1921                        TensorDtype::Q8Row | TensorDtype::Q8_2f => {
1922                            let pre: Vec<std::borrow::Cow<'_, [f32]>> = (0..b)
1923                                .map(|bi| {
1924                                    prescale(
1925                                        &transformed[bi * cols..(bi + 1) * cols],
1926                                        col_field,
1927                                        *dtype,
1928                                    )
1929                                })
1930                                .collect();
1931                            qmatmat(self.quant_bytes(), row_scale, &pre, rows, cols, out, pool)
1932                        }
1933                        _ => unreachable!("unsupported mapped Prism dtype {dtype:?}"),
1934                    }
1935                    return;
1936                }
1937                if *dtype == TensorDtype::Q4Block {
1938                    q4matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1939                    return;
1940                }
1941                if *dtype == TensorDtype::Q4TiledP {
1942                    // GPU batched q4tp GEMM (dequant + f32nt mul_mm on the
1943                    // device); the probe keeps whichever beats the CPU arm.
1944                    // Narrow (prompt-encode) and wide (DiT) batches probe
1945                    // as separate classes — the regimes have opposite
1946                    // winners and one shared verdict locked the wrong arm.
1947                    // Kill switch (gpu::mm_kill): one grossly slow GPU op
1948                    // (a fair-condition op is ≤~100 ms even at 1024px)
1949                    // means the device is contended by another process
1950                    // (e.g. a simulator) — verdicts are per-process, so
1951                    // without the bail the whole render crawls behind
1952                    // someone else's queue.
1953                    if b >= 32
1954                        && b * rows * cols >= 128_000_000
1955                        && cols % 32 == 0
1956                        && !crate::gpu::mm_killed()
1957                        && crate::gpu::enabled_here()
1958                    {
1959                        let class = if b >= 128 {
1960                            crate::gpu::OpClass::MatmatWide
1961                        } else {
1962                            crate::gpu::OpClass::Matmat
1963                        };
1964                        if let Self::Mapped { model, idx, .. } = self {
1965                            // In-process A/B (`CMF_MM_AB=1`). Three
1966                            // wall-clock A/Bs on a shared stand disagreed
1967                            // with each other by 25% on the same change,
1968                            // because the machine drifts between processes
1969                            // and interleaving whole renders does not fix
1970                            // that. Here both arms run back to back on the
1971                            // SAME data inside one call, so whatever the
1972                            // machine is doing, it does to both — and the
1973                            // disagreement between their outputs falls out
1974                            // for free. Doubles the work; a diagnostic,
1975                            // not a mode.
1976                            if crate::mm_ab::on() {
1977                                let mut g = vec![0f32; b * rows];
1978                                let t = std::time::Instant::now();
1979                                let took = crate::gpu::q4tp_matmat(
1980                                    model, *idx, xs_all, b, rows, cols, &mut g,
1981                                );
1982                                let dg = t.elapsed();
1983                                let t = std::time::Instant::now();
1984                                q4tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1985                                let dc = t.elapsed();
1986                                crate::mm_ab::record(b, rows, cols, took, dg, dc, &g, out);
1987                                return;
1988                            }
1989                            let t0 = std::time::Instant::now();
1990                            // A cold call takes the device arm: its sample
1991                            // is discarded either way, and the upload is
1992                            // what the next step needs.
1993                            let resident = crate::gpu::weight_is_resident(model, *idx);
1994                            match crate::gpu::probe_arm_cold_prefers_gpu(class, resident) {
1995                                crate::gpu::ProbeArm::Gpu => {
1996                                    if crate::gpu::q4tp_matmat(
1997                                        model, *idx, xs_all, b, rows, cols, out,
1998                                    ) {
1999                                        let el = t0.elapsed();
2000                                        // Work-proportional budget: ~8× the
2001                                        // fair-device estimate (+20 ms slack).
2002                                        // An absolute cap missed the worst
2003                                        // case — contended ops sit at
2004                                        // 100–240 ms each and still bury a
2005                                        // render whose fair op is 3–9 ms.
2006                                        // Cold ops (first PSO build, buffer
2007                                        // alloc) are exempt: a one-off
2008                                        // ~50 ms compile is not contention.
2009                                        let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
2010                                        let budget = std::time::Duration::from_secs_f64(
2011                                            flops / 1.5e12 * 8.0 + 0.020,
2012                                        );
2013                                        crate::gpu::mm_budget_check(
2014                                            "q4tp matmat",
2015                                            el,
2016                                            budget,
2017                                            crate::gpu::probe_was_cold() || !resident,
2018                                        );
2019                                        crate::gpu::probe_record(class, true, el);
2020                                        return;
2021                                    }
2022                                }
2023                                crate::gpu::ProbeArm::CpuTimed => {
2024                                    q4tp_matmat(
2025                                        self.quant_bytes(),
2026                                        xs_all,
2027                                        b,
2028                                        rows,
2029                                        cols,
2030                                        out,
2031                                        pool,
2032                                    );
2033                                    crate::gpu::probe_record(class, false, t0.elapsed());
2034                                    return;
2035                                }
2036                                crate::gpu::ProbeArm::Cpu => {}
2037                            }
2038                        }
2039                    }
2040                    q4tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2041                    return;
2042                }
2043                if *dtype == TensorDtype::Q2TiledP {
2044                    // Same device arm as q4tp, behind the same probe:
2045                    // the planes differ, the dispatch does not. Without
2046                    // this a q2tp file ran its widest projections on the
2047                    // host while the 4-bit one had the card, which is a
2048                    // codec paying for its size twice.
2049                    if b >= 32
2050                        && b * rows * cols >= 128_000_000
2051                        && cols % 32 == 0
2052                        && !crate::gpu::mm_killed()
2053                        && crate::gpu::enabled_here()
2054                    {
2055                        let class = if b >= 128 {
2056                            crate::gpu::OpClass::MatmatWide
2057                        } else {
2058                            crate::gpu::OpClass::Matmat
2059                        };
2060                        if let Self::Mapped { model, idx, .. } = self {
2061                            let t0 = std::time::Instant::now();
2062                            match crate::gpu::probe_arm(class) {
2063                                crate::gpu::ProbeArm::Gpu => {
2064                                    if crate::gpu::q2tp_matmat(
2065                                        model, *idx, xs_all, b, rows, cols, out,
2066                                    ) {
2067                                        crate::gpu::probe_record(class, true, t0.elapsed());
2068                                        return;
2069                                    }
2070                                }
2071                                crate::gpu::ProbeArm::CpuTimed => {
2072                                    q2tp_matmat(
2073                                        self.quant_bytes(),
2074                                        xs_all,
2075                                        b,
2076                                        rows,
2077                                        cols,
2078                                        out,
2079                                        pool,
2080                                    );
2081                                    crate::gpu::probe_record(class, false, t0.elapsed());
2082                                    return;
2083                                }
2084                                crate::gpu::ProbeArm::Cpu => {}
2085                            }
2086                        }
2087                    }
2088                    // Without a host arm a q2tp tensor falls through to
2089                    // the q8 fallback, which reads it at one BYTE per
2090                    // weight — a 2x overrun that killed pool workers
2091                    // mid-prefill while the dispatcher waited forever.
2092                    q2tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2093                    return;
2094                }
2095                if *dtype == TensorDtype::Q4Tiled {
2096                    // GPU batched q4t GEMM (dequant + f32nt mul_mm on the
2097                    // device); the probe keeps whichever beats the CPU arm.
2098                    // Narrow (prompt-encode) and wide (DiT) batches probe
2099                    // as separate classes — the regimes have opposite
2100                    // winners and one shared verdict locked the wrong arm.
2101                    // Kill switch (gpu::mm_kill): one grossly slow GPU op
2102                    // (a fair-condition op is ≤~100 ms even at 1024px)
2103                    // means the device is contended by another process
2104                    // (e.g. a simulator) — verdicts are per-process, so
2105                    // without the bail the whole render crawls behind
2106                    // someone else's queue.
2107                    if b >= 32
2108                        && b * rows * cols >= 128_000_000
2109                        && cols % 32 == 0
2110                        && !crate::gpu::mm_killed()
2111                        && crate::gpu::enabled_here()
2112                    {
2113                        let class = if b >= 128 {
2114                            crate::gpu::OpClass::MatmatWide
2115                        } else {
2116                            crate::gpu::OpClass::Matmat
2117                        };
2118                        if let Self::Mapped { model, idx, .. } = self {
2119                            let t0 = std::time::Instant::now();
2120                            match crate::gpu::probe_arm(class) {
2121                                crate::gpu::ProbeArm::Gpu => {
2122                                    if crate::gpu::q4t_matmat(
2123                                        model, *idx, xs_all, b, rows, cols, out,
2124                                    ) {
2125                                        let el = t0.elapsed();
2126                                        // Work-proportional budget: ~8× the
2127                                        // fair-device estimate (+20 ms slack).
2128                                        // An absolute cap missed the worst
2129                                        // case — contended ops sit at
2130                                        // 100–240 ms each and still bury a
2131                                        // render whose fair op is 3–9 ms.
2132                                        // Cold ops (first PSO build, buffer
2133                                        // alloc) are exempt: a one-off
2134                                        // ~50 ms compile is not contention.
2135                                        let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
2136                                        let budget = std::time::Duration::from_secs_f64(
2137                                            flops / 1.5e12 * 8.0 + 0.020,
2138                                        );
2139                                        crate::gpu::mm_budget_check(
2140                                            "q4t matmat",
2141                                            el,
2142                                            budget,
2143                                            crate::gpu::probe_was_cold(),
2144                                        );
2145                                        crate::gpu::probe_record(class, true, el);
2146                                        return;
2147                                    }
2148                                }
2149                                crate::gpu::ProbeArm::CpuTimed => {
2150                                    q4t_matmat(
2151                                        self.quant_bytes(),
2152                                        xs_all,
2153                                        b,
2154                                        rows,
2155                                        cols,
2156                                        out,
2157                                        pool,
2158                                    );
2159                                    crate::gpu::probe_record(class, false, t0.elapsed());
2160                                    return;
2161                                }
2162                                crate::gpu::ProbeArm::Cpu => {}
2163                            }
2164                        }
2165                    }
2166                    q4t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2167                    return;
2168                }
2169                if *dtype == TensorDtype::Q1 {
2170                    // GPU batched q1 GEMM for wide prefill (q1_mul_mm on the
2171                    // device); the probe keeps whichever beats the CPU matmat.
2172                    if b >= 32
2173                        && b * rows * cols >= 128_000_000
2174                        && cols % 64 == 0
2175                        && crate::gpu::enabled_here()
2176                    {
2177                        if let Self::Mapped { model, idx, .. } = self {
2178                            let t0 = std::time::Instant::now();
2179                            match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
2180                                crate::gpu::ProbeArm::Gpu => {
2181                                    if crate::gpu::q1_matmat(
2182                                        model, *idx, xs_all, b, rows, cols, out,
2183                                    ) {
2184                                        crate::gpu::probe_record(
2185                                            crate::gpu::OpClass::Matmat,
2186                                            true,
2187                                            t0.elapsed(),
2188                                        );
2189                                        return;
2190                                    }
2191                                }
2192                                crate::gpu::ProbeArm::CpuTimed => {
2193                                    q1_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2194                                    crate::gpu::probe_record(
2195                                        crate::gpu::OpClass::Matmat,
2196                                        false,
2197                                        t0.elapsed(),
2198                                    );
2199                                    return;
2200                                }
2201                                crate::gpu::ProbeArm::Cpu => {}
2202                            }
2203                        }
2204                    }
2205                    q1_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2206                    return;
2207                }
2208                if *dtype == TensorDtype::Q1T {
2209                    // GPU batched GEMM for wide prefill (base + overlay on the
2210                    // device); probe keeps the winner vs the CPU matmat.
2211                    if b >= 32 && b * rows * cols >= 128_000_000 && crate::gpu::enabled_here() {
2212                        if let Self::Mapped { model, idx, .. } = self {
2213                            let t0 = std::time::Instant::now();
2214                            match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
2215                                crate::gpu::ProbeArm::Gpu => {
2216                                    if crate::gpu::q1t_matmat(
2217                                        model, *idx, xs_all, b, rows, cols, out,
2218                                    ) {
2219                                        crate::gpu::probe_record(
2220                                            crate::gpu::OpClass::Matmat,
2221                                            true,
2222                                            t0.elapsed(),
2223                                        );
2224                                        return;
2225                                    }
2226                                }
2227                                crate::gpu::ProbeArm::CpuTimed => {
2228                                    q1t_matmat(
2229                                        self.quant_bytes(),
2230                                        xs_all,
2231                                        b,
2232                                        rows,
2233                                        cols,
2234                                        out,
2235                                        pool,
2236                                    );
2237                                    crate::gpu::probe_record(
2238                                        crate::gpu::OpClass::Matmat,
2239                                        false,
2240                                        t0.elapsed(),
2241                                    );
2242                                    return;
2243                                }
2244                                crate::gpu::ProbeArm::Cpu => {}
2245                            }
2246                        }
2247                    }
2248                    q1t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2249                    return;
2250                }
2251                if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
2252                    vbitmatmat(
2253                        self.quant_bytes(),
2254                        vbit_offsets,
2255                        xs_all,
2256                        b,
2257                        rows,
2258                        cols,
2259                        out,
2260                        pool,
2261                    );
2262                    return;
2263                }
2264                let pre: Vec<std::borrow::Cow<'_, [f32]>> = (0..b)
2265                    .map(|bi| prescale(&xs_all[bi * cols..(bi + 1) * cols], col_field, *dtype))
2266                    .collect();
2267                // D5: large prefill-batch GEMMs — on the GPU (threshold by
2268                // work volume: submission carries b×rows×cols MACs).
2269                // Runtime probe: the naive GEMM shader + sync readback
2270                // lose to the CPU GEMM on slow driver stacks — alternate
2271                // both arms and keep the winner.
2272                if b >= 8 && b * rows * cols >= 128_000_000 && crate::gpu::enabled_here() {
2273                    if let Self::Mapped { model, idx, .. } = self {
2274                        let t0 = std::time::Instant::now();
2275                        match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
2276                            crate::gpu::ProbeArm::Gpu
2277                                if crate::gpu::probe_deciding(crate::gpu::OpClass::Matmat)
2278                                    && !crate::gpu::q8_resident_or_upload(model, *idx) =>
2279                            {
2280                                // Cold weights during probing: the upload
2281                                // has started, the count runs on the CPU —
2282                                // the GPU arm samples on the next touch.
2283                                let q = self.quant_bytes();
2284                                qmatmat(q, row_scale, &pre, rows, cols, out, pool);
2285                                return;
2286                            }
2287                            crate::gpu::ProbeArm::Gpu => {
2288                                let flat: Vec<f32> =
2289                                    pre.iter().flat_map(|v| v.iter().copied()).collect();
2290                                if crate::gpu::q8_matmat(
2291                                    model, *idx, row_scale, &flat, b, rows, cols, out,
2292                                ) {
2293                                    crate::gpu::probe_record(
2294                                        crate::gpu::OpClass::Matmat,
2295                                        true,
2296                                        t0.elapsed(),
2297                                    );
2298                                    return;
2299                                }
2300                            }
2301                            crate::gpu::ProbeArm::CpuTimed => {
2302                                let q = self.quant_bytes();
2303                                qmatmat(q, row_scale, &pre, rows, cols, out, pool);
2304                                crate::gpu::probe_record(
2305                                    crate::gpu::OpClass::Matmat,
2306                                    false,
2307                                    t0.elapsed(),
2308                                );
2309                                return;
2310                            }
2311                            crate::gpu::ProbeArm::Cpu => {}
2312                        }
2313                    }
2314                }
2315                let q = self.quant_bytes();
2316                qmatmat(q, row_scale, &pre, rows, cols, out, pool);
2317            }
2318        }
2319    }
2320}
2321
2322impl QTensor {
2323    /// The device GEMM this tensor would take, run once on the caller's
2324    /// data — the startup parity probe's arm, and the one place that knows
2325    /// which entry point each codec has.
2326    ///
2327    /// It exists because the probe used to look for a `q4tp` weight by
2328    /// name AND dtype, and a container packed any other way was declared
2329    /// "host path" for the whole render even though its codec had a device
2330    /// GEMM of its own. A gate that only recognizes one codec is a gate
2331    /// that silently downgrades every other one.
2332    pub fn device_matmat(&self, xs: &[f32], b: usize, out: &mut [f32]) -> bool {
2333        let (rows, cols) = (self.rows(), self.cols());
2334        let Self::Mapped {
2335            model,
2336            idx,
2337            dtype,
2338            row_scale,
2339            col_field,
2340            ..
2341        } = self
2342        else {
2343            return false;
2344        };
2345        if crate::prism::has_contract(model) {
2346            return false;
2347        }
2348        match *dtype {
2349            TensorDtype::Q4TiledP => crate::gpu::q4tp_matmat(model, *idx, xs, b, rows, cols, out),
2350            // The two-field codec folds its column field into the
2351            // activation, which leaves a plain per-row int8 GEMM — the
2352            // same kernel `q8_row` uses, on both backends.
2353            TensorDtype::Q8Row | TensorDtype::Q8_2f => {
2354                // The field belongs to the weight; only a backend that cannot
2355                // apply it there makes a scaled copy of the activation.
2356                if *dtype == TensorDtype::Q8_2f
2357                    && std::env::var("CMF_Q8_2F_DEV").as_deref() != Ok("0")
2358                    && crate::gpu::q8_matmat_2f(
2359                        model, *idx, row_scale, col_field, xs, b, rows, cols, out,
2360                    )
2361                {
2362                    return true;
2363                }
2364                let flat: Vec<f32> = (0..b)
2365                    .flat_map(|bi| {
2366                        prescale(&xs[bi * cols..(bi + 1) * cols], col_field, *dtype).into_owned()
2367                    })
2368                    .collect();
2369                crate::gpu::q8_matmat(model, *idx, row_scale, &flat, b, rows, cols, out)
2370            }
2371            _ => false,
2372        }
2373    }
2374
2375    /// Multi-matrix job (roadmap §3 P0): N tensors sharing one input
2376    /// run under a SINGLE pool dispatch — QKV or gate+up cost one
2377    /// barrier instead of N. Per-row math is the exact same kernel as
2378    /// `matvec` (bit-identical outputs); only the dispatch is fused.
2379    /// Falls back to N sequential matvecs when the set is not a uniform
2380    /// q8-family/F32 group or there is no pool.
2381    pub fn matvec_many<const N: usize>(
2382        ts: [&QTensor; N],
2383        x: &[f32],
2384        mut outs: [&mut [f32]; N],
2385        pool: Option<&Pool>,
2386    ) {
2387        let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
2388        if ts.iter().any(|t| t.has_prism_contract()) {
2389            // The fused range kernels have no transform descriptor.  Let
2390            // each tensor's ordinary matvec dispatch perform the explicit
2391            // signed FWHT (and retain CPU fallback for mixed q2tp/q4tp).
2392            for (t, o) in ts.iter().zip(outs.iter_mut()) {
2393                t.matvec(x, o, pool);
2394            }
2395            return;
2396        }
2397        let uniform_q8 = ts.iter().all(|t| {
2398            matches!(
2399                t,
2400                Self::Mapped {
2401                    dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f,
2402                    ..
2403                }
2404            )
2405        });
2406        let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
2407        let uniform_q4 = ts.iter().all(|t| {
2408            matches!(
2409                t,
2410                Self::Mapped {
2411                    dtype: TensorDtype::Q4Block,
2412                    ..
2413                }
2414            )
2415        });
2416        let uniform_vbit = ts.iter().all(|t| {
2417            matches!(
2418                t,
2419                Self::Mapped {
2420                    dtype: TensorDtype::Vbit | TensorDtype::VbitRo,
2421                    ..
2422                }
2423            )
2424        });
2425        let uniform_q1 = ts.iter().all(|t| {
2426            matches!(
2427                t,
2428                Self::Mapped {
2429                    dtype: TensorDtype::Q1,
2430                    ..
2431                }
2432            )
2433        });
2434        let uniform_q1t = ts.iter().all(|t| {
2435            matches!(
2436                t,
2437                Self::Mapped {
2438                    dtype: TensorDtype::Q1T,
2439                    ..
2440                }
2441            )
2442        });
2443        // q4tp is the skeleton dtype of the big MoE files, and without an arm
2444        // here every projection that shares an input paid its own pool
2445        // barrier: DeepSeek-V4's attention step alone hands this function
2446        // wq_a, wkv and both compressors' pairs off the same hidden state.
2447        let uniform_q4tp = ts.iter().all(|t| {
2448            matches!(
2449                t,
2450                Self::Mapped {
2451                    dtype: TensorDtype::Q4TiledP,
2452                    ..
2453                }
2454            )
2455        }) && ts
2456            .iter()
2457            .all(|t| t.cols() == ts[0].cols() && t.cols() % GROUP_SIZE == 0);
2458        let Some(pool) = pool else {
2459            for (t, o) in ts.iter().zip(outs.iter_mut()) {
2460                t.matvec(x, o, None);
2461            }
2462            return;
2463        };
2464        if total_rows < 256
2465            || !(uniform_q8
2466                || uniform_f32
2467                || uniform_q4
2468                || uniform_vbit
2469                || uniform_q1
2470                || uniform_q1t
2471                || uniform_q4tp)
2472        {
2473            for (t, o) in ts.iter().zip(outs.iter_mut()) {
2474                t.matvec(x, o, Some(pool));
2475            }
2476            return;
2477        }
2478
2479        if uniform_q4tp {
2480            // Every tensor's rows laid end to end in one virtual row space,
2481            // so the whole set is ONE dispatch. The per-row body is the
2482            // `q4tp_matvec` arm verbatim — same activation split, same
2483            // accumulation order — so the outputs are bit-identical to the
2484            // sequential calls this replaces.
2485            let cols = ts[0].cols();
2486            let gpr = cols / GROUP_SIZE;
2487            let views: Vec<Q4tpView> = ts
2488                .iter()
2489                .map(|t| Q4tpView::new(t.quant_bytes(), t.rows(), cols))
2490                .collect();
2491            let rows_of: Vec<usize> = ts.iter().map(|t| t.rows()).collect();
2492            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2493            // flat index -> (which tensor, which of its rows)
2494            let locate = |flat: usize| -> (usize, usize) {
2495                let mut acc = 0;
2496                for (i, &r) in rows_of.iter().enumerate() {
2497                    if flat < acc + r {
2498                        return (i, flat - acc);
2499                    }
2500                    acc += r;
2501                }
2502                (rows_of.len() - 1, 0)
2503            };
2504            let (views, outs_addr) = (&views, &outs_addr);
2505            if a8w8_enabled() {
2506                let act = split_act(x);
2507                let act = &act;
2508                let run = |start: usize, end: usize| {
2509                    let mut sc = vec![0f32; gpr];
2510                    for flat in start..end {
2511                        let (t, r) = locate(flat);
2512                        let v = &views[t];
2513                        v.scales_into(r, gpr, &mut sc);
2514                        let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, &sc) * act.sx;
2515                        for &(j, xv) in &act.outliers {
2516                            let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
2517                            acc += w * s * xv;
2518                        }
2519                        // SAFETY: one worker owns each (tensor, row) pair.
2520                        unsafe { *outs_addr[t].at(r) = acc };
2521                    }
2522                };
2523                pool.run_rows(total_rows, &run);
2524            } else {
2525                let run = |start: usize, end: usize| {
2526                    let mut sc = vec![0f32; gpr];
2527                    for flat in start..end {
2528                        let (t, r) = locate(flat);
2529                        let v = &views[t];
2530                        v.scales_into(r, gpr, &mut sc);
2531                        // SAFETY: one worker owns each (tensor, row) pair.
2532                        unsafe { *outs_addr[t].at(r) = q4tp_row_exact(v.nib, r, gpr, x, &sc) };
2533                    }
2534                };
2535                pool.run_rows(total_rows, &run);
2536            }
2537            return;
2538        }
2539
2540        if uniform_q1 {
2541            // One shared activation split + group sums (q1 has no col
2542            // field; the same input feeds every tensor).
2543            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2544            if a8w8_enabled() {
2545                let act = split_act(x);
2546                let gsum = q1_group_sums(&act.xq, ts[0].cols() / GROUP_SIZE);
2547                let (act, gsum) = (&act, &gsum);
2548                let closures: [_; N] = std::array::from_fn(|i| {
2549                    let (bytes, gpr, out) =
2550                        (ts[i].quant_bytes(), ts[i].cols() / GROUP_SIZE, outs_addr[i]);
2551                    move |s: usize, e: usize| q1_range_a8w8(bytes, gpr, act, gsum, out, s, e)
2552                });
2553                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2554                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2555                pool.run_many(&parts);
2556            } else {
2557                let closures: [_; N] = std::array::from_fn(|i| {
2558                    let (bytes, gpr, out) =
2559                        (ts[i].quant_bytes(), ts[i].cols() / GROUP_SIZE, outs_addr[i]);
2560                    move |s: usize, e: usize| q1_range_f32(bytes, gpr, x, out, s, e)
2561                });
2562                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2563                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2564                pool.run_many(&parts);
2565            }
2566            return;
2567        }
2568
2569        if uniform_q1t {
2570            // Q1T batched: one shared activation split + overlay decode,
2571            // all tensors' rows in ONE pool dispatch (saves N−1 dispatches
2572            // and N−1 redundant split_act calls per layer).
2573            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2574            const TILE: usize = cortiq_core::quant::Q1T_TILE;
2575            if a8w8_enabled() {
2576                let act = split_act(x);
2577                let act = &act;
2578                let x_ref = x;
2579                let closures: [_; N] = std::array::from_fn(|i| {
2580                    let bytes = ts[i].quant_bytes();
2581                    let (rows, cols) = (ts[i].rows(), ts[i].cols());
2582                    let gpr = cols / GROUP_SIZE;
2583                    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
2584                    let out = outs_addr[i];
2585                    move |s: usize, e: usize| {
2586                        q1t_range_a8w8(bytes, gpr, rp_off, ent_off, has_ov, act, x_ref, out, s, e)
2587                    }
2588                });
2589                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2590                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2591                pool.run_many(&parts);
2592            } else {
2593                let x_ref = x;
2594                let closures: [_; N] = std::array::from_fn(|i| {
2595                    let bytes = ts[i].quant_bytes();
2596                    let (rows, cols) = (ts[i].rows(), ts[i].cols());
2597                    let gpr = cols / GROUP_SIZE;
2598                    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
2599                    let out = outs_addr[i];
2600                    move |s: usize, e: usize| {
2601                        q1t_range_f32_batch(bytes, gpr, rp_off, ent_off, has_ov, x_ref, out, s, e)
2602                    }
2603                });
2604                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2605                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2606                pool.run_many(&parts);
2607            }
2608            return;
2609        }
2610
2611        if uniform_q4 || uniform_vbit {
2612            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2613            // q4/vbit share one activation split — no per-tensor col field.
2614            if a8w8_enabled() {
2615                let act = split_act(x);
2616                let act = &act;
2617                if uniform_q4 {
2618                    let closures: [_; N] = std::array::from_fn(|i| {
2619                        let (packed, scales) =
2620                            q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
2621                        let (gpr, cols, out) =
2622                            (ts[i].cols() / GROUP_SIZE, ts[i].cols(), outs_addr[i]);
2623                        move |s: usize, e: usize| {
2624                            q4_range_a8w8(packed, scales, gpr, cols, act, out, s, e)
2625                        }
2626                    });
2627                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2628                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2629                    pool.run_many(&parts);
2630                } else {
2631                    let closures: [_; N] = std::array::from_fn(|i| {
2632                        let Self::Mapped { vbit_offsets, .. } = ts[i] else {
2633                            unreachable!()
2634                        };
2635                        let (bytes, rows, cols, out) = (
2636                            ts[i].quant_bytes(),
2637                            ts[i].rows(),
2638                            ts[i].cols(),
2639                            outs_addr[i],
2640                        );
2641                        move |s: usize, e: usize| {
2642                            vbit_range_a8w8(bytes, vbit_offsets, x, act, rows, cols, out, s, e)
2643                        }
2644                    });
2645                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2646                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2647                    pool.run_many(&parts);
2648                }
2649                return;
2650            }
2651            if uniform_q4 {
2652                let closures: [_; N] = std::array::from_fn(|i| {
2653                    let (packed, scales) =
2654                        q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
2655                    let (gpr, out) = (ts[i].cols() / GROUP_SIZE, outs_addr[i]);
2656                    move |s: usize, e: usize| q4_range_f32(packed, scales, gpr, x, out, s, e)
2657                });
2658                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2659                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2660                pool.run_many(&parts);
2661            } else {
2662                let closures: [_; N] = std::array::from_fn(|i| {
2663                    let Self::Mapped { vbit_offsets, .. } = ts[i] else {
2664                        unreachable!()
2665                    };
2666                    let (bytes, rows, cols, out) = (
2667                        ts[i].quant_bytes(),
2668                        ts[i].rows(),
2669                        ts[i].cols(),
2670                        outs_addr[i],
2671                    );
2672                    move |s: usize, e: usize| {
2673                        vbit_range_f32(bytes, vbit_offsets, x, rows, cols, out, s, e)
2674                    }
2675                });
2676                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2677                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2678                pool.run_many(&parts);
2679            }
2680            return;
2681        }
2682
2683        if uniform_f32 {
2684            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2685            let closures: [_; N] = std::array::from_fn(|i| {
2686                let Self::F32 { data, cols, .. } = ts[i] else {
2687                    unreachable!()
2688                };
2689                let out = outs_addr[i];
2690                move |start: usize, end: usize| {
2691                    for o in start..end {
2692                        let row = &data[o * cols..(o + 1) * cols];
2693                        let mut sum = 0.0f32;
2694                        for j in 0..*cols {
2695                            sum += row[j] * x[j];
2696                        }
2697                        // SAFETY: disjoint (tensor, row) cells per worker.
2698                        unsafe { *out.at(o) = sum };
2699                    }
2700                }
2701            });
2702            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2703                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2704            pool.run_many(&parts);
2705            return;
2706        }
2707
2708        // Uniform q8-family: per-tensor prescale (q8_2f col fields
2709        // differ per tensor) + the shared range kernels.
2710        struct Ctx<'a> {
2711            bytes: &'a [u8],
2712            #[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
2713            rep: &'a [u8],
2714            row_scale: &'a [f32],
2715            cols: usize,
2716            xs: std::borrow::Cow<'a, [f32]>,
2717        }
2718        let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
2719            let Self::Mapped {
2720                dtype,
2721                cols,
2722                row_scale,
2723                col_field,
2724                repack,
2725                ..
2726            } = ts[i]
2727            else {
2728                unreachable!()
2729            };
2730            Ctx {
2731                bytes: ts[i].quant_bytes(),
2732                rep: repack,
2733                row_scale,
2734                cols: *cols,
2735                xs: prescale(x, col_field, *dtype),
2736            }
2737        });
2738        let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2739        #[cfg(target_arch = "aarch64")]
2740        if sdot_enabled() {
2741            let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
2742            let closures: [_; N] = std::array::from_fn(|i| {
2743                let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
2744                move |start: usize, end: usize| {
2745                    q8_range_sdot(c.bytes, c.rep, c.row_scale, act, c.cols, out, start, end)
2746                }
2747            });
2748            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2749                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2750            pool.run_many(&parts);
2751            return;
2752        }
2753        #[cfg(target_arch = "x86_64")]
2754        if avx2_a8w8_enabled() {
2755            let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
2756            let closures: [_; N] = std::array::from_fn(|i| {
2757                let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
2758                move |start: usize, end: usize| {
2759                    q8_range_avx2(c.bytes, c.row_scale, act, c.cols, out, start, end)
2760                }
2761            });
2762            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2763                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2764            pool.run_many(&parts);
2765            return;
2766        }
2767        let closures: [_; N] = std::array::from_fn(|i| {
2768            let (c, out) = (&ctxs[i], outs_addr[i]);
2769            move |start: usize, end: usize| {
2770                q8_range_f32(c.bytes, c.row_scale, &c.xs, c.cols, out, start, end)
2771            }
2772        });
2773        let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2774            std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2775        pool.run_many(&parts);
2776    }
2777}
2778
2779impl QTensor {
2780    /// Pair-input multi-matrix job: N tensors × 2 shared inputs under a
2781    /// single pool dispatch — the MTP/pair decode path publishes one job
2782    /// for Q/K/V (and one for gate+up) instead of one per tensor.
2783    /// Per-row math is exactly `matvec2`'s kernels; bit-identical.
2784    #[allow(clippy::needless_range_loop)]
2785    pub fn matvec2_many<const N: usize>(
2786        ts: [&QTensor; N],
2787        x1: &[f32],
2788        x2: &[f32],
2789        mut o1s: [&mut [f32]; N],
2790        mut o2s: [&mut [f32]; N],
2791        pool: Option<&Pool>,
2792    ) {
2793        let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
2794        if ts.iter().any(|t| t.has_prism_contract()) {
2795            for i in 0..N {
2796                ts[i].matvec2(x1, x2, o1s[i], o2s[i], pool);
2797            }
2798            return;
2799        }
2800        let uniform_q8 = ts.iter().all(|t| {
2801            matches!(
2802                t,
2803                Self::Mapped {
2804                    dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f,
2805                    ..
2806                }
2807            )
2808        });
2809        let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
2810        let uniform_q4 = ts.iter().all(|t| {
2811            matches!(
2812                t,
2813                Self::Mapped {
2814                    dtype: TensorDtype::Q4Block,
2815                    ..
2816                }
2817            )
2818        });
2819        let uniform_vbit = ts.iter().all(|t| {
2820            matches!(
2821                t,
2822                Self::Mapped {
2823                    dtype: TensorDtype::Vbit | TensorDtype::VbitRo,
2824                    ..
2825                }
2826            )
2827        });
2828        let fusable = pool.is_some()
2829            && total_rows >= 256
2830            && (uniform_q8 || uniform_f32 || uniform_q4 || uniform_vbit);
2831        if !fusable {
2832            for i in 0..N {
2833                ts[i].matvec2(x1, x2, o1s[i], o2s[i], pool);
2834            }
2835            return;
2836        }
2837        let pool = pool.unwrap();
2838
2839        if uniform_q4 || uniform_vbit {
2840            let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
2841            let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
2842            // q4/vbit share activation splits — no per-tensor col field.
2843            if a8w8_enabled() {
2844                let a1 = split_act(x1);
2845                let a2 = split_act(x2);
2846                let (a1, a2) = (&a1, &a2);
2847                if uniform_q4 {
2848                    let closures: [_; N] = std::array::from_fn(|i| {
2849                        let (packed, scales) =
2850                            q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
2851                        let (gpr, cols, o1, o2) =
2852                            (ts[i].cols() / GROUP_SIZE, ts[i].cols(), p1[i], p2[i]);
2853                        move |s: usize, e: usize| {
2854                            q4_range2_a8w8(packed, scales, gpr, cols, a1, a2, o1, o2, s, e)
2855                        }
2856                    });
2857                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2858                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2859                    pool.run_many(&parts);
2860                } else {
2861                    let closures: [_; N] = std::array::from_fn(|i| {
2862                        let Self::Mapped { vbit_offsets, .. } = ts[i] else {
2863                            unreachable!()
2864                        };
2865                        let (bytes, rows, cols, o1, o2) = (
2866                            ts[i].quant_bytes(),
2867                            ts[i].rows(),
2868                            ts[i].cols(),
2869                            p1[i],
2870                            p2[i],
2871                        );
2872                        move |s: usize, e: usize| {
2873                            vbit_range2_a8w8(
2874                                bytes,
2875                                vbit_offsets,
2876                                x1,
2877                                x2,
2878                                a1,
2879                                a2,
2880                                rows,
2881                                cols,
2882                                o1,
2883                                o2,
2884                                s,
2885                                e,
2886                            )
2887                        }
2888                    });
2889                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2890                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2891                    pool.run_many(&parts);
2892                }
2893                return;
2894            }
2895            if uniform_q4 {
2896                let closures: [_; N] = std::array::from_fn(|i| {
2897                    let (packed, scales) =
2898                        q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
2899                    let (gpr, o1, o2) = (ts[i].cols() / GROUP_SIZE, p1[i], p2[i]);
2900                    move |s: usize, e: usize| {
2901                        q4_range2_f32(packed, scales, gpr, x1, x2, o1, o2, s, e)
2902                    }
2903                });
2904                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2905                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2906                pool.run_many(&parts);
2907            } else {
2908                let closures: [_; N] = std::array::from_fn(|i| {
2909                    let Self::Mapped { vbit_offsets, .. } = ts[i] else {
2910                        unreachable!()
2911                    };
2912                    let (bytes, rows, cols, o1, o2) = (
2913                        ts[i].quant_bytes(),
2914                        ts[i].rows(),
2915                        ts[i].cols(),
2916                        p1[i],
2917                        p2[i],
2918                    );
2919                    move |s: usize, e: usize| {
2920                        vbit_range2_f32(bytes, vbit_offsets, x1, x2, rows, cols, o1, o2, s, e)
2921                    }
2922                });
2923                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2924                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2925                pool.run_many(&parts);
2926            }
2927            return;
2928        }
2929
2930        if uniform_f32 {
2931            let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
2932            let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
2933            let closures: [_; N] = std::array::from_fn(|i| {
2934                let Self::F32 { data, cols, .. } = ts[i] else {
2935                    unreachable!()
2936                };
2937                let (o1, o2) = (p1[i], p2[i]);
2938                move |start: usize, end: usize| {
2939                    for o in start..end {
2940                        let row = &data[o * cols..(o + 1) * cols];
2941                        let (mut s1, mut s2) = (0.0f32, 0.0f32);
2942                        for j in 0..*cols {
2943                            s1 += row[j] * x1[j];
2944                            s2 += row[j] * x2[j];
2945                        }
2946                        // SAFETY: disjoint (tensor, row) cells per worker.
2947                        unsafe {
2948                            *o1.at(o) = s1;
2949                            *o2.at(o) = s2;
2950                        }
2951                    }
2952                }
2953            });
2954            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2955                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2956            pool.run_many(&parts);
2957            return;
2958        }
2959
2960        struct Ctx<'a> {
2961            bytes: &'a [u8],
2962            row_scale: &'a [f32],
2963            cols: usize,
2964            xs1: std::borrow::Cow<'a, [f32]>,
2965            xs2: std::borrow::Cow<'a, [f32]>,
2966        }
2967        let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
2968            let Self::Mapped {
2969                dtype,
2970                cols,
2971                row_scale,
2972                col_field,
2973                ..
2974            } = ts[i]
2975            else {
2976                unreachable!()
2977            };
2978            Ctx {
2979                bytes: ts[i].quant_bytes(),
2980                row_scale,
2981                cols: *cols,
2982                xs1: prescale(x1, col_field, *dtype),
2983                xs2: prescale(x2, col_field, *dtype),
2984            }
2985        });
2986        let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
2987        let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
2988        #[cfg(target_arch = "aarch64")]
2989        if sdot_enabled() {
2990            let acts: [(SplitAct, SplitAct); N] =
2991                std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
2992            let closures: [_; N] = std::array::from_fn(|i| {
2993                let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
2994                move |start: usize, end: usize| {
2995                    q8_range2_sdot(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
2996                }
2997            });
2998            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2999                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3000            pool.run_many(&parts);
3001            return;
3002        }
3003        #[cfg(target_arch = "x86_64")]
3004        if avx2_a8w8_enabled() {
3005            let acts: [(SplitAct, SplitAct); N] =
3006                std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
3007            let closures: [_; N] = std::array::from_fn(|i| {
3008                let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
3009                move |start: usize, end: usize| {
3010                    q8_range2_avx2(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
3011                }
3012            });
3013            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3014                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3015            pool.run_many(&parts);
3016            return;
3017        }
3018        let closures: [_; N] = std::array::from_fn(|i| {
3019            let (c, o1, o2) = (&ctxs[i], p1[i], p2[i]);
3020            move |start: usize, end: usize| {
3021                q8_range2_f32(
3022                    c.bytes,
3023                    c.row_scale,
3024                    &c.xs1,
3025                    &c.xs2,
3026                    c.cols,
3027                    o1,
3028                    o2,
3029                    start,
3030                    end,
3031                )
3032            }
3033        });
3034        let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3035            std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3036        pool.run_many(&parts);
3037    }
3038
3039    /// Fused gate+up matvec with SiLU·mul: for each row r, computes
3040    /// `silu(gate·x) * (up·x)` and writes to `out[r]`. ONE pool dispatch,
3041    /// no intermediate g/u buffers, no separate silu pass. Falls back
3042    /// (returns false) for unsupported dtype combos.
3043    pub fn matvec_silu_mul(
3044        gate: &QTensor,
3045        up: &QTensor,
3046        x: &[f32],
3047        out: &mut [f32],
3048        pool: Option<&Pool>,
3049    ) -> bool {
3050        Self::matvec_silu_mul_limited(gate, up, x, out, 0.0, pool)
3051    }
3052
3053    /// Fused gate+up+SiLU with the GLM asymmetrical clamp.  `limit == 0`
3054    /// preserves the historical unclamped helper; a positive limit clamps
3055    /// `up` to both sides and `gate` only from above, matching the GLM
3056    /// SwiGLU reference.  Keeping the limit in the row kernel avoids the two
3057    /// intermediate vectors and the extra combine pass on the Q2TP experts.
3058    pub fn matvec_silu_mul_limited(
3059        gate: &QTensor,
3060        up: &QTensor,
3061        x: &[f32],
3062        out: &mut [f32],
3063        limit: f32,
3064        pool: Option<&Pool>,
3065    ) -> bool {
3066        if gate.has_prism_contract() || up.has_prism_contract() {
3067            // The fused gate/up kernels consume x directly.  Prism requires
3068            // a per-matrix signed FWHT, so the caller must use two ordinary
3069            // descriptor-aware matvecs instead of an unrotated fast path.
3070            return false;
3071        }
3072        let inter = gate.rows();
3073        debug_assert_eq!(up.rows(), inter);
3074        debug_assert_eq!(out.len(), inter);
3075        debug_assert_eq!(gate.cols(), up.cols());
3076        if !a8w8_enabled() {
3077            return false;
3078        }
3079        let act = split_act(x);
3080        let act = &act;
3081        let x_ref = x;
3082        let out_addr = SendMut(out.as_mut_ptr());
3083
3084        match (gate, up) {
3085            // Q4Block gate + Q4Block up (most common mobile q4 models)
3086            (
3087                Self::Mapped {
3088                    dtype: TensorDtype::Q4Block,
3089                    ..
3090                },
3091                Self::Mapped {
3092                    dtype: TensorDtype::Q4Block,
3093                    ..
3094                },
3095            ) => {
3096                let (gp, gs) = q4_split(gate.quant_bytes(), gate.rows(), gate.cols());
3097                let (up_p, up_s) = q4_split(up.quant_bytes(), up.rows(), up.cols());
3098                let gpr = gate.cols() / GROUP_SIZE;
3099                let cols = gate.cols();
3100                let run = move |start: usize, end: usize| {
3101                    for r in start..end {
3102                        let mut gv = dot_q4_row_i8(gp, gs, r * gpr, gpr, &act.xq) * act.sx;
3103                        let mut uv = dot_q4_row_i8(up_p, up_s, r * gpr, gpr, &act.xq) * act.sx;
3104                        for &(j, xv) in &act.outliers {
3105                            let flat = r * cols + j;
3106                            let gb = gp[flat / 2];
3107                            let gn = if flat & 1 == 0 { gb & 0x0F } else { gb >> 4 };
3108                            let gsc = f16_to_f32(u16::from_le_bytes([
3109                                gs[(flat / GROUP_SIZE) * 2],
3110                                gs[(flat / GROUP_SIZE) * 2 + 1],
3111                            ]));
3112                            gv += ((gn as i32 - 8) as f32) * gsc * xv;
3113                            let ub = up_p[flat / 2];
3114                            let un = if flat & 1 == 0 { ub & 0x0F } else { ub >> 4 };
3115                            let usc = f16_to_f32(u16::from_le_bytes([
3116                                up_s[(flat / GROUP_SIZE) * 2],
3117                                up_s[(flat / GROUP_SIZE) * 2 + 1],
3118                            ]));
3119                            uv += ((un as i32 - 8) as f32) * usc * xv;
3120                        }
3121                        let silu_g = gv / (1.0 + (-gv).exp());
3122                        // SAFETY: disjoint row ranges per worker.
3123                        unsafe { *out_addr.at(r) = silu_g * uv };
3124                    }
3125                };
3126                dispatch_rows(pool, inter, &run);
3127                true
3128            }
3129            // Q4Tiled gate + Q4Tiled up — one row pass, both tile
3130            // streams sequential, silu·mul fused (same per-row math as
3131            // `q4t_matvec`).
3132            (
3133                Self::Mapped {
3134                    dtype: TensorDtype::Q4Tiled,
3135                    ..
3136                },
3137                Self::Mapped {
3138                    dtype: TensorDtype::Q4Tiled,
3139                    ..
3140                },
3141            ) => {
3142                let g_bytes = gate.quant_bytes();
3143                let u_bytes = up.quant_bytes();
3144                let gpr = gate.cols() / GROUP_SIZE;
3145                let run = move |start: usize, end: usize| {
3146                    for r in start..end {
3147                        let mut gv = dot_q4t_row_i8(g_bytes, r, gpr, &act.xq) * act.sx;
3148                        let mut uv = dot_q4t_row_i8(u_bytes, r, gpr, &act.xq) * act.sx;
3149                        for &(j, xv) in &act.outliers {
3150                            let (w, s) = q4t_outlier(g_bytes, r, gpr, j);
3151                            gv += w * s * xv;
3152                            let (w, s) = q4t_outlier(u_bytes, r, gpr, j);
3153                            uv += w * s * xv;
3154                        }
3155                        let silu_g = gv / (1.0 + (-gv).exp());
3156                        // SAFETY: disjoint row ranges per worker.
3157                        unsafe { *out_addr.at(r) = silu_g * uv };
3158                    }
3159                };
3160                dispatch_rows(pool, inter, &run);
3161                true
3162            }
3163            // Q4TiledP gate + Q4TiledP up — the same fused row pass, with
3164            // each row's two ladders built once and spent on both streams.
3165            (
3166                Self::Mapped {
3167                    dtype: TensorDtype::Q4TiledP,
3168                    ..
3169                },
3170                Self::Mapped {
3171                    dtype: TensorDtype::Q4TiledP,
3172                    ..
3173                },
3174            ) => {
3175                let cols = gate.cols();
3176                let gpr = cols / GROUP_SIZE;
3177                let gv_view = Q4tpView::new(gate.quant_bytes(), inter, cols);
3178                let uv_view = Q4tpView::new(up.quant_bytes(), inter, cols);
3179                let run = |start: usize, end: usize| {
3180                    let (mut gsc, mut usc) = (vec![0f32; gpr], vec![0f32; gpr]);
3181                    for r in start..end {
3182                        gv_view.scales_into(r, gpr, &mut gsc);
3183                        uv_view.scales_into(r, gpr, &mut usc);
3184                        let mut gv = dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, &gsc) * act.sx;
3185                        let mut uv = dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, &usc) * act.sx;
3186                        for &(j, xv) in &act.outliers {
3187                            let (w, s) = q4tp_outlier(gv_view.nib, r, gpr, j, &gsc);
3188                            gv += w * s * xv;
3189                            let (w, s) = q4tp_outlier(uv_view.nib, r, gpr, j, &usc);
3190                            uv += w * s * xv;
3191                        }
3192                        let silu_g = gv / (1.0 + (-gv).exp());
3193                        // SAFETY: disjoint row ranges per worker.
3194                        unsafe { *out_addr.at(r) = silu_g * uv };
3195                    }
3196                };
3197                dispatch_rows(pool, inter, &run);
3198                true
3199            }
3200            // Q1 gate + Q1 up — one row pass over both sign streams,
3201            // silu·mul fused (the per-row math of `q1_range_a8w8`); the
3202            // activation group sums are shared by both streams. Without
3203            // this arm a q1 dense FFN paid two dispatches + a combine
3204            // loop — the exact barrier this function exists to remove.
3205            (
3206                Self::Mapped {
3207                    dtype: TensorDtype::Q1,
3208                    ..
3209                },
3210                Self::Mapped {
3211                    dtype: TensorDtype::Q1,
3212                    ..
3213                },
3214            ) => {
3215                let g_bytes = gate.quant_bytes();
3216                let u_bytes = up.quant_bytes();
3217                let gpr = gate.cols() / GROUP_SIZE;
3218                let gsum = q1_group_sums(&act.xq, gpr);
3219                let gsum = &gsum;
3220                let run = move |start: usize, end: usize| {
3221                    for r in start..end {
3222                        let mut gv = dot_q1_row_i8(g_bytes, r, gpr, &act.xq, gsum) * act.sx;
3223                        let mut uv = dot_q1_row_i8(u_bytes, r, gpr, &act.xq, gsum) * act.sx;
3224                        for &(j, xv) in &act.outliers {
3225                            let (w, s) = q1_outlier(g_bytes, r, gpr, j);
3226                            gv += w * s * xv;
3227                            let (w, s) = q1_outlier(u_bytes, r, gpr, j);
3228                            uv += w * s * xv;
3229                        }
3230                        let silu_g = gv / (1.0 + (-gv).exp());
3231                        // SAFETY: disjoint row ranges per worker.
3232                        unsafe { *out_addr.at(r) = silu_g * uv };
3233                    }
3234                };
3235                dispatch_rows(pool, inter, &run);
3236                true
3237            }
3238            // Q2TiledP gate + Q2TiledP up — the 2-bit expert pair (MoE
3239            // FFNs of the W2 class): one row pass, both ladders built
3240            // once, integer code dots with shared group sums.
3241            (
3242                Self::Mapped {
3243                    dtype: TensorDtype::Q2TiledP,
3244                    ..
3245                },
3246                Self::Mapped {
3247                    dtype: TensorDtype::Q2TiledP,
3248                    ..
3249                },
3250            ) => {
3251                let cols = gate.cols();
3252                let gpr = cols / GROUP_SIZE;
3253                let gv_view = Q4tpView::new_q2(gate.quant_bytes(), inter, cols);
3254                let uv_view = Q4tpView::new_q2(up.quant_bytes(), inter, cols);
3255                let gsum = q1_group_sums(&act.xq, gpr);
3256                let gsum = &gsum;
3257                let run = move |start: usize, end: usize| {
3258                    let (mut gsc, mut usc) = (vec![0f32; gpr], vec![0f32; gpr]);
3259                    for r in start..end {
3260                        gv_view.scales_into(r, gpr, &mut gsc);
3261                        uv_view.scales_into(r, gpr, &mut usc);
3262                        let mut gv =
3263                            dot_q2tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsum, &gsc) * act.sx;
3264                        let mut uv =
3265                            dot_q2tp_row_i8(uv_view.nib, r, gpr, &act.xq, gsum, &usc) * act.sx;
3266                        for &(j, xv) in &act.outliers {
3267                            let (w, s) = q2tp_outlier(gv_view.nib, r, gpr, j, &gsc);
3268                            gv += w * s * xv;
3269                            let (w, s) = q2tp_outlier(uv_view.nib, r, gpr, j, &usc);
3270                            uv += w * s * xv;
3271                        }
3272                        let silu_g = gv / (1.0 + (-gv).exp());
3273                        // SAFETY: disjoint row ranges per worker.
3274                        unsafe { *out_addr.at(r) = silu_g * uv };
3275                    }
3276                };
3277                dispatch_rows(pool, inter, &run);
3278                true
3279            }
3280            // Q8Row gate + Q8Row up — one row pass over both i8 streams.
3281            // Q8_2f stays out on purpose: its column field prescales the
3282            // activations PER TENSOR, which breaks this fn's shared
3283            // split_act contract — it keeps the two-dispatch path.
3284            (
3285                Self::Mapped {
3286                    dtype: TensorDtype::Q8Row,
3287                    row_scale: g_rs,
3288                    ..
3289                },
3290                Self::Mapped {
3291                    dtype: TensorDtype::Q8Row,
3292                    row_scale: u_rs,
3293                    ..
3294                },
3295            ) => {
3296                let g_bytes = gate.quant_bytes();
3297                let u_bytes = up.quant_bytes();
3298                let cols = gate.cols();
3299                let run = move |start: usize, end: usize| {
3300                    for r in start..end {
3301                        let gv = q8_row_dot(&g_bytes[r * cols..(r + 1) * cols], act) * g_rs[r];
3302                        let uv = q8_row_dot(&u_bytes[r * cols..(r + 1) * cols], act) * u_rs[r];
3303                        let silu_g = gv / (1.0 + (-gv).exp());
3304                        // SAFETY: disjoint row ranges per worker.
3305                        unsafe { *out_addr.at(r) = silu_g * uv };
3306                    }
3307                };
3308                dispatch_rows(pool, inter, &run);
3309                true
3310            }
3311            // Q1T gate + Q1T up
3312            (
3313                Self::Mapped {
3314                    dtype: TensorDtype::Q1T,
3315                    ..
3316                },
3317                Self::Mapped {
3318                    dtype: TensorDtype::Q1T,
3319                    ..
3320                },
3321            ) => {
3322                const TILE: usize = cortiq_core::quant::Q1T_TILE;
3323                let g_bytes = gate.quant_bytes();
3324                let u_bytes = up.quant_bytes();
3325                let gpr = gate.cols() / GROUP_SIZE;
3326                let (g_rp, g_ent, g_ov) = q1t_overlay(g_bytes, inter * gpr * TILE, inter);
3327                let (u_rp, u_ent, u_ov) = q1t_overlay(u_bytes, inter * gpr * TILE, inter);
3328                let run = move |start: usize, end: usize| {
3329                    for r in start..end {
3330                        let mut gv = q1t_dot_row_i8(g_bytes, r, gpr, &act.xq) * act.sx;
3331                        let mut uv = q1t_dot_row_i8(u_bytes, r, gpr, &act.xq) * act.sx;
3332                        for &(j, xv) in &act.outliers {
3333                            gv += q1t_base_weight(g_bytes, r, gpr, j) * xv;
3334                            uv += q1t_base_weight(u_bytes, r, gpr, j) * xv;
3335                        }
3336                        gv += q1t_row_outlier_correction(g_bytes, r, g_rp, g_ent, g_ov, x_ref);
3337                        uv += q1t_row_outlier_correction(u_bytes, r, u_rp, u_ent, u_ov, x_ref);
3338                        let silu_g = gv / (1.0 + (-gv).exp());
3339                        // SAFETY: disjoint row ranges per worker.
3340                        unsafe { *out_addr.at(r) = silu_g * uv };
3341                    }
3342                };
3343                dispatch_rows(pool, inter, &run);
3344                true
3345            }
3346            _ => false,
3347        }
3348    }
3349
3350    /// Every routed expert's fused gate/up/SiLU under ONE pool dispatch.
3351    ///
3352    /// The per-expert path pays a pool barrier per expert per stage: at 9
3353    /// experts over 40 layers that is ~720 barriers a token, and a decode
3354    /// profile of Qwen3.6-35B-A3B showed the pool parked in
3355    /// `psynch_cvwait` about twice as long as it spent computing. Laying
3356    /// every expert's rows end-to-end in one virtual row space collapses
3357    /// the stage to a single dispatch. The per-row body is the
3358    /// single-expert q4tp arm verbatim, so outputs are bit-identical.
3359    ///
3360    /// `false` = something is outside the fused q4tp kernel (dtype, shape,
3361    /// or the `CMF_SDOT=0` exact contract); the caller walks the ordinary
3362    /// per-expert path.
3363    pub fn moe_gate_up_many(
3364        pairs: &[(&QTensor, &QTensor)],
3365        x: &[f32],
3366        outs: &mut [Vec<f32>],
3367        pool: Option<&Pool>,
3368    ) -> bool {
3369        if pairs.is_empty() || pairs.len() != outs.len() || !a8w8_enabled() {
3370            return false;
3371        }
3372        let inter = pairs[0].0.rows();
3373        let cols = pairs[0].0.cols();
3374        if cols % GROUP_SIZE != 0 {
3375            return false;
3376        }
3377        let gpr = cols / GROUP_SIZE;
3378        // Uniform layout across every routed pair: q4tp, or the 2-bit
3379        // profile's q2tp gate/up (the W2 class). Mixed sets refuse.
3380        let q2 = matches!(
3381            pairs[0].0,
3382            Self::Mapped {
3383                dtype: TensorDtype::Q2TiledP,
3384                ..
3385            }
3386        );
3387        let want = if q2 {
3388            TensorDtype::Q2TiledP
3389        } else {
3390            TensorDtype::Q4TiledP
3391        };
3392        let mut views = Vec::with_capacity(pairs.len() * 2);
3393        for ((g, u), o) in pairs.iter().zip(outs.iter()) {
3394            let both = matches!(g, Self::Mapped { dtype, .. } if *dtype == want)
3395                && matches!(u, Self::Mapped { dtype, .. } if *dtype == want);
3396            if !both
3397                || g.rows() != inter
3398                || u.rows() != inter
3399                || g.cols() != cols
3400                || u.cols() != cols
3401                || o.len() != inter
3402            {
3403                return false;
3404            }
3405            let mk = if q2 { Q4tpView::new_q2 } else { Q4tpView::new };
3406            views.push(mk(g.quant_bytes(), inter, cols));
3407            views.push(mk(u.quant_bytes(), inter, cols));
3408        }
3409        let act = split_act(x);
3410        let gsum = if q2 {
3411            q1_group_sums(&act.xq, gpr)
3412        } else {
3413            Vec::new()
3414        };
3415        let (act, gsum) = (&act, &gsum);
3416        let ptrs: Vec<SendMut> = outs.iter_mut().map(|o| SendMut(o.as_mut_ptr())).collect();
3417        let (views, ptrs) = (&views, &ptrs);
3418        let run = |start: usize, end: usize| {
3419            let (mut gsc, mut usc) = (vec![0f32; gpr], vec![0f32; gpr]);
3420            for flat in start..end {
3421                let (e, r) = (flat / inter, flat % inter);
3422                let gv_view = &views[e * 2];
3423                let uv_view = &views[e * 2 + 1];
3424                gv_view.scales_into(r, gpr, &mut gsc);
3425                uv_view.scales_into(r, gpr, &mut usc);
3426                let (mut gv, mut uv) = if q2 {
3427                    (
3428                        dot_q2tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsum, &gsc) * act.sx,
3429                        dot_q2tp_row_i8(uv_view.nib, r, gpr, &act.xq, gsum, &usc) * act.sx,
3430                    )
3431                } else {
3432                    (
3433                        dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, &gsc) * act.sx,
3434                        dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, &usc) * act.sx,
3435                    )
3436                };
3437                for &(j, xv) in &act.outliers {
3438                    let (og, ou) = if q2 {
3439                        (
3440                            q2tp_outlier(gv_view.nib, r, gpr, j, &gsc),
3441                            q2tp_outlier(uv_view.nib, r, gpr, j, &usc),
3442                        )
3443                    } else {
3444                        (
3445                            q4tp_outlier(gv_view.nib, r, gpr, j, &gsc),
3446                            q4tp_outlier(uv_view.nib, r, gpr, j, &usc),
3447                        )
3448                    };
3449                    gv += og.0 * og.1 * xv;
3450                    uv += ou.0 * ou.1 * xv;
3451                }
3452                let silu_g = gv / (1.0 + (-gv).exp());
3453                // SAFETY: one worker owns each (expert, row) pair.
3454                unsafe { *ptrs[e].at(r) = silu_g * uv };
3455            }
3456        };
3457        dispatch_rows(pool, pairs.len() * inter, &run);
3458        true
3459    }
3460
3461    /// Every routed expert's down projection, weighted and summed into
3462    /// `out`, under ONE pool dispatch.
3463    ///
3464    /// Partitioned by OUTPUT row rather than by expert: each row is owned
3465    /// by a single worker, so the experts are summed in the caller's order
3466    /// — the same sequence of f32 adds the serial `out[i] += w·eo[i]` loop
3467    /// performs, hence bit-identical. Partitioning by expert instead would
3468    /// race on the shared accumulator.
3469    pub fn moe_down_many(
3470        downs: &[&QTensor],
3471        gs: &[Vec<f32>],
3472        weights: &[f32],
3473        out: &mut [f32],
3474        pool: Option<&Pool>,
3475    ) -> bool {
3476        if downs.is_empty()
3477            || downs.len() != gs.len()
3478            || downs.len() != weights.len()
3479            || !a8w8_enabled()
3480        {
3481            return false;
3482        }
3483        let rows = out.len();
3484        let cols = downs[0].cols();
3485        if cols % GROUP_SIZE != 0 {
3486            return false;
3487        }
3488        let gpr = cols / GROUP_SIZE;
3489        let mut views = Vec::with_capacity(downs.len());
3490        for (d, g) in downs.iter().zip(gs.iter()) {
3491            if !matches!(
3492                d,
3493                Self::Mapped {
3494                    dtype: TensorDtype::Q4TiledP,
3495                    ..
3496                }
3497            ) || d.rows() != rows
3498                || d.cols() != cols
3499                || g.len() != cols
3500            {
3501                return false;
3502            }
3503            views.push(Q4tpView::new(d.quant_bytes(), rows, cols));
3504        }
3505        // One int8 split per expert — the activation vectors differ.
3506        let acts: Vec<SplitAct> = gs.iter().map(|g| split_act(g)).collect();
3507        // Partitioned by OUTPUT row, with the experts folded inside: each
3508        // row is owned by one worker, so they are summed in the caller's
3509        // order — the same f32 sequence the serial `out[i] += w·eo[i]`
3510        // loop produces. Partitioning by expert instead would either race
3511        // on the accumulator or need a scratch plane and a second pass;
3512        // measured, that variant was a wash, so this keeps the simpler
3513        // shape.
3514        let out_addr = SendMut(out.as_mut_ptr());
3515        let (views, acts, weights) = (&views, &acts, &weights);
3516        let run = |start: usize, end: usize| {
3517            let mut sc = vec![0f32; gpr];
3518            for r in start..end {
3519                let mut acc = 0f32;
3520                for (e, v) in views.iter().enumerate() {
3521                    v.scales_into(r, gpr, &mut sc);
3522                    let a = &acts[e];
3523                    let mut d = dot_q4tp_row_i8(v.nib, r, gpr, &a.xq, &sc) * a.sx;
3524                    for &(j, xv) in &a.outliers {
3525                        let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
3526                        d += w * s * xv;
3527                    }
3528                    acc += weights[e] * d;
3529                }
3530                // SAFETY: disjoint row ranges per worker.
3531                unsafe { *out_addr.at(r) = acc };
3532            }
3533        };
3534        dispatch_rows(pool, rows, &run);
3535        true
3536    }
3537}
3538
3539/// Batched q8 kernel: same math as qmatvec, the row makes a single
3540/// pass from memory for the whole batch.
3541/// Accelerate CBLAS — the Apple AMX matrix units, the same engine
3542/// llama.cpp's `-ngl 0` prefill rides via ggml-blas.
3543#[cfg(target_os = "macos")]
3544mod accel_blas {
3545    #[link(name = "Accelerate", kind = "framework")]
3546    unsafe extern "C" {
3547        pub fn cblas_sgemm(
3548            order: i32,
3549            trans_a: i32,
3550            trans_b: i32,
3551            m: i32,
3552            n: i32,
3553            k: i32,
3554            alpha: f32,
3555            a: *const f32,
3556            lda: i32,
3557            b: *const f32,
3558            ldb: i32,
3559            beta: f32,
3560            c: *mut f32,
3561            ldc: i32,
3562        );
3563    }
3564}
3565
3566#[cfg(target_os = "macos")]
3567pub(crate) fn accel_gemm_enabled() -> bool {
3568    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
3569    *ON.get_or_init(|| std::env::var("CMF_ACCEL").map(|v| v != "0").unwrap_or(true))
3570}
3571
3572/// Off macOS the "accel" GEMM is the portable NEON micro-kernel below —
3573/// same entry point, so the batched-attention path opens on mobile.
3574#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
3575pub(crate) fn accel_gemm_enabled() -> bool {
3576    true
3577}
3578
3579/// Portable NEON f32 GEMM (row-major, optional Bᵀ): a 4×8 fmla
3580/// micro-kernel with A broadcast against B panels — the mobile stand-in
3581/// for Accelerate in the batched causal attention (QKᵀ and P·V). Not a
3582/// BLAS: shapes here are the attention panels (m ≤ heads·chunk,
3583/// k = head_dim or context), and the goal is removing the per-position
3584/// quadratic wall, not peak GEMM.
3585#[cfg(target_arch = "aarch64")]
3586#[allow(clippy::too_many_arguments)]
3587pub(crate) fn neon_gemm_rm(
3588    m: usize,
3589    n: usize,
3590    k: usize,
3591    alpha: f32,
3592    a: &[f32],
3593    lda: usize,
3594    b_mat: &[f32],
3595    ldb: usize,
3596    b_rows_are_n: bool,
3597    c: &mut [f32],
3598    ldc: usize,
3599) {
3600    debug_assert!(a.len() >= (m - 1) * lda + k);
3601    debug_assert!(c.len() >= (m - 1) * ldc + n);
3602    // SAFETY: bounds asserted above; NEON is baseline on aarch64.
3603    unsafe {
3604        use core::arch::aarch64::*;
3605        let mut i = 0usize;
3606        while i < m {
3607            let mi = (m - i).min(4);
3608            let mut j = 0usize;
3609            while j < n {
3610                let nj = (n - j).min(8);
3611                if mi == 4 && nj == 8 {
3612                    let (mut c0a, mut c0b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
3613                    let (mut c1a, mut c1b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
3614                    let (mut c2a, mut c2b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
3615                    let (mut c3a, mut c3b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
3616                    for p in 0..k {
3617                        let (b0, b1) = if b_rows_are_n {
3618                            // B is [n, k]: column p of Bᵀ = element p of
3619                            // eight consecutive B rows — gathered.
3620                            let base = b_mat.as_ptr().add(j * ldb + p);
3621                            let g = |o: usize| *base.add(o * ldb);
3622                            ([g(0), g(1), g(2), g(3)], [g(4), g(5), g(6), g(7)])
3623                        } else {
3624                            let base = b_mat.as_ptr().add(p * ldb + j);
3625                            (
3626                                [*base, *base.add(1), *base.add(2), *base.add(3)],
3627                                [*base.add(4), *base.add(5), *base.add(6), *base.add(7)],
3628                            )
3629                        };
3630                        let bv0 = vld1q_f32(b0.as_ptr());
3631                        let bv1 = vld1q_f32(b1.as_ptr());
3632                        let a0 = vdupq_n_f32(*a.as_ptr().add(i * lda + p));
3633                        let a1 = vdupq_n_f32(*a.as_ptr().add((i + 1) * lda + p));
3634                        let a2 = vdupq_n_f32(*a.as_ptr().add((i + 2) * lda + p));
3635                        let a3 = vdupq_n_f32(*a.as_ptr().add((i + 3) * lda + p));
3636                        c0a = vfmaq_f32(c0a, a0, bv0);
3637                        c0b = vfmaq_f32(c0b, a0, bv1);
3638                        c1a = vfmaq_f32(c1a, a1, bv0);
3639                        c1b = vfmaq_f32(c1b, a1, bv1);
3640                        c2a = vfmaq_f32(c2a, a2, bv0);
3641                        c2b = vfmaq_f32(c2b, a2, bv1);
3642                        c3a = vfmaq_f32(c3a, a3, bv0);
3643                        c3b = vfmaq_f32(c3b, a3, bv1);
3644                    }
3645                    let al = vdupq_n_f32(alpha);
3646                    for (r, (ca, cb)) in [(c0a, c0b), (c1a, c1b), (c2a, c2b), (c3a, c3b)]
3647                        .iter()
3648                        .enumerate()
3649                    {
3650                        let dst = c.as_mut_ptr().add((i + r) * ldc + j);
3651                        vst1q_f32(dst, vmulq_f32(*ca, al));
3652                        vst1q_f32(dst.add(4), vmulq_f32(*cb, al));
3653                    }
3654                } else {
3655                    for r in 0..mi {
3656                        for q in 0..nj {
3657                            let mut acc = 0f32;
3658                            for p in 0..k {
3659                                let bv = if b_rows_are_n {
3660                                    b_mat[(j + q) * ldb + p]
3661                                } else {
3662                                    b_mat[p * ldb + j + q]
3663                                };
3664                                acc += a[(i + r) * lda + p] * bv;
3665                            }
3666                            c[(i + r) * ldc + j + q] = acc * alpha;
3667                        }
3668                    }
3669                }
3670                j += nj;
3671            }
3672            i += mi;
3673        }
3674    }
3675}
3676
3677/// Off-macOS aarch64: the batched attention rides the NEON micro-GEMM.
3678#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
3679#[allow(clippy::too_many_arguments)]
3680pub(crate) fn sgemm_rm(
3681    m: usize,
3682    n: usize,
3683    k: usize,
3684    alpha: f32,
3685    a: &[f32],
3686    lda: usize,
3687    b_mat: &[f32],
3688    ldb: usize,
3689    b_rows_are_n: bool,
3690    c: &mut [f32],
3691    ldc: usize,
3692) {
3693    neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
3694}
3695
3696/// Row-major f32 GEMM, exposed for offline tools (the AWNP pass builds a
3697/// per-layer projection and applies it to every expert; a naive triple loop
3698/// would turn a two-minute job into half an hour).
3699#[allow(clippy::too_many_arguments)]
3700pub fn sgemm_public(
3701    m: usize,
3702    n: usize,
3703    k: usize,
3704    alpha: f32,
3705    a: &[f32],
3706    lda: usize,
3707    b_mat: &[f32],
3708    ldb: usize,
3709    b_rows_are_n: bool,
3710    c: &mut [f32],
3711    ldc: usize,
3712) {
3713    #[cfg(any(target_os = "macos", target_arch = "aarch64"))]
3714    {
3715        sgemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
3716    }
3717    // x86 without Accelerate has no sgemm_rm: the specialized paths there are
3718    // quantized kernels, not an f32 GEMM. Only the offline AWNP pass reaches
3719    // this, so correctness matters and throughput does not — a triple loop is
3720    // the honest fallback rather than a reason to make the tool macOS-only.
3721    #[cfg(not(any(target_os = "macos", target_arch = "aarch64")))]
3722    {
3723        for i in 0..m {
3724            for j in 0..n {
3725                let mut acc = 0f32;
3726                for p in 0..k {
3727                    let bv = if b_rows_are_n {
3728                        b_mat[j * ldb + p]
3729                    } else {
3730                        b_mat[p * ldb + j]
3731                    };
3732                    acc += a[i * lda + p] * bv;
3733                }
3734                c[i * ldc + j] = alpha * acc;
3735            }
3736        }
3737    }
3738}
3739
3740/// Row-major f32 GEMM on Accelerate: C[m,n] = alpha·A[m,k] × B(ᵀ).
3741/// `b_rows_are_n` = true multiplies by Bᵀ where B is stored [n, k].
3742#[cfg(target_os = "macos")]
3743#[allow(clippy::too_many_arguments)]
3744pub(crate) fn sgemm_rm(
3745    m: usize,
3746    n: usize,
3747    k: usize,
3748    alpha: f32,
3749    a: &[f32],
3750    lda: usize,
3751    b_mat: &[f32],
3752    ldb: usize,
3753    b_rows_are_n: bool,
3754    c: &mut [f32],
3755    ldc: usize,
3756) {
3757    debug_assert!(a.len() >= (m - 1) * lda + k);
3758    debug_assert!(c.len() >= (m - 1) * ldc + n);
3759    // Test hook: route the attention GEMMs through the portable NEON
3760    // micro-kernel ON APPLE SILICON — how the mobile batched attend is
3761    // measured without a phone in the loop. (Intel macOS has no NEON —
3762    // the hook is a no-op there, Accelerate continues below.)
3763    #[cfg(target_arch = "aarch64")]
3764    if std::env::var("CMF_FORCE_NEON_GEMM")
3765        .map(|v| v == "1")
3766        .unwrap_or(false)
3767    {
3768        return neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
3769    }
3770    unsafe {
3771        accel_blas::cblas_sgemm(
3772            101, // RowMajor
3773            111, // NoTrans A
3774            if b_rows_are_n { 112 } else { 111 },
3775            m as i32,
3776            n as i32,
3777            k as i32,
3778            alpha,
3779            a.as_ptr(),
3780            lda as i32,
3781            b_mat.as_ptr(),
3782            ldb as i32,
3783            0.0,
3784            c.as_mut_ptr(),
3785            ldc as i32,
3786        );
3787    }
3788}
3789
3790/// Prefill GEMM through Accelerate (macOS): dequantize q8 rows into
3791/// f32 tiles (scale folded in, pool-parallel) and multiply each tile
3792/// on the AMX with one row-major sgemm. Tiles live in cache, weights
3793/// stream once. Numerics are f32-GEMM (not the int8 dot): prefill
3794/// logits shift within f32 rounding — tolerance-class, like every
3795/// reduction-order change; decode (M=1) never takes this path.
3796#[cfg(target_os = "macos")]
3797fn qmatmat_accel(
3798    q: &[u8],
3799    row_scale: &[f32],
3800    pre: &[std::borrow::Cow<'_, [f32]>],
3801    rows: usize,
3802    cols: usize,
3803    out: &mut [f32],
3804    pool: Option<&Pool>,
3805) {
3806    // NOTE: double-buffering the dequant against the sgemm (a scoped
3807    // thread driving the pool on tile k+1 while the caller multiplies
3808    // tile k) was tried and LOST ~6%: Accelerate's sgemm is itself
3809    // multithreaded, and the dequant workers just steal its cores.
3810    const TR: usize = 2048;
3811    let b = pre.len();
3812    thread_local! {
3813        static XPANEL: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
3814        static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
3815    }
3816    XPANEL.with(|xp| {
3817        WTILE.with(|wt| {
3818            let mut xpanel = xp.borrow_mut();
3819            xpanel.clear();
3820            for x in pre {
3821                xpanel.extend_from_slice(x);
3822            }
3823            let mut wtile = wt.borrow_mut();
3824            wtile.resize(TR * cols, 0.0);
3825            let mut r0 = 0usize;
3826            while r0 < rows {
3827                let tr = TR.min(rows - r0);
3828                // Dequant the tile (scale folded) — pool-parallel.
3829                let wt_addr = SendMut(wtile.as_mut_ptr());
3830                let run = |start: usize, end: usize| {
3831                    for r in start..end {
3832                        let row = &q[(r0 + r) * cols..(r0 + r + 1) * cols];
3833                        let s = row_scale[r0 + r];
3834                        // SAFETY: workers cover disjoint r ranges.
3835                        let dst =
3836                            unsafe { std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols) };
3837                        for (d, &v) in dst.iter_mut().zip(row) {
3838                            *d = (v as i8) as f32 * s;
3839                        }
3840                    }
3841                };
3842                dispatch_rows(pool, tr, &run);
3843                // C[b, tr] (at column r0 of out[b, rows]) = X · Wtileᵀ
3844                unsafe {
3845                    accel_blas::cblas_sgemm(
3846                        101, // RowMajor
3847                        111, // NoTrans A
3848                        112, // Trans B
3849                        b as i32,
3850                        tr as i32,
3851                        cols as i32,
3852                        1.0,
3853                        xpanel.as_ptr(),
3854                        cols as i32,
3855                        wtile.as_ptr(),
3856                        cols as i32,
3857                        0.0,
3858                        out.as_mut_ptr().add(r0),
3859                        rows as i32,
3860                    );
3861                }
3862                r0 += tr;
3863            }
3864        })
3865    });
3866}
3867
3868fn qmatmat(
3869    q: &[u8],
3870    row_scale: &[f32],
3871    pre: &[std::borrow::Cow<'_, [f32]>],
3872    rows: usize,
3873    cols: usize,
3874    out: &mut [f32],
3875    pool: Option<&Pool>,
3876) {
3877    let b = pre.len();
3878    debug_assert_eq!(out.len(), b * rows);
3879    // Big prefill batches ride the AMX (roadmap PR3): the row×batch
3880    // SDOT loop below peaks near the CPU's dot throughput, an order
3881    // below the matrix units. Small tensors and tiny test models stay
3882    // on the exact integer path.
3883    #[cfg(target_os = "macos")]
3884    if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
3885        qmatmat_accel(q, row_scale, pre, rows, cols, out, pool);
3886        return;
3887    }
3888    #[cfg(target_arch = "aarch64")]
3889    if sdot_enabled() {
3890        let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
3891        let out_addr = SendMut(out.as_mut_ptr());
3892        // Blocked 2×4 (mobile prefill: no AMX to fall back on — this
3893        // path IS the ARM prefill GEMM off Apple silicon).
3894        let blocked_ok = blocked_enabled();
3895        let use_i8mm = i8mm_enabled();
3896        if blocked_ok {
3897            let run = |start: usize, end: usize| {
3898                let mut o = start;
3899                while o < end {
3900                    if o + 2 <= end {
3901                        let r0 = &q[o * cols..(o + 1) * cols];
3902                        let r1 = &q[(o + 1) * cols..(o + 2) * cols];
3903                        let mut bi = 0usize;
3904                        while bi + 4 <= acts.len() {
3905                            let xs = [
3906                                acts[bi].xq.as_slice(),
3907                                acts[bi + 1].xq.as_slice(),
3908                                acts[bi + 2].xq.as_slice(),
3909                                acts[bi + 3].xq.as_slice(),
3910                            ];
3911                            let d = if use_i8mm {
3912                                unsafe { dot_i8_smmla_2x4(r0, r1, xs) }
3913                            } else {
3914                                unsafe { dot_i8_sdot_2x4(r0, r1, xs) }
3915                            };
3916                            for (r, row) in [r0, r1].into_iter().enumerate() {
3917                                for k in 0..4 {
3918                                    let act = &acts[bi + k];
3919                                    let mut v = d[r][k] as f32 * act.sx;
3920                                    for &(j, xv) in &act.outliers {
3921                                        v += (row[j] as i8) as f32 * xv;
3922                                    }
3923                                    unsafe {
3924                                        *out_addr.at((bi + k) * rows + o + r) = v * row_scale[o + r]
3925                                    };
3926                                }
3927                            }
3928                            bi += 4;
3929                        }
3930                        while bi < acts.len() {
3931                            for (r, row) in [r0, r1].into_iter().enumerate() {
3932                                let v = row_dot_sdot(row, &acts[bi]) * row_scale[o + r];
3933                                unsafe { *out_addr.at(bi * rows + o + r) = v };
3934                            }
3935                            bi += 1;
3936                        }
3937                        o += 2;
3938                    } else {
3939                        let row = &q[o * cols..(o + 1) * cols];
3940                        for (bi, act) in acts.iter().enumerate() {
3941                            let v = row_dot_sdot(row, act) * row_scale[o];
3942                            unsafe { *out_addr.at(bi * rows + o) = v };
3943                        }
3944                        o += 1;
3945                    }
3946                }
3947            };
3948            dispatch_rows(pool, rows, &run);
3949            return;
3950        }
3951        let run = |start: usize, end: usize| {
3952            for o in start..end {
3953                let row = &q[o * cols..(o + 1) * cols];
3954                for (bi, act) in acts.iter().enumerate() {
3955                    let v = row_dot_sdot(row, act) * row_scale[o];
3956                    unsafe { *out_addr.at(bi * rows + o) = v };
3957                }
3958            }
3959        };
3960        dispatch_rows(pool, rows, &run);
3961        return;
3962    }
3963    // x86 A8W8 batch. Non-VNNI parts take the BLOCKED 2×4 kernel
3964    // (roadmap P0: two weight rows' abs() stay in registers across four
3965    // activation streams); VNNI machines keep the per-row bias-trick
3966    // dot, which is already throughput-bound there.
3967    #[cfg(target_arch = "x86_64")]
3968    if avx2_a8w8_enabled() {
3969        let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
3970        let out_addr = SendMut(out.as_mut_ptr());
3971        // CMF_X86_BLOCKED=0 forces the per-row path (paired in-process
3972        // A/B on noisy shared-vCPU hosts).
3973        let blocked_ok = blocked_enabled();
3974        if !avx512vnni_enabled() && blocked_ok {
3975            let run = |start: usize, end: usize| {
3976                let mut o = start;
3977                while o < end {
3978                    if o + 2 <= end {
3979                        let r0 = &q[o * cols..(o + 1) * cols];
3980                        let r1 = &q[(o + 1) * cols..(o + 2) * cols];
3981                        let mut bi = 0usize;
3982                        while bi + 4 <= acts.len() {
3983                            let xs = [
3984                                acts[bi].xq.as_slice(),
3985                                acts[bi + 1].xq.as_slice(),
3986                                acts[bi + 2].xq.as_slice(),
3987                                acts[bi + 3].xq.as_slice(),
3988                            ];
3989                            let d = unsafe { dot_i8_i8_avx2_2x4(r0, r1, xs) };
3990                            for (r, row) in [r0, r1].into_iter().enumerate() {
3991                                for k in 0..4 {
3992                                    let act = &acts[bi + k];
3993                                    let mut v = d[r][k] as f32 * act.sx;
3994                                    for &(j, xv) in &act.outliers {
3995                                        v += (row[j] as i8) as f32 * xv;
3996                                    }
3997                                    unsafe {
3998                                        *out_addr.at((bi + k) * rows + o + r) = v * row_scale[o + r]
3999                                    };
4000                                }
4001                            }
4002                            bi += 4;
4003                        }
4004                        while bi < acts.len() {
4005                            for (r, row) in [r0, r1].into_iter().enumerate() {
4006                                let v = row_dot_avx2(row, &acts[bi]) * row_scale[o + r];
4007                                unsafe { *out_addr.at(bi * rows + o + r) = v };
4008                            }
4009                            bi += 1;
4010                        }
4011                        o += 2;
4012                    } else {
4013                        let row = &q[o * cols..(o + 1) * cols];
4014                        for (bi, act) in acts.iter().enumerate() {
4015                            let v = row_dot_avx2(row, act) * row_scale[o];
4016                            unsafe { *out_addr.at(bi * rows + o) = v };
4017                        }
4018                        o += 1;
4019                    }
4020                }
4021            };
4022            dispatch_rows(pool, rows, &run);
4023            return;
4024        }
4025        let run = |start: usize, end: usize| {
4026            for o in start..end {
4027                let row = &q[o * cols..(o + 1) * cols];
4028                for (bi, act) in acts.iter().enumerate() {
4029                    let v = row_dot_avx2(row, act) * row_scale[o];
4030                    unsafe { *out_addr.at(bi * rows + o) = v };
4031                }
4032            }
4033        };
4034        dispatch_rows(pool, rows, &run);
4035        return;
4036    }
4037    let out_addr = SendMut(out.as_mut_ptr());
4038    let run = |start: usize, end: usize| {
4039        for o in start..end {
4040            let row = &q[o * cols..(o + 1) * cols];
4041            for (bi, x) in pre.iter().enumerate() {
4042                let mut acc = 0f32;
4043                for j in 0..cols {
4044                    acc += (row[j] as i8) as f32 * x[j];
4045                }
4046                unsafe { *out_addr.at(bi * rows + o) = acc * row_scale[o] };
4047            }
4048        }
4049    };
4050    dispatch_rows(pool, rows, &run);
4051}
4052
4053/// Split rows across pool workers (shared qmatvec pattern). Self-balancing
4054/// — see `Pool::run_rows` for why a static 1/n split is wrong here.
4055fn dispatch_rows(pool: Option<&Pool>, rows: usize, run: &(dyn Fn(usize, usize) + Sync)) {
4056    match pool {
4057        Some(pool) if rows >= 256 => pool.run_rows(rows, run),
4058        _ => run(0, rows),
4059    }
4060}
4061
4062/// Split a q4_block blob into (packed nibbles, f16 group scales).
4063fn q4_split(bytes: &[u8], rows: usize, cols: usize) -> (&[u8], &[u8]) {
4064    let groups = rows * cols / GROUP_SIZE;
4065    bytes.split_at(groups * 16)
4066}
4067
4068/// SIMD unpack for the dominant vbit width B=4 (94% of rows on the
4069/// log2-shape calibration): 16 packed bytes -> 32 centered i8 values.
4070/// vbit packs MSB-first, so the HIGH nibble is the even element
4071/// (opposite of q4_block's lo-first interleave). Centering is u-7.
4072#[inline]
4073fn vbit_fill4(data: &[u8], buf: &mut [u8]) {
4074    #[cfg(target_arch = "aarch64")]
4075    unsafe {
4076        return vbit_fill4_neon(data, buf);
4077    }
4078    #[cfg(target_arch = "x86_64")]
4079    if avx2_enabled() {
4080        return unsafe { vbit_fill4_avx2(data, buf) };
4081    }
4082    #[allow(unreachable_code)]
4083    for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
4084        let u = unpack8::<4>(&data[blk * 4..]);
4085        for k in 0..8 {
4086            chunk[k] = (u[k] - 7) as i8 as u8;
4087        }
4088    }
4089}
4090
4091#[cfg(target_arch = "aarch64")]
4092#[target_feature(enable = "neon")]
4093unsafe fn vbit_fill4_neon(data: &[u8], buf: &mut [u8]) {
4094    // SAFETY: buf.len() is a multiple of GROUP_SIZE=32; data holds
4095    // buf.len()/2 packed bytes (validated at load).
4096    unsafe {
4097        use core::arch::aarch64::*;
4098        let n = buf.len();
4099        let mask = vdupq_n_u8(0x0F);
4100        let seven = vdupq_n_s8(7);
4101        let mut g = 0usize;
4102        while g * 32 + 32 <= n {
4103            let b = vld1q_u8(data.as_ptr().add(g * 16));
4104            let hi = vshrq_n_u8::<4>(b);
4105            let lo = vandq_u8(b, mask);
4106            let z0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(hi, lo)), seven);
4107            let z1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(hi, lo)), seven);
4108            vst1q_u8(buf.as_mut_ptr().add(g * 32), vreinterpretq_u8_s8(z0));
4109            vst1q_u8(buf.as_mut_ptr().add(g * 32 + 16), vreinterpretq_u8_s8(z1));
4110            g += 1;
4111        }
4112    }
4113}
4114
4115#[cfg(target_arch = "x86_64")]
4116#[target_feature(enable = "avx2")]
4117unsafe fn vbit_fill4_avx2(data: &[u8], buf: &mut [u8]) {
4118    // SAFETY: see vbit_fill4_neon.
4119    unsafe {
4120        use core::arch::x86_64::*;
4121        let n = buf.len();
4122        let mask = _mm_set1_epi8(0x0F);
4123        let seven = _mm256_set1_epi8(7);
4124        let mut g = 0usize;
4125        while g * 32 + 32 <= n {
4126            let b = _mm_loadu_si128(data.as_ptr().add(g * 16) as *const __m128i);
4127            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), mask);
4128            let lo = _mm_and_si128(b, mask);
4129            let z = _mm256_sub_epi8(
4130                _mm256_set_m128i(_mm_unpackhi_epi8(hi, lo), _mm_unpacklo_epi8(hi, lo)),
4131                seven,
4132            );
4133            _mm256_storeu_si256(buf.as_mut_ptr().add(g * 32) as *mut __m256i, z);
4134            g += 1;
4135        }
4136    }
4137}
4138
4139/// Unpack 8 MSB-first B-bit values from exactly B bytes (fixed shifts —
4140/// no serial bit-buffer, auto-vectorizable). Every 32-value group starts
4141/// byte-aligned (32·B/8 is integral for B∈3..8), so groups decompose
4142/// into 4 such blocks.
4143#[inline(always)]
4144fn unpack8<const B: usize>(data: &[u8]) -> [i32; 8] {
4145    let mut acc = 0u64;
4146    for i in 0..B {
4147        acc = (acc << 8) | data[i] as u64;
4148    }
4149    let mask = (1u64 << B) - 1;
4150    let mut out = [0i32; 8];
4151    for (k, o) in out.iter_mut().enumerate() {
4152        *o = ((acc >> ((7 - k) * B)) & mask) as i32;
4153    }
4154    out
4155}
4156
4157/// Fused vbit matvec straight from the mapped bytes (spec §3, P13
4158/// FIG.3): [u8 bits: rows][f16 scales: rows·cols/32][bit-packed rows,
4159/// MSB-first, byte-padded]. Row data offsets are precomputed at load
4160/// (`vbit_row_offsets`) — the per-call prefix scan was O(rows) pure
4161/// overhead on every matvec.
4162#[allow(clippy::too_many_arguments)]
4163fn vbitmatvec(
4164    bytes: &[u8],
4165    offsets: &[usize],
4166    x: &[f32],
4167    rows: usize,
4168    cols: usize,
4169    out: &mut [f32],
4170    pool: Option<&Pool>,
4171) {
4172    debug_assert_eq!(out.len(), rows);
4173    debug_assert_eq!(offsets.len(), rows + 1);
4174
4175    // SDOT path: unpack the row to centered i8 once, then per-group
4176    // int8 dot against the quantized activations — same A8W8 contract
4177    // as q8 (bounded noise; CMF_SDOT=0 keeps the exact scalar path).
4178    if a8w8_enabled() {
4179        let act = split_act(x);
4180        let out_addr = SendMut(out.as_mut_ptr());
4181        let run = move |start: usize, end: usize| {
4182            vbit_range_a8w8(bytes, offsets, x, &act, rows, cols, out_addr, start, end)
4183        };
4184        dispatch_rows(pool, rows, &run);
4185        return;
4186    }
4187
4188    let out_addr = SendMut(out.as_mut_ptr());
4189    let run = move |start: usize, end: usize| {
4190        vbit_range_f32(bytes, offsets, x, rows, cols, out_addr, start, end)
4191    };
4192    dispatch_rows(pool, rows, &run);
4193}
4194
4195/// One vbit row range via the A8W8 int8 path — kernel body of
4196/// `vbitmatvec`, extracted so multi-matrix jobs can drive it for
4197/// several tensors in one dispatch (b=8 rows go exact f32).
4198#[allow(clippy::too_many_arguments)]
4199fn vbit_range_a8w8(
4200    bytes: &[u8],
4201    offsets: &[usize],
4202    x: &[f32],
4203    act: &SplitAct,
4204    rows: usize,
4205    cols: usize,
4206    out: SendMut,
4207    start: usize,
4208    end: usize,
4209) {
4210    let ng = cols / GROUP_SIZE;
4211    let bits = &bytes[..rows];
4212    let sc_off = rows;
4213    let row_dot = |r: usize| -> f32 {
4214        let b = bits[r] as usize;
4215        let l = (1i32 << (b - 1)) - 1;
4216        let mask = (1u64 << b) - 1;
4217        let data = &bytes[offsets[r]..offsets[r + 1]];
4218        if b == 8 {
4219            // u−L reaches 128 → does not fit i8; exact f32 path.
4220            let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
4221            let mut dot = 0f32;
4222            for g in 0..ng {
4223                let so = (r * ng + g) * 2;
4224                let sgf = f16_to_f32(u16::from_le_bytes([
4225                    bytes[sc_off + so],
4226                    bytes[sc_off + so + 1],
4227                ]));
4228                let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
4229                let mut gd = 0f32;
4230                for &xv in xg.iter() {
4231                    if nbits < 8 {
4232                        acc = (acc << 8) | data[idx] as u64;
4233                        idx += 1;
4234                        nbits += 8;
4235                    }
4236                    let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
4237                    nbits -= 8;
4238                    gd += (u - l) as f32 * xv;
4239                }
4240                dot += gd * sgf;
4241            }
4242            return dot;
4243        }
4244        // Per-worker scratch: this closure runs for every row of the
4245        // tensor (lm_head ≈ 150k rows/token) — a heap allocation per
4246        // row was measurable pure overhead.
4247        thread_local! {
4248            static VBIT_SCRATCH: std::cell::RefCell<Vec<u8>> =
4249                const { std::cell::RefCell::new(Vec::new()) };
4250        }
4251        #[inline(always)]
4252        fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
4253            for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
4254                let u = unpack8::<B>(&data[blk * B..]);
4255                for k in 0..8 {
4256                    chunk[k] = (u[k] - l) as i8 as u8;
4257                }
4258            }
4259        }
4260        let _ = mask;
4261        VBIT_SCRATCH.with(|scratch| {
4262            let mut buf = scratch.borrow_mut();
4263            buf.resize(cols, 0);
4264            match b {
4265                3 => fill::<3>(data, l, &mut buf),
4266                4 => vbit_fill4(data, &mut buf),
4267                5 => fill::<5>(data, l, &mut buf),
4268                6 => fill::<6>(data, l, &mut buf),
4269                _ => unreachable!(),
4270            }
4271            let mut dot = 0f32;
4272            for g in 0..ng {
4273                let so = (r * ng + g) * 2;
4274                let s = f16_to_f32(u16::from_le_bytes([
4275                    bytes[sc_off + so],
4276                    bytes[sc_off + so + 1],
4277                ]));
4278                let d = dot_i8_i8(
4279                    &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
4280                    &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
4281                ) as f32
4282                    * act.sx;
4283                dot += d * s;
4284            }
4285            for &(j, xv) in &act.outliers {
4286                let so = (r * ng + j / GROUP_SIZE) * 2;
4287                let s = f16_to_f32(u16::from_le_bytes([
4288                    bytes[sc_off + so],
4289                    bytes[sc_off + so + 1],
4290                ]));
4291                // xq is zeroed at outlier slots — add the exact term.
4292                dot += (buf[j] as i8) as f32 * s * xv;
4293            }
4294            dot
4295        })
4296    };
4297    for r in start..end {
4298        // SAFETY: disjoint row ranges per worker.
4299        unsafe { *out.at(r) = row_dot(r) };
4300    }
4301}
4302
4303/// Exact scalar vbit row range (same extraction, non-SDOT path).
4304#[allow(clippy::too_many_arguments)]
4305fn vbit_range_f32(
4306    bytes: &[u8],
4307    offsets: &[usize],
4308    x: &[f32],
4309    rows: usize,
4310    cols: usize,
4311    out: SendMut,
4312    start: usize,
4313    end: usize,
4314) {
4315    let ng = cols / GROUP_SIZE;
4316    let bits = &bytes[..rows];
4317    let sc_off = rows;
4318    // Per-bit-width specialized inner loops: the compiler unrolls the
4319    // constant shifts (the generic bit-buffer loop was branch-bound —
4320    // 5.6 vs 13.2 tok/s q4 on the 0.8B).
4321    #[inline(always)]
4322    fn dot_row<const B: usize>(
4323        data: &[u8],
4324        bytes: &[u8],
4325        sc_off: usize,
4326        r: usize,
4327        ng: usize,
4328        x: &[f32],
4329    ) -> f32 {
4330        let l = ((1i32 << (B - 1)) - 1) as f32;
4331        let gbytes = GROUP_SIZE * B / 8;
4332        let mut dot = 0f32;
4333        for g in 0..ng {
4334            let so = (r * ng + g) * 2;
4335            let s = f16_to_f32(u16::from_le_bytes([
4336                bytes[sc_off + so],
4337                bytes[sc_off + so + 1],
4338            ]));
4339            let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
4340            let gd0 = &data[g * gbytes..(g + 1) * gbytes];
4341            let mut gd = 0f32;
4342            for blk in 0..GROUP_SIZE / 8 {
4343                let u = unpack8::<B>(&gd0[blk * B..]);
4344                let xb = &xg[blk * 8..blk * 8 + 8];
4345                for k in 0..8 {
4346                    gd += (u[k] as f32 - l) * xb[k];
4347                }
4348            }
4349            dot += gd * s;
4350        }
4351        dot
4352    }
4353    for r in start..end {
4354        let data = &bytes[offsets[r]..offsets[r + 1]];
4355        let v = match bits[r] {
4356            3 => dot_row::<3>(data, bytes, sc_off, r, ng, x),
4357            4 => dot_row::<4>(data, bytes, sc_off, r, ng, x),
4358            5 => dot_row::<5>(data, bytes, sc_off, r, ng, x),
4359            6 => dot_row::<6>(data, bytes, sc_off, r, ng, x),
4360            8 => dot_row::<8>(data, bytes, sc_off, r, ng, x),
4361            b => unreachable!("vbit bit-width {b} (validated at load)"),
4362        };
4363        // SAFETY: disjoint row ranges per worker.
4364        unsafe { *out.at(r) = v };
4365    }
4366}
4367
4368/// Fused two-input vbit matvec: each row is unpacked from the mmap ONCE
4369/// and dotted against BOTH activations (MTP verify / pair prefill used
4370/// to run two full matvecs — double weight traffic and double unpack).
4371/// Per-input math is identical to `vbitmatvec` → same accuracy contract.
4372#[allow(clippy::too_many_arguments)]
4373fn vbitmatvec2(
4374    bytes: &[u8],
4375    offsets: &[usize],
4376    x1: &[f32],
4377    x2: &[f32],
4378    rows: usize,
4379    cols: usize,
4380    o1: &mut [f32],
4381    o2: &mut [f32],
4382    pool: Option<&Pool>,
4383) {
4384    debug_assert_eq!(o1.len(), rows);
4385    debug_assert_eq!(o2.len(), rows);
4386
4387    if a8w8_enabled() {
4388        let a1 = split_act(x1);
4389        let a2 = split_act(x2);
4390        let p1 = SendMut(o1.as_mut_ptr());
4391        let p2 = SendMut(o2.as_mut_ptr());
4392        let run = move |start: usize, end: usize| {
4393            vbit_range2_a8w8(
4394                bytes, offsets, x1, x2, &a1, &a2, rows, cols, p1, p2, start, end,
4395            )
4396        };
4397        dispatch_rows(pool, rows, &run);
4398        return;
4399    }
4400
4401    let p1 = SendMut(o1.as_mut_ptr());
4402    let p2 = SendMut(o2.as_mut_ptr());
4403    let run = move |start: usize, end: usize| {
4404        vbit_range2_f32(bytes, offsets, x1, x2, rows, cols, p1, p2, start, end)
4405    };
4406    dispatch_rows(pool, rows, &run);
4407}
4408
4409/// Two-input vbit row range via the A8W8 int8 path — kernel body of
4410/// `vbitmatvec2`, extracted for pair multi-matrix jobs (b=8 rows go
4411/// exact f32 for both lanes, bits streamed once).
4412#[allow(clippy::too_many_arguments)]
4413fn vbit_range2_a8w8(
4414    bytes: &[u8],
4415    offsets: &[usize],
4416    x1: &[f32],
4417    x2: &[f32],
4418    a1: &SplitAct,
4419    a2: &SplitAct,
4420    rows: usize,
4421    cols: usize,
4422    p1: SendMut,
4423    p2: SendMut,
4424    start: usize,
4425    end: usize,
4426) {
4427    let ng = cols / GROUP_SIZE;
4428    let bits = &bytes[..rows];
4429    let sc_off = rows;
4430    let row_dots = |r: usize| -> (f32, f32) {
4431        let b = bits[r] as usize;
4432        let l = (1i32 << (b - 1)) - 1;
4433        let data = &bytes[offsets[r]..offsets[r + 1]];
4434        if b == 8 {
4435            // u−L reaches 128 → does not fit i8; exact f32 path,
4436            // bits still streamed once for both lanes.
4437            let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
4438            let (mut d1, mut d2) = (0f32, 0f32);
4439            for g in 0..ng {
4440                let so = (r * ng + g) * 2;
4441                let sgf = f16_to_f32(u16::from_le_bytes([
4442                    bytes[sc_off + so],
4443                    bytes[sc_off + so + 1],
4444                ]));
4445                let (mut g1, mut g2) = (0f32, 0f32);
4446                for k in 0..GROUP_SIZE {
4447                    if nbits < 8 {
4448                        acc = (acc << 8) | data[idx] as u64;
4449                        idx += 1;
4450                        nbits += 8;
4451                    }
4452                    let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
4453                    nbits -= 8;
4454                    let w = (u - l) as f32;
4455                    g1 += w * x1[g * GROUP_SIZE + k];
4456                    g2 += w * x2[g * GROUP_SIZE + k];
4457                }
4458                d1 += g1 * sgf;
4459                d2 += g2 * sgf;
4460            }
4461            return (d1, d2);
4462        }
4463        thread_local! {
4464            static VBIT_SCRATCH2: std::cell::RefCell<Vec<u8>> =
4465                const { std::cell::RefCell::new(Vec::new()) };
4466        }
4467        #[inline(always)]
4468        fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
4469            for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
4470                let u = unpack8::<B>(&data[blk * B..]);
4471                for k in 0..8 {
4472                    chunk[k] = (u[k] - l) as i8 as u8;
4473                }
4474            }
4475        }
4476        VBIT_SCRATCH2.with(|scratch| {
4477            let mut buf = scratch.borrow_mut();
4478            buf.resize(cols, 0);
4479            match b {
4480                3 => fill::<3>(data, l, &mut buf),
4481                4 => vbit_fill4(data, &mut buf),
4482                5 => fill::<5>(data, l, &mut buf),
4483                6 => fill::<6>(data, l, &mut buf),
4484                _ => unreachable!(),
4485            }
4486            let (mut d1, mut d2) = (0f32, 0f32);
4487            for g in 0..ng {
4488                let so = (r * ng + g) * 2;
4489                let s = f16_to_f32(u16::from_le_bytes([
4490                    bytes[sc_off + so],
4491                    bytes[sc_off + so + 1],
4492                ]));
4493                let wg = &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
4494                let v1 = dot_i8_i8(wg, &a1.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a1.sx;
4495                let v2 = dot_i8_i8(wg, &a2.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a2.sx;
4496                d1 += v1 * s;
4497                d2 += v2 * s;
4498            }
4499            for &(j, xv) in &a1.outliers {
4500                let so = (r * ng + j / GROUP_SIZE) * 2;
4501                let s = f16_to_f32(u16::from_le_bytes([
4502                    bytes[sc_off + so],
4503                    bytes[sc_off + so + 1],
4504                ]));
4505                d1 += (buf[j] as i8) as f32 * s * xv;
4506            }
4507            for &(j, xv) in &a2.outliers {
4508                let so = (r * ng + j / GROUP_SIZE) * 2;
4509                let s = f16_to_f32(u16::from_le_bytes([
4510                    bytes[sc_off + so],
4511                    bytes[sc_off + so + 1],
4512                ]));
4513                d2 += (buf[j] as i8) as f32 * s * xv;
4514            }
4515            (d1, d2)
4516        })
4517    };
4518    for r in start..end {
4519        let (v1, v2) = row_dots(r);
4520        // SAFETY: disjoint row ranges per worker.
4521        unsafe {
4522            *p1.at(r) = v1;
4523            *p2.at(r) = v2;
4524        }
4525    }
4526}
4527
4528/// Two-input exact scalar vbit row range (same extraction) —
4529/// per-bit-width specialized, two accumulators per row; per-lane
4530/// accumulation order matches `vbitmatvec` exactly.
4531#[allow(clippy::too_many_arguments)]
4532fn vbit_range2_f32(
4533    bytes: &[u8],
4534    offsets: &[usize],
4535    x1: &[f32],
4536    x2: &[f32],
4537    rows: usize,
4538    cols: usize,
4539    p1: SendMut,
4540    p2: SendMut,
4541    start: usize,
4542    end: usize,
4543) {
4544    let ng = cols / GROUP_SIZE;
4545    let bits = &bytes[..rows];
4546    let sc_off = rows;
4547    #[inline(always)]
4548    #[allow(clippy::too_many_arguments)]
4549    fn dot_row2<const B: usize>(
4550        data: &[u8],
4551        bytes: &[u8],
4552        sc_off: usize,
4553        r: usize,
4554        ng: usize,
4555        x1: &[f32],
4556        x2: &[f32],
4557    ) -> (f32, f32) {
4558        let l = ((1i32 << (B - 1)) - 1) as f32;
4559        let gbytes = GROUP_SIZE * B / 8;
4560        let (mut d1, mut d2) = (0f32, 0f32);
4561        for g in 0..ng {
4562            let so = (r * ng + g) * 2;
4563            let s = f16_to_f32(u16::from_le_bytes([
4564                bytes[sc_off + so],
4565                bytes[sc_off + so + 1],
4566            ]));
4567            let x1g = &x1[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
4568            let x2g = &x2[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
4569            let gd0 = &data[g * gbytes..(g + 1) * gbytes];
4570            let (mut g1, mut g2) = (0f32, 0f32);
4571            for blk in 0..GROUP_SIZE / 8 {
4572                let u = unpack8::<B>(&gd0[blk * B..]);
4573                for k in 0..8 {
4574                    let w = u[k] as f32 - l;
4575                    g1 += w * x1g[blk * 8 + k];
4576                    g2 += w * x2g[blk * 8 + k];
4577                }
4578            }
4579            d1 += g1 * s;
4580            d2 += g2 * s;
4581        }
4582        (d1, d2)
4583    }
4584    for r in start..end {
4585        let data = &bytes[offsets[r]..offsets[r + 1]];
4586        let (v1, v2) = match bits[r] {
4587            3 => dot_row2::<3>(data, bytes, sc_off, r, ng, x1, x2),
4588            4 => dot_row2::<4>(data, bytes, sc_off, r, ng, x1, x2),
4589            5 => dot_row2::<5>(data, bytes, sc_off, r, ng, x1, x2),
4590            6 => dot_row2::<6>(data, bytes, sc_off, r, ng, x1, x2),
4591            8 => dot_row2::<8>(data, bytes, sc_off, r, ng, x1, x2),
4592            b => unreachable!("vbit bit-width {b} (validated at load)"),
4593        };
4594        // SAFETY: disjoint row ranges per worker.
4595        unsafe {
4596            *p1.at(r) = v1;
4597            *p2.at(r) = v2;
4598        }
4599    }
4600}
4601
4602// ───────────────────── q4_tiled kernels (§4.3) ─────────────────────
4603
4604/// One q4_tiled row dot on the A8W8 int8 path: per 32-group the tile
4605/// is ONE sequential read — [f16 scale][16B nibbles] — versus the two
4606/// distant streams of the split layout. Values/order identical to the
4607/// split kernels.
4608#[inline]
4609#[allow(unreachable_code)]
4610fn dot_q4t_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
4611    #[cfg(target_arch = "aarch64")]
4612    unsafe {
4613        return dot_q4t_row_sdot(bytes, r, gpr, xq);
4614    }
4615    #[cfg(target_arch = "x86_64")]
4616    unsafe {
4617        if vnni_tiles_enabled() {
4618            return dot_q4t_row_vnni(bytes, r, gpr, xq);
4619        }
4620        return dot_q4t_row_avx2(bytes, r, gpr, xq);
4621    }
4622    let mut acc = 0f32;
4623    for gi in 0..gpr {
4624        let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
4625        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
4626        let mut d = 0i32;
4627        for (k, &b) in tile[2..].iter().enumerate() {
4628            d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
4629                + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
4630        }
4631        acc += d as f32 * s;
4632    }
4633    acc
4634}
4635
4636#[cfg(target_arch = "aarch64")]
4637#[target_feature(enable = "neon,dotprod")]
4638unsafe fn dot_q4t_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
4639    // SAFETY: callers uphold slice-length contracts (18B tile per group,
4640    // xq.len() == gpr·GROUP_SIZE).
4641    unsafe {
4642        use core::arch::aarch64::*;
4643        use core::arch::asm;
4644        let lomask = vdupq_n_u8(0x0F);
4645        let eight = vdupq_n_s8(8);
4646        let mut acc = 0f32;
4647        for gi in 0..gpr {
4648            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4649            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4650            let b = vld1q_u8(t.add(2));
4651            let lo = vandq_u8(b, lomask);
4652            let hi = vshrq_n_u8::<4>(b);
4653            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
4654            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
4655            let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
4656            let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
4657            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
4658            asm!(
4659                "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
4660                "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
4661                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4662                e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
4663                options(pure, nomem, nostack),
4664            );
4665            acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
4666        }
4667        acc
4668    }
4669}
4670
4671#[cfg(target_arch = "x86_64")]
4672#[target_feature(enable = "avx2")]
4673unsafe fn dot_q4t_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
4674    // SAFETY: see dot_q4t_row_sdot.
4675    unsafe {
4676        use core::arch::x86_64::*;
4677        let lomask = _mm_set1_epi8(0x0F);
4678        let eight = _mm256_set1_epi8(8);
4679        let ones = _mm256_set1_epi16(1);
4680        let mut acc = 0f32;
4681        for gi in 0..gpr {
4682            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4683            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4684            let b = _mm_loadu_si128(t.add(2) as *const __m128i);
4685            let lo = _mm_and_si128(b, lomask);
4686            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
4687            let w = _mm256_sub_epi8(
4688                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4689                eight,
4690            );
4691            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4692            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
4693            let d = _mm256_madd_epi16(p16, ones);
4694            let hi128 = _mm256_extracti128_si256::<1>(d);
4695            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
4696            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
4697            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
4698            acc += _mm_cvtsi128_si32(s32) as f32 * s;
4699        }
4700        acc
4701    }
4702}
4703
4704/// VNNI twin of `dot_q4t_row_avx2`: same unpack, `vpdpbusd` replaces
4705/// the maddubs+madd pair (see `dpbusd_hsum` — sums are bit-identical).
4706/// 256-bit VL encoding, so the VEX `vpsignb` stays usable.
4707#[cfg(target_arch = "x86_64")]
4708#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
4709unsafe fn dot_q4t_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
4710    // SAFETY: see dot_q4t_row_sdot.
4711    unsafe {
4712        use core::arch::x86_64::*;
4713        let lomask = _mm_set1_epi8(0x0F);
4714        let eight = _mm256_set1_epi8(8);
4715        let mut acc = 0f32;
4716        for gi in 0..gpr {
4717            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4718            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4719            let b = _mm_loadu_si128(t.add(2) as *const __m128i);
4720            let lo = _mm_and_si128(b, lomask);
4721            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
4722            let w = _mm256_sub_epi8(
4723                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4724                eight,
4725            );
4726            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4727            let d = dpbusd_hsum(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
4728            acc += d as f32 * s;
4729        }
4730        acc
4731    }
4732}
4733
4734/// One q4_tiled row against FOUR activation streams: the nibble unpack
4735/// and abs() happen once per group instead of once per (group,
4736/// activation) — the unpack is the dominant per-element cost of the
4737/// tiled format (roadmap P0 portable blocking, q4t leg).
4738#[cfg(target_arch = "x86_64")]
4739// `fma` is NOT implied by `avx2`: without it LLVM lowers _mm256_fmadd_ps
4740// to a libm call per lane — measured 2x slower than the reduction this
4741// kernel replaces. The runtime gate (`avx2_enabled`) already requires
4742// both features, so declaring it here is safe.
4743#[target_feature(enable = "avx2,fma")]
4744unsafe fn dot_q4t_row_1x4_avx2(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
4745    // SAFETY: callers uphold the 18B-tile and xq-length contracts.
4746    unsafe {
4747        use core::arch::x86_64::*;
4748        let lomask = _mm_set1_epi8(0x0F);
4749        let eight = _mm256_set1_epi8(8);
4750        let ones = _mm256_set1_epi16(1);
4751        // One f32 accumulator VECTOR per activation, reduced once at the
4752        // end. Folding each group's i32 lanes to a scalar inside the loop
4753        // costs an extracti128 + three shift/add + a movd — a cross-lane
4754        // dependency chain per (group, activation), 288 of them per row at
4755        // cols=2304. The per-group scale is what forces a float
4756        // accumulator; it does not force a horizontal sum.
4757        //
4758        // The four accumulators are NAMED, not an array: as `[__m256; 4]`
4759        // indexed by a loop variable LLVM keeps them in memory and every
4760        // group pays four 32-byte loads and stores. That alone made this
4761        // kernel 2x SLOWER than the per-group reduction it replaces
4762        // (measured on the EPYC box: 150 s vs 71 s for two 256² steps).
4763        let mut f0 = _mm256_setzero_ps();
4764        let mut f1 = _mm256_setzero_ps();
4765        let mut f2 = _mm256_setzero_ps();
4766        let mut f3 = _mm256_setzero_ps();
4767        for gi in 0..gpr {
4768            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4769            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4770            let sv = _mm256_set1_ps(s);
4771            let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
4772            let lo = _mm_and_si128(bb, lomask);
4773            let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
4774            let w = _mm256_sub_epi8(
4775                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4776                eight,
4777            );
4778            let aw = _mm256_abs_epi8(w);
4779            let off = gi * GROUP_SIZE;
4780            let dot = |xq: &[i8]| {
4781                let x = _mm256_loadu_si256(xq.as_ptr().add(off) as *const __m256i);
4782                let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
4783                _mm256_cvtepi32_ps(_mm256_madd_epi16(p16, ones))
4784            };
4785            f0 = _mm256_fmadd_ps(dot(xs[0]), sv, f0);
4786            f1 = _mm256_fmadd_ps(dot(xs[1]), sv, f1);
4787            f2 = _mm256_fmadd_ps(dot(xs[2]), sv, f2);
4788            f3 = _mm256_fmadd_ps(dot(xs[3]), sv, f3);
4789        }
4790        [
4791            hsum256_ps(f0),
4792            hsum256_ps(f1),
4793            hsum256_ps(f2),
4794            hsum256_ps(f3),
4795        ]
4796    }
4797}
4798
4799/// Horizontal sum of eight f32 lanes — the one cross-lane reduction the
4800/// blocked kernels pay, once per row instead of once per group.
4801#[cfg(target_arch = "x86_64")]
4802#[target_feature(enable = "avx2")]
4803#[inline]
4804unsafe fn hsum256_ps(v: core::arch::x86_64::__m256) -> f32 {
4805    // SAFETY: pure register arithmetic on the caller's vector.
4806    unsafe {
4807        use core::arch::x86_64::*;
4808        let hi = _mm256_extractf128_ps::<1>(v);
4809        let s = _mm_add_ps(_mm256_castps256_ps128(v), hi);
4810        let s = _mm_add_ps(s, _mm_movehl_ps(s, s));
4811        let s = _mm_add_ss(s, _mm_shuffle_ps::<0x55>(s, s));
4812        _mm_cvtss_f32(s)
4813    }
4814}
4815
4816/// VNNI twin of `dot_q4t_row_1x4_avx2` (see `dpbusd_hsum`).
4817#[cfg(target_arch = "x86_64")]
4818#[target_feature(enable = "avx2,fma,avx512f,avx512bw,avx512vl,avx512vnni")]
4819unsafe fn dot_q4t_row_1x4_vnni(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
4820    // SAFETY: callers uphold the 18B-tile and xq-length contracts.
4821    unsafe {
4822        use core::arch::x86_64::*;
4823        let lomask = _mm_set1_epi8(0x0F);
4824        let eight = _mm256_set1_epi8(8);
4825        // Same shape as the AVX2 twin: accumulate in f32 vectors and pay
4826        // one cross-lane reduction per row, not per (group, activation).
4827        let mut f0 = _mm256_setzero_ps();
4828        let mut f1 = _mm256_setzero_ps();
4829        let mut f2 = _mm256_setzero_ps();
4830        let mut f3 = _mm256_setzero_ps();
4831        for gi in 0..gpr {
4832            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4833            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4834            let sv = _mm256_set1_ps(s);
4835            let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
4836            let lo = _mm_and_si128(bb, lomask);
4837            let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
4838            let w = _mm256_sub_epi8(
4839                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4840                eight,
4841            );
4842            let aw = _mm256_abs_epi8(w);
4843            let off = gi * GROUP_SIZE;
4844            let dot = |xq: &[i8]| {
4845                let x = _mm256_loadu_si256(xq.as_ptr().add(off) as *const __m256i);
4846                _mm256_cvtepi32_ps(_mm256_dpbusd_epi32(
4847                    _mm256_setzero_si256(),
4848                    aw,
4849                    _mm256_sign_epi8(x, w),
4850                ))
4851            };
4852            f0 = _mm256_fmadd_ps(dot(xs[0]), sv, f0);
4853            f1 = _mm256_fmadd_ps(dot(xs[1]), sv, f1);
4854            f2 = _mm256_fmadd_ps(dot(xs[2]), sv, f2);
4855            f3 = _mm256_fmadd_ps(dot(xs[3]), sv, f3);
4856        }
4857        let acc = [
4858            hsum256_ps(f0),
4859            hsum256_ps(f1),
4860            hsum256_ps(f2),
4861            hsum256_ps(f3),
4862        ];
4863        acc
4864    }
4865}
4866
4867/// ARM twin of `dot_q4t_row_1x4_avx2`: one nibble unpack per group
4868/// serves FOUR activation streams. Per stream the group order and f32
4869/// accumulation match `dot_q4t_row_sdot` exactly — batch == matvec
4870/// bit-for-bit.
4871#[cfg(target_arch = "aarch64")]
4872#[target_feature(enable = "neon,dotprod")]
4873unsafe fn dot_q4t_row_1x4_sdot(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
4874    // SAFETY: callers uphold the 18B-tile and xq-length contracts.
4875    unsafe {
4876        use core::arch::aarch64::*;
4877        use core::arch::asm;
4878        let lomask = vdupq_n_u8(0x0F);
4879        let eight = vdupq_n_s8(8);
4880        let mut acc = [0f32; 4];
4881        for gi in 0..gpr {
4882            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4883            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4884            let b = vld1q_u8(t.add(2));
4885            let lo = vandq_u8(b, lomask);
4886            let hi = vshrq_n_u8::<4>(b);
4887            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
4888            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
4889            for (k, xq) in xs.iter().enumerate() {
4890                let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
4891                let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
4892                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
4893                asm!(
4894                    "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
4895                    "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
4896                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4897                    e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
4898                    options(pure, nomem, nostack),
4899                );
4900                acc[k] += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
4901            }
4902        }
4903        acc
4904    }
4905}
4906
4907/// Exact-term correction for A8W8 outliers on a tiled row.
4908#[inline]
4909fn q4t_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
4910    let gi = j / GROUP_SIZE;
4911    let k = j % GROUP_SIZE;
4912    let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
4913    let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
4914    let byte = tile[2 + k / 2];
4915    let nib = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
4916    ((nib as i32 - 8) as f32, s)
4917}
4918
4919/// Exact scalar q4_tiled row (CMF_SDOT=0 contract) — same pairwise
4920/// accumulation shape as `q4_range_f32`.
4921#[inline]
4922fn q4t_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
4923    let mut acc = 0f32;
4924    for gi in 0..gpr {
4925        let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
4926        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
4927        let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
4928        let mut ga = 0f32;
4929        for (k, &b) in tile[2..].iter().enumerate() {
4930            ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
4931                + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
4932        }
4933        acc += ga * s;
4934    }
4935    acc
4936}
4937
4938/// Split view of a `q4tp` payload. The three planes are resolved once per
4939/// matvec instead of per row — `q4tp_sections` is cheap, but doing it inside
4940/// the row loop would put a division on the hot path for nothing.
4941struct Q4tpView<'a> {
4942    nib: &'a [u8],
4943    params: &'a [u8],
4944    codes: &'a [u8],
4945    stride: usize,
4946    /// q2tp reads the ladder with rung 0 = exact zero.
4947    zero_rung: bool,
4948}
4949
4950impl<'a> Q4tpView<'a> {
4951    fn new(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
4952        let (params_off, codes_off, stride) = q4tp_sections(rows, cols);
4953        Self {
4954            nib: &bytes[..params_off],
4955            params: &bytes[params_off..codes_off],
4956            codes: &bytes[codes_off..],
4957            stride,
4958            zero_rung: false,
4959        }
4960    }
4961
4962    /// The q2tp view: identical params/codes planes, 8 B weight chunks.
4963    fn new_q2(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
4964        let (params_off, codes_off, stride) = q2tp_sections(rows, cols);
4965        Self {
4966            nib: &bytes[..params_off],
4967            params: &bytes[params_off..codes_off],
4968            codes: &bytes[codes_off..],
4969            stride,
4970            zero_rung: true,
4971        }
4972    }
4973
4974    /// Expand row `r`'s per-tile scales into `out` (length `gpr`).
4975    ///
4976    /// Doing this once per row — rather than decoding a 5-bit code inside the
4977    /// tile loop — is what makes the format free at runtime. Random access to
4978    /// a packed 5-bit field costs a division, two bounds checks and a branch;
4979    /// the tile's actual work is two `sdot`s, so per-tile decoding dominated
4980    /// the kernel and cost 5x (measured: 1.4 vs 6.9 tok/s on Nanbeige-3B).
4981    /// Walking the plane sequentially with a bit accumulator is ~3 ops.
4982    /// Eight 5-bit codes are exactly five bytes, so a whole group of
4983    /// eight decodes from one little-endian word at fixed shifts. The
4984    /// bit-accumulator this replaces carried a data-dependent `while
4985    /// have < 5` refill whose branch sat in the innermost loop of every
4986    /// q4tp row; a decode profile put this function above the dot
4987    /// products it feeds. Same bitstream, same codes — just no branch
4988    /// and eight independent extractions.
4989    #[inline]
4990    fn scales_into(&self, r: usize, gpr: usize, out: &mut [f32]) {
4991        let tab = if self.zero_rung {
4992            q2tp_ladder(self.params, r)
4993        } else {
4994            q4tp_ladder(self.params, r)
4995        };
4996        let codes = &self.codes[r * self.stride..(r + 1) * self.stride];
4997        let out = &mut out[..gpr];
4998        let mut chunks = out.chunks_exact_mut(8);
4999        let mut ci = 0usize;
5000        for c in &mut chunks {
5001            let w = u64::from(codes[ci])
5002                | u64::from(codes[ci + 1]) << 8
5003                | u64::from(codes[ci + 2]) << 16
5004                | u64::from(codes[ci + 3]) << 24
5005                | u64::from(codes[ci + 4]) << 32;
5006            for (k, o) in c.iter_mut().enumerate() {
5007                *o = tab[((w >> (5 * k)) & 31) as usize];
5008            }
5009            ci += 5;
5010        }
5011        // Fewer than eight codes left: the shared total accessor, which
5012        // tolerates a 5-bit field whose spill byte is past the stride.
5013        let tail = &codes[ci..];
5014        for (k, o) in chunks.into_remainder().iter_mut().enumerate() {
5015            *o = tab[q4tp_code(tail, k)];
5016        }
5017    }
5018}
5019
5020#[inline]
5021fn dot_q4tp_row_i8(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
5022    #[cfg(target_arch = "aarch64")]
5023    unsafe {
5024        return dot_q4tp_row_sdot(nib, r, gpr, xq, scales);
5025    }
5026    #[cfg(target_arch = "x86_64")]
5027    unsafe {
5028        if vnni_tiles_enabled() {
5029            return dot_q4tp_row_vnni(nib, r, gpr, xq, scales);
5030        }
5031        return dot_q4tp_row_avx2(nib, r, gpr, xq, scales);
5032    }
5033    #[allow(unreachable_code)]
5034    {
5035        let mut acc = 0f32;
5036        for gi in 0..gpr {
5037            let tile = &nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
5038            let s = scales[gi];
5039            let mut d = 0i32;
5040            for (k, &b) in tile.iter().enumerate() {
5041                d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
5042                    + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
5043            }
5044            acc += d as f32 * s;
5045        }
5046        acc
5047    }
5048}
5049
5050/// q4tp twin of `dot_q4t_row_sdot`: identical nibble math, but the tile
5051/// stride is 16 B (no inline scale) and the scale is a ladder lookup.
5052#[cfg(target_arch = "aarch64")]
5053#[target_feature(enable = "neon,dotprod")]
5054unsafe fn dot_q4tp_row_sdot(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
5055    // SAFETY: callers uphold slice-length contracts (16B tile per group,
5056    // xq.len() == gpr·GROUP_SIZE, codes covering gpr 5-bit fields).
5057    unsafe {
5058        use core::arch::aarch64::*;
5059        use core::arch::asm;
5060        let lomask = vdupq_n_u8(0x0F);
5061        let eight = vdupq_n_s8(8);
5062        let mut acc = 0f32;
5063        for gi in 0..gpr {
5064            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5065            let s = *scales.get_unchecked(gi);
5066            let b = vld1q_u8(t);
5067            let lo = vandq_u8(b, lomask);
5068            let hi = vshrq_n_u8::<4>(b);
5069            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
5070            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
5071            let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
5072            let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
5073            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5074            asm!(
5075                "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
5076                "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
5077                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5078                e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
5079                options(pure, nomem, nostack),
5080            );
5081            acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
5082        }
5083        acc
5084    }
5085}
5086
5087#[cfg(target_arch = "x86_64")]
5088#[target_feature(enable = "avx2")]
5089unsafe fn dot_q4tp_row_avx2(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
5090    // SAFETY: see dot_q4tp_row_sdot.
5091    unsafe {
5092        use core::arch::x86_64::*;
5093        let lomask = _mm_set1_epi8(0x0F);
5094        let eight = _mm256_set1_epi8(8);
5095        let ones = _mm256_set1_epi16(1);
5096        let mut acc = 0f32;
5097        for gi in 0..gpr {
5098            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5099            let s = *scales.get_unchecked(gi);
5100            let b = _mm_loadu_si128(t as *const __m128i);
5101            let lo = _mm_and_si128(b, lomask);
5102            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
5103            let w = _mm256_sub_epi8(
5104                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5105                eight,
5106            );
5107            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5108            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
5109            let d = _mm256_madd_epi16(p16, ones);
5110            let hi128 = _mm256_extracti128_si256::<1>(d);
5111            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
5112            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
5113            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
5114            acc += _mm_cvtsi128_si32(s32) as f32 * s;
5115        }
5116        acc
5117    }
5118}
5119
5120/// VNNI twin of `dot_q4tp_row_avx2` (see `dot_q4t_row_vnni` for why the
5121/// 256-bit VL encoding is the one to use here).
5122#[cfg(target_arch = "x86_64")]
5123#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
5124unsafe fn dot_q4tp_row_vnni(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
5125    // SAFETY: see dot_q4tp_row_sdot.
5126    unsafe {
5127        use core::arch::x86_64::*;
5128        let lomask = _mm_set1_epi8(0x0F);
5129        let eight = _mm256_set1_epi8(8);
5130        let mut acc = 0f32;
5131        for gi in 0..gpr {
5132            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5133            let s = *scales.get_unchecked(gi);
5134            let b = _mm_loadu_si128(t as *const __m128i);
5135            let lo = _mm_and_si128(b, lomask);
5136            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
5137            let w = _mm256_sub_epi8(
5138                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5139                eight,
5140            );
5141            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5142            acc += dpbusd_hsum(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w)) as f32 * s;
5143        }
5144        acc
5145    }
5146}
5147
5148/// Exact scalar q4tp row — the `CMF_SDOT=0` contract, same pairwise
5149/// accumulation shape as `q4t_row_exact`.
5150#[inline]
5151fn q4tp_row_exact(nib: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
5152    let mut acc = 0f32;
5153    for gi in 0..gpr {
5154        let tile = &nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
5155        let s = scales[gi];
5156        let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
5157        let mut ga = 0f32;
5158        for (k, &b) in tile.iter().enumerate() {
5159            ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
5160                + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
5161        }
5162        acc += ga * s;
5163    }
5164    acc
5165}
5166
5167/// Single weight of a q4tp tensor — the a8w8 outlier path, which restores
5168/// activation outliers at full precision after the int8 pass.
5169#[inline]
5170fn q4tp_outlier(nib: &[u8], r: usize, gpr: usize, j: usize, scales: &[f32]) -> (f32, f32) {
5171    let (gi, k) = (j / GROUP_SIZE, j % GROUP_SIZE);
5172    let byte = nib[(r * gpr + gi) * Q4TP_NIB + k / 2];
5173    let n = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
5174    ((n as i32 - 8) as f32, scales[gi])
5175}
5176
5177/// Fused q4tp matvec (dispatch mirrors `q4t_matvec`).
5178fn q4tp_matvec(
5179    bytes: &[u8],
5180    x: &[f32],
5181    rows: usize,
5182    cols: usize,
5183    out: &mut [f32],
5184    pool: Option<&Pool>,
5185) {
5186    debug_assert_eq!(out.len(), rows);
5187    let gpr = cols / GROUP_SIZE;
5188    let v = Q4tpView::new(bytes, rows, cols);
5189    let out_addr = SendMut(out.as_mut_ptr());
5190    if a8w8_enabled() {
5191        let act = split_act(x);
5192        let run = |start: usize, end: usize| {
5193            // One scratch row of scales per worker — borrowed, not minted.
5194            with_krow(gpr, |sc| {
5195                for r in start..end {
5196                    v.scales_into(r, gpr, sc);
5197                    let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, sc) * act.sx;
5198                    for &(j, xv) in &act.outliers {
5199                        let (w, s) = q4tp_outlier(v.nib, r, gpr, j, sc);
5200                        acc += w * s * xv;
5201                    }
5202                    // SAFETY: disjoint row ranges per worker.
5203                    unsafe { *out_addr.at(r) = acc };
5204                }
5205            })
5206        };
5207        dispatch_rows(pool, rows, &run);
5208        return;
5209    }
5210    let run = |start: usize, end: usize| {
5211        with_krow(gpr, |sc| {
5212            for r in start..end {
5213                v.scales_into(r, gpr, sc);
5214                // SAFETY: disjoint row ranges per worker.
5215                unsafe { *out_addr.at(r) = q4tp_row_exact(v.nib, r, gpr, x, sc) };
5216            }
5217        })
5218    };
5219    dispatch_rows(pool, rows, &run);
5220}
5221
5222/// Fused two-input q4tp matvec — the SwiGLU gate/up pair. Weights and the
5223/// row ladder are read once and spent on both activation streams.
5224#[allow(clippy::too_many_arguments)]
5225fn q4tp_matvec2(
5226    bytes: &[u8],
5227    x1: &[f32],
5228    x2: &[f32],
5229    rows: usize,
5230    cols: usize,
5231    o1: &mut [f32],
5232    o2: &mut [f32],
5233    pool: Option<&Pool>,
5234) {
5235    let gpr = cols / GROUP_SIZE;
5236    let v = Q4tpView::new(bytes, rows, cols);
5237    let (p1, p2) = (SendMut(o1.as_mut_ptr()), SendMut(o2.as_mut_ptr()));
5238    let run = |start: usize, end: usize| {
5239        let mut sc = vec![0f32; gpr];
5240        for r in start..end {
5241            v.scales_into(r, gpr, &mut sc);
5242            // SAFETY: disjoint row ranges per worker.
5243            unsafe {
5244                *p1.at(r) = q4tp_row_exact(v.nib, r, gpr, x1, &sc);
5245                *p2.at(r) = q4tp_row_exact(v.nib, r, gpr, x2, &sc);
5246            }
5247        }
5248    };
5249    dispatch_rows(pool, rows, &run);
5250}
5251
5252/// One q2tp outlier weight at column `j` of row `r`: the 2-bit code and
5253/// its group scale, mirrored on `q4tp_outlier`.
5254#[inline]
5255fn q2tp_outlier(chunks: &[u8], r: usize, gpr: usize, j: usize, scales: &[f32]) -> (f32, f32) {
5256    let (gi, k) = (j / GROUP_SIZE, j % GROUP_SIZE);
5257    let byte = chunks[(r * gpr + gi) * Q2TP_CHUNK + k / 4];
5258    let c = (byte >> (2 * (k % 4))) & 3;
5259    (c as f32 - 1.5, scales[gi])
5260}
5261
5262#[cfg(target_arch = "x86_64")]
5263const Q2TP_DECODE_U32: [u32; 256] = {
5264    let mut tab = [0u32; 256];
5265    let mut b = 0usize;
5266    while b < 256 {
5267        tab[b] = ((b as u32) & 3)
5268            | ((((b as u32) >> 2) & 3) << 8)
5269            | ((((b as u32) >> 4) & 3) << 16)
5270            | ((((b as u32) >> 6) & 3) << 24);
5271        b += 1;
5272    }
5273    tab
5274};
5275
5276/// Eight packed q2tp bytes against 32 signed activation bytes. `maddubs`
5277/// exactly computes unsigned 2-bit code × signed i8; its pair sums cannot
5278/// saturate (2 × 3 × 127 < i16::MAX), and the second madd widens to i32.
5279#[cfg(target_arch = "x86_64")]
5280#[target_feature(enable = "avx2")]
5281unsafe fn q2tp_code_dot_avx2(ch: &[u8], x: &[i8]) -> i32 {
5282    use core::arch::x86_64::*;
5283    debug_assert!(ch.len() >= Q2TP_CHUNK && x.len() >= GROUP_SIZE);
5284    let codes = _mm256_setr_epi32(
5285        Q2TP_DECODE_U32[ch[0] as usize] as i32,
5286        Q2TP_DECODE_U32[ch[1] as usize] as i32,
5287        Q2TP_DECODE_U32[ch[2] as usize] as i32,
5288        Q2TP_DECODE_U32[ch[3] as usize] as i32,
5289        Q2TP_DECODE_U32[ch[4] as usize] as i32,
5290        Q2TP_DECODE_U32[ch[5] as usize] as i32,
5291        Q2TP_DECODE_U32[ch[6] as usize] as i32,
5292        Q2TP_DECODE_U32[ch[7] as usize] as i32,
5293    );
5294    let xv = unsafe { _mm256_loadu_si256(x.as_ptr().cast()) };
5295    let pair = _mm256_maddubs_epi16(codes, xv);
5296    let quad = _mm256_madd_epi16(pair, _mm256_set1_epi16(1));
5297    let sum128 = _mm_add_epi32(
5298        _mm256_castsi256_si128(quad),
5299        _mm256_extracti128_si256(quad, 1),
5300    );
5301    let sum64 = _mm_hadd_epi32(sum128, sum128);
5302    _mm_cvtsi128_si32(_mm_hadd_epi32(sum64, sum64))
5303}
5304
5305/// Integer dot of one q2tp row against pre-quantized activations:
5306/// Σ_g s_g · (Σ c·xq − 1.5·Σ xq). The half-integer grid (c − 1.5)
5307/// becomes exact integer math through the group sums — the same trick
5308/// every a8w8 kernel in this file rides. The codes decode into a
5309/// 32-byte scratch in natural order and the dot itself is the shared
5310/// SDOT primitive; elsewhere a scalar integer loop.
5311#[inline]
5312fn dot_q2tp_row_i8(
5313    chunks: &[u8],
5314    r: usize,
5315    gpr: usize,
5316    xq: &[i8],
5317    gsum: &[i32],
5318    scales: &[f32],
5319) -> f32 {
5320    let mut acc = 0f32;
5321    let base = r * gpr * Q2TP_CHUNK;
5322    #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
5323    let mut codes = [0i8; GROUP_SIZE];
5324    #[cfg(target_arch = "x86_64")]
5325    let avx2 = std::arch::is_x86_feature_detected!("avx2");
5326    for gi in 0..gpr {
5327        let ch = &chunks[base + gi * Q2TP_CHUNK..base + (gi + 1) * Q2TP_CHUNK];
5328        let xg = &xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
5329        #[cfg(target_arch = "aarch64")]
5330        // NEON: the byte's four 2-bit fields land in four lane vectors
5331        // (shift+mask), vld4 de-interleaves xq to match (xj[k] =
5332        // xq[4k+j]), widening MACs accumulate exactly in i32. A scalar
5333        // decode here cost as much as the dot it fed — the profile put
5334        // it at the top of the whole W2 decode.
5335        let dot = unsafe {
5336            use core::arch::aarch64::*;
5337            let b = vld1_u8(ch.as_ptr());
5338            let three = vdup_n_u8(3);
5339            let c0 = vreinterpret_s8_u8(vand_u8(b, three));
5340            let c1 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 2), three));
5341            let c2 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 4), three));
5342            let c3 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 6), three));
5343            let x4 = vld4_s8(xg.as_ptr());
5344            let mut acc4 = vdupq_n_s32(0);
5345            acc4 = vpadalq_s16(acc4, vmull_s8(c0, x4.0));
5346            acc4 = vpadalq_s16(acc4, vmull_s8(c1, x4.1));
5347            acc4 = vpadalq_s16(acc4, vmull_s8(c2, x4.2));
5348            acc4 = vpadalq_s16(acc4, vmull_s8(c3, x4.3));
5349            vaddvq_s32(acc4)
5350        };
5351        #[cfg(target_arch = "x86_64")]
5352        let dot: i32 = if avx2 {
5353            // SAFETY: the runtime feature check gates the target-feature body;
5354            // the group slices above are exactly 8 and 32 bytes long.
5355            unsafe { q2tp_code_dot_avx2(ch, xg) }
5356        } else {
5357            ch.iter()
5358                .enumerate()
5359                .map(|(k, &b)| {
5360                    ((b & 3) as i32) * xg[k * 4] as i32
5361                        + (((b >> 2) & 3) as i32) * xg[k * 4 + 1] as i32
5362                        + (((b >> 4) & 3) as i32) * xg[k * 4 + 2] as i32
5363                        + (((b >> 6) & 3) as i32) * xg[k * 4 + 3] as i32
5364                })
5365                .sum()
5366        };
5367        #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
5368        let dot: i32 = {
5369            for (k, &b) in ch.iter().enumerate() {
5370                codes[k * 4] = (b & 3) as i8;
5371                codes[k * 4 + 1] = ((b >> 2) & 3) as i8;
5372                codes[k * 4 + 2] = ((b >> 4) & 3) as i8;
5373                codes[k * 4 + 3] = ((b >> 6) & 3) as i8;
5374            }
5375            codes
5376                .iter()
5377                .zip(xg)
5378                .map(|(&c, &x)| c as i32 * x as i32)
5379                .sum()
5380        };
5381        acc += scales[gi] * (dot as f32 - 1.5 * gsum[gi] as f32);
5382    }
5383    acc
5384}
5385
5386/// Exact f32 dot of one q2tp row: 2-bit fields LSB-first, (c − 1.5)·s.
5387/// Scalar on purpose — the 2-bit class targets the GPU graph; the CPU
5388/// path exists for parity gates and small-machine fallback.
5389fn q2tp_row_exact(chunks: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
5390    q2tp_row_exact_center(chunks, r, gpr, x, scales, 1.5)
5391}
5392
5393/// Fused Prism affine row: the derived correction is applied inside the
5394/// decoded code, avoiding a second accumulated dot and avoiding cancellation
5395/// between `B=(c-1.5)s` and `+.5s` for long 5120/17408 rows.
5396#[inline]
5397fn q2tp_affine_row_exact(chunks: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
5398    q2tp_row_exact_center(chunks, r, gpr, x, scales, 1.0)
5399}
5400
5401#[inline]
5402fn q2tp_row_exact_center(
5403    chunks: &[u8],
5404    r: usize,
5405    gpr: usize,
5406    x: &[f32],
5407    scales: &[f32],
5408    center: f32,
5409) -> f32 {
5410    let mut acc = 0f32;
5411    for gi in 0..gpr {
5412        let ch = &chunks[(r * gpr + gi) * Q2TP_CHUNK..(r * gpr + gi + 1) * Q2TP_CHUNK];
5413        let s = scales[gi];
5414        let xb = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
5415        let mut g = 0f32;
5416        for (k, &b) in ch.iter().enumerate() {
5417            g += ((b & 3) as f32 - center) * xb[k * 4]
5418                + (((b >> 2) & 3) as f32 - center) * xb[k * 4 + 1]
5419                + (((b >> 4) & 3) as f32 - center) * xb[k * 4 + 2]
5420                + (((b >> 6) & 3) as f32 - center) * xb[k * 4 + 3];
5421        }
5422        acc += s * g;
5423    }
5424    acc
5425}
5426
5427fn q2tp_matvec(
5428    bytes: &[u8],
5429    x: &[f32],
5430    rows: usize,
5431    cols: usize,
5432    out: &mut [f32],
5433    pool: Option<&Pool>,
5434) {
5435    q2tp_matvec_mode(bytes, x, rows, cols, out, pool, false);
5436}
5437
5438fn q2tp_affine_matvec(
5439    bytes: &[u8],
5440    x: &[f32],
5441    rows: usize,
5442    cols: usize,
5443    out: &mut [f32],
5444    pool: Option<&Pool>,
5445) {
5446    q2tp_matvec_mode(bytes, x, rows, cols, out, pool, true);
5447}
5448
5449fn q2tp_matvec_mode(
5450    bytes: &[u8],
5451    x: &[f32],
5452    rows: usize,
5453    cols: usize,
5454    out: &mut [f32],
5455    pool: Option<&Pool>,
5456    affine: bool,
5457) {
5458    debug_assert_eq!(out.len(), rows);
5459    let gpr = cols / GROUP_SIZE;
5460    let v = Q4tpView::new_q2(bytes, rows, cols);
5461    let out_addr = SendMut(out.as_mut_ptr());
5462    // a8w8 fast path (CMF_SDOT=0 keeps the exact scalar walk): integer
5463    // code dots + group sums, exact outlier correction — the same
5464    // contract as every sibling kernel; measured 2-bit rows were the
5465    // only scalar holdout in the family.
5466    if !affine && a8w8_enabled() {
5467        let act = split_act(x);
5468        let gsum = q1_group_sums(&act.xq, gpr);
5469        let (act, gsum) = (&act, &gsum);
5470        let run = move |start: usize, end: usize| {
5471            with_krow(gpr, |sc| {
5472                for r in start..end {
5473                    v.scales_into(r, gpr, sc);
5474                    let mut acc = dot_q2tp_row_i8(v.nib, r, gpr, &act.xq, gsum, sc) * act.sx;
5475                    for &(j, xv) in &act.outliers {
5476                        let (w, s) = q2tp_outlier(v.nib, r, gpr, j, sc);
5477                        acc += w * s * xv;
5478                    }
5479                    // SAFETY: disjoint row ranges per worker.
5480                    unsafe { *out_addr.at(r) = acc };
5481                }
5482            })
5483        };
5484        dispatch_rows(pool, rows, &run);
5485        return;
5486    }
5487    let run = |start: usize, end: usize| {
5488        with_krow(gpr, |sc| {
5489            for r in start..end {
5490                v.scales_into(r, gpr, sc);
5491                // SAFETY: disjoint row ranges per worker.
5492                unsafe {
5493                    *out_addr.at(r) = if affine {
5494                        q2tp_affine_row_exact(v.nib, r, gpr, x, sc)
5495                    } else {
5496                        q2tp_row_exact(v.nib, r, gpr, x, sc)
5497                    }
5498                };
5499            }
5500        })
5501    };
5502    dispatch_rows(pool, rows, &run);
5503}
5504
5505/// Fused two-input q2tp matvec — the SwiGLU gate/up pair.
5506#[allow(clippy::too_many_arguments)]
5507fn q2tp_matvec2(
5508    bytes: &[u8],
5509    x1: &[f32],
5510    x2: &[f32],
5511    rows: usize,
5512    cols: usize,
5513    o1: &mut [f32],
5514    o2: &mut [f32],
5515    pool: Option<&Pool>,
5516) {
5517    q2tp_matvec2_mode(bytes, x1, x2, rows, cols, o1, o2, pool, false);
5518}
5519
5520#[allow(clippy::too_many_arguments)]
5521fn q2tp_affine_matvec2(
5522    bytes: &[u8],
5523    x1: &[f32],
5524    x2: &[f32],
5525    rows: usize,
5526    cols: usize,
5527    o1: &mut [f32],
5528    o2: &mut [f32],
5529    pool: Option<&Pool>,
5530) {
5531    q2tp_matvec2_mode(bytes, x1, x2, rows, cols, o1, o2, pool, true);
5532}
5533
5534#[allow(clippy::too_many_arguments)]
5535fn q2tp_matvec2_mode(
5536    bytes: &[u8],
5537    x1: &[f32],
5538    x2: &[f32],
5539    rows: usize,
5540    cols: usize,
5541    o1: &mut [f32],
5542    o2: &mut [f32],
5543    pool: Option<&Pool>,
5544    affine: bool,
5545) {
5546    let gpr = cols / GROUP_SIZE;
5547    let v = Q4tpView::new_q2(bytes, rows, cols);
5548    let (p1, p2) = (SendMut(o1.as_mut_ptr()), SendMut(o2.as_mut_ptr()));
5549    let run = |start: usize, end: usize| {
5550        let mut sc = vec![0f32; gpr];
5551        for r in start..end {
5552            v.scales_into(r, gpr, &mut sc);
5553            // SAFETY: disjoint row ranges per worker.
5554            unsafe {
5555                *p1.at(r) = if affine {
5556                    q2tp_affine_row_exact(v.nib, r, gpr, x1, &sc)
5557                } else {
5558                    q2tp_row_exact(v.nib, r, gpr, x1, &sc)
5559                };
5560                *p2.at(r) = if affine {
5561                    q2tp_affine_row_exact(v.nib, r, gpr, x2, &sc)
5562                } else {
5563                    q2tp_row_exact(v.nib, r, gpr, x2, &sc)
5564                };
5565            }
5566        }
5567    };
5568    dispatch_rows(pool, rows, &run);
5569}
5570
5571/// Batched q2tp matmat: scalar row kernel over every batch column. CPU
5572/// prefill only — decode rides the graph, so plain and correct beats
5573/// clever here.
5574/// Test doors into the host 2-bit kernels: the stand's heap corruption
5575/// pointed at down-shaped tensors, and the private fns need a way to be
5576/// held to a reference without a model file around them.
5577pub fn q2tp_matvec_for_test(bytes: &[u8], x: &[f32], rows: usize, cols: usize, out: &mut [f32]) {
5578    // The facade IS the reference: encoder oracles hold requant output
5579    // to the exact scalar walk. The production dispatch may take the i8
5580    // fast path, whose error scale is the ACTIVATIONS' — a different
5581    // claim than the encoder correctness these tests pin.
5582    let gpr = cols / GROUP_SIZE;
5583    let v = Q4tpView::new_q2(bytes, rows, cols);
5584    with_krow(gpr, |sc| {
5585        for r in 0..rows {
5586            v.scales_into(r, gpr, sc);
5587            out[r] = q2tp_row_exact(v.nib, r, gpr, x, sc);
5588        }
5589    });
5590}
5591
5592/// Test door for the descriptor-specific fused affine decode.  Production
5593/// callers select this through a validated Prism header, never by dtype alone.
5594pub fn q2tp_affine_matvec_for_test(
5595    bytes: &[u8],
5596    x: &[f32],
5597    rows: usize,
5598    cols: usize,
5599    out: &mut [f32],
5600) {
5601    q2tp_affine_matvec(bytes, x, rows, cols, out, None);
5602}
5603
5604pub fn q2tp_matmat_for_test(
5605    bytes: &[u8],
5606    xs_all: &[f32],
5607    b: usize,
5608    rows: usize,
5609    cols: usize,
5610    out: &mut [f32],
5611) {
5612    q2tp_matmat(bytes, xs_all, b, rows, cols, out, None);
5613}
5614
5615fn q2tp_matmat(
5616    bytes: &[u8],
5617    xs_all: &[f32],
5618    b: usize,
5619    rows: usize,
5620    cols: usize,
5621    out: &mut [f32],
5622    pool: Option<&Pool>,
5623) {
5624    q2tp_matmat_mode(bytes, xs_all, b, rows, cols, out, pool, false);
5625}
5626
5627fn q2tp_affine_matmat(
5628    bytes: &[u8],
5629    xs_all: &[f32],
5630    b: usize,
5631    rows: usize,
5632    cols: usize,
5633    out: &mut [f32],
5634    pool: Option<&Pool>,
5635) {
5636    q2tp_matmat_mode(bytes, xs_all, b, rows, cols, out, pool, true);
5637}
5638
5639fn q2tp_matmat_mode(
5640    bytes: &[u8],
5641    xs_all: &[f32],
5642    b: usize,
5643    rows: usize,
5644    cols: usize,
5645    out: &mut [f32],
5646    pool: Option<&Pool>,
5647    affine: bool,
5648) {
5649    debug_assert_eq!(out.len(), b * rows);
5650    let gpr = cols / GROUP_SIZE;
5651    let v = Q4tpView::new_q2(bytes, rows, cols);
5652    let out_addr = SendMut(out.as_mut_ptr());
5653    let run = |start: usize, end: usize| {
5654        let mut sc = vec![0f32; gpr];
5655        for r in start..end {
5656            v.scales_into(r, gpr, &mut sc);
5657            for bi in 0..b {
5658                let x = &xs_all[bi * cols..(bi + 1) * cols];
5659                // SAFETY: disjoint row ranges per worker.
5660                unsafe {
5661                    *out_addr.at(bi * rows + r) = if affine {
5662                        q2tp_affine_row_exact(v.nib, r, gpr, x, &sc)
5663                    } else {
5664                        q2tp_row_exact(v.nib, r, gpr, x, &sc)
5665                    }
5666                };
5667            }
5668        }
5669    };
5670    dispatch_rows(pool, rows, &run);
5671}
5672
5673/// The pre-vectorised shape, kept for A/B (`CMF_Q4TP_V1=1`): the
5674/// horizontal add lands once per group per column instead of once per
5675/// row. Same weights, same activations — only the reduction differs.
5676#[cfg(target_arch = "aarch64")]
5677#[target_feature(enable = "neon,dotprod")]
5678unsafe fn dot_q4tp_row_1x4_sdot_v1(
5679    nib: &[u8],
5680    r: usize,
5681    gpr: usize,
5682    xs: [&[i8]; 4],
5683    scales: &[f32],
5684) -> [f32; 4] {
5685    unsafe {
5686        use core::arch::aarch64::*;
5687        use core::arch::asm;
5688        let lomask = vdupq_n_u8(0x0F);
5689        let eight = vdupq_n_s8(8);
5690        let (mut f0, mut f1, mut f2, mut f3) = (0f32, 0f32, 0f32, 0f32);
5691        for gi in 0..gpr {
5692            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5693            let s = *scales.get_unchecked(gi);
5694            let bb = vld1q_u8(t);
5695            let lo = vandq_u8(bb, lomask);
5696            let hi = vshrq_n_u8::<4>(bb);
5697            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
5698            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
5699            let mut d = [0f32; 4];
5700            for (k, dk) in d.iter_mut().enumerate() {
5701                let x0 = vld1q_s8(xs[k].as_ptr().add(gi * GROUP_SIZE));
5702                let x1 = vld1q_s8(xs[k].as_ptr().add(gi * GROUP_SIZE + 16));
5703                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5704                asm!(
5705                    "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
5706                    "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
5707                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5708                    e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
5709                    options(pure, nomem, nostack),
5710                );
5711                *dk = vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
5712            }
5713            f0 += d[0];
5714            f1 += d[1];
5715            f2 += d[2];
5716            f3 += d[3];
5717        }
5718        [f0, f1, f2, f3]
5719    }
5720}
5721
5722/// Which q4tp batch kernel to run: 1 = the previous one, 2 = the tuned
5723/// one, 0 = decide from the CPU. An atomic rather than a `OnceLock` so a
5724/// benchmark can alternate the two inside one process, where the machine's
5725/// mood — a shared box drifts ±25% between runs — is the same for both.
5726/// What the two mean is per-architecture: on x86 the blocked AVX-512 path
5727/// against the per-column one, on ARM the two reduction shapes.
5728#[allow(dead_code)]
5729static Q4TP_ALT: std::sync::atomic::AtomicU8 = std::sync::atomic::AtomicU8::new(0);
5730
5731/// Blocking pays on x86 only with 512-bit VNNI. With AVX2 alone, four
5732/// columns sharing an unpack still measured slower than the per-column
5733/// path (23.2 ms against 19.4 on a 48-thread EPYC), because that path
5734/// already dequantizes the row once — so the blocked kernel bought a
5735/// second unpack-free pass at the price of half the vector width.
5736#[cfg(target_arch = "x86_64")]
5737fn q4tp_blocked_x86() -> bool {
5738    match Q4TP_ALT.load(std::sync::atomic::Ordering::Relaxed) {
5739        1 => false,
5740        // A forced ON still asks the CPU. The switch exists so a bench can
5741        // pick a kernel, not so it can promise instructions the machine
5742        // does not have — CI caught that as a SIGILL on a runner without
5743        // AVX-512, where the parity test had turned the path on by hand.
5744        2 => avx512vnni_enabled(),
5745        // Deliberately not cached back into the switch: both gates below
5746        // hold their own `OnceLock`, and latching their answer here would
5747        // make a test's override outlive the test that set it.
5748        _ => blocked_enabled() && avx512vnni_enabled(),
5749    }
5750}
5751
5752/// `CMF_Q4TP_V1=1` picks the old reduction shape (A/B only).
5753#[cfg(target_arch = "aarch64")]
5754#[allow(dead_code)]
5755fn q4tp_v1() -> bool {
5756    match Q4TP_ALT.load(std::sync::atomic::Ordering::Relaxed) {
5757        1 => true,
5758        2 => false,
5759        _ => {
5760            static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
5761            *ON.get_or_init(|| std::env::var("CMF_Q4TP_V1").is_ok_and(|v| v != "0"))
5762        }
5763    }
5764}
5765
5766/// Two weight rows against eight columns. The activation load is the
5767/// same for both rows, so it is paid once for twice the arithmetic, and
5768/// sixteen accumulator chains run where eight did — which is what a kernel
5769/// retiring 0.29 instructions a cycle is short of. Register pressure is
5770/// the limit: sixteen `zmm` accumulators, two weight tiles, one
5771/// activation, of thirty-two.
5772///
5773/// Four rows by four columns spends the same sixteen accumulators the
5774/// other way and measured worse — 1488 GFLOP/s against 1644 — so the
5775/// unpack, which four rows pay twice as often, costs more than the extra
5776/// sharing of one activation load buys.
5777#[cfg(target_arch = "x86_64")]
5778#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
5779unsafe fn dot_q4tp_2x8_avx512(
5780    nib: &[u8],
5781    r0: usize,
5782    gpr: usize,
5783    xs: [&[i8]; 8],
5784    sc0: &[f32],
5785    sc1: &[f32],
5786) -> [[f32; 8]; 2] {
5787    // SAFETY: as dot_q4tp_row_1x8_avx512, two adjacent rows at once; the
5788    // caller guarantees r0 + 1 < rows and the ISA.
5789    unsafe {
5790        use core::arch::x86_64::*;
5791        let lomask = _mm256_set1_epi8(0x0F);
5792        let eight = _mm256_set1_epi8(8);
5793        let zero = _mm512_setzero_si512();
5794        let mut v0 = [_mm512_setzero_ps(); 8];
5795        let mut v1 = [_mm512_setzero_ps(); 8];
5796        let pairs = gpr / 2;
5797        let unpack = |r: usize, gi: usize| -> (__m512i, __mmask64) {
5798            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5799            let bb = _mm256_loadu_si256(t as *const __m256i);
5800            let lo = _mm256_and_si256(bb, lomask);
5801            let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
5802            let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
5803            let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
5804            let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
5805            let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
5806            (_mm512_abs_epi8(w), _mm512_movepi8_mask(w))
5807        };
5808        for gp in 0..pairs {
5809            let gi = gp * 2;
5810            let (wa0, neg0) = unpack(r0, gi);
5811            let (wa1, neg1) = unpack(r0 + 1, gi);
5812            let off = gi * GROUP_SIZE;
5813            let sv = |sc: &[f32]| {
5814                _mm512_insertf32x8::<1>(
5815                    _mm512_castps256_ps512(_mm256_set1_ps(*sc.get_unchecked(gi))),
5816                    _mm256_set1_ps(*sc.get_unchecked(gi + 1)),
5817                )
5818            };
5819            let s0 = sv(sc0);
5820            let s1 = sv(sc1);
5821            for k in 0..8 {
5822                let xv = _mm512_loadu_si512(xs[k].as_ptr().add(off) as *const __m512i);
5823                let d0 = _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(
5824                    zero,
5825                    wa0,
5826                    _mm512_mask_sub_epi8(xv, neg0, zero, xv),
5827                ));
5828                let d1 = _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(
5829                    zero,
5830                    wa1,
5831                    _mm512_mask_sub_epi8(xv, neg1, zero, xv),
5832                ));
5833                v0[k] = _mm512_fmadd_ps(d0, s0, v0[k]);
5834                v1[k] = _mm512_fmadd_ps(d1, s1, v1[k]);
5835            }
5836        }
5837        let mut acc = [[0f32; 8]; 2];
5838        for k in 0..8 {
5839            acc[0][k] = _mm512_reduce_add_ps(v0[k]);
5840            acc[1][k] = _mm512_reduce_add_ps(v1[k]);
5841        }
5842        if gpr % 2 == 1 {
5843            let off = (gpr - 1) * GROUP_SIZE;
5844            for j in off..off + GROUP_SIZE {
5845                let (w0, sa) = q4tp_outlier(nib, r0, gpr, j, sc0);
5846                let (w1, sb) = q4tp_outlier(nib, r0 + 1, gpr, j, sc1);
5847                for k in 0..8 {
5848                    let x = *xs[k].get_unchecked(j) as f32;
5849                    acc[0][k] += w0 * sa * x;
5850                    acc[1][k] += w1 * sb * x;
5851                }
5852            }
5853        }
5854        acc
5855    }
5856}
5857
5858/// The same, eight columns at a time. One unpack then feeds twice as many
5859/// activation streams, so a wide batch reads the weight tile half as
5860/// often; the price is eight accumulators live at once. Measured 9.0 ->
5861/// 8.3 ms at 9216x2304, b=296 on a 48-thread EPYC 9B45.
5862#[cfg(target_arch = "x86_64")]
5863#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
5864unsafe fn dot_q4tp_row_1x8_avx512(
5865    nib: &[u8],
5866    r: usize,
5867    gpr: usize,
5868    xs: [&[i8]; 8],
5869    scales: &[f32],
5870) -> [f32; 8] {
5871    // SAFETY: as dot_q4tp_row_1x4_avx2; caller guarantees the ISA.
5872    unsafe {
5873        use core::arch::x86_64::*;
5874        let lomask = _mm256_set1_epi8(0x0F);
5875        let eight = _mm256_set1_epi8(8);
5876        let zero = _mm512_setzero_si512();
5877        let (mut v0, mut v1, mut v2, mut v3) = (
5878            _mm512_setzero_ps(),
5879            _mm512_setzero_ps(),
5880            _mm512_setzero_ps(),
5881            _mm512_setzero_ps(),
5882        );
5883        let (mut v4, mut v5, mut v6, mut v7) = (
5884            _mm512_setzero_ps(),
5885            _mm512_setzero_ps(),
5886            _mm512_setzero_ps(),
5887            _mm512_setzero_ps(),
5888        );
5889        let pairs = gpr / 2;
5890        for gp in 0..pairs {
5891            let gi = gp * 2;
5892            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5893            let bb = _mm256_loadu_si256(t as *const __m256i);
5894            let lo = _mm256_and_si256(bb, lomask);
5895            let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
5896            // `unpack` works per 128-bit lane, so the halves come out as
5897            // [A.lo, B.lo] and [A.hi, B.hi]; the shuffle reorders the four
5898            // 128-bit lanes into the weights' natural order, which is what
5899            // the straight activation load expects.
5900            let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
5901            let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
5902            let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
5903            let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
5904            let wabs = _mm512_abs_epi8(w);
5905            let neg = _mm512_movepi8_mask(w);
5906            let off = gi * GROUP_SIZE;
5907            let sv = _mm512_insertf32x8::<1>(
5908                _mm512_castps256_ps512(_mm256_set1_ps(*scales.get_unchecked(gi))),
5909                _mm256_set1_ps(*scales.get_unchecked(gi + 1)),
5910            );
5911            let dot = |x: &[i8]| -> __m512 {
5912                let xv = _mm512_loadu_si512(x.as_ptr().add(off) as *const __m512i);
5913                let sx = _mm512_mask_sub_epi8(xv, neg, zero, xv);
5914                _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(zero, wabs, sx))
5915            };
5916            v0 = _mm512_fmadd_ps(dot(xs[0]), sv, v0);
5917            v1 = _mm512_fmadd_ps(dot(xs[1]), sv, v1);
5918            v2 = _mm512_fmadd_ps(dot(xs[2]), sv, v2);
5919            v3 = _mm512_fmadd_ps(dot(xs[3]), sv, v3);
5920            v4 = _mm512_fmadd_ps(dot(xs[4]), sv, v4);
5921            v5 = _mm512_fmadd_ps(dot(xs[5]), sv, v5);
5922            v6 = _mm512_fmadd_ps(dot(xs[6]), sv, v6);
5923            v7 = _mm512_fmadd_ps(dot(xs[7]), sv, v7);
5924        }
5925        let mut acc = [
5926            _mm512_reduce_add_ps(v0),
5927            _mm512_reduce_add_ps(v1),
5928            _mm512_reduce_add_ps(v2),
5929            _mm512_reduce_add_ps(v3),
5930            _mm512_reduce_add_ps(v4),
5931            _mm512_reduce_add_ps(v5),
5932            _mm512_reduce_add_ps(v6),
5933            _mm512_reduce_add_ps(v7),
5934        ];
5935        // An odd group count leaves one group over; the narrow kernel
5936        // finishes it rather than the tail being a special case here.
5937        if gpr % 2 == 1 {
5938            let off = (gpr - 1) * GROUP_SIZE;
5939            for j in off..off + GROUP_SIZE {
5940                let (w, s) = q4tp_outlier(nib, r, gpr, j, scales);
5941                let ws = w * s;
5942                for k in 0..8 {
5943                    acc[k] += ws * *xs[k].get_unchecked(j) as f32;
5944                }
5945            }
5946        }
5947        acc
5948    }
5949}
5950
5951/// The same four columns, 512 bits wide. Two groups (64 weights) ride one
5952/// unpack and one `vpdpbusd`, where AVX2 needs two unpacks and four
5953/// `maddubs`/`madd` pairs — about 2.3x fewer instructions for the same
5954/// arithmetic. The two groups carry different scales, so the fma takes a
5955/// vector whose halves hold each group's scale rather than a broadcast.
5956///
5957/// There is no 512-bit `vpsignb`, so the activation's sign is applied by
5958/// negating under a mask taken from the weight's sign bits. That mask is
5959/// per-tile, so it is hoisted out of the column loop and the per-column
5960/// cost stays exactly one instruction, as with `sign_epi8`. Weights of
5961/// zero are not zeroed by the mask trick and do not need to be: their
5962/// magnitude is zero, so the product is.
5963#[cfg(target_arch = "x86_64")]
5964#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
5965unsafe fn dot_q4tp_row_1x4_avx512(
5966    nib: &[u8],
5967    r: usize,
5968    gpr: usize,
5969    xs: [&[i8]; 4],
5970    scales: &[f32],
5971) -> [f32; 4] {
5972    // SAFETY: as dot_q4tp_row_1x4_avx2; caller guarantees the ISA.
5973    unsafe {
5974        use core::arch::x86_64::*;
5975        let lomask = _mm256_set1_epi8(0x0F);
5976        let eight = _mm256_set1_epi8(8);
5977        let zero = _mm512_setzero_si512();
5978        let (mut v0, mut v1, mut v2, mut v3) = (
5979            _mm512_setzero_ps(),
5980            _mm512_setzero_ps(),
5981            _mm512_setzero_ps(),
5982            _mm512_setzero_ps(),
5983        );
5984        let pairs = gpr / 2;
5985        for gp in 0..pairs {
5986            let gi = gp * 2;
5987            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5988            let bb = _mm256_loadu_si256(t as *const __m256i);
5989            let lo = _mm256_and_si256(bb, lomask);
5990            let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
5991            // `unpack` works per 128-bit lane, so the halves come out as
5992            // [A.lo, B.lo] and [A.hi, B.hi]; the shuffle reorders the four
5993            // 128-bit lanes into the weights' natural order, which is what
5994            // the straight activation load expects.
5995            let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
5996            let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
5997            let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
5998            let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
5999            let wabs = _mm512_abs_epi8(w);
6000            let neg = _mm512_movepi8_mask(w);
6001            let off = gi * GROUP_SIZE;
6002            let sv = _mm512_insertf32x8::<1>(
6003                _mm512_castps256_ps512(_mm256_set1_ps(*scales.get_unchecked(gi))),
6004                _mm256_set1_ps(*scales.get_unchecked(gi + 1)),
6005            );
6006            let dot = |x: &[i8]| -> __m512 {
6007                let xv = _mm512_loadu_si512(x.as_ptr().add(off) as *const __m512i);
6008                let sx = _mm512_mask_sub_epi8(xv, neg, zero, xv);
6009                _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(zero, wabs, sx))
6010            };
6011            v0 = _mm512_fmadd_ps(dot(xs[0]), sv, v0);
6012            v1 = _mm512_fmadd_ps(dot(xs[1]), sv, v1);
6013            v2 = _mm512_fmadd_ps(dot(xs[2]), sv, v2);
6014            v3 = _mm512_fmadd_ps(dot(xs[3]), sv, v3);
6015        }
6016        let mut acc = [
6017            _mm512_reduce_add_ps(v0),
6018            _mm512_reduce_add_ps(v1),
6019            _mm512_reduce_add_ps(v2),
6020            _mm512_reduce_add_ps(v3),
6021        ];
6022        // An odd group count leaves one group over; the narrow kernel
6023        // finishes it rather than the tail being a special case here.
6024        if gpr % 2 == 1 {
6025            let off = (gpr - 1) * GROUP_SIZE;
6026            for j in off..off + GROUP_SIZE {
6027                let (w, s) = q4tp_outlier(nib, r, gpr, j, scales);
6028                let ws = w * s;
6029                for k in 0..4 {
6030                    acc[k] += ws * *xs[k].get_unchecked(j) as f32;
6031                }
6032            }
6033        }
6034        acc
6035    }
6036}
6037
6038/// Four batch columns against one q4tp row: the tile is unpacked ONCE and
6039/// spent on four activation streams, which is where a prefill batch stops
6040/// being weight-bandwidth-bound. Twin of `dot_q4t_row_1x4_sdot`.
6041#[cfg(target_arch = "aarch64")]
6042#[target_feature(enable = "neon,dotprod")]
6043unsafe fn dot_q4tp_row_1x4_sdot(
6044    nib: &[u8],
6045    r: usize,
6046    gpr: usize,
6047    xs: [&[i8]; 4],
6048    scales: &[f32],
6049) -> [f32; 4] {
6050    // SAFETY: see dot_q4tp_row_sdot; every xs[k] is gpr·GROUP_SIZE long.
6051    unsafe {
6052        use core::arch::aarch64::*;
6053        use core::arch::asm;
6054        let lomask = vdupq_n_u8(0x0F);
6055        let eight = vdupq_n_s8(8);
6056        // Named accumulators, NOT an array indexed by a loop variable: the
6057        // latter does not stay in registers (the same defect cost 2x in the
6058        // AVX2 q4t kernel and again in WGSL).
6059        //
6060        // They are VECTORS, and the horizontal add happens once at the end
6061        // instead of once per group per column. `vaddvq` is a cross-lane
6062        // reduction — with 72 groups and four columns the old shape paid
6063        // 288 of them per row, each one a dependency stall the pipeline
6064        // cannot hide, to save four float adds. The group's scale now
6065        // rides an fma into the lane accumulators, so the arithmetic per
6066        // group is one convert and one fma. Summation order changes (the
6067        // lanes carry independent partial sums), which is the same
6068        // round-off class the SDOT path already lives in — the strict
6069        // kernel (`CMF_SDOT=0`, what `cortiq ppl` runs) is unchanged and
6070        // stays the reference.
6071        let (mut v0, mut v1, mut v2, mut v3) = (
6072            vdupq_n_f32(0.0),
6073            vdupq_n_f32(0.0),
6074            vdupq_n_f32(0.0),
6075            vdupq_n_f32(0.0),
6076        );
6077        for gi in 0..gpr {
6078            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
6079            let s = *scales.get_unchecked(gi);
6080            let bb = vld1q_u8(t);
6081            let lo = vandq_u8(bb, lomask);
6082            let hi = vshrq_n_u8::<4>(bb);
6083            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
6084            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
6085            let off = gi * GROUP_SIZE;
6086            let dot4 = |x: &[i8]| -> int32x4_t {
6087                let x0 = vld1q_s8(x.as_ptr().add(off));
6088                let x1 = vld1q_s8(x.as_ptr().add(off + 16));
6089                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6090                asm!(
6091                    "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
6092                    "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
6093                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6094                    e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
6095                    options(pure, nomem, nostack),
6096                );
6097                vaddq_s32(a0, a1)
6098            };
6099            v0 = vfmaq_n_f32(v0, vcvtq_f32_s32(dot4(xs[0])), s);
6100            v1 = vfmaq_n_f32(v1, vcvtq_f32_s32(dot4(xs[1])), s);
6101            v2 = vfmaq_n_f32(v2, vcvtq_f32_s32(dot4(xs[2])), s);
6102            v3 = vfmaq_n_f32(v3, vcvtq_f32_s32(dot4(xs[3])), s);
6103        }
6104        [
6105            vaddvq_f32(v0),
6106            vaddvq_f32(v1),
6107            vaddvq_f32(v2),
6108            vaddvq_f32(v3),
6109        ]
6110    }
6111}
6112
6113/// Fused q4tp matmat — the same three arms `q4t_matmat` has. Shipping only
6114/// the scalar one made Nanbeige-3B decode at 1.2 tok/s against q4t's 5.9:
6115/// the format was fine, the missing arms were the whole regression.
6116fn q4tp_matmat(
6117    bytes: &[u8],
6118    xs_all: &[f32],
6119    b: usize,
6120    rows: usize,
6121    cols: usize,
6122    out: &mut [f32],
6123    pool: Option<&Pool>,
6124) {
6125    debug_assert_eq!(out.len(), b * rows);
6126    let gpr = cols / GROUP_SIZE;
6127    let v = Q4tpView::new(bytes, rows, cols);
6128
6129    // Wide batches ride the AMX through a dequant-tile sgemm, as in q4t.
6130    #[cfg(target_os = "macos")]
6131    if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
6132        dequant_matmat_accel(
6133            &|r, dst| {
6134                let mut sc = [0f32; 32];
6135                let mut scv;
6136                let s: &[f32] = if gpr <= 32 {
6137                    v.scales_into(r, gpr, &mut sc);
6138                    &sc[..gpr]
6139                } else {
6140                    scv = vec![0f32; gpr];
6141                    v.scales_into(r, gpr, &mut scv);
6142                    &scv
6143                };
6144                for gi in 0..gpr {
6145                    let tile = &v.nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
6146                    for (k, &bb) in tile.iter().enumerate() {
6147                        dst[gi * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s[gi];
6148                        dst[gi * GROUP_SIZE + k * 2 + 1] =
6149                            (((bb >> 4) & 0x0F) as f32 - 8.0) * s[gi];
6150                    }
6151                }
6152            },
6153            xs_all,
6154            b,
6155            rows,
6156            cols,
6157            out,
6158            pool,
6159        );
6160        return;
6161    }
6162
6163    let out_addr = SendMut(out.as_mut_ptr());
6164    if a8w8_enabled() {
6165        let acts: Vec<SplitAct> = (0..b)
6166            .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
6167            .collect();
6168        let acts = &acts;
6169        #[cfg(target_arch = "aarch64")]
6170        let blocked_ok = sdot_enabled() && blocked_enabled();
6171        // x86 gets the same blocking: one tile unpack spent on four
6172        // columns. Without it every column re-decoded the row, which is
6173        // why a 48-core EPYC measured a sixth of an M4's per-core rate.
6174        // The gate is `avx2_enabled`, as in q4t — `sdot_enabled` answers
6175        // for ARM's dotprod and is hard-wired false everywhere else, so
6176        // asking it here left the whole blocked path unreachable on x86.
6177        #[cfg(target_arch = "x86_64")]
6178        let blocked_ok = q4tp_blocked_x86();
6179        #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
6180        let blocked_ok = false;
6181        // Columns are swept in panels that fit L2. Without this a
6182        // row-pair walks every activation in the batch — 4.8 MB at
6183        // 512x512 — and does it again for the next pair, so the whole
6184        // batch streams out of the shared cache once per row. Measured
6185        // 800 GB/s of it, flat across batch sizes, which is the signature
6186        // of a loop bound by traffic rather than by arithmetic. A panel of
6187        // 256 columns is 590 KB beside 221 KB of this worker's weights:
6188        // both stay resident and the batch crosses L3 once instead of
6189        // once per row.
6190        let panel_cols: usize = std::env::var("CMF_Q4TP_PANEL")
6191            .ok()
6192            .and_then(|v| v.parse().ok())
6193            .filter(|v| *v > 0)
6194            .unwrap_or(256);
6195        let run = |start: usize, end: usize| {
6196            for abase in (0..acts.len()).step_by(panel_cols) {
6197                let alen = (acts.len() - abase).min(panel_cols);
6198                let mut sc = vec![0f32; gpr];
6199                #[cfg(target_arch = "x86_64")]
6200                let mut r_lo = start;
6201                #[cfg(target_arch = "x86_64")]
6202                if blocked_ok && alen >= 8 {
6203                    let mut sc1 = vec![0f32; gpr];
6204                    while r_lo + 2 <= end {
6205                        v.scales_into(r_lo, gpr, &mut sc);
6206                        v.scales_into(r_lo + 1, gpr, &mut sc1);
6207                        let mut bi = 0usize;
6208                        while bi + 8 <= alen {
6209                            let xs = [
6210                                acts[abase + bi].xq.as_slice(),
6211                                acts[abase + bi + 1].xq.as_slice(),
6212                                acts[abase + bi + 2].xq.as_slice(),
6213                                acts[abase + bi + 3].xq.as_slice(),
6214                                acts[abase + bi + 4].xq.as_slice(),
6215                                acts[abase + bi + 5].xq.as_slice(),
6216                                acts[abase + bi + 6].xq.as_slice(),
6217                                acts[abase + bi + 7].xq.as_slice(),
6218                            ];
6219                            let d = unsafe { dot_q4tp_2x8_avx512(v.nib, r_lo, gpr, xs, &sc, &sc1) };
6220                            for (row, dr, scr) in [(r_lo, &d[0], &sc), (r_lo + 1, &d[1], &sc1)] {
6221                                for k in 0..8 {
6222                                    let act = &acts[abase + bi + k];
6223                                    let mut acc = dr[k] * act.sx;
6224                                    for &(j, xv) in &act.outliers {
6225                                        let (w, s) = q4tp_outlier(v.nib, row, gpr, j, scr);
6226                                        acc += w * s * xv;
6227                                    }
6228                                    // SAFETY: disjoint (bi, r) cells per worker.
6229                                    unsafe { *out_addr.at((abase + bi + k) * rows + row) = acc };
6230                                }
6231                            }
6232                            bi += 8;
6233                        }
6234                        // Columns past the last group of eight, both rows —
6235                        // the same single-row kernel the tail below uses.
6236                        for row in [r_lo, r_lo + 1] {
6237                            let scr: &[f32] = if row == r_lo { &sc } else { &sc1 };
6238                            for b2 in bi..alen {
6239                                let act = &acts[abase + b2];
6240                                let xs4 = [
6241                                    act.xq.as_slice(),
6242                                    act.xq.as_slice(),
6243                                    act.xq.as_slice(),
6244                                    act.xq.as_slice(),
6245                                ];
6246                                let d =
6247                                    unsafe { dot_q4tp_row_1x4_avx512(v.nib, row, gpr, xs4, scr) };
6248                                let mut acc = d[0] * act.sx;
6249                                for &(j, xv) in &act.outliers {
6250                                    let (w, s) = q4tp_outlier(v.nib, row, gpr, j, scr);
6251                                    acc += w * s * xv;
6252                                }
6253                                // SAFETY: disjoint (bi, r) cells per worker.
6254                                unsafe { *out_addr.at((abase + b2) * rows + row) = acc };
6255                            }
6256                        }
6257                        r_lo += 2;
6258                    }
6259                }
6260                #[cfg(target_arch = "x86_64")]
6261                let row_start = r_lo;
6262                #[cfg(not(target_arch = "x86_64"))]
6263                let row_start = start;
6264                for r in row_start..end {
6265                    v.scales_into(r, gpr, &mut sc);
6266                    let mut bi = 0usize;
6267                    #[cfg(target_arch = "x86_64")]
6268                    if blocked_ok {
6269                        while bi + 8 <= alen {
6270                            let xs = [
6271                                acts[abase + bi].xq.as_slice(),
6272                                acts[abase + bi + 1].xq.as_slice(),
6273                                acts[abase + bi + 2].xq.as_slice(),
6274                                acts[abase + bi + 3].xq.as_slice(),
6275                                acts[abase + bi + 4].xq.as_slice(),
6276                                acts[abase + bi + 5].xq.as_slice(),
6277                                acts[abase + bi + 6].xq.as_slice(),
6278                                acts[abase + bi + 7].xq.as_slice(),
6279                            ];
6280                            let d = unsafe { dot_q4tp_row_1x8_avx512(v.nib, r, gpr, xs, &sc) };
6281                            for k in 0..8 {
6282                                let act = &acts[abase + bi + k];
6283                                let mut acc = d[k] * act.sx;
6284                                for &(j, xv) in &act.outliers {
6285                                    let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
6286                                    acc += w * s * xv;
6287                                }
6288                                // SAFETY: disjoint (bi, r) cells per worker.
6289                                unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
6290                            }
6291                            bi += 8;
6292                        }
6293                        while bi + 4 <= alen {
6294                            let xs = [
6295                                acts[abase + bi].xq.as_slice(),
6296                                acts[abase + bi + 1].xq.as_slice(),
6297                                acts[abase + bi + 2].xq.as_slice(),
6298                                acts[abase + bi + 3].xq.as_slice(),
6299                            ];
6300                            let d = unsafe { dot_q4tp_row_1x4_avx512(v.nib, r, gpr, xs, &sc) };
6301                            for k in 0..4 {
6302                                let act = &acts[abase + bi + k];
6303                                let mut acc = d[k] * act.sx;
6304                                for &(j, xv) in &act.outliers {
6305                                    let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
6306                                    acc += w * s * xv;
6307                                }
6308                                // SAFETY: disjoint (bi, r) cells per worker.
6309                                unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
6310                            }
6311                            bi += 4;
6312                        }
6313                    }
6314                    #[cfg(target_arch = "aarch64")]
6315                    if blocked_ok {
6316                        while bi + 4 <= alen {
6317                            let xs = [
6318                                acts[abase + bi].xq.as_slice(),
6319                                acts[abase + bi + 1].xq.as_slice(),
6320                                acts[abase + bi + 2].xq.as_slice(),
6321                                acts[abase + bi + 3].xq.as_slice(),
6322                            ];
6323                            let d = unsafe {
6324                                if q4tp_v1() {
6325                                    dot_q4tp_row_1x4_sdot_v1(v.nib, r, gpr, xs, &sc)
6326                                } else {
6327                                    dot_q4tp_row_1x4_sdot(v.nib, r, gpr, xs, &sc)
6328                                }
6329                            };
6330                            for k in 0..4 {
6331                                let act = &acts[abase + bi + k];
6332                                let mut acc = d[k] * act.sx;
6333                                for &(j, xv) in &act.outliers {
6334                                    let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
6335                                    acc += w * s * xv;
6336                                }
6337                                // SAFETY: disjoint (bi, r) cells per worker.
6338                                unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
6339                            }
6340                            bi += 4;
6341                        }
6342                    }
6343                    let _ = blocked_ok;
6344                    while bi < alen {
6345                        let act = &acts[abase + bi];
6346                        let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, &sc) * act.sx;
6347                        for &(j, xv) in &act.outliers {
6348                            let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
6349                            acc += w * s * xv;
6350                        }
6351                        // SAFETY: disjoint (bi, r) cells per worker range.
6352                        unsafe { *out_addr.at((abase + bi) * rows + r) = acc };
6353                        bi += 1;
6354                    }
6355                }
6356            }
6357        };
6358        dispatch_rows(pool, rows, &run);
6359        return;
6360    }
6361
6362    let run = |start: usize, end: usize| {
6363        let mut sc = vec![0f32; gpr];
6364        for r in start..end {
6365            v.scales_into(r, gpr, &mut sc);
6366            for bi in 0..b {
6367                let x = &xs_all[bi * cols..(bi + 1) * cols];
6368                // SAFETY: disjoint (bi, r) cells per worker range.
6369                unsafe { *out_addr.at(bi * rows + r) = q4tp_row_exact(v.nib, r, gpr, x, &sc) };
6370            }
6371        }
6372    };
6373    dispatch_rows(pool, rows, &run);
6374}
6375
6376/// Fused q4_tiled matvec (dispatch mirrors `q4matvec`).
6377fn q4t_matvec(
6378    bytes: &[u8],
6379    x: &[f32],
6380    rows: usize,
6381    cols: usize,
6382    out: &mut [f32],
6383    pool: Option<&Pool>,
6384) {
6385    debug_assert_eq!(out.len(), rows);
6386    let gpr = cols / GROUP_SIZE;
6387    let out_addr = SendMut(out.as_mut_ptr());
6388    if a8w8_enabled() {
6389        let act = split_act(x);
6390        let run = move |start: usize, end: usize| {
6391            for r in start..end {
6392                let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
6393                for &(j, xv) in &act.outliers {
6394                    let (w, s) = q4t_outlier(bytes, r, gpr, j);
6395                    acc += w * s * xv;
6396                }
6397                // SAFETY: disjoint row ranges per worker.
6398                unsafe { *out_addr.at(r) = acc };
6399            }
6400        };
6401        dispatch_rows(pool, rows, &run);
6402        return;
6403    }
6404    let run = move |start: usize, end: usize| {
6405        for r in start..end {
6406            // SAFETY: disjoint row ranges per worker.
6407            unsafe { *out_addr.at(r) = q4t_row_exact(bytes, r, gpr, x) };
6408        }
6409    };
6410    dispatch_rows(pool, rows, &run);
6411}
6412
6413/// Fused two-input q4_tiled matvec (weights read once per pair).
6414#[allow(clippy::too_many_arguments)]
6415fn q4t_matvec2(
6416    bytes: &[u8],
6417    x1: &[f32],
6418    x2: &[f32],
6419    rows: usize,
6420    cols: usize,
6421    o1: &mut [f32],
6422    o2: &mut [f32],
6423    pool: Option<&Pool>,
6424) {
6425    let gpr = cols / GROUP_SIZE;
6426    let p1 = SendMut(o1.as_mut_ptr());
6427    let p2 = SendMut(o2.as_mut_ptr());
6428    if a8w8_enabled() {
6429        let a1 = split_act(x1);
6430        let a2 = split_act(x2);
6431        let run = move |start: usize, end: usize| {
6432            for r in start..end {
6433                let mut v1 = dot_q4t_row_i8(bytes, r, gpr, &a1.xq) * a1.sx;
6434                let mut v2 = dot_q4t_row_i8(bytes, r, gpr, &a2.xq) * a2.sx;
6435                for &(j, xv) in &a1.outliers {
6436                    let (w, s) = q4t_outlier(bytes, r, gpr, j);
6437                    v1 += w * s * xv;
6438                }
6439                for &(j, xv) in &a2.outliers {
6440                    let (w, s) = q4t_outlier(bytes, r, gpr, j);
6441                    v2 += w * s * xv;
6442                }
6443                // SAFETY: disjoint row ranges per worker.
6444                unsafe {
6445                    *p1.at(r) = v1;
6446                    *p2.at(r) = v2;
6447                }
6448            }
6449        };
6450        dispatch_rows(pool, rows, &run);
6451        return;
6452    }
6453    let run = move |start: usize, end: usize| {
6454        for r in start..end {
6455            // SAFETY: disjoint row ranges per worker.
6456            unsafe {
6457                *p1.at(r) = q4t_row_exact(bytes, r, gpr, x1);
6458                *p2.at(r) = q4t_row_exact(bytes, r, gpr, x2);
6459            }
6460        }
6461    };
6462    dispatch_rows(pool, rows, &run);
6463}
6464
6465/// Batched q4_tiled matmat: each row's tiles stream once per microbatch.
6466#[allow(clippy::too_many_arguments)]
6467/// Prefill GEMM through Accelerate for group-quantized codecs: a
6468/// caller-supplied row dequantizer fills f32 tiles (pool-parallel) and
6469/// each tile rides the AMX with one sgemm — the generic sibling of
6470/// `qmatmat_accel` (q8). Numerics are f32-GEMM (tolerance class);
6471/// decode (b=1) never takes this path.
6472#[cfg(target_os = "macos")]
6473fn dequant_matmat_accel(
6474    dequant_row: &(dyn Fn(usize, &mut [f32]) + Sync),
6475    xs_all: &[f32],
6476    b: usize,
6477    rows: usize,
6478    cols: usize,
6479    out: &mut [f32],
6480    pool: Option<&Pool>,
6481) {
6482    const TR: usize = 2048;
6483    thread_local! {
6484        static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
6485    }
6486    WTILE.with(|wt| {
6487        let mut wtile = wt.borrow_mut();
6488        wtile.resize(TR * cols, 0.0);
6489        let mut r0 = 0usize;
6490        while r0 < rows {
6491            let tr = TR.min(rows - r0);
6492            let wt_addr = SendMut(wtile.as_mut_ptr());
6493            let run = |start: usize, end: usize| {
6494                for r in start..end {
6495                    // SAFETY: workers cover disjoint r ranges.
6496                    let dst = unsafe { std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols) };
6497                    dequant_row(r0 + r, dst);
6498                }
6499            };
6500            dispatch_rows(pool, tr, &run);
6501            unsafe {
6502                accel_blas::cblas_sgemm(
6503                    101, // RowMajor
6504                    111, // NoTrans A
6505                    112, // Trans B
6506                    b as i32,
6507                    tr as i32,
6508                    cols as i32,
6509                    1.0,
6510                    xs_all.as_ptr(),
6511                    cols as i32,
6512                    wtile.as_ptr(),
6513                    cols as i32,
6514                    0.0,
6515                    out.as_mut_ptr().add(r0),
6516                    rows as i32,
6517                );
6518            }
6519            r0 += tr;
6520        }
6521    });
6522}
6523
6524fn q4t_matmat(
6525    bytes: &[u8],
6526    xs_all: &[f32],
6527    b: usize,
6528    rows: usize,
6529    cols: usize,
6530    out: &mut [f32],
6531    pool: Option<&Pool>,
6532) {
6533    debug_assert_eq!(out.len(), b * rows);
6534    let gpr = cols / GROUP_SIZE;
6535    // Wide batches ride the AMX like q8's qmatmat: on Apple silicon
6536    // the dequant-tile sgemm is an order above the SDOT row loop for
6537    // prefill shapes (imagegen DiT forwards are exactly this).
6538    #[cfg(target_os = "macos")]
6539    if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
6540        dequant_matmat_accel(
6541            &|r, dst| {
6542                for gi in 0..gpr {
6543                    let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
6544                    let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
6545                    for (k, &bb) in tile[2..].iter().enumerate() {
6546                        dst[gi * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s;
6547                        dst[gi * GROUP_SIZE + k * 2 + 1] = (((bb >> 4) & 0x0F) as f32 - 8.0) * s;
6548                    }
6549                }
6550            },
6551            xs_all,
6552            b,
6553            rows,
6554            cols,
6555            out,
6556            pool,
6557        );
6558        return;
6559    }
6560    let out_addr = SendMut(out.as_mut_ptr());
6561    if a8w8_enabled() {
6562        let acts: Vec<SplitAct> = (0..b)
6563            .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
6564            .collect();
6565        let acts = &acts;
6566        #[cfg(target_arch = "x86_64")]
6567        let blocked_ok = avx2_enabled() && blocked_enabled();
6568        #[cfg(target_arch = "aarch64")]
6569        let blocked_ok = sdot_enabled() && blocked_enabled();
6570        #[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))]
6571        let blocked_ok = false;
6572        let run = move |start: usize, end: usize| {
6573            for r in start..end {
6574                let mut bi = 0usize;
6575                #[cfg(target_arch = "aarch64")]
6576                if blocked_ok {
6577                    while bi + 4 <= acts.len() {
6578                        let xs = [
6579                            acts[bi].xq.as_slice(),
6580                            acts[bi + 1].xq.as_slice(),
6581                            acts[bi + 2].xq.as_slice(),
6582                            acts[bi + 3].xq.as_slice(),
6583                        ];
6584                        let d = unsafe { dot_q4t_row_1x4_sdot(bytes, r, gpr, xs) };
6585                        for k in 0..4 {
6586                            let act = &acts[bi + k];
6587                            let mut acc = d[k] * act.sx;
6588                            for &(j, xv) in &act.outliers {
6589                                let (w, sc) = q4t_outlier(bytes, r, gpr, j);
6590                                acc += w * sc * xv;
6591                            }
6592                            // SAFETY: disjoint (bi, r) cells per worker.
6593                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
6594                        }
6595                        bi += 4;
6596                    }
6597                }
6598                #[cfg(target_arch = "x86_64")]
6599                if blocked_ok {
6600                    while bi + 4 <= acts.len() {
6601                        let xs = [
6602                            acts[bi].xq.as_slice(),
6603                            acts[bi + 1].xq.as_slice(),
6604                            acts[bi + 2].xq.as_slice(),
6605                            acts[bi + 3].xq.as_slice(),
6606                        ];
6607                        let d = unsafe {
6608                            if vnni_tiles_enabled() {
6609                                dot_q4t_row_1x4_vnni(bytes, r, gpr, xs)
6610                            } else {
6611                                dot_q4t_row_1x4_avx2(bytes, r, gpr, xs)
6612                            }
6613                        };
6614                        for k in 0..4 {
6615                            let act = &acts[bi + k];
6616                            let mut acc = d[k] * act.sx;
6617                            for &(j, xv) in &act.outliers {
6618                                let (w, sc) = q4t_outlier(bytes, r, gpr, j);
6619                                acc += w * sc * xv;
6620                            }
6621                            // SAFETY: disjoint (bi, r) cells per worker.
6622                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
6623                        }
6624                        bi += 4;
6625                    }
6626                }
6627                let _ = blocked_ok;
6628                while bi < acts.len() {
6629                    let act = &acts[bi];
6630                    let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
6631                    for &(j, xv) in &act.outliers {
6632                        let (w, s) = q4t_outlier(bytes, r, gpr, j);
6633                        acc += w * s * xv;
6634                    }
6635                    // SAFETY: disjoint (bi, r) cells per worker range.
6636                    unsafe { *out_addr.at(bi * rows + r) = acc };
6637                    bi += 1;
6638                }
6639            }
6640        };
6641        dispatch_rows(pool, rows, &run);
6642        return;
6643    }
6644    let run = move |start: usize, end: usize| {
6645        for r in start..end {
6646            for bi in 0..b {
6647                let x = &xs_all[bi * cols..(bi + 1) * cols];
6648                // SAFETY: disjoint (bi, r) cells per worker range.
6649                unsafe { *out_addr.at(bi * rows + r) = q4t_row_exact(bytes, r, gpr, x) };
6650            }
6651        }
6652    };
6653    dispatch_rows(pool, rows, &run);
6654}
6655
6656// ── q1 (dtype 12): binary weights, [f16 scale][4B sign bits] per
6657// 32-group tile. The kernel family mirrors q4_tiled: one sequential
6658// stream of 6-byte tiles, per-tile integer dot × scale, exact outlier
6659// correction (A8W8 contract), exact scalar path under CMF_SDOT=0. ──
6660
6661/// Per-32-group sums of the quantized activation — the ±1 identity's
6662/// shared half: `dot = −2·sdot(mask, x) − gsum[g]`, computed ONCE per
6663/// matvec and reused by every row.
6664fn q1_group_sums(xq: &[i8], gpr: usize) -> Vec<i32> {
6665    (0..gpr)
6666        .map(|gi| {
6667            xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE]
6668                .iter()
6669                .map(|&v| v as i32)
6670                .sum()
6671        })
6672        .collect()
6673}
6674
6675/// One q1 row via the A8W8 int8 path — mask-SDOT on ARM (no ±1
6676/// expansion at all), scalar bit loop elsewhere (AVX2 queued with the
6677/// x86 pass).
6678#[inline]
6679#[allow(unreachable_code)]
6680/// AVX2 q1 row via the same ±1 identity as the ARM sdot kernel: the
6681/// sign bits expand to a {0, −1} byte mask through shuffle+cmpeq, the
6682/// masked activation sums through maddubs(1, x&mask), and
6683/// `dot = −(2·masked_sum + Σx_group)` — bit-identical integer math.
6684#[cfg(target_arch = "x86_64")]
6685#[target_feature(enable = "avx2")]
6686unsafe fn dot_q1_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
6687    // SAFETY: callers uphold the 6B-tile and xq/gsum length contracts.
6688    unsafe {
6689        use core::arch::x86_64::*;
6690        // Byte j of the mask must replicate bits-byte j/8.
6691        let expand = _mm256_setr_epi8(
6692            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,
6693            3, 3, 3,
6694        );
6695        let bitsel = _mm256_setr_epi8(
6696            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
6697            -128, 1, 2, 4, 8, 16, 32, 64, -128,
6698        );
6699        let ones8 = _mm256_set1_epi8(1);
6700        let ones16 = _mm256_set1_epi16(1);
6701        let mut acc = 0f32;
6702        for gi in 0..gpr {
6703            let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
6704            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
6705            let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
6706            let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
6707            let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
6708            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6709            let sel = _mm256_and_si256(x, mask);
6710            // Σ of selected i8 lanes: maddubs(1u8, sel_i8) pairs → madd.
6711            let p16 = _mm256_maddubs_epi16(ones8, sel);
6712            let d32 = _mm256_madd_epi16(p16, ones16);
6713            let hi128 = _mm256_extracti128_si256::<1>(d32);
6714            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
6715            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
6716            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
6717            let msum = _mm_cvtsi128_si32(s32);
6718            // The and-select keeps x UN-negated (unlike ARM's −1-mask
6719            // sdot): d = Σ_set − Σ_unset = 2·Σ_set − Σ_all.
6720            let d = 2 * msum - gsum[gi];
6721            acc += d as f32 * s;
6722        }
6723        acc
6724    }
6725}
6726
6727/// VNNI twin of `dot_q1_row_avx2`: the masked-select sum goes through
6728/// one `vpdpbusd(1u8, sel)` (see `dpbusd_hsum` — bit-identical).
6729#[cfg(target_arch = "x86_64")]
6730#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
6731unsafe fn dot_q1_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
6732    // SAFETY: callers uphold the 6B-tile and xq/gsum length contracts.
6733    unsafe {
6734        use core::arch::x86_64::*;
6735        let expand = _mm256_setr_epi8(
6736            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,
6737            3, 3, 3,
6738        );
6739        let bitsel = _mm256_setr_epi8(
6740            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
6741            -128, 1, 2, 4, 8, 16, 32, 64, -128,
6742        );
6743        let ones8 = _mm256_set1_epi8(1);
6744        let mut acc = 0f32;
6745        for gi in 0..gpr {
6746            let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
6747            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
6748            let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
6749            let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
6750            let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
6751            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6752            let msum = dpbusd_hsum(ones8, _mm256_and_si256(x, mask));
6753            let d = 2 * msum - gsum[gi];
6754            acc += d as f32 * s;
6755        }
6756        acc
6757    }
6758}
6759
6760/// VNNI twin of `dot_q1_row_1x4_avx2` (see `dpbusd_hsum`).
6761#[cfg(target_arch = "x86_64")]
6762#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
6763unsafe fn dot_q1_row_1x4_vnni(
6764    bytes: &[u8],
6765    r: usize,
6766    gpr: usize,
6767    xs: [&[i8]; 4],
6768    gsums: [&[i32]; 4],
6769) -> [f32; 4] {
6770    // SAFETY: callers uphold the 6B-tile and xq/gsum length contracts.
6771    unsafe {
6772        use core::arch::x86_64::*;
6773        let expand = _mm256_setr_epi8(
6774            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,
6775            3, 3, 3,
6776        );
6777        let bitsel = _mm256_setr_epi8(
6778            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
6779            -128, 1, 2, 4, 8, 16, 32, 64, -128,
6780        );
6781        let ones8 = _mm256_set1_epi8(1);
6782        let mut acc = [0f32; 4];
6783        for gi in 0..gpr {
6784            let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
6785            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
6786            let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
6787            let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
6788            let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
6789            for (k, xq) in xs.iter().enumerate() {
6790                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6791                let msum = dpbusd_hsum(ones8, _mm256_and_si256(x, mask));
6792                let d = 2 * msum - gsums[k][gi];
6793                acc[k] += d as f32 * s;
6794            }
6795        }
6796        acc
6797    }
6798}
6799
6800/// The blocked 1×4 flavor: the expanded bit mask serves four activation
6801/// streams per group (mask build once, four select+reduce chains).
6802#[cfg(target_arch = "x86_64")]
6803#[target_feature(enable = "avx2")]
6804unsafe fn dot_q1_row_1x4_avx2(
6805    bytes: &[u8],
6806    r: usize,
6807    gpr: usize,
6808    xs: [&[i8]; 4],
6809    gsums: [&[i32]; 4],
6810) -> [f32; 4] {
6811    // SAFETY: callers uphold the 6B-tile and xq/gsum length contracts.
6812    unsafe {
6813        use core::arch::x86_64::*;
6814        let expand = _mm256_setr_epi8(
6815            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,
6816            3, 3, 3,
6817        );
6818        let bitsel = _mm256_setr_epi8(
6819            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
6820            -128, 1, 2, 4, 8, 16, 32, 64, -128,
6821        );
6822        let ones8 = _mm256_set1_epi8(1);
6823        let ones16 = _mm256_set1_epi16(1);
6824        let mut acc = [0f32; 4];
6825        for gi in 0..gpr {
6826            let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
6827            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
6828            let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
6829            let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
6830            let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
6831            for (k, xq) in xs.iter().enumerate() {
6832                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6833                let sel = _mm256_and_si256(x, mask);
6834                let p16 = _mm256_maddubs_epi16(ones8, sel);
6835                let d32 = _mm256_madd_epi16(p16, ones16);
6836                let hi128 = _mm256_extracti128_si256::<1>(d32);
6837                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
6838                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
6839                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
6840                let msum = _mm_cvtsi128_si32(s32);
6841                let d = 2 * msum - gsums[k][gi];
6842                acc[k] += d as f32 * s;
6843            }
6844        }
6845        acc
6846    }
6847}
6848
6849#[allow(unreachable_code)]
6850fn dot_q1_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
6851    #[cfg(target_arch = "aarch64")]
6852    unsafe {
6853        return dot_q1_row_sdot(bytes, r, gpr, xq, gsum);
6854    }
6855    #[cfg(target_arch = "x86_64")]
6856    if avx2_enabled() {
6857        unsafe {
6858            if vnni_tiles_enabled() {
6859                return dot_q1_row_vnni(bytes, r, gpr, xq, gsum);
6860            }
6861            return dot_q1_row_avx2(bytes, r, gpr, xq, gsum);
6862        }
6863    }
6864    let _ = gsum;
6865    let mut acc = 0f32;
6866    for gi in 0..gpr {
6867        let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
6868        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
6869        let mut d = 0i32;
6870        for (j, &b) in tile[2..].iter().enumerate() {
6871            for k in 0..8 {
6872                let w = ((b >> k) & 1) as i32 * 2 - 1;
6873                d += w * xq[gi * GROUP_SIZE + j * 8 + k] as i32;
6874            }
6875        }
6876        acc += d as f32 * s;
6877    }
6878    acc
6879}
6880
6881/// SDOT q1 row via the ±1 identity: the vtst mask (0xFF where the bit
6882/// is set, i.e. −1 as i8) feeds `sdot` DIRECTLY — no expansion to ±1
6883/// lanes at all — and `dot = −(2·sdot(mask, x) + Σx_group)`, with the
6884/// per-group activation sums shared across every row of the matvec.
6885/// Four tiles (128 weights) per iteration: integer dots reduce through
6886/// a vpaddq tree into ONE i32x4 that meets its four scales in a single
6887/// fused f32 multiply-add. Integer math throughout — bit-identical to
6888/// the scalar ±1 reference.
6889#[cfg(target_arch = "aarch64")]
6890#[target_feature(enable = "neon,dotprod")]
6891unsafe fn dot_q1_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
6892    // SAFETY: callers uphold slice-length contracts (6B tile per group,
6893    // xq.len() == gpr·GROUP_SIZE, gsum.len() == gpr).
6894    unsafe {
6895        use core::arch::aarch64::*;
6896        use core::arch::asm;
6897        const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
6898        let m = vld1q_u8(MASKS.as_ptr());
6899        // One tile's −Σ_set(x) as an UNREDUCED i32x4 (two mask-sdots).
6900        macro_rules! tile_dot {
6901            ($t:expr, $x:expr) => {{
6902                let v0 = vcombine_u8(vdup_n_u8(*$t.add(2)), vdup_n_u8(*$t.add(3)));
6903                let v1 = vcombine_u8(vdup_n_u8(*$t.add(4)), vdup_n_u8(*$t.add(5)));
6904                let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
6905                let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
6906                let x0 = vld1q_s8($x);
6907                let x1 = vld1q_s8($x.add(16));
6908                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6909                asm!(
6910                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
6911                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
6912                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6913                    w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
6914                    options(pure, nomem, nostack),
6915                );
6916                vaddq_s32(a0, a1)
6917            }};
6918        }
6919        // TBL unpack over PAIR loads: one vld1q covers two 6B tiles
6920        // ([s s b b b b][s s b b b b] + 4B slack), TBL replicates each
6921        // bit-byte across 8 lanes for vtst, and the four scales gather
6922        // through tbl2 into one fcvtl — the 16 ld1r broadcast loads and
6923        // 4 branchy software f16 conversions per 128 weights (the
6924        // measured load-port wall of this kernel) become 2 vector
6925        // loads + 9 table lookups. Integer math order is unchanged —
6926        // bit-identical results (FCVTL is exact on every f16).
6927        const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
6928        const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
6929        const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
6930        const IW11: [u8; 16] = [
6931            10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
6932        ];
6933        const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
6934        let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
6935        let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
6936        let isc = vld1_u8(ISC.as_ptr());
6937        // One tile's −Σ_set(x) from a TBL-unpacked pair load.
6938        macro_rules! tile_dot_tbl {
6939            ($ld:expr, $i0:expr, $i1:expr, $x:expr) => {{
6940                let w0 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i0), m));
6941                let w1 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i1), m));
6942                let x0 = vld1q_s8($x);
6943                let x1 = vld1q_s8($x.add(16));
6944                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6945                asm!(
6946                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
6947                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
6948                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6949                    w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
6950                    options(pure, nomem, nostack),
6951                );
6952                vaddq_s32(a0, a1)
6953            }};
6954        }
6955        let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
6956        let row_base = r * gpr * Q1_TILE;
6957        let abs_end = bytes.len();
6958        let xp = xq.as_ptr();
6959        let gp = gsum.as_ptr();
6960        let mut accv = vdupq_n_f32(0.0);
6961        let mut gi = 0;
6962        // The second pair load reads 4B past tile gi+3 — stay inside
6963        // the payload slice (only the file's final tiles fall back).
6964        while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
6965            let t0 = base.add(gi * Q1_TILE);
6966            let ld_a = vld1q_u8(t0);
6967            let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
6968            let d0 = tile_dot_tbl!(ld_a, iw00, iw01, xp.add(gi * GROUP_SIZE));
6969            let d1 = tile_dot_tbl!(ld_a, iw10, iw11, xp.add((gi + 1) * GROUP_SIZE));
6970            let d2 = tile_dot_tbl!(ld_b, iw00, iw01, xp.add((gi + 2) * GROUP_SIZE));
6971            let d3 = tile_dot_tbl!(ld_b, iw10, iw11, xp.add((gi + 3) * GROUP_SIZE));
6972            // [−Σ0, −Σ1, −Σ2, −Σ3] → dots = −(2·Σset_neg + gsum)
6973            let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
6974            let g = vld1q_s32(gp.add(gi));
6975            let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
6976            let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
6977            let scf: float32x4_t;
6978            asm!(
6979                "fcvtl {o:v}.4s, {i:v}.4h",
6980                o = out(vreg) scf, i = in(vreg) sc16,
6981                options(pure, nomem, nostack),
6982            );
6983            accv = vfmaq_f32(accv, vcvtq_f32_s32(dots), scf);
6984            gi += 4;
6985        }
6986        let mut acc = vaddvq_f32(accv);
6987        while gi < gpr {
6988            let t = base.add(gi * Q1_TILE);
6989            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
6990            let d = vaddvq_s32(tile_dot!(t, xp.add(gi * GROUP_SIZE)));
6991            acc += (-(2 * d + *gp.add(gi))) as f32 * s;
6992            gi += 1;
6993        }
6994        acc
6995    }
6996}
6997
6998/// Blocked q1 1×4: one TBL unpack of the tile pair serves FOUR
6999/// activation streams (prefill amortization — the same idea as the
7000/// AVX2 twin; per stream the group order, fma order and tail match the
7001/// single-row kernel exactly, so batch == matvec bit-for-bit).
7002#[cfg(target_arch = "aarch64")]
7003#[target_feature(enable = "neon,dotprod")]
7004unsafe fn dot_q1_row_1x4_sdot(
7005    bytes: &[u8],
7006    r: usize,
7007    gpr: usize,
7008    xs: [&[i8]; 4],
7009    gs: [&[i32]; 4],
7010) -> [f32; 4] {
7011    // SAFETY: same slice-length contracts as `dot_q1_row_sdot`, ×4.
7012    unsafe {
7013        use core::arch::aarch64::*;
7014        use core::arch::asm;
7015        const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
7016        const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
7017        const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
7018        const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
7019        const IW11: [u8; 16] = [
7020            10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
7021        ];
7022        const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
7023        let m = vld1q_u8(MASKS.as_ptr());
7024        let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
7025        let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
7026        let isc = vld1_u8(ISC.as_ptr());
7027        macro_rules! sdot2 {
7028            ($w0:expr, $w1:expr, $x:expr) => {{
7029                let x0 = vld1q_s8($x);
7030                let x1 = vld1q_s8($x.add(16));
7031                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
7032                asm!(
7033                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
7034                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
7035                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
7036                    w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
7037                    options(pure, nomem, nostack),
7038                );
7039                vaddq_s32(a0, a1)
7040            }};
7041        }
7042        let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
7043        let row_base = r * gpr * Q1_TILE;
7044        let abs_end = bytes.len();
7045        let mut accv = [vdupq_n_f32(0.0); 4];
7046        let mut gi = 0;
7047        while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
7048            let t0 = base.add(gi * Q1_TILE);
7049            let ld_a = vld1q_u8(t0);
7050            let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
7051            // Unpack ONCE — eight ±mask vectors serve all four streams.
7052            let w00 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw00), m));
7053            let w01 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw01), m));
7054            let w10 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw10), m));
7055            let w11 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw11), m));
7056            let w20 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw00), m));
7057            let w21 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw01), m));
7058            let w30 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw10), m));
7059            let w31 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw11), m));
7060            let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
7061            let scf: float32x4_t;
7062            asm!(
7063                "fcvtl {o:v}.4s, {i:v}.4h",
7064                o = out(vreg) scf, i = in(vreg) sc16,
7065                options(pure, nomem, nostack),
7066            );
7067            for k in 0..4 {
7068                let xp = xs[k].as_ptr();
7069                let d0 = sdot2!(w00, w01, xp.add(gi * GROUP_SIZE));
7070                let d1 = sdot2!(w10, w11, xp.add((gi + 1) * GROUP_SIZE));
7071                let d2 = sdot2!(w20, w21, xp.add((gi + 2) * GROUP_SIZE));
7072                let d3 = sdot2!(w30, w31, xp.add((gi + 3) * GROUP_SIZE));
7073                let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
7074                let g = vld1q_s32(gs[k].as_ptr().add(gi));
7075                let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
7076                accv[k] = vfmaq_f32(accv[k], vcvtq_f32_s32(dots), scf);
7077            }
7078            gi += 4;
7079        }
7080        let mut acc = [
7081            vaddvq_f32(accv[0]),
7082            vaddvq_f32(accv[1]),
7083            vaddvq_f32(accv[2]),
7084            vaddvq_f32(accv[3]),
7085        ];
7086        while gi < gpr {
7087            let t = base.add(gi * Q1_TILE);
7088            let sc = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
7089            let v0 = vcombine_u8(vdup_n_u8(*t.add(2)), vdup_n_u8(*t.add(3)));
7090            let v1 = vcombine_u8(vdup_n_u8(*t.add(4)), vdup_n_u8(*t.add(5)));
7091            let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
7092            let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
7093            for k in 0..4 {
7094                let d = vaddvq_s32(sdot2!(w0, w1, xs[k].as_ptr().add(gi * GROUP_SIZE)));
7095                acc[k] += (-(2 * d + *gs[k].as_ptr().add(gi))) as f32 * sc;
7096            }
7097            gi += 1;
7098        }
7099        acc
7100    }
7101}
7102
7103/// (weight ±1, scale) of one q1 element — the exact outlier term.
7104#[inline]
7105fn q1_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
7106    let gi = j / GROUP_SIZE;
7107    let k = j % GROUP_SIZE;
7108    let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
7109    let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
7110    let bit = (tile[2 + k / 8] >> (k % 8)) & 1;
7111    ((bit as i32 * 2 - 1) as f32, s)
7112}
7113
7114/// Exact scalar q1 row (CMF_SDOT=0 contract).
7115#[inline]
7116fn q1_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
7117    let mut acc = 0f32;
7118    for gi in 0..gpr {
7119        let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
7120        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
7121        let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
7122        let mut ga = 0f32;
7123        for (j, &b) in tile[2..].iter().enumerate() {
7124            for k in 0..8 {
7125                ga += (((b >> k) & 1) as f32 * 2.0 - 1.0) * xg[j * 8 + k];
7126            }
7127        }
7128        acc += ga * s;
7129    }
7130    acc
7131}
7132
7133/// One q1 row range via A8W8 (the body of `q1_matvec`'s hot loop,
7134/// extracted so multi-matrix jobs drive the same kernel).
7135#[allow(clippy::too_many_arguments)]
7136fn q1_range_a8w8(
7137    bytes: &[u8],
7138    gpr: usize,
7139    act: &SplitAct,
7140    gsum: &[i32],
7141    out: SendMut,
7142    start: usize,
7143    end: usize,
7144) {
7145    for r in start..end {
7146        let mut acc = dot_q1_row_i8(bytes, r, gpr, &act.xq, gsum) * act.sx;
7147        for &(j, xv) in &act.outliers {
7148            let (w, s) = q1_outlier(bytes, r, gpr, j);
7149            acc += w * s * xv;
7150        }
7151        // SAFETY: disjoint row ranges per worker.
7152        unsafe { *out.at(r) = acc };
7153    }
7154}
7155
7156/// Exact-scalar q1 row range (CMF_SDOT=0 contract).
7157fn q1_range_f32(bytes: &[u8], gpr: usize, x: &[f32], out: SendMut, start: usize, end: usize) {
7158    for r in start..end {
7159        // SAFETY: disjoint row ranges per worker.
7160        unsafe { *out.at(r) = q1_row_exact(bytes, r, gpr, x) };
7161    }
7162}
7163
7164/// q1t per-row overlay locator. After the base (`base_len`) come
7165/// `[u32 row_ptr[rows+1]]` then `[(u16 col, f16 val)]` grouped by row (row
7166/// `r`'s entries are `[row_ptr[r], row_ptr[r+1])`). Returns
7167/// `(row_ptr offset, entries offset, present)`.
7168fn q1t_overlay(bytes: &[u8], base_len: usize, rows: usize) -> (usize, usize, bool) {
7169    let entries = base_len + (rows + 1) * 4;
7170    (base_len, entries, entries <= bytes.len())
7171}
7172
7173/// Read `row_ptr[r]` from the overlay's prefix-sum table.
7174#[inline]
7175fn q1t_rowptr(bytes: &[u8], rp_off: usize, r: usize) -> usize {
7176    let o = rp_off + r * 4;
7177    u32::from_le_bytes([bytes[o], bytes[o + 1], bytes[o + 2], bytes[o + 3]]) as usize
7178}
7179
7180/// Byte → the 5 ternary signs it packs `{−1,0,+1}` as f32, precomputed so
7181/// decoding a q1t code is a table load, not the base-3 divide/modulo per
7182/// weight (division is ~20–40× the cost of a load). Built at compile time.
7183const SIGN5: [[f32; 5]; 256] = {
7184    let mut lut = [[0.0f32; 5]; 256];
7185    let pow3 = [1u16, 3, 9, 27, 81];
7186    let mut byte = 0usize;
7187    while byte < 256 {
7188        let mut i = 0usize;
7189        while i < 5 {
7190            let code = (byte as u16 / pow3[i]) % 3;
7191            lut[byte][i] = if code == 1 {
7192                1.0
7193            } else if code == 2 {
7194                -1.0
7195            } else {
7196                0.0
7197            };
7198            i += 1;
7199        }
7200        byte += 1;
7201    }
7202    lut
7203};
7204
7205/// Same table, as i8 signs — the operand for the int8 SDOT base kernel.
7206const SIGN5_I8: [[i8; 5]; 256] = {
7207    let mut lut = [[0i8; 5]; 256];
7208    let pow3 = [1u16, 3, 9, 27, 81];
7209    let mut byte = 0usize;
7210    while byte < 256 {
7211        let mut i = 0usize;
7212        while i < 5 {
7213            let code = (byte as u16 / pow3[i]) % 3;
7214            lut[byte][i] = if code == 1 {
7215                1
7216            } else if code == 2 {
7217                -1
7218            } else {
7219                0
7220            };
7221            i += 1;
7222        }
7223        byte += 1;
7224    }
7225    lut
7226};
7227
7228/// The same 5 i8 signs packed into a u64 (`[s0 s1 s2 s3 s4 0 0 0]`, LE) so the
7229/// group unpack is 7 unaligned u64 stores at offsets 0,5,10,…,30 instead of
7230/// six 5-byte copies + LUT indexing — each store's trailing zeros are fixed by
7231/// the next store, and the last one runs 6 B past the 32nd weight (the unpack
7232/// buffer is padded to 40). This is the decode/prefill hot inner op.
7233const SIGN5_U64: [u64; 256] = {
7234    let mut lut = [0u64; 256];
7235    let pow3 = [1u16, 3, 9, 27, 81];
7236    let mut byte = 0usize;
7237    while byte < 256 {
7238        let mut v = 0u64;
7239        let mut i = 0usize;
7240        while i < 5 {
7241            let code = (byte as u16 / pow3[i]) % 3;
7242            let s: u8 = if code == 1 {
7243                1
7244            } else if code == 2 {
7245                0xFF
7246            } else {
7247                0
7248            };
7249            v |= (s as u64) << (i * 8);
7250            i += 1;
7251        }
7252        lut[byte] = v;
7253        byte += 1;
7254    }
7255    lut
7256};
7257
7258/// Ternary base weight at `(row r, col j)` = `sign(code)·s_group`. Used to add
7259/// back activation-outlier columns, whose `x` was zeroed for the int8 bulk dot
7260/// (`split_act`). At a weight-outlier position the code is 0, so this is 0 and
7261/// the overlay correction owns that column — no double counting.
7262#[inline]
7263fn q1t_base_weight(bytes: &[u8], r: usize, gpr: usize, j: usize) -> f32 {
7264    const TILE: usize = cortiq_core::quant::Q1T_TILE;
7265    let off = (r * gpr + j / GROUP_SIZE) * TILE;
7266    let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
7267    let within = j % GROUP_SIZE;
7268    SIGN5[bytes[off + 2 + within / 5] as usize][within % 5] * s
7269}
7270
7271/// One 32-group int8 dot via two SDOTs. Bit-exact vs the scalar i8 sum
7272/// (integer accumulation is order-independent).
7273#[cfg(target_arch = "aarch64")]
7274#[target_feature(enable = "neon,dotprod")]
7275#[inline]
7276unsafe fn sdot32_i8(w: *const i8, x: *const i8) -> i32 {
7277    // SAFETY: caller guarantees 32 readable i8 at each pointer.
7278    unsafe {
7279        use core::arch::aarch64::*;
7280        use core::arch::asm;
7281        let w0 = vld1q_s8(w);
7282        let w1 = vld1q_s8(w.add(16));
7283        let x0 = vld1q_s8(x);
7284        let x1 = vld1q_s8(x.add(16));
7285        let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
7286        asm!(
7287            "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
7288            "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
7289            a0 = inout(vreg) a0, a1 = inout(vreg) a1,
7290            w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
7291            options(pure, nomem, nostack),
7292        );
7293        vaddvq_s32(vaddq_s32(a0, a1))
7294    }
7295}
7296
7297/// One 32-group int8 dot via AVX2: signed·signed as `maddubs(|w|, sign(x,w))`
7298/// then `madd` and a horizontal reduce (the same idiom as `dot_q4t_row_avx2`).
7299#[cfg(target_arch = "x86_64")]
7300#[target_feature(enable = "avx2")]
7301#[inline]
7302unsafe fn i8dot32_avx2(w: *const i8, x: *const i8) -> i32 {
7303    // SAFETY: caller guarantees 32 readable i8 at each pointer.
7304    unsafe {
7305        use core::arch::x86_64::*;
7306        let wv = _mm256_loadu_si256(w as *const __m256i);
7307        let xv = _mm256_loadu_si256(x as *const __m256i);
7308        let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
7309        let d = _mm256_madd_epi16(p16, _mm256_set1_epi16(1));
7310        let hi128 = _mm256_extracti128_si256::<1>(d);
7311        let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
7312        let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7313        let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7314        _mm_cvtsi128_si32(s32)
7315    }
7316}
7317
7318/// Unpack one q1t group's base-3 codes into 32 i8 signs via 7 unaligned u64
7319/// stores (see `SIGN5_U64`). `dst` MUST have ≥ 40 bytes: the 7th store writes
7320/// `dst[30..38]`. Stores go in order so each one's trailing zeros are
7321/// overwritten by the next; the final 6 padding bytes are unused by the dot.
7322#[inline]
7323fn q1t_unpack_group_i8(codes: *const u8, dst: &mut [i8]) {
7324    debug_assert!(dst.len() >= 40);
7325    // SAFETY: codes points at 7 readable bytes; dst has ≥ 40 bytes so every
7326    // 8-byte store at offset bi*5 (bi ≤ 6 → ≤ 30) stays in bounds.
7327    unsafe {
7328        let p = dst.as_mut_ptr();
7329        for bi in 0..7 {
7330            core::ptr::write_unaligned(
7331                p.add(bi * 5) as *mut u64,
7332                SIGN5_U64[*codes.add(bi) as usize],
7333            );
7334        }
7335    }
7336}
7337
7338/// One 32-group int8 dot, arch-dispatched (the matmat inner loop, where the
7339/// row's signs are unpacked once and dotted against every batch input).
7340/// Callers are gated by `a8w8_enabled()`, so the target-feature arms are
7341/// reachable; the scalar arm is a non-SIMD-arch fallback.
7342#[inline]
7343fn q1t_i8dot32(w: *const i8, x: *const i8) -> i32 {
7344    #[cfg(target_arch = "aarch64")]
7345    unsafe {
7346        return sdot32_i8(w, x);
7347    }
7348    #[cfg(target_arch = "x86_64")]
7349    unsafe {
7350        return i8dot32_avx2(w, x);
7351    }
7352    #[allow(unreachable_code)]
7353    unsafe {
7354        let mut s = 0i32;
7355        for k in 0..GROUP_SIZE {
7356            s += *w.add(k) as i32 * *x.add(k) as i32;
7357        }
7358        s
7359    }
7360}
7361
7362#[inline]
7363unsafe fn q1t_unpack_reg_u64s(codes: *const u8) -> (u64, u64, u64, u64) {
7364    let (s0, s1, s2, s3, s4, s5, s6) = unsafe {
7365        (
7366            SIGN5_U64[*codes as usize],
7367            SIGN5_U64[*codes.add(1) as usize],
7368            SIGN5_U64[*codes.add(2) as usize],
7369            SIGN5_U64[*codes.add(3) as usize],
7370            SIGN5_U64[*codes.add(4) as usize],
7371            SIGN5_U64[*codes.add(5) as usize],
7372            SIGN5_U64[*codes.add(6) as usize],
7373        )
7374    };
7375
7376    let u0 = s0 | (s1 << 40);
7377    let u1 = (s1 >> 24) | (s2 << 16) | (s3 << 56);
7378    let u2 = (s3 >> 8) | (s4 << 32);
7379    let u3 = (s4 >> 32) | (s5 << 8) | (s6 << 48);
7380
7381    (u0, u1, u2, u3)
7382}
7383
7384/// One q1t row's int8 base dot: `Σ_group s·dot(signs, xq)` (before the shared
7385/// `sx`). Direct register unpacking (zero stack stores/loads, no STLF stalls).
7386/// ARM SDOT.
7387#[cfg(target_arch = "aarch64")]
7388#[target_feature(enable = "neon,dotprod")]
7389unsafe fn q1t_dot_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
7390    use core::arch::aarch64::*;
7391    use core::arch::asm;
7392    unsafe {
7393        const TILE: usize = cortiq_core::quant::Q1T_TILE;
7394        let mut acc = 0f32;
7395        let bytes_ptr = bytes.as_ptr();
7396        let xq_ptr = xq.as_ptr();
7397        let row_off = r * gpr * TILE;
7398
7399        let gpr2 = gpr & !1;
7400        let mut gi = 0;
7401        while gi < gpr2 {
7402            let off0 = row_off + gi * TILE;
7403            let off1 = off0 + TILE;
7404            let s0 = f16_to_f32(u16::from_le_bytes([
7405                *bytes_ptr.add(off0),
7406                *bytes_ptr.add(off0 + 1),
7407            ]));
7408            let s1 = f16_to_f32(u16::from_le_bytes([
7409                *bytes_ptr.add(off1),
7410                *bytes_ptr.add(off1 + 1),
7411            ]));
7412
7413            let (u0_0, u1_0, u2_0, u3_0) = q1t_unpack_reg_u64s(bytes_ptr.add(off0 + 2));
7414            let (u0_1, u1_1, u2_1, u3_1) = q1t_unpack_reg_u64s(bytes_ptr.add(off1 + 2));
7415
7416            let w0_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_0), vcreate_u64(u1_0)));
7417            let w1_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_0), vcreate_u64(u3_0)));
7418            let w0_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_1), vcreate_u64(u1_1)));
7419            let w1_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_1), vcreate_u64(u3_1)));
7420
7421            let x0_0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE));
7422            let x1_0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE + 16));
7423            let x0_1 = vld1q_s8(xq_ptr.add((gi + 1) * GROUP_SIZE));
7424            let x1_1 = vld1q_s8(xq_ptr.add((gi + 1) * GROUP_SIZE + 16));
7425
7426            let (mut a0_0, mut a1_0) = (vdupq_n_s32(0), vdupq_n_s32(0));
7427            let (mut a0_1, mut a1_1) = (vdupq_n_s32(0), vdupq_n_s32(0));
7428            asm!(
7429                "sdot {a0_0:v}.4s, {w0_0:v}.16b, {x0_0:v}.16b",
7430                "sdot {a1_0:v}.4s, {w1_0:v}.16b, {x1_0:v}.16b",
7431                "sdot {a0_1:v}.4s, {w0_1:v}.16b, {x0_1:v}.16b",
7432                "sdot {a1_1:v}.4s, {w1_1:v}.16b, {x1_1:v}.16b",
7433                a0_0 = inout(vreg) a0_0, a1_0 = inout(vreg) a1_0,
7434                a0_1 = inout(vreg) a0_1, a1_1 = inout(vreg) a1_1,
7435                w0_0 = in(vreg) w0_0, x0_0 = in(vreg) x0_0, w1_0 = in(vreg) w1_0, x1_0 = in(vreg) x1_0,
7436                w0_1 = in(vreg) w0_1, x0_1 = in(vreg) x0_1, w1_1 = in(vreg) w1_1, x1_1 = in(vreg) x1_1,
7437                options(pure, nomem, nostack),
7438            );
7439            let d0 = vaddvq_s32(vaddq_s32(a0_0, a1_0));
7440            let d1 = vaddvq_s32(vaddq_s32(a0_1, a1_1));
7441            acc += d0 as f32 * s0 + d1 as f32 * s1;
7442            gi += 2;
7443        }
7444
7445        if gi < gpr {
7446            let off = row_off + gi * TILE;
7447            let s = f16_to_f32(u16::from_le_bytes([
7448                *bytes_ptr.add(off),
7449                *bytes_ptr.add(off + 1),
7450            ]));
7451            let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
7452            let w0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0), vcreate_u64(u1)));
7453            let w1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2), vcreate_u64(u3)));
7454            let x0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE));
7455            let x1 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE + 16));
7456            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
7457            asm!(
7458                "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
7459                "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
7460                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
7461                w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
7462                options(pure, nomem, nostack),
7463            );
7464            let d = vaddvq_s32(vaddq_s32(a0, a1));
7465            acc += d as f32 * s;
7466        }
7467        acc
7468    }
7469}
7470
7471/// x86 AVX2 mirror of `q1t_dot_row_sdot` (maddubs int8 dot per group).
7472#[cfg(target_arch = "x86_64")]
7473#[target_feature(enable = "avx2")]
7474unsafe fn q1t_dot_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
7475    use core::arch::x86_64::*;
7476    unsafe {
7477        const TILE: usize = cortiq_core::quant::Q1T_TILE;
7478        let mut acc = 0f32;
7479        let bytes_ptr = bytes.as_ptr();
7480        let xq_ptr = xq.as_ptr();
7481        let row_off = r * gpr * TILE;
7482
7483        let ones = _mm256_set1_epi16(1);
7484        for gi in 0..gpr {
7485            let off = row_off + gi * TILE;
7486            let s = f16_to_f32(u16::from_le_bytes([
7487                *bytes_ptr.add(off),
7488                *bytes_ptr.add(off + 1),
7489            ]));
7490            let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
7491            let wv = _mm256_set_epi64x(u3 as i64, u2 as i64, u1 as i64, u0 as i64);
7492            let xv = _mm256_loadu_si256(xq_ptr.add(gi * GROUP_SIZE) as *const __m256i);
7493            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
7494            let d256 = _mm256_madd_epi16(p16, ones);
7495            let d128 = _mm_add_epi32(
7496                _mm256_castsi256_si128(d256),
7497                _mm256_extracti128_si256(d256, 1),
7498            );
7499            let d64 = _mm_add_epi32(d128, _mm_shuffle_epi32(d128, 0xee));
7500            let d32 = _mm_cvtsi128_si32(_mm_add_epi32(d64, _mm_shuffle_epi32(d64, 0x55)));
7501            acc += d32 as f32 * s;
7502        }
7503        acc
7504    }
7505}
7506
7507/// VNNI twin of `q1t_dot_row_avx2` (see `dpbusd_hsum`).
7508#[cfg(target_arch = "x86_64")]
7509#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
7510unsafe fn q1t_dot_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
7511    use core::arch::x86_64::*;
7512    // SAFETY: same tile/xq contracts as `q1t_dot_row_avx2`.
7513    unsafe {
7514        const TILE: usize = cortiq_core::quant::Q1T_TILE;
7515        let mut acc = 0f32;
7516        let bytes_ptr = bytes.as_ptr();
7517        let xq_ptr = xq.as_ptr();
7518        let row_off = r * gpr * TILE;
7519        for gi in 0..gpr {
7520            let off = row_off + gi * TILE;
7521            let s = f16_to_f32(u16::from_le_bytes([
7522                *bytes_ptr.add(off),
7523                *bytes_ptr.add(off + 1),
7524            ]));
7525            let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
7526            let wv = _mm256_set_epi64x(u3 as i64, u2 as i64, u1 as i64, u0 as i64);
7527            let xv = _mm256_loadu_si256(xq_ptr.add(gi * GROUP_SIZE) as *const __m256i);
7528            let d = dpbusd_hsum(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
7529            acc += d as f32 * s;
7530        }
7531        acc
7532    }
7533}
7534
7535/// Per-row int8 base dot, dispatched once per row (matvec decode hot path).
7536/// Callers are gated by `a8w8_enabled()`, so the target-feature kernels are
7537/// reachable.
7538#[inline]
7539fn q1t_dot_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
7540    #[cfg(target_arch = "aarch64")]
7541    unsafe {
7542        return q1t_dot_row_sdot(bytes, r, gpr, xq);
7543    }
7544    #[cfg(target_arch = "x86_64")]
7545    unsafe {
7546        if vnni_tiles_enabled() {
7547            return q1t_dot_row_vnni(bytes, r, gpr, xq);
7548        }
7549        return q1t_dot_row_avx2(bytes, r, gpr, xq);
7550    }
7551    #[allow(unreachable_code)]
7552    {
7553        const TILE: usize = cortiq_core::quant::Q1T_TILE;
7554        let mut acc = 0f32;
7555        let mut sg = [0i8; GROUP_SIZE + 8]; // +8 slack for the u64-store unpack
7556        for gi in 0..gpr {
7557            let off = (r * gpr + gi) * TILE;
7558            let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
7559            q1t_unpack_group_i8(bytes.as_ptr().wrapping_add(off + 2), &mut sg);
7560            let mut d = 0i32;
7561            for k in 0..GROUP_SIZE {
7562                d += sg[k] as i32 * xq[gi * GROUP_SIZE + k] as i32;
7563            }
7564            acc += d as f32 * s;
7565        }
7566        acc
7567    }
7568}
7569
7570/// Σ over a row's outliers of `value·x[col]` — the correction that adds the
7571/// overlay's exact weights on top of the base dot. INVARIANT: the encoder
7572/// writes ternary code 0 at every outlier position (`quantize_q1t`), so the
7573/// base contributes nothing there and this is a plain `value·x`, not
7574/// `(value − base)·x` — no scattered per-outlier scale read. Row `r`'s entries
7575/// are the contiguous slice `[row_ptr[r], row_ptr[r+1])`, so no binary search.
7576fn q1t_row_outlier_correction(
7577    bytes: &[u8],
7578    r: usize,
7579    rp_off: usize,
7580    entries_off: usize,
7581    has_ov: bool,
7582    x: &[f32],
7583) -> f32 {
7584    if !has_ov {
7585        return 0.0;
7586    }
7587    let (c0, c1) = (
7588        q1t_rowptr(bytes, rp_off, r),
7589        q1t_rowptr(bytes, rp_off, r + 1),
7590    );
7591    let mut corr = 0f32;
7592    for p in c0..c1 {
7593        let e = entries_off + p * 4;
7594        let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
7595        let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
7596        corr += val * x[col];
7597    }
7598    corr
7599}
7600
7601/// Dequantize one q1t row into `buf[..cols]` via the sign LUT (no division),
7602/// then apply the row's outliers (its `[row_ptr[r], row_ptr[r+1])` slice).
7603/// Used by the batched (prefill) path where the decode amortizes over the batch.
7604fn q1t_dequant_row(
7605    bytes: &[u8],
7606    r: usize,
7607    gpr: usize,
7608    rp_off: usize,
7609    entries_off: usize,
7610    has_ov: bool,
7611    buf: &mut [f32],
7612) {
7613    const TILE: usize = cortiq_core::quant::Q1T_TILE;
7614    for g in 0..gpr {
7615        let off = (r * gpr + g) * TILE;
7616        let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
7617        let codes = &bytes[off + 2..off + TILE];
7618        let bc = g * GROUP_SIZE;
7619        // 6 full bytes (30 codes) + a 7th byte holding the last 2.
7620        for bi in 0..6 {
7621            let lut = &SIGN5[codes[bi] as usize];
7622            let d = &mut buf[bc + bi * 5..bc + bi * 5 + 5];
7623            for i in 0..5 {
7624                d[i] = lut[i] * s;
7625            }
7626        }
7627        let lut = &SIGN5[codes[6] as usize];
7628        buf[bc + 30] = lut[0] * s;
7629        buf[bc + 31] = lut[1] * s;
7630    }
7631    if !has_ov {
7632        return;
7633    }
7634    let (c0, c1) = (
7635        q1t_rowptr(bytes, rp_off, r),
7636        q1t_rowptr(bytes, rp_off, r + 1),
7637    );
7638    for p in c0..c1 {
7639        let e = entries_off + p * 4;
7640        let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
7641        buf[col] = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
7642    }
7643}
7644
7645/// Add the sparse outlier overlay onto a base dot already in `out` (the GPU
7646/// computes the ternary base; the overlay stays on the CPU — its entries are
7647/// few and its per-row gather doesn't vectorize on the GPU). Row-parallel.
7648fn q1t_add_overlay(
7649    bytes: &[u8],
7650    x: &[f32],
7651    rows: usize,
7652    cols: usize,
7653    out: &mut [f32],
7654    pool: Option<&Pool>,
7655) {
7656    const TILE: usize = cortiq_core::quant::Q1T_TILE;
7657    let gpr = cols / GROUP_SIZE;
7658    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
7659    if !has_ov {
7660        return;
7661    }
7662    let out_addr = SendMut(out.as_mut_ptr());
7663    let run = move |start: usize, end: usize| {
7664        for r in start..end {
7665            let corr = q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
7666            // SAFETY: disjoint rows; add onto the base the GPU already wrote.
7667            unsafe { *out_addr.at(r) += corr };
7668        }
7669    };
7670    dispatch_rows(pool, rows, &run);
7671}
7672
7673/// Q1T row range via the A8W8 int8 path — shared activation split,
7674/// per-row: base SDOT dot + outlier correction + overlay.
7675#[allow(clippy::too_many_arguments)]
7676fn q1t_range_a8w8(
7677    bytes: &[u8],
7678    gpr: usize,
7679    rp_off: usize,
7680    ent_off: usize,
7681    has_ov: bool,
7682    act: &SplitAct,
7683    x: &[f32],
7684    out: SendMut,
7685    start: usize,
7686    end: usize,
7687) {
7688    for r in start..end {
7689        let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
7690        for &(j, xv) in &act.outliers {
7691            acc += q1t_base_weight(bytes, r, gpr, j) * xv;
7692        }
7693        acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
7694        // SAFETY: disjoint row ranges per worker.
7695        unsafe { *out.at(r) = acc };
7696    }
7697}
7698
7699/// Q1T row range via the f32 path (no SDOT) — for matvec_many batched
7700/// dispatch when a8w8 is unavailable.
7701#[allow(clippy::too_many_arguments)]
7702fn q1t_range_f32_batch(
7703    bytes: &[u8],
7704    gpr: usize,
7705    rp_off: usize,
7706    ent_off: usize,
7707    has_ov: bool,
7708    x: &[f32],
7709    out: SendMut,
7710    start: usize,
7711    end: usize,
7712) {
7713    const TILE: usize = cortiq_core::quant::Q1T_TILE;
7714    let mut sg = [0f32; GROUP_SIZE];
7715    for r in start..end {
7716        let mut acc = 0f32;
7717        for g in 0..gpr {
7718            let off = (r * gpr + g) * TILE;
7719            let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
7720            let codes = &bytes[off + 2..off + TILE];
7721            let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
7722            for bi in 0..6 {
7723                sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
7724            }
7725            let lut = &SIGN5[codes[6] as usize];
7726            sg[30] = lut[0];
7727            sg[31] = lut[1];
7728            let mut gsum = 0f32;
7729            for k in 0..GROUP_SIZE {
7730                gsum += sg[k] * xg[k];
7731            }
7732            acc += s * gsum;
7733        }
7734        acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
7735        // SAFETY: disjoint row ranges per worker.
7736        unsafe { *out.at(r) = acc };
7737    }
7738}
7739
7740/// Ternary (q1t) matvec — decode+dot straight from mmap, one group at a time:
7741/// no per-ROW buffer, no division (the sign LUT), and a tiny per-group sign
7742/// buffer so the 32-wide dot vectorizes. This is the decode hot path.
7743fn q1t_matvec(
7744    bytes: &[u8],
7745    x: &[f32],
7746    rows: usize,
7747    cols: usize,
7748    out: &mut [f32],
7749    pool: Option<&Pool>,
7750) {
7751    debug_assert_eq!(out.len(), rows);
7752    const TILE: usize = cortiq_core::quant::Q1T_TILE;
7753    let gpr = cols / GROUP_SIZE;
7754    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
7755    let out_addr = SendMut(out.as_mut_ptr());
7756    // int8 SDOT base dot (ARM dotprod): ~4× the f32 arithmetic. x → i8 once
7757    // (`split_act`), activation outliers added back exactly in f32, weight
7758    // overlay on top. ARM SDOT / x86 AVX2; CMF_SDOT=0 keeps the exact f32 path.
7759    if a8w8_enabled() {
7760        let act = split_act(x);
7761        let act = &act;
7762        let run = move |start: usize, end: usize| {
7763            for r in start..end {
7764                let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
7765                for &(j, xv) in &act.outliers {
7766                    acc += q1t_base_weight(bytes, r, gpr, j) * xv;
7767                }
7768                acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
7769                // SAFETY: disjoint row ranges per worker.
7770                unsafe { *out_addr.at(r) = acc };
7771            }
7772        };
7773        dispatch_rows(pool, rows, &run);
7774        return;
7775    }
7776    let run = move |start: usize, end: usize| {
7777        // Per-group signs, unpacked contiguously so the dot below is a clean
7778        // 32-wide reduction the autovectorizer turns into f32x4 FMAs — the
7779        // 5-values-per-byte base-3 layout won't SIMD in place.
7780        let mut sg = [0f32; GROUP_SIZE];
7781        for r in start..end {
7782            let mut acc = 0f32;
7783            for g in 0..gpr {
7784                let off = (r * gpr + g) * TILE;
7785                let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
7786                let codes = &bytes[off + 2..off + TILE];
7787                let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
7788                for bi in 0..6 {
7789                    sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
7790                }
7791                let lut = &SIGN5[codes[6] as usize];
7792                sg[30] = lut[0];
7793                sg[31] = lut[1];
7794                let mut gsum = 0f32;
7795                for k in 0..GROUP_SIZE {
7796                    gsum += sg[k] * xg[k];
7797                }
7798                acc += s * gsum;
7799            }
7800            acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
7801            unsafe { *out_addr.at(r) = acc };
7802        }
7803    };
7804    dispatch_rows(pool, rows, &run);
7805}
7806
7807/// Fused-pair twin of `q1t_dot_row_sdot`: ONE register unpack of the
7808/// ternary codes serves BOTH activation streams (the unpack chain is
7809/// the dominant per-row cost — MTP verify pairs paid it twice). Per
7810/// stream the group order and f32 accumulation match the single-row
7811/// kernel exactly, so pair == 2×matvec bit-for-bit.
7812#[cfg(target_arch = "aarch64")]
7813#[target_feature(enable = "neon,dotprod")]
7814unsafe fn q1t_dot_row_sdot2(bytes: &[u8], r: usize, gpr: usize, xa: &[i8], xb: &[i8]) -> [f32; 2] {
7815    use core::arch::aarch64::*;
7816    use core::arch::asm;
7817    // SAFETY: same slice-length contracts as `q1t_dot_row_sdot`, ×2.
7818    unsafe {
7819        const TILE: usize = cortiq_core::quant::Q1T_TILE;
7820        let bytes_ptr = bytes.as_ptr();
7821        let row_off = r * gpr * TILE;
7822        let xp = [xa.as_ptr(), xb.as_ptr()];
7823        let mut acc = [0f32; 2];
7824        macro_rules! sdot2 {
7825            ($w0:expr, $w1:expr, $x:expr) => {{
7826                let x0 = vld1q_s8($x);
7827                let x1 = vld1q_s8($x.add(16));
7828                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
7829                asm!(
7830                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
7831                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
7832                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
7833                    w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
7834                    options(pure, nomem, nostack),
7835                );
7836                vaddvq_s32(vaddq_s32(a0, a1))
7837            }};
7838        }
7839        let gpr2 = gpr & !1;
7840        let mut gi = 0;
7841        while gi < gpr2 {
7842            let off0 = row_off + gi * TILE;
7843            let off1 = off0 + TILE;
7844            let s0 = f16_to_f32(u16::from_le_bytes([
7845                *bytes_ptr.add(off0),
7846                *bytes_ptr.add(off0 + 1),
7847            ]));
7848            let s1 = f16_to_f32(u16::from_le_bytes([
7849                *bytes_ptr.add(off1),
7850                *bytes_ptr.add(off1 + 1),
7851            ]));
7852            let (u0_0, u1_0, u2_0, u3_0) = q1t_unpack_reg_u64s(bytes_ptr.add(off0 + 2));
7853            let (u0_1, u1_1, u2_1, u3_1) = q1t_unpack_reg_u64s(bytes_ptr.add(off1 + 2));
7854            let w0_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_0), vcreate_u64(u1_0)));
7855            let w1_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_0), vcreate_u64(u3_0)));
7856            let w0_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_1), vcreate_u64(u1_1)));
7857            let w1_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_1), vcreate_u64(u3_1)));
7858            for k in 0..2 {
7859                let d0 = sdot2!(w0_0, w1_0, xp[k].add(gi * GROUP_SIZE));
7860                let d1 = sdot2!(w0_1, w1_1, xp[k].add((gi + 1) * GROUP_SIZE));
7861                acc[k] += d0 as f32 * s0 + d1 as f32 * s1;
7862            }
7863            gi += 2;
7864        }
7865        if gi < gpr {
7866            let off = row_off + gi * TILE;
7867            let s = f16_to_f32(u16::from_le_bytes([
7868                *bytes_ptr.add(off),
7869                *bytes_ptr.add(off + 1),
7870            ]));
7871            let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
7872            let w0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0), vcreate_u64(u1)));
7873            let w1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2), vcreate_u64(u3)));
7874            for k in 0..2 {
7875                let d = sdot2!(w0, w1, xp[k].add(gi * GROUP_SIZE));
7876                acc[k] += d as f32 * s;
7877            }
7878        }
7879        acc
7880    }
7881}
7882
7883/// Fused Q1T pair matvec: ONE pass over the rows serves both
7884/// activation streams — on ARM the ternary register unpack happens
7885/// once per tile pair (`q1t_dot_row_sdot2`); elsewhere the second dot
7886/// rides the row's L1-warm tile bytes. Per stream the math matches
7887/// `q1t_matvec` exactly.
7888fn q1t_matvec2(
7889    bytes: &[u8],
7890    x1: &[f32],
7891    x2: &[f32],
7892    rows: usize,
7893    cols: usize,
7894    o1: &mut [f32],
7895    o2: &mut [f32],
7896    pool: Option<&Pool>,
7897) {
7898    debug_assert_eq!(o1.len(), rows);
7899    debug_assert_eq!(o2.len(), rows);
7900    const TILE: usize = cortiq_core::quant::Q1T_TILE;
7901    let gpr = cols / GROUP_SIZE;
7902    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
7903    let out1 = SendMut(o1.as_mut_ptr());
7904    let out2 = SendMut(o2.as_mut_ptr());
7905    if a8w8_enabled() {
7906        let a1 = split_act(x1);
7907        let a2 = split_act(x2);
7908        let (a1, a2) = (&a1, &a2);
7909        let run = move |start: usize, end: usize| {
7910            for r in start..end {
7911                #[cfg(target_arch = "aarch64")]
7912                // a8w8 on aarch64 ⇔ sdot_enabled(), so the kernel's
7913                // target features are present.
7914                let ds = unsafe { q1t_dot_row_sdot2(bytes, r, gpr, &a1.xq, &a2.xq) };
7915                #[cfg(not(target_arch = "aarch64"))]
7916                let ds = [
7917                    q1t_dot_row_i8(bytes, r, gpr, &a1.xq),
7918                    q1t_dot_row_i8(bytes, r, gpr, &a2.xq),
7919                ];
7920                let mut acc1 = ds[0] * a1.sx;
7921                for &(j, xv) in &a1.outliers {
7922                    acc1 += q1t_base_weight(bytes, r, gpr, j) * xv;
7923                }
7924                acc1 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x1);
7925                let mut acc2 = ds[1] * a2.sx;
7926                for &(j, xv) in &a2.outliers {
7927                    acc2 += q1t_base_weight(bytes, r, gpr, j) * xv;
7928                }
7929                acc2 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x2);
7930                // SAFETY: disjoint row ranges per worker.
7931                unsafe {
7932                    *out1.at(r) = acc1;
7933                    *out2.at(r) = acc2;
7934                }
7935            }
7936        };
7937        dispatch_rows(pool, rows, &run);
7938        return;
7939    }
7940    let run = move |start: usize, end: usize| {
7941        // Exact path (CMF_SDOT=0): unpack the sign LUT once per group,
7942        // dot both streams — same op order per stream as `q1t_matvec`.
7943        let mut sg = [0f32; GROUP_SIZE];
7944        for r in start..end {
7945            let mut acc1 = 0f32;
7946            let mut acc2 = 0f32;
7947            for g in 0..gpr {
7948                let off = (r * gpr + g) * TILE;
7949                let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
7950                let codes = &bytes[off + 2..off + TILE];
7951                for bi in 0..6 {
7952                    sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
7953                }
7954                let lut = &SIGN5[codes[6] as usize];
7955                sg[30] = lut[0];
7956                sg[31] = lut[1];
7957                let xg1 = &x1[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
7958                let xg2 = &x2[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
7959                let mut gsum1 = 0f32;
7960                for k in 0..GROUP_SIZE {
7961                    gsum1 += sg[k] * xg1[k];
7962                }
7963                acc1 += s * gsum1;
7964                let mut gsum2 = 0f32;
7965                for k in 0..GROUP_SIZE {
7966                    gsum2 += sg[k] * xg2[k];
7967                }
7968                acc2 += s * gsum2;
7969            }
7970            acc1 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x1);
7971            acc2 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x2);
7972            // SAFETY: disjoint row ranges per worker.
7973            unsafe {
7974                *out1.at(r) = acc1;
7975                *out2.at(r) = acc2;
7976            }
7977        }
7978    };
7979    dispatch_rows(pool, rows, &run);
7980}
7981
7982/// Ternary (q1t) matmat (prefill) — dequant each row once, dot the whole
7983/// batch against it (amortizes the per-row decode).
7984fn q1t_matmat(
7985    bytes: &[u8],
7986    xs: &[f32],
7987    b: usize,
7988    rows: usize,
7989    cols: usize,
7990    out: &mut [f32],
7991    pool: Option<&Pool>,
7992) {
7993    debug_assert_eq!(out.len(), b * rows);
7994    const TILE: usize = cortiq_core::quant::Q1T_TILE;
7995    let gpr = cols / GROUP_SIZE;
7996    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
7997    let out_addr = SendMut(out.as_mut_ptr());
7998    // int8 prefill (ARM SDOT / x86 AVX2): quantize the B inputs once, unpack
7999    // each weight row's signs to i8 ONCE, then int8-dot against every input —
8000    // the row sign-decode amortizes over the whole batch. CMF_SDOT=0 → f32.
8001    if a8w8_enabled() {
8002        let acts: Vec<SplitAct> = (0..b)
8003            .map(|bi| split_act(&xs[bi * cols..(bi + 1) * cols]))
8004            .collect();
8005        let acts = &acts;
8006        let run = move |start: usize, end: usize| {
8007            let mut sg = vec![0i8; cols + 8]; // row signs, i8 (+8 unpack slack)
8008            let mut sc = vec![0f32; gpr]; // per-group scales
8009            let mut accs = vec![0f32; b]; // per-batch accumulators, reused per row
8010            for r in start..end {
8011                for g in 0..gpr {
8012                    let off = (r * gpr + g) * TILE;
8013                    sc[g] = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
8014                    q1t_unpack_group_i8(
8015                        bytes.as_ptr().wrapping_add(off + 2),
8016                        &mut sg[g * GROUP_SIZE..],
8017                    );
8018                }
8019                for bi in 0..b {
8020                    let act = &acts[bi];
8021                    let mut isum = 0f32;
8022                    for g in 0..gpr {
8023                        let d = q1t_i8dot32(
8024                            sg.as_ptr().wrapping_add(g * GROUP_SIZE),
8025                            act.xq.as_ptr().wrapping_add(g * GROUP_SIZE),
8026                        );
8027                        isum += d as f32 * sc[g];
8028                    }
8029                    let mut acc = isum * act.sx;
8030                    for &(j, xv) in &act.outliers {
8031                        acc += q1t_base_weight(bytes, r, gpr, j) * xv;
8032                    }
8033                    accs[bi] = acc;
8034                }
8035                // Overlay ONCE per row for the whole batch: read each (col, val)
8036                // from mmap a single time (was b× — the re-read dominated prefill)
8037                // and fan it out over the batch via the cached inputs.
8038                if has_ov {
8039                    let (c0, c1) = (
8040                        q1t_rowptr(bytes, rp_off, r),
8041                        q1t_rowptr(bytes, rp_off, r + 1),
8042                    );
8043                    for p in c0..c1 {
8044                        let e = ent_off + p * 4;
8045                        let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
8046                        let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
8047                        for bi in 0..b {
8048                            accs[bi] += val * xs[bi * cols + col];
8049                        }
8050                    }
8051                }
8052                for bi in 0..b {
8053                    unsafe { *out_addr.at(bi * rows + r) = accs[bi] };
8054                }
8055            }
8056        };
8057        dispatch_rows(pool, rows, &run);
8058        return;
8059    }
8060    let run = move |start: usize, end: usize| {
8061        let mut buf = vec![0f32; cols];
8062        for r in start..end {
8063            q1t_dequant_row(bytes, r, gpr, rp_off, ent_off, has_ov, &mut buf);
8064            for bi in 0..b {
8065                let xr = &xs[bi * cols..(bi + 1) * cols];
8066                let mut acc = 0f32;
8067                for j in 0..cols {
8068                    acc += buf[j] * xr[j];
8069                }
8070                unsafe { *out_addr.at(bi * rows + r) = acc };
8071            }
8072        }
8073    };
8074    dispatch_rows(pool, rows, &run);
8075}
8076
8077fn q1_matvec(
8078    bytes: &[u8],
8079    x: &[f32],
8080    rows: usize,
8081    cols: usize,
8082    out: &mut [f32],
8083    pool: Option<&Pool>,
8084) {
8085    debug_assert_eq!(out.len(), rows);
8086    let gpr = cols / GROUP_SIZE;
8087    let out_addr = SendMut(out.as_mut_ptr());
8088    if a8w8_enabled() {
8089        let act = split_act(x);
8090        let gsum = q1_group_sums(&act.xq, gpr);
8091        let (act, gsum) = (&act, &gsum);
8092        let run = move |start: usize, end: usize| {
8093            q1_range_a8w8(bytes, gpr, act, gsum, out_addr, start, end)
8094        };
8095        dispatch_rows(pool, rows, &run);
8096        return;
8097    }
8098    let run = move |start: usize, end: usize| q1_range_f32(bytes, gpr, x, out_addr, start, end);
8099    dispatch_rows(pool, rows, &run);
8100}
8101
8102/// Fused two-input q1 matvec (weights read once per pair).
8103#[allow(clippy::too_many_arguments)]
8104fn q1_matvec2(
8105    bytes: &[u8],
8106    x1: &[f32],
8107    x2: &[f32],
8108    rows: usize,
8109    cols: usize,
8110    o1: &mut [f32],
8111    o2: &mut [f32],
8112    pool: Option<&Pool>,
8113) {
8114    let gpr = cols / GROUP_SIZE;
8115    let p1 = SendMut(o1.as_mut_ptr());
8116    let p2 = SendMut(o2.as_mut_ptr());
8117    if a8w8_enabled() {
8118        let a1 = split_act(x1);
8119        let a2 = split_act(x2);
8120        let g1 = q1_group_sums(&a1.xq, gpr);
8121        let g2 = q1_group_sums(&a2.xq, gpr);
8122        let (a1, a2, g1, g2) = (&a1, &a2, &g1, &g2);
8123        let run = move |start: usize, end: usize| {
8124            for r in start..end {
8125                let mut v1 = dot_q1_row_i8(bytes, r, gpr, &a1.xq, g1) * a1.sx;
8126                let mut v2 = dot_q1_row_i8(bytes, r, gpr, &a2.xq, g2) * a2.sx;
8127                for &(j, xv) in &a1.outliers {
8128                    let (w, s) = q1_outlier(bytes, r, gpr, j);
8129                    v1 += w * s * xv;
8130                }
8131                for &(j, xv) in &a2.outliers {
8132                    let (w, s) = q1_outlier(bytes, r, gpr, j);
8133                    v2 += w * s * xv;
8134                }
8135                // SAFETY: disjoint row ranges per worker.
8136                unsafe {
8137                    *p1.at(r) = v1;
8138                    *p2.at(r) = v2;
8139                }
8140            }
8141        };
8142        dispatch_rows(pool, rows, &run);
8143        return;
8144    }
8145    let run = move |start: usize, end: usize| {
8146        for r in start..end {
8147            // SAFETY: disjoint row ranges per worker.
8148            unsafe {
8149                *p1.at(r) = q1_row_exact(bytes, r, gpr, x1);
8150                *p2.at(r) = q1_row_exact(bytes, r, gpr, x2);
8151            }
8152        }
8153    };
8154    dispatch_rows(pool, rows, &run);
8155}
8156
8157/// Batched q1 matmat: each row's tiles stream once per microbatch.
8158#[allow(clippy::too_many_arguments)]
8159fn q1_matmat(
8160    bytes: &[u8],
8161    xs_all: &[f32],
8162    b: usize,
8163    rows: usize,
8164    cols: usize,
8165    out: &mut [f32],
8166    pool: Option<&Pool>,
8167) {
8168    debug_assert_eq!(out.len(), b * rows);
8169    let gpr = cols / GROUP_SIZE;
8170    let out_addr = SendMut(out.as_mut_ptr());
8171    if a8w8_enabled() {
8172        let acts: Vec<(SplitAct, Vec<i32>)> = (0..b)
8173            .map(|bi| {
8174                let act = split_act(&xs_all[bi * cols..(bi + 1) * cols]);
8175                let gsum = q1_group_sums(&act.xq, gpr);
8176                (act, gsum)
8177            })
8178            .collect();
8179        let acts = &acts;
8180        #[cfg(target_arch = "x86_64")]
8181        let blocked_ok = avx2_enabled() && blocked_enabled();
8182        #[cfg(target_arch = "aarch64")]
8183        let blocked_ok = sdot_enabled() && blocked_enabled();
8184        let run = move |start: usize, end: usize| {
8185            for r in start..end {
8186                let mut bi = 0usize;
8187                // Blocked 1×4: the unpacked bit mask serves four
8188                // activation streams per group.
8189                #[cfg(target_arch = "aarch64")]
8190                if blocked_ok {
8191                    while bi + 4 <= acts.len() {
8192                        let xs = [
8193                            acts[bi].0.xq.as_slice(),
8194                            acts[bi + 1].0.xq.as_slice(),
8195                            acts[bi + 2].0.xq.as_slice(),
8196                            acts[bi + 3].0.xq.as_slice(),
8197                        ];
8198                        let gs = [
8199                            acts[bi].1.as_slice(),
8200                            acts[bi + 1].1.as_slice(),
8201                            acts[bi + 2].1.as_slice(),
8202                            acts[bi + 3].1.as_slice(),
8203                        ];
8204                        let d = unsafe { dot_q1_row_1x4_sdot(bytes, r, gpr, xs, gs) };
8205                        for k in 0..4 {
8206                            let (act, _) = &acts[bi + k];
8207                            let mut acc = d[k] * act.sx;
8208                            for &(j, xv) in &act.outliers {
8209                                let (w, sc) = q1_outlier(bytes, r, gpr, j);
8210                                acc += w * sc * xv;
8211                            }
8212                            // SAFETY: disjoint (bi, r) cells per worker.
8213                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
8214                        }
8215                        bi += 4;
8216                    }
8217                }
8218                #[cfg(target_arch = "x86_64")]
8219                if blocked_ok {
8220                    while bi + 4 <= acts.len() {
8221                        let xs = [
8222                            acts[bi].0.xq.as_slice(),
8223                            acts[bi + 1].0.xq.as_slice(),
8224                            acts[bi + 2].0.xq.as_slice(),
8225                            acts[bi + 3].0.xq.as_slice(),
8226                        ];
8227                        let gs = [
8228                            acts[bi].1.as_slice(),
8229                            acts[bi + 1].1.as_slice(),
8230                            acts[bi + 2].1.as_slice(),
8231                            acts[bi + 3].1.as_slice(),
8232                        ];
8233                        let d = unsafe {
8234                            if vnni_tiles_enabled() {
8235                                dot_q1_row_1x4_vnni(bytes, r, gpr, xs, gs)
8236                            } else {
8237                                dot_q1_row_1x4_avx2(bytes, r, gpr, xs, gs)
8238                            }
8239                        };
8240                        for k in 0..4 {
8241                            let (act, _) = &acts[bi + k];
8242                            let mut acc = d[k] * act.sx;
8243                            for &(j, xv) in &act.outliers {
8244                                let (w, sc) = q1_outlier(bytes, r, gpr, j);
8245                                acc += w * sc * xv;
8246                            }
8247                            // SAFETY: disjoint (bi, r) cells per worker.
8248                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
8249                        }
8250                        bi += 4;
8251                    }
8252                }
8253                while bi < acts.len() {
8254                    let (act, gsum) = &acts[bi];
8255                    let mut acc = dot_q1_row_i8(bytes, r, gpr, &act.xq, gsum) * act.sx;
8256                    for &(j, xv) in &act.outliers {
8257                        let (w, s) = q1_outlier(bytes, r, gpr, j);
8258                        acc += w * s * xv;
8259                    }
8260                    // SAFETY: disjoint (bi, r) cells per worker range.
8261                    unsafe { *out_addr.at(bi * rows + r) = acc };
8262                    bi += 1;
8263                }
8264            }
8265        };
8266        dispatch_rows(pool, rows, &run);
8267        return;
8268    }
8269    let run = move |start: usize, end: usize| {
8270        for r in start..end {
8271            for bi in 0..b {
8272                let x = &xs_all[bi * cols..(bi + 1) * cols];
8273                // SAFETY: disjoint (bi, r) cells per worker range.
8274                unsafe { *out_addr.at(bi * rows + r) = q1_row_exact(bytes, r, gpr, x) };
8275            }
8276        }
8277    };
8278    dispatch_rows(pool, rows, &run);
8279}
8280
8281/// Fused q4_block matvec straight from the mapped bytes. SDOT path when
8282/// dotprod is available (port of vmfcore `dot_q4_block_sdot`, measured
8283/// +23% on q4 decode): nibbles → centered i8, int8×int8 `sdot` per
8284/// 32-group, exact outlier correction — the same A8W8 contract as q8.
8285/// `CMF_SDOT=0` keeps the exact scalar path.
8286fn q4matvec(
8287    bytes: &[u8],
8288    x: &[f32],
8289    rows: usize,
8290    cols: usize,
8291    out: &mut [f32],
8292    pool: Option<&Pool>,
8293) {
8294    debug_assert_eq!(out.len(), rows);
8295    let (packed, scales) = q4_split(bytes, rows, cols);
8296    let gpr = cols / GROUP_SIZE;
8297    let out_addr = SendMut(out.as_mut_ptr());
8298
8299    if a8w8_enabled() {
8300        let act = split_act(x);
8301        let run = move |start: usize, end: usize| {
8302            q4_range_a8w8(packed, scales, gpr, cols, &act, out_addr, start, end)
8303        };
8304        dispatch_rows(pool, rows, &run);
8305        return;
8306    }
8307
8308    let run =
8309        move |start: usize, end: usize| q4_range_f32(packed, scales, gpr, x, out_addr, start, end);
8310    dispatch_rows(pool, rows, &run);
8311}
8312
8313/// One q4 row via the A8W8 int8 path — SDOT on ARM, AVX2 maddubs on
8314/// x86 (scalar fallback is unreachable: callers gate on a8w8_enabled).
8315#[inline]
8316#[allow(unreachable_code)]
8317/// One UNPACKED q4 row (centered i8 in `buf`) against four activation
8318/// streams: the 32-byte weight chunk and its abs() load once per group,
8319/// the per-group f16 scale decodes once — four maddubs+reduce chains
8320/// instead of four full (load, abs, dot) rounds.
8321#[cfg(target_arch = "x86_64")]
8322#[target_feature(enable = "avx2")]
8323unsafe fn dot_q4b_row_1x4_avx2(
8324    buf: &[u8],
8325    scales: &[u8],
8326    g0: usize,
8327    gpr: usize,
8328    xs: [&[i8]; 4],
8329) -> [f32; 4] {
8330    // SAFETY: callers uphold buffer contracts (buf.len() == gpr·32).
8331    unsafe {
8332        use core::arch::x86_64::*;
8333        let ones = _mm256_set1_epi16(1);
8334        let mut acc = [0f32; 4];
8335        for gi in 0..gpr {
8336            let s = f16_to_f32(u16::from_le_bytes([
8337                scales[(g0 + gi) * 2],
8338                scales[(g0 + gi) * 2 + 1],
8339            ]));
8340            let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8341            let aw = _mm256_abs_epi8(w);
8342            for (k, xq) in xs.iter().enumerate() {
8343                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8344                let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
8345                let d = _mm256_madd_epi16(p16, ones);
8346                let hi128 = _mm256_extracti128_si256::<1>(d);
8347                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
8348                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8349                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8350                acc[k] += _mm_cvtsi128_si32(s32) as f32 * s;
8351            }
8352        }
8353        acc
8354    }
8355}
8356
8357/// VNNI twin of `dot_q4b_row_1x4_avx2` (see `dpbusd_hsum`).
8358#[cfg(target_arch = "x86_64")]
8359#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
8360unsafe fn dot_q4b_row_1x4_vnni(
8361    buf: &[u8],
8362    scales: &[u8],
8363    g0: usize,
8364    gpr: usize,
8365    xs: [&[i8]; 4],
8366) -> [f32; 4] {
8367    // SAFETY: callers uphold buffer contracts (buf.len() == gpr·32).
8368    unsafe {
8369        use core::arch::x86_64::*;
8370        let mut acc = [0f32; 4];
8371        for gi in 0..gpr {
8372            let s = f16_to_f32(u16::from_le_bytes([
8373                scales[(g0 + gi) * 2],
8374                scales[(g0 + gi) * 2 + 1],
8375            ]));
8376            let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8377            let aw = _mm256_abs_epi8(w);
8378            for (k, xq) in xs.iter().enumerate() {
8379                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8380                let d = dpbusd_hsum(aw, _mm256_sign_epi8(x, w));
8381                acc[k] += d as f32 * s;
8382            }
8383        }
8384        acc
8385    }
8386}
8387
8388/// The vbit flavor of the blocked 1×4: the per-activation A8W8 scale
8389/// folds in PER GROUP as `(d·sx)·s` — bit-matching the single-matvec
8390/// accumulation order (the q4_block flavor applies sx once at the end,
8391/// matching ITS single path; the two conventions are historical and
8392/// each blocked leg must mirror its own).
8393#[cfg(target_arch = "x86_64")]
8394#[target_feature(enable = "avx2")]
8395unsafe fn dot_q4b_row_1x4_sx_avx2(
8396    buf: &[u8],
8397    scales: &[u8],
8398    g0: usize,
8399    gpr: usize,
8400    xs: [&[i8]; 4],
8401    sxs: [f32; 4],
8402) -> [f32; 4] {
8403    // SAFETY: callers uphold buffer contracts (buf.len() == gpr·32).
8404    unsafe {
8405        use core::arch::x86_64::*;
8406        let ones = _mm256_set1_epi16(1);
8407        let mut acc = [0f32; 4];
8408        for gi in 0..gpr {
8409            let s = f16_to_f32(u16::from_le_bytes([
8410                scales[(g0 + gi) * 2],
8411                scales[(g0 + gi) * 2 + 1],
8412            ]));
8413            let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8414            let aw = _mm256_abs_epi8(w);
8415            for (k, xq) in xs.iter().enumerate() {
8416                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8417                let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
8418                let d = _mm256_madd_epi16(p16, ones);
8419                let hi128 = _mm256_extracti128_si256::<1>(d);
8420                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
8421                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8422                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8423                acc[k] += (_mm_cvtsi128_si32(s32) as f32 * sxs[k]) * s;
8424            }
8425        }
8426        acc
8427    }
8428}
8429
8430/// VNNI twin of `dot_q4b_row_1x4_sx_avx2` (see `dpbusd_hsum`; the
8431/// per-group `(d·sx)·s` fold mirrors the vbit single path).
8432#[cfg(target_arch = "x86_64")]
8433#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
8434unsafe fn dot_q4b_row_1x4_sx_vnni(
8435    buf: &[u8],
8436    scales: &[u8],
8437    g0: usize,
8438    gpr: usize,
8439    xs: [&[i8]; 4],
8440    sxs: [f32; 4],
8441) -> [f32; 4] {
8442    // SAFETY: callers uphold buffer contracts (buf.len() == gpr·32).
8443    unsafe {
8444        use core::arch::x86_64::*;
8445        let mut acc = [0f32; 4];
8446        for gi in 0..gpr {
8447            let s = f16_to_f32(u16::from_le_bytes([
8448                scales[(g0 + gi) * 2],
8449                scales[(g0 + gi) * 2 + 1],
8450            ]));
8451            let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8452            let aw = _mm256_abs_epi8(w);
8453            for (k, xq) in xs.iter().enumerate() {
8454                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8455                let d = dpbusd_hsum(aw, _mm256_sign_epi8(x, w));
8456                acc[k] += (d as f32 * sxs[k]) * s;
8457            }
8458        }
8459        acc
8460    }
8461}
8462
8463#[allow(unreachable_code)]
8464fn dot_q4_row_i8(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
8465    #[cfg(target_arch = "aarch64")]
8466    unsafe {
8467        return dot_q4_row_sdot(packed, scales, g0, gpr, xq);
8468    }
8469    #[cfg(target_arch = "x86_64")]
8470    unsafe {
8471        return dot_q4_row_avx2(packed, scales, g0, gpr, xq);
8472    }
8473    let mut acc = 0f32;
8474    for gi in 0..gpr {
8475        let g = g0 + gi;
8476        let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
8477        let mut d = 0i32;
8478        for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
8479            d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
8480                + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
8481        }
8482        acc += d as f32 * s;
8483    }
8484    acc
8485}
8486
8487/// Two-activation q4 row via the A8W8 int8 path (see `dot_q4_row_i8`).
8488#[inline]
8489#[allow(unreachable_code)]
8490fn dot_q4_row_i8_2(
8491    packed: &[u8],
8492    scales: &[u8],
8493    g0: usize,
8494    gpr: usize,
8495    xq1: &[i8],
8496    xq2: &[i8],
8497) -> (f32, f32) {
8498    #[cfg(target_arch = "aarch64")]
8499    unsafe {
8500        return dot_q4_row_sdot2(packed, scales, g0, gpr, xq1, xq2);
8501    }
8502    #[cfg(target_arch = "x86_64")]
8503    unsafe {
8504        return dot_q4_row_avx2_2(packed, scales, g0, gpr, xq1, xq2);
8505    }
8506    (
8507        dot_q4_row_i8(packed, scales, g0, gpr, xq1),
8508        dot_q4_row_i8(packed, scales, g0, gpr, xq2),
8509    )
8510}
8511
8512/// One q4 row range via SDOT (kernel body of `q4matvec`, extracted so
8513/// multi-matrix jobs can drive it for several tensors in one dispatch).
8514#[allow(clippy::too_many_arguments)]
8515fn q4_range_a8w8(
8516    packed: &[u8],
8517    scales: &[u8],
8518    gpr: usize,
8519    cols: usize,
8520    act: &SplitAct,
8521    out: SendMut,
8522    start: usize,
8523    end: usize,
8524) {
8525    for r in start..end {
8526        let mut acc = dot_q4_row_i8(packed, scales, r * gpr, gpr, &act.xq) * act.sx;
8527        // xq is zeroed at outlier slots — add the exact terms.
8528        for &(j, xv) in &act.outliers {
8529            let flat = r * cols + j;
8530            let byte = packed[flat / 2];
8531            let nib = if flat & 1 == 0 {
8532                byte & 0x0F
8533            } else {
8534                byte >> 4
8535            };
8536            let s = f16_to_f32(u16::from_le_bytes([
8537                scales[(flat / GROUP_SIZE) * 2],
8538                scales[(flat / GROUP_SIZE) * 2 + 1],
8539            ]));
8540            acc += ((nib as i32 - 8) as f32) * s * xv;
8541        }
8542        // SAFETY: disjoint row ranges per worker.
8543        unsafe { *out.at(r) = acc };
8544    }
8545}
8546
8547/// Two-input q4 row range via the A8W8 int8 path — kernel body of
8548/// `q4matvec2`, extracted for pair multi-matrix jobs.
8549#[allow(clippy::too_many_arguments)]
8550fn q4_range2_a8w8(
8551    packed: &[u8],
8552    scales: &[u8],
8553    gpr: usize,
8554    cols: usize,
8555    a1: &SplitAct,
8556    a2: &SplitAct,
8557    p1: SendMut,
8558    p2: SendMut,
8559    start: usize,
8560    end: usize,
8561) {
8562    for r in start..end {
8563        let (s1, s2) = dot_q4_row_i8_2(packed, scales, r * gpr, gpr, &a1.xq, &a2.xq);
8564        let mut acc1 = s1 * a1.sx;
8565        let mut acc2 = s2 * a2.sx;
8566        // xq is zeroed at outlier slots — add the exact terms.
8567        let fix = |outliers: &[(usize, f32)], acc: &mut f32| {
8568            for &(j, xv) in outliers {
8569                let flat = r * cols + j;
8570                let byte = packed[flat / 2];
8571                let nib = if flat & 1 == 0 {
8572                    byte & 0x0F
8573                } else {
8574                    byte >> 4
8575                };
8576                let s = f16_to_f32(u16::from_le_bytes([
8577                    scales[(flat / GROUP_SIZE) * 2],
8578                    scales[(flat / GROUP_SIZE) * 2 + 1],
8579                ]));
8580                *acc += ((nib as i32 - 8) as f32) * s * xv;
8581            }
8582        };
8583        fix(&a1.outliers, &mut acc1);
8584        fix(&a2.outliers, &mut acc2);
8585        // SAFETY: disjoint row ranges per worker.
8586        unsafe {
8587            *p1.at(r) = acc1;
8588            *p2.at(r) = acc2;
8589        }
8590    }
8591}
8592
8593/// Exact scalar q4 row range (same extraction, non-SDOT path).
8594fn q4_range_f32(
8595    packed: &[u8],
8596    scales: &[u8],
8597    gpr: usize,
8598    x: &[f32],
8599    out: SendMut,
8600    start: usize,
8601    end: usize,
8602) {
8603    for r in start..end {
8604        let mut acc = 0f32;
8605        for gi in 0..gpr {
8606            let g = r * gpr + gi;
8607            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
8608            let pk = &packed[g * 16..(g + 1) * 16];
8609            let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
8610            let mut ga = 0f32;
8611            for (k, &b) in pk.iter().enumerate() {
8612                ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
8613                    + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
8614            }
8615            acc += ga * s;
8616        }
8617        // SAFETY: disjoint row ranges per worker.
8618        unsafe { *out.at(r) = acc };
8619    }
8620}
8621
8622/// Fused two-input q4 matvec: nibbles are unpacked ONCE per group and
8623/// dotted against both activations (was: two full matvecs — double
8624/// weight traffic). Per-lane math matches `q4matvec` exactly.
8625#[allow(clippy::too_many_arguments)]
8626fn q4matvec2(
8627    bytes: &[u8],
8628    x1: &[f32],
8629    x2: &[f32],
8630    rows: usize,
8631    cols: usize,
8632    o1: &mut [f32],
8633    o2: &mut [f32],
8634    pool: Option<&Pool>,
8635) {
8636    debug_assert_eq!(o1.len(), rows);
8637    debug_assert_eq!(o2.len(), rows);
8638    let (packed, scales) = q4_split(bytes, rows, cols);
8639    let gpr = cols / GROUP_SIZE;
8640
8641    if a8w8_enabled() {
8642        let a1 = split_act(x1);
8643        let a2 = split_act(x2);
8644        let p1 = SendMut(o1.as_mut_ptr());
8645        let p2 = SendMut(o2.as_mut_ptr());
8646        let run = move |start: usize, end: usize| {
8647            q4_range2_a8w8(packed, scales, gpr, cols, &a1, &a2, p1, p2, start, end)
8648        };
8649        dispatch_rows(pool, rows, &run);
8650        return;
8651    }
8652
8653    let p1 = SendMut(o1.as_mut_ptr());
8654    let p2 = SendMut(o2.as_mut_ptr());
8655    let run = move |start: usize, end: usize| {
8656        q4_range2_f32(packed, scales, gpr, x1, x2, p1, p2, start, end)
8657    };
8658    dispatch_rows(pool, rows, &run);
8659}
8660
8661/// Two-input exact scalar q4 row range (same extraction).
8662#[allow(clippy::too_many_arguments)]
8663fn q4_range2_f32(
8664    packed: &[u8],
8665    scales: &[u8],
8666    gpr: usize,
8667    x1: &[f32],
8668    x2: &[f32],
8669    p1: SendMut,
8670    p2: SendMut,
8671    start: usize,
8672    end: usize,
8673) {
8674    for r in start..end {
8675        let (mut acc1, mut acc2) = (0f32, 0f32);
8676        for gi in 0..gpr {
8677            let g = r * gpr + gi;
8678            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
8679            let pk = &packed[g * 16..(g + 1) * 16];
8680            let x1g = &x1[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
8681            let x2g = &x2[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
8682            let (mut g1, mut g2) = (0f32, 0f32);
8683            for (k, &b) in pk.iter().enumerate() {
8684                let wl = (b & 0x0F) as f32 - 8.0;
8685                let wh = ((b >> 4) & 0x0F) as f32 - 8.0;
8686                g1 += wl * x1g[k * 2] + wh * x1g[k * 2 + 1];
8687                g2 += wl * x2g[k * 2] + wh * x2g[k * 2 + 1];
8688            }
8689            acc1 += g1 * s;
8690            acc2 += g2 * s;
8691        }
8692        // SAFETY: disjoint row ranges per worker.
8693        unsafe {
8694            *p1.at(r) = acc1;
8695            *p2.at(r) = acc2;
8696        }
8697    }
8698}
8699
8700thread_local! {
8701    /// Per-worker decoded-row scratch for the batched q4/vbit kernels
8702    /// (centered i8 for SDOT, f32 for the exact/scalar paths).
8703    static ROW_I8: std::cell::RefCell<Vec<u8>> = const { std::cell::RefCell::new(Vec::new()) };
8704    static ROW_F32: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
8705}
8706
8707/// Batched q4 matmat: each weight row is unpacked from the mmap ONCE
8708/// and dotted against ALL b activations (prefill used to fall back to b
8709/// full matvecs — b× weight traffic and b× nibble decode). Per-position
8710/// math matches `q4matvec` exactly: same group order, same accumulation.
8711/// `out` is row-major [b, rows] like `qmatmat`.
8712#[allow(clippy::too_many_arguments)]
8713fn q4matmat(
8714    bytes: &[u8],
8715    xs_all: &[f32],
8716    b: usize,
8717    rows: usize,
8718    cols: usize,
8719    out: &mut [f32],
8720    pool: Option<&Pool>,
8721) {
8722    debug_assert_eq!(xs_all.len(), b * cols);
8723    debug_assert_eq!(out.len(), b * rows);
8724    let (packed, scales) = q4_split(bytes, rows, cols);
8725    let gpr = cols / GROUP_SIZE;
8726    let gscale = |g: usize| f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
8727
8728    if a8w8_enabled() {
8729        let acts: Vec<SplitAct> = (0..b)
8730            .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
8731            .collect();
8732        let acts = &acts;
8733        let out_addr = SendMut(out.as_mut_ptr());
8734        let run = move |start: usize, end: usize| {
8735            ROW_I8.with(|rb| {
8736                let mut buf = rb.borrow_mut();
8737                buf.resize(cols, 0);
8738                for r in start..end {
8739                    // Unpack the row's nibbles to centered i8 once
8740                    // (element 2k = low nibble, 2k+1 = high — flat order,
8741                    // same as dot_q4_row_sdot's zip).
8742                    for gi in 0..gpr {
8743                        let g = r * gpr + gi;
8744                        for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
8745                            buf[gi * GROUP_SIZE + k * 2] = ((bt & 0x0F) as i32 - 8) as i8 as u8;
8746                            buf[gi * GROUP_SIZE + k * 2 + 1] =
8747                                (((bt >> 4) & 0x0F) as i32 - 8) as i8 as u8;
8748                        }
8749                    }
8750                    let mut bi = 0usize;
8751                    #[cfg(target_arch = "x86_64")]
8752                    if avx2_enabled() && blocked_enabled() {
8753                        while bi + 4 <= acts.len() {
8754                            let xs = [
8755                                acts[bi].xq.as_slice(),
8756                                acts[bi + 1].xq.as_slice(),
8757                                acts[bi + 2].xq.as_slice(),
8758                                acts[bi + 3].xq.as_slice(),
8759                            ];
8760                            let d = unsafe {
8761                                if vnni_tiles_enabled() {
8762                                    dot_q4b_row_1x4_vnni(&buf, scales, r * gpr, gpr, xs)
8763                                } else {
8764                                    dot_q4b_row_1x4_avx2(&buf, scales, r * gpr, gpr, xs)
8765                                }
8766                            };
8767                            for k in 0..4 {
8768                                let act = &acts[bi + k];
8769                                let mut acc = d[k] * act.sx;
8770                                for &(j, xv) in &act.outliers {
8771                                    acc += (buf[j] as i8) as f32
8772                                        * gscale((r * cols + j) / GROUP_SIZE)
8773                                        * xv;
8774                                }
8775                                // SAFETY: disjoint (bi, r) cells per worker.
8776                                unsafe { *out_addr.at((bi + k) * rows + r) = acc };
8777                            }
8778                            bi += 4;
8779                        }
8780                    }
8781                    while bi < acts.len() {
8782                        let act = &acts[bi];
8783                        let mut acc = 0f32;
8784                        for gi in 0..gpr {
8785                            let d = dot_i8_i8(
8786                                &buf[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
8787                                &act.xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
8788                            );
8789                            acc += d as f32 * gscale(r * gpr + gi);
8790                        }
8791                        acc *= act.sx;
8792                        // xq is zeroed at outlier slots — exact terms.
8793                        for &(j, xv) in &act.outliers {
8794                            acc += (buf[j] as i8) as f32 * gscale((r * cols + j) / GROUP_SIZE) * xv;
8795                        }
8796                        // SAFETY: disjoint (bi, r) cells per worker row range.
8797                        unsafe { *out_addr.at(bi * rows + r) = acc };
8798                        bi += 1;
8799                    }
8800                }
8801            })
8802        };
8803        dispatch_rows(pool, rows, &run);
8804        return;
8805    }
8806
8807    let out_addr = SendMut(out.as_mut_ptr());
8808    let run = move |start: usize, end: usize| {
8809        ROW_F32.with(|rb| {
8810            let mut buf = rb.borrow_mut();
8811            buf.resize(cols, 0.0);
8812            for r in start..end {
8813                // Decode raw (nib − 8) values once; scales stay per-group
8814                // so the accumulation order matches q4matvec bit-for-bit.
8815                for gi in 0..gpr {
8816                    let g = r * gpr + gi;
8817                    for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
8818                        buf[gi * GROUP_SIZE + k * 2] = (bt & 0x0F) as f32 - 8.0;
8819                        buf[gi * GROUP_SIZE + k * 2 + 1] = ((bt >> 4) & 0x0F) as f32 - 8.0;
8820                    }
8821                }
8822                for bi in 0..b {
8823                    let x = &xs_all[bi * cols..(bi + 1) * cols];
8824                    let mut acc = 0f32;
8825                    for gi in 0..gpr {
8826                        let mut ga = 0f32;
8827                        // Pairwise (lo + hi) addition, matching
8828                        // q4matvec's `ga += lo·x + hi·x` shape exactly —
8829                        // a flat one-per-element loop rounds differently
8830                        // and broke bit-parity on the scalar (x86) path.
8831                        for k in 0..GROUP_SIZE / 2 {
8832                            let e = gi * GROUP_SIZE + k * 2;
8833                            ga += buf[e] * x[e] + buf[e + 1] * x[e + 1];
8834                        }
8835                        acc += ga * gscale(r * gpr + gi);
8836                    }
8837                    // SAFETY: disjoint (bi, r) cells per worker row range.
8838                    unsafe { *out_addr.at(bi * rows + r) = acc };
8839                }
8840            }
8841        })
8842    };
8843    dispatch_rows(pool, rows, &run);
8844}
8845
8846/// Batched vbit matmat: each variable-bit row is decoded from the mmap
8847/// ONCE for the whole microbatch. Same per-position math as
8848/// `vbitmatvec` (SDOT A8W8 with exact outliers / exact f32 for b=8 rows
8849/// and the scalar path).
8850#[allow(clippy::too_many_arguments)]
8851fn vbitmatmat(
8852    bytes: &[u8],
8853    offsets: &[usize],
8854    xs_all: &[f32],
8855    b: usize,
8856    rows: usize,
8857    cols: usize,
8858    out: &mut [f32],
8859    pool: Option<&Pool>,
8860) {
8861    debug_assert_eq!(xs_all.len(), b * cols);
8862    debug_assert_eq!(out.len(), b * rows);
8863    debug_assert_eq!(offsets.len(), rows + 1);
8864    let ng = cols / GROUP_SIZE;
8865    let bits = &bytes[..rows];
8866    let sc_off = rows;
8867    let gscale = |r: usize, g: usize| {
8868        let so = (r * ng + g) * 2;
8869        f16_to_f32(u16::from_le_bytes([
8870            bytes[sc_off + so],
8871            bytes[sc_off + so + 1],
8872        ]))
8873    };
8874
8875    // Decode row r's raw (u − L) values into `dst` (f32, unscaled).
8876    let decode_f32 = |r: usize, dst: &mut [f32]| {
8877        let bw = bits[r] as usize;
8878        let l = ((1i32 << (bw - 1)) - 1) as f32;
8879        let data = &bytes[offsets[r]..offsets[r + 1]];
8880        let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
8881        for d in dst.iter_mut() {
8882            while nbits < bw {
8883                acc = (acc << 8) | data[idx] as u64;
8884                idx += 1;
8885                nbits += 8;
8886            }
8887            let u = ((acc >> (nbits - bw)) & ((1u64 << bw) - 1)) as f32;
8888            nbits -= bw;
8889            *d = u - l;
8890        }
8891    };
8892
8893    if a8w8_enabled() {
8894        let acts: Vec<SplitAct> = (0..b)
8895            .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
8896            .collect();
8897        let acts = &acts;
8898        let out_addr = SendMut(out.as_mut_ptr());
8899        let run = move |start: usize, end: usize| {
8900            for r in start..end {
8901                let bw = bits[r] as usize;
8902                if bw == 8 {
8903                    // u−L reaches 128 → no i8 path; decode once, exact
8904                    // f32 dots for every position (same as vbitmatvec).
8905                    ROW_F32.with(|rb| {
8906                        let mut buf = rb.borrow_mut();
8907                        buf.resize(cols, 0.0);
8908                        decode_f32(r, &mut buf);
8909                        for bi in 0..b {
8910                            let x = &xs_all[bi * cols..(bi + 1) * cols];
8911                            let mut dot = 0f32;
8912                            for g in 0..ng {
8913                                let mut gd = 0f32;
8914                                for k in 0..GROUP_SIZE {
8915                                    gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
8916                                }
8917                                dot += gd * gscale(r, g);
8918                            }
8919                            // SAFETY: disjoint (bi, r) cells per worker range.
8920                            unsafe { *out_addr.at(bi * rows + r) = dot };
8921                        }
8922                    });
8923                    continue;
8924                }
8925                let l = (1i32 << (bw - 1)) - 1;
8926                let data = &bytes[offsets[r]..offsets[r + 1]];
8927                ROW_I8.with(|rb| {
8928                    let mut buf = rb.borrow_mut();
8929                    buf.resize(cols, 0);
8930                    #[inline(always)]
8931                    fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
8932                        for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
8933                            let u = unpack8::<B>(&data[blk * B..]);
8934                            for k in 0..8 {
8935                                chunk[k] = (u[k] - l) as i8 as u8;
8936                            }
8937                        }
8938                    }
8939                    match bw {
8940                        3 => fill::<3>(data, l, &mut buf),
8941                        4 => vbit_fill4(data, &mut buf),
8942                        5 => fill::<5>(data, l, &mut buf),
8943                        6 => fill::<6>(data, l, &mut buf),
8944                        _ => unreachable!("vbit bit-width {bw} (validated at load)"),
8945                    }
8946                    let mut bi = 0usize;
8947                    // The vbit scale table shares q4_block's layout
8948                    // (contiguous f16 per (row·ng + g)), so the same
8949                    // blocked 1×4 kernel serves the decoded row.
8950                    #[cfg(target_arch = "x86_64")]
8951                    if avx2_enabled() && blocked_enabled() {
8952                        while bi + 4 <= acts.len() {
8953                            let xs = [
8954                                acts[bi].xq.as_slice(),
8955                                acts[bi + 1].xq.as_slice(),
8956                                acts[bi + 2].xq.as_slice(),
8957                                acts[bi + 3].xq.as_slice(),
8958                            ];
8959                            let sxs = [
8960                                acts[bi].sx,
8961                                acts[bi + 1].sx,
8962                                acts[bi + 2].sx,
8963                                acts[bi + 3].sx,
8964                            ];
8965                            let d = unsafe {
8966                                if vnni_tiles_enabled() {
8967                                    dot_q4b_row_1x4_sx_vnni(
8968                                        &buf,
8969                                        &bytes[sc_off..],
8970                                        r * ng,
8971                                        ng,
8972                                        xs,
8973                                        sxs,
8974                                    )
8975                                } else {
8976                                    dot_q4b_row_1x4_sx_avx2(
8977                                        &buf,
8978                                        &bytes[sc_off..],
8979                                        r * ng,
8980                                        ng,
8981                                        xs,
8982                                        sxs,
8983                                    )
8984                                }
8985                            };
8986                            for k in 0..4 {
8987                                let act = &acts[bi + k];
8988                                let mut dot = d[k];
8989                                for &(j, xv) in &act.outliers {
8990                                    dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
8991                                }
8992                                // SAFETY: disjoint (bi, r) cells per worker.
8993                                unsafe { *out_addr.at((bi + k) * rows + r) = dot };
8994                            }
8995                            bi += 4;
8996                        }
8997                    }
8998                    while bi < acts.len() {
8999                        let act = &acts[bi];
9000                        let mut dot = 0f32;
9001                        for g in 0..ng {
9002                            let d = dot_i8_i8(
9003                                &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
9004                                &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
9005                            ) as f32
9006                                * act.sx;
9007                            dot += d * gscale(r, g);
9008                        }
9009                        for &(j, xv) in &act.outliers {
9010                            dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
9011                        }
9012                        // SAFETY: disjoint (bi, r) cells per worker range.
9013                        unsafe { *out_addr.at(bi * rows + r) = dot };
9014                        bi += 1;
9015                    }
9016                });
9017            }
9018        };
9019        dispatch_rows(pool, rows, &run);
9020        return;
9021    }
9022
9023    let out_addr = SendMut(out.as_mut_ptr());
9024    let run = move |start: usize, end: usize| {
9025        ROW_F32.with(|rb| {
9026            let mut buf = rb.borrow_mut();
9027            buf.resize(cols, 0.0);
9028            for r in start..end {
9029                decode_f32(r, &mut buf);
9030                for bi in 0..b {
9031                    let x = &xs_all[bi * cols..(bi + 1) * cols];
9032                    let mut dot = 0f32;
9033                    for g in 0..ng {
9034                        let mut gd = 0f32;
9035                        for k in 0..GROUP_SIZE {
9036                            gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
9037                        }
9038                        dot += gd * gscale(r, g);
9039                    }
9040                    // SAFETY: disjoint (bi, r) cells per worker range.
9041                    unsafe { *out_addr.at(bi * rows + r) = dot };
9042                }
9043            }
9044        })
9045    };
9046    dispatch_rows(pool, rows, &run);
9047}
9048
9049/// Build a GPU batch job for a q8-family mapped tensor (primary
9050/// shard): prescaled input + directory coordinates. None → not
9051/// GPU-eligible, caller stays on the CPU.
9052pub(crate) fn gpu_batch_job<'a>(
9053    t: &'a QTensor,
9054    x: &[f32],
9055) -> Option<(std::sync::Arc<CmfModel>, crate::gpu::BatchJob<'a>)> {
9056    match t {
9057        QTensor::Mapped {
9058            model,
9059            idx,
9060            dtype: dt @ (TensorDtype::Q8Row | TensorDtype::Q8_2f),
9061            rows,
9062            cols,
9063            row_scale,
9064            col_field,
9065            ..
9066        } => Some((
9067            model.clone(),
9068            crate::gpu::BatchJob {
9069                idx: *idx,
9070                rows: *rows,
9071                cols: *cols,
9072                row_scale,
9073                xs: prescale(x, col_field, *dt).into_owned(),
9074                layout: crate::gpu::BatchLayout::Q8,
9075            },
9076        )),
9077        // q1: raw f32 activations, tile-embedded scales.
9078        QTensor::Mapped {
9079            model,
9080            idx,
9081            dtype: TensorDtype::Q1,
9082            rows,
9083            cols,
9084            ..
9085        } => Some((
9086            model.clone(),
9087            crate::gpu::BatchJob {
9088                idx: *idx,
9089                rows: *rows,
9090                cols: *cols,
9091                row_scale: &[],
9092                xs: x.to_vec(),
9093                layout: crate::gpu::BatchLayout::Q1,
9094            },
9095        )),
9096        // q4_tiled / q4tp: raw f32 activations; the scales live in the
9097        // payload (inline tiles / row ladder), so row_scale stays empty.
9098        // The GDN projection batch already runs these layouts on Metal —
9099        // this arm lets the attention QKV batch reach the same kernels.
9100        QTensor::Mapped {
9101            model,
9102            idx,
9103            dtype: dt @ (TensorDtype::Q4Tiled | TensorDtype::Q4TiledP),
9104            rows,
9105            cols,
9106            ..
9107        } => Some((
9108            model.clone(),
9109            crate::gpu::BatchJob {
9110                idx: *idx,
9111                rows: *rows,
9112                cols: *cols,
9113                row_scale: &[],
9114                xs: x.to_vec(),
9115                layout: if *dt == TensorDtype::Q4Tiled {
9116                    crate::gpu::BatchLayout::Q4t
9117                } else {
9118                    crate::gpu::BatchLayout::Q4tp
9119                },
9120            },
9121        )),
9122        _ => None,
9123    }
9124}
9125
9126thread_local! {
9127    static PRESCALE_BUF1: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
9128    static PRESCALE_BUF2: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
9129}
9130
9131pub(crate) fn prescale<'a>(
9132    x: &'a [f32],
9133    col_field: &[f32],
9134    dtype: TensorDtype,
9135) -> std::borrow::Cow<'a, [f32]> {
9136    if dtype == TensorDtype::Q8_2f {
9137        x.iter().zip(col_field).map(|(a, c)| a * c).collect()
9138    } else {
9139        std::borrow::Cow::Borrowed(x)
9140    }
9141}
9142
9143/// θ col-field fold for q8_2f activations. Borrowed pass-through for
9144/// every other dtype, using thread-local buffers to eliminate per-matvec allocations.
9145pub(crate) fn prescale_with<R, F: FnOnce(&[f32]) -> R>(
9146    x: &[f32],
9147    col_field: &[f32],
9148    dtype: TensorDtype,
9149    buf_id: u8,
9150    f: F,
9151) -> R {
9152    if dtype == TensorDtype::Q8_2f {
9153        if buf_id == 1 {
9154            PRESCALE_BUF1.with(|b| {
9155                let mut buf = b.borrow_mut();
9156                buf.clear();
9157                buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
9158                f(&buf)
9159            })
9160        } else {
9161            PRESCALE_BUF2.with(|b| {
9162                let mut buf = b.borrow_mut();
9163                buf.clear();
9164                buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
9165                f(&buf)
9166            })
9167        }
9168    } else {
9169        f(x)
9170    }
9171}
9172
9173// ───────────────────── x86-64 AVX2 kernels (roadmap этап 2) ─────────────────────
9174
9175/// AVX2+FMA available? Default ON when the CPU supports both;
9176/// `CMF_AVX2=0` disables (falls back to the autovectorized loops).
9177#[cfg(target_arch = "x86_64")]
9178pub(crate) fn avx2_enabled() -> bool {
9179    use std::sync::OnceLock;
9180    static ON: OnceLock<bool> = OnceLock::new();
9181    *ON.get_or_init(|| {
9182        std::env::var("CMF_AVX2").map(|v| v != "0").unwrap_or(true)
9183            && std::arch::is_x86_feature_detected!("avx2")
9184            && std::arch::is_x86_feature_detected!("fma")
9185    })
9186}
9187
9188/// AVX2 A8W8 allowed? The quantized-activation contract is switched by
9189/// the SAME env as the ARM SDOT path: `CMF_SDOT=0` keeps exact kernels
9190/// (the golden-parity exact gate relies on it) — AVX2 f32 kernels stay
9191/// active either way, they are exact (regrouped sums only).
9192#[cfg(target_arch = "x86_64")]
9193fn avx2_a8w8_enabled() -> bool {
9194    use std::sync::OnceLock;
9195    static ON: OnceLock<bool> = OnceLock::new();
9196    *ON.get_or_init(|| {
9197        avx2_enabled() && std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true)
9198    })
9199}
9200
9201/// A8W8 quantized-activation path available on THIS machine? One
9202/// switch across architectures: ARM dotprod (CMF_SDOT) or x86 AVX2
9203/// (CMF_AVX2 + the same CMF_SDOT exact-contract override).
9204#[inline]
9205pub(crate) fn a8w8_enabled() -> bool {
9206    #[cfg(target_arch = "aarch64")]
9207    {
9208        sdot_enabled()
9209    }
9210    #[cfg(target_arch = "x86_64")]
9211    {
9212        avx2_a8w8_enabled()
9213    }
9214    #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
9215    {
9216        false
9217    }
9218}
9219
9220/// int8·int8 dot dispatch: SDOT on ARM; AVX-512 VNNI (vpdpbusd) or AVX2
9221/// maddubs on x86. Callers are gated by `a8w8_enabled()`.
9222#[inline]
9223#[allow(unreachable_code)]
9224fn dot_i8_i8(w: &[u8], xq: &[i8]) -> i32 {
9225    #[cfg(target_arch = "aarch64")]
9226    unsafe {
9227        return dot_i8_sdot(w, xq);
9228    }
9229    #[cfg(target_arch = "x86_64")]
9230    unsafe {
9231        if avx512vnni_enabled() {
9232            return dot_i8_i8_vnni(w, xq);
9233        }
9234        return dot_i8_i8_avx2(w, xq);
9235    }
9236    w.iter()
9237        .zip(xq)
9238        .map(|(&a, &b)| (a as i8) as i32 * b as i32)
9239        .sum()
9240}
9241
9242/// AVX-512 VNNI available? (F+BW+VL+VNNI; `CMF_AVX512=0` falls back to
9243/// AVX2.) VL matters: short 32-byte groups (q4/vbit) ride the 256-bit
9244/// `vpdpbusd` encoding.
9245#[cfg(target_arch = "x86_64")]
9246fn avx512vnni_enabled() -> bool {
9247    use std::sync::OnceLock;
9248    static ON: OnceLock<bool> = OnceLock::new();
9249    *ON.get_or_init(|| {
9250        std::env::var("CMF_AVX512")
9251            .map(|v| v != "0")
9252            .unwrap_or(true)
9253            && std::arch::is_x86_feature_detected!("avx512f")
9254            && std::arch::is_x86_feature_detected!("avx512bw")
9255            && std::arch::is_x86_feature_detected!("avx512vl")
9256            && std::arch::is_x86_feature_detected!("avx512vnni")
9257    })
9258}
9259
9260/// Grouped-codec VNNI arms (the q4t/q4b/q1/q1t tile kernels): default
9261/// ON where AVX-512 VNNI exists (`CMF_VNNI_TILES=0` opt-out). Measured
9262/// on Ryzen 7950X (Zen4, 3 alternating process pairs, blocked GEMM
9263/// 4864×896 b=256): q4t 63→68 GF/s (+8%), q1 53→56 (+6%), q4b 72→75
9264/// (+4%) — consistent, no leg regressed. The tile kernels keep a
9265/// horizontal reduce per 32-weight group, so the `vpdpbusd` saving is
9266/// smaller than the long-dot q8 win (+13%), but it is real and free.
9267#[cfg(target_arch = "x86_64")]
9268fn vnni_tiles_enabled() -> bool {
9269    use std::sync::OnceLock;
9270    static ON: OnceLock<bool> = OnceLock::new();
9271    *ON.get_or_init(|| {
9272        std::env::var("CMF_VNNI_TILES")
9273            .map(|v| v != "0")
9274            .unwrap_or(true)
9275            && avx512vnni_enabled()
9276    })
9277}
9278
9279/// One 256-bit u8×i8 dot → i32 via `vpdpbusd` into a fresh accumulator
9280/// plus the same horizontal reduce the AVX2 kernels use. Products are
9281/// bounded (|w| ≤ 8 or ≤ 1), so maddubs never saturated — the i32 sum
9282/// is bit-identical to the maddubs+madd pair it replaces.
9283#[cfg(target_arch = "x86_64")]
9284#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
9285#[inline]
9286unsafe fn dpbusd_hsum(aw: core::arch::x86_64::__m256i, xs: core::arch::x86_64::__m256i) -> i32 {
9287    // SAFETY: pure register math.
9288    unsafe {
9289        use core::arch::x86_64::*;
9290        let d = _mm256_dpbusd_epi32(_mm256_setzero_si256(), aw, xs);
9291        let hi128 = _mm256_extracti128_si256::<1>(d);
9292        let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
9293        let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
9294        let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
9295        _mm_cvtsi128_si32(s32)
9296    }
9297}
9298
9299/// int8·int8 via AVX-512 VNNI: `vpdpbusd` fuses the maddubs+madd+add
9300/// triple into one u8×i8 dot-accumulate. AVX-512 has no vpsignb, so the
9301/// |w|·sign(x,w) trick becomes |w| × (x negated where w<0) via a mask
9302/// subtract — w==0 lanes contribute 0 through |w|=0 either way.
9303#[cfg(target_arch = "x86_64")]
9304#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
9305unsafe fn dot_i8_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
9306    // SAFETY: callers uphold slice-length contracts (see call sites).
9307    unsafe {
9308        use core::arch::x86_64::*;
9309        let n = w.len();
9310        let mut j = 0usize;
9311        let mut total: i32;
9312        // 4 independent accumulators: vpdpbusd is its own loop-carried
9313        // dependency (~5-cycle latency) — a single-acc loop runs
9314        // latency-bound and LOSES to the AVX2 maddubs kernel, measured
9315        // on Granite Rapids.
9316        {
9317            #[inline(always)]
9318            unsafe fn step(
9319                w: *const u8,
9320                x: *const i8,
9321                acc: core::arch::x86_64::__m512i,
9322            ) -> core::arch::x86_64::__m512i {
9323                unsafe {
9324                    use core::arch::x86_64::*;
9325                    let wv = _mm512_loadu_si512(w as *const _);
9326                    let xv = _mm512_loadu_si512(x as *const _);
9327                    let aw = _mm512_abs_epi8(wv);
9328                    let neg = _mm512_movepi8_mask(wv);
9329                    let sx = _mm512_mask_sub_epi8(xv, neg, _mm512_setzero_si512(), xv);
9330                    _mm512_dpbusd_epi32(acc, aw, sx)
9331                }
9332            }
9333            let (mut a0, mut a1, mut a2, mut a3) = (
9334                _mm512_setzero_si512(),
9335                _mm512_setzero_si512(),
9336                _mm512_setzero_si512(),
9337                _mm512_setzero_si512(),
9338            );
9339            while j + 256 <= n {
9340                a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
9341                a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), a1);
9342                a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), a2);
9343                a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), a3);
9344                j += 256;
9345            }
9346            while j + 64 <= n {
9347                a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
9348                j += 64;
9349            }
9350            let s01 = _mm512_add_epi32(a0, a1);
9351            let s23 = _mm512_add_epi32(a2, a3);
9352            total = _mm512_reduce_add_epi32(_mm512_add_epi32(s01, s23));
9353        }
9354        // 32-wide (q4/vbit groups are exactly 32 bytes).
9355        if j + 32 <= n {
9356            let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
9357            let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
9358            let d = _mm256_dpbusd_epi32(
9359                _mm256_setzero_si256(),
9360                _mm256_abs_epi8(wv),
9361                _mm256_sign_epi8(xv, wv),
9362            );
9363            let hi128 = _mm256_extracti128_si256::<1>(d);
9364            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
9365            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
9366            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
9367            total += _mm_cvtsi128_si32(s32);
9368            j += 32;
9369        }
9370        while j < n {
9371            total += (w[j] as i8) as i32 * xq[j] as i32;
9372            j += 1;
9373        }
9374        total
9375    }
9376}
9377
9378/// i8 row · f32 x via AVX2/FMA (x86 mirror of `dot_i8_f32_neon`).
9379#[cfg(target_arch = "x86_64")]
9380#[target_feature(enable = "avx2,fma")]
9381unsafe fn dot_i8_f32_avx2(w: &[u8], x: &[f32]) -> f32 {
9382    // SAFETY: callers uphold slice-length contracts (see call sites).
9383    unsafe {
9384        use core::arch::x86_64::*;
9385        let n = x.len();
9386        let wp = w.as_ptr();
9387        let xp = x.as_ptr();
9388        let (mut a0, mut a1) = (_mm256_setzero_ps(), _mm256_setzero_ps());
9389        let mut j = 0usize;
9390        while j + 16 <= n {
9391            let wb = _mm_loadu_si128(wp.add(j) as *const __m128i);
9392            let lo = _mm256_cvtepi8_epi32(wb);
9393            let hi = _mm256_cvtepi8_epi32(_mm_srli_si128::<8>(wb));
9394            a0 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(lo), _mm256_loadu_ps(xp.add(j)), a0);
9395            a1 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(hi), _mm256_loadu_ps(xp.add(j + 8)), a1);
9396            j += 16;
9397        }
9398        let acc = _mm256_add_ps(a0, a1);
9399        let hi128 = _mm256_extractf128_ps::<1>(acc);
9400        let s128 = _mm_add_ps(_mm256_castps256_ps128(acc), hi128);
9401        let s64 = _mm_add_ps(s128, _mm_movehl_ps(s128, s128));
9402        let s32 = _mm_add_ss(s64, _mm_shuffle_ps::<1>(s64, s64));
9403        let mut sum = _mm_cvtss_f32(s32);
9404        while j < n {
9405            sum += (*wp.add(j) as i8) as f32 * *xp.add(j);
9406            j += 1;
9407        }
9408        sum
9409    }
9410}
9411
9412/// int8(weight)·int8(activation) → i32 via AVX2 maddubs — the x86
9413/// analogue of the SDOT A8W8 path. `maddubs` takes u8×i8, so the
9414/// standard sign trick applies: |w| × sign(x, w) ≡ w × x per lane.
9415/// Pair saturation is safe: |w|≤128, |x|≤127 → 2·128·127 < 32767.
9416#[cfg(target_arch = "x86_64")]
9417#[target_feature(enable = "avx2")]
9418unsafe fn dot_i8_i8_avx2(w: &[u8], xq: &[i8]) -> i32 {
9419    // SAFETY: callers uphold slice-length contracts (see call sites).
9420    unsafe {
9421        use core::arch::x86_64::*;
9422        let n = w.len();
9423        let ones = _mm256_set1_epi16(1);
9424        let mut acc = _mm256_setzero_si256();
9425        let mut j = 0usize;
9426        while j + 32 <= n {
9427            let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
9428            let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
9429            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
9430            acc = _mm256_add_epi32(acc, _mm256_madd_epi16(p16, ones));
9431            j += 32;
9432        }
9433        let hi128 = _mm256_extracti128_si256::<1>(acc);
9434        let s128 = _mm_add_epi32(_mm256_castsi256_si128(acc), hi128);
9435        let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
9436        let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
9437        let mut s = _mm_cvtsi128_si32(s32);
9438        while j < n {
9439            s += (w[j] as i8) as i32 * xq[j] as i32;
9440            j += 1;
9441        }
9442        s
9443    }
9444}
9445
9446/// smmla 2×4: one instruction covers a 2-row × 2-activation × 8-deep
9447/// tile (32 MACs vs sdot's 16) — the weight pair loads once per 8-k
9448/// slice as a combined 2×8 register and meets two activation pairs.
9449#[cfg(target_arch = "aarch64")]
9450#[target_feature(enable = "neon,i8mm")]
9451unsafe fn dot_i8_smmla_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
9452    // SAFETY: callers uphold slice-length contracts.
9453    unsafe {
9454        use core::arch::aarch64::*;
9455        use core::arch::asm;
9456        let n = w0.len();
9457        let w0p = w0.as_ptr() as *const i8;
9458        let w1p = w1.as_ptr() as *const i8;
9459        // acc01 holds [c(r0,x0) c(r0,x1) c(r1,x0) c(r1,x1)]; acc23 the
9460        // same for x2/x3.
9461        let mut acc01 = vdupq_n_s32(0);
9462        let mut acc23 = vdupq_n_s32(0);
9463        let mut i = 0usize;
9464        while i + 8 <= n {
9465            let wa = vcombine_s8(vld1_s8(w0p.add(i)), vld1_s8(w1p.add(i)));
9466            let xb01 = vcombine_s8(
9467                vld1_s8(xs[0].as_ptr().add(i)),
9468                vld1_s8(xs[1].as_ptr().add(i)),
9469            );
9470            let xb23 = vcombine_s8(
9471                vld1_s8(xs[2].as_ptr().add(i)),
9472                vld1_s8(xs[3].as_ptr().add(i)),
9473            );
9474            asm!(
9475                "smmla {a01:v}.4s, {w:v}.16b, {x01:v}.16b",
9476                "smmla {a23:v}.4s, {w:v}.16b, {x23:v}.16b",
9477                a01 = inout(vreg) acc01, a23 = inout(vreg) acc23,
9478                w = in(vreg) wa, x01 = in(vreg) xb01, x23 = in(vreg) xb23,
9479                options(pure, nomem, nostack),
9480            );
9481            i += 8;
9482        }
9483        let mut out = [[0i32; 4]; 2];
9484        let a01: [i32; 4] = core::mem::transmute(acc01);
9485        let a23: [i32; 4] = core::mem::transmute(acc23);
9486        out[0][0] = a01[0];
9487        out[0][1] = a01[1];
9488        out[1][0] = a01[2];
9489        out[1][1] = a01[3];
9490        out[0][2] = a23[0];
9491        out[0][3] = a23[1];
9492        out[1][2] = a23[2];
9493        out[1][3] = a23[3];
9494        if i < n {
9495            for (k, x) in xs.iter().enumerate() {
9496                for j in i..n {
9497                    out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
9498                    out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
9499                }
9500            }
9501        }
9502        out
9503    }
9504}
9505
9506/// ARM twin of the x86 blocked prefill GEMM: two weight rows stay in
9507/// registers across four activation streams, eight sdot accumulators.
9508/// (The per-row form re-read each W row once per activation.)
9509#[cfg(target_arch = "aarch64")]
9510#[target_feature(enable = "neon,dotprod")]
9511unsafe fn dot_i8_sdot_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
9512    // SAFETY: callers uphold slice-length contracts.
9513    unsafe {
9514        use core::arch::aarch64::*;
9515        use core::arch::asm;
9516        let n = w0.len();
9517        let w0p = w0.as_ptr() as *const i8;
9518        let w1p = w1.as_ptr() as *const i8;
9519        let mut acc = [[vdupq_n_s32(0); 4]; 2];
9520        let mut i = 0usize;
9521        while i + 16 <= n {
9522            let wv0 = vld1q_s8(w0p.add(i));
9523            let wv1 = vld1q_s8(w1p.add(i));
9524            for (k, x) in xs.iter().enumerate() {
9525                let xv = vld1q_s8(x.as_ptr().add(i));
9526                let (mut a0, mut a1) = (acc[0][k], acc[1][k]);
9527                asm!(
9528                    "sdot {a0:v}.4s, {w0:v}.16b, {x:v}.16b",
9529                    "sdot {a1:v}.4s, {w1:v}.16b, {x:v}.16b",
9530                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
9531                    w0 = in(vreg) wv0, w1 = in(vreg) wv1, x = in(vreg) xv,
9532                    options(pure, nomem, nostack),
9533                );
9534                acc[0][k] = a0;
9535                acc[1][k] = a1;
9536            }
9537            i += 16;
9538        }
9539        let mut out = [[0i32; 4]; 2];
9540        for r in 0..2 {
9541            for k in 0..4 {
9542                out[r][k] = vaddvq_s32(acc[r][k]);
9543            }
9544        }
9545        if i < n {
9546            for (k, x) in xs.iter().enumerate() {
9547                for j in i..n {
9548                    out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
9549                    out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
9550                }
9551            }
9552        }
9553        out
9554    }
9555}
9556
9557/// Blocked 2 weight rows × 4 activations for the prefill GEMM
9558/// (roadmap P0: packed panels + multi-row accumulators). The two rows'
9559/// abs() live in registers across all four activation streams; the
9560/// sign-fixup is recomputed per pair (the price of the maddubs trick).
9561/// Returns raw i8·i8 dots; the caller applies scales and outliers.
9562#[cfg(target_arch = "x86_64")]
9563#[target_feature(enable = "avx2")]
9564unsafe fn dot_i8_i8_avx2_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
9565    // SAFETY: callers uphold slice-length contracts.
9566    unsafe {
9567        use core::arch::x86_64::*;
9568        let n = w0.len();
9569        let ones = _mm256_set1_epi16(1);
9570        let mut acc = [[_mm256_setzero_si256(); 4]; 2];
9571        let mut j = 0usize;
9572        while j + 32 <= n {
9573            let wv0 = _mm256_loadu_si256(w0.as_ptr().add(j) as *const __m256i);
9574            let wv1 = _mm256_loadu_si256(w1.as_ptr().add(j) as *const __m256i);
9575            let aw0 = _mm256_abs_epi8(wv0);
9576            let aw1 = _mm256_abs_epi8(wv1);
9577            for (k, x) in xs.iter().enumerate() {
9578                let xv = _mm256_loadu_si256(x.as_ptr().add(j) as *const __m256i);
9579                let p0 = _mm256_maddubs_epi16(aw0, _mm256_sign_epi8(xv, wv0));
9580                acc[0][k] = _mm256_add_epi32(acc[0][k], _mm256_madd_epi16(p0, ones));
9581                let p1 = _mm256_maddubs_epi16(aw1, _mm256_sign_epi8(xv, wv1));
9582                acc[1][k] = _mm256_add_epi32(acc[1][k], _mm256_madd_epi16(p1, ones));
9583            }
9584            j += 32;
9585        }
9586        let mut out = [[0i32; 4]; 2];
9587        for r in 0..2 {
9588            for k in 0..4 {
9589                let a = acc[r][k];
9590                let hi128 = _mm256_extracti128_si256::<1>(a);
9591                let s128 = _mm_add_epi32(_mm256_castsi256_si128(a), hi128);
9592                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
9593                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
9594                out[r][k] = _mm_cvtsi128_si32(s32);
9595            }
9596        }
9597        if j < n {
9598            for (k, x) in xs.iter().enumerate() {
9599                for i in j..n {
9600                    out[0][k] += (w0[i] as i8) as i32 * x[i] as i32;
9601                    out[1][k] += (w1[i] as i8) as i32 * x[i] as i32;
9602                }
9603            }
9604        }
9605        out
9606    }
9607}
9608
9609/// AVX2/VNNI q8 row dot with exact outlier correction (x86 mirror of
9610/// `row_dot_sdot` — same A8W8 contract). With AVX-512 VNNI the row goes
9611/// through the bias trick: Σ(w+128)·x via pure `vpdpbusd` (no per-lane
9612/// sign fixups), corrected by −128·Σx with Σx precomputed per split.
9613#[cfg(target_arch = "x86_64")]
9614#[inline]
9615fn row_dot_avx2(row: &[u8], act: &SplitAct) -> f32 {
9616    let dot = if avx512vnni_enabled() && row.len() >= 64 {
9617        (unsafe { dot_u8p128_i8_vnni(row, &act.xq) }) - 128 * act.xsum
9618    } else {
9619        unsafe { dot_i8_i8_avx2(row, &act.xq) }
9620    };
9621    let mut acc = dot as f32 * act.sx;
9622    for &(j, xv) in &act.outliers {
9623        acc += (row[j] as i8) as f32 * xv;
9624    }
9625    acc
9626}
9627
9628/// Σ (w[i]+128)·x[i] via pure `vpdpbusd` — the caller subtracts
9629/// 128·Σx. Four independent accumulators (dpbusd is ~5-cycle latency;
9630/// a single-acc loop runs latency-bound, measured on Granite Rapids).
9631#[cfg(target_arch = "x86_64")]
9632#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
9633unsafe fn dot_u8p128_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
9634    // SAFETY: callers uphold slice-length contracts (see call sites).
9635    unsafe {
9636        use core::arch::x86_64::*;
9637        let n = w.len();
9638        let flip = _mm512_set1_epi8(-128); // XOR 0x80: i8 w → u8 (w+128)
9639        #[inline(always)]
9640        unsafe fn step(
9641            w: *const u8,
9642            x: *const i8,
9643            flip: core::arch::x86_64::__m512i,
9644            acc: core::arch::x86_64::__m512i,
9645        ) -> core::arch::x86_64::__m512i {
9646            unsafe {
9647                use core::arch::x86_64::*;
9648                let wv = _mm512_xor_si512(_mm512_loadu_si512(w as *const _), flip);
9649                _mm512_dpbusd_epi32(acc, wv, _mm512_loadu_si512(x as *const _))
9650            }
9651        }
9652        let (mut a0, mut a1, mut a2, mut a3) = (
9653            _mm512_setzero_si512(),
9654            _mm512_setzero_si512(),
9655            _mm512_setzero_si512(),
9656            _mm512_setzero_si512(),
9657        );
9658        let mut j = 0usize;
9659        while j + 256 <= n {
9660            a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
9661            a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), flip, a1);
9662            a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), flip, a2);
9663            a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), flip, a3);
9664            j += 256;
9665        }
9666        while j + 64 <= n {
9667            a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
9668            j += 64;
9669        }
9670        let mut total = _mm512_reduce_add_epi32(_mm512_add_epi32(
9671            _mm512_add_epi32(a0, a1),
9672            _mm512_add_epi32(a2, a3),
9673        ));
9674        // Scalar tail: (w as i8) + 128 ≡ (w as u8) ^ 0x80.
9675        while j < n {
9676            total += ((w[j] ^ 0x80) as i32) * xq[j] as i32;
9677            j += 1;
9678        }
9679        total
9680    }
9681}
9682
9683/// One q4 row via AVX2: nibbles → centered i8 (unpacklo/hi restores the
9684/// writer's flat order, same as the NEON vzip pair), maddubs against
9685/// the pre-quantized activation group, × the group's f16 scale. Pair
9686/// saturation safe: |w|≤8, |x|≤127 → 2·8·127 ≪ 32767. Mirror of
9687/// `dot_q4_row_sdot`.
9688#[cfg(target_arch = "x86_64")]
9689#[target_feature(enable = "avx2")]
9690unsafe fn dot_q4_row_avx2(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
9691    // SAFETY: callers uphold slice-length contracts (16 packed bytes and
9692    // 2 scale bytes per group; xq.len() == gpr·GROUP_SIZE).
9693    unsafe {
9694        use core::arch::x86_64::*;
9695        let lomask = _mm_set1_epi8(0x0F);
9696        let eight = _mm256_set1_epi8(8);
9697        let ones = _mm256_set1_epi16(1);
9698        let mut acc = 0f32;
9699        for gi in 0..gpr {
9700            let g = g0 + gi;
9701            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
9702            let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
9703            let lo = _mm_and_si128(b, lomask);
9704            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
9705            let w = _mm256_sub_epi8(
9706                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
9707                eight,
9708            );
9709            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9710            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
9711            let d = _mm256_madd_epi16(p16, ones);
9712            let hi128 = _mm256_extracti128_si256::<1>(d);
9713            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
9714            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
9715            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
9716            acc += _mm_cvtsi128_si32(s32) as f32 * s;
9717        }
9718        acc
9719    }
9720}
9721
9722/// Two-activation q4 row via AVX2: nibbles unpacked ONCE per group,
9723/// both activations dotted against the same centered i8 register.
9724#[cfg(target_arch = "x86_64")]
9725#[target_feature(enable = "avx2")]
9726unsafe fn dot_q4_row_avx2_2(
9727    packed: &[u8],
9728    scales: &[u8],
9729    g0: usize,
9730    gpr: usize,
9731    xq1: &[i8],
9732    xq2: &[i8],
9733) -> (f32, f32) {
9734    // SAFETY: callers uphold slice-length contracts (see dot_q4_row_avx2).
9735    unsafe {
9736        use core::arch::x86_64::*;
9737        let lomask = _mm_set1_epi8(0x0F);
9738        let eight = _mm256_set1_epi8(8);
9739        let ones = _mm256_set1_epi16(1);
9740        let (mut acc1, mut acc2) = (0f32, 0f32);
9741        #[inline(always)]
9742        unsafe fn hsum(d: core::arch::x86_64::__m256i) -> i32 {
9743            unsafe {
9744                use core::arch::x86_64::*;
9745                let hi128 = _mm256_extracti128_si256::<1>(d);
9746                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
9747                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
9748                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
9749                _mm_cvtsi128_si32(s32)
9750            }
9751        }
9752        for gi in 0..gpr {
9753            let g = g0 + gi;
9754            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
9755            let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
9756            let lo = _mm_and_si128(b, lomask);
9757            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
9758            let w = _mm256_sub_epi8(
9759                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
9760                eight,
9761            );
9762            let aw = _mm256_abs_epi8(w);
9763            let x1 = _mm256_loadu_si256(xq1.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9764            let x2 = _mm256_loadu_si256(xq2.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9765            let d1 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x1, w)), ones);
9766            let d2 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x2, w)), ones);
9767            acc1 += hsum(d1) as f32 * s;
9768            acc2 += hsum(d2) as f32 * s;
9769        }
9770        (acc1, acc2)
9771    }
9772}
9773
9774/// One q8 row range via AVX2 (x86 mirror of `q8_range_sdot`).
9775#[cfg(target_arch = "x86_64")]
9776fn q8_range_avx2(
9777    q: &[u8],
9778    row_scale: &[f32],
9779    act: &SplitAct,
9780    cols: usize,
9781    out_addr: SendMut,
9782    start: usize,
9783    end: usize,
9784) {
9785    for o in start..end {
9786        let v = row_dot_avx2(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
9787        // SAFETY: disjoint row ranges per worker.
9788        unsafe { *out_addr.at(o) = v };
9789    }
9790}
9791
9792/// Two-input q8 row range via AVX2 (x86 mirror of `q8_range2_sdot`).
9793#[cfg(target_arch = "x86_64")]
9794#[allow(clippy::too_many_arguments)]
9795fn q8_range2_avx2(
9796    q: &[u8],
9797    row_scale: &[f32],
9798    a1: &SplitAct,
9799    a2: &SplitAct,
9800    cols: usize,
9801    p1: SendMut,
9802    p2: SendMut,
9803    start: usize,
9804    end: usize,
9805) {
9806    for o in start..end {
9807        let row = &q[o * cols..(o + 1) * cols];
9808        // SAFETY: disjoint row ranges per worker.
9809        unsafe {
9810            *p1.at(o) = row_dot_avx2(row, a1) * row_scale[o];
9811            *p2.at(o) = row_dot_avx2(row, a2) * row_scale[o];
9812        }
9813    }
9814}
9815
9816// ───────────────────── A8W8 SDOT path (port of vmfcore, ×1.78 decode) ─────────────────────
9817
9818/// ARMv8.6 i8mm (smmla): 32 int8 MACs per instruction vs sdot's 16 —
9819/// yet MEASURED 2.4× SLOWER than the blocked sdot on Apple silicon
9820/// (108 vs 264 GF/s): the on-the-fly vcombine packing and the two-
9821/// accumulator dependency chain swamp the MAC advantage, and Apple's
9822/// four SIMD pipes already keep sdot fed. OPT-IN (CMF_I8MM=1) for
9823/// field trials on Cortex-A710/X-class parts with two pipes, where the
9824/// balance may differ; a pre-interleaved weight layout (repack infra)
9825/// is the known path if it ever earns its keep.
9826#[cfg(target_arch = "aarch64")]
9827fn i8mm_enabled() -> bool {
9828    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
9829    *ON.get_or_init(|| {
9830        std::env::var("CMF_I8MM").map(|v| v == "1").unwrap_or(false)
9831            && std::arch::is_aarch64_feature_detected!("i8mm")
9832    })
9833}
9834
9835/// SDOT enabled? Default ON when the CPU has ARMv8.2 dotprod;
9836/// `CMF_SDOT=0` disables (falls back to i8×f32 NEON).
9837/// (On non-ARM release builds only the test tolerance switch calls it.)
9838#[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
9839fn sdot_enabled() -> bool {
9840    use std::sync::OnceLock;
9841    static ON: OnceLock<bool> = OnceLock::new();
9842    *ON.get_or_init(|| {
9843        let want = std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true);
9844        if !want {
9845            return false;
9846        }
9847
9848        #[cfg(target_arch = "aarch64")]
9849        {
9850            if std::arch::is_aarch64_feature_detected!("dotprod") {
9851                return true;
9852            }
9853            #[cfg(target_os = "android")]
9854            {
9855                if let Ok(cpuinfo) = std::fs::read_to_string("/proc/cpuinfo") {
9856                    if cpuinfo.lines().any(|l| {
9857                        (l.starts_with("Features") || l.starts_with("features"))
9858                            && l.contains("asimddp")
9859                    }) {
9860                        return true;
9861                    }
9862                }
9863            }
9864            false
9865        }
9866        #[cfg(not(target_arch = "aarch64"))]
9867        {
9868            false
9869        }
9870    })
9871}
9872
9873/// Two-field activation split (≡ vmfcore `q8_split_prep`): outlier
9874/// channels (>8·rms) are computed exactly in f32; the bulk (outliers
9875/// zeroed → clean absmax) goes through int8 SDOT. Computed ONCE per
9876/// matvec, shared by all rows/workers.
9877struct SplitAct {
9878    xq: Vec<i8>,
9879    sx: f32,
9880    outliers: Vec<(usize, f32)>,
9881    /// Σ xq — the VNNI bias-trick correction (`(w+128)·x` sums need
9882    /// `−128·Σx`); one i32 per split, computed once per matvec.
9883    #[cfg_attr(not(target_arch = "x86_64"), allow(dead_code))]
9884    xsum: i32,
9885}
9886
9887thread_local! {
9888    /// Recycled xq buffers: split_act runs for every matvec (~200/token)
9889    /// and its hidden-size allocation was steady-state heap churn.
9890    static XQ_FREE: std::cell::RefCell<Vec<Vec<i8>>> =
9891        const { std::cell::RefCell::new(Vec::new()) };
9892}
9893
9894impl Drop for SplitAct {
9895    fn drop(&mut self) {
9896        let buf = std::mem::take(&mut self.xq);
9897        if buf.capacity() > 0 {
9898            XQ_FREE.with(|f| {
9899                let mut f = f.borrow_mut();
9900                if f.len() < 16 {
9901                    f.push(buf);
9902                }
9903            });
9904        }
9905    }
9906}
9907
9908thread_local! {
9909    /// One scratch row per WORKER, kept for the life of the thread.
9910    ///
9911    /// The kernels take a row of group scales per dispatch, and a fresh
9912    /// `vec![0f32; gpr]` inside the closure is one allocation per worker per
9913    /// dispatch — on the release checkpoint about six thousand a token, a
9914    /// quarter of everything the benchmark counts.
9915    static KROW: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
9916}
9917
9918/// Borrow `n` floats of the calling worker's scratch. Nothing inside a
9919/// kernel body borrows it again, which is what keeps the RefCell honest.
9920#[inline]
9921fn with_krow<R>(n: usize, f: impl FnOnce(&mut [f32]) -> R) -> R {
9922    KROW.with(|s| {
9923        let mut b = s.borrow_mut();
9924        if b.len() < n {
9925            b.resize(n, 0.0);
9926        }
9927        f(&mut b[..n])
9928    })
9929}
9930
9931fn split_act(x: &[f32]) -> SplitAct {
9932    let n = x.len();
9933    let rms = (x.iter().map(|&v| (v * v) as f64).sum::<f64>() / n.max(1) as f64).sqrt() as f32;
9934    let thr = 8.0 * rms;
9935    // One pass: collect outliers and the bulk absmax (outliers excluded —
9936    // identical to the old zero-then-fold over a copied buffer, minus the
9937    // full-vector copy).
9938    let mut outliers: Vec<(usize, f32)> = Vec::new();
9939    let mut amax = 0f32;
9940    for (j, &v) in x.iter().enumerate() {
9941        let a = v.abs();
9942        if a > thr {
9943            outliers.push((j, v));
9944        } else if a > amax {
9945            amax = a;
9946        }
9947    }
9948    let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
9949    let inv = 1.0 / sx;
9950    let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
9951    xq.clear();
9952    xq.reserve(n);
9953    if outliers.is_empty() {
9954        xq.extend(
9955            x.iter()
9956                .map(|&v| (v * inv).round().clamp(-127.0, 127.0) as i8),
9957        );
9958    } else {
9959        // Outlier slots quantize to 0 (their exact term is added later).
9960        xq.extend(x.iter().map(|&v| {
9961            if v.abs() > thr {
9962                0
9963            } else {
9964                (v * inv).round().clamp(-127.0, 127.0) as i8
9965            }
9966        }));
9967    }
9968    let xsum = xq.iter().map(|&v| v as i32).sum();
9969    SplitAct {
9970        xq,
9971        sx,
9972        outliers,
9973        xsum,
9974    }
9975}
9976
9977fn split_act_q8_2f(x: &[f32], col: &[f32]) -> SplitAct {
9978    let n = x.len();
9979    let rms = (x
9980        .iter()
9981        .zip(col)
9982        .map(|(&a, &c)| {
9983            let v = a * c;
9984            (v * v) as f64
9985        })
9986        .sum::<f64>()
9987        / n.max(1) as f64)
9988        .sqrt() as f32;
9989    let thr = 8.0 * rms;
9990
9991    let mut outliers = Vec::new();
9992    let mut amax = 0f32;
9993    for (j, (&a, &c)) in x.iter().zip(col).enumerate() {
9994        let v = a * c;
9995        let s = v.abs();
9996        if s > thr {
9997            outliers.push((j, v));
9998        } else if s > amax {
9999            amax = s;
10000        }
10001    }
10002
10003    let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
10004    let inv = 1.0 / sx;
10005    let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
10006    xq.clear();
10007    xq.reserve(n);
10008    if outliers.is_empty() {
10009        xq.extend(
10010            x.iter()
10011                .zip(col)
10012                .map(|(&a, &c)| ((a * c) * inv).round().clamp(-127.0, 127.0) as i8),
10013        );
10014    } else {
10015        xq.extend(x.iter().zip(col).map(|(&a, &c)| {
10016            let v = a * c;
10017            if v.abs() > thr {
10018                0
10019            } else {
10020                (v * inv).round().clamp(-127.0, 127.0) as i8
10021            }
10022        }));
10023    }
10024    let xsum = xq.iter().map(|&v| v as i32).sum();
10025    SplitAct {
10026        xq,
10027        sx,
10028        outliers,
10029        xsum,
10030    }
10031}
10032
10033/// int8(weight)·int8(activation) → i32 via `sdot` (inline asm — the
10034/// vdotq intrinsic is unstable; port of vmfcore `dot_i8_sdot`).
10035#[cfg(target_arch = "aarch64")]
10036#[target_feature(enable = "neon,dotprod")]
10037unsafe fn dot_i8_sdot(w: &[u8], xq: &[i8]) -> i32 {
10038    // SAFETY: callers uphold slice-length contracts (see call sites).
10039    unsafe {
10040        use core::arch::aarch64::*;
10041        use core::arch::asm;
10042        let wp = w.as_ptr() as *const i8;
10043        let n = w.len();
10044        let (mut a0, mut a1, mut a2, mut a3) = (
10045            vdupq_n_s32(0),
10046            vdupq_n_s32(0),
10047            vdupq_n_s32(0),
10048            vdupq_n_s32(0),
10049        );
10050        let mut i = 0;
10051        while i + 64 <= n {
10052            let (w0, x0) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
10053            let (w1, x1) = (vld1q_s8(wp.add(i + 16)), vld1q_s8(xq.as_ptr().add(i + 16)));
10054            let (w2, x2) = (vld1q_s8(wp.add(i + 32)), vld1q_s8(xq.as_ptr().add(i + 32)));
10055            let (w3, x3) = (vld1q_s8(wp.add(i + 48)), vld1q_s8(xq.as_ptr().add(i + 48)));
10056            asm!(
10057                "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
10058                "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
10059                "sdot {a2:v}.4s, {w2:v}.16b, {x2:v}.16b",
10060                "sdot {a3:v}.4s, {w3:v}.16b, {x3:v}.16b",
10061                a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
10062                w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
10063                w2 = in(vreg) w2, x2 = in(vreg) x2, w3 = in(vreg) w3, x3 = in(vreg) x3,
10064                options(pure, nomem, nostack),
10065            );
10066            i += 64;
10067        }
10068        while i + 16 <= n {
10069            let (wv, xv) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
10070            asm!("sdot {a:v}.4s, {w:v}.16b, {x:v}.16b",
10071                 a = inout(vreg) a0, w = in(vreg) wv, x = in(vreg) xv, options(pure, nomem, nostack));
10072            i += 16;
10073        }
10074        let mut s = vaddvq_s32(vaddq_s32(vaddq_s32(a0, a1), vaddq_s32(a2, a3)));
10075        while i < n {
10076            s += (*wp.add(i)) as i32 * xq[i] as i32;
10077            i += 1;
10078        }
10079        s
10080    }
10081}
10082
10083/// Row-blocked SDOT: 4 output rows per pass — the activation chunk is
10084/// loaded once and reused, 4 independent accumulators hide sdot latency
10085/// (port of vmfcore `dot_i8_sdot_4rows`).
10086#[cfg(target_arch = "aarch64")]
10087#[target_feature(enable = "neon,dotprod")]
10088unsafe fn dot_i8_sdot_4rows(w0: &[u8], w1: &[u8], w2: &[u8], w3: &[u8], xq: &[i8]) -> [i32; 4] {
10089    // SAFETY: callers uphold slice-length contracts (see call sites).
10090    unsafe {
10091        use core::arch::aarch64::*;
10092        use core::arch::asm;
10093        let n = xq.len();
10094        let px = xq.as_ptr();
10095        let (p0, p1, p2, p3) = (
10096            w0.as_ptr() as *const i8,
10097            w1.as_ptr() as *const i8,
10098            w2.as_ptr() as *const i8,
10099            w3.as_ptr() as *const i8,
10100        );
10101        let (mut a0, mut a1, mut a2, mut a3) = (
10102            vdupq_n_s32(0),
10103            vdupq_n_s32(0),
10104            vdupq_n_s32(0),
10105            vdupq_n_s32(0),
10106        );
10107        let mut i = 0;
10108        while i + 16 <= n {
10109            let x = vld1q_s8(px.add(i));
10110            let v0 = vld1q_s8(p0.add(i));
10111            let v1 = vld1q_s8(p1.add(i));
10112            let v2 = vld1q_s8(p2.add(i));
10113            let v3 = vld1q_s8(p3.add(i));
10114            asm!(
10115                "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
10116                "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
10117                "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
10118                "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
10119                a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
10120                v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
10121                options(pure, nomem, nostack),
10122            );
10123            i += 16;
10124        }
10125        let mut r = [
10126            vaddvq_s32(a0),
10127            vaddvq_s32(a1),
10128            vaddvq_s32(a2),
10129            vaddvq_s32(a3),
10130        ];
10131        while i < n {
10132            let xi = *px.add(i) as i32;
10133            r[0] += (*p0.add(i)) as i32 * xi;
10134            r[1] += (*p1.add(i)) as i32 * xi;
10135            r[2] += (*p2.add(i)) as i32 * xi;
10136            r[3] += (*p3.add(i)) as i32 * xi;
10137            i += 1;
10138        }
10139        r
10140    }
10141}
10142
10143/// 4 interleaved rows in one pass: the repacked group is [r0[c], r1[c],
10144/// r2[c], r3[c]] per 16-byte chunk, so each iteration reads ONE 64-byte
10145/// line plus the shared activation chunk — a single sequential weight
10146/// stream per worker. Per-row accumulation is the same one-accumulator
10147/// scheme as `dot_i8_sdot_4rows`; integer sums are exact, so outputs
10148/// are bit-identical to the mmap-layout kernel.
10149#[cfg(target_arch = "aarch64")]
10150#[target_feature(enable = "neon,dotprod")]
10151unsafe fn dot_i8_sdot_4rows_il(g: &[u8], xq: &[i8]) -> [i32; 4] {
10152    // SAFETY: callers uphold slice-length contracts (g.len() == 4·n,
10153    // n % 16 == 0 — guaranteed by the repack gate).
10154    unsafe {
10155        use core::arch::aarch64::*;
10156        use core::arch::asm;
10157        let n = xq.len();
10158        let px = xq.as_ptr();
10159        let pg = g.as_ptr() as *const i8;
10160        let (mut a0, mut a1, mut a2, mut a3) = (
10161            vdupq_n_s32(0),
10162            vdupq_n_s32(0),
10163            vdupq_n_s32(0),
10164            vdupq_n_s32(0),
10165        );
10166        let mut i = 0;
10167        while i + 16 <= n {
10168            let x = vld1q_s8(px.add(i));
10169            let base = pg.add(4 * i);
10170            let v0 = vld1q_s8(base);
10171            let v1 = vld1q_s8(base.add(16));
10172            let v2 = vld1q_s8(base.add(32));
10173            let v3 = vld1q_s8(base.add(48));
10174            asm!(
10175                "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
10176                "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
10177                "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
10178                "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
10179                a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
10180                v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
10181                options(pure, nomem, nostack),
10182            );
10183            i += 16;
10184        }
10185        [
10186            vaddvq_s32(a0),
10187            vaddvq_s32(a1),
10188            vaddvq_s32(a2),
10189            vaddvq_s32(a3),
10190        ]
10191    }
10192}
10193
10194/// One q8 row range via SDOT (4-row blocks + tail) — the body of
10195/// `qmatvec`'s hot loop, extracted so multi-matrix jobs can drive the
10196/// SAME kernel for several tensors under one pool dispatch. `rep` — the
10197/// load-time interleaved repack (empty = mmap layout only); rows outside
10198/// full 4-row groups always come from the mmap layout.
10199#[cfg(target_arch = "aarch64")]
10200fn q8_range_sdot(
10201    q: &[u8],
10202    rep: &[u8],
10203    row_scale: &[f32],
10204    act: &SplitAct,
10205    cols: usize,
10206    out_addr: SendMut,
10207    start: usize,
10208    end: usize,
10209) {
10210    let mut o = start;
10211    // Leading rows to the group boundary (repack path only): the pool
10212    // splits row ranges arbitrarily, groups are absolute.
10213    if !rep.is_empty() {
10214        while o < end && o % 4 != 0 {
10215            let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
10216            unsafe { *out_addr.at(o) = v };
10217            o += 1;
10218        }
10219    }
10220    while o + 4 <= end {
10221        let r = if rep.is_empty() {
10222            unsafe {
10223                dot_i8_sdot_4rows(
10224                    &q[o * cols..(o + 1) * cols],
10225                    &q[(o + 1) * cols..(o + 2) * cols],
10226                    &q[(o + 2) * cols..(o + 3) * cols],
10227                    &q[(o + 3) * cols..(o + 4) * cols],
10228                    &act.xq,
10229                )
10230            }
10231        } else {
10232            unsafe { dot_i8_sdot_4rows_il(&rep[o * cols..(o + 4) * cols], &act.xq) }
10233        };
10234        for k in 0..4 {
10235            let mut acc = r[k] as f32 * act.sx;
10236            for &(j, xv) in &act.outliers {
10237                acc += (q[(o + k) * cols + j] as i8) as f32 * xv;
10238            }
10239            // SAFETY: disjoint row ranges per worker.
10240            unsafe { *out_addr.at(o + k) = acc * row_scale[o + k] };
10241        }
10242        o += 4;
10243    }
10244    while o < end {
10245        let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
10246        unsafe { *out_addr.at(o) = v };
10247        o += 1;
10248    }
10249}
10250
10251/// Two-input q8 row range via SDOT — `qmatvec2`'s hot loop, extracted
10252/// for the fused pair multi-matrix job (`matvec2_many`).
10253#[cfg(target_arch = "aarch64")]
10254#[allow(clippy::too_many_arguments)]
10255fn q8_range2_sdot(
10256    q: &[u8],
10257    row_scale: &[f32],
10258    a1: &SplitAct,
10259    a2: &SplitAct,
10260    cols: usize,
10261    p1: SendMut,
10262    p2: SendMut,
10263    start: usize,
10264    end: usize,
10265) {
10266    for o in start..end {
10267        let row = &q[o * cols..(o + 1) * cols];
10268        // SAFETY: disjoint row ranges per worker.
10269        unsafe {
10270            *p1.at(o) = row_dot_sdot(row, a1) * row_scale[o];
10271            *p2.at(o) = row_dot_sdot(row, a2) * row_scale[o];
10272        }
10273    }
10274}
10275
10276/// Two-input q8 row range, f32 kernel (non-SDOT) — same extraction.
10277#[allow(clippy::too_many_arguments)]
10278fn q8_range2_f32(
10279    q: &[u8],
10280    row_scale: &[f32],
10281    x1: &[f32],
10282    x2: &[f32],
10283    cols: usize,
10284    p1: SendMut,
10285    p2: SendMut,
10286    start: usize,
10287    end: usize,
10288) {
10289    for o in start..end {
10290        let row = &q[o * cols..(o + 1) * cols];
10291        // SAFETY: disjoint row ranges per worker.
10292        unsafe {
10293            *p1.at(o) = dot_i8_f32(row, x1) * row_scale[o];
10294            *p2.at(o) = dot_i8_f32(row, x2) * row_scale[o];
10295        }
10296    }
10297}
10298
10299/// Scalar/NEON-f32 q8 row range (non-SDOT platforms) — same extraction.
10300fn q8_range_f32(
10301    q: &[u8],
10302    row_scale: &[f32],
10303    xs: &[f32],
10304    cols: usize,
10305    out_addr: SendMut,
10306    start: usize,
10307    end: usize,
10308) {
10309    for o in start..end {
10310        let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
10311        // SAFETY: disjoint row ranges per worker.
10312        unsafe { *out_addr.at(o) = v };
10313    }
10314}
10315
10316/// One q8 row against a split activation, portable: the per-arch fast
10317/// dots where they exist, the exact scalar loop elsewhere. The scalar
10318/// arm is also the test oracle for both fast arms.
10319#[inline]
10320fn q8_row_dot(row: &[u8], act: &SplitAct) -> f32 {
10321    #[cfg(target_arch = "aarch64")]
10322    return row_dot_sdot(row, act);
10323    #[cfg(target_arch = "x86_64")]
10324    return row_dot_avx2(row, act);
10325    #[allow(unreachable_code)]
10326    q8_row_dot_scalar(row, act)
10327}
10328
10329#[allow(dead_code)]
10330fn q8_row_dot_scalar(row: &[u8], act: &SplitAct) -> f32 {
10331    let mut acc = 0i32;
10332    for (k, &b) in row.iter().enumerate() {
10333        acc += (b as i8) as i32 * act.xq[k] as i32;
10334    }
10335    let mut acc = acc as f32 * act.sx;
10336    for &(j, xv) in &act.outliers {
10337        acc += (row[j] as i8) as f32 * xv;
10338    }
10339    acc
10340}
10341
10342/// SDOT row dot with exact outlier correction:
10343/// `dot = sdot(w, xq)·sx + Σ_outl w[j]·x[j]` (then × row_scale by caller).
10344#[cfg(target_arch = "aarch64")]
10345#[inline]
10346fn row_dot_sdot(row: &[u8], act: &SplitAct) -> f32 {
10347    let mut acc = unsafe { dot_i8_sdot(row, &act.xq) } as f32 * act.sx;
10348    for &(j, xv) in &act.outliers {
10349        acc += (row[j] as i8) as f32 * xv;
10350    }
10351    acc
10352}
10353
10354/// One q4 row via SDOT: each 32-group's nibbles unpack to centered i8
10355/// (nib−8 ∈ [−8,7]), int8×int8 `sdot` against the pre-quantized
10356/// activation group, × the group's f16 scale. Returns Σ_g dot_g·s_g;
10357/// the caller multiplies by the activation scale and adds the exact
10358/// outlier terms (port of vmfcore `dot_q4_block_sdot`, +23% measured).
10359/// Nibble order matches the writer: element 2k = low nibble, 2k+1 = high
10360/// → zip(lo,hi) restores flat order.
10361#[cfg(target_arch = "aarch64")]
10362#[target_feature(enable = "neon,dotprod")]
10363unsafe fn dot_q4_row_sdot(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
10364    // SAFETY: callers uphold slice-length contracts (16 packed bytes and
10365    // 2 scale bytes per group; xq.len() == gpr·GROUP_SIZE).
10366    unsafe {
10367        use core::arch::aarch64::*;
10368        use core::arch::asm;
10369        let lomask = vdupq_n_u8(0x0F);
10370        let eight = vdupq_n_s8(8);
10371        let mut acc = 0f32;
10372        for gi in 0..gpr {
10373            let g = g0 + gi;
10374            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
10375            let b = vld1q_u8(packed.as_ptr().add(g * 16));
10376            let lo = vandq_u8(b, lomask);
10377            let hi = vshrq_n_u8::<4>(b);
10378            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
10379            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
10380            let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
10381            let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
10382            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
10383            asm!(
10384                "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
10385                "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
10386                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
10387                e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
10388                options(pure, nomem, nostack),
10389            );
10390            acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
10391        }
10392        acc
10393    }
10394}
10395
10396/// Two-activation q4 row via SDOT: the nibble unpack (the expensive
10397/// part) happens ONCE per group; both pre-quantized activations are
10398/// dotted against the same centered i8 registers. Per-lane math matches
10399/// `dot_q4_row_sdot` exactly.
10400#[cfg(target_arch = "aarch64")]
10401#[target_feature(enable = "neon,dotprod")]
10402unsafe fn dot_q4_row_sdot2(
10403    packed: &[u8],
10404    scales: &[u8],
10405    g0: usize,
10406    gpr: usize,
10407    xq1: &[i8],
10408    xq2: &[i8],
10409) -> (f32, f32) {
10410    // SAFETY: callers uphold slice-length contracts (16 packed bytes and
10411    // 2 scale bytes per group; xq*.len() == gpr·GROUP_SIZE).
10412    unsafe {
10413        use core::arch::aarch64::*;
10414        use core::arch::asm;
10415        let lomask = vdupq_n_u8(0x0F);
10416        let eight = vdupq_n_s8(8);
10417        let (mut acc1, mut acc2) = (0f32, 0f32);
10418        for gi in 0..gpr {
10419            let g = g0 + gi;
10420            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
10421            let b = vld1q_u8(packed.as_ptr().add(g * 16));
10422            let lo = vandq_u8(b, lomask);
10423            let hi = vshrq_n_u8::<4>(b);
10424            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
10425            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
10426            let x10 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE));
10427            let x11 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE + 16));
10428            let x20 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE));
10429            let x21 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE + 16));
10430            let (mut a0, mut a1, mut b0, mut b1) = (
10431                vdupq_n_s32(0),
10432                vdupq_n_s32(0),
10433                vdupq_n_s32(0),
10434                vdupq_n_s32(0),
10435            );
10436            asm!(
10437                "sdot {a0:v}.4s, {e0:v}.16b, {x10:v}.16b",
10438                "sdot {a1:v}.4s, {e1:v}.16b, {x11:v}.16b",
10439                "sdot {b0:v}.4s, {e0:v}.16b, {x20:v}.16b",
10440                "sdot {b1:v}.4s, {e1:v}.16b, {x21:v}.16b",
10441                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
10442                b0 = inout(vreg) b0, b1 = inout(vreg) b1,
10443                e0 = in(vreg) e0, e1 = in(vreg) e1,
10444                x10 = in(vreg) x10, x11 = in(vreg) x11,
10445                x20 = in(vreg) x20, x21 = in(vreg) x21,
10446                options(pure, nomem, nostack),
10447            );
10448            acc1 += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
10449            acc2 += vaddvq_s32(vaddq_s32(b0, b1)) as f32 * s;
10450        }
10451        (acc1, acc2)
10452    }
10453}
10454
10455// ───────────────────── fused int8 kernels ─────────────────────
10456
10457/// `acc += w · row` where the row is centered i8 — NEON widen+fma on
10458/// aarch64, scalar elsewhere. The KV-cache q8 value path rides on this.
10459#[inline]
10460pub(crate) fn axpy_i8_f32(acc: &mut [f32], row: &[i8], w: f32) {
10461    #[cfg(target_arch = "aarch64")]
10462    unsafe {
10463        return axpy_i8_f32_neon(acc, row, w);
10464    }
10465    #[cfg(target_arch = "x86_64")]
10466    if avx2_enabled() {
10467        return unsafe { axpy_i8_f32_avx2(acc, row, w) };
10468    }
10469    #[allow(unreachable_code)]
10470    {
10471        for (a, &b) in acc.iter_mut().zip(row) {
10472            *a += w * b as f32;
10473        }
10474    }
10475}
10476
10477/// i8→f32 axpy via AVX2/FMA (x86 mirror of `axpy_i8_f32_neon`).
10478#[cfg(target_arch = "x86_64")]
10479#[target_feature(enable = "avx2,fma")]
10480unsafe fn axpy_i8_f32_avx2(acc: &mut [f32], row: &[i8], w: f32) {
10481    // SAFETY: callers uphold slice-length contracts (see call sites).
10482    unsafe {
10483        use core::arch::x86_64::*;
10484        let n = acc.len().min(row.len());
10485        let ap = acc.as_mut_ptr();
10486        let rp = row.as_ptr();
10487        let wv = _mm256_set1_ps(w);
10488        let mut j = 0usize;
10489        while j + 16 <= n {
10490            let rb = _mm_loadu_si128(rp.add(j) as *const __m128i);
10491            let lo = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(rb));
10492            let hi = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(rb)));
10493            let v0 = _mm256_fmadd_ps(wv, lo, _mm256_loadu_ps(ap.add(j)));
10494            let v1 = _mm256_fmadd_ps(wv, hi, _mm256_loadu_ps(ap.add(j + 8)));
10495            _mm256_storeu_ps(ap.add(j), v0);
10496            _mm256_storeu_ps(ap.add(j + 8), v1);
10497            j += 16;
10498        }
10499        while j < n {
10500            *ap.add(j) += w * (*rp.add(j)) as f32;
10501            j += 1;
10502        }
10503    }
10504}
10505
10506#[cfg(target_arch = "aarch64")]
10507#[target_feature(enable = "neon")]
10508unsafe fn axpy_i8_f32_neon(acc: &mut [f32], row: &[i8], w: f32) {
10509    // SAFETY: callers uphold slice-length contracts (see call sites).
10510    unsafe {
10511        use core::arch::aarch64::*;
10512        let n = acc.len().min(row.len());
10513        let ap = acc.as_mut_ptr();
10514        let rp = row.as_ptr();
10515        let wv = vdupq_n_f32(w);
10516        let mut j = 0usize;
10517        while j + 16 <= n {
10518            let rb = vld1q_s8(rp.add(j));
10519            let lo = vmovl_s8(vget_low_s8(rb));
10520            let hi = vmovl_s8(vget_high_s8(rb));
10521            for (off, half) in [(0, lo), (8, hi)] {
10522                let f0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(half)));
10523                let f1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(half)));
10524                let o = j + off;
10525                vst1q_f32(ap.add(o), vfmaq_f32(vld1q_f32(ap.add(o)), wv, f0));
10526                vst1q_f32(ap.add(o + 4), vfmaq_f32(vld1q_f32(ap.add(o + 4)), wv, f1));
10527            }
10528            j += 16;
10529        }
10530        while j < n {
10531            *ap.add(j) += w * (*rp.add(j)) as f32;
10532            j += 1;
10533        }
10534    }
10535}
10536
10537/// i8 row · f32 x. NEON on aarch64 (ported from vmfcore `dot_i8_f32_neon`,
10538/// ≈9× scalar), scalar elsewhere.
10539#[inline]
10540pub(crate) fn dot_i8_f32(w: &[u8], x: &[f32]) -> f32 {
10541    #[cfg(target_arch = "aarch64")]
10542    unsafe {
10543        return dot_i8_f32_neon(w, x);
10544    }
10545    #[cfg(target_arch = "x86_64")]
10546    if avx2_enabled() {
10547        return unsafe { dot_i8_f32_avx2(w, x) };
10548    }
10549    #[allow(unreachable_code)]
10550    {
10551        let mut sum = 0.0f32;
10552        for (j, &b) in w.iter().enumerate() {
10553            sum += (b as i8) as f32 * x[j];
10554        }
10555        sum
10556    }
10557}
10558
10559/// i8 row · (x ⊙ col_field) — the q8_2f row dot with the θ col-field
10560/// folded into the product (no prescaled copy of x). NEON on aarch64,
10561/// scalar elsewhere. Used by the active-neuron path `row_dot`.
10562#[inline]
10563fn dot_i8_col_f32(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
10564    #[cfg(target_arch = "aarch64")]
10565    unsafe {
10566        return dot_i8_col_f32_neon(w, x, col);
10567    }
10568    #[allow(unreachable_code)]
10569    {
10570        let mut sum = 0.0f32;
10571        for (j, &b) in w.iter().enumerate() {
10572            sum += (b as i8) as f32 * x[j] * col[j];
10573        }
10574        sum
10575    }
10576}
10577
10578#[cfg(target_arch = "aarch64")]
10579#[target_feature(enable = "neon")]
10580unsafe fn dot_i8_col_f32_neon(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
10581    // SAFETY: callers uphold slice-length contracts (see call sites).
10582    unsafe {
10583        use core::arch::aarch64::*;
10584        let n = x.len();
10585        let wp = w.as_ptr() as *const i8;
10586        let xp = x.as_ptr();
10587        let cp = col.as_ptr();
10588        let (mut a0, mut a1, mut a2, mut a3) = (
10589            vdupq_n_f32(0.0),
10590            vdupq_n_f32(0.0),
10591            vdupq_n_f32(0.0),
10592            vdupq_n_f32(0.0),
10593        );
10594        let mut j = 0usize;
10595        while j + 16 <= n {
10596            let wb = vld1q_s8(wp.add(j));
10597            let lo = vmovl_s8(vget_low_s8(wb));
10598            let hi = vmovl_s8(vget_high_s8(wb));
10599            let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
10600            let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
10601            let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
10602            let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
10603            a0 = vfmaq_f32(
10604                a0,
10605                w0,
10606                vmulq_f32(vld1q_f32(xp.add(j)), vld1q_f32(cp.add(j))),
10607            );
10608            a1 = vfmaq_f32(
10609                a1,
10610                w1,
10611                vmulq_f32(vld1q_f32(xp.add(j + 4)), vld1q_f32(cp.add(j + 4))),
10612            );
10613            a2 = vfmaq_f32(
10614                a2,
10615                w2,
10616                vmulq_f32(vld1q_f32(xp.add(j + 8)), vld1q_f32(cp.add(j + 8))),
10617            );
10618            a3 = vfmaq_f32(
10619                a3,
10620                w3,
10621                vmulq_f32(vld1q_f32(xp.add(j + 12)), vld1q_f32(cp.add(j + 12))),
10622            );
10623            j += 16;
10624        }
10625        let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
10626        while j < n {
10627            sum += (*wp.add(j)) as f32 * *xp.add(j) * *cp.add(j);
10628            j += 1;
10629        }
10630        sum
10631    }
10632}
10633
10634#[cfg(target_arch = "aarch64")]
10635#[target_feature(enable = "neon")]
10636unsafe fn dot_i8_f32_neon(w: &[u8], x: &[f32]) -> f32 {
10637    // SAFETY: callers uphold slice-length contracts (see call sites).
10638    unsafe {
10639        use core::arch::aarch64::*;
10640        let n = x.len();
10641        let wp = w.as_ptr() as *const i8;
10642        let xp = x.as_ptr();
10643        let (mut a0, mut a1, mut a2, mut a3) = (
10644            vdupq_n_f32(0.0),
10645            vdupq_n_f32(0.0),
10646            vdupq_n_f32(0.0),
10647            vdupq_n_f32(0.0),
10648        );
10649        let mut j = 0usize;
10650        while j + 16 <= n {
10651            let wb = vld1q_s8(wp.add(j));
10652            let lo = vmovl_s8(vget_low_s8(wb));
10653            let hi = vmovl_s8(vget_high_s8(wb));
10654            let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
10655            let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
10656            let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
10657            let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
10658            a0 = vfmaq_f32(a0, w0, vld1q_f32(xp.add(j)));
10659            a1 = vfmaq_f32(a1, w1, vld1q_f32(xp.add(j + 4)));
10660            a2 = vfmaq_f32(a2, w2, vld1q_f32(xp.add(j + 8)));
10661            a3 = vfmaq_f32(a3, w3, vld1q_f32(xp.add(j + 12)));
10662            j += 16;
10663        }
10664        let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
10665        while j < n {
10666            sum += (*wp.add(j)) as f32 * *xp.add(j);
10667            j += 1;
10668        }
10669        sum
10670    }
10671}
10672
10673#[allow(clippy::too_many_arguments)]
10674fn qmatvec(
10675    q: &[u8],
10676    rep: &[u8],
10677    row_scale: &[f32],
10678    x: &[f32],
10679    col_field: &[f32],
10680    dtype: TensorDtype,
10681    rows: usize,
10682    cols: usize,
10683    out: &mut [f32],
10684    pool: Option<&Pool>,
10685) {
10686    debug_assert_eq!(out.len(), rows);
10687    #[cfg(not(target_arch = "aarch64"))]
10688    let _ = rep;
10689
10690    #[cfg(target_arch = "aarch64")]
10691    if sdot_enabled() {
10692        let act = if dtype == TensorDtype::Q8_2f {
10693            split_act_q8_2f(x, col_field)
10694        } else {
10695            split_act(x)
10696        };
10697        let out_addr = SendMut(out.as_mut_ptr());
10698        let run_range = |start: usize, end: usize| {
10699            q8_range_sdot(q, rep, row_scale, &act, cols, out_addr, start, end)
10700        };
10701        match pool {
10702            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
10703            _ => run_range(0, rows),
10704        }
10705        return;
10706    }
10707    // x86 A8W8 via AVX2 maddubs — same quantized-activation contract as
10708    // the SDOT path (CMF_AVX2=0 keeps the exact i8×f32 loop).
10709    #[cfg(target_arch = "x86_64")]
10710    if avx2_a8w8_enabled() {
10711        let act = if dtype == TensorDtype::Q8_2f {
10712            split_act_q8_2f(x, col_field)
10713        } else {
10714            split_act(x)
10715        };
10716        let out_addr = SendMut(out.as_mut_ptr());
10717        let run_range = |start: usize, end: usize| {
10718            q8_range_avx2(q, row_scale, &act, cols, out_addr, start, end)
10719        };
10720        match pool {
10721            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
10722            _ => run_range(0, rows),
10723        }
10724        return;
10725    }
10726
10727    prescale_with(x, col_field, dtype, 1, |xs| {
10728        let out_addr = SendMut(out.as_mut_ptr());
10729        let run_range = move |start: usize, end: usize| {
10730            for o in start..end {
10731                let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
10732                // SAFETY: disjoint row ranges per worker.
10733                unsafe { *out_addr.at(o) = v };
10734            }
10735        };
10736        match pool {
10737            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
10738            _ => run_range(0, rows),
10739        }
10740    });
10741}
10742
10743#[allow(clippy::too_many_arguments)]
10744fn qmatvec2(
10745    q: &[u8],
10746    row_scale: &[f32],
10747    x1: &[f32],
10748    x2: &[f32],
10749    col_field: &[f32],
10750    dtype: TensorDtype,
10751    rows: usize,
10752    cols: usize,
10753    o1: &mut [f32],
10754    o2: &mut [f32],
10755    pool: Option<&Pool>,
10756) {
10757    #[cfg(target_arch = "aarch64")]
10758    if sdot_enabled() {
10759        let a1s = if dtype == TensorDtype::Q8_2f {
10760            split_act_q8_2f(x1, col_field)
10761        } else {
10762            split_act(x1)
10763        };
10764        let a2s = if dtype == TensorDtype::Q8_2f {
10765            split_act_q8_2f(x2, col_field)
10766        } else {
10767            split_act(x2)
10768        };
10769        let p1 = SendMut(o1.as_mut_ptr());
10770        let p2 = SendMut(o2.as_mut_ptr());
10771        let run_range = |start: usize, end: usize| {
10772            q8_range2_sdot(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
10773        };
10774        match pool {
10775            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
10776            _ => run_range(0, rows),
10777        }
10778        return;
10779    }
10780    #[cfg(target_arch = "x86_64")]
10781    if avx2_a8w8_enabled() {
10782        let a1s = if dtype == TensorDtype::Q8_2f {
10783            split_act_q8_2f(x1, col_field)
10784        } else {
10785            split_act(x1)
10786        };
10787        let a2s = if dtype == TensorDtype::Q8_2f {
10788            split_act_q8_2f(x2, col_field)
10789        } else {
10790            split_act(x2)
10791        };
10792        let p1 = SendMut(o1.as_mut_ptr());
10793        let p2 = SendMut(o2.as_mut_ptr());
10794        let run_range = |start: usize, end: usize| {
10795            q8_range2_avx2(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
10796        };
10797        match pool {
10798            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
10799            _ => run_range(0, rows),
10800        }
10801        return;
10802    }
10803
10804    prescale_with(x1, col_field, dtype, 1, |x1s| {
10805        prescale_with(x2, col_field, dtype, 2, |x2s| {
10806            let p1 = SendMut(o1.as_mut_ptr());
10807            let p2 = SendMut(o2.as_mut_ptr());
10808            let run_range = move |start: usize, end: usize| {
10809                for o in start..end {
10810                    let row = &q[o * cols..(o + 1) * cols];
10811                    let s1 = dot_i8_f32(row, x1s) * row_scale[o];
10812                    let s2 = dot_i8_f32(row, x2s) * row_scale[o];
10813                    // SAFETY: disjoint row ranges per worker.
10814                    unsafe {
10815                        *p1.at(o) = s1;
10816                        *p2.at(o) = s2;
10817                    }
10818                }
10819            };
10820            match pool {
10821                Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
10822                _ => run_range(0, rows),
10823            }
10824        });
10825    });
10826}
10827
10828#[derive(Clone, Copy)]
10829struct SendMut(*mut f32);
10830unsafe impl Send for SendMut {}
10831unsafe impl Sync for SendMut {}
10832
10833impl SendMut {
10834    #[inline]
10835    fn at(self, i: usize) -> *mut f32 {
10836        unsafe { self.0.add(i) }
10837    }
10838}
10839
10840#[cfg(test)]
10841mod tests {
10842    use super::*;
10843
10844    #[test]
10845    fn q2tp_i8_dot_matches_exact_on_grid() {
10846        // On-grid activations (±1 → sx=1/127, xq=±127 dequantizes
10847        // exactly, no outliers) must make the integer path agree with
10848        // the exact scalar walk to f32 rounding.
10849        let (rows, cols) = (5, 64);
10850        let gpr = cols / GROUP_SIZE;
10851        // Synthetic codes plane + a flat ladder: scales_into is not under
10852        // test here, so drive dot_q2tp_row_i8 / q2tp_row_exact directly
10853        // with hand-made scales.
10854        let chunks: Vec<u8> = (0..rows * gpr * Q2TP_CHUNK)
10855            .map(|i| (i as u32).wrapping_mul(2654435761) as u8)
10856            .collect();
10857        let scales: Vec<f32> = (0..gpr).map(|g| 0.5 + g as f32 * 0.25).collect();
10858        let x: Vec<f32> = (0..cols)
10859            .map(|i| if i % 3 == 0 { -1.0 } else { 1.0 })
10860            .collect();
10861        let act = split_act(&x);
10862        assert!(
10863            act.outliers.is_empty(),
10864            "on-grid input must have no outliers"
10865        );
10866        let gsum = q1_group_sums(&act.xq, gpr);
10867        for r in 0..rows {
10868            let exact = q2tp_row_exact(&chunks, r, gpr, &x, &scales);
10869            let fast = dot_q2tp_row_i8(&chunks, r, gpr, &act.xq, &gsum, &scales) * act.sx;
10870            assert!(
10871                (exact - fast).abs() <= exact.abs() * 1e-5 + 1e-5,
10872                "row {r}: exact {exact} vs i8 {fast}"
10873            );
10874        }
10875    }
10876
10877    #[test]
10878    fn q2tp_affine_fuses_half_scale_correction_without_changing_raw_decode() {
10879        let (rows, cols) = (1usize, GROUP_SIZE);
10880        let mut bytes = vec![0u8; Q2TP_CHUNK + 4 + 1];
10881        // Repeating symbols 0,1,2,0 at unit scale.  q2tp's raw B is
10882        // (c-1.5), while the affine Prism operator is (c-1.0).
10883        bytes[..Q2TP_CHUNK].fill(0x24); // codes 0,1,2,0 in LSB-first order
10884        bytes[Q2TP_CHUNK..Q2TP_CHUNK + 2].copy_from_slice(&0u16.to_le_bytes());
10885        bytes[Q2TP_CHUNK + 2..Q2TP_CHUNK + 4].copy_from_slice(&0u16.to_le_bytes());
10886        bytes[Q2TP_CHUNK + 4] = 1; // dtype16 rung 1 = 1.0
10887        let x = vec![1.0f32; cols];
10888        let mut raw = vec![0.0f32; rows];
10889        let mut affine = vec![0.0f32; rows];
10890        q2tp_matvec_for_test(&bytes, &x, rows, cols, &mut raw);
10891        q2tp_affine_matvec_for_test(&bytes, &x, rows, cols, &mut affine);
10892        assert_eq!(raw, vec![-24.0]);
10893        assert_eq!(affine, vec![-8.0]);
10894        assert!((affine[0] - (raw[0] + 0.5 * cols as f32)).abs() < 1e-6);
10895    }
10896
10897    #[cfg(target_arch = "x86_64")]
10898    #[test]
10899    fn q2tp_avx2_dot_matches_scalar_for_random_patterns() {
10900        // Compare the release AVX2 integer dot against the scalar oracle over
10901        // arbitrary packed bytes/activation signs.  This guards the exact
10902        // table-load path used after rejecting a faster-looking decoder whose
10903        // full-checkpoint greedy output drifted.
10904        if !std::arch::is_x86_feature_detected!("avx2") {
10905            return;
10906        }
10907        let mut seed = 0x9e3779b9u32;
10908        let mut next = || {
10909            seed = seed.wrapping_mul(1664525).wrapping_add(1013904223);
10910            seed
10911        };
10912        for _ in 0..20_000 {
10913            let mut ch = [0u8; Q2TP_CHUNK];
10914            let mut x = [0i8; GROUP_SIZE];
10915            for b in &mut ch {
10916                *b = next() as u8;
10917            }
10918            for v in &mut x {
10919                *v = (next() >> 24) as i8;
10920            }
10921            let mut reference = 0i32;
10922            for (k, &b) in ch.iter().enumerate() {
10923                reference += (b & 3) as i32 * x[k * 4] as i32;
10924                reference += ((b >> 2) & 3) as i32 * x[k * 4 + 1] as i32;
10925                reference += ((b >> 4) & 3) as i32 * x[k * 4 + 2] as i32;
10926                reference += ((b >> 6) & 3) as i32 * x[k * 4 + 3] as i32;
10927            }
10928            // SAFETY: guarded by the runtime AVX2 feature check and fixed
10929            // 8-byte/32-byte slice lengths above.
10930            let got = unsafe { q2tp_code_dot_avx2(&ch, &x) };
10931            assert_eq!(got, reference, "packed q2 lane mismatch");
10932        }
10933    }
10934
10935    #[test]
10936    fn q8_row_dot_fast_matches_scalar() {
10937        // The per-arch fast dot must agree with the exact scalar oracle
10938        // (same contract the fused q8 FFN arm rides on).
10939        let cols = 96;
10940        let row: Vec<u8> = (0..cols)
10941            .map(|i| ((i * 37 % 251) - 125) as i8 as u8)
10942            .collect();
10943        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.13).sin()).collect();
10944        let act = split_act(&x);
10945        let fast = q8_row_dot(&row, &act);
10946        let scalar = q8_row_dot_scalar(&row, &act);
10947        assert!(
10948            (fast - scalar).abs() <= scalar.abs() * 1e-5 + 1e-5,
10949            "fast {fast} vs scalar {scalar}"
10950        );
10951    }
10952
10953    #[test]
10954    fn f32_matvec_matches_matvec_rows_bitexact() {
10955        let (rows, cols) = (300, 40);
10956        let w: Vec<f32> = (0..rows * cols).map(|i| (i as f32 * 0.017).sin()).collect();
10957        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.05).cos()).collect();
10958        let qt = QTensor::from_f32(w.clone(), rows, cols);
10959
10960        let mut a = vec![0.0f32; rows];
10961        matvec_rows(None, &w, &x, &mut a);
10962        let mut b = vec![0.0f32; rows];
10963        qt.matvec(&x, &mut b, None);
10964        assert_eq!(a, b);
10965    }
10966
10967    #[test]
10968    fn sdot_kernel_exact_on_grid() {
10969        // Activations already on the i8 grid (±1 with amax=1 → sx=1/127,
10970        // xq=±127 dequantizes EXACTLY) → the SDOT path must match the
10971        // exact f32 dot to float rounding. This isolates kernel
10972        // correctness from quantization noise.
10973        eprintln!("sdot_enabled = {}", sdot_enabled());
10974        let (rows, cols) = (9, 80); // odd rows → exercises 4-row + tail
10975        let w: Vec<u8> = (0..rows * cols)
10976            .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
10977            .collect();
10978        let scales: Vec<f32> = (0..rows).map(|o| 0.005 + o as f32 * 0.001).collect();
10979        let x: Vec<f32> = (0..cols)
10980            .map(|i| match i % 3 {
10981                0 => 1.0,
10982                1 => -1.0,
10983                _ => 0.0,
10984            })
10985            .collect();
10986        let mut a = vec![0.0f32; rows];
10987        qmatvec(
10988            &w,
10989            &[],
10990            &scales,
10991            &x,
10992            &[],
10993            TensorDtype::Q8Row,
10994            rows,
10995            cols,
10996            &mut a,
10997            None,
10998        );
10999        for o in 0..rows {
11000            let mut acc = 0.0f32;
11001            for j in 0..cols {
11002                acc += (w[o * cols + j] as i8) as f32 * x[j];
11003            }
11004            let expect = acc * scales[o];
11005            assert!(
11006                (a[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
11007                "row {o}: {} vs {expect}",
11008                a[o]
11009            );
11010        }
11011    }
11012
11013    #[test]
11014    fn q1_tbl_fast_path_matches_reference() {
11015        // gpr = 8 exercises the TBL pair-load fast loop, and the LAST
11016        // row's final 4-tile window trips the 4B-overread guard (the
11017        // payload ends exactly at the last tile) — both paths must
11018        // agree with the dequant reference.
11019        let (rows, cols) = (5, 256);
11020        let gpr = cols / GROUP_SIZE;
11021        let mut bytes = Vec::new();
11022        for t in 0..rows * gpr {
11023            let s = 0.007 + (t % 11) as f32 * 0.004;
11024            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
11025            for j in 0..4 {
11026                bytes.push(((t * 53 + j * 89 + 7) % 249) as u8);
11027            }
11028        }
11029        let x: Vec<f32> = (0..cols)
11030            .map(|i| if (i * 5) % 7 < 3 { 1.0 } else { -1.0 })
11031            .collect();
11032        let mut w = vec![0.0f32; rows * cols];
11033        cortiq_core::quant::dequant_q1(&bytes, &mut w);
11034        let mut got = vec![0.0f32; rows];
11035        q1_matvec(&bytes, &x, rows, cols, &mut got, None);
11036        for o in 0..rows {
11037            let expect: f32 = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
11038            assert!(
11039                (got[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
11040                "row {o}: {} vs {expect}",
11041                got[o]
11042            );
11043        }
11044        // Blocked 1×4 batch (b=5: one quad + remainder) must equal the
11045        // single-matvec path bit-for-bit.
11046        let b = 5usize;
11047        let mut xs_all = Vec::new();
11048        for bi in 0..b {
11049            xs_all.extend(x.iter().map(|v| if bi % 2 == 0 { *v } else { -*v }));
11050        }
11051        let mut mm = vec![0.0f32; b * rows];
11052        q1_matmat(&bytes, &xs_all, b, rows, cols, &mut mm, None);
11053        for bi in 0..b {
11054            let mut single = vec![0.0f32; rows];
11055            q1_matvec(
11056                &bytes,
11057                &xs_all[bi * cols..(bi + 1) * cols],
11058                rows,
11059                cols,
11060                &mut single,
11061                None,
11062            );
11063            assert_eq!(&mm[bi * rows..(bi + 1) * rows], &single[..], "stream {bi}");
11064        }
11065    }
11066
11067    #[test]
11068    fn q1_kernels_match_exact_reference() {
11069        // Synthetic q1 payload: 6-byte tiles [f16 scale][4B bits].
11070        let (rows, cols) = (7, 96);
11071        let gpr = cols / GROUP_SIZE;
11072        let mut bytes = Vec::new();
11073        for t in 0..rows * gpr {
11074            let s = 0.01 + (t % 13) as f32 * 0.003;
11075            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
11076            for j in 0..4 {
11077                bytes.push(((t * 31 + j * 97) % 251) as u8);
11078            }
11079        }
11080        // On-grid activations (±1, amax 1) → the SDOT path is exact.
11081        let x: Vec<f32> = (0..cols)
11082            .map(|i| if i % 3 == 0 { 1.0 } else { -1.0 })
11083            .collect();
11084        // Reference through the core dequant.
11085        let mut w = vec![0.0f32; rows * cols];
11086        cortiq_core::quant::dequant_q1(&bytes, &mut w);
11087        let mut expect = vec![0.0f32; rows];
11088        for o in 0..rows {
11089            expect[o] = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
11090        }
11091        let mut got = vec![0.0f32; rows];
11092        q1_matvec(&bytes, &x, rows, cols, &mut got, None);
11093        for o in 0..rows {
11094            assert!(
11095                (got[o] - expect[o]).abs() < 1e-3 * expect[o].abs().max(1e-3),
11096                "row {o}: {} vs {}",
11097                got[o],
11098                expect[o]
11099            );
11100        }
11101        // Pair and batch paths agree with the single path.
11102        let x2: Vec<f32> = x.iter().map(|v| -v).collect();
11103        let (mut a1, mut a2) = (vec![0.0f32; rows], vec![0.0f32; rows]);
11104        q1_matvec2(&bytes, &x, &x2, rows, cols, &mut a1, &mut a2, None);
11105        assert_eq!(a1, got);
11106        let mut xs = x.clone();
11107        xs.extend_from_slice(&x2);
11108        let mut mm = vec![0.0f32; 2 * rows];
11109        q1_matmat(&bytes, &xs, 2, rows, cols, &mut mm, None);
11110        assert_eq!(&mm[..rows], got.as_slice());
11111        assert_eq!(&mm[rows..], a2.as_slice());
11112    }
11113
11114    #[test]
11115    fn repack_is_bit_identical() {
11116        // The interleaved-repack kernel must produce EXACTLY the same
11117        // bits as the mmap-layout kernel: integer accumulation is order-
11118        // exact, the f32 epilogue is identical. Odd rows exercise the
11119        // tail; direct range calls exercise unaligned pool splits.
11120        let (rows, cols) = (267, 96); // 66 groups + 3 tail rows, cols % 16 == 0
11121        let w: Vec<u8> = (0..rows * cols)
11122            .map(|i| (((i * 89) % 253) as i32 - 126) as i8 as u8)
11123            .collect();
11124        let scales: Vec<f32> = (0..rows).map(|o| 0.003 + o as f32 * 0.0007).collect();
11125        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.37).sin() * 2.0).collect();
11126        let rep = q8_repack_layout(&w, rows, cols);
11127        // Group interleave round-trips.
11128        for g in 0..rows / 4 {
11129            for c in 0..cols / 16 {
11130                for lane in 0..4 {
11131                    assert_eq!(
11132                        &rep[g * 4 * cols + c * 64 + lane * 16
11133                            ..g * 4 * cols + c * 64 + lane * 16 + 16],
11134                        &w[(g * 4 + lane) * cols + c * 16..(g * 4 + lane) * cols + c * 16 + 16],
11135                    );
11136                }
11137            }
11138        }
11139        let mut a = vec![0.0f32; rows];
11140        qmatvec(
11141            &w,
11142            &[],
11143            &scales,
11144            &x,
11145            &[],
11146            TensorDtype::Q8Row,
11147            rows,
11148            cols,
11149            &mut a,
11150            None,
11151        );
11152        let mut b = vec![0.0f32; rows];
11153        qmatvec(
11154            &w,
11155            &rep,
11156            &scales,
11157            &x,
11158            &[],
11159            TensorDtype::Q8Row,
11160            rows,
11161            cols,
11162            &mut b,
11163            None,
11164        );
11165        assert_eq!(a, b, "full-range repack output diverged");
11166
11167        #[cfg(target_arch = "aarch64")]
11168        if sdot_enabled() {
11169            // Unaligned range split (pool workers get arbitrary bounds).
11170            let act = split_act(&x);
11171            let mut c1 = vec![0.0f32; rows];
11172            let mut c2 = vec![0.0f32; rows];
11173            q8_range_sdot(
11174                &w,
11175                &[],
11176                &scales,
11177                &act,
11178                cols,
11179                SendMut(c1.as_mut_ptr()),
11180                3,
11181                rows - 2,
11182            );
11183            q8_range_sdot(
11184                &w,
11185                &rep,
11186                &scales,
11187                &act,
11188                cols,
11189                SendMut(c2.as_mut_ptr()),
11190                3,
11191                rows - 2,
11192            );
11193            assert_eq!(c1, c2, "unaligned-range repack output diverged");
11194        }
11195    }
11196
11197    #[test]
11198    fn sdot_a8w8_noise_is_bounded() {
11199        // Off-grid activations: A8 quantization noise must stay small in
11200        // relative L2 over the whole output (realistic accuracy contract;
11201        // vmfcore measured argmax-identical decode on real models).
11202        let (rows, cols) = (16, 512);
11203        let w: Vec<u8> = (0..rows * cols)
11204            .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
11205            .collect();
11206        let scales = vec![0.01f32; rows];
11207        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.21).sin()).collect();
11208        let mut a = vec![0.0f32; rows];
11209        qmatvec(
11210            &w,
11211            &[],
11212            &scales,
11213            &x,
11214            &[],
11215            TensorDtype::Q8Row,
11216            rows,
11217            cols,
11218            &mut a,
11219            None,
11220        );
11221        let (mut num, mut den) = (0f64, 0f64);
11222        for o in 0..rows {
11223            let mut acc = 0.0f32;
11224            for j in 0..cols {
11225                acc += (w[o * cols + j] as i8) as f32 * x[j];
11226            }
11227            let expect = acc * scales[o];
11228            num += ((a[o] - expect) as f64).powi(2);
11229            den += (expect as f64).powi(2);
11230        }
11231        let rel = (num / den.max(1e-12)).sqrt();
11232        assert!(rel < 0.05, "A8W8 relative L2 error too high: {rel}");
11233    }
11234
11235    #[test]
11236    fn i8_dot_neon_matches_scalar() {
11237        let n = 100;
11238        let w: Vec<u8> = (0..n).map(|i| ((i * 37 + 11) % 251) as u8).collect();
11239        let x: Vec<f32> = (0..n).map(|i| (i as f32 * 0.13).sin()).collect();
11240        let mut scalar = 0.0f32;
11241        for j in 0..n {
11242            scalar += (w[j] as i8) as f32 * x[j];
11243        }
11244        let fast = dot_i8_f32(&w, &x);
11245        assert!((scalar - fast).abs() < 1e-3 * scalar.abs().max(1.0));
11246    }
11247
11248    /// Fused vbit matvec must match full dequant_vbit + dense matvec.
11249    #[test]
11250    fn vbitmatvec_matches_full_dequant() {
11251        let (rows, cols) = (6, 64);
11252        let ng = cols / GROUP_SIZE;
11253        // Hand-craft: bits per row, f16 scales, packed rows.
11254        let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
11255        let mut bytes = bits.clone();
11256        for g in 0..rows * ng {
11257            let s = 0.02 + 0.001 * g as f32;
11258            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
11259        }
11260        for r in 0..rows {
11261            let b = bits[r] as usize;
11262            let (mut acc, mut nb) = (0u64, 0usize);
11263            let mut rowbytes = Vec::new();
11264            for i in 0..cols {
11265                let v = ((i * 7 + r * 13) % (1 << b)) as u64;
11266                acc = (acc << b) | v;
11267                nb += b;
11268                while nb >= 8 {
11269                    nb -= 8;
11270                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
11271                }
11272            }
11273            if nb > 0 {
11274                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
11275            }
11276            bytes.extend_from_slice(&rowbytes);
11277        }
11278        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
11279
11280        let mut reference = vec![0f32; rows * cols];
11281        cortiq_core::quant::dequant_vbit(&bytes, rows, cols, &mut reference).unwrap();
11282        let mut expect = vec![0f32; rows];
11283        for r in 0..rows {
11284            expect[r] = reference[r * cols..(r + 1) * cols]
11285                .iter()
11286                .zip(&x)
11287                .map(|(w, xv)| w * xv)
11288                .sum();
11289        }
11290        let mut got = vec![0f32; rows];
11291        let offsets = vbit_row_offsets(&bytes, rows, cols);
11292        vbitmatvec(&bytes, &offsets, &x, rows, cols, &mut got, None);
11293        // SDOT path quantizes activations to i8 (A8W8): bounded noise,
11294        // same contract as q8 (exact path is pinned by CMF_SDOT=0 in
11295        // the golden-parity gate).
11296        let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
11297        let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
11298        for r in 0..rows {
11299            assert!(
11300                (got[r] - expect[r]).abs() < tol * scale,
11301                "row {r}: {} vs {}",
11302                got[r],
11303                expect[r]
11304            );
11305        }
11306    }
11307
11308    /// Fused q4 matvec must match the reference full-dequant + dense
11309    /// matvec bit-for-bit in structure (same f32 math, group order).
11310    /// vbit matmat: the blocked 1×4 leg must match the per-row path
11311    /// (paired env toggle; larger shape so both code paths engage).
11312    #[test]
11313    #[cfg(target_arch = "x86_64")]
11314    fn vbit_matmat_blocked_matches_per_row() {
11315        let (rows, cols, b) = (64usize, 128usize, 9usize);
11316        let ng = cols / GROUP_SIZE;
11317        let bits: Vec<u8> = (0..rows).map(|r| [3u8, 4, 5, 6][r % 4]).collect();
11318        let mut bytes = bits.clone();
11319        for g in 0..rows * ng {
11320            let sc = 0.02 + 0.0005 * g as f32;
11321            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
11322        }
11323        for r in 0..rows {
11324            let bw = bits[r] as usize;
11325            let (mut acc, mut nb) = (0u64, 0usize);
11326            let mut rowbytes = Vec::new();
11327            for i in 0..cols {
11328                let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
11329                acc = (acc << bw) | v;
11330                nb += bw;
11331                while nb >= 8 {
11332                    nb -= 8;
11333                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
11334                }
11335            }
11336            if nb > 0 {
11337                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
11338            }
11339            bytes.extend_from_slice(&rowbytes);
11340        }
11341        let x: Vec<f32> = (0..b * cols)
11342            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
11343            .collect();
11344        let offsets = vbit_row_offsets(&bytes, rows, cols);
11345        let mut y_a = vec![0f32; b * rows];
11346        let mut y_b = vec![0f32; b * rows];
11347        unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
11348        vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_a, None);
11349        unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
11350        vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_b, None);
11351        unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
11352        let max_d = y_a
11353            .iter()
11354            .zip(&y_b)
11355            .map(|(p, q)| (p - q).abs())
11356            .fold(0.0f32, f32::max);
11357        assert!(max_d < 1e-4, "vbit blocked ≠ per-row: max|Δ| = {max_d}");
11358    }
11359
11360    /// q4t blocked 1×4 (SDOT on ARM, AVX2 on x86) must equal the
11361    /// per-row path exactly: same nibble unpack, same group order,
11362    /// same f32 accumulation — batch == matvec bit-for-bit. b=9 covers
11363    /// two full 1×4 blocks plus a remainder through the single-row
11364    /// kernel. (Both paths produce identical output, so the shared
11365    /// CMF_X86_BLOCKED env var racing with other tests cannot flip
11366    /// the verdict — worst case both sides take the same path.)
11367    #[test]
11368    fn q4t_matmat_blocked_matches_per_row() {
11369        let (rows, cols, b) = (16usize, 64usize, 9usize);
11370        let gpr = cols / GROUP_SIZE;
11371        let mut bytes = vec![0u8; rows * gpr * Q4_TILE];
11372        for r in 0..rows {
11373            for g in 0..gpr {
11374                let t = (r * gpr + g) * Q4_TILE;
11375                let sc = 0.02 + 0.001 * (r * gpr + g) as f32;
11376                bytes[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
11377                for k in 0..16 {
11378                    bytes[t + 2 + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
11379                }
11380            }
11381        }
11382        let x: Vec<f32> = (0..b * cols)
11383            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
11384            .collect();
11385        let mut y_blk = vec![0f32; b * rows];
11386        let mut y_row = vec![0f32; b * rows];
11387        unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
11388        q4t_matmat(&bytes, &x, b, rows, cols, &mut y_blk, None);
11389        unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
11390        q4t_matmat(&bytes, &x, b, rows, cols, &mut y_row, None);
11391        unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
11392        assert_eq!(y_blk, y_row, "q4t blocked 1x4 ≠ per-row");
11393    }
11394
11395    /// The wide-batch Accelerate arm of q4t_matmat vs a brute-force
11396    /// f32 dequant matmul: both are f32 GEMMs, so only reduction
11397    /// order differs — tight tolerance.
11398    /// A synthetic q4tp payload: random nibbles plus a per-row ladder whose
11399    /// span varies row to row, so the codes actually exercise the full 0..31
11400    /// range rather than clustering on one rung.
11401    fn synth_q4tp(rows: usize, cols: usize) -> Vec<u8> {
11402        use cortiq_core::quant::{f32_to_f16, q4tp_code_stride, q4tp_put_code};
11403        let gpr = cols / GROUP_SIZE;
11404        let stride = q4tp_code_stride(gpr);
11405        let (params_off, codes_off, _) = q4tp_sections(rows, cols);
11406        let mut b = vec![0u8; codes_off + rows * stride];
11407        for r in 0..rows {
11408            for g in 0..gpr {
11409                let t = (r * gpr + g) * Q4TP_NIB;
11410                for k in 0..16 {
11411                    b[t + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
11412                }
11413            }
11414            let lo = -6.0 - 0.03 * (r % 17) as f32;
11415            let step = 0.01 + 0.004 * (r % 11) as f32;
11416            let p = params_off + r * 4;
11417            b[p..p + 2].copy_from_slice(&f32_to_f16(lo).to_le_bytes());
11418            b[p + 2..p + 4].copy_from_slice(&f32_to_f16(step).to_le_bytes());
11419            let crow = &mut b[codes_off + r * stride..codes_off + (r + 1) * stride];
11420            for g in 0..gpr {
11421                q4tp_put_code(crow, g, (r * 5 + g * 3) % 32);
11422            }
11423        }
11424        b
11425    }
11426
11427    /// The same weights re-expressed as q4_tiled, so the proven kernel can
11428    /// be the reference: each tile stores the ladder scale its code selects.
11429    /// Only the f16 rounding of that scale separates the two payloads.
11430    fn q4tp_as_q4t(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
11431        let gpr = cols / GROUP_SIZE;
11432        let v = Q4tpView::new(bytes, rows, cols);
11433        let mut out = vec![0u8; rows * gpr * Q4_TILE];
11434        let mut sc = vec![0f32; gpr];
11435        for r in 0..rows {
11436            v.scales_into(r, gpr, &mut sc);
11437            for g in 0..gpr {
11438                let t = (r * gpr + g) * Q4_TILE;
11439                let s = sc[g];
11440                out[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
11441                let src = (r * gpr + g) * Q4TP_NIB;
11442                out[t + 2..t + Q4_TILE].copy_from_slice(&v.nib[src..src + Q4TP_NIB]);
11443            }
11444        }
11445        out
11446    }
11447
11448    /// The exact (`CMF_SDOT=0`) path must reproduce `dequant_q4tp` to f32
11449    /// rounding — that scalar routine is the format's definition, and the
11450    /// kernels re-derive the scale from the ladder independently. Call the
11451    /// row kernel directly: `matmat` picks the int8 arm when a8w8 is on,
11452    /// so routing through it would test the other path by accident.
11453    #[test]
11454    fn q4tp_exact_path_matches_dequant_reference() {
11455        let (rows, cols) = (256usize, 512usize);
11456        let gpr = cols / GROUP_SIZE;
11457        let bytes = synth_q4tp(rows, cols);
11458        let mut w = vec![0f32; rows * cols];
11459        cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
11460
11461        let x: Vec<f32> = (0..cols)
11462            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
11463            .collect();
11464        let v = Q4tpView::new(&bytes, rows, cols);
11465        let mut sc = vec![0f32; gpr];
11466        for r in 0..rows {
11467            v.scales_into(r, gpr, &mut sc);
11468            let got = q4tp_row_exact(v.nib, r, gpr, &x, &sc);
11469            let want: f32 = (0..cols).map(|c| w[r * cols + c] * x[c]).sum();
11470            // These dot products cancel down to ~1e-3 from terms of ~5e-2, so
11471            // the meaningful yardstick is the summed magnitude, not the result:
11472            // against the result any reordering of a 512-term f32 sum "fails".
11473            let mag: f32 = (0..cols).map(|c| (w[r * cols + c] * x[c]).abs()).sum();
11474            assert!(
11475                (got - want).abs() <= 1e-5 * mag,
11476                "row {r}: kernel {got} vs dequant {want}"
11477            );
11478        }
11479    }
11480
11481    /// The int8 (a8w8) path can't be checked against an f32 reference — the
11482    /// activation quantization dominates. Check it against the q4t kernel it
11483    /// was ported from instead, on payloads holding the same weights: that
11484    /// isolates exactly what the port could break (16 B stride, ladder
11485    /// lookup, nibble unpack) from what it deliberately shares.
11486    #[test]
11487    fn q4tp_matvec_matches_the_q4t_kernel_it_was_ported_from() {
11488        let (rows, cols) = (256usize, 512usize);
11489        let bytes = synth_q4tp(rows, cols);
11490        let twin = q4tp_as_q4t(&bytes, rows, cols);
11491        let x: Vec<f32> = (0..cols)
11492            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
11493            .collect();
11494
11495        let mut got = vec![0f32; rows];
11496        q4tp_matvec(&bytes, &x, rows, cols, &mut got, None);
11497        let mut want = vec![0f32; rows];
11498        q4t_matvec(&twin, &x, rows, cols, &mut want, None);
11499
11500        // Scale is f16 in the twin and f32 here, so allow that rounding on
11501        // top of the summed magnitude (same cancellation argument as above).
11502        let mut w = vec![0f32; rows * cols];
11503        cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
11504        for r in 0..rows {
11505            let mag: f32 = (0..cols).map(|c| (w[r * cols + c] * x[c]).abs()).sum();
11506            assert!(
11507                (got[r] - want[r]).abs() <= 1e-3 * mag,
11508                "row {r}: q4tp {} vs q4t {}",
11509                got[r],
11510                want[r]
11511            );
11512        }
11513    }
11514
11515    /// `matmat` carries three arms (Accelerate, blocked int8 1x4, scalar).
11516    /// Batch 5 crosses the blocked kernel's stride, so this exercises the
11517    /// 1x4 path AND its scalar tail in one run — the blocked kernel is new
11518    /// code and its four accumulators are exactly what tends to go wrong.
11519    #[test]
11520    fn q4tp_matmat_matches_the_q4t_kernel_it_was_ported_from() {
11521        let (rows, cols, b) = (256usize, 512usize, 5usize);
11522        let bytes = synth_q4tp(rows, cols);
11523        let twin = q4tp_as_q4t(&bytes, rows, cols);
11524        let xs: Vec<f32> = (0..b * cols)
11525            .map(|i| ((i * 29 + 11) % 89) as f32 / 89.0 - 0.5)
11526            .collect();
11527
11528        let mut got = vec![0f32; b * rows];
11529        q4tp_matmat(&bytes, &xs, b, rows, cols, &mut got, None);
11530        let mut want = vec![0f32; b * rows];
11531        q4t_matmat(&twin, &xs, b, rows, cols, &mut want, None);
11532
11533        let mut w = vec![0f32; rows * cols];
11534        cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
11535        for t in 0..b {
11536            for r in 0..rows {
11537                let mag: f32 = (0..cols)
11538                    .map(|c| (w[r * cols + c] * xs[t * cols + c]).abs())
11539                    .sum();
11540                let (g, wa) = (got[t * rows + r], want[t * rows + r]);
11541                assert!(
11542                    (g - wa).abs() <= 1e-3 * mag,
11543                    "batch {t} row {r}: q4tp {g} vs q4t {wa}"
11544                );
11545            }
11546        }
11547    }
11548
11549    #[test]
11550    fn q4tp_matvec2_matches_the_single_stream_kernel() {
11551        let (rows, cols) = (128usize, 256usize);
11552        let gpr = cols / GROUP_SIZE;
11553        let bytes = synth_q4tp(rows, cols);
11554        let xs: Vec<f32> = (0..2 * cols)
11555            .map(|i| ((i * 29 + 11) % 89) as f32 / 89.0 - 0.5)
11556            .collect();
11557
11558        let (mut o1, mut o2) = (vec![0f32; rows], vec![0f32; rows]);
11559        q4tp_matvec2(
11560            &bytes,
11561            &xs[..cols],
11562            &xs[cols..],
11563            rows,
11564            cols,
11565            &mut o1,
11566            &mut o2,
11567            None,
11568        );
11569
11570        // matvec2 takes the exact path for both streams, so the single-row
11571        // kernel is an exact reference — no tolerance for path differences.
11572        let v = Q4tpView::new(&bytes, rows, cols);
11573        let mut sc = vec![0f32; gpr];
11574        for r in 0..rows {
11575            v.scales_into(r, gpr, &mut sc);
11576            assert_eq!(o1[r], q4tp_row_exact(v.nib, r, gpr, &xs[..cols], &sc));
11577            assert_eq!(o2[r], q4tp_row_exact(v.nib, r, gpr, &xs[cols..], &sc));
11578        }
11579    }
11580
11581    /// q4tp must not COST speed — it exists to save bytes, and a format that
11582    /// trades 7% of a file for a slower model is a bad trade. This guard is
11583    /// here because correctness tests happily passed while `q4tp_matmat` was
11584    /// missing its int8 and Accelerate arms and the model ran 5x slower.
11585    /// Measured on M-series: 0.97-1.04x, i.e. parity (16 B tiles are better
11586    /// aligned than q4t's 18 B, which pays for the scale indirection).
11587    #[test]
11588    fn q4tp_matvec_keeps_pace_with_q4t() {
11589        let (rows, cols) = (4096usize, 3072usize);
11590        let bytes = synth_q4tp(rows, cols);
11591        let twin = q4tp_as_q4t(&bytes, rows, cols);
11592        let x: Vec<f32> = (0..cols).map(|i| (i % 97) as f32 / 97.0 - 0.5).collect();
11593        let mut o = vec![0f32; rows];
11594        let n = 12;
11595        let mut best = (f64::MAX, f64::MAX);
11596        // Interleaved A/B, minimum statistic: this machine throttles, and a
11597        // mean over a thermal ramp reliably indicts whichever ran second.
11598        for _ in 0..3 {
11599            let t0 = std::time::Instant::now();
11600            for _ in 0..n {
11601                q4t_matvec(&twin, &x, rows, cols, &mut o, None);
11602            }
11603            best.0 = best.0.min(t0.elapsed().as_secs_f64());
11604            let t0 = std::time::Instant::now();
11605            for _ in 0..n {
11606                q4tp_matvec(&bytes, &x, rows, cols, &mut o, None);
11607            }
11608            best.1 = best.1.min(t0.elapsed().as_secs_f64());
11609        }
11610        let ratio = best.1 / best.0;
11611        println!(
11612            "q4t {:.3} ms | q4tp {:.3} ms | {ratio:.2}x",
11613            best.0 * 1e3 / n as f64,
11614            best.1 * 1e3 / n as f64
11615        );
11616        assert!(ratio < 2.0, "q4tp matvec {ratio:.2}x slower than q4t");
11617    }
11618
11619    #[cfg(target_os = "macos")]
11620    #[test]
11621    fn q4t_matmat_accel_matches_dequant_reference() {
11622        if !accel_gemm_enabled() {
11623            return; // CMF_ACCEL=0
11624        }
11625        let (rows, cols, b) = (512usize, 1024usize, 8usize); // ≥500K → accel arm
11626        let gpr = cols / GROUP_SIZE;
11627        let mut bytes = vec![0u8; rows * gpr * Q4_TILE];
11628        for r in 0..rows {
11629            for g in 0..gpr {
11630                let t = (r * gpr + g) * Q4_TILE;
11631                let sc = 0.02 + 0.0005 * ((r * gpr + g) % 64) as f32;
11632                bytes[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
11633                for k in 0..16 {
11634                    bytes[t + 2 + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
11635                }
11636            }
11637        }
11638        let x: Vec<f32> = (0..b * cols)
11639            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
11640            .collect();
11641        let mut got = vec![0f32; b * rows];
11642        q4t_matmat(&bytes, &x, b, rows, cols, &mut got, None);
11643        // Brute-force reference off the same tiles.
11644        let mut w = vec![0f32; rows * cols];
11645        for r in 0..rows {
11646            for g in 0..gpr {
11647                let t = (r * gpr + g) * Q4_TILE;
11648                let s = f16_to_f32(u16::from_le_bytes([bytes[t], bytes[t + 1]]));
11649                for (k, &bb) in bytes[t + 2..t + Q4_TILE].iter().enumerate() {
11650                    w[r * cols + g * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s;
11651                    w[r * cols + g * GROUP_SIZE + k * 2 + 1] =
11652                        (((bb >> 4) & 0x0F) as f32 - 8.0) * s;
11653                }
11654            }
11655        }
11656        for bi in 0..b {
11657            for r in 0..rows {
11658                let want: f32 = (0..cols).map(|j| x[bi * cols + j] * w[r * cols + j]).sum();
11659                let d = (got[bi * rows + r] - want).abs();
11660                assert!(
11661                    d <= want.abs().max(1.0) * 1e-4,
11662                    "accel q4t GEMM diverged at ({bi},{r}): {} vs {want}",
11663                    got[bi * rows + r]
11664                );
11665            }
11666        }
11667    }
11668
11669    #[test]
11670    fn q4matvec_matches_full_dequant() {
11671        let (rows, cols) = (8, 64);
11672        let groups = rows * cols / GROUP_SIZE;
11673        // Hand-craft a q4_block blob: nibbles then f16 scales.
11674        let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
11675        for i in 0..groups * 16 {
11676            bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
11677        }
11678        for g in 0..groups {
11679            let s = 0.01 + 0.003 * g as f32;
11680            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
11681        }
11682        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
11683
11684        let mut reference = vec![0.0f32; rows * cols];
11685        cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
11686        let mut expect = vec![0.0f32; rows];
11687        for r in 0..rows {
11688            expect[r] = reference[r * cols..(r + 1) * cols]
11689                .iter()
11690                .zip(&x)
11691                .map(|(w, xv)| w * xv)
11692                .sum();
11693        }
11694
11695        let mut got = vec![0.0f32; rows];
11696        q4matvec(&bytes, &x, rows, cols, &mut got, None);
11697        // SDOT path quantizes activations to i8 (A8W8): bounded noise,
11698        // same contract as q8/vbit (exact path is pinned by CMF_SDOT=0
11699        // in the golden-parity gate).
11700        let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
11701        let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
11702        for r in 0..rows {
11703            assert!(
11704                (got[r] - expect[r]).abs() < tol * scale,
11705                "row {r}: {} vs {}",
11706                got[r],
11707                expect[r]
11708            );
11709        }
11710    }
11711
11712    /// Fused two-input vbit matvec must equal two single matvecs exactly
11713    /// (same per-lane accumulation order on both scalar and SDOT paths).
11714    #[test]
11715    fn vbitmatvec2_equals_two_singles() {
11716        let (rows, cols) = (6, 64);
11717        let ng = cols / GROUP_SIZE;
11718        let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
11719        let mut bytes = bits.clone();
11720        for g in 0..rows * ng {
11721            let s = 0.02 + 0.001 * g as f32;
11722            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
11723        }
11724        for r in 0..rows {
11725            let b = bits[r] as usize;
11726            let (mut acc, mut nb) = (0u64, 0usize);
11727            let mut rowbytes = Vec::new();
11728            for i in 0..cols {
11729                let v = ((i * 7 + r * 13) % (1 << b)) as u64;
11730                acc = (acc << b) | v;
11731                nb += b;
11732                while nb >= 8 {
11733                    nb -= 8;
11734                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
11735                }
11736            }
11737            if nb > 0 {
11738                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
11739            }
11740            bytes.extend_from_slice(&rowbytes);
11741        }
11742        let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
11743        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).cos()).collect();
11744        let offsets = vbit_row_offsets(&bytes, rows, cols);
11745
11746        let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
11747        vbitmatvec(&bytes, &offsets, &x1, rows, cols, &mut a1, None);
11748        vbitmatvec(&bytes, &offsets, &x2, rows, cols, &mut a2, None);
11749        let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
11750        vbitmatvec2(
11751            &bytes, &offsets, &x1, &x2, rows, cols, &mut b1, &mut b2, None,
11752        );
11753        assert_eq!(a1, b1, "fused vbit lane 1 must be bit-identical");
11754        assert_eq!(a2, b2, "fused vbit lane 2 must be bit-identical");
11755    }
11756
11757    /// Fused two-input q4 matvec must equal two single matvecs exactly.
11758    #[test]
11759    fn q4matvec2_equals_two_singles() {
11760        let (rows, cols) = (8, 128);
11761        let groups = rows * cols / GROUP_SIZE;
11762        let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
11763        for i in 0..groups * 16 {
11764            bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
11765        }
11766        for g in 0..groups {
11767            let s = 0.01 + 0.003 * g as f32;
11768            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
11769        }
11770        // Include an outlier channel so the SDOT correction path is
11771        // exercised in the pair kernel too.
11772        let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
11773        x1[9] = 250.0;
11774        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
11775
11776        let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
11777        q4matvec(&bytes, &x1, rows, cols, &mut a1, None);
11778        q4matvec(&bytes, &x2, rows, cols, &mut a2, None);
11779        let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
11780        q4matvec2(&bytes, &x1, &x2, rows, cols, &mut b1, &mut b2, None);
11781        assert_eq!(a1, b1, "fused q4 lane 1 must be bit-identical");
11782        assert_eq!(a2, b2, "fused q4 lane 2 must be bit-identical");
11783    }
11784
11785    /// Multi-matrix job must equal separate matvecs exactly — same
11786    /// kernels, only the dispatch is fused.
11787    #[test]
11788    fn matvec_many_equals_separate_matvecs() {
11789        use crate::pool::Pool;
11790        let (r1, r2, cols) = (300, 200, 64);
11791        let mk = |salt: usize, rows: usize| {
11792            QTensor::from_f32(
11793                (0..rows * cols)
11794                    .map(|i| ((i * 7 + salt) % 97) as f32 / 97.0 - 0.5)
11795                    .collect(),
11796                rows,
11797                cols,
11798            )
11799        };
11800        let (a, b) = (mk(1, r1), mk(5, r2));
11801        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).sin()).collect();
11802        let pool = Pool::new(3);
11803
11804        let (mut ea, mut eb) = (vec![0f32; r1], vec![0f32; r2]);
11805        a.matvec(&x, &mut ea, Some(&pool));
11806        b.matvec(&x, &mut eb, Some(&pool));
11807        let (mut ga, mut gb) = (vec![0f32; r1], vec![0f32; r2]);
11808        QTensor::matvec_many([&a, &b], &x, [&mut ga, &mut gb], Some(&pool));
11809        assert_eq!(ea, ga, "fused multi-matrix lane 1 must be bit-identical");
11810        assert_eq!(eb, gb, "fused multi-matrix lane 2 must be bit-identical");
11811    }
11812
11813    /// The public Q4TP operator must take the real mapped matvec_many arm,
11814    /// rather than the F32 fallback above.  Build a tiny valid CMF so both
11815    /// handles retain their mmap payloads, then compare the fused dispatch
11816    /// with two ordinary mapped matvec calls bit-for-bit.
11817    #[test]
11818    fn q4tp_matvec_many_equals_separate_matvecs() {
11819        use crate::pool::Pool;
11820        use cortiq_core::{CMF_VERSION, CmfHeader, CmfModel, QuantType, TensorSpec};
11821
11822        let (r1, r2, cols) = (300usize, 200usize, 64usize);
11823        let arch: cortiq_core::ModelArch = serde_json::from_value(serde_json::json!({
11824            "arch_name": "tiny-q4tp",
11825            "hidden_size": cols,
11826            "intermediate_size": cols * 2,
11827            "num_layers": 1,
11828            "num_attention_heads": 2,
11829            "num_kv_heads": 1,
11830            "head_dim": 32,
11831            "vocab_size": r1,
11832            "layer_types": ["FullAttention"],
11833            "rms_norm_eps": 1e-6,
11834            "max_position_embeddings": 8,
11835            "linear_conv_kernel_dim": 0,
11836            "linear_num_key_heads": 0,
11837            "linear_num_value_heads": 0
11838        }))
11839        .unwrap();
11840        let header = CmfHeader {
11841            format: "cmf".into(),
11842            version: CMF_VERSION,
11843            arch,
11844            quant_type: QuantType::Q4Block,
11845            provenance: None,
11846            tokenizer_config: None,
11847            section_hashes: None,
11848            skills: Vec::new(),
11849            shard: None,
11850            calibration: None,
11851            routing: None,
11852        };
11853        let specs = [
11854            TensorSpec {
11855                name: "q".into(),
11856                dtype: TensorDtype::Q4TiledP,
11857                shape: vec![r1, cols],
11858                data: synth_q4tp(r1, cols),
11859            },
11860            TensorSpec {
11861                name: "kv".into(),
11862                dtype: TensorDtype::Q4TiledP,
11863                shape: vec![r2, cols],
11864                data: synth_q4tp(r2, cols),
11865            },
11866        ];
11867        let dir = std::env::temp_dir().join(format!("cmf-q4tp-many-{}", std::process::id()));
11868        std::fs::create_dir_all(&dir).unwrap();
11869        let path = dir.join("m.cmf");
11870        CmfModel::write(&path, &header, &specs, None, None).unwrap();
11871        let model = Arc::new(CmfModel::open(&path).unwrap());
11872        let (a, b) = (
11873            QTensor::from_model(&model, "q").unwrap(),
11874            QTensor::from_model(&model, "kv").unwrap(),
11875        );
11876        assert_eq!(a.model_dtype(), Some(TensorDtype::Q4TiledP));
11877        assert_eq!(b.model_dtype(), Some(TensorDtype::Q4TiledP));
11878        let x: Vec<f32> = (0..cols)
11879            .map(|i| ((i * 17 + 3) % 97) as f32 / 97.0 - 0.5)
11880            .collect();
11881        let pool = Pool::new(3);
11882        let (mut ea, mut eb) = (vec![0.0f32; r1], vec![0.0f32; r2]);
11883        a.matvec(&x, &mut ea, Some(&pool));
11884        b.matvec(&x, &mut eb, Some(&pool));
11885        let (mut ga, mut gb) = (vec![0.0f32; r1], vec![0.0f32; r2]);
11886        QTensor::matvec_many([&a, &b], &x, [&mut ga, &mut gb], Some(&pool));
11887        assert_eq!(ea, ga, "Q4TP fused lane 1 must be bit-identical");
11888        assert_eq!(eb, gb, "Q4TP fused lane 2 must be bit-identical");
11889        let _ = std::fs::remove_dir_all(&dir);
11890    }
11891
11892    /// Batched q4/vbit matmat must equal per-position matvec calls
11893    /// exactly (the fallback it replaced) — same kernels, same order.
11894    #[test]
11895    fn batched_matmat_equals_per_position_matvec() {
11896        let (rows, cols, b) = (8, 64, 5);
11897        // q4 blob.
11898        let groups = rows * cols / GROUP_SIZE;
11899        let mut q4 = Vec::new();
11900        for i in 0..groups * 16 {
11901            q4.push((((i * 7 + 3) % 256) & 0xFF) as u8);
11902        }
11903        for g in 0..groups {
11904            q4.extend_from_slice(
11905                &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
11906            );
11907        }
11908        // vbit blob (mixed widths incl. 8).
11909        let ng = cols / GROUP_SIZE;
11910        let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4, 5, 3];
11911        let mut vb = bits.clone();
11912        for g in 0..rows * ng {
11913            vb.extend_from_slice(
11914                &cortiq_core::quant::f32_to_f16(0.02 + 0.001 * g as f32).to_le_bytes(),
11915            );
11916        }
11917        for r in 0..rows {
11918            let bw = bits[r] as usize;
11919            let (mut acc, mut nb) = (0u64, 0usize);
11920            let mut rowbytes = Vec::new();
11921            for i in 0..cols {
11922                let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
11923                acc = (acc << bw) | v;
11924                nb += bw;
11925                while nb >= 8 {
11926                    nb -= 8;
11927                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
11928                }
11929            }
11930            if nb > 0 {
11931                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
11932            }
11933            vb.extend_from_slice(&rowbytes);
11934        }
11935        let offsets = vbit_row_offsets(&vb, rows, cols);
11936
11937        let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
11938
11939        // q4: batch vs singles.
11940        let mut got = vec![0f32; b * rows];
11941        q4matmat(&q4, &xs, b, rows, cols, &mut got, None);
11942        for bi in 0..b {
11943            let mut expect = vec![0f32; rows];
11944            q4matvec(
11945                &q4,
11946                &xs[bi * cols..(bi + 1) * cols],
11947                rows,
11948                cols,
11949                &mut expect,
11950                None,
11951            );
11952            assert_eq!(
11953                &got[bi * rows..(bi + 1) * rows],
11954                &expect[..],
11955                "q4 batch pos {bi}"
11956            );
11957        }
11958
11959        // vbit: batch vs singles.
11960        let mut got = vec![0f32; b * rows];
11961        vbitmatmat(&vb, &offsets, &xs, b, rows, cols, &mut got, None);
11962        for bi in 0..b {
11963            let mut expect = vec![0f32; rows];
11964            vbitmatvec(
11965                &vb,
11966                &offsets,
11967                &xs[bi * cols..(bi + 1) * cols],
11968                rows,
11969                cols,
11970                &mut expect,
11971                None,
11972            );
11973            assert_eq!(
11974                &got[bi * rows..(bi + 1) * rows],
11975                &expect[..],
11976                "vbit batch pos {bi}"
11977            );
11978        }
11979    }
11980
11981    /// q4_tiled kernels must produce BIT-identical outputs to the q4
11982    /// split kernels on the same values (same ints, same order — only
11983    /// the byte placement differs).
11984    #[test]
11985    fn q4_tiled_matches_q4_block_bitexact() {
11986        let (rows, cols, b) = (8usize, 128usize, 3usize);
11987        let groups = rows * cols / GROUP_SIZE;
11988        let mut split = Vec::with_capacity(groups * 18);
11989        for i in 0..groups * 16 {
11990            split.push((((i * 7 + 3) % 256) & 0xFF) as u8);
11991        }
11992        for g in 0..groups {
11993            split.extend_from_slice(
11994                &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
11995            );
11996        }
11997        // Re-tile: [scale][nibbles] per group.
11998        let (packed, scales) = split.split_at(groups * 16);
11999        let mut tiled = Vec::with_capacity(groups * Q4_TILE);
12000        for g in 0..groups {
12001            tiled.extend_from_slice(&scales[g * 2..g * 2 + 2]);
12002            tiled.extend_from_slice(&packed[g * 16..(g + 1) * 16]);
12003        }
12004
12005        let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
12006        x1[9] = 250.0; // exercise the outlier path
12007        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
12008
12009        let (mut a, mut t) = (vec![0f32; rows], vec![0f32; rows]);
12010        q4matvec(&split, &x1, rows, cols, &mut a, None);
12011        q4t_matvec(&tiled, &x1, rows, cols, &mut t, None);
12012        assert_eq!(a, t, "q4t matvec must match q4 bit-for-bit");
12013
12014        let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
12015        let (mut t1, mut t2) = (vec![0f32; rows], vec![0f32; rows]);
12016        q4matvec2(&split, &x1, &x2, rows, cols, &mut a1, &mut a2, None);
12017        q4t_matvec2(&tiled, &x1, &x2, rows, cols, &mut t1, &mut t2, None);
12018        assert_eq!(a1, t1);
12019        assert_eq!(a2, t2);
12020
12021        let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
12022        let (mut am, mut tm) = (vec![0f32; b * rows], vec![0f32; b * rows]);
12023        q4matmat(&split, &xs, b, rows, cols, &mut am, None);
12024        q4t_matmat(&tiled, &xs, b, rows, cols, &mut tm, None);
12025        assert_eq!(am, tm, "q4t matmat must match q4 bit-for-bit");
12026    }
12027
12028    /// q4 SDOT outlier correction: a single huge activation channel
12029    /// (>8·rms → outlier, zeroed in xq) must still contribute its EXACT
12030    /// term. On-grid bulk (±1/0 → xq dequantizes exactly) isolates the
12031    /// correction from A8W8 noise. cols must exceed 64: at n=64 the
12032    /// 8·rms threshold equals sqrt(v²+rest) ≥ v, so a single outlier
12033    /// can never qualify (8² = n).
12034    #[test]
12035    fn q4matvec_sdot_outlier_exact() {
12036        let (rows, cols) = (4, 128);
12037        let groups = rows * cols / GROUP_SIZE;
12038        let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
12039        for i in 0..groups * 16 {
12040            bytes.push(((i * 11 + 5) % 256) as u8);
12041        }
12042        for g in 0..groups {
12043            let s = 0.02 + 0.002 * g as f32;
12044            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
12045        }
12046        let mut x: Vec<f32> = (0..cols)
12047            .map(|i| match i % 3 {
12048                0 => 1.0,
12049                1 => -1.0,
12050                _ => 0.0,
12051            })
12052            .collect();
12053        x[17] = 300.0; // ≫ 8·rms → outlier channel
12054
12055        let mut reference = vec![0.0f32; rows * cols];
12056        cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
12057        let mut expect = vec![0.0f32; rows];
12058        for r in 0..rows {
12059            expect[r] = reference[r * cols..(r + 1) * cols]
12060                .iter()
12061                .zip(&x)
12062                .map(|(w, xv)| w * xv)
12063                .sum();
12064        }
12065        let mut got = vec![0.0f32; rows];
12066        q4matvec(&bytes, &x, rows, cols, &mut got, None);
12067        let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
12068        for r in 0..rows {
12069            assert!(
12070                (got[r] - expect[r]).abs() < 2e-3 * scale,
12071                "row {r}: {} vs {} (outlier term must be exact)",
12072                got[r],
12073                expect[r]
12074            );
12075        }
12076    }
12077
12078    /// The fused q1t matvec must equal the reference (dequant_q1t → dot),
12079    /// including the ternary zero level and the binary-searched outlier
12080    /// overlay. Guards the mmap kernel that makes a 12B q1t runnable.
12081    #[test]
12082    fn q1t_matvec_matches_reference() {
12083        use cortiq_core::quant::{dequant_q1t, f32_to_f16};
12084        let (rows, cols) = (3usize, 64usize); // gpr = 2
12085        let gpr = cols / GROUP_SIZE;
12086        let scales = [0.5f32, 0.3, 0.7, 0.2, 0.6, 0.15];
12087        // Overlay (must be sorted by flat index): a few spikes across rows.
12088        let outliers: [(u32, f32); 3] = [(5, 9.0), (70, -4.5), (150, 3.25)];
12089        let is_out = |flat: usize| outliers.iter().any(|&(i, _)| i as usize == flat);
12090        let mut bytes = Vec::new();
12091        for r in 0..rows {
12092            for g in 0..gpr {
12093                bytes.extend_from_slice(&f32_to_f16(scales[r * gpr + g]).to_le_bytes());
12094                let mut c = [0u8; 7];
12095                for k in 0..GROUP_SIZE {
12096                    // Encoder invariant: code 0 at outlier positions.
12097                    let code = if is_out(r * cols + g * GROUP_SIZE + k) {
12098                        0
12099                    } else {
12100                        ((k + r * 3 + g) % 3) as u8 // 0,1,2
12101                    };
12102                    cortiq_core::quant::q1t_pack(&mut c, k, code);
12103                }
12104                bytes.extend_from_slice(&c);
12105            }
12106        }
12107        // Per-row overlay: [u32 row_ptr[rows+1]] then [(u16 col, f16 val)] by
12108        // row (outliers are sorted by flat index → already grouped by row).
12109        let mut row_ptr = vec![0u32; rows + 1];
12110        for &(idx, _) in &outliers {
12111            row_ptr[idx as usize / cols + 1] += 1;
12112        }
12113        for r in 0..rows {
12114            row_ptr[r + 1] += row_ptr[r];
12115        }
12116        for &p in &row_ptr {
12117            bytes.extend_from_slice(&p.to_le_bytes());
12118        }
12119        for &(idx, v) in &outliers {
12120            bytes.extend_from_slice(&((idx as usize % cols) as u16).to_le_bytes());
12121            bytes.extend_from_slice(&f32_to_f16(v).to_le_bytes());
12122        }
12123
12124        let mut refw = vec![0f32; rows * cols];
12125        dequant_q1t(&bytes, rows, cols, &mut refw);
12126        // On-grid activations (±1, amax 1) so the int8 SDOT path reconstructs
12127        // x exactly and matches the f32 reference (same trick as the q1 test).
12128        let x: Vec<f32> = (0..cols)
12129            .map(|j| if j % 3 == 0 { 1.0 } else { -1.0 })
12130            .collect();
12131        let mut expect = vec![0f32; rows];
12132        for r in 0..rows {
12133            let mut a = 0.0f32;
12134            for j in 0..cols {
12135                a += refw[r * cols + j] * x[j];
12136            }
12137            expect[r] = a;
12138        }
12139        let tol = |e: f32| 1e-3 * e.abs().max(1e-3);
12140        let mut got = vec![0f32; rows];
12141        q1t_matvec(&bytes, &x, rows, cols, &mut got, None);
12142        for r in 0..rows {
12143            assert!(
12144                (got[r] - expect[r]).abs() < tol(expect[r]),
12145                "row {r}: {} vs {}",
12146                got[r],
12147                expect[r]
12148            );
12149        }
12150        // matmat (b=2, f32 decode path) must agree too.
12151        let x2: Vec<f32> = x.iter().chain(x.iter()).copied().collect();
12152        let mut gm = vec![0f32; 2 * rows];
12153        q1t_matmat(&bytes, &x2, 2, rows, cols, &mut gm, None);
12154        for r in 0..rows {
12155            assert!((gm[r] - expect[r]).abs() < tol(expect[r]));
12156            assert!((gm[rows + r] - expect[r]).abs() < tol(expect[r]));
12157        }
12158        // Fused pair (q1t_matvec2) must equal two single matvecs
12159        // bit-for-bit: same unpack, same group order, same f32
12160        // accumulation per stream. Distinct x2 exercises both lanes.
12161        let xb: Vec<f32> = (0..cols)
12162            .map(|j| if j % 5 == 0 { -1.0 } else { 1.0 })
12163            .collect();
12164        let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
12165        q1t_matvec(&bytes, &x, rows, cols, &mut s1, None);
12166        q1t_matvec(&bytes, &xb, rows, cols, &mut s2, None);
12167        let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
12168        q1t_matvec2(&bytes, &x, &xb, rows, cols, &mut p1, &mut p2, None);
12169        assert_eq!(p1, s1, "q1t pair lane 1 ≠ single matvec");
12170        assert_eq!(p2, s2, "q1t pair lane 2 ≠ single matvec");
12171    }
12172
12173    /// Pair == 2×matvec with an ODD group count (the kernel's tail
12174    /// group) and no overlay section.
12175    #[test]
12176    fn q1t_matvec2_odd_gpr_matches_singles() {
12177        use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_pack};
12178        let (rows, cols) = (5usize, 96usize); // gpr = 3 → paired + tail
12179        let gpr = cols / GROUP_SIZE;
12180        let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE);
12181        for r in 0..rows {
12182            for g in 0..gpr {
12183                bytes.extend_from_slice(&f32_to_f16(0.1 + 0.05 * (r + g) as f32).to_le_bytes());
12184                let mut c = [0u8; 7];
12185                for k in 0..GROUP_SIZE {
12186                    q1t_pack(&mut c, k, ((k * 7 + r * 5 + g * 3) % 3) as u8);
12187                }
12188                bytes.extend_from_slice(&c);
12189            }
12190        }
12191        let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.31).sin()).collect();
12192        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).cos()).collect();
12193        let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
12194        q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
12195        q1t_matvec(&bytes, &x2, rows, cols, &mut s2, None);
12196        let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
12197        q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
12198        assert_eq!(p1, s1, "odd-gpr pair lane 1 ≠ single");
12199        assert_eq!(p2, s2, "odd-gpr pair lane 2 ≠ single");
12200    }
12201
12202    // Speed A/B: fused pair (one unpack, two streams) vs two single
12203    // matvecs. Single-threaded, FFN-sized, min-of paired in-process.
12204    //   cargo test -p cortiq-engine --release q1t_matvec2_speed -- --ignored --nocapture
12205    #[test]
12206    #[ignore]
12207    fn q1t_matvec2_speed() {
12208        use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_pack};
12209        use std::time::Instant;
12210        let (rows, cols) = (8192usize, 4096usize);
12211        let gpr = cols / GROUP_SIZE;
12212        let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE);
12213        for r in 0..rows {
12214            for g in 0..gpr {
12215                let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
12216                bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
12217                let mut c = [0u8; 7];
12218                for k in 0..GROUP_SIZE {
12219                    q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
12220                }
12221                bytes.extend_from_slice(&c);
12222            }
12223        }
12224        let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.31).sin()).collect();
12225        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).cos()).collect();
12226        let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
12227        let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
12228        // Warm both paths once.
12229        q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
12230        q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
12231        let (mut t_pair, mut t_two) = (f64::MAX, f64::MAX);
12232        for _ in 0..8 {
12233            let t0 = Instant::now();
12234            q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
12235            t_pair = t_pair.min(t0.elapsed().as_secs_f64() * 1000.0);
12236            let t1 = Instant::now();
12237            q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
12238            q1t_matvec(&bytes, &x2, rows, cols, &mut s2, None);
12239            t_two = t_two.min(t1.elapsed().as_secs_f64() * 1000.0);
12240        }
12241        assert_eq!(p1, s1);
12242        assert_eq!(p2, s2);
12243        println!("q1t pair {rows}x{cols}: fused {t_pair:.2} ms | two singles {t_two:.2} ms");
12244    }
12245
12246    // Speed A/B: the base-3-division decode (what the packing commit left in
12247    // place) vs the fused sign-LUT matvec. Both single-threaded, same bytes.
12248    //   cargo test -p cortiq-engine q1t_matvec_speed -- --ignored --nocapture
12249    #[test]
12250    #[ignore]
12251    fn q1t_matvec_speed() {
12252        use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_code, q1t_pack};
12253        use std::time::Instant;
12254        let (rows, cols) = (8192usize, 4096usize); // FFN-sized
12255        let gpr = cols / GROUP_SIZE;
12256        let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE + 16);
12257        for r in 0..rows {
12258            for g in 0..gpr {
12259                let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
12260                bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
12261                let mut c = [0u8; 7];
12262                for k in 0..GROUP_SIZE {
12263                    q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
12264                }
12265                bytes.extend_from_slice(&c);
12266            }
12267        }
12268        let (n, stride) = (rows * cols, 40usize); // ~2.5% outliers, per-row overlay
12269        let mut row_ptr = vec![0u32; rows + 1];
12270        let mut idx = 0usize;
12271        while idx < n {
12272            row_ptr[idx / cols + 1] += 1;
12273            idx += stride;
12274        }
12275        for r in 0..rows {
12276            row_ptr[r + 1] += row_ptr[r];
12277        }
12278        for &p in &row_ptr {
12279            bytes.extend_from_slice(&p.to_le_bytes());
12280        }
12281        let mut idx = 0usize;
12282        while idx < n {
12283            bytes.extend_from_slice(&((idx % cols) as u16).to_le_bytes());
12284            bytes.extend_from_slice(&f32_to_f16((idx % 13) as f32 * 0.1 - 0.6).to_le_bytes());
12285            idx += stride;
12286        }
12287        // On-grid ±1 so the fast path's int8 SDOT is exact vs the f32 "slow"
12288        // reference (the A/B is a timing check; values must still agree).
12289        let x: Vec<f32> = (0..cols)
12290            .map(|j| if j % 3 == 0 { 1.0 } else { -1.0 })
12291            .collect();
12292        let (rp_off, ent_off, has_ov) = q1t_overlay(&bytes, rows * gpr * Q1T_TILE, rows);
12293
12294        // "before": base-3 division decode into a buffer, then dot.
12295        let slow = |out: &mut [f32]| {
12296            let mut buf = vec![0f32; cols];
12297            for r in 0..rows {
12298                for g in 0..gpr {
12299                    let off = (r * gpr + g) * Q1T_TILE;
12300                    let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
12301                    let codes = &bytes[off + 2..off + Q1T_TILE];
12302                    for k in 0..GROUP_SIZE {
12303                        buf[g * GROUP_SIZE + k] = match q1t_code(codes, k) {
12304                            1 => s,
12305                            2 => -s,
12306                            _ => 0.0,
12307                        };
12308                    }
12309                }
12310                out[r] = q1t_row_outlier_correction(&bytes, r, rp_off, ent_off, has_ov, &x)
12311                    + (0..cols).map(|j| buf[j] * x[j]).sum::<f32>();
12312            }
12313        };
12314        let iters = 5;
12315        let mut a = vec![0f32; rows];
12316        slow(&mut a); // warm
12317        let t = Instant::now();
12318        for _ in 0..iters {
12319            slow(&mut a);
12320        }
12321        let slow_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
12322
12323        let mut b = vec![0f32; rows];
12324        q1t_matvec(&bytes, &x, rows, cols, &mut b, None); // warm
12325        let t = Instant::now();
12326        for _ in 0..iters {
12327            q1t_matvec(&bytes, &x, rows, cols, &mut b, None);
12328        }
12329        let fast_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
12330
12331        for r in 0..rows {
12332            assert!((a[r] - b[r]).abs() < 1e-2, "mismatch row {r}");
12333        }
12334        println!(
12335            "q1t matvec {rows}x{cols} (1 thread): div-decode {slow_ms:.2} ms  fused-LUT {fast_ms:.2} ms  => {:.2}x",
12336            slow_ms / fast_ms
12337        );
12338    }
12339}
12340
12341#[cfg(test)]
12342mod gemm_bench {
12343    /// `cargo test -p cortiq-engine --release q4tp_matmat_throughput -- --ignored --nocapture`
12344    /// Times the batched q4tp GEMM at the shapes the image DiT runs
12345    /// (b=296 tokens, 2304 -> 9216), on synthetic bytes: no model, no
12346    /// mmap, no thermal drift over minutes — a kernel change shows up
12347    /// here in seconds where a full render hides it in noise.
12348    ///
12349    /// On macOS add `CMF_ACCEL=0`: this shape is over the 500k-cell mark
12350    /// where the matmat hands off to Accelerate's dequant sgemm, and
12351    /// without the opt-out both rows below measure the AMX, not the
12352    /// kernel under test.
12353    #[test]
12354    #[ignore]
12355    fn q4tp_matmat_throughput() {
12356        // 296 is a prompt-encode batch; the image DiT runs 2085 at
12357        // 512x512, where the activation panel stops fitting L2 and the
12358        // loop's shape starts to matter more than its instructions.
12359        let b: usize = std::env::var("CMF_BENCH_B")
12360            .ok()
12361            .and_then(|v| v.parse().ok())
12362            .unwrap_or(296);
12363        let (rows, cols) = (9216usize, 2304usize);
12364        let (_, _, _) = (rows, cols, b);
12365        let total =
12366            cortiq_core::quant::expected_nbytes(cortiq_core::TensorDtype::Q4TiledP, &[rows, cols])
12367                .unwrap();
12368        // Random nibbles are fine, but the row params are f16 (lo, step)
12369        // of a geometric ladder: garbage there gives exp2 of a huge
12370        // exponent, the scales come back inf, and the whole bench times
12371        // NaN arithmetic instead of the kernel.
12372        let (params_off, codes_off, _) = cortiq_core::quant::q4tp_sections(rows, cols);
12373        let mut bytes: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
12374        let lo = cortiq_core::quant::f32_to_f16(-4.0);
12375        let step = cortiq_core::quant::f32_to_f16(0.1);
12376        for r in 0..rows {
12377            let o = params_off + r * 4;
12378            bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
12379            bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
12380        }
12381        let _ = codes_off;
12382        let xs: Vec<f32> = (0..b * cols)
12383            .map(|i| ((i % 97) as f32 - 48.0) / 48.0)
12384            .collect();
12385        let mut out = vec![0f32; b * rows];
12386        let pool = crate::pool::Pool::from_env();
12387        // A shared 48-core stand drifts ±25% run to run, which is wider
12388        // than any kernel change worth making. So: alternate the two
12389        // kernels inside one process and keep the BEST time for
12390        // each. Interleaving makes both see the same interference, and a
12391        // minimum is the one statistic another tenant cannot inflate.
12392        super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
12393        let reps: usize = std::env::var("CMF_BENCH_REPS")
12394            .ok()
12395            .and_then(|v| v.parse().ok())
12396            .unwrap_or(10);
12397        let mut best = [f64::MAX; 2];
12398        let mut sums = [0f32; 2];
12399        for _ in 0..reps {
12400            for (k, w) in [(0usize, 1u8), (1usize, 2u8)] {
12401                super::Q4TP_ALT.store(w, std::sync::atomic::Ordering::Relaxed);
12402                let t = std::time::Instant::now();
12403                super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
12404                best[k] = best[k].min(t.elapsed().as_secs_f64());
12405                sums[k] = out.iter().take(64).sum::<f32>();
12406            }
12407        }
12408        let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
12409        for (k, name) in ["previous", "tuned   "].iter().enumerate() {
12410            println!(
12411                "q4tp matmat {rows}x{cols} b={b} {name}: {:.1} ms  {:.1} GFLOP/s  (checksum {:.3})",
12412                best[k] * 1e3,
12413                flops / best[k] / 1e9,
12414                sums[k]
12415            );
12416        }
12417        assert!(
12418            (sums[0] - sums[1]).abs() < 1e-2,
12419            "the tuned kernel changed the result: {} vs {}",
12420            sums[0],
12421            sums[1]
12422        );
12423    }
12424
12425    /// The blocked kernel must agree with the per-column path exactly —
12426    /// same weights, same activation split, only a different instruction
12427    /// mix. Shapes are chosen to hit the awkward cases: a column count
12428    /// that leaves an odd group (the 512-bit kernel does two at a time),
12429    /// and a batch that does not divide by four.
12430    #[test]
12431    fn q4tp_matmat_blocked_matches_scalar() {
12432        use std::sync::atomic::Ordering::Relaxed;
12433        // The last shape carries the image DiT's column count — 2304, so
12434        // 72 groups of accumulation, which is where a reordered sum can
12435        // actually drift — and runs through the thread pool, since the
12436        // blocked path splits rows across workers. Its row count stays
12437        // under 500k cells on purpose: above that, macOS diverts the whole
12438        // matmat to the Accelerate/AMX dequant sgemm and neither kernel
12439        // here would run.
12440        for &(rows, cols, b) in &[
12441            (64usize, 128usize, 7usize),
12442            (33, 96, 4),
12443            (16, 256, 9),
12444            (192, 2304, 37),
12445        ] {
12446            let total = cortiq_core::quant::expected_nbytes(
12447                cortiq_core::TensorDtype::Q4TiledP,
12448                &[rows, cols],
12449            )
12450            .unwrap();
12451            let (params_off, _, _) = cortiq_core::quant::q4tp_sections(rows, cols);
12452            let mut bytes: Vec<u8> = (0..total).map(|i| (i * 61 % 251) as u8).collect();
12453            let lo = cortiq_core::quant::f32_to_f16(-4.0);
12454            let step = cortiq_core::quant::f32_to_f16(0.1);
12455            for r in 0..rows {
12456                let o = params_off + r * 4;
12457                bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
12458                bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
12459            }
12460            let xs: Vec<f32> = (0..b * cols)
12461                .map(|i| ((i % 89) as f32 - 44.0) / 44.0)
12462                .collect();
12463            let mut got = vec![0f32; b * rows];
12464            let mut want = vec![0f32; b * rows];
12465            let gpr = cols / 32;
12466            let view = super::Q4tpView::new(&bytes, rows, cols);
12467            let pool = crate::pool::Pool::from_env();
12468            super::Q4TP_ALT.store(2, Relaxed);
12469            super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut got, pool.as_deref());
12470            super::Q4TP_ALT.store(1, Relaxed);
12471            super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut want, pool.as_deref());
12472            super::Q4TP_ALT.store(0, Relaxed);
12473            // Measured against the output's scale, not cell by cell: a
12474            // dot product of 2304 terms lands near zero wherever the row
12475            // and the activation nearly cancel, and there a per-cell
12476            // ratio reports 1e-3 for an absolute error of 5e-6 — f32's
12477            // own rounding, reordered. What must stay small is the error
12478            // relative to what the layer actually outputs.
12479            let scale = want.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
12480            let (mut worst, mut at) = (0f32, 0usize);
12481            for (i, (g, w)) in got.iter().zip(&want).enumerate() {
12482                if (g - w).abs() > worst {
12483                    worst = (g - w).abs();
12484                    at = i;
12485                }
12486            }
12487            assert!(
12488                worst <= 1e-4 * scale,
12489                "{rows}x{cols} b={b}: blocked and scalar disagree by {worst:.3e} \
12490                 (scale {scale:.3e}) at cell {at}: {} vs {}",
12491                got[at],
12492                want[at]
12493            );
12494
12495            // "Same speed, no quality loss" is a claim about which answer
12496            // is RIGHT, not about which two agree. Both paths sum the same
12497            // 2304 products in different orders, so f64 decides: the
12498            // blocked kernel keeps sixteen partial sums and folds them at
12499            // the end, which is a shallower addition tree than the
12500            // per-column path's running scalar, and it must not be worse.
12501            let (mut e_blocked, mut e_scalar) = (0f64, 0f64);
12502            for bi in 0..b {
12503                let act = super::split_act(&xs[bi * cols..(bi + 1) * cols]);
12504                for r in 0..rows {
12505                    let mut sc = vec![0f32; gpr];
12506                    view.scales_into(r, gpr, &mut sc);
12507                    let mut exact = 0f64;
12508                    for j in 0..cols {
12509                        let (w, sq) = super::q4tp_outlier(view.nib, r, gpr, j, &sc);
12510                        exact += w as f64 * sq as f64 * act.xq[j] as f64;
12511                    }
12512                    exact *= act.sx as f64;
12513                    for &(j, xv) in &act.outliers {
12514                        let (w, sq) = super::q4tp_outlier(view.nib, r, gpr, j, &sc);
12515                        exact += w as f64 * sq as f64 * xv as f64;
12516                    }
12517                    let i = bi * rows + r;
12518                    e_blocked = e_blocked.max((got[i] as f64 - exact).abs());
12519                    e_scalar = e_scalar.max((want[i] as f64 - exact).abs());
12520                }
12521            }
12522            println!(
12523                "{rows}x{cols} b={b}: worst error vs f64 — blocked {e_blocked:.3e}, \
12524                 per-column {e_scalar:.3e}"
12525            );
12526            // An absolute bar, not a race between the two: at these
12527            // magnitudes both sit in f32's last bits, and on a small shape
12528            // whichever one happens to round the unluckiest cell "wins" by
12529            // a factor the next seed reverses.
12530            assert!(
12531                e_blocked <= 1e-5 * scale as f64 && e_scalar <= 1e-5 * scale as f64,
12532                "{rows}x{cols} b={b}: error against f64 too large — blocked \
12533                 {e_blocked:.3e}, per-column {e_scalar:.3e}, scale {scale:.3e}"
12534            );
12535        }
12536    }
12537
12538    /// The q4t twin of the throughput bench, same shape and rules, so the
12539    /// two quantisations' batch kernels can be read against each other.
12540    /// `cargo test -p cortiq-engine --release q4t_matmat_throughput -- --ignored --nocapture`
12541    #[test]
12542    #[ignore]
12543    fn q4t_matmat_throughput() {
12544        let (rows, cols, b) = (9216usize, 2304usize, 296usize);
12545        let total =
12546            cortiq_core::quant::expected_nbytes(cortiq_core::TensorDtype::Q4Tiled, &[rows, cols])
12547                .unwrap();
12548        // q4t carries a per-group f16 scale in the tile's first two bytes;
12549        // random bytes there decode to inf and the bench would time NaNs.
12550        let mut bytes: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
12551        let sc = cortiq_core::quant::f32_to_f16(0.02);
12552        for t in bytes.chunks_mut(super::Q4_TILE) {
12553            t[..2].copy_from_slice(&sc.to_le_bytes());
12554        }
12555        let xs: Vec<f32> = (0..b * cols)
12556            .map(|i| ((i % 97) as f32 - 48.0) / 48.0)
12557            .collect();
12558        let mut out = vec![0f32; b * rows];
12559        let pool = crate::pool::Pool::from_env();
12560        super::q4t_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
12561        let reps: usize = std::env::var("CMF_BENCH_REPS")
12562            .ok()
12563            .and_then(|v| v.parse().ok())
12564            .unwrap_or(10);
12565        let mut best = f64::MAX;
12566        for _ in 0..reps {
12567            let t = std::time::Instant::now();
12568            super::q4t_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
12569            best = best.min(t.elapsed().as_secs_f64());
12570        }
12571        let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
12572        println!(
12573            "q4t matmat {rows}x{cols} b={b}: {:.1} ms  {:.1} GFLOP/s  (checksum {:.3})",
12574            best * 1e3,
12575            flops / best / 1e9,
12576            out.iter().take(64).sum::<f32>()
12577        );
12578    }
12579}