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        let _prof = crate::cpuprof::time(crate::cpuprof::Slot::Matmat);
1769        // GPTQ calibration: fold this layer's inputs into its Hessian. Only
1770        // Mapped tensors carry a directory name; the check is a relaxed
1771        // atomic load, free when not calibrating.
1772        if crate::gptq_capture::capturing() {
1773            if let Self::Mapped { model, idx, .. } = self {
1774                crate::gptq_capture::accumulate(&model.tensors[*idx].name, xs_all, b, cols);
1775            }
1776        }
1777        match self {
1778            Self::F32 { data, .. } => {
1779                let out_addr = SendMut(out.as_mut_ptr());
1780                let run = |start: usize, end: usize| {
1781                    for o in start..end {
1782                        let row = &data[o * cols..(o + 1) * cols];
1783                        for bi in 0..b {
1784                            let x = &xs_all[bi * cols..(bi + 1) * cols];
1785                            let mut acc = 0f32;
1786                            for j in 0..cols {
1787                                acc += row[j] * x[j];
1788                            }
1789                            unsafe { *out_addr.at(bi * rows + o) = acc };
1790                        }
1791                    }
1792                };
1793                dispatch_rows(pool, rows, &run);
1794            }
1795            Self::Mapped {
1796                model,
1797                idx,
1798                dtype,
1799                row_scale,
1800                col_field,
1801                vbit_offsets,
1802                ..
1803            } => {
1804                if crate::prism::is_forward_weight(model, &model.tensors[*idx].name) {
1805                    let mut transformed = Vec::with_capacity(xs_all.len());
1806                    for bi in 0..b {
1807                        transformed.extend_from_slice(&crate::prism::forward(
1808                            model,
1809                            &xs_all[bi * cols..(bi + 1) * cols],
1810                        ));
1811                    }
1812                    match dtype {
1813                        TensorDtype::Q4Block => {
1814                            q4matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
1815                        }
1816                        TensorDtype::Q4Tiled => {
1817                            q4t_matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
1818                        }
1819                        TensorDtype::Q4TiledP => {
1820                            q4tp_matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
1821                        }
1822                        TensorDtype::Q2TiledP => {
1823                            let affine =
1824                                crate::prism::is_affine_target(model, &model.tensors[*idx].name);
1825                            // Affine Prism Q2TP has a descriptor-aware GPU
1826                            // kernel for short/tail batches too.  Unlike the
1827                            // ordinary Q2TP path, don't force b<32 back to a
1828                            // scalar CPU matmat: prefill chunks and the final
1829                            // tail both need to stay on the tested GPU arm.
1830                            let gpu_batch_ok = if affine {
1831                                b >= 2
1832                            } else {
1833                                b >= 32 && b * rows * cols >= 128_000_000
1834                            };
1835                            if gpu_batch_ok
1836                                && cols % 32 == 0
1837                                && crate::gpu::enabled_here()
1838                                && crate::gpu::q2tp_gpu_opt_in()
1839                            {
1840                                let gpu_ok = if affine {
1841                                    crate::gpu::q2tp_affine_matmat(
1842                                        model,
1843                                        *idx,
1844                                        &transformed,
1845                                        b,
1846                                        rows,
1847                                        cols,
1848                                        out,
1849                                    )
1850                                } else {
1851                                    crate::gpu::q2tp_matmat(
1852                                        model,
1853                                        *idx,
1854                                        &transformed,
1855                                        b,
1856                                        rows,
1857                                        cols,
1858                                        out,
1859                                    )
1860                                };
1861                                if gpu_ok {
1862                                    return;
1863                                }
1864                            }
1865                            // A one-token Prism decode is the other short
1866                            // case.  Use the descriptor-aware matvec kernel
1867                            // before falling back to the exact CPU path.
1868                            if affine
1869                                && b == 1
1870                                && cols % 32 == 0
1871                                && crate::gpu::enabled_here()
1872                                && crate::gpu::q2tp_gpu_opt_in()
1873                                && crate::gpu::q2tp_affine_matvec(
1874                                    model,
1875                                    *idx,
1876                                    &transformed[..cols],
1877                                    rows,
1878                                    cols,
1879                                    &mut out[..rows],
1880                                )
1881                            {
1882                                return;
1883                            }
1884                            if affine {
1885                                q2tp_affine_matmat(
1886                                    self.quant_bytes(),
1887                                    &transformed,
1888                                    b,
1889                                    rows,
1890                                    cols,
1891                                    out,
1892                                    pool,
1893                                )
1894                            } else {
1895                                q2tp_matmat(
1896                                    self.quant_bytes(),
1897                                    &transformed,
1898                                    b,
1899                                    rows,
1900                                    cols,
1901                                    out,
1902                                    pool,
1903                                )
1904                            }
1905                        }
1906                        TensorDtype::Q1 => {
1907                            q1_matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
1908                        }
1909                        TensorDtype::Q1T => {
1910                            q1t_matmat(self.quant_bytes(), &transformed, b, rows, cols, out, pool)
1911                        }
1912                        TensorDtype::Vbit | TensorDtype::VbitRo => vbitmatmat(
1913                            self.quant_bytes(),
1914                            vbit_offsets,
1915                            &transformed,
1916                            b,
1917                            rows,
1918                            cols,
1919                            out,
1920                            pool,
1921                        ),
1922                        TensorDtype::Q8Row | TensorDtype::Q8_2f => {
1923                            let pre: Vec<std::borrow::Cow<'_, [f32]>> = (0..b)
1924                                .map(|bi| {
1925                                    prescale(
1926                                        &transformed[bi * cols..(bi + 1) * cols],
1927                                        col_field,
1928                                        *dtype,
1929                                    )
1930                                })
1931                                .collect();
1932                            qmatmat(self.quant_bytes(), row_scale, &pre, rows, cols, out, pool)
1933                        }
1934                        _ => unreachable!("unsupported mapped Prism dtype {dtype:?}"),
1935                    }
1936                    return;
1937                }
1938                if *dtype == TensorDtype::Q4Block {
1939                    q4matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1940                    return;
1941                }
1942                if *dtype == TensorDtype::Q4TiledP {
1943                    // GPU batched q4tp GEMM (dequant + f32nt mul_mm on the
1944                    // device); the probe keeps whichever beats the CPU arm.
1945                    // Narrow (prompt-encode) and wide (DiT) batches probe
1946                    // as separate classes — the regimes have opposite
1947                    // winners and one shared verdict locked the wrong arm.
1948                    // Kill switch (gpu::mm_kill): one grossly slow GPU op
1949                    // (a fair-condition op is ≤~100 ms even at 1024px)
1950                    // means the device is contended by another process
1951                    // (e.g. a simulator) — verdicts are per-process, so
1952                    // without the bail the whole render crawls behind
1953                    // someone else's queue.
1954                    if b >= 32
1955                        && b * rows * cols >= 128_000_000
1956                        && cols % 32 == 0
1957                        && !crate::gpu::mm_killed()
1958                        && crate::gpu::enabled_here()
1959                    {
1960                        let class = if b >= 128 {
1961                            crate::gpu::OpClass::MatmatWide
1962                        } else {
1963                            crate::gpu::OpClass::Matmat
1964                        };
1965                        if let Self::Mapped { model, idx, .. } = self {
1966                            // In-process A/B (`CMF_MM_AB=1`). Three
1967                            // wall-clock A/Bs on a shared stand disagreed
1968                            // with each other by 25% on the same change,
1969                            // because the machine drifts between processes
1970                            // and interleaving whole renders does not fix
1971                            // that. Here both arms run back to back on the
1972                            // SAME data inside one call, so whatever the
1973                            // machine is doing, it does to both — and the
1974                            // disagreement between their outputs falls out
1975                            // for free. Doubles the work; a diagnostic,
1976                            // not a mode.
1977                            if crate::mm_ab::on() {
1978                                let mut g = vec![0f32; b * rows];
1979                                let t = std::time::Instant::now();
1980                                let took = crate::gpu::q4tp_matmat(
1981                                    model, *idx, xs_all, b, rows, cols, &mut g,
1982                                );
1983                                let dg = t.elapsed();
1984                                let t = std::time::Instant::now();
1985                                q4tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1986                                let dc = t.elapsed();
1987                                crate::mm_ab::record(b, rows, cols, took, dg, dc, &g, out);
1988                                return;
1989                            }
1990                            let t0 = std::time::Instant::now();
1991                            // A cold call takes the device arm: its sample
1992                            // is discarded either way, and the upload is
1993                            // what the next step needs.
1994                            let resident = crate::gpu::weight_is_resident(model, *idx);
1995                            match crate::gpu::probe_arm_cold_prefers_gpu(class, resident) {
1996                                crate::gpu::ProbeArm::Gpu => {
1997                                    if crate::gpu::q4tp_matmat(
1998                                        model, *idx, xs_all, b, rows, cols, out,
1999                                    ) {
2000                                        let el = t0.elapsed();
2001                                        // Work-proportional budget: ~8× the
2002                                        // fair-device estimate (+20 ms slack).
2003                                        // An absolute cap missed the worst
2004                                        // case — contended ops sit at
2005                                        // 100–240 ms each and still bury a
2006                                        // render whose fair op is 3–9 ms.
2007                                        // Cold ops (first PSO build, buffer
2008                                        // alloc) are exempt: a one-off
2009                                        // ~50 ms compile is not contention.
2010                                        let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
2011                                        let budget = std::time::Duration::from_secs_f64(
2012                                            flops / 1.5e12 * 8.0 + 0.020,
2013                                        );
2014                                        crate::gpu::mm_budget_check(
2015                                            "q4tp matmat",
2016                                            el,
2017                                            budget,
2018                                            crate::gpu::probe_was_cold() || !resident,
2019                                        );
2020                                        crate::gpu::probe_record(class, true, el);
2021                                        return;
2022                                    }
2023                                }
2024                                crate::gpu::ProbeArm::CpuTimed => {
2025                                    q4tp_matmat(
2026                                        self.quant_bytes(),
2027                                        xs_all,
2028                                        b,
2029                                        rows,
2030                                        cols,
2031                                        out,
2032                                        pool,
2033                                    );
2034                                    crate::gpu::probe_record(class, false, t0.elapsed());
2035                                    return;
2036                                }
2037                                crate::gpu::ProbeArm::Cpu => {}
2038                            }
2039                        }
2040                    }
2041                    q4tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2042                    return;
2043                }
2044                if *dtype == TensorDtype::Q2TiledP {
2045                    // Same device arm as q4tp, behind the same probe:
2046                    // the planes differ, the dispatch does not. Without
2047                    // this a q2tp file ran its widest projections on the
2048                    // host while the 4-bit one had the card, which is a
2049                    // codec paying for its size twice.
2050                    if b >= 32
2051                        && b * rows * cols >= 128_000_000
2052                        && cols % 32 == 0
2053                        && !crate::gpu::mm_killed()
2054                        && crate::gpu::enabled_here()
2055                    {
2056                        let class = if b >= 128 {
2057                            crate::gpu::OpClass::MatmatWide
2058                        } else {
2059                            crate::gpu::OpClass::Matmat
2060                        };
2061                        if let Self::Mapped { model, idx, .. } = self {
2062                            let t0 = std::time::Instant::now();
2063                            match crate::gpu::probe_arm(class) {
2064                                crate::gpu::ProbeArm::Gpu => {
2065                                    if crate::gpu::q2tp_matmat(
2066                                        model, *idx, xs_all, b, rows, cols, out,
2067                                    ) {
2068                                        crate::gpu::probe_record(class, true, t0.elapsed());
2069                                        return;
2070                                    }
2071                                }
2072                                crate::gpu::ProbeArm::CpuTimed => {
2073                                    q2tp_matmat(
2074                                        self.quant_bytes(),
2075                                        xs_all,
2076                                        b,
2077                                        rows,
2078                                        cols,
2079                                        out,
2080                                        pool,
2081                                    );
2082                                    crate::gpu::probe_record(class, false, t0.elapsed());
2083                                    return;
2084                                }
2085                                crate::gpu::ProbeArm::Cpu => {}
2086                            }
2087                        }
2088                    }
2089                    // Without a host arm a q2tp tensor falls through to
2090                    // the q8 fallback, which reads it at one BYTE per
2091                    // weight — a 2x overrun that killed pool workers
2092                    // mid-prefill while the dispatcher waited forever.
2093                    q2tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2094                    return;
2095                }
2096                if *dtype == TensorDtype::Q4Tiled {
2097                    // GPU batched q4t GEMM (dequant + f32nt mul_mm on the
2098                    // device); the probe keeps whichever beats the CPU arm.
2099                    // Narrow (prompt-encode) and wide (DiT) batches probe
2100                    // as separate classes — the regimes have opposite
2101                    // winners and one shared verdict locked the wrong arm.
2102                    // Kill switch (gpu::mm_kill): one grossly slow GPU op
2103                    // (a fair-condition op is ≤~100 ms even at 1024px)
2104                    // means the device is contended by another process
2105                    // (e.g. a simulator) — verdicts are per-process, so
2106                    // without the bail the whole render crawls behind
2107                    // someone else's queue.
2108                    if b >= 32
2109                        && b * rows * cols >= 128_000_000
2110                        && cols % 32 == 0
2111                        && !crate::gpu::mm_killed()
2112                        && crate::gpu::enabled_here()
2113                    {
2114                        let class = if b >= 128 {
2115                            crate::gpu::OpClass::MatmatWide
2116                        } else {
2117                            crate::gpu::OpClass::Matmat
2118                        };
2119                        if let Self::Mapped { model, idx, .. } = self {
2120                            let t0 = std::time::Instant::now();
2121                            match crate::gpu::probe_arm(class) {
2122                                crate::gpu::ProbeArm::Gpu => {
2123                                    if crate::gpu::q4t_matmat(
2124                                        model, *idx, xs_all, b, rows, cols, out,
2125                                    ) {
2126                                        let el = t0.elapsed();
2127                                        // Work-proportional budget: ~8× the
2128                                        // fair-device estimate (+20 ms slack).
2129                                        // An absolute cap missed the worst
2130                                        // case — contended ops sit at
2131                                        // 100–240 ms each and still bury a
2132                                        // render whose fair op is 3–9 ms.
2133                                        // Cold ops (first PSO build, buffer
2134                                        // alloc) are exempt: a one-off
2135                                        // ~50 ms compile is not contention.
2136                                        let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
2137                                        let budget = std::time::Duration::from_secs_f64(
2138                                            flops / 1.5e12 * 8.0 + 0.020,
2139                                        );
2140                                        crate::gpu::mm_budget_check(
2141                                            "q4t matmat",
2142                                            el,
2143                                            budget,
2144                                            crate::gpu::probe_was_cold(),
2145                                        );
2146                                        crate::gpu::probe_record(class, true, el);
2147                                        return;
2148                                    }
2149                                }
2150                                crate::gpu::ProbeArm::CpuTimed => {
2151                                    q4t_matmat(
2152                                        self.quant_bytes(),
2153                                        xs_all,
2154                                        b,
2155                                        rows,
2156                                        cols,
2157                                        out,
2158                                        pool,
2159                                    );
2160                                    crate::gpu::probe_record(class, false, t0.elapsed());
2161                                    return;
2162                                }
2163                                crate::gpu::ProbeArm::Cpu => {}
2164                            }
2165                        }
2166                    }
2167                    q4t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2168                    return;
2169                }
2170                if *dtype == TensorDtype::Q1 {
2171                    // GPU batched q1 GEMM for wide prefill (q1_mul_mm on the
2172                    // device); the probe keeps whichever beats the CPU matmat.
2173                    if b >= 32
2174                        && b * rows * cols >= 128_000_000
2175                        && cols % 64 == 0
2176                        && crate::gpu::enabled_here()
2177                    {
2178                        if let Self::Mapped { model, idx, .. } = self {
2179                            let t0 = std::time::Instant::now();
2180                            match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
2181                                crate::gpu::ProbeArm::Gpu => {
2182                                    if crate::gpu::q1_matmat(
2183                                        model, *idx, xs_all, b, rows, cols, out,
2184                                    ) {
2185                                        crate::gpu::probe_record(
2186                                            crate::gpu::OpClass::Matmat,
2187                                            true,
2188                                            t0.elapsed(),
2189                                        );
2190                                        return;
2191                                    }
2192                                }
2193                                crate::gpu::ProbeArm::CpuTimed => {
2194                                    q1_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2195                                    crate::gpu::probe_record(
2196                                        crate::gpu::OpClass::Matmat,
2197                                        false,
2198                                        t0.elapsed(),
2199                                    );
2200                                    return;
2201                                }
2202                                crate::gpu::ProbeArm::Cpu => {}
2203                            }
2204                        }
2205                    }
2206                    q1_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2207                    return;
2208                }
2209                if *dtype == TensorDtype::Q1T {
2210                    // GPU batched GEMM for wide prefill (base + overlay on the
2211                    // device); probe keeps the winner vs the CPU matmat.
2212                    if b >= 32 && b * rows * cols >= 128_000_000 && crate::gpu::enabled_here() {
2213                        if let Self::Mapped { model, idx, .. } = self {
2214                            let t0 = std::time::Instant::now();
2215                            match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
2216                                crate::gpu::ProbeArm::Gpu => {
2217                                    if crate::gpu::q1t_matmat(
2218                                        model, *idx, xs_all, b, rows, cols, out,
2219                                    ) {
2220                                        crate::gpu::probe_record(
2221                                            crate::gpu::OpClass::Matmat,
2222                                            true,
2223                                            t0.elapsed(),
2224                                        );
2225                                        return;
2226                                    }
2227                                }
2228                                crate::gpu::ProbeArm::CpuTimed => {
2229                                    q1t_matmat(
2230                                        self.quant_bytes(),
2231                                        xs_all,
2232                                        b,
2233                                        rows,
2234                                        cols,
2235                                        out,
2236                                        pool,
2237                                    );
2238                                    crate::gpu::probe_record(
2239                                        crate::gpu::OpClass::Matmat,
2240                                        false,
2241                                        t0.elapsed(),
2242                                    );
2243                                    return;
2244                                }
2245                                crate::gpu::ProbeArm::Cpu => {}
2246                            }
2247                        }
2248                    }
2249                    q1t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
2250                    return;
2251                }
2252                if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
2253                    vbitmatmat(
2254                        self.quant_bytes(),
2255                        vbit_offsets,
2256                        xs_all,
2257                        b,
2258                        rows,
2259                        cols,
2260                        out,
2261                        pool,
2262                    );
2263                    return;
2264                }
2265                let pre: Vec<std::borrow::Cow<'_, [f32]>> = (0..b)
2266                    .map(|bi| prescale(&xs_all[bi * cols..(bi + 1) * cols], col_field, *dtype))
2267                    .collect();
2268                // D5: large prefill-batch GEMMs — on the GPU (threshold by
2269                // work volume: submission carries b×rows×cols MACs).
2270                // Runtime probe: the naive GEMM shader + sync readback
2271                // lose to the CPU GEMM on slow driver stacks — alternate
2272                // both arms and keep the winner.
2273                if b >= 8 && b * rows * cols >= 128_000_000 && crate::gpu::enabled_here() {
2274                    if let Self::Mapped { model, idx, .. } = self {
2275                        let t0 = std::time::Instant::now();
2276                        match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
2277                            crate::gpu::ProbeArm::Gpu
2278                                if crate::gpu::probe_deciding(crate::gpu::OpClass::Matmat)
2279                                    && !crate::gpu::q8_resident_or_upload(model, *idx) =>
2280                            {
2281                                // Cold weights during probing: the upload
2282                                // has started, the count runs on the CPU —
2283                                // the GPU arm samples on the next touch.
2284                                let q = self.quant_bytes();
2285                                qmatmat(q, row_scale, &pre, rows, cols, out, pool);
2286                                return;
2287                            }
2288                            crate::gpu::ProbeArm::Gpu => {
2289                                let flat: Vec<f32> =
2290                                    pre.iter().flat_map(|v| v.iter().copied()).collect();
2291                                if crate::gpu::q8_matmat(
2292                                    model, *idx, row_scale, &flat, b, rows, cols, out,
2293                                ) {
2294                                    crate::gpu::probe_record(
2295                                        crate::gpu::OpClass::Matmat,
2296                                        true,
2297                                        t0.elapsed(),
2298                                    );
2299                                    return;
2300                                }
2301                            }
2302                            crate::gpu::ProbeArm::CpuTimed => {
2303                                let q = self.quant_bytes();
2304                                qmatmat(q, row_scale, &pre, rows, cols, out, pool);
2305                                crate::gpu::probe_record(
2306                                    crate::gpu::OpClass::Matmat,
2307                                    false,
2308                                    t0.elapsed(),
2309                                );
2310                                return;
2311                            }
2312                            crate::gpu::ProbeArm::Cpu => {}
2313                        }
2314                    }
2315                }
2316                let q = self.quant_bytes();
2317                qmatmat(q, row_scale, &pre, rows, cols, out, pool);
2318            }
2319        }
2320    }
2321}
2322
2323impl QTensor {
2324    /// The device GEMM this tensor would take, run once on the caller's
2325    /// data — the startup parity probe's arm, and the one place that knows
2326    /// which entry point each codec has.
2327    ///
2328    /// It exists because the probe used to look for a `q4tp` weight by
2329    /// name AND dtype, and a container packed any other way was declared
2330    /// "host path" for the whole render even though its codec had a device
2331    /// GEMM of its own. A gate that only recognizes one codec is a gate
2332    /// that silently downgrades every other one.
2333    pub fn device_matmat(&self, xs: &[f32], b: usize, out: &mut [f32]) -> bool {
2334        let (rows, cols) = (self.rows(), self.cols());
2335        let Self::Mapped {
2336            model,
2337            idx,
2338            dtype,
2339            row_scale,
2340            col_field,
2341            ..
2342        } = self
2343        else {
2344            return false;
2345        };
2346        if crate::prism::has_contract(model) {
2347            return false;
2348        }
2349        match *dtype {
2350            TensorDtype::Q4TiledP => crate::gpu::q4tp_matmat(model, *idx, xs, b, rows, cols, out),
2351            // The two-field codec folds its column field into the
2352            // activation, which leaves a plain per-row int8 GEMM — the
2353            // same kernel `q8_row` uses, on both backends.
2354            TensorDtype::Q8Row | TensorDtype::Q8_2f => {
2355                // The field belongs to the weight; only a backend that cannot
2356                // apply it there makes a scaled copy of the activation.
2357                if *dtype == TensorDtype::Q8_2f
2358                    && std::env::var("CMF_Q8_2F_DEV").as_deref() != Ok("0")
2359                    && crate::gpu::q8_matmat_2f(
2360                        model, *idx, row_scale, col_field, xs, b, rows, cols, out,
2361                    )
2362                {
2363                    return true;
2364                }
2365                let flat: Vec<f32> = (0..b)
2366                    .flat_map(|bi| {
2367                        prescale(&xs[bi * cols..(bi + 1) * cols], col_field, *dtype).into_owned()
2368                    })
2369                    .collect();
2370                crate::gpu::q8_matmat(model, *idx, row_scale, &flat, b, rows, cols, out)
2371            }
2372            _ => false,
2373        }
2374    }
2375
2376    /// Multi-matrix job (roadmap §3 P0): N tensors sharing one input
2377    /// run under a SINGLE pool dispatch — QKV or gate+up cost one
2378    /// barrier instead of N. Per-row math is the exact same kernel as
2379    /// `matvec` (bit-identical outputs); only the dispatch is fused.
2380    /// Falls back to N sequential matvecs when the set is not a uniform
2381    /// q8-family/F32 group or there is no pool.
2382    pub fn matvec_many<const N: usize>(
2383        ts: [&QTensor; N],
2384        x: &[f32],
2385        mut outs: [&mut [f32]; N],
2386        pool: Option<&Pool>,
2387    ) {
2388        let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
2389        if ts.iter().any(|t| t.has_prism_contract()) {
2390            // The fused range kernels have no transform descriptor.  Let
2391            // each tensor's ordinary matvec dispatch perform the explicit
2392            // signed FWHT (and retain CPU fallback for mixed q2tp/q4tp).
2393            for (t, o) in ts.iter().zip(outs.iter_mut()) {
2394                t.matvec(x, o, pool);
2395            }
2396            return;
2397        }
2398        let uniform_q8 = ts.iter().all(|t| {
2399            matches!(
2400                t,
2401                Self::Mapped {
2402                    dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f,
2403                    ..
2404                }
2405            )
2406        });
2407        let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
2408        let uniform_q4 = ts.iter().all(|t| {
2409            matches!(
2410                t,
2411                Self::Mapped {
2412                    dtype: TensorDtype::Q4Block,
2413                    ..
2414                }
2415            )
2416        });
2417        let uniform_vbit = ts.iter().all(|t| {
2418            matches!(
2419                t,
2420                Self::Mapped {
2421                    dtype: TensorDtype::Vbit | TensorDtype::VbitRo,
2422                    ..
2423                }
2424            )
2425        });
2426        let uniform_q1 = ts.iter().all(|t| {
2427            matches!(
2428                t,
2429                Self::Mapped {
2430                    dtype: TensorDtype::Q1,
2431                    ..
2432                }
2433            )
2434        });
2435        let uniform_q1t = ts.iter().all(|t| {
2436            matches!(
2437                t,
2438                Self::Mapped {
2439                    dtype: TensorDtype::Q1T,
2440                    ..
2441                }
2442            )
2443        });
2444        // q4tp is the skeleton dtype of the big MoE files, and without an arm
2445        // here every projection that shares an input paid its own pool
2446        // barrier: DeepSeek-V4's attention step alone hands this function
2447        // wq_a, wkv and both compressors' pairs off the same hidden state.
2448        let uniform_q4tp = ts.iter().all(|t| {
2449            matches!(
2450                t,
2451                Self::Mapped {
2452                    dtype: TensorDtype::Q4TiledP,
2453                    ..
2454                }
2455            )
2456        }) && ts
2457            .iter()
2458            .all(|t| t.cols() == ts[0].cols() && t.cols() % GROUP_SIZE == 0);
2459        let Some(pool) = pool else {
2460            for (t, o) in ts.iter().zip(outs.iter_mut()) {
2461                t.matvec(x, o, None);
2462            }
2463            return;
2464        };
2465        if total_rows < 256
2466            || !(uniform_q8
2467                || uniform_f32
2468                || uniform_q4
2469                || uniform_vbit
2470                || uniform_q1
2471                || uniform_q1t
2472                || uniform_q4tp)
2473        {
2474            for (t, o) in ts.iter().zip(outs.iter_mut()) {
2475                t.matvec(x, o, Some(pool));
2476            }
2477            return;
2478        }
2479
2480        if uniform_q4tp {
2481            // Every tensor's rows laid end to end in one virtual row space,
2482            // so the whole set is ONE dispatch. The per-row body is the
2483            // `q4tp_matvec` arm verbatim — same activation split, same
2484            // accumulation order — so the outputs are bit-identical to the
2485            // sequential calls this replaces.
2486            let cols = ts[0].cols();
2487            let gpr = cols / GROUP_SIZE;
2488            let views: Vec<Q4tpView> = ts
2489                .iter()
2490                .map(|t| Q4tpView::new(t.quant_bytes(), t.rows(), cols))
2491                .collect();
2492            let rows_of: Vec<usize> = ts.iter().map(|t| t.rows()).collect();
2493            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2494            // flat index -> (which tensor, which of its rows)
2495            let locate = |flat: usize| -> (usize, usize) {
2496                let mut acc = 0;
2497                for (i, &r) in rows_of.iter().enumerate() {
2498                    if flat < acc + r {
2499                        return (i, flat - acc);
2500                    }
2501                    acc += r;
2502                }
2503                (rows_of.len() - 1, 0)
2504            };
2505            let (views, outs_addr) = (&views, &outs_addr);
2506            if a8w8_enabled() {
2507                let act = split_act(x);
2508                let act = &act;
2509                let run = |start: usize, end: usize| {
2510                    let mut sc = vec![0f32; gpr];
2511                    for flat in start..end {
2512                        let (t, r) = locate(flat);
2513                        let v = &views[t];
2514                        v.scales_into(r, gpr, &mut sc);
2515                        let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, &sc) * act.sx;
2516                        for &(j, xv) in &act.outliers {
2517                            let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
2518                            acc += w * s * xv;
2519                        }
2520                        // SAFETY: one worker owns each (tensor, row) pair.
2521                        unsafe { *outs_addr[t].at(r) = acc };
2522                    }
2523                };
2524                pool.run_rows(total_rows, &run);
2525            } else {
2526                let run = |start: usize, end: usize| {
2527                    let mut sc = vec![0f32; gpr];
2528                    for flat in start..end {
2529                        let (t, r) = locate(flat);
2530                        let v = &views[t];
2531                        v.scales_into(r, gpr, &mut sc);
2532                        // SAFETY: one worker owns each (tensor, row) pair.
2533                        unsafe { *outs_addr[t].at(r) = q4tp_row_exact(v.nib, r, gpr, x, &sc) };
2534                    }
2535                };
2536                pool.run_rows(total_rows, &run);
2537            }
2538            return;
2539        }
2540
2541        if uniform_q1 {
2542            // One shared activation split + group sums (q1 has no col
2543            // field; the same input feeds every tensor).
2544            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2545            if a8w8_enabled() {
2546                let act = split_act(x);
2547                let gsum = q1_group_sums(&act.xq, ts[0].cols() / GROUP_SIZE);
2548                let (act, gsum) = (&act, &gsum);
2549                let closures: [_; N] = std::array::from_fn(|i| {
2550                    let (bytes, gpr, out) =
2551                        (ts[i].quant_bytes(), ts[i].cols() / GROUP_SIZE, outs_addr[i]);
2552                    move |s: usize, e: usize| q1_range_a8w8(bytes, gpr, act, gsum, out, s, e)
2553                });
2554                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2555                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2556                pool.run_many(&parts);
2557            } else {
2558                let closures: [_; N] = std::array::from_fn(|i| {
2559                    let (bytes, gpr, out) =
2560                        (ts[i].quant_bytes(), ts[i].cols() / GROUP_SIZE, outs_addr[i]);
2561                    move |s: usize, e: usize| q1_range_f32(bytes, gpr, x, out, s, e)
2562                });
2563                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2564                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2565                pool.run_many(&parts);
2566            }
2567            return;
2568        }
2569
2570        if uniform_q1t {
2571            // Q1T batched: one shared activation split + overlay decode,
2572            // all tensors' rows in ONE pool dispatch (saves N−1 dispatches
2573            // and N−1 redundant split_act calls per layer).
2574            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2575            const TILE: usize = cortiq_core::quant::Q1T_TILE;
2576            if a8w8_enabled() {
2577                let act = split_act(x);
2578                let act = &act;
2579                let x_ref = x;
2580                let closures: [_; N] = std::array::from_fn(|i| {
2581                    let bytes = ts[i].quant_bytes();
2582                    let (rows, cols) = (ts[i].rows(), ts[i].cols());
2583                    let gpr = cols / GROUP_SIZE;
2584                    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
2585                    let out = outs_addr[i];
2586                    move |s: usize, e: usize| {
2587                        q1t_range_a8w8(bytes, gpr, rp_off, ent_off, has_ov, act, x_ref, out, s, e)
2588                    }
2589                });
2590                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2591                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2592                pool.run_many(&parts);
2593            } else {
2594                let x_ref = x;
2595                let closures: [_; N] = std::array::from_fn(|i| {
2596                    let bytes = ts[i].quant_bytes();
2597                    let (rows, cols) = (ts[i].rows(), ts[i].cols());
2598                    let gpr = cols / GROUP_SIZE;
2599                    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
2600                    let out = outs_addr[i];
2601                    move |s: usize, e: usize| {
2602                        q1t_range_f32_batch(bytes, gpr, rp_off, ent_off, has_ov, x_ref, out, s, e)
2603                    }
2604                });
2605                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2606                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2607                pool.run_many(&parts);
2608            }
2609            return;
2610        }
2611
2612        if uniform_q4 || uniform_vbit {
2613            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2614            // q4/vbit share one activation split — no per-tensor col field.
2615            if a8w8_enabled() {
2616                let act = split_act(x);
2617                let act = &act;
2618                if uniform_q4 {
2619                    let closures: [_; N] = std::array::from_fn(|i| {
2620                        let (packed, scales) =
2621                            q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
2622                        let (gpr, cols, out) =
2623                            (ts[i].cols() / GROUP_SIZE, ts[i].cols(), outs_addr[i]);
2624                        move |s: usize, e: usize| {
2625                            q4_range_a8w8(packed, scales, gpr, cols, act, out, s, e)
2626                        }
2627                    });
2628                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2629                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2630                    pool.run_many(&parts);
2631                } else {
2632                    let closures: [_; N] = std::array::from_fn(|i| {
2633                        let Self::Mapped { vbit_offsets, .. } = ts[i] else {
2634                            unreachable!()
2635                        };
2636                        let (bytes, rows, cols, out) = (
2637                            ts[i].quant_bytes(),
2638                            ts[i].rows(),
2639                            ts[i].cols(),
2640                            outs_addr[i],
2641                        );
2642                        move |s: usize, e: usize| {
2643                            vbit_range_a8w8(bytes, vbit_offsets, x, act, rows, cols, out, s, e)
2644                        }
2645                    });
2646                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2647                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2648                    pool.run_many(&parts);
2649                }
2650                return;
2651            }
2652            if uniform_q4 {
2653                let closures: [_; N] = std::array::from_fn(|i| {
2654                    let (packed, scales) =
2655                        q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
2656                    let (gpr, out) = (ts[i].cols() / GROUP_SIZE, outs_addr[i]);
2657                    move |s: usize, e: usize| q4_range_f32(packed, scales, gpr, x, out, s, e)
2658                });
2659                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2660                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2661                pool.run_many(&parts);
2662            } else {
2663                let closures: [_; N] = std::array::from_fn(|i| {
2664                    let Self::Mapped { vbit_offsets, .. } = ts[i] else {
2665                        unreachable!()
2666                    };
2667                    let (bytes, rows, cols, out) = (
2668                        ts[i].quant_bytes(),
2669                        ts[i].rows(),
2670                        ts[i].cols(),
2671                        outs_addr[i],
2672                    );
2673                    move |s: usize, e: usize| {
2674                        vbit_range_f32(bytes, vbit_offsets, x, rows, cols, out, s, e)
2675                    }
2676                });
2677                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2678                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2679                pool.run_many(&parts);
2680            }
2681            return;
2682        }
2683
2684        if uniform_f32 {
2685            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2686            let closures: [_; N] = std::array::from_fn(|i| {
2687                let Self::F32 { data, cols, .. } = ts[i] else {
2688                    unreachable!()
2689                };
2690                let out = outs_addr[i];
2691                move |start: usize, end: usize| {
2692                    for o in start..end {
2693                        let row = &data[o * cols..(o + 1) * cols];
2694                        let mut sum = 0.0f32;
2695                        for j in 0..*cols {
2696                            sum += row[j] * x[j];
2697                        }
2698                        // SAFETY: disjoint (tensor, row) cells per worker.
2699                        unsafe { *out.at(o) = sum };
2700                    }
2701                }
2702            });
2703            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2704                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2705            pool.run_many(&parts);
2706            return;
2707        }
2708
2709        // Uniform q8-family: per-tensor prescale (q8_2f col fields
2710        // differ per tensor) + the shared range kernels.
2711        struct Ctx<'a> {
2712            bytes: &'a [u8],
2713            #[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
2714            rep: &'a [u8],
2715            row_scale: &'a [f32],
2716            cols: usize,
2717            xs: std::borrow::Cow<'a, [f32]>,
2718        }
2719        let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
2720            let Self::Mapped {
2721                dtype,
2722                cols,
2723                row_scale,
2724                col_field,
2725                repack,
2726                ..
2727            } = ts[i]
2728            else {
2729                unreachable!()
2730            };
2731            Ctx {
2732                bytes: ts[i].quant_bytes(),
2733                rep: repack,
2734                row_scale,
2735                cols: *cols,
2736                xs: prescale(x, col_field, *dtype),
2737            }
2738        });
2739        let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2740        #[cfg(target_arch = "aarch64")]
2741        if sdot_enabled() {
2742            let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
2743            let closures: [_; N] = std::array::from_fn(|i| {
2744                let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
2745                move |start: usize, end: usize| {
2746                    q8_range_sdot(c.bytes, c.rep, c.row_scale, act, c.cols, out, start, end)
2747                }
2748            });
2749            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2750                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2751            pool.run_many(&parts);
2752            return;
2753        }
2754        #[cfg(target_arch = "x86_64")]
2755        if avx2_a8w8_enabled() {
2756            let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
2757            let closures: [_; N] = std::array::from_fn(|i| {
2758                let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
2759                move |start: usize, end: usize| {
2760                    q8_range_avx2(c.bytes, c.row_scale, act, c.cols, out, start, end)
2761                }
2762            });
2763            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2764                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2765            pool.run_many(&parts);
2766            return;
2767        }
2768        let closures: [_; N] = std::array::from_fn(|i| {
2769            let (c, out) = (&ctxs[i], outs_addr[i]);
2770            move |start: usize, end: usize| {
2771                q8_range_f32(c.bytes, c.row_scale, &c.xs, c.cols, out, start, end)
2772            }
2773        });
2774        let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2775            std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2776        pool.run_many(&parts);
2777    }
2778}
2779
2780impl QTensor {
2781    /// Pair-input multi-matrix job: N tensors × 2 shared inputs under a
2782    /// single pool dispatch — the MTP/pair decode path publishes one job
2783    /// for Q/K/V (and one for gate+up) instead of one per tensor.
2784    /// Per-row math is exactly `matvec2`'s kernels; bit-identical.
2785    #[allow(clippy::needless_range_loop)]
2786    pub fn matvec2_many<const N: usize>(
2787        ts: [&QTensor; N],
2788        x1: &[f32],
2789        x2: &[f32],
2790        mut o1s: [&mut [f32]; N],
2791        mut o2s: [&mut [f32]; N],
2792        pool: Option<&Pool>,
2793    ) {
2794        let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
2795        if ts.iter().any(|t| t.has_prism_contract()) {
2796            for i in 0..N {
2797                ts[i].matvec2(x1, x2, o1s[i], o2s[i], pool);
2798            }
2799            return;
2800        }
2801        let uniform_q8 = ts.iter().all(|t| {
2802            matches!(
2803                t,
2804                Self::Mapped {
2805                    dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f,
2806                    ..
2807                }
2808            )
2809        });
2810        let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
2811        let uniform_q4 = ts.iter().all(|t| {
2812            matches!(
2813                t,
2814                Self::Mapped {
2815                    dtype: TensorDtype::Q4Block,
2816                    ..
2817                }
2818            )
2819        });
2820        let uniform_vbit = ts.iter().all(|t| {
2821            matches!(
2822                t,
2823                Self::Mapped {
2824                    dtype: TensorDtype::Vbit | TensorDtype::VbitRo,
2825                    ..
2826                }
2827            )
2828        });
2829        let fusable = pool.is_some()
2830            && total_rows >= 256
2831            && (uniform_q8 || uniform_f32 || uniform_q4 || uniform_vbit);
2832        if !fusable {
2833            for i in 0..N {
2834                ts[i].matvec2(x1, x2, o1s[i], o2s[i], pool);
2835            }
2836            return;
2837        }
2838        let pool = pool.unwrap();
2839
2840        if uniform_q4 || uniform_vbit {
2841            let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
2842            let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
2843            // q4/vbit share activation splits — no per-tensor col field.
2844            if a8w8_enabled() {
2845                let a1 = split_act(x1);
2846                let a2 = split_act(x2);
2847                let (a1, a2) = (&a1, &a2);
2848                if uniform_q4 {
2849                    let closures: [_; N] = std::array::from_fn(|i| {
2850                        let (packed, scales) =
2851                            q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
2852                        let (gpr, cols, o1, o2) =
2853                            (ts[i].cols() / GROUP_SIZE, ts[i].cols(), p1[i], p2[i]);
2854                        move |s: usize, e: usize| {
2855                            q4_range2_a8w8(packed, scales, gpr, cols, a1, a2, o1, o2, s, e)
2856                        }
2857                    });
2858                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2859                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2860                    pool.run_many(&parts);
2861                } else {
2862                    let closures: [_; N] = std::array::from_fn(|i| {
2863                        let Self::Mapped { vbit_offsets, .. } = ts[i] else {
2864                            unreachable!()
2865                        };
2866                        let (bytes, rows, cols, o1, o2) = (
2867                            ts[i].quant_bytes(),
2868                            ts[i].rows(),
2869                            ts[i].cols(),
2870                            p1[i],
2871                            p2[i],
2872                        );
2873                        move |s: usize, e: usize| {
2874                            vbit_range2_a8w8(
2875                                bytes,
2876                                vbit_offsets,
2877                                x1,
2878                                x2,
2879                                a1,
2880                                a2,
2881                                rows,
2882                                cols,
2883                                o1,
2884                                o2,
2885                                s,
2886                                e,
2887                            )
2888                        }
2889                    });
2890                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2891                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2892                    pool.run_many(&parts);
2893                }
2894                return;
2895            }
2896            if uniform_q4 {
2897                let closures: [_; N] = std::array::from_fn(|i| {
2898                    let (packed, scales) =
2899                        q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
2900                    let (gpr, o1, o2) = (ts[i].cols() / GROUP_SIZE, p1[i], p2[i]);
2901                    move |s: usize, e: usize| {
2902                        q4_range2_f32(packed, scales, gpr, x1, x2, o1, o2, s, e)
2903                    }
2904                });
2905                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2906                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2907                pool.run_many(&parts);
2908            } else {
2909                let closures: [_; N] = std::array::from_fn(|i| {
2910                    let Self::Mapped { vbit_offsets, .. } = ts[i] else {
2911                        unreachable!()
2912                    };
2913                    let (bytes, rows, cols, o1, o2) = (
2914                        ts[i].quant_bytes(),
2915                        ts[i].rows(),
2916                        ts[i].cols(),
2917                        p1[i],
2918                        p2[i],
2919                    );
2920                    move |s: usize, e: usize| {
2921                        vbit_range2_f32(bytes, vbit_offsets, x1, x2, rows, cols, o1, o2, s, e)
2922                    }
2923                });
2924                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2925                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2926                pool.run_many(&parts);
2927            }
2928            return;
2929        }
2930
2931        if uniform_f32 {
2932            let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
2933            let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
2934            let closures: [_; N] = std::array::from_fn(|i| {
2935                let Self::F32 { data, cols, .. } = ts[i] else {
2936                    unreachable!()
2937                };
2938                let (o1, o2) = (p1[i], p2[i]);
2939                move |start: usize, end: usize| {
2940                    for o in start..end {
2941                        let row = &data[o * cols..(o + 1) * cols];
2942                        let (mut s1, mut s2) = (0.0f32, 0.0f32);
2943                        for j in 0..*cols {
2944                            s1 += row[j] * x1[j];
2945                            s2 += row[j] * x2[j];
2946                        }
2947                        // SAFETY: disjoint (tensor, row) cells per worker.
2948                        unsafe {
2949                            *o1.at(o) = s1;
2950                            *o2.at(o) = s2;
2951                        }
2952                    }
2953                }
2954            });
2955            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2956                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2957            pool.run_many(&parts);
2958            return;
2959        }
2960
2961        struct Ctx<'a> {
2962            bytes: &'a [u8],
2963            row_scale: &'a [f32],
2964            cols: usize,
2965            xs1: std::borrow::Cow<'a, [f32]>,
2966            xs2: std::borrow::Cow<'a, [f32]>,
2967        }
2968        let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
2969            let Self::Mapped {
2970                dtype,
2971                cols,
2972                row_scale,
2973                col_field,
2974                ..
2975            } = ts[i]
2976            else {
2977                unreachable!()
2978            };
2979            Ctx {
2980                bytes: ts[i].quant_bytes(),
2981                row_scale,
2982                cols: *cols,
2983                xs1: prescale(x1, col_field, *dtype),
2984                xs2: prescale(x2, col_field, *dtype),
2985            }
2986        });
2987        let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
2988        let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
2989        #[cfg(target_arch = "aarch64")]
2990        if sdot_enabled() {
2991            let acts: [(SplitAct, SplitAct); N] =
2992                std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
2993            let closures: [_; N] = std::array::from_fn(|i| {
2994                let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
2995                move |start: usize, end: usize| {
2996                    q8_range2_sdot(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
2997                }
2998            });
2999            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3000                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3001            pool.run_many(&parts);
3002            return;
3003        }
3004        #[cfg(target_arch = "x86_64")]
3005        if avx2_a8w8_enabled() {
3006            let acts: [(SplitAct, SplitAct); N] =
3007                std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
3008            let closures: [_; N] = std::array::from_fn(|i| {
3009                let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
3010                move |start: usize, end: usize| {
3011                    q8_range2_avx2(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
3012                }
3013            });
3014            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3015                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3016            pool.run_many(&parts);
3017            return;
3018        }
3019        let closures: [_; N] = std::array::from_fn(|i| {
3020            let (c, o1, o2) = (&ctxs[i], p1[i], p2[i]);
3021            move |start: usize, end: usize| {
3022                q8_range2_f32(
3023                    c.bytes,
3024                    c.row_scale,
3025                    &c.xs1,
3026                    &c.xs2,
3027                    c.cols,
3028                    o1,
3029                    o2,
3030                    start,
3031                    end,
3032                )
3033            }
3034        });
3035        let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
3036            std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
3037        pool.run_many(&parts);
3038    }
3039
3040    /// Fused gate+up matvec with SiLU·mul: for each row r, computes
3041    /// `silu(gate·x) * (up·x)` and writes to `out[r]`. ONE pool dispatch,
3042    /// no intermediate g/u buffers, no separate silu pass. Falls back
3043    /// (returns false) for unsupported dtype combos.
3044    pub fn matvec_silu_mul(
3045        gate: &QTensor,
3046        up: &QTensor,
3047        x: &[f32],
3048        out: &mut [f32],
3049        pool: Option<&Pool>,
3050    ) -> bool {
3051        Self::matvec_silu_mul_limited(gate, up, x, out, 0.0, pool)
3052    }
3053
3054    /// Fused gate+up+SiLU with the GLM asymmetrical clamp.  `limit == 0`
3055    /// preserves the historical unclamped helper; a positive limit clamps
3056    /// `up` to both sides and `gate` only from above, matching the GLM
3057    /// SwiGLU reference.  Keeping the limit in the row kernel avoids the two
3058    /// intermediate vectors and the extra combine pass on the Q2TP experts.
3059    pub fn matvec_silu_mul_limited(
3060        gate: &QTensor,
3061        up: &QTensor,
3062        x: &[f32],
3063        out: &mut [f32],
3064        limit: f32,
3065        pool: Option<&Pool>,
3066    ) -> bool {
3067        if gate.has_prism_contract() || up.has_prism_contract() {
3068            // The fused gate/up kernels consume x directly.  Prism requires
3069            // a per-matrix signed FWHT, so the caller must use two ordinary
3070            // descriptor-aware matvecs instead of an unrotated fast path.
3071            return false;
3072        }
3073        let inter = gate.rows();
3074        debug_assert_eq!(up.rows(), inter);
3075        debug_assert_eq!(out.len(), inter);
3076        debug_assert_eq!(gate.cols(), up.cols());
3077        if !a8w8_enabled() {
3078            return false;
3079        }
3080        let act = split_act(x);
3081        let act = &act;
3082        let x_ref = x;
3083        let out_addr = SendMut(out.as_mut_ptr());
3084
3085        match (gate, up) {
3086            // Q4Block gate + Q4Block up (most common mobile q4 models)
3087            (
3088                Self::Mapped {
3089                    dtype: TensorDtype::Q4Block,
3090                    ..
3091                },
3092                Self::Mapped {
3093                    dtype: TensorDtype::Q4Block,
3094                    ..
3095                },
3096            ) => {
3097                let (gp, gs) = q4_split(gate.quant_bytes(), gate.rows(), gate.cols());
3098                let (up_p, up_s) = q4_split(up.quant_bytes(), up.rows(), up.cols());
3099                let gpr = gate.cols() / GROUP_SIZE;
3100                let cols = gate.cols();
3101                let run = move |start: usize, end: usize| {
3102                    for r in start..end {
3103                        let mut gv = dot_q4_row_i8(gp, gs, r * gpr, gpr, &act.xq) * act.sx;
3104                        let mut uv = dot_q4_row_i8(up_p, up_s, r * gpr, gpr, &act.xq) * act.sx;
3105                        for &(j, xv) in &act.outliers {
3106                            let flat = r * cols + j;
3107                            let gb = gp[flat / 2];
3108                            let gn = if flat & 1 == 0 { gb & 0x0F } else { gb >> 4 };
3109                            let gsc = f16_to_f32(u16::from_le_bytes([
3110                                gs[(flat / GROUP_SIZE) * 2],
3111                                gs[(flat / GROUP_SIZE) * 2 + 1],
3112                            ]));
3113                            gv += ((gn as i32 - 8) as f32) * gsc * xv;
3114                            let ub = up_p[flat / 2];
3115                            let un = if flat & 1 == 0 { ub & 0x0F } else { ub >> 4 };
3116                            let usc = f16_to_f32(u16::from_le_bytes([
3117                                up_s[(flat / GROUP_SIZE) * 2],
3118                                up_s[(flat / GROUP_SIZE) * 2 + 1],
3119                            ]));
3120                            uv += ((un as i32 - 8) as f32) * usc * xv;
3121                        }
3122                        let silu_g = gv / (1.0 + (-gv).exp());
3123                        // SAFETY: disjoint row ranges per worker.
3124                        unsafe { *out_addr.at(r) = silu_g * uv };
3125                    }
3126                };
3127                dispatch_rows(pool, inter, &run);
3128                true
3129            }
3130            // Q4Tiled gate + Q4Tiled up — one row pass, both tile
3131            // streams sequential, silu·mul fused (same per-row math as
3132            // `q4t_matvec`).
3133            (
3134                Self::Mapped {
3135                    dtype: TensorDtype::Q4Tiled,
3136                    ..
3137                },
3138                Self::Mapped {
3139                    dtype: TensorDtype::Q4Tiled,
3140                    ..
3141                },
3142            ) => {
3143                let g_bytes = gate.quant_bytes();
3144                let u_bytes = up.quant_bytes();
3145                let gpr = gate.cols() / GROUP_SIZE;
3146                let run = move |start: usize, end: usize| {
3147                    for r in start..end {
3148                        let mut gv = dot_q4t_row_i8(g_bytes, r, gpr, &act.xq) * act.sx;
3149                        let mut uv = dot_q4t_row_i8(u_bytes, r, gpr, &act.xq) * act.sx;
3150                        for &(j, xv) in &act.outliers {
3151                            let (w, s) = q4t_outlier(g_bytes, r, gpr, j);
3152                            gv += w * s * xv;
3153                            let (w, s) = q4t_outlier(u_bytes, r, gpr, j);
3154                            uv += w * s * xv;
3155                        }
3156                        let silu_g = gv / (1.0 + (-gv).exp());
3157                        // SAFETY: disjoint row ranges per worker.
3158                        unsafe { *out_addr.at(r) = silu_g * uv };
3159                    }
3160                };
3161                dispatch_rows(pool, inter, &run);
3162                true
3163            }
3164            // Q4TiledP gate + Q4TiledP up — the same fused row pass, with
3165            // each row's two ladders built once and spent on both streams.
3166            (
3167                Self::Mapped {
3168                    dtype: TensorDtype::Q4TiledP,
3169                    ..
3170                },
3171                Self::Mapped {
3172                    dtype: TensorDtype::Q4TiledP,
3173                    ..
3174                },
3175            ) => {
3176                let cols = gate.cols();
3177                let gpr = cols / GROUP_SIZE;
3178                let gv_view = Q4tpView::new(gate.quant_bytes(), inter, cols);
3179                let uv_view = Q4tpView::new(up.quant_bytes(), inter, cols);
3180                let run = |start: usize, end: usize| {
3181                    let (mut gsc, mut usc) = (vec![0f32; gpr], vec![0f32; gpr]);
3182                    for r in start..end {
3183                        gv_view.scales_into(r, gpr, &mut gsc);
3184                        uv_view.scales_into(r, gpr, &mut usc);
3185                        let mut gv = dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, &gsc) * act.sx;
3186                        let mut uv = dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, &usc) * act.sx;
3187                        for &(j, xv) in &act.outliers {
3188                            let (w, s) = q4tp_outlier(gv_view.nib, r, gpr, j, &gsc);
3189                            gv += w * s * xv;
3190                            let (w, s) = q4tp_outlier(uv_view.nib, r, gpr, j, &usc);
3191                            uv += w * s * xv;
3192                        }
3193                        let silu_g = gv / (1.0 + (-gv).exp());
3194                        // SAFETY: disjoint row ranges per worker.
3195                        unsafe { *out_addr.at(r) = silu_g * uv };
3196                    }
3197                };
3198                dispatch_rows(pool, inter, &run);
3199                true
3200            }
3201            // Q1 gate + Q1 up — one row pass over both sign streams,
3202            // silu·mul fused (the per-row math of `q1_range_a8w8`); the
3203            // activation group sums are shared by both streams. Without
3204            // this arm a q1 dense FFN paid two dispatches + a combine
3205            // loop — the exact barrier this function exists to remove.
3206            (
3207                Self::Mapped {
3208                    dtype: TensorDtype::Q1,
3209                    ..
3210                },
3211                Self::Mapped {
3212                    dtype: TensorDtype::Q1,
3213                    ..
3214                },
3215            ) => {
3216                let g_bytes = gate.quant_bytes();
3217                let u_bytes = up.quant_bytes();
3218                let gpr = gate.cols() / GROUP_SIZE;
3219                let gsum = q1_group_sums(&act.xq, gpr);
3220                let gsum = &gsum;
3221                let run = move |start: usize, end: usize| {
3222                    for r in start..end {
3223                        let mut gv = dot_q1_row_i8(g_bytes, r, gpr, &act.xq, gsum) * act.sx;
3224                        let mut uv = dot_q1_row_i8(u_bytes, r, gpr, &act.xq, gsum) * act.sx;
3225                        for &(j, xv) in &act.outliers {
3226                            let (w, s) = q1_outlier(g_bytes, r, gpr, j);
3227                            gv += w * s * xv;
3228                            let (w, s) = q1_outlier(u_bytes, r, gpr, j);
3229                            uv += w * s * xv;
3230                        }
3231                        let silu_g = gv / (1.0 + (-gv).exp());
3232                        // SAFETY: disjoint row ranges per worker.
3233                        unsafe { *out_addr.at(r) = silu_g * uv };
3234                    }
3235                };
3236                dispatch_rows(pool, inter, &run);
3237                true
3238            }
3239            // Q2TiledP gate + Q2TiledP up — the 2-bit expert pair (MoE
3240            // FFNs of the W2 class): one row pass, both ladders built
3241            // once, integer code dots with shared group sums.
3242            (
3243                Self::Mapped {
3244                    dtype: TensorDtype::Q2TiledP,
3245                    ..
3246                },
3247                Self::Mapped {
3248                    dtype: TensorDtype::Q2TiledP,
3249                    ..
3250                },
3251            ) => {
3252                let cols = gate.cols();
3253                let gpr = cols / GROUP_SIZE;
3254                let gv_view = Q4tpView::new_q2(gate.quant_bytes(), inter, cols);
3255                let uv_view = Q4tpView::new_q2(up.quant_bytes(), inter, cols);
3256                let gsum = q1_group_sums(&act.xq, gpr);
3257                let gsum = &gsum;
3258                let run = move |start: usize, end: usize| {
3259                    let (mut gsc, mut usc) = (vec![0f32; gpr], vec![0f32; gpr]);
3260                    for r in start..end {
3261                        gv_view.scales_into(r, gpr, &mut gsc);
3262                        uv_view.scales_into(r, gpr, &mut usc);
3263                        let mut gv =
3264                            dot_q2tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsum, &gsc) * act.sx;
3265                        let mut uv =
3266                            dot_q2tp_row_i8(uv_view.nib, r, gpr, &act.xq, gsum, &usc) * act.sx;
3267                        for &(j, xv) in &act.outliers {
3268                            let (w, s) = q2tp_outlier(gv_view.nib, r, gpr, j, &gsc);
3269                            gv += w * s * xv;
3270                            let (w, s) = q2tp_outlier(uv_view.nib, r, gpr, j, &usc);
3271                            uv += w * s * xv;
3272                        }
3273                        let silu_g = gv / (1.0 + (-gv).exp());
3274                        // SAFETY: disjoint row ranges per worker.
3275                        unsafe { *out_addr.at(r) = silu_g * uv };
3276                    }
3277                };
3278                dispatch_rows(pool, inter, &run);
3279                true
3280            }
3281            // Q8Row gate + Q8Row up — one row pass over both i8 streams.
3282            // Q8_2f stays out on purpose: its column field prescales the
3283            // activations PER TENSOR, which breaks this fn's shared
3284            // split_act contract — it keeps the two-dispatch path.
3285            (
3286                Self::Mapped {
3287                    dtype: TensorDtype::Q8Row,
3288                    row_scale: g_rs,
3289                    ..
3290                },
3291                Self::Mapped {
3292                    dtype: TensorDtype::Q8Row,
3293                    row_scale: u_rs,
3294                    ..
3295                },
3296            ) => {
3297                let g_bytes = gate.quant_bytes();
3298                let u_bytes = up.quant_bytes();
3299                let cols = gate.cols();
3300                let run = move |start: usize, end: usize| {
3301                    for r in start..end {
3302                        let gv = q8_row_dot(&g_bytes[r * cols..(r + 1) * cols], act) * g_rs[r];
3303                        let uv = q8_row_dot(&u_bytes[r * cols..(r + 1) * cols], act) * u_rs[r];
3304                        let silu_g = gv / (1.0 + (-gv).exp());
3305                        // SAFETY: disjoint row ranges per worker.
3306                        unsafe { *out_addr.at(r) = silu_g * uv };
3307                    }
3308                };
3309                dispatch_rows(pool, inter, &run);
3310                true
3311            }
3312            // Q1T gate + Q1T up
3313            (
3314                Self::Mapped {
3315                    dtype: TensorDtype::Q1T,
3316                    ..
3317                },
3318                Self::Mapped {
3319                    dtype: TensorDtype::Q1T,
3320                    ..
3321                },
3322            ) => {
3323                const TILE: usize = cortiq_core::quant::Q1T_TILE;
3324                let g_bytes = gate.quant_bytes();
3325                let u_bytes = up.quant_bytes();
3326                let gpr = gate.cols() / GROUP_SIZE;
3327                let (g_rp, g_ent, g_ov) = q1t_overlay(g_bytes, inter * gpr * TILE, inter);
3328                let (u_rp, u_ent, u_ov) = q1t_overlay(u_bytes, inter * gpr * TILE, inter);
3329                let run = move |start: usize, end: usize| {
3330                    for r in start..end {
3331                        let mut gv = q1t_dot_row_i8(g_bytes, r, gpr, &act.xq) * act.sx;
3332                        let mut uv = q1t_dot_row_i8(u_bytes, r, gpr, &act.xq) * act.sx;
3333                        for &(j, xv) in &act.outliers {
3334                            gv += q1t_base_weight(g_bytes, r, gpr, j) * xv;
3335                            uv += q1t_base_weight(u_bytes, r, gpr, j) * xv;
3336                        }
3337                        gv += q1t_row_outlier_correction(g_bytes, r, g_rp, g_ent, g_ov, x_ref);
3338                        uv += q1t_row_outlier_correction(u_bytes, r, u_rp, u_ent, u_ov, x_ref);
3339                        let silu_g = gv / (1.0 + (-gv).exp());
3340                        // SAFETY: disjoint row ranges per worker.
3341                        unsafe { *out_addr.at(r) = silu_g * uv };
3342                    }
3343                };
3344                dispatch_rows(pool, inter, &run);
3345                true
3346            }
3347            _ => false,
3348        }
3349    }
3350
3351    /// Every routed expert's fused gate/up/SiLU under ONE pool dispatch.
3352    ///
3353    /// The per-expert path pays a pool barrier per expert per stage: at 9
3354    /// experts over 40 layers that is ~720 barriers a token, and a decode
3355    /// profile of Qwen3.6-35B-A3B showed the pool parked in
3356    /// `psynch_cvwait` about twice as long as it spent computing. Laying
3357    /// every expert's rows end-to-end in one virtual row space collapses
3358    /// the stage to a single dispatch. The per-row body is the
3359    /// single-expert q4tp arm verbatim, so outputs are bit-identical.
3360    ///
3361    /// `false` = something is outside the fused q4tp kernel (dtype, shape,
3362    /// or the `CMF_SDOT=0` exact contract); the caller walks the ordinary
3363    /// per-expert path.
3364    pub fn moe_gate_up_many(
3365        pairs: &[(&QTensor, &QTensor)],
3366        x: &[f32],
3367        outs: &mut [Vec<f32>],
3368        pool: Option<&Pool>,
3369    ) -> bool {
3370        if pairs.is_empty() || pairs.len() != outs.len() || !a8w8_enabled() {
3371            return false;
3372        }
3373        let inter = pairs[0].0.rows();
3374        let cols = pairs[0].0.cols();
3375        if cols % GROUP_SIZE != 0 {
3376            return false;
3377        }
3378        let gpr = cols / GROUP_SIZE;
3379        // Uniform layout across every routed pair: q4tp, or the 2-bit
3380        // profile's q2tp gate/up (the W2 class). Mixed sets refuse.
3381        let q2 = matches!(
3382            pairs[0].0,
3383            Self::Mapped {
3384                dtype: TensorDtype::Q2TiledP,
3385                ..
3386            }
3387        );
3388        let want = if q2 {
3389            TensorDtype::Q2TiledP
3390        } else {
3391            TensorDtype::Q4TiledP
3392        };
3393        let mut views = Vec::with_capacity(pairs.len() * 2);
3394        for ((g, u), o) in pairs.iter().zip(outs.iter()) {
3395            let both = matches!(g, Self::Mapped { dtype, .. } if *dtype == want)
3396                && matches!(u, Self::Mapped { dtype, .. } if *dtype == want);
3397            if !both
3398                || g.rows() != inter
3399                || u.rows() != inter
3400                || g.cols() != cols
3401                || u.cols() != cols
3402                || o.len() != inter
3403            {
3404                return false;
3405            }
3406            let mk = if q2 { Q4tpView::new_q2 } else { Q4tpView::new };
3407            views.push(mk(g.quant_bytes(), inter, cols));
3408            views.push(mk(u.quant_bytes(), inter, cols));
3409        }
3410        let act = split_act(x);
3411        let gsum = if q2 {
3412            q1_group_sums(&act.xq, gpr)
3413        } else {
3414            Vec::new()
3415        };
3416        let (act, gsum) = (&act, &gsum);
3417        let ptrs: Vec<SendMut> = outs.iter_mut().map(|o| SendMut(o.as_mut_ptr())).collect();
3418        let (views, ptrs) = (&views, &ptrs);
3419        let run = |start: usize, end: usize| {
3420            let (mut gsc, mut usc) = (vec![0f32; gpr], vec![0f32; gpr]);
3421            for flat in start..end {
3422                let (e, r) = (flat / inter, flat % inter);
3423                let gv_view = &views[e * 2];
3424                let uv_view = &views[e * 2 + 1];
3425                gv_view.scales_into(r, gpr, &mut gsc);
3426                uv_view.scales_into(r, gpr, &mut usc);
3427                let (mut gv, mut uv) = if q2 {
3428                    (
3429                        dot_q2tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsum, &gsc) * act.sx,
3430                        dot_q2tp_row_i8(uv_view.nib, r, gpr, &act.xq, gsum, &usc) * act.sx,
3431                    )
3432                } else {
3433                    (
3434                        dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, &gsc) * act.sx,
3435                        dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, &usc) * act.sx,
3436                    )
3437                };
3438                for &(j, xv) in &act.outliers {
3439                    let (og, ou) = if q2 {
3440                        (
3441                            q2tp_outlier(gv_view.nib, r, gpr, j, &gsc),
3442                            q2tp_outlier(uv_view.nib, r, gpr, j, &usc),
3443                        )
3444                    } else {
3445                        (
3446                            q4tp_outlier(gv_view.nib, r, gpr, j, &gsc),
3447                            q4tp_outlier(uv_view.nib, r, gpr, j, &usc),
3448                        )
3449                    };
3450                    gv += og.0 * og.1 * xv;
3451                    uv += ou.0 * ou.1 * xv;
3452                }
3453                let silu_g = gv / (1.0 + (-gv).exp());
3454                // SAFETY: one worker owns each (expert, row) pair.
3455                unsafe { *ptrs[e].at(r) = silu_g * uv };
3456            }
3457        };
3458        dispatch_rows(pool, pairs.len() * inter, &run);
3459        true
3460    }
3461
3462    /// Every routed expert's down projection, weighted and summed into
3463    /// `out`, under ONE pool dispatch.
3464    ///
3465    /// Partitioned by OUTPUT row rather than by expert: each row is owned
3466    /// by a single worker, so the experts are summed in the caller's order
3467    /// — the same sequence of f32 adds the serial `out[i] += w·eo[i]` loop
3468    /// performs, hence bit-identical. Partitioning by expert instead would
3469    /// race on the shared accumulator.
3470    pub fn moe_down_many(
3471        downs: &[&QTensor],
3472        gs: &[Vec<f32>],
3473        weights: &[f32],
3474        out: &mut [f32],
3475        pool: Option<&Pool>,
3476    ) -> bool {
3477        if downs.is_empty()
3478            || downs.len() != gs.len()
3479            || downs.len() != weights.len()
3480            || !a8w8_enabled()
3481        {
3482            return false;
3483        }
3484        let rows = out.len();
3485        let cols = downs[0].cols();
3486        if cols % GROUP_SIZE != 0 {
3487            return false;
3488        }
3489        let gpr = cols / GROUP_SIZE;
3490        let mut views = Vec::with_capacity(downs.len());
3491        for (d, g) in downs.iter().zip(gs.iter()) {
3492            if !matches!(
3493                d,
3494                Self::Mapped {
3495                    dtype: TensorDtype::Q4TiledP,
3496                    ..
3497                }
3498            ) || d.rows() != rows
3499                || d.cols() != cols
3500                || g.len() != cols
3501            {
3502                return false;
3503            }
3504            views.push(Q4tpView::new(d.quant_bytes(), rows, cols));
3505        }
3506        // One int8 split per expert — the activation vectors differ.
3507        let acts: Vec<SplitAct> = gs.iter().map(|g| split_act(g)).collect();
3508        // Partitioned by OUTPUT row, with the experts folded inside: each
3509        // row is owned by one worker, so they are summed in the caller's
3510        // order — the same f32 sequence the serial `out[i] += w·eo[i]`
3511        // loop produces. Partitioning by expert instead would either race
3512        // on the accumulator or need a scratch plane and a second pass;
3513        // measured, that variant was a wash, so this keeps the simpler
3514        // shape.
3515        let out_addr = SendMut(out.as_mut_ptr());
3516        let (views, acts, weights) = (&views, &acts, &weights);
3517        let run = |start: usize, end: usize| {
3518            let mut sc = vec![0f32; gpr];
3519            for r in start..end {
3520                let mut acc = 0f32;
3521                for (e, v) in views.iter().enumerate() {
3522                    v.scales_into(r, gpr, &mut sc);
3523                    let a = &acts[e];
3524                    let mut d = dot_q4tp_row_i8(v.nib, r, gpr, &a.xq, &sc) * a.sx;
3525                    for &(j, xv) in &a.outliers {
3526                        let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
3527                        d += w * s * xv;
3528                    }
3529                    acc += weights[e] * d;
3530                }
3531                // SAFETY: disjoint row ranges per worker.
3532                unsafe { *out_addr.at(r) = acc };
3533            }
3534        };
3535        dispatch_rows(pool, rows, &run);
3536        true
3537    }
3538}
3539
3540/// Batched q8 kernel: same math as qmatvec, the row makes a single
3541/// pass from memory for the whole batch.
3542/// Accelerate CBLAS — the Apple AMX matrix units, the same engine
3543/// llama.cpp's `-ngl 0` prefill rides via ggml-blas.
3544#[cfg(target_os = "macos")]
3545mod accel_blas {
3546    #[link(name = "Accelerate", kind = "framework")]
3547    unsafe extern "C" {
3548        pub fn cblas_sgemm(
3549            order: i32,
3550            trans_a: i32,
3551            trans_b: i32,
3552            m: i32,
3553            n: i32,
3554            k: i32,
3555            alpha: f32,
3556            a: *const f32,
3557            lda: i32,
3558            b: *const f32,
3559            ldb: i32,
3560            beta: f32,
3561            c: *mut f32,
3562            ldc: i32,
3563        );
3564    }
3565}
3566
3567#[cfg(target_os = "macos")]
3568pub(crate) fn accel_gemm_enabled() -> bool {
3569    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
3570    *ON.get_or_init(|| std::env::var("CMF_ACCEL").map(|v| v != "0").unwrap_or(true))
3571}
3572
3573/// Off macOS the "accel" GEMM is the portable NEON micro-kernel below —
3574/// same entry point, so the batched-attention path opens on mobile.
3575#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
3576pub(crate) fn accel_gemm_enabled() -> bool {
3577    true
3578}
3579
3580/// Portable NEON f32 GEMM (row-major, optional Bᵀ): a 4×8 fmla
3581/// micro-kernel with A broadcast against B panels — the mobile stand-in
3582/// for Accelerate in the batched causal attention (QKᵀ and P·V). Not a
3583/// BLAS: shapes here are the attention panels (m ≤ heads·chunk,
3584/// k = head_dim or context), and the goal is removing the per-position
3585/// quadratic wall, not peak GEMM.
3586#[cfg(target_arch = "aarch64")]
3587#[allow(clippy::too_many_arguments)]
3588pub(crate) fn neon_gemm_rm(
3589    m: usize,
3590    n: usize,
3591    k: usize,
3592    alpha: f32,
3593    a: &[f32],
3594    lda: usize,
3595    b_mat: &[f32],
3596    ldb: usize,
3597    b_rows_are_n: bool,
3598    c: &mut [f32],
3599    ldc: usize,
3600) {
3601    debug_assert!(a.len() >= (m - 1) * lda + k);
3602    debug_assert!(c.len() >= (m - 1) * ldc + n);
3603    // SAFETY: bounds asserted above; NEON is baseline on aarch64.
3604    unsafe {
3605        use core::arch::aarch64::*;
3606        let mut i = 0usize;
3607        while i < m {
3608            let mi = (m - i).min(4);
3609            let mut j = 0usize;
3610            while j < n {
3611                let nj = (n - j).min(8);
3612                if mi == 4 && nj == 8 {
3613                    let (mut c0a, mut c0b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
3614                    let (mut c1a, mut c1b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
3615                    let (mut c2a, mut c2b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
3616                    let (mut c3a, mut c3b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
3617                    for p in 0..k {
3618                        let (b0, b1) = if b_rows_are_n {
3619                            // B is [n, k]: column p of Bᵀ = element p of
3620                            // eight consecutive B rows — gathered.
3621                            let base = b_mat.as_ptr().add(j * ldb + p);
3622                            let g = |o: usize| *base.add(o * ldb);
3623                            ([g(0), g(1), g(2), g(3)], [g(4), g(5), g(6), g(7)])
3624                        } else {
3625                            let base = b_mat.as_ptr().add(p * ldb + j);
3626                            (
3627                                [*base, *base.add(1), *base.add(2), *base.add(3)],
3628                                [*base.add(4), *base.add(5), *base.add(6), *base.add(7)],
3629                            )
3630                        };
3631                        let bv0 = vld1q_f32(b0.as_ptr());
3632                        let bv1 = vld1q_f32(b1.as_ptr());
3633                        let a0 = vdupq_n_f32(*a.as_ptr().add(i * lda + p));
3634                        let a1 = vdupq_n_f32(*a.as_ptr().add((i + 1) * lda + p));
3635                        let a2 = vdupq_n_f32(*a.as_ptr().add((i + 2) * lda + p));
3636                        let a3 = vdupq_n_f32(*a.as_ptr().add((i + 3) * lda + p));
3637                        c0a = vfmaq_f32(c0a, a0, bv0);
3638                        c0b = vfmaq_f32(c0b, a0, bv1);
3639                        c1a = vfmaq_f32(c1a, a1, bv0);
3640                        c1b = vfmaq_f32(c1b, a1, bv1);
3641                        c2a = vfmaq_f32(c2a, a2, bv0);
3642                        c2b = vfmaq_f32(c2b, a2, bv1);
3643                        c3a = vfmaq_f32(c3a, a3, bv0);
3644                        c3b = vfmaq_f32(c3b, a3, bv1);
3645                    }
3646                    let al = vdupq_n_f32(alpha);
3647                    for (r, (ca, cb)) in [(c0a, c0b), (c1a, c1b), (c2a, c2b), (c3a, c3b)]
3648                        .iter()
3649                        .enumerate()
3650                    {
3651                        let dst = c.as_mut_ptr().add((i + r) * ldc + j);
3652                        vst1q_f32(dst, vmulq_f32(*ca, al));
3653                        vst1q_f32(dst.add(4), vmulq_f32(*cb, al));
3654                    }
3655                } else {
3656                    for r in 0..mi {
3657                        for q in 0..nj {
3658                            let mut acc = 0f32;
3659                            for p in 0..k {
3660                                let bv = if b_rows_are_n {
3661                                    b_mat[(j + q) * ldb + p]
3662                                } else {
3663                                    b_mat[p * ldb + j + q]
3664                                };
3665                                acc += a[(i + r) * lda + p] * bv;
3666                            }
3667                            c[(i + r) * ldc + j + q] = acc * alpha;
3668                        }
3669                    }
3670                }
3671                j += nj;
3672            }
3673            i += mi;
3674        }
3675    }
3676}
3677
3678/// Off-macOS aarch64: the batched attention rides the NEON micro-GEMM.
3679#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
3680#[allow(clippy::too_many_arguments)]
3681pub(crate) fn sgemm_rm(
3682    m: usize,
3683    n: usize,
3684    k: usize,
3685    alpha: f32,
3686    a: &[f32],
3687    lda: usize,
3688    b_mat: &[f32],
3689    ldb: usize,
3690    b_rows_are_n: bool,
3691    c: &mut [f32],
3692    ldc: usize,
3693) {
3694    neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
3695}
3696
3697/// Row-major f32 GEMM, exposed for offline tools (the AWNP pass builds a
3698/// per-layer projection and applies it to every expert; a naive triple loop
3699/// would turn a two-minute job into half an hour).
3700#[allow(clippy::too_many_arguments)]
3701pub fn sgemm_public(
3702    m: usize,
3703    n: usize,
3704    k: usize,
3705    alpha: f32,
3706    a: &[f32],
3707    lda: usize,
3708    b_mat: &[f32],
3709    ldb: usize,
3710    b_rows_are_n: bool,
3711    c: &mut [f32],
3712    ldc: usize,
3713) {
3714    #[cfg(any(target_os = "macos", target_arch = "aarch64"))]
3715    {
3716        sgemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
3717    }
3718    // x86 without Accelerate has no sgemm_rm: the specialized paths there are
3719    // quantized kernels, not an f32 GEMM. Only the offline AWNP pass reaches
3720    // this, so correctness matters and throughput does not — a triple loop is
3721    // the honest fallback rather than a reason to make the tool macOS-only.
3722    #[cfg(not(any(target_os = "macos", target_arch = "aarch64")))]
3723    {
3724        for i in 0..m {
3725            for j in 0..n {
3726                let mut acc = 0f32;
3727                for p in 0..k {
3728                    let bv = if b_rows_are_n {
3729                        b_mat[j * ldb + p]
3730                    } else {
3731                        b_mat[p * ldb + j]
3732                    };
3733                    acc += a[i * lda + p] * bv;
3734                }
3735                c[i * ldc + j] = alpha * acc;
3736            }
3737        }
3738    }
3739}
3740
3741/// Row-major f32 GEMM on Accelerate: C[m,n] = alpha·A[m,k] × B(ᵀ).
3742/// `b_rows_are_n` = true multiplies by Bᵀ where B is stored [n, k].
3743#[cfg(target_os = "macos")]
3744#[allow(clippy::too_many_arguments)]
3745pub(crate) fn sgemm_rm(
3746    m: usize,
3747    n: usize,
3748    k: usize,
3749    alpha: f32,
3750    a: &[f32],
3751    lda: usize,
3752    b_mat: &[f32],
3753    ldb: usize,
3754    b_rows_are_n: bool,
3755    c: &mut [f32],
3756    ldc: usize,
3757) {
3758    debug_assert!(a.len() >= (m - 1) * lda + k);
3759    debug_assert!(c.len() >= (m - 1) * ldc + n);
3760    // Test hook: route the attention GEMMs through the portable NEON
3761    // micro-kernel ON APPLE SILICON — how the mobile batched attend is
3762    // measured without a phone in the loop. (Intel macOS has no NEON —
3763    // the hook is a no-op there, Accelerate continues below.)
3764    #[cfg(target_arch = "aarch64")]
3765    if std::env::var("CMF_FORCE_NEON_GEMM")
3766        .map(|v| v == "1")
3767        .unwrap_or(false)
3768    {
3769        return neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
3770    }
3771    unsafe {
3772        accel_blas::cblas_sgemm(
3773            101, // RowMajor
3774            111, // NoTrans A
3775            if b_rows_are_n { 112 } else { 111 },
3776            m as i32,
3777            n as i32,
3778            k as i32,
3779            alpha,
3780            a.as_ptr(),
3781            lda as i32,
3782            b_mat.as_ptr(),
3783            ldb as i32,
3784            0.0,
3785            c.as_mut_ptr(),
3786            ldc as i32,
3787        );
3788    }
3789}
3790
3791/// Prefill GEMM through Accelerate (macOS): dequantize q8 rows into
3792/// f32 tiles (scale folded in, pool-parallel) and multiply each tile
3793/// on the AMX with one row-major sgemm. Tiles live in cache, weights
3794/// stream once. Numerics are f32-GEMM (not the int8 dot): prefill
3795/// logits shift within f32 rounding — tolerance-class, like every
3796/// reduction-order change; decode (M=1) never takes this path.
3797#[cfg(target_os = "macos")]
3798fn qmatmat_accel(
3799    q: &[u8],
3800    row_scale: &[f32],
3801    pre: &[std::borrow::Cow<'_, [f32]>],
3802    rows: usize,
3803    cols: usize,
3804    out: &mut [f32],
3805    pool: Option<&Pool>,
3806) {
3807    // NOTE: double-buffering the dequant against the sgemm (a scoped
3808    // thread driving the pool on tile k+1 while the caller multiplies
3809    // tile k) was tried and LOST ~6%: Accelerate's sgemm is itself
3810    // multithreaded, and the dequant workers just steal its cores.
3811    const TR: usize = 2048;
3812    let b = pre.len();
3813    thread_local! {
3814        static XPANEL: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
3815        static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
3816    }
3817    XPANEL.with(|xp| {
3818        WTILE.with(|wt| {
3819            let mut xpanel = xp.borrow_mut();
3820            xpanel.clear();
3821            for x in pre {
3822                xpanel.extend_from_slice(x);
3823            }
3824            let mut wtile = wt.borrow_mut();
3825            wtile.resize(TR * cols, 0.0);
3826            let mut r0 = 0usize;
3827            while r0 < rows {
3828                let tr = TR.min(rows - r0);
3829                // Dequant the tile (scale folded) — pool-parallel.
3830                let wt_addr = SendMut(wtile.as_mut_ptr());
3831                let run = |start: usize, end: usize| {
3832                    for r in start..end {
3833                        let row = &q[(r0 + r) * cols..(r0 + r + 1) * cols];
3834                        let s = row_scale[r0 + r];
3835                        // SAFETY: workers cover disjoint r ranges.
3836                        let dst =
3837                            unsafe { std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols) };
3838                        for (d, &v) in dst.iter_mut().zip(row) {
3839                            *d = (v as i8) as f32 * s;
3840                        }
3841                    }
3842                };
3843                dispatch_rows(pool, tr, &run);
3844                // C[b, tr] (at column r0 of out[b, rows]) = X · Wtileᵀ
3845                unsafe {
3846                    accel_blas::cblas_sgemm(
3847                        101, // RowMajor
3848                        111, // NoTrans A
3849                        112, // Trans B
3850                        b as i32,
3851                        tr as i32,
3852                        cols as i32,
3853                        1.0,
3854                        xpanel.as_ptr(),
3855                        cols as i32,
3856                        wtile.as_ptr(),
3857                        cols as i32,
3858                        0.0,
3859                        out.as_mut_ptr().add(r0),
3860                        rows as i32,
3861                    );
3862                }
3863                r0 += tr;
3864            }
3865        })
3866    });
3867}
3868
3869fn qmatmat(
3870    q: &[u8],
3871    row_scale: &[f32],
3872    pre: &[std::borrow::Cow<'_, [f32]>],
3873    rows: usize,
3874    cols: usize,
3875    out: &mut [f32],
3876    pool: Option<&Pool>,
3877) {
3878    let b = pre.len();
3879    debug_assert_eq!(out.len(), b * rows);
3880    // Big prefill batches ride the AMX (roadmap PR3): the row×batch
3881    // SDOT loop below peaks near the CPU's dot throughput, an order
3882    // below the matrix units. Small tensors and tiny test models stay
3883    // on the exact integer path.
3884    #[cfg(target_os = "macos")]
3885    if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
3886        qmatmat_accel(q, row_scale, pre, rows, cols, out, pool);
3887        return;
3888    }
3889    #[cfg(target_arch = "aarch64")]
3890    if sdot_enabled() {
3891        let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
3892        let out_addr = SendMut(out.as_mut_ptr());
3893        // Blocked 2×4 (mobile prefill: no AMX to fall back on — this
3894        // path IS the ARM prefill GEMM off Apple silicon).
3895        let blocked_ok = blocked_enabled();
3896        let use_i8mm = i8mm_enabled();
3897        if blocked_ok {
3898            let run = |start: usize, end: usize| {
3899                let mut o = start;
3900                while o < end {
3901                    if o + 2 <= end {
3902                        let r0 = &q[o * cols..(o + 1) * cols];
3903                        let r1 = &q[(o + 1) * cols..(o + 2) * cols];
3904                        let mut bi = 0usize;
3905                        while bi + 4 <= acts.len() {
3906                            let xs = [
3907                                acts[bi].xq.as_slice(),
3908                                acts[bi + 1].xq.as_slice(),
3909                                acts[bi + 2].xq.as_slice(),
3910                                acts[bi + 3].xq.as_slice(),
3911                            ];
3912                            let d = if use_i8mm {
3913                                unsafe { dot_i8_smmla_2x4(r0, r1, xs) }
3914                            } else {
3915                                unsafe { dot_i8_sdot_2x4(r0, r1, xs) }
3916                            };
3917                            for (r, row) in [r0, r1].into_iter().enumerate() {
3918                                for k in 0..4 {
3919                                    let act = &acts[bi + k];
3920                                    let mut v = d[r][k] as f32 * act.sx;
3921                                    for &(j, xv) in &act.outliers {
3922                                        v += (row[j] as i8) as f32 * xv;
3923                                    }
3924                                    unsafe {
3925                                        *out_addr.at((bi + k) * rows + o + r) = v * row_scale[o + r]
3926                                    };
3927                                }
3928                            }
3929                            bi += 4;
3930                        }
3931                        while bi < acts.len() {
3932                            for (r, row) in [r0, r1].into_iter().enumerate() {
3933                                let v = row_dot_sdot(row, &acts[bi]) * row_scale[o + r];
3934                                unsafe { *out_addr.at(bi * rows + o + r) = v };
3935                            }
3936                            bi += 1;
3937                        }
3938                        o += 2;
3939                    } else {
3940                        let row = &q[o * cols..(o + 1) * cols];
3941                        for (bi, act) in acts.iter().enumerate() {
3942                            let v = row_dot_sdot(row, act) * row_scale[o];
3943                            unsafe { *out_addr.at(bi * rows + o) = v };
3944                        }
3945                        o += 1;
3946                    }
3947                }
3948            };
3949            dispatch_rows(pool, rows, &run);
3950            return;
3951        }
3952        let run = |start: usize, end: usize| {
3953            for o in start..end {
3954                let row = &q[o * cols..(o + 1) * cols];
3955                for (bi, act) in acts.iter().enumerate() {
3956                    let v = row_dot_sdot(row, act) * row_scale[o];
3957                    unsafe { *out_addr.at(bi * rows + o) = v };
3958                }
3959            }
3960        };
3961        dispatch_rows(pool, rows, &run);
3962        return;
3963    }
3964    // x86 A8W8 batch. Non-VNNI parts take the BLOCKED 2×4 kernel
3965    // (roadmap P0: two weight rows' abs() stay in registers across four
3966    // activation streams); VNNI machines keep the per-row bias-trick
3967    // dot, which is already throughput-bound there.
3968    #[cfg(target_arch = "x86_64")]
3969    if avx2_a8w8_enabled() {
3970        let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
3971        let out_addr = SendMut(out.as_mut_ptr());
3972        // CMF_X86_BLOCKED=0 forces the per-row path (paired in-process
3973        // A/B on noisy shared-vCPU hosts).
3974        let blocked_ok = blocked_enabled();
3975        if !avx512vnni_enabled() && blocked_ok {
3976            let run = |start: usize, end: usize| {
3977                let mut o = start;
3978                while o < end {
3979                    if o + 2 <= end {
3980                        let r0 = &q[o * cols..(o + 1) * cols];
3981                        let r1 = &q[(o + 1) * cols..(o + 2) * cols];
3982                        let mut bi = 0usize;
3983                        while bi + 4 <= acts.len() {
3984                            let xs = [
3985                                acts[bi].xq.as_slice(),
3986                                acts[bi + 1].xq.as_slice(),
3987                                acts[bi + 2].xq.as_slice(),
3988                                acts[bi + 3].xq.as_slice(),
3989                            ];
3990                            let d = unsafe { dot_i8_i8_avx2_2x4(r0, r1, xs) };
3991                            for (r, row) in [r0, r1].into_iter().enumerate() {
3992                                for k in 0..4 {
3993                                    let act = &acts[bi + k];
3994                                    let mut v = d[r][k] as f32 * act.sx;
3995                                    for &(j, xv) in &act.outliers {
3996                                        v += (row[j] as i8) as f32 * xv;
3997                                    }
3998                                    unsafe {
3999                                        *out_addr.at((bi + k) * rows + o + r) = v * row_scale[o + r]
4000                                    };
4001                                }
4002                            }
4003                            bi += 4;
4004                        }
4005                        while bi < acts.len() {
4006                            for (r, row) in [r0, r1].into_iter().enumerate() {
4007                                let v = row_dot_avx2(row, &acts[bi]) * row_scale[o + r];
4008                                unsafe { *out_addr.at(bi * rows + o + r) = v };
4009                            }
4010                            bi += 1;
4011                        }
4012                        o += 2;
4013                    } else {
4014                        let row = &q[o * cols..(o + 1) * cols];
4015                        for (bi, act) in acts.iter().enumerate() {
4016                            let v = row_dot_avx2(row, act) * row_scale[o];
4017                            unsafe { *out_addr.at(bi * rows + o) = v };
4018                        }
4019                        o += 1;
4020                    }
4021                }
4022            };
4023            dispatch_rows(pool, rows, &run);
4024            return;
4025        }
4026        let run = |start: usize, end: usize| {
4027            for o in start..end {
4028                let row = &q[o * cols..(o + 1) * cols];
4029                for (bi, act) in acts.iter().enumerate() {
4030                    let v = row_dot_avx2(row, act) * row_scale[o];
4031                    unsafe { *out_addr.at(bi * rows + o) = v };
4032                }
4033            }
4034        };
4035        dispatch_rows(pool, rows, &run);
4036        return;
4037    }
4038    let out_addr = SendMut(out.as_mut_ptr());
4039    let run = |start: usize, end: usize| {
4040        for o in start..end {
4041            let row = &q[o * cols..(o + 1) * cols];
4042            for (bi, x) in pre.iter().enumerate() {
4043                let mut acc = 0f32;
4044                for j in 0..cols {
4045                    acc += (row[j] as i8) as f32 * x[j];
4046                }
4047                unsafe { *out_addr.at(bi * rows + o) = acc * row_scale[o] };
4048            }
4049        }
4050    };
4051    dispatch_rows(pool, rows, &run);
4052}
4053
4054/// Split rows across pool workers (shared qmatvec pattern). Self-balancing
4055/// — see `Pool::run_rows` for why a static 1/n split is wrong here.
4056fn dispatch_rows(pool: Option<&Pool>, rows: usize, run: &(dyn Fn(usize, usize) + Sync)) {
4057    match pool {
4058        Some(pool) if rows >= 256 => pool.run_rows(rows, run),
4059        _ => run(0, rows),
4060    }
4061}
4062
4063/// Split a q4_block blob into (packed nibbles, f16 group scales).
4064fn q4_split(bytes: &[u8], rows: usize, cols: usize) -> (&[u8], &[u8]) {
4065    let groups = rows * cols / GROUP_SIZE;
4066    bytes.split_at(groups * 16)
4067}
4068
4069/// SIMD unpack for the dominant vbit width B=4 (94% of rows on the
4070/// log2-shape calibration): 16 packed bytes -> 32 centered i8 values.
4071/// vbit packs MSB-first, so the HIGH nibble is the even element
4072/// (opposite of q4_block's lo-first interleave). Centering is u-7.
4073#[inline]
4074fn vbit_fill4(data: &[u8], buf: &mut [u8]) {
4075    #[cfg(target_arch = "aarch64")]
4076    unsafe {
4077        return vbit_fill4_neon(data, buf);
4078    }
4079    #[cfg(target_arch = "x86_64")]
4080    if avx2_enabled() {
4081        return unsafe { vbit_fill4_avx2(data, buf) };
4082    }
4083    #[allow(unreachable_code)]
4084    for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
4085        let u = unpack8::<4>(&data[blk * 4..]);
4086        for k in 0..8 {
4087            chunk[k] = (u[k] - 7) as i8 as u8;
4088        }
4089    }
4090}
4091
4092#[cfg(target_arch = "aarch64")]
4093#[target_feature(enable = "neon")]
4094unsafe fn vbit_fill4_neon(data: &[u8], buf: &mut [u8]) {
4095    // SAFETY: buf.len() is a multiple of GROUP_SIZE=32; data holds
4096    // buf.len()/2 packed bytes (validated at load).
4097    unsafe {
4098        use core::arch::aarch64::*;
4099        let n = buf.len();
4100        let mask = vdupq_n_u8(0x0F);
4101        let seven = vdupq_n_s8(7);
4102        let mut g = 0usize;
4103        while g * 32 + 32 <= n {
4104            let b = vld1q_u8(data.as_ptr().add(g * 16));
4105            let hi = vshrq_n_u8::<4>(b);
4106            let lo = vandq_u8(b, mask);
4107            let z0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(hi, lo)), seven);
4108            let z1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(hi, lo)), seven);
4109            vst1q_u8(buf.as_mut_ptr().add(g * 32), vreinterpretq_u8_s8(z0));
4110            vst1q_u8(buf.as_mut_ptr().add(g * 32 + 16), vreinterpretq_u8_s8(z1));
4111            g += 1;
4112        }
4113    }
4114}
4115
4116#[cfg(target_arch = "x86_64")]
4117#[target_feature(enable = "avx2")]
4118unsafe fn vbit_fill4_avx2(data: &[u8], buf: &mut [u8]) {
4119    // SAFETY: see vbit_fill4_neon.
4120    unsafe {
4121        use core::arch::x86_64::*;
4122        let n = buf.len();
4123        let mask = _mm_set1_epi8(0x0F);
4124        let seven = _mm256_set1_epi8(7);
4125        let mut g = 0usize;
4126        while g * 32 + 32 <= n {
4127            let b = _mm_loadu_si128(data.as_ptr().add(g * 16) as *const __m128i);
4128            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), mask);
4129            let lo = _mm_and_si128(b, mask);
4130            let z = _mm256_sub_epi8(
4131                _mm256_set_m128i(_mm_unpackhi_epi8(hi, lo), _mm_unpacklo_epi8(hi, lo)),
4132                seven,
4133            );
4134            _mm256_storeu_si256(buf.as_mut_ptr().add(g * 32) as *mut __m256i, z);
4135            g += 1;
4136        }
4137    }
4138}
4139
4140/// Unpack 8 MSB-first B-bit values from exactly B bytes (fixed shifts —
4141/// no serial bit-buffer, auto-vectorizable). Every 32-value group starts
4142/// byte-aligned (32·B/8 is integral for B∈3..8), so groups decompose
4143/// into 4 such blocks.
4144#[inline(always)]
4145fn unpack8<const B: usize>(data: &[u8]) -> [i32; 8] {
4146    let mut acc = 0u64;
4147    for i in 0..B {
4148        acc = (acc << 8) | data[i] as u64;
4149    }
4150    let mask = (1u64 << B) - 1;
4151    let mut out = [0i32; 8];
4152    for (k, o) in out.iter_mut().enumerate() {
4153        *o = ((acc >> ((7 - k) * B)) & mask) as i32;
4154    }
4155    out
4156}
4157
4158/// Fused vbit matvec straight from the mapped bytes (spec §3, P13
4159/// FIG.3): [u8 bits: rows][f16 scales: rows·cols/32][bit-packed rows,
4160/// MSB-first, byte-padded]. Row data offsets are precomputed at load
4161/// (`vbit_row_offsets`) — the per-call prefix scan was O(rows) pure
4162/// overhead on every matvec.
4163#[allow(clippy::too_many_arguments)]
4164fn vbitmatvec(
4165    bytes: &[u8],
4166    offsets: &[usize],
4167    x: &[f32],
4168    rows: usize,
4169    cols: usize,
4170    out: &mut [f32],
4171    pool: Option<&Pool>,
4172) {
4173    debug_assert_eq!(out.len(), rows);
4174    debug_assert_eq!(offsets.len(), rows + 1);
4175
4176    // SDOT path: unpack the row to centered i8 once, then per-group
4177    // int8 dot against the quantized activations — same A8W8 contract
4178    // as q8 (bounded noise; CMF_SDOT=0 keeps the exact scalar path).
4179    if a8w8_enabled() {
4180        let act = split_act(x);
4181        let out_addr = SendMut(out.as_mut_ptr());
4182        let run = move |start: usize, end: usize| {
4183            vbit_range_a8w8(bytes, offsets, x, &act, rows, cols, out_addr, start, end)
4184        };
4185        dispatch_rows(pool, rows, &run);
4186        return;
4187    }
4188
4189    let out_addr = SendMut(out.as_mut_ptr());
4190    let run = move |start: usize, end: usize| {
4191        vbit_range_f32(bytes, offsets, x, rows, cols, out_addr, start, end)
4192    };
4193    dispatch_rows(pool, rows, &run);
4194}
4195
4196/// One vbit row range via the A8W8 int8 path — kernel body of
4197/// `vbitmatvec`, extracted so multi-matrix jobs can drive it for
4198/// several tensors in one dispatch (b=8 rows go exact f32).
4199#[allow(clippy::too_many_arguments)]
4200fn vbit_range_a8w8(
4201    bytes: &[u8],
4202    offsets: &[usize],
4203    x: &[f32],
4204    act: &SplitAct,
4205    rows: usize,
4206    cols: usize,
4207    out: SendMut,
4208    start: usize,
4209    end: usize,
4210) {
4211    let ng = cols / GROUP_SIZE;
4212    let bits = &bytes[..rows];
4213    let sc_off = rows;
4214    let row_dot = |r: usize| -> f32 {
4215        let b = bits[r] as usize;
4216        let l = (1i32 << (b - 1)) - 1;
4217        let mask = (1u64 << b) - 1;
4218        let data = &bytes[offsets[r]..offsets[r + 1]];
4219        if b == 8 {
4220            // u−L reaches 128 → does not fit i8; exact f32 path.
4221            let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
4222            let mut dot = 0f32;
4223            for g in 0..ng {
4224                let so = (r * ng + g) * 2;
4225                let sgf = f16_to_f32(u16::from_le_bytes([
4226                    bytes[sc_off + so],
4227                    bytes[sc_off + so + 1],
4228                ]));
4229                let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
4230                let mut gd = 0f32;
4231                for &xv in xg.iter() {
4232                    if nbits < 8 {
4233                        acc = (acc << 8) | data[idx] as u64;
4234                        idx += 1;
4235                        nbits += 8;
4236                    }
4237                    let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
4238                    nbits -= 8;
4239                    gd += (u - l) as f32 * xv;
4240                }
4241                dot += gd * sgf;
4242            }
4243            return dot;
4244        }
4245        // Per-worker scratch: this closure runs for every row of the
4246        // tensor (lm_head ≈ 150k rows/token) — a heap allocation per
4247        // row was measurable pure overhead.
4248        thread_local! {
4249            static VBIT_SCRATCH: std::cell::RefCell<Vec<u8>> =
4250                const { std::cell::RefCell::new(Vec::new()) };
4251        }
4252        #[inline(always)]
4253        fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
4254            for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
4255                let u = unpack8::<B>(&data[blk * B..]);
4256                for k in 0..8 {
4257                    chunk[k] = (u[k] - l) as i8 as u8;
4258                }
4259            }
4260        }
4261        let _ = mask;
4262        VBIT_SCRATCH.with(|scratch| {
4263            let mut buf = scratch.borrow_mut();
4264            buf.resize(cols, 0);
4265            match b {
4266                3 => fill::<3>(data, l, &mut buf),
4267                4 => vbit_fill4(data, &mut buf),
4268                5 => fill::<5>(data, l, &mut buf),
4269                6 => fill::<6>(data, l, &mut buf),
4270                _ => unreachable!(),
4271            }
4272            let mut dot = 0f32;
4273            for g in 0..ng {
4274                let so = (r * ng + g) * 2;
4275                let s = f16_to_f32(u16::from_le_bytes([
4276                    bytes[sc_off + so],
4277                    bytes[sc_off + so + 1],
4278                ]));
4279                let d = dot_i8_i8(
4280                    &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
4281                    &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
4282                ) as f32
4283                    * act.sx;
4284                dot += d * s;
4285            }
4286            for &(j, xv) in &act.outliers {
4287                let so = (r * ng + j / GROUP_SIZE) * 2;
4288                let s = f16_to_f32(u16::from_le_bytes([
4289                    bytes[sc_off + so],
4290                    bytes[sc_off + so + 1],
4291                ]));
4292                // xq is zeroed at outlier slots — add the exact term.
4293                dot += (buf[j] as i8) as f32 * s * xv;
4294            }
4295            dot
4296        })
4297    };
4298    for r in start..end {
4299        // SAFETY: disjoint row ranges per worker.
4300        unsafe { *out.at(r) = row_dot(r) };
4301    }
4302}
4303
4304/// Exact scalar vbit row range (same extraction, non-SDOT path).
4305#[allow(clippy::too_many_arguments)]
4306fn vbit_range_f32(
4307    bytes: &[u8],
4308    offsets: &[usize],
4309    x: &[f32],
4310    rows: usize,
4311    cols: usize,
4312    out: SendMut,
4313    start: usize,
4314    end: usize,
4315) {
4316    let ng = cols / GROUP_SIZE;
4317    let bits = &bytes[..rows];
4318    let sc_off = rows;
4319    // Per-bit-width specialized inner loops: the compiler unrolls the
4320    // constant shifts (the generic bit-buffer loop was branch-bound —
4321    // 5.6 vs 13.2 tok/s q4 on the 0.8B).
4322    #[inline(always)]
4323    fn dot_row<const B: usize>(
4324        data: &[u8],
4325        bytes: &[u8],
4326        sc_off: usize,
4327        r: usize,
4328        ng: usize,
4329        x: &[f32],
4330    ) -> f32 {
4331        let l = ((1i32 << (B - 1)) - 1) as f32;
4332        let gbytes = GROUP_SIZE * B / 8;
4333        let mut dot = 0f32;
4334        for g in 0..ng {
4335            let so = (r * ng + g) * 2;
4336            let s = f16_to_f32(u16::from_le_bytes([
4337                bytes[sc_off + so],
4338                bytes[sc_off + so + 1],
4339            ]));
4340            let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
4341            let gd0 = &data[g * gbytes..(g + 1) * gbytes];
4342            let mut gd = 0f32;
4343            for blk in 0..GROUP_SIZE / 8 {
4344                let u = unpack8::<B>(&gd0[blk * B..]);
4345                let xb = &xg[blk * 8..blk * 8 + 8];
4346                for k in 0..8 {
4347                    gd += (u[k] as f32 - l) * xb[k];
4348                }
4349            }
4350            dot += gd * s;
4351        }
4352        dot
4353    }
4354    for r in start..end {
4355        let data = &bytes[offsets[r]..offsets[r + 1]];
4356        let v = match bits[r] {
4357            3 => dot_row::<3>(data, bytes, sc_off, r, ng, x),
4358            4 => dot_row::<4>(data, bytes, sc_off, r, ng, x),
4359            5 => dot_row::<5>(data, bytes, sc_off, r, ng, x),
4360            6 => dot_row::<6>(data, bytes, sc_off, r, ng, x),
4361            8 => dot_row::<8>(data, bytes, sc_off, r, ng, x),
4362            b => unreachable!("vbit bit-width {b} (validated at load)"),
4363        };
4364        // SAFETY: disjoint row ranges per worker.
4365        unsafe { *out.at(r) = v };
4366    }
4367}
4368
4369/// Fused two-input vbit matvec: each row is unpacked from the mmap ONCE
4370/// and dotted against BOTH activations (MTP verify / pair prefill used
4371/// to run two full matvecs — double weight traffic and double unpack).
4372/// Per-input math is identical to `vbitmatvec` → same accuracy contract.
4373#[allow(clippy::too_many_arguments)]
4374fn vbitmatvec2(
4375    bytes: &[u8],
4376    offsets: &[usize],
4377    x1: &[f32],
4378    x2: &[f32],
4379    rows: usize,
4380    cols: usize,
4381    o1: &mut [f32],
4382    o2: &mut [f32],
4383    pool: Option<&Pool>,
4384) {
4385    debug_assert_eq!(o1.len(), rows);
4386    debug_assert_eq!(o2.len(), rows);
4387
4388    if a8w8_enabled() {
4389        let a1 = split_act(x1);
4390        let a2 = split_act(x2);
4391        let p1 = SendMut(o1.as_mut_ptr());
4392        let p2 = SendMut(o2.as_mut_ptr());
4393        let run = move |start: usize, end: usize| {
4394            vbit_range2_a8w8(
4395                bytes, offsets, x1, x2, &a1, &a2, rows, cols, p1, p2, start, end,
4396            )
4397        };
4398        dispatch_rows(pool, rows, &run);
4399        return;
4400    }
4401
4402    let p1 = SendMut(o1.as_mut_ptr());
4403    let p2 = SendMut(o2.as_mut_ptr());
4404    let run = move |start: usize, end: usize| {
4405        vbit_range2_f32(bytes, offsets, x1, x2, rows, cols, p1, p2, start, end)
4406    };
4407    dispatch_rows(pool, rows, &run);
4408}
4409
4410/// Two-input vbit row range via the A8W8 int8 path — kernel body of
4411/// `vbitmatvec2`, extracted for pair multi-matrix jobs (b=8 rows go
4412/// exact f32 for both lanes, bits streamed once).
4413#[allow(clippy::too_many_arguments)]
4414fn vbit_range2_a8w8(
4415    bytes: &[u8],
4416    offsets: &[usize],
4417    x1: &[f32],
4418    x2: &[f32],
4419    a1: &SplitAct,
4420    a2: &SplitAct,
4421    rows: usize,
4422    cols: usize,
4423    p1: SendMut,
4424    p2: SendMut,
4425    start: usize,
4426    end: usize,
4427) {
4428    let ng = cols / GROUP_SIZE;
4429    let bits = &bytes[..rows];
4430    let sc_off = rows;
4431    let row_dots = |r: usize| -> (f32, f32) {
4432        let b = bits[r] as usize;
4433        let l = (1i32 << (b - 1)) - 1;
4434        let data = &bytes[offsets[r]..offsets[r + 1]];
4435        if b == 8 {
4436            // u−L reaches 128 → does not fit i8; exact f32 path,
4437            // bits still streamed once for both lanes.
4438            let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
4439            let (mut d1, mut d2) = (0f32, 0f32);
4440            for g in 0..ng {
4441                let so = (r * ng + g) * 2;
4442                let sgf = f16_to_f32(u16::from_le_bytes([
4443                    bytes[sc_off + so],
4444                    bytes[sc_off + so + 1],
4445                ]));
4446                let (mut g1, mut g2) = (0f32, 0f32);
4447                for k in 0..GROUP_SIZE {
4448                    if nbits < 8 {
4449                        acc = (acc << 8) | data[idx] as u64;
4450                        idx += 1;
4451                        nbits += 8;
4452                    }
4453                    let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
4454                    nbits -= 8;
4455                    let w = (u - l) as f32;
4456                    g1 += w * x1[g * GROUP_SIZE + k];
4457                    g2 += w * x2[g * GROUP_SIZE + k];
4458                }
4459                d1 += g1 * sgf;
4460                d2 += g2 * sgf;
4461            }
4462            return (d1, d2);
4463        }
4464        thread_local! {
4465            static VBIT_SCRATCH2: std::cell::RefCell<Vec<u8>> =
4466                const { std::cell::RefCell::new(Vec::new()) };
4467        }
4468        #[inline(always)]
4469        fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
4470            for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
4471                let u = unpack8::<B>(&data[blk * B..]);
4472                for k in 0..8 {
4473                    chunk[k] = (u[k] - l) as i8 as u8;
4474                }
4475            }
4476        }
4477        VBIT_SCRATCH2.with(|scratch| {
4478            let mut buf = scratch.borrow_mut();
4479            buf.resize(cols, 0);
4480            match b {
4481                3 => fill::<3>(data, l, &mut buf),
4482                4 => vbit_fill4(data, &mut buf),
4483                5 => fill::<5>(data, l, &mut buf),
4484                6 => fill::<6>(data, l, &mut buf),
4485                _ => unreachable!(),
4486            }
4487            let (mut d1, mut d2) = (0f32, 0f32);
4488            for g in 0..ng {
4489                let so = (r * ng + g) * 2;
4490                let s = f16_to_f32(u16::from_le_bytes([
4491                    bytes[sc_off + so],
4492                    bytes[sc_off + so + 1],
4493                ]));
4494                let wg = &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
4495                let v1 = dot_i8_i8(wg, &a1.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a1.sx;
4496                let v2 = dot_i8_i8(wg, &a2.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a2.sx;
4497                d1 += v1 * s;
4498                d2 += v2 * s;
4499            }
4500            for &(j, xv) in &a1.outliers {
4501                let so = (r * ng + j / GROUP_SIZE) * 2;
4502                let s = f16_to_f32(u16::from_le_bytes([
4503                    bytes[sc_off + so],
4504                    bytes[sc_off + so + 1],
4505                ]));
4506                d1 += (buf[j] as i8) as f32 * s * xv;
4507            }
4508            for &(j, xv) in &a2.outliers {
4509                let so = (r * ng + j / GROUP_SIZE) * 2;
4510                let s = f16_to_f32(u16::from_le_bytes([
4511                    bytes[sc_off + so],
4512                    bytes[sc_off + so + 1],
4513                ]));
4514                d2 += (buf[j] as i8) as f32 * s * xv;
4515            }
4516            (d1, d2)
4517        })
4518    };
4519    for r in start..end {
4520        let (v1, v2) = row_dots(r);
4521        // SAFETY: disjoint row ranges per worker.
4522        unsafe {
4523            *p1.at(r) = v1;
4524            *p2.at(r) = v2;
4525        }
4526    }
4527}
4528
4529/// Two-input exact scalar vbit row range (same extraction) —
4530/// per-bit-width specialized, two accumulators per row; per-lane
4531/// accumulation order matches `vbitmatvec` exactly.
4532#[allow(clippy::too_many_arguments)]
4533fn vbit_range2_f32(
4534    bytes: &[u8],
4535    offsets: &[usize],
4536    x1: &[f32],
4537    x2: &[f32],
4538    rows: usize,
4539    cols: usize,
4540    p1: SendMut,
4541    p2: SendMut,
4542    start: usize,
4543    end: usize,
4544) {
4545    let ng = cols / GROUP_SIZE;
4546    let bits = &bytes[..rows];
4547    let sc_off = rows;
4548    #[inline(always)]
4549    #[allow(clippy::too_many_arguments)]
4550    fn dot_row2<const B: usize>(
4551        data: &[u8],
4552        bytes: &[u8],
4553        sc_off: usize,
4554        r: usize,
4555        ng: usize,
4556        x1: &[f32],
4557        x2: &[f32],
4558    ) -> (f32, f32) {
4559        let l = ((1i32 << (B - 1)) - 1) as f32;
4560        let gbytes = GROUP_SIZE * B / 8;
4561        let (mut d1, mut d2) = (0f32, 0f32);
4562        for g in 0..ng {
4563            let so = (r * ng + g) * 2;
4564            let s = f16_to_f32(u16::from_le_bytes([
4565                bytes[sc_off + so],
4566                bytes[sc_off + so + 1],
4567            ]));
4568            let x1g = &x1[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
4569            let x2g = &x2[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
4570            let gd0 = &data[g * gbytes..(g + 1) * gbytes];
4571            let (mut g1, mut g2) = (0f32, 0f32);
4572            for blk in 0..GROUP_SIZE / 8 {
4573                let u = unpack8::<B>(&gd0[blk * B..]);
4574                for k in 0..8 {
4575                    let w = u[k] as f32 - l;
4576                    g1 += w * x1g[blk * 8 + k];
4577                    g2 += w * x2g[blk * 8 + k];
4578                }
4579            }
4580            d1 += g1 * s;
4581            d2 += g2 * s;
4582        }
4583        (d1, d2)
4584    }
4585    for r in start..end {
4586        let data = &bytes[offsets[r]..offsets[r + 1]];
4587        let (v1, v2) = match bits[r] {
4588            3 => dot_row2::<3>(data, bytes, sc_off, r, ng, x1, x2),
4589            4 => dot_row2::<4>(data, bytes, sc_off, r, ng, x1, x2),
4590            5 => dot_row2::<5>(data, bytes, sc_off, r, ng, x1, x2),
4591            6 => dot_row2::<6>(data, bytes, sc_off, r, ng, x1, x2),
4592            8 => dot_row2::<8>(data, bytes, sc_off, r, ng, x1, x2),
4593            b => unreachable!("vbit bit-width {b} (validated at load)"),
4594        };
4595        // SAFETY: disjoint row ranges per worker.
4596        unsafe {
4597            *p1.at(r) = v1;
4598            *p2.at(r) = v2;
4599        }
4600    }
4601}
4602
4603// ───────────────────── q4_tiled kernels (§4.3) ─────────────────────
4604
4605/// One q4_tiled row dot on the A8W8 int8 path: per 32-group the tile
4606/// is ONE sequential read — [f16 scale][16B nibbles] — versus the two
4607/// distant streams of the split layout. Values/order identical to the
4608/// split kernels.
4609#[inline]
4610#[allow(unreachable_code)]
4611fn dot_q4t_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
4612    #[cfg(target_arch = "aarch64")]
4613    unsafe {
4614        return dot_q4t_row_sdot(bytes, r, gpr, xq);
4615    }
4616    #[cfg(target_arch = "x86_64")]
4617    unsafe {
4618        if vnni_tiles_enabled() {
4619            return dot_q4t_row_vnni(bytes, r, gpr, xq);
4620        }
4621        return dot_q4t_row_avx2(bytes, r, gpr, xq);
4622    }
4623    let mut acc = 0f32;
4624    for gi in 0..gpr {
4625        let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
4626        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
4627        let mut d = 0i32;
4628        for (k, &b) in tile[2..].iter().enumerate() {
4629            d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
4630                + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
4631        }
4632        acc += d as f32 * s;
4633    }
4634    acc
4635}
4636
4637#[cfg(target_arch = "aarch64")]
4638#[target_feature(enable = "neon,dotprod")]
4639unsafe fn dot_q4t_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
4640    // SAFETY: callers uphold slice-length contracts (18B tile per group,
4641    // xq.len() == gpr·GROUP_SIZE).
4642    unsafe {
4643        use core::arch::aarch64::*;
4644        use core::arch::asm;
4645        let lomask = vdupq_n_u8(0x0F);
4646        let eight = vdupq_n_s8(8);
4647        let mut acc = 0f32;
4648        for gi in 0..gpr {
4649            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4650            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4651            let b = vld1q_u8(t.add(2));
4652            let lo = vandq_u8(b, lomask);
4653            let hi = vshrq_n_u8::<4>(b);
4654            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
4655            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
4656            let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
4657            let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
4658            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
4659            asm!(
4660                "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
4661                "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
4662                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4663                e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
4664                options(pure, nomem, nostack),
4665            );
4666            acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
4667        }
4668        acc
4669    }
4670}
4671
4672#[cfg(target_arch = "x86_64")]
4673#[target_feature(enable = "avx2")]
4674unsafe fn dot_q4t_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
4675    // SAFETY: see dot_q4t_row_sdot.
4676    unsafe {
4677        use core::arch::x86_64::*;
4678        let lomask = _mm_set1_epi8(0x0F);
4679        let eight = _mm256_set1_epi8(8);
4680        let ones = _mm256_set1_epi16(1);
4681        let mut acc = 0f32;
4682        for gi in 0..gpr {
4683            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4684            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4685            let b = _mm_loadu_si128(t.add(2) as *const __m128i);
4686            let lo = _mm_and_si128(b, lomask);
4687            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
4688            let w = _mm256_sub_epi8(
4689                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4690                eight,
4691            );
4692            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4693            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
4694            let d = _mm256_madd_epi16(p16, ones);
4695            let hi128 = _mm256_extracti128_si256::<1>(d);
4696            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
4697            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
4698            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
4699            acc += _mm_cvtsi128_si32(s32) as f32 * s;
4700        }
4701        acc
4702    }
4703}
4704
4705/// VNNI twin of `dot_q4t_row_avx2`: same unpack, `vpdpbusd` replaces
4706/// the maddubs+madd pair (see `dpbusd_hsum` — sums are bit-identical).
4707/// 256-bit VL encoding, so the VEX `vpsignb` stays usable.
4708#[cfg(target_arch = "x86_64")]
4709#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
4710unsafe fn dot_q4t_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
4711    // SAFETY: see dot_q4t_row_sdot.
4712    unsafe {
4713        use core::arch::x86_64::*;
4714        let lomask = _mm_set1_epi8(0x0F);
4715        let eight = _mm256_set1_epi8(8);
4716        let mut acc = 0f32;
4717        for gi in 0..gpr {
4718            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4719            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4720            let b = _mm_loadu_si128(t.add(2) as *const __m128i);
4721            let lo = _mm_and_si128(b, lomask);
4722            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
4723            let w = _mm256_sub_epi8(
4724                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4725                eight,
4726            );
4727            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4728            let d = dpbusd_hsum(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
4729            acc += d as f32 * s;
4730        }
4731        acc
4732    }
4733}
4734
4735/// One q4_tiled row against FOUR activation streams: the nibble unpack
4736/// and abs() happen once per group instead of once per (group,
4737/// activation) — the unpack is the dominant per-element cost of the
4738/// tiled format (roadmap P0 portable blocking, q4t leg).
4739#[cfg(target_arch = "x86_64")]
4740// `fma` is NOT implied by `avx2`: without it LLVM lowers _mm256_fmadd_ps
4741// to a libm call per lane — measured 2x slower than the reduction this
4742// kernel replaces. The runtime gate (`avx2_enabled`) already requires
4743// both features, so declaring it here is safe.
4744#[target_feature(enable = "avx2,fma")]
4745unsafe fn dot_q4t_row_1x4_avx2(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
4746    // SAFETY: callers uphold the 18B-tile and xq-length contracts.
4747    unsafe {
4748        use core::arch::x86_64::*;
4749        let lomask = _mm_set1_epi8(0x0F);
4750        let eight = _mm256_set1_epi8(8);
4751        let ones = _mm256_set1_epi16(1);
4752        // One f32 accumulator VECTOR per activation, reduced once at the
4753        // end. Folding each group's i32 lanes to a scalar inside the loop
4754        // costs an extracti128 + three shift/add + a movd — a cross-lane
4755        // dependency chain per (group, activation), 288 of them per row at
4756        // cols=2304. The per-group scale is what forces a float
4757        // accumulator; it does not force a horizontal sum.
4758        //
4759        // The four accumulators are NAMED, not an array: as `[__m256; 4]`
4760        // indexed by a loop variable LLVM keeps them in memory and every
4761        // group pays four 32-byte loads and stores. That alone made this
4762        // kernel 2x SLOWER than the per-group reduction it replaces
4763        // (measured on the EPYC box: 150 s vs 71 s for two 256² steps).
4764        let mut f0 = _mm256_setzero_ps();
4765        let mut f1 = _mm256_setzero_ps();
4766        let mut f2 = _mm256_setzero_ps();
4767        let mut f3 = _mm256_setzero_ps();
4768        for gi in 0..gpr {
4769            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4770            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4771            let sv = _mm256_set1_ps(s);
4772            let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
4773            let lo = _mm_and_si128(bb, lomask);
4774            let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
4775            let w = _mm256_sub_epi8(
4776                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4777                eight,
4778            );
4779            let aw = _mm256_abs_epi8(w);
4780            let off = gi * GROUP_SIZE;
4781            let dot = |xq: &[i8]| {
4782                let x = _mm256_loadu_si256(xq.as_ptr().add(off) as *const __m256i);
4783                let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
4784                _mm256_cvtepi32_ps(_mm256_madd_epi16(p16, ones))
4785            };
4786            f0 = _mm256_fmadd_ps(dot(xs[0]), sv, f0);
4787            f1 = _mm256_fmadd_ps(dot(xs[1]), sv, f1);
4788            f2 = _mm256_fmadd_ps(dot(xs[2]), sv, f2);
4789            f3 = _mm256_fmadd_ps(dot(xs[3]), sv, f3);
4790        }
4791        [
4792            hsum256_ps(f0),
4793            hsum256_ps(f1),
4794            hsum256_ps(f2),
4795            hsum256_ps(f3),
4796        ]
4797    }
4798}
4799
4800/// Horizontal sum of eight f32 lanes — the one cross-lane reduction the
4801/// blocked kernels pay, once per row instead of once per group.
4802#[cfg(target_arch = "x86_64")]
4803#[target_feature(enable = "avx2")]
4804#[inline]
4805unsafe fn hsum256_ps(v: core::arch::x86_64::__m256) -> f32 {
4806    // SAFETY: pure register arithmetic on the caller's vector.
4807    unsafe {
4808        use core::arch::x86_64::*;
4809        let hi = _mm256_extractf128_ps::<1>(v);
4810        let s = _mm_add_ps(_mm256_castps256_ps128(v), hi);
4811        let s = _mm_add_ps(s, _mm_movehl_ps(s, s));
4812        let s = _mm_add_ss(s, _mm_shuffle_ps::<0x55>(s, s));
4813        _mm_cvtss_f32(s)
4814    }
4815}
4816
4817/// VNNI twin of `dot_q4t_row_1x4_avx2` (see `dpbusd_hsum`).
4818#[cfg(target_arch = "x86_64")]
4819#[target_feature(enable = "avx2,fma,avx512f,avx512bw,avx512vl,avx512vnni")]
4820unsafe fn dot_q4t_row_1x4_vnni(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
4821    // SAFETY: callers uphold the 18B-tile and xq-length contracts.
4822    unsafe {
4823        use core::arch::x86_64::*;
4824        let lomask = _mm_set1_epi8(0x0F);
4825        let eight = _mm256_set1_epi8(8);
4826        // Same shape as the AVX2 twin: accumulate in f32 vectors and pay
4827        // one cross-lane reduction per row, not per (group, activation).
4828        let mut f0 = _mm256_setzero_ps();
4829        let mut f1 = _mm256_setzero_ps();
4830        let mut f2 = _mm256_setzero_ps();
4831        let mut f3 = _mm256_setzero_ps();
4832        for gi in 0..gpr {
4833            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4834            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4835            let sv = _mm256_set1_ps(s);
4836            let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
4837            let lo = _mm_and_si128(bb, lomask);
4838            let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
4839            let w = _mm256_sub_epi8(
4840                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4841                eight,
4842            );
4843            let aw = _mm256_abs_epi8(w);
4844            let off = gi * GROUP_SIZE;
4845            let dot = |xq: &[i8]| {
4846                let x = _mm256_loadu_si256(xq.as_ptr().add(off) as *const __m256i);
4847                _mm256_cvtepi32_ps(_mm256_dpbusd_epi32(
4848                    _mm256_setzero_si256(),
4849                    aw,
4850                    _mm256_sign_epi8(x, w),
4851                ))
4852            };
4853            f0 = _mm256_fmadd_ps(dot(xs[0]), sv, f0);
4854            f1 = _mm256_fmadd_ps(dot(xs[1]), sv, f1);
4855            f2 = _mm256_fmadd_ps(dot(xs[2]), sv, f2);
4856            f3 = _mm256_fmadd_ps(dot(xs[3]), sv, f3);
4857        }
4858        let acc = [
4859            hsum256_ps(f0),
4860            hsum256_ps(f1),
4861            hsum256_ps(f2),
4862            hsum256_ps(f3),
4863        ];
4864        acc
4865    }
4866}
4867
4868/// ARM twin of `dot_q4t_row_1x4_avx2`: one nibble unpack per group
4869/// serves FOUR activation streams. Per stream the group order and f32
4870/// accumulation match `dot_q4t_row_sdot` exactly — batch == matvec
4871/// bit-for-bit.
4872#[cfg(target_arch = "aarch64")]
4873#[target_feature(enable = "neon,dotprod")]
4874unsafe fn dot_q4t_row_1x4_sdot(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
4875    // SAFETY: callers uphold the 18B-tile and xq-length contracts.
4876    unsafe {
4877        use core::arch::aarch64::*;
4878        use core::arch::asm;
4879        let lomask = vdupq_n_u8(0x0F);
4880        let eight = vdupq_n_s8(8);
4881        let mut acc = [0f32; 4];
4882        for gi in 0..gpr {
4883            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4884            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4885            let b = vld1q_u8(t.add(2));
4886            let lo = vandq_u8(b, lomask);
4887            let hi = vshrq_n_u8::<4>(b);
4888            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
4889            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
4890            for (k, xq) in xs.iter().enumerate() {
4891                let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
4892                let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
4893                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
4894                asm!(
4895                    "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
4896                    "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
4897                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4898                    e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
4899                    options(pure, nomem, nostack),
4900                );
4901                acc[k] += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
4902            }
4903        }
4904        acc
4905    }
4906}
4907
4908/// Exact-term correction for A8W8 outliers on a tiled row.
4909#[inline]
4910fn q4t_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
4911    let gi = j / GROUP_SIZE;
4912    let k = j % GROUP_SIZE;
4913    let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
4914    let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
4915    let byte = tile[2 + k / 2];
4916    let nib = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
4917    ((nib as i32 - 8) as f32, s)
4918}
4919
4920/// Exact scalar q4_tiled row (CMF_SDOT=0 contract) — same pairwise
4921/// accumulation shape as `q4_range_f32`.
4922#[inline]
4923fn q4t_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
4924    let mut acc = 0f32;
4925    for gi in 0..gpr {
4926        let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
4927        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
4928        let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
4929        let mut ga = 0f32;
4930        for (k, &b) in tile[2..].iter().enumerate() {
4931            ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
4932                + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
4933        }
4934        acc += ga * s;
4935    }
4936    acc
4937}
4938
4939/// Split view of a `q4tp` payload. The three planes are resolved once per
4940/// matvec instead of per row — `q4tp_sections` is cheap, but doing it inside
4941/// the row loop would put a division on the hot path for nothing.
4942struct Q4tpView<'a> {
4943    nib: &'a [u8],
4944    params: &'a [u8],
4945    codes: &'a [u8],
4946    stride: usize,
4947    /// q2tp reads the ladder with rung 0 = exact zero.
4948    zero_rung: bool,
4949}
4950
4951impl<'a> Q4tpView<'a> {
4952    fn new(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
4953        let (params_off, codes_off, stride) = q4tp_sections(rows, cols);
4954        Self {
4955            nib: &bytes[..params_off],
4956            params: &bytes[params_off..codes_off],
4957            codes: &bytes[codes_off..],
4958            stride,
4959            zero_rung: false,
4960        }
4961    }
4962
4963    /// The q2tp view: identical params/codes planes, 8 B weight chunks.
4964    fn new_q2(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
4965        let (params_off, codes_off, stride) = q2tp_sections(rows, cols);
4966        Self {
4967            nib: &bytes[..params_off],
4968            params: &bytes[params_off..codes_off],
4969            codes: &bytes[codes_off..],
4970            stride,
4971            zero_rung: true,
4972        }
4973    }
4974
4975    /// Expand row `r`'s per-tile scales into `out` (length `gpr`).
4976    ///
4977    /// Doing this once per row — rather than decoding a 5-bit code inside the
4978    /// tile loop — is what makes the format free at runtime. Random access to
4979    /// a packed 5-bit field costs a division, two bounds checks and a branch;
4980    /// the tile's actual work is two `sdot`s, so per-tile decoding dominated
4981    /// the kernel and cost 5x (measured: 1.4 vs 6.9 tok/s on Nanbeige-3B).
4982    /// Walking the plane sequentially with a bit accumulator is ~3 ops.
4983    /// Eight 5-bit codes are exactly five bytes, so a whole group of
4984    /// eight decodes from one little-endian word at fixed shifts. The
4985    /// bit-accumulator this replaces carried a data-dependent `while
4986    /// have < 5` refill whose branch sat in the innermost loop of every
4987    /// q4tp row; a decode profile put this function above the dot
4988    /// products it feeds. Same bitstream, same codes — just no branch
4989    /// and eight independent extractions.
4990    #[inline]
4991    fn scales_into(&self, r: usize, gpr: usize, out: &mut [f32]) {
4992        let tab = if self.zero_rung {
4993            q2tp_ladder(self.params, r)
4994        } else {
4995            q4tp_ladder(self.params, r)
4996        };
4997        let codes = &self.codes[r * self.stride..(r + 1) * self.stride];
4998        let out = &mut out[..gpr];
4999        let mut chunks = out.chunks_exact_mut(8);
5000        let mut ci = 0usize;
5001        for c in &mut chunks {
5002            let w = u64::from(codes[ci])
5003                | u64::from(codes[ci + 1]) << 8
5004                | u64::from(codes[ci + 2]) << 16
5005                | u64::from(codes[ci + 3]) << 24
5006                | u64::from(codes[ci + 4]) << 32;
5007            for (k, o) in c.iter_mut().enumerate() {
5008                *o = tab[((w >> (5 * k)) & 31) as usize];
5009            }
5010            ci += 5;
5011        }
5012        // Fewer than eight codes left: the shared total accessor, which
5013        // tolerates a 5-bit field whose spill byte is past the stride.
5014        let tail = &codes[ci..];
5015        for (k, o) in chunks.into_remainder().iter_mut().enumerate() {
5016            *o = tab[q4tp_code(tail, k)];
5017        }
5018    }
5019}
5020
5021#[inline]
5022fn dot_q4tp_row_i8(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
5023    #[cfg(target_arch = "aarch64")]
5024    unsafe {
5025        return dot_q4tp_row_sdot(nib, r, gpr, xq, scales);
5026    }
5027    #[cfg(target_arch = "x86_64")]
5028    unsafe {
5029        if vnni_tiles_enabled() {
5030            return dot_q4tp_row_vnni(nib, r, gpr, xq, scales);
5031        }
5032        return dot_q4tp_row_avx2(nib, r, gpr, xq, scales);
5033    }
5034    #[allow(unreachable_code)]
5035    {
5036        let mut acc = 0f32;
5037        for gi in 0..gpr {
5038            let tile = &nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
5039            let s = scales[gi];
5040            let mut d = 0i32;
5041            for (k, &b) in tile.iter().enumerate() {
5042                d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
5043                    + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
5044            }
5045            acc += d as f32 * s;
5046        }
5047        acc
5048    }
5049}
5050
5051/// q4tp twin of `dot_q4t_row_sdot`: identical nibble math, but the tile
5052/// stride is 16 B (no inline scale) and the scale is a ladder lookup.
5053#[cfg(target_arch = "aarch64")]
5054#[target_feature(enable = "neon,dotprod")]
5055unsafe fn dot_q4tp_row_sdot(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
5056    // SAFETY: callers uphold slice-length contracts (16B tile per group,
5057    // xq.len() == gpr·GROUP_SIZE, codes covering gpr 5-bit fields).
5058    unsafe {
5059        use core::arch::aarch64::*;
5060        use core::arch::asm;
5061        let lomask = vdupq_n_u8(0x0F);
5062        let eight = vdupq_n_s8(8);
5063        let mut acc = 0f32;
5064        for gi in 0..gpr {
5065            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5066            let s = *scales.get_unchecked(gi);
5067            let b = vld1q_u8(t);
5068            let lo = vandq_u8(b, lomask);
5069            let hi = vshrq_n_u8::<4>(b);
5070            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
5071            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
5072            let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
5073            let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
5074            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5075            asm!(
5076                "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
5077                "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
5078                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5079                e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
5080                options(pure, nomem, nostack),
5081            );
5082            acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
5083        }
5084        acc
5085    }
5086}
5087
5088#[cfg(target_arch = "x86_64")]
5089#[target_feature(enable = "avx2")]
5090unsafe fn dot_q4tp_row_avx2(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
5091    // SAFETY: see dot_q4tp_row_sdot.
5092    unsafe {
5093        use core::arch::x86_64::*;
5094        let lomask = _mm_set1_epi8(0x0F);
5095        let eight = _mm256_set1_epi8(8);
5096        let ones = _mm256_set1_epi16(1);
5097        let mut acc = 0f32;
5098        for gi in 0..gpr {
5099            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5100            let s = *scales.get_unchecked(gi);
5101            let b = _mm_loadu_si128(t as *const __m128i);
5102            let lo = _mm_and_si128(b, lomask);
5103            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
5104            let w = _mm256_sub_epi8(
5105                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5106                eight,
5107            );
5108            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5109            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
5110            let d = _mm256_madd_epi16(p16, ones);
5111            let hi128 = _mm256_extracti128_si256::<1>(d);
5112            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
5113            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
5114            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
5115            acc += _mm_cvtsi128_si32(s32) as f32 * s;
5116        }
5117        acc
5118    }
5119}
5120
5121/// VNNI twin of `dot_q4tp_row_avx2` (see `dot_q4t_row_vnni` for why the
5122/// 256-bit VL encoding is the one to use here).
5123#[cfg(target_arch = "x86_64")]
5124#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
5125unsafe fn dot_q4tp_row_vnni(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
5126    // SAFETY: see dot_q4tp_row_sdot.
5127    unsafe {
5128        use core::arch::x86_64::*;
5129        let lomask = _mm_set1_epi8(0x0F);
5130        let eight = _mm256_set1_epi8(8);
5131        let mut acc = 0f32;
5132        for gi in 0..gpr {
5133            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5134            let s = *scales.get_unchecked(gi);
5135            let b = _mm_loadu_si128(t as *const __m128i);
5136            let lo = _mm_and_si128(b, lomask);
5137            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
5138            let w = _mm256_sub_epi8(
5139                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
5140                eight,
5141            );
5142            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5143            acc += dpbusd_hsum(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w)) as f32 * s;
5144        }
5145        acc
5146    }
5147}
5148
5149/// Exact scalar q4tp row — the `CMF_SDOT=0` contract, same pairwise
5150/// accumulation shape as `q4t_row_exact`.
5151#[inline]
5152fn q4tp_row_exact(nib: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
5153    let mut acc = 0f32;
5154    for gi in 0..gpr {
5155        let tile = &nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
5156        let s = scales[gi];
5157        let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
5158        let mut ga = 0f32;
5159        for (k, &b) in tile.iter().enumerate() {
5160            ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
5161                + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
5162        }
5163        acc += ga * s;
5164    }
5165    acc
5166}
5167
5168/// Single weight of a q4tp tensor — the a8w8 outlier path, which restores
5169/// activation outliers at full precision after the int8 pass.
5170#[inline]
5171fn q4tp_outlier(nib: &[u8], r: usize, gpr: usize, j: usize, scales: &[f32]) -> (f32, f32) {
5172    let (gi, k) = (j / GROUP_SIZE, j % GROUP_SIZE);
5173    let byte = nib[(r * gpr + gi) * Q4TP_NIB + k / 2];
5174    let n = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
5175    ((n as i32 - 8) as f32, scales[gi])
5176}
5177
5178/// Fused q4tp matvec (dispatch mirrors `q4t_matvec`).
5179fn q4tp_matvec(
5180    bytes: &[u8],
5181    x: &[f32],
5182    rows: usize,
5183    cols: usize,
5184    out: &mut [f32],
5185    pool: Option<&Pool>,
5186) {
5187    debug_assert_eq!(out.len(), rows);
5188    let gpr = cols / GROUP_SIZE;
5189    let v = Q4tpView::new(bytes, rows, cols);
5190    let out_addr = SendMut(out.as_mut_ptr());
5191    if a8w8_enabled() {
5192        let act = split_act(x);
5193        let run = |start: usize, end: usize| {
5194            // One scratch row of scales per worker — borrowed, not minted.
5195            with_krow(gpr, |sc| {
5196                for r in start..end {
5197                    v.scales_into(r, gpr, sc);
5198                    let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, sc) * act.sx;
5199                    for &(j, xv) in &act.outliers {
5200                        let (w, s) = q4tp_outlier(v.nib, r, gpr, j, sc);
5201                        acc += w * s * xv;
5202                    }
5203                    // SAFETY: disjoint row ranges per worker.
5204                    unsafe { *out_addr.at(r) = acc };
5205                }
5206            })
5207        };
5208        dispatch_rows(pool, rows, &run);
5209        return;
5210    }
5211    let run = |start: usize, end: usize| {
5212        with_krow(gpr, |sc| {
5213            for r in start..end {
5214                v.scales_into(r, gpr, sc);
5215                // SAFETY: disjoint row ranges per worker.
5216                unsafe { *out_addr.at(r) = q4tp_row_exact(v.nib, r, gpr, x, sc) };
5217            }
5218        })
5219    };
5220    dispatch_rows(pool, rows, &run);
5221}
5222
5223/// Fused two-input q4tp matvec — the SwiGLU gate/up pair. Weights and the
5224/// row ladder are read once and spent on both activation streams.
5225#[allow(clippy::too_many_arguments)]
5226fn q4tp_matvec2(
5227    bytes: &[u8],
5228    x1: &[f32],
5229    x2: &[f32],
5230    rows: usize,
5231    cols: usize,
5232    o1: &mut [f32],
5233    o2: &mut [f32],
5234    pool: Option<&Pool>,
5235) {
5236    let gpr = cols / GROUP_SIZE;
5237    let v = Q4tpView::new(bytes, rows, cols);
5238    let (p1, p2) = (SendMut(o1.as_mut_ptr()), SendMut(o2.as_mut_ptr()));
5239    let run = |start: usize, end: usize| {
5240        let mut sc = vec![0f32; gpr];
5241        for r in start..end {
5242            v.scales_into(r, gpr, &mut sc);
5243            // SAFETY: disjoint row ranges per worker.
5244            unsafe {
5245                *p1.at(r) = q4tp_row_exact(v.nib, r, gpr, x1, &sc);
5246                *p2.at(r) = q4tp_row_exact(v.nib, r, gpr, x2, &sc);
5247            }
5248        }
5249    };
5250    dispatch_rows(pool, rows, &run);
5251}
5252
5253/// One q2tp outlier weight at column `j` of row `r`: the 2-bit code and
5254/// its group scale, mirrored on `q4tp_outlier`.
5255#[inline]
5256fn q2tp_outlier(chunks: &[u8], r: usize, gpr: usize, j: usize, scales: &[f32]) -> (f32, f32) {
5257    let (gi, k) = (j / GROUP_SIZE, j % GROUP_SIZE);
5258    let byte = chunks[(r * gpr + gi) * Q2TP_CHUNK + k / 4];
5259    let c = (byte >> (2 * (k % 4))) & 3;
5260    (c as f32 - 1.5, scales[gi])
5261}
5262
5263#[cfg(target_arch = "x86_64")]
5264const Q2TP_DECODE_U32: [u32; 256] = {
5265    let mut tab = [0u32; 256];
5266    let mut b = 0usize;
5267    while b < 256 {
5268        tab[b] = ((b as u32) & 3)
5269            | ((((b as u32) >> 2) & 3) << 8)
5270            | ((((b as u32) >> 4) & 3) << 16)
5271            | ((((b as u32) >> 6) & 3) << 24);
5272        b += 1;
5273    }
5274    tab
5275};
5276
5277/// Eight packed q2tp bytes against 32 signed activation bytes. `maddubs`
5278/// exactly computes unsigned 2-bit code × signed i8; its pair sums cannot
5279/// saturate (2 × 3 × 127 < i16::MAX), and the second madd widens to i32.
5280#[cfg(target_arch = "x86_64")]
5281#[target_feature(enable = "avx2")]
5282unsafe fn q2tp_code_dot_avx2(ch: &[u8], x: &[i8]) -> i32 {
5283    use core::arch::x86_64::*;
5284    debug_assert!(ch.len() >= Q2TP_CHUNK && x.len() >= GROUP_SIZE);
5285    let codes = _mm256_setr_epi32(
5286        Q2TP_DECODE_U32[ch[0] as usize] as i32,
5287        Q2TP_DECODE_U32[ch[1] as usize] as i32,
5288        Q2TP_DECODE_U32[ch[2] as usize] as i32,
5289        Q2TP_DECODE_U32[ch[3] as usize] as i32,
5290        Q2TP_DECODE_U32[ch[4] as usize] as i32,
5291        Q2TP_DECODE_U32[ch[5] as usize] as i32,
5292        Q2TP_DECODE_U32[ch[6] as usize] as i32,
5293        Q2TP_DECODE_U32[ch[7] as usize] as i32,
5294    );
5295    let xv = unsafe { _mm256_loadu_si256(x.as_ptr().cast()) };
5296    let pair = _mm256_maddubs_epi16(codes, xv);
5297    let quad = _mm256_madd_epi16(pair, _mm256_set1_epi16(1));
5298    let sum128 = _mm_add_epi32(
5299        _mm256_castsi256_si128(quad),
5300        _mm256_extracti128_si256(quad, 1),
5301    );
5302    let sum64 = _mm_hadd_epi32(sum128, sum128);
5303    _mm_cvtsi128_si32(_mm_hadd_epi32(sum64, sum64))
5304}
5305
5306/// Integer dot of one q2tp row against pre-quantized activations:
5307/// Σ_g s_g · (Σ c·xq − 1.5·Σ xq). The half-integer grid (c − 1.5)
5308/// becomes exact integer math through the group sums — the same trick
5309/// every a8w8 kernel in this file rides. The codes decode into a
5310/// 32-byte scratch in natural order and the dot itself is the shared
5311/// SDOT primitive; elsewhere a scalar integer loop.
5312#[inline]
5313fn dot_q2tp_row_i8(
5314    chunks: &[u8],
5315    r: usize,
5316    gpr: usize,
5317    xq: &[i8],
5318    gsum: &[i32],
5319    scales: &[f32],
5320) -> f32 {
5321    let mut acc = 0f32;
5322    let base = r * gpr * Q2TP_CHUNK;
5323    #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
5324    let mut codes = [0i8; GROUP_SIZE];
5325    #[cfg(target_arch = "x86_64")]
5326    let avx2 = std::arch::is_x86_feature_detected!("avx2");
5327    for gi in 0..gpr {
5328        let ch = &chunks[base + gi * Q2TP_CHUNK..base + (gi + 1) * Q2TP_CHUNK];
5329        let xg = &xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
5330        #[cfg(target_arch = "aarch64")]
5331        // NEON: the byte's four 2-bit fields land in four lane vectors
5332        // (shift+mask), vld4 de-interleaves xq to match (xj[k] =
5333        // xq[4k+j]), widening MACs accumulate exactly in i32. A scalar
5334        // decode here cost as much as the dot it fed — the profile put
5335        // it at the top of the whole W2 decode.
5336        let dot = unsafe {
5337            use core::arch::aarch64::*;
5338            let b = vld1_u8(ch.as_ptr());
5339            let three = vdup_n_u8(3);
5340            let c0 = vreinterpret_s8_u8(vand_u8(b, three));
5341            let c1 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 2), three));
5342            let c2 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 4), three));
5343            let c3 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 6), three));
5344            let x4 = vld4_s8(xg.as_ptr());
5345            let mut acc4 = vdupq_n_s32(0);
5346            acc4 = vpadalq_s16(acc4, vmull_s8(c0, x4.0));
5347            acc4 = vpadalq_s16(acc4, vmull_s8(c1, x4.1));
5348            acc4 = vpadalq_s16(acc4, vmull_s8(c2, x4.2));
5349            acc4 = vpadalq_s16(acc4, vmull_s8(c3, x4.3));
5350            vaddvq_s32(acc4)
5351        };
5352        #[cfg(target_arch = "x86_64")]
5353        let dot: i32 = if avx2 {
5354            // SAFETY: the runtime feature check gates the target-feature body;
5355            // the group slices above are exactly 8 and 32 bytes long.
5356            unsafe { q2tp_code_dot_avx2(ch, xg) }
5357        } else {
5358            ch.iter()
5359                .enumerate()
5360                .map(|(k, &b)| {
5361                    ((b & 3) as i32) * xg[k * 4] as i32
5362                        + (((b >> 2) & 3) as i32) * xg[k * 4 + 1] as i32
5363                        + (((b >> 4) & 3) as i32) * xg[k * 4 + 2] as i32
5364                        + (((b >> 6) & 3) as i32) * xg[k * 4 + 3] as i32
5365                })
5366                .sum()
5367        };
5368        #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
5369        let dot: i32 = {
5370            for (k, &b) in ch.iter().enumerate() {
5371                codes[k * 4] = (b & 3) as i8;
5372                codes[k * 4 + 1] = ((b >> 2) & 3) as i8;
5373                codes[k * 4 + 2] = ((b >> 4) & 3) as i8;
5374                codes[k * 4 + 3] = ((b >> 6) & 3) as i8;
5375            }
5376            codes
5377                .iter()
5378                .zip(xg)
5379                .map(|(&c, &x)| c as i32 * x as i32)
5380                .sum()
5381        };
5382        acc += scales[gi] * (dot as f32 - 1.5 * gsum[gi] as f32);
5383    }
5384    acc
5385}
5386
5387/// Exact f32 dot of one q2tp row: 2-bit fields LSB-first, (c − 1.5)·s.
5388/// Scalar on purpose — the 2-bit class targets the GPU graph; the CPU
5389/// path exists for parity gates and small-machine fallback.
5390fn q2tp_row_exact(chunks: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
5391    q2tp_row_exact_center(chunks, r, gpr, x, scales, 1.5)
5392}
5393
5394/// Fused Prism affine row: the derived correction is applied inside the
5395/// decoded code, avoiding a second accumulated dot and avoiding cancellation
5396/// between `B=(c-1.5)s` and `+.5s` for long 5120/17408 rows.
5397#[inline]
5398fn q2tp_affine_row_exact(chunks: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
5399    q2tp_row_exact_center(chunks, r, gpr, x, scales, 1.0)
5400}
5401
5402#[inline]
5403fn q2tp_row_exact_center(
5404    chunks: &[u8],
5405    r: usize,
5406    gpr: usize,
5407    x: &[f32],
5408    scales: &[f32],
5409    center: f32,
5410) -> f32 {
5411    let mut acc = 0f32;
5412    for gi in 0..gpr {
5413        let ch = &chunks[(r * gpr + gi) * Q2TP_CHUNK..(r * gpr + gi + 1) * Q2TP_CHUNK];
5414        let s = scales[gi];
5415        let xb = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
5416        let mut g = 0f32;
5417        for (k, &b) in ch.iter().enumerate() {
5418            g += ((b & 3) as f32 - center) * xb[k * 4]
5419                + (((b >> 2) & 3) as f32 - center) * xb[k * 4 + 1]
5420                + (((b >> 4) & 3) as f32 - center) * xb[k * 4 + 2]
5421                + (((b >> 6) & 3) as f32 - center) * xb[k * 4 + 3];
5422        }
5423        acc += s * g;
5424    }
5425    acc
5426}
5427
5428fn q2tp_matvec(
5429    bytes: &[u8],
5430    x: &[f32],
5431    rows: usize,
5432    cols: usize,
5433    out: &mut [f32],
5434    pool: Option<&Pool>,
5435) {
5436    q2tp_matvec_mode(bytes, x, rows, cols, out, pool, false);
5437}
5438
5439fn q2tp_affine_matvec(
5440    bytes: &[u8],
5441    x: &[f32],
5442    rows: usize,
5443    cols: usize,
5444    out: &mut [f32],
5445    pool: Option<&Pool>,
5446) {
5447    q2tp_matvec_mode(bytes, x, rows, cols, out, pool, true);
5448}
5449
5450fn q2tp_matvec_mode(
5451    bytes: &[u8],
5452    x: &[f32],
5453    rows: usize,
5454    cols: usize,
5455    out: &mut [f32],
5456    pool: Option<&Pool>,
5457    affine: bool,
5458) {
5459    debug_assert_eq!(out.len(), rows);
5460    let gpr = cols / GROUP_SIZE;
5461    let v = Q4tpView::new_q2(bytes, rows, cols);
5462    let out_addr = SendMut(out.as_mut_ptr());
5463    // a8w8 fast path (CMF_SDOT=0 keeps the exact scalar walk): integer
5464    // code dots + group sums, exact outlier correction — the same
5465    // contract as every sibling kernel; measured 2-bit rows were the
5466    // only scalar holdout in the family.
5467    if !affine && a8w8_enabled() {
5468        let act = split_act(x);
5469        let gsum = q1_group_sums(&act.xq, gpr);
5470        let (act, gsum) = (&act, &gsum);
5471        let run = move |start: usize, end: usize| {
5472            with_krow(gpr, |sc| {
5473                for r in start..end {
5474                    v.scales_into(r, gpr, sc);
5475                    let mut acc = dot_q2tp_row_i8(v.nib, r, gpr, &act.xq, gsum, sc) * act.sx;
5476                    for &(j, xv) in &act.outliers {
5477                        let (w, s) = q2tp_outlier(v.nib, r, gpr, j, sc);
5478                        acc += w * s * xv;
5479                    }
5480                    // SAFETY: disjoint row ranges per worker.
5481                    unsafe { *out_addr.at(r) = acc };
5482                }
5483            })
5484        };
5485        dispatch_rows(pool, rows, &run);
5486        return;
5487    }
5488    let run = |start: usize, end: usize| {
5489        with_krow(gpr, |sc| {
5490            for r in start..end {
5491                v.scales_into(r, gpr, sc);
5492                // SAFETY: disjoint row ranges per worker.
5493                unsafe {
5494                    *out_addr.at(r) = if affine {
5495                        q2tp_affine_row_exact(v.nib, r, gpr, x, sc)
5496                    } else {
5497                        q2tp_row_exact(v.nib, r, gpr, x, sc)
5498                    }
5499                };
5500            }
5501        })
5502    };
5503    dispatch_rows(pool, rows, &run);
5504}
5505
5506/// Fused two-input q2tp matvec — the SwiGLU gate/up pair.
5507#[allow(clippy::too_many_arguments)]
5508fn q2tp_matvec2(
5509    bytes: &[u8],
5510    x1: &[f32],
5511    x2: &[f32],
5512    rows: usize,
5513    cols: usize,
5514    o1: &mut [f32],
5515    o2: &mut [f32],
5516    pool: Option<&Pool>,
5517) {
5518    q2tp_matvec2_mode(bytes, x1, x2, rows, cols, o1, o2, pool, false);
5519}
5520
5521#[allow(clippy::too_many_arguments)]
5522fn q2tp_affine_matvec2(
5523    bytes: &[u8],
5524    x1: &[f32],
5525    x2: &[f32],
5526    rows: usize,
5527    cols: usize,
5528    o1: &mut [f32],
5529    o2: &mut [f32],
5530    pool: Option<&Pool>,
5531) {
5532    q2tp_matvec2_mode(bytes, x1, x2, rows, cols, o1, o2, pool, true);
5533}
5534
5535#[allow(clippy::too_many_arguments)]
5536fn q2tp_matvec2_mode(
5537    bytes: &[u8],
5538    x1: &[f32],
5539    x2: &[f32],
5540    rows: usize,
5541    cols: usize,
5542    o1: &mut [f32],
5543    o2: &mut [f32],
5544    pool: Option<&Pool>,
5545    affine: bool,
5546) {
5547    let gpr = cols / GROUP_SIZE;
5548    let v = Q4tpView::new_q2(bytes, rows, cols);
5549    let (p1, p2) = (SendMut(o1.as_mut_ptr()), SendMut(o2.as_mut_ptr()));
5550    let run = |start: usize, end: usize| {
5551        let mut sc = vec![0f32; gpr];
5552        for r in start..end {
5553            v.scales_into(r, gpr, &mut sc);
5554            // SAFETY: disjoint row ranges per worker.
5555            unsafe {
5556                *p1.at(r) = if affine {
5557                    q2tp_affine_row_exact(v.nib, r, gpr, x1, &sc)
5558                } else {
5559                    q2tp_row_exact(v.nib, r, gpr, x1, &sc)
5560                };
5561                *p2.at(r) = if affine {
5562                    q2tp_affine_row_exact(v.nib, r, gpr, x2, &sc)
5563                } else {
5564                    q2tp_row_exact(v.nib, r, gpr, x2, &sc)
5565                };
5566            }
5567        }
5568    };
5569    dispatch_rows(pool, rows, &run);
5570}
5571
5572/// Batched q2tp matmat: scalar row kernel over every batch column. CPU
5573/// prefill only — decode rides the graph, so plain and correct beats
5574/// clever here.
5575/// Test doors into the host 2-bit kernels: the stand's heap corruption
5576/// pointed at down-shaped tensors, and the private fns need a way to be
5577/// held to a reference without a model file around them.
5578pub fn q2tp_matvec_for_test(bytes: &[u8], x: &[f32], rows: usize, cols: usize, out: &mut [f32]) {
5579    // The facade IS the reference: encoder oracles hold requant output
5580    // to the exact scalar walk. The production dispatch may take the i8
5581    // fast path, whose error scale is the ACTIVATIONS' — a different
5582    // claim than the encoder correctness these tests pin.
5583    let gpr = cols / GROUP_SIZE;
5584    let v = Q4tpView::new_q2(bytes, rows, cols);
5585    with_krow(gpr, |sc| {
5586        for r in 0..rows {
5587            v.scales_into(r, gpr, sc);
5588            out[r] = q2tp_row_exact(v.nib, r, gpr, x, sc);
5589        }
5590    });
5591}
5592
5593/// Test door for the descriptor-specific fused affine decode.  Production
5594/// callers select this through a validated Prism header, never by dtype alone.
5595pub fn q2tp_affine_matvec_for_test(
5596    bytes: &[u8],
5597    x: &[f32],
5598    rows: usize,
5599    cols: usize,
5600    out: &mut [f32],
5601) {
5602    q2tp_affine_matvec(bytes, x, rows, cols, out, None);
5603}
5604
5605pub fn q2tp_matmat_for_test(
5606    bytes: &[u8],
5607    xs_all: &[f32],
5608    b: usize,
5609    rows: usize,
5610    cols: usize,
5611    out: &mut [f32],
5612) {
5613    q2tp_matmat(bytes, xs_all, b, rows, cols, out, None);
5614}
5615
5616fn q2tp_matmat(
5617    bytes: &[u8],
5618    xs_all: &[f32],
5619    b: usize,
5620    rows: usize,
5621    cols: usize,
5622    out: &mut [f32],
5623    pool: Option<&Pool>,
5624) {
5625    q2tp_matmat_mode(bytes, xs_all, b, rows, cols, out, pool, false);
5626}
5627
5628fn q2tp_affine_matmat(
5629    bytes: &[u8],
5630    xs_all: &[f32],
5631    b: usize,
5632    rows: usize,
5633    cols: usize,
5634    out: &mut [f32],
5635    pool: Option<&Pool>,
5636) {
5637    q2tp_matmat_mode(bytes, xs_all, b, rows, cols, out, pool, true);
5638}
5639
5640fn q2tp_matmat_mode(
5641    bytes: &[u8],
5642    xs_all: &[f32],
5643    b: usize,
5644    rows: usize,
5645    cols: usize,
5646    out: &mut [f32],
5647    pool: Option<&Pool>,
5648    affine: bool,
5649) {
5650    debug_assert_eq!(out.len(), b * rows);
5651    let gpr = cols / GROUP_SIZE;
5652    let v = Q4tpView::new_q2(bytes, rows, cols);
5653    let out_addr = SendMut(out.as_mut_ptr());
5654    let run = |start: usize, end: usize| {
5655        let mut sc = vec![0f32; gpr];
5656        for r in start..end {
5657            v.scales_into(r, gpr, &mut sc);
5658            for bi in 0..b {
5659                let x = &xs_all[bi * cols..(bi + 1) * cols];
5660                // SAFETY: disjoint row ranges per worker.
5661                unsafe {
5662                    *out_addr.at(bi * rows + r) = if affine {
5663                        q2tp_affine_row_exact(v.nib, r, gpr, x, &sc)
5664                    } else {
5665                        q2tp_row_exact(v.nib, r, gpr, x, &sc)
5666                    }
5667                };
5668            }
5669        }
5670    };
5671    dispatch_rows(pool, rows, &run);
5672}
5673
5674/// The pre-vectorised shape, kept for A/B (`CMF_Q4TP_V1=1`): the
5675/// horizontal add lands once per group per column instead of once per
5676/// row. Same weights, same activations — only the reduction differs.
5677#[cfg(target_arch = "aarch64")]
5678#[target_feature(enable = "neon,dotprod")]
5679unsafe fn dot_q4tp_row_1x4_sdot_v1(
5680    nib: &[u8],
5681    r: usize,
5682    gpr: usize,
5683    xs: [&[i8]; 4],
5684    scales: &[f32],
5685) -> [f32; 4] {
5686    unsafe {
5687        use core::arch::aarch64::*;
5688        use core::arch::asm;
5689        let lomask = vdupq_n_u8(0x0F);
5690        let eight = vdupq_n_s8(8);
5691        let (mut f0, mut f1, mut f2, mut f3) = (0f32, 0f32, 0f32, 0f32);
5692        for gi in 0..gpr {
5693            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5694            let s = *scales.get_unchecked(gi);
5695            let bb = vld1q_u8(t);
5696            let lo = vandq_u8(bb, lomask);
5697            let hi = vshrq_n_u8::<4>(bb);
5698            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
5699            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
5700            let mut d = [0f32; 4];
5701            for (k, dk) in d.iter_mut().enumerate() {
5702                let x0 = vld1q_s8(xs[k].as_ptr().add(gi * GROUP_SIZE));
5703                let x1 = vld1q_s8(xs[k].as_ptr().add(gi * GROUP_SIZE + 16));
5704                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5705                asm!(
5706                    "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
5707                    "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
5708                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5709                    e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
5710                    options(pure, nomem, nostack),
5711                );
5712                *dk = vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
5713            }
5714            f0 += d[0];
5715            f1 += d[1];
5716            f2 += d[2];
5717            f3 += d[3];
5718        }
5719        [f0, f1, f2, f3]
5720    }
5721}
5722
5723/// Which q4tp batch kernel to run: 1 = the previous one, 2 = the tuned
5724/// one, 0 = decide from the CPU. An atomic rather than a `OnceLock` so a
5725/// benchmark can alternate the two inside one process, where the machine's
5726/// mood — a shared box drifts ±25% between runs — is the same for both.
5727/// What the two mean is per-architecture: on x86 the blocked AVX-512 path
5728/// against the per-column one, on ARM the two reduction shapes.
5729#[allow(dead_code)]
5730static Q4TP_ALT: std::sync::atomic::AtomicU8 = std::sync::atomic::AtomicU8::new(0);
5731
5732/// Blocking pays on x86 only with 512-bit VNNI. With AVX2 alone, four
5733/// columns sharing an unpack still measured slower than the per-column
5734/// path (23.2 ms against 19.4 on a 48-thread EPYC), because that path
5735/// already dequantizes the row once — so the blocked kernel bought a
5736/// second unpack-free pass at the price of half the vector width.
5737#[cfg(target_arch = "x86_64")]
5738fn q4tp_blocked_x86() -> bool {
5739    match Q4TP_ALT.load(std::sync::atomic::Ordering::Relaxed) {
5740        1 => false,
5741        // A forced ON still asks the CPU. The switch exists so a bench can
5742        // pick a kernel, not so it can promise instructions the machine
5743        // does not have — CI caught that as a SIGILL on a runner without
5744        // AVX-512, where the parity test had turned the path on by hand.
5745        2 => avx512vnni_enabled(),
5746        // Deliberately not cached back into the switch: both gates below
5747        // hold their own `OnceLock`, and latching their answer here would
5748        // make a test's override outlive the test that set it.
5749        _ => blocked_enabled() && avx512vnni_enabled(),
5750    }
5751}
5752
5753/// `CMF_Q4TP_V1=1` picks the old reduction shape (A/B only).
5754#[cfg(target_arch = "aarch64")]
5755#[allow(dead_code)]
5756fn q4tp_v1() -> bool {
5757    match Q4TP_ALT.load(std::sync::atomic::Ordering::Relaxed) {
5758        1 => true,
5759        2 => false,
5760        _ => {
5761            static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
5762            *ON.get_or_init(|| std::env::var("CMF_Q4TP_V1").is_ok_and(|v| v != "0"))
5763        }
5764    }
5765}
5766
5767/// Two weight rows against eight columns. The activation load is the
5768/// same for both rows, so it is paid once for twice the arithmetic, and
5769/// sixteen accumulator chains run where eight did — which is what a kernel
5770/// retiring 0.29 instructions a cycle is short of. Register pressure is
5771/// the limit: sixteen `zmm` accumulators, two weight tiles, one
5772/// activation, of thirty-two.
5773///
5774/// Four rows by four columns spends the same sixteen accumulators the
5775/// other way and measured worse — 1488 GFLOP/s against 1644 — so the
5776/// unpack, which four rows pay twice as often, costs more than the extra
5777/// sharing of one activation load buys.
5778#[cfg(target_arch = "x86_64")]
5779#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
5780unsafe fn dot_q4tp_2x8_avx512(
5781    nib: &[u8],
5782    r0: usize,
5783    gpr: usize,
5784    xs: [&[i8]; 8],
5785    sc0: &[f32],
5786    sc1: &[f32],
5787) -> [[f32; 8]; 2] {
5788    // SAFETY: as dot_q4tp_row_1x8_avx512, two adjacent rows at once; the
5789    // caller guarantees r0 + 1 < rows and the ISA.
5790    unsafe {
5791        use core::arch::x86_64::*;
5792        let lomask = _mm256_set1_epi8(0x0F);
5793        let eight = _mm256_set1_epi8(8);
5794        let zero = _mm512_setzero_si512();
5795        let mut v0 = [_mm512_setzero_ps(); 8];
5796        let mut v1 = [_mm512_setzero_ps(); 8];
5797        let pairs = gpr / 2;
5798        let unpack = |r: usize, gi: usize| -> (__m512i, __mmask64) {
5799            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5800            let bb = _mm256_loadu_si256(t as *const __m256i);
5801            let lo = _mm256_and_si256(bb, lomask);
5802            let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
5803            let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
5804            let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
5805            let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
5806            let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
5807            (_mm512_abs_epi8(w), _mm512_movepi8_mask(w))
5808        };
5809        for gp in 0..pairs {
5810            let gi = gp * 2;
5811            let (wa0, neg0) = unpack(r0, gi);
5812            let (wa1, neg1) = unpack(r0 + 1, gi);
5813            let off = gi * GROUP_SIZE;
5814            let sv = |sc: &[f32]| {
5815                _mm512_insertf32x8::<1>(
5816                    _mm512_castps256_ps512(_mm256_set1_ps(*sc.get_unchecked(gi))),
5817                    _mm256_set1_ps(*sc.get_unchecked(gi + 1)),
5818                )
5819            };
5820            let s0 = sv(sc0);
5821            let s1 = sv(sc1);
5822            for k in 0..8 {
5823                let xv = _mm512_loadu_si512(xs[k].as_ptr().add(off) as *const __m512i);
5824                let d0 = _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(
5825                    zero,
5826                    wa0,
5827                    _mm512_mask_sub_epi8(xv, neg0, zero, xv),
5828                ));
5829                let d1 = _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(
5830                    zero,
5831                    wa1,
5832                    _mm512_mask_sub_epi8(xv, neg1, zero, xv),
5833                ));
5834                v0[k] = _mm512_fmadd_ps(d0, s0, v0[k]);
5835                v1[k] = _mm512_fmadd_ps(d1, s1, v1[k]);
5836            }
5837        }
5838        let mut acc = [[0f32; 8]; 2];
5839        for k in 0..8 {
5840            acc[0][k] = _mm512_reduce_add_ps(v0[k]);
5841            acc[1][k] = _mm512_reduce_add_ps(v1[k]);
5842        }
5843        if gpr % 2 == 1 {
5844            let off = (gpr - 1) * GROUP_SIZE;
5845            for j in off..off + GROUP_SIZE {
5846                let (w0, sa) = q4tp_outlier(nib, r0, gpr, j, sc0);
5847                let (w1, sb) = q4tp_outlier(nib, r0 + 1, gpr, j, sc1);
5848                for k in 0..8 {
5849                    let x = *xs[k].get_unchecked(j) as f32;
5850                    acc[0][k] += w0 * sa * x;
5851                    acc[1][k] += w1 * sb * x;
5852                }
5853            }
5854        }
5855        acc
5856    }
5857}
5858
5859/// The same, eight columns at a time. One unpack then feeds twice as many
5860/// activation streams, so a wide batch reads the weight tile half as
5861/// often; the price is eight accumulators live at once. Measured 9.0 ->
5862/// 8.3 ms at 9216x2304, b=296 on a 48-thread EPYC 9B45.
5863#[cfg(target_arch = "x86_64")]
5864#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
5865unsafe fn dot_q4tp_row_1x8_avx512(
5866    nib: &[u8],
5867    r: usize,
5868    gpr: usize,
5869    xs: [&[i8]; 8],
5870    scales: &[f32],
5871) -> [f32; 8] {
5872    // SAFETY: as dot_q4tp_row_1x4_avx2; caller guarantees the ISA.
5873    unsafe {
5874        use core::arch::x86_64::*;
5875        let lomask = _mm256_set1_epi8(0x0F);
5876        let eight = _mm256_set1_epi8(8);
5877        let zero = _mm512_setzero_si512();
5878        let (mut v0, mut v1, mut v2, mut v3) = (
5879            _mm512_setzero_ps(),
5880            _mm512_setzero_ps(),
5881            _mm512_setzero_ps(),
5882            _mm512_setzero_ps(),
5883        );
5884        let (mut v4, mut v5, mut v6, mut v7) = (
5885            _mm512_setzero_ps(),
5886            _mm512_setzero_ps(),
5887            _mm512_setzero_ps(),
5888            _mm512_setzero_ps(),
5889        );
5890        let pairs = gpr / 2;
5891        for gp in 0..pairs {
5892            let gi = gp * 2;
5893            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5894            let bb = _mm256_loadu_si256(t as *const __m256i);
5895            let lo = _mm256_and_si256(bb, lomask);
5896            let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
5897            // `unpack` works per 128-bit lane, so the halves come out as
5898            // [A.lo, B.lo] and [A.hi, B.hi]; the shuffle reorders the four
5899            // 128-bit lanes into the weights' natural order, which is what
5900            // the straight activation load expects.
5901            let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
5902            let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
5903            let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
5904            let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
5905            let wabs = _mm512_abs_epi8(w);
5906            let neg = _mm512_movepi8_mask(w);
5907            let off = gi * GROUP_SIZE;
5908            let sv = _mm512_insertf32x8::<1>(
5909                _mm512_castps256_ps512(_mm256_set1_ps(*scales.get_unchecked(gi))),
5910                _mm256_set1_ps(*scales.get_unchecked(gi + 1)),
5911            );
5912            let dot = |x: &[i8]| -> __m512 {
5913                let xv = _mm512_loadu_si512(x.as_ptr().add(off) as *const __m512i);
5914                let sx = _mm512_mask_sub_epi8(xv, neg, zero, xv);
5915                _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(zero, wabs, sx))
5916            };
5917            v0 = _mm512_fmadd_ps(dot(xs[0]), sv, v0);
5918            v1 = _mm512_fmadd_ps(dot(xs[1]), sv, v1);
5919            v2 = _mm512_fmadd_ps(dot(xs[2]), sv, v2);
5920            v3 = _mm512_fmadd_ps(dot(xs[3]), sv, v3);
5921            v4 = _mm512_fmadd_ps(dot(xs[4]), sv, v4);
5922            v5 = _mm512_fmadd_ps(dot(xs[5]), sv, v5);
5923            v6 = _mm512_fmadd_ps(dot(xs[6]), sv, v6);
5924            v7 = _mm512_fmadd_ps(dot(xs[7]), sv, v7);
5925        }
5926        let mut acc = [
5927            _mm512_reduce_add_ps(v0),
5928            _mm512_reduce_add_ps(v1),
5929            _mm512_reduce_add_ps(v2),
5930            _mm512_reduce_add_ps(v3),
5931            _mm512_reduce_add_ps(v4),
5932            _mm512_reduce_add_ps(v5),
5933            _mm512_reduce_add_ps(v6),
5934            _mm512_reduce_add_ps(v7),
5935        ];
5936        // An odd group count leaves one group over; the narrow kernel
5937        // finishes it rather than the tail being a special case here.
5938        if gpr % 2 == 1 {
5939            let off = (gpr - 1) * GROUP_SIZE;
5940            for j in off..off + GROUP_SIZE {
5941                let (w, s) = q4tp_outlier(nib, r, gpr, j, scales);
5942                let ws = w * s;
5943                for k in 0..8 {
5944                    acc[k] += ws * *xs[k].get_unchecked(j) as f32;
5945                }
5946            }
5947        }
5948        acc
5949    }
5950}
5951
5952/// The same four columns, 512 bits wide. Two groups (64 weights) ride one
5953/// unpack and one `vpdpbusd`, where AVX2 needs two unpacks and four
5954/// `maddubs`/`madd` pairs — about 2.3x fewer instructions for the same
5955/// arithmetic. The two groups carry different scales, so the fma takes a
5956/// vector whose halves hold each group's scale rather than a broadcast.
5957///
5958/// There is no 512-bit `vpsignb`, so the activation's sign is applied by
5959/// negating under a mask taken from the weight's sign bits. That mask is
5960/// per-tile, so it is hoisted out of the column loop and the per-column
5961/// cost stays exactly one instruction, as with `sign_epi8`. Weights of
5962/// zero are not zeroed by the mask trick and do not need to be: their
5963/// magnitude is zero, so the product is.
5964#[cfg(target_arch = "x86_64")]
5965#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
5966unsafe fn dot_q4tp_row_1x4_avx512(
5967    nib: &[u8],
5968    r: usize,
5969    gpr: usize,
5970    xs: [&[i8]; 4],
5971    scales: &[f32],
5972) -> [f32; 4] {
5973    // SAFETY: as dot_q4tp_row_1x4_avx2; caller guarantees the ISA.
5974    unsafe {
5975        use core::arch::x86_64::*;
5976        let lomask = _mm256_set1_epi8(0x0F);
5977        let eight = _mm256_set1_epi8(8);
5978        let zero = _mm512_setzero_si512();
5979        let (mut v0, mut v1, mut v2, mut v3) = (
5980            _mm512_setzero_ps(),
5981            _mm512_setzero_ps(),
5982            _mm512_setzero_ps(),
5983            _mm512_setzero_ps(),
5984        );
5985        let pairs = gpr / 2;
5986        for gp in 0..pairs {
5987            let gi = gp * 2;
5988            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5989            let bb = _mm256_loadu_si256(t as *const __m256i);
5990            let lo = _mm256_and_si256(bb, lomask);
5991            let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
5992            // `unpack` works per 128-bit lane, so the halves come out as
5993            // [A.lo, B.lo] and [A.hi, B.hi]; the shuffle reorders the four
5994            // 128-bit lanes into the weights' natural order, which is what
5995            // the straight activation load expects.
5996            let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
5997            let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
5998            let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
5999            let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
6000            let wabs = _mm512_abs_epi8(w);
6001            let neg = _mm512_movepi8_mask(w);
6002            let off = gi * GROUP_SIZE;
6003            let sv = _mm512_insertf32x8::<1>(
6004                _mm512_castps256_ps512(_mm256_set1_ps(*scales.get_unchecked(gi))),
6005                _mm256_set1_ps(*scales.get_unchecked(gi + 1)),
6006            );
6007            let dot = |x: &[i8]| -> __m512 {
6008                let xv = _mm512_loadu_si512(x.as_ptr().add(off) as *const __m512i);
6009                let sx = _mm512_mask_sub_epi8(xv, neg, zero, xv);
6010                _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(zero, wabs, sx))
6011            };
6012            v0 = _mm512_fmadd_ps(dot(xs[0]), sv, v0);
6013            v1 = _mm512_fmadd_ps(dot(xs[1]), sv, v1);
6014            v2 = _mm512_fmadd_ps(dot(xs[2]), sv, v2);
6015            v3 = _mm512_fmadd_ps(dot(xs[3]), sv, v3);
6016        }
6017        let mut acc = [
6018            _mm512_reduce_add_ps(v0),
6019            _mm512_reduce_add_ps(v1),
6020            _mm512_reduce_add_ps(v2),
6021            _mm512_reduce_add_ps(v3),
6022        ];
6023        // An odd group count leaves one group over; the narrow kernel
6024        // finishes it rather than the tail being a special case here.
6025        if gpr % 2 == 1 {
6026            let off = (gpr - 1) * GROUP_SIZE;
6027            for j in off..off + GROUP_SIZE {
6028                let (w, s) = q4tp_outlier(nib, r, gpr, j, scales);
6029                let ws = w * s;
6030                for k in 0..4 {
6031                    acc[k] += ws * *xs[k].get_unchecked(j) as f32;
6032                }
6033            }
6034        }
6035        acc
6036    }
6037}
6038
6039/// Four batch columns against one q4tp row: the tile is unpacked ONCE and
6040/// spent on four activation streams, which is where a prefill batch stops
6041/// being weight-bandwidth-bound. Twin of `dot_q4t_row_1x4_sdot`.
6042#[cfg(target_arch = "aarch64")]
6043#[target_feature(enable = "neon,dotprod")]
6044unsafe fn dot_q4tp_row_1x4_sdot(
6045    nib: &[u8],
6046    r: usize,
6047    gpr: usize,
6048    xs: [&[i8]; 4],
6049    scales: &[f32],
6050) -> [f32; 4] {
6051    // SAFETY: see dot_q4tp_row_sdot; every xs[k] is gpr·GROUP_SIZE long.
6052    unsafe {
6053        use core::arch::aarch64::*;
6054        use core::arch::asm;
6055        let lomask = vdupq_n_u8(0x0F);
6056        let eight = vdupq_n_s8(8);
6057        // Named accumulators, NOT an array indexed by a loop variable: the
6058        // latter does not stay in registers (the same defect cost 2x in the
6059        // AVX2 q4t kernel and again in WGSL).
6060        //
6061        // They are VECTORS, and the horizontal add happens once at the end
6062        // instead of once per group per column. `vaddvq` is a cross-lane
6063        // reduction — with 72 groups and four columns the old shape paid
6064        // 288 of them per row, each one a dependency stall the pipeline
6065        // cannot hide, to save four float adds. The group's scale now
6066        // rides an fma into the lane accumulators, so the arithmetic per
6067        // group is one convert and one fma. Summation order changes (the
6068        // lanes carry independent partial sums), which is the same
6069        // round-off class the SDOT path already lives in — the strict
6070        // kernel (`CMF_SDOT=0`, what `cortiq ppl` runs) is unchanged and
6071        // stays the reference.
6072        let (mut v0, mut v1, mut v2, mut v3) = (
6073            vdupq_n_f32(0.0),
6074            vdupq_n_f32(0.0),
6075            vdupq_n_f32(0.0),
6076            vdupq_n_f32(0.0),
6077        );
6078        for gi in 0..gpr {
6079            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
6080            let s = *scales.get_unchecked(gi);
6081            let bb = vld1q_u8(t);
6082            let lo = vandq_u8(bb, lomask);
6083            let hi = vshrq_n_u8::<4>(bb);
6084            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
6085            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
6086            let off = gi * GROUP_SIZE;
6087            let dot4 = |x: &[i8]| -> int32x4_t {
6088                let x0 = vld1q_s8(x.as_ptr().add(off));
6089                let x1 = vld1q_s8(x.as_ptr().add(off + 16));
6090                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6091                asm!(
6092                    "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
6093                    "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
6094                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6095                    e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
6096                    options(pure, nomem, nostack),
6097                );
6098                vaddq_s32(a0, a1)
6099            };
6100            v0 = vfmaq_n_f32(v0, vcvtq_f32_s32(dot4(xs[0])), s);
6101            v1 = vfmaq_n_f32(v1, vcvtq_f32_s32(dot4(xs[1])), s);
6102            v2 = vfmaq_n_f32(v2, vcvtq_f32_s32(dot4(xs[2])), s);
6103            v3 = vfmaq_n_f32(v3, vcvtq_f32_s32(dot4(xs[3])), s);
6104        }
6105        [
6106            vaddvq_f32(v0),
6107            vaddvq_f32(v1),
6108            vaddvq_f32(v2),
6109            vaddvq_f32(v3),
6110        ]
6111    }
6112}
6113
6114/// Fused q4tp matmat — the same three arms `q4t_matmat` has. Shipping only
6115/// the scalar one made Nanbeige-3B decode at 1.2 tok/s against q4t's 5.9:
6116/// the format was fine, the missing arms were the whole regression.
6117fn q4tp_matmat(
6118    bytes: &[u8],
6119    xs_all: &[f32],
6120    b: usize,
6121    rows: usize,
6122    cols: usize,
6123    out: &mut [f32],
6124    pool: Option<&Pool>,
6125) {
6126    debug_assert_eq!(out.len(), b * rows);
6127    let gpr = cols / GROUP_SIZE;
6128    let v = Q4tpView::new(bytes, rows, cols);
6129
6130    // Wide batches ride the AMX through a dequant-tile sgemm, as in q4t.
6131    #[cfg(target_os = "macos")]
6132    if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
6133        dequant_matmat_accel(
6134            &|r, dst| {
6135                let mut sc = [0f32; 32];
6136                let mut scv;
6137                let s: &[f32] = if gpr <= 32 {
6138                    v.scales_into(r, gpr, &mut sc);
6139                    &sc[..gpr]
6140                } else {
6141                    scv = vec![0f32; gpr];
6142                    v.scales_into(r, gpr, &mut scv);
6143                    &scv
6144                };
6145                for gi in 0..gpr {
6146                    let tile = &v.nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
6147                    for (k, &bb) in tile.iter().enumerate() {
6148                        dst[gi * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s[gi];
6149                        dst[gi * GROUP_SIZE + k * 2 + 1] =
6150                            (((bb >> 4) & 0x0F) as f32 - 8.0) * s[gi];
6151                    }
6152                }
6153            },
6154            xs_all,
6155            b,
6156            rows,
6157            cols,
6158            out,
6159            pool,
6160        );
6161        return;
6162    }
6163
6164    let out_addr = SendMut(out.as_mut_ptr());
6165    if a8w8_enabled() {
6166        let acts: Vec<SplitAct> = (0..b)
6167            .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
6168            .collect();
6169        let acts = &acts;
6170        #[cfg(target_arch = "aarch64")]
6171        let blocked_ok = sdot_enabled() && blocked_enabled();
6172        // x86 gets the same blocking: one tile unpack spent on four
6173        // columns. Without it every column re-decoded the row, which is
6174        // why a 48-core EPYC measured a sixth of an M4's per-core rate.
6175        // The gate is `avx2_enabled`, as in q4t — `sdot_enabled` answers
6176        // for ARM's dotprod and is hard-wired false everywhere else, so
6177        // asking it here left the whole blocked path unreachable on x86.
6178        #[cfg(target_arch = "x86_64")]
6179        let blocked_ok = q4tp_blocked_x86();
6180        #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
6181        let blocked_ok = false;
6182        // Columns are swept in panels that fit L2. Without this a
6183        // row-pair walks every activation in the batch — 4.8 MB at
6184        // 512x512 — and does it again for the next pair, so the whole
6185        // batch streams out of the shared cache once per row. Measured
6186        // 800 GB/s of it, flat across batch sizes, which is the signature
6187        // of a loop bound by traffic rather than by arithmetic. A panel of
6188        // 256 columns is 590 KB beside 221 KB of this worker's weights:
6189        // both stay resident and the batch crosses L3 once instead of
6190        // once per row.
6191        let panel_cols: usize = std::env::var("CMF_Q4TP_PANEL")
6192            .ok()
6193            .and_then(|v| v.parse().ok())
6194            .filter(|v| *v > 0)
6195            .unwrap_or(256);
6196        let run = |start: usize, end: usize| {
6197            for abase in (0..acts.len()).step_by(panel_cols) {
6198                let alen = (acts.len() - abase).min(panel_cols);
6199                let mut sc = vec![0f32; gpr];
6200                #[cfg(target_arch = "x86_64")]
6201                let mut r_lo = start;
6202                #[cfg(target_arch = "x86_64")]
6203                if blocked_ok && alen >= 8 {
6204                    let mut sc1 = vec![0f32; gpr];
6205                    while r_lo + 2 <= end {
6206                        v.scales_into(r_lo, gpr, &mut sc);
6207                        v.scales_into(r_lo + 1, gpr, &mut sc1);
6208                        let mut bi = 0usize;
6209                        while bi + 8 <= alen {
6210                            let xs = [
6211                                acts[abase + bi].xq.as_slice(),
6212                                acts[abase + bi + 1].xq.as_slice(),
6213                                acts[abase + bi + 2].xq.as_slice(),
6214                                acts[abase + bi + 3].xq.as_slice(),
6215                                acts[abase + bi + 4].xq.as_slice(),
6216                                acts[abase + bi + 5].xq.as_slice(),
6217                                acts[abase + bi + 6].xq.as_slice(),
6218                                acts[abase + bi + 7].xq.as_slice(),
6219                            ];
6220                            let d = unsafe { dot_q4tp_2x8_avx512(v.nib, r_lo, gpr, xs, &sc, &sc1) };
6221                            for (row, dr, scr) in [(r_lo, &d[0], &sc), (r_lo + 1, &d[1], &sc1)] {
6222                                for k in 0..8 {
6223                                    let act = &acts[abase + bi + k];
6224                                    let mut acc = dr[k] * act.sx;
6225                                    for &(j, xv) in &act.outliers {
6226                                        let (w, s) = q4tp_outlier(v.nib, row, gpr, j, scr);
6227                                        acc += w * s * xv;
6228                                    }
6229                                    // SAFETY: disjoint (bi, r) cells per worker.
6230                                    unsafe { *out_addr.at((abase + bi + k) * rows + row) = acc };
6231                                }
6232                            }
6233                            bi += 8;
6234                        }
6235                        // Columns past the last group of eight, both rows —
6236                        // the same single-row kernel the tail below uses.
6237                        for row in [r_lo, r_lo + 1] {
6238                            let scr: &[f32] = if row == r_lo { &sc } else { &sc1 };
6239                            for b2 in bi..alen {
6240                                let act = &acts[abase + b2];
6241                                let xs4 = [
6242                                    act.xq.as_slice(),
6243                                    act.xq.as_slice(),
6244                                    act.xq.as_slice(),
6245                                    act.xq.as_slice(),
6246                                ];
6247                                let d =
6248                                    unsafe { dot_q4tp_row_1x4_avx512(v.nib, row, gpr, xs4, scr) };
6249                                let mut acc = d[0] * act.sx;
6250                                for &(j, xv) in &act.outliers {
6251                                    let (w, s) = q4tp_outlier(v.nib, row, gpr, j, scr);
6252                                    acc += w * s * xv;
6253                                }
6254                                // SAFETY: disjoint (bi, r) cells per worker.
6255                                unsafe { *out_addr.at((abase + b2) * rows + row) = acc };
6256                            }
6257                        }
6258                        r_lo += 2;
6259                    }
6260                }
6261                #[cfg(target_arch = "x86_64")]
6262                let row_start = r_lo;
6263                #[cfg(not(target_arch = "x86_64"))]
6264                let row_start = start;
6265                for r in row_start..end {
6266                    v.scales_into(r, gpr, &mut sc);
6267                    let mut bi = 0usize;
6268                    #[cfg(target_arch = "x86_64")]
6269                    if blocked_ok {
6270                        while bi + 8 <= alen {
6271                            let xs = [
6272                                acts[abase + bi].xq.as_slice(),
6273                                acts[abase + bi + 1].xq.as_slice(),
6274                                acts[abase + bi + 2].xq.as_slice(),
6275                                acts[abase + bi + 3].xq.as_slice(),
6276                                acts[abase + bi + 4].xq.as_slice(),
6277                                acts[abase + bi + 5].xq.as_slice(),
6278                                acts[abase + bi + 6].xq.as_slice(),
6279                                acts[abase + bi + 7].xq.as_slice(),
6280                            ];
6281                            let d = unsafe { dot_q4tp_row_1x8_avx512(v.nib, r, gpr, xs, &sc) };
6282                            for k in 0..8 {
6283                                let act = &acts[abase + bi + k];
6284                                let mut acc = d[k] * act.sx;
6285                                for &(j, xv) in &act.outliers {
6286                                    let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
6287                                    acc += w * s * xv;
6288                                }
6289                                // SAFETY: disjoint (bi, r) cells per worker.
6290                                unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
6291                            }
6292                            bi += 8;
6293                        }
6294                        while bi + 4 <= alen {
6295                            let xs = [
6296                                acts[abase + bi].xq.as_slice(),
6297                                acts[abase + bi + 1].xq.as_slice(),
6298                                acts[abase + bi + 2].xq.as_slice(),
6299                                acts[abase + bi + 3].xq.as_slice(),
6300                            ];
6301                            let d = unsafe { dot_q4tp_row_1x4_avx512(v.nib, r, gpr, xs, &sc) };
6302                            for k in 0..4 {
6303                                let act = &acts[abase + bi + k];
6304                                let mut acc = d[k] * act.sx;
6305                                for &(j, xv) in &act.outliers {
6306                                    let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
6307                                    acc += w * s * xv;
6308                                }
6309                                // SAFETY: disjoint (bi, r) cells per worker.
6310                                unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
6311                            }
6312                            bi += 4;
6313                        }
6314                    }
6315                    #[cfg(target_arch = "aarch64")]
6316                    if blocked_ok {
6317                        while bi + 4 <= alen {
6318                            let xs = [
6319                                acts[abase + bi].xq.as_slice(),
6320                                acts[abase + bi + 1].xq.as_slice(),
6321                                acts[abase + bi + 2].xq.as_slice(),
6322                                acts[abase + bi + 3].xq.as_slice(),
6323                            ];
6324                            let d = unsafe {
6325                                if q4tp_v1() {
6326                                    dot_q4tp_row_1x4_sdot_v1(v.nib, r, gpr, xs, &sc)
6327                                } else {
6328                                    dot_q4tp_row_1x4_sdot(v.nib, r, gpr, xs, &sc)
6329                                }
6330                            };
6331                            for k in 0..4 {
6332                                let act = &acts[abase + bi + k];
6333                                let mut acc = d[k] * act.sx;
6334                                for &(j, xv) in &act.outliers {
6335                                    let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
6336                                    acc += w * s * xv;
6337                                }
6338                                // SAFETY: disjoint (bi, r) cells per worker.
6339                                unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
6340                            }
6341                            bi += 4;
6342                        }
6343                    }
6344                    let _ = blocked_ok;
6345                    while bi < alen {
6346                        let act = &acts[abase + bi];
6347                        let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, &sc) * act.sx;
6348                        for &(j, xv) in &act.outliers {
6349                            let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
6350                            acc += w * s * xv;
6351                        }
6352                        // SAFETY: disjoint (bi, r) cells per worker range.
6353                        unsafe { *out_addr.at((abase + bi) * rows + r) = acc };
6354                        bi += 1;
6355                    }
6356                }
6357            }
6358        };
6359        dispatch_rows(pool, rows, &run);
6360        return;
6361    }
6362
6363    let run = |start: usize, end: usize| {
6364        let mut sc = vec![0f32; gpr];
6365        for r in start..end {
6366            v.scales_into(r, gpr, &mut sc);
6367            for bi in 0..b {
6368                let x = &xs_all[bi * cols..(bi + 1) * cols];
6369                // SAFETY: disjoint (bi, r) cells per worker range.
6370                unsafe { *out_addr.at(bi * rows + r) = q4tp_row_exact(v.nib, r, gpr, x, &sc) };
6371            }
6372        }
6373    };
6374    dispatch_rows(pool, rows, &run);
6375}
6376
6377/// Fused q4_tiled matvec (dispatch mirrors `q4matvec`).
6378fn q4t_matvec(
6379    bytes: &[u8],
6380    x: &[f32],
6381    rows: usize,
6382    cols: usize,
6383    out: &mut [f32],
6384    pool: Option<&Pool>,
6385) {
6386    debug_assert_eq!(out.len(), rows);
6387    let gpr = cols / GROUP_SIZE;
6388    let out_addr = SendMut(out.as_mut_ptr());
6389    if a8w8_enabled() {
6390        let act = split_act(x);
6391        let run = move |start: usize, end: usize| {
6392            for r in start..end {
6393                let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
6394                for &(j, xv) in &act.outliers {
6395                    let (w, s) = q4t_outlier(bytes, r, gpr, j);
6396                    acc += w * s * xv;
6397                }
6398                // SAFETY: disjoint row ranges per worker.
6399                unsafe { *out_addr.at(r) = acc };
6400            }
6401        };
6402        dispatch_rows(pool, rows, &run);
6403        return;
6404    }
6405    let run = move |start: usize, end: usize| {
6406        for r in start..end {
6407            // SAFETY: disjoint row ranges per worker.
6408            unsafe { *out_addr.at(r) = q4t_row_exact(bytes, r, gpr, x) };
6409        }
6410    };
6411    dispatch_rows(pool, rows, &run);
6412}
6413
6414/// Fused two-input q4_tiled matvec (weights read once per pair).
6415#[allow(clippy::too_many_arguments)]
6416fn q4t_matvec2(
6417    bytes: &[u8],
6418    x1: &[f32],
6419    x2: &[f32],
6420    rows: usize,
6421    cols: usize,
6422    o1: &mut [f32],
6423    o2: &mut [f32],
6424    pool: Option<&Pool>,
6425) {
6426    let gpr = cols / GROUP_SIZE;
6427    let p1 = SendMut(o1.as_mut_ptr());
6428    let p2 = SendMut(o2.as_mut_ptr());
6429    if a8w8_enabled() {
6430        let a1 = split_act(x1);
6431        let a2 = split_act(x2);
6432        let run = move |start: usize, end: usize| {
6433            for r in start..end {
6434                let mut v1 = dot_q4t_row_i8(bytes, r, gpr, &a1.xq) * a1.sx;
6435                let mut v2 = dot_q4t_row_i8(bytes, r, gpr, &a2.xq) * a2.sx;
6436                for &(j, xv) in &a1.outliers {
6437                    let (w, s) = q4t_outlier(bytes, r, gpr, j);
6438                    v1 += w * s * xv;
6439                }
6440                for &(j, xv) in &a2.outliers {
6441                    let (w, s) = q4t_outlier(bytes, r, gpr, j);
6442                    v2 += w * s * xv;
6443                }
6444                // SAFETY: disjoint row ranges per worker.
6445                unsafe {
6446                    *p1.at(r) = v1;
6447                    *p2.at(r) = v2;
6448                }
6449            }
6450        };
6451        dispatch_rows(pool, rows, &run);
6452        return;
6453    }
6454    let run = move |start: usize, end: usize| {
6455        for r in start..end {
6456            // SAFETY: disjoint row ranges per worker.
6457            unsafe {
6458                *p1.at(r) = q4t_row_exact(bytes, r, gpr, x1);
6459                *p2.at(r) = q4t_row_exact(bytes, r, gpr, x2);
6460            }
6461        }
6462    };
6463    dispatch_rows(pool, rows, &run);
6464}
6465
6466/// Batched q4_tiled matmat: each row's tiles stream once per microbatch.
6467#[allow(clippy::too_many_arguments)]
6468/// Prefill GEMM through Accelerate for group-quantized codecs: a
6469/// caller-supplied row dequantizer fills f32 tiles (pool-parallel) and
6470/// each tile rides the AMX with one sgemm — the generic sibling of
6471/// `qmatmat_accel` (q8). Numerics are f32-GEMM (tolerance class);
6472/// decode (b=1) never takes this path.
6473#[cfg(target_os = "macos")]
6474fn dequant_matmat_accel(
6475    dequant_row: &(dyn Fn(usize, &mut [f32]) + Sync),
6476    xs_all: &[f32],
6477    b: usize,
6478    rows: usize,
6479    cols: usize,
6480    out: &mut [f32],
6481    pool: Option<&Pool>,
6482) {
6483    const TR: usize = 2048;
6484    thread_local! {
6485        static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
6486    }
6487    WTILE.with(|wt| {
6488        let mut wtile = wt.borrow_mut();
6489        wtile.resize(TR * cols, 0.0);
6490        let mut r0 = 0usize;
6491        while r0 < rows {
6492            let tr = TR.min(rows - r0);
6493            let wt_addr = SendMut(wtile.as_mut_ptr());
6494            let run = |start: usize, end: usize| {
6495                for r in start..end {
6496                    // SAFETY: workers cover disjoint r ranges.
6497                    let dst = unsafe { std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols) };
6498                    dequant_row(r0 + r, dst);
6499                }
6500            };
6501            dispatch_rows(pool, tr, &run);
6502            unsafe {
6503                accel_blas::cblas_sgemm(
6504                    101, // RowMajor
6505                    111, // NoTrans A
6506                    112, // Trans B
6507                    b as i32,
6508                    tr as i32,
6509                    cols as i32,
6510                    1.0,
6511                    xs_all.as_ptr(),
6512                    cols as i32,
6513                    wtile.as_ptr(),
6514                    cols as i32,
6515                    0.0,
6516                    out.as_mut_ptr().add(r0),
6517                    rows as i32,
6518                );
6519            }
6520            r0 += tr;
6521        }
6522    });
6523}
6524
6525fn q4t_matmat(
6526    bytes: &[u8],
6527    xs_all: &[f32],
6528    b: usize,
6529    rows: usize,
6530    cols: usize,
6531    out: &mut [f32],
6532    pool: Option<&Pool>,
6533) {
6534    debug_assert_eq!(out.len(), b * rows);
6535    let gpr = cols / GROUP_SIZE;
6536    // Wide batches ride the AMX like q8's qmatmat: on Apple silicon
6537    // the dequant-tile sgemm is an order above the SDOT row loop for
6538    // prefill shapes (imagegen DiT forwards are exactly this).
6539    #[cfg(target_os = "macos")]
6540    if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
6541        dequant_matmat_accel(
6542            &|r, dst| {
6543                for gi in 0..gpr {
6544                    let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
6545                    let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
6546                    for (k, &bb) in tile[2..].iter().enumerate() {
6547                        dst[gi * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s;
6548                        dst[gi * GROUP_SIZE + k * 2 + 1] = (((bb >> 4) & 0x0F) as f32 - 8.0) * s;
6549                    }
6550                }
6551            },
6552            xs_all,
6553            b,
6554            rows,
6555            cols,
6556            out,
6557            pool,
6558        );
6559        return;
6560    }
6561    let out_addr = SendMut(out.as_mut_ptr());
6562    if a8w8_enabled() {
6563        let acts: Vec<SplitAct> = (0..b)
6564            .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
6565            .collect();
6566        let acts = &acts;
6567        #[cfg(target_arch = "x86_64")]
6568        let blocked_ok = avx2_enabled() && blocked_enabled();
6569        #[cfg(target_arch = "aarch64")]
6570        let blocked_ok = sdot_enabled() && blocked_enabled();
6571        #[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))]
6572        let blocked_ok = false;
6573        let run = move |start: usize, end: usize| {
6574            for r in start..end {
6575                let mut bi = 0usize;
6576                #[cfg(target_arch = "aarch64")]
6577                if blocked_ok {
6578                    while bi + 4 <= acts.len() {
6579                        let xs = [
6580                            acts[bi].xq.as_slice(),
6581                            acts[bi + 1].xq.as_slice(),
6582                            acts[bi + 2].xq.as_slice(),
6583                            acts[bi + 3].xq.as_slice(),
6584                        ];
6585                        let d = unsafe { dot_q4t_row_1x4_sdot(bytes, r, gpr, xs) };
6586                        for k in 0..4 {
6587                            let act = &acts[bi + k];
6588                            let mut acc = d[k] * act.sx;
6589                            for &(j, xv) in &act.outliers {
6590                                let (w, sc) = q4t_outlier(bytes, r, gpr, j);
6591                                acc += w * sc * xv;
6592                            }
6593                            // SAFETY: disjoint (bi, r) cells per worker.
6594                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
6595                        }
6596                        bi += 4;
6597                    }
6598                }
6599                #[cfg(target_arch = "x86_64")]
6600                if blocked_ok {
6601                    while bi + 4 <= acts.len() {
6602                        let xs = [
6603                            acts[bi].xq.as_slice(),
6604                            acts[bi + 1].xq.as_slice(),
6605                            acts[bi + 2].xq.as_slice(),
6606                            acts[bi + 3].xq.as_slice(),
6607                        ];
6608                        let d = unsafe {
6609                            if vnni_tiles_enabled() {
6610                                dot_q4t_row_1x4_vnni(bytes, r, gpr, xs)
6611                            } else {
6612                                dot_q4t_row_1x4_avx2(bytes, r, gpr, xs)
6613                            }
6614                        };
6615                        for k in 0..4 {
6616                            let act = &acts[bi + k];
6617                            let mut acc = d[k] * act.sx;
6618                            for &(j, xv) in &act.outliers {
6619                                let (w, sc) = q4t_outlier(bytes, r, gpr, j);
6620                                acc += w * sc * xv;
6621                            }
6622                            // SAFETY: disjoint (bi, r) cells per worker.
6623                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
6624                        }
6625                        bi += 4;
6626                    }
6627                }
6628                let _ = blocked_ok;
6629                while bi < acts.len() {
6630                    let act = &acts[bi];
6631                    let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
6632                    for &(j, xv) in &act.outliers {
6633                        let (w, s) = q4t_outlier(bytes, r, gpr, j);
6634                        acc += w * s * xv;
6635                    }
6636                    // SAFETY: disjoint (bi, r) cells per worker range.
6637                    unsafe { *out_addr.at(bi * rows + r) = acc };
6638                    bi += 1;
6639                }
6640            }
6641        };
6642        dispatch_rows(pool, rows, &run);
6643        return;
6644    }
6645    let run = move |start: usize, end: usize| {
6646        for r in start..end {
6647            for bi in 0..b {
6648                let x = &xs_all[bi * cols..(bi + 1) * cols];
6649                // SAFETY: disjoint (bi, r) cells per worker range.
6650                unsafe { *out_addr.at(bi * rows + r) = q4t_row_exact(bytes, r, gpr, x) };
6651            }
6652        }
6653    };
6654    dispatch_rows(pool, rows, &run);
6655}
6656
6657// ── q1 (dtype 12): binary weights, [f16 scale][4B sign bits] per
6658// 32-group tile. The kernel family mirrors q4_tiled: one sequential
6659// stream of 6-byte tiles, per-tile integer dot × scale, exact outlier
6660// correction (A8W8 contract), exact scalar path under CMF_SDOT=0. ──
6661
6662/// Per-32-group sums of the quantized activation — the ±1 identity's
6663/// shared half: `dot = −2·sdot(mask, x) − gsum[g]`, computed ONCE per
6664/// matvec and reused by every row.
6665fn q1_group_sums(xq: &[i8], gpr: usize) -> Vec<i32> {
6666    (0..gpr)
6667        .map(|gi| {
6668            xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE]
6669                .iter()
6670                .map(|&v| v as i32)
6671                .sum()
6672        })
6673        .collect()
6674}
6675
6676/// One q1 row via the A8W8 int8 path — mask-SDOT on ARM (no ±1
6677/// expansion at all), scalar bit loop elsewhere (AVX2 queued with the
6678/// x86 pass).
6679#[inline]
6680#[allow(unreachable_code)]
6681/// AVX2 q1 row via the same ±1 identity as the ARM sdot kernel: the
6682/// sign bits expand to a {0, −1} byte mask through shuffle+cmpeq, the
6683/// masked activation sums through maddubs(1, x&mask), and
6684/// `dot = −(2·masked_sum + Σx_group)` — bit-identical integer math.
6685#[cfg(target_arch = "x86_64")]
6686#[target_feature(enable = "avx2")]
6687unsafe fn dot_q1_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
6688    // SAFETY: callers uphold the 6B-tile and xq/gsum length contracts.
6689    unsafe {
6690        use core::arch::x86_64::*;
6691        // Byte j of the mask must replicate bits-byte j/8.
6692        let expand = _mm256_setr_epi8(
6693            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,
6694            3, 3, 3,
6695        );
6696        let bitsel = _mm256_setr_epi8(
6697            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
6698            -128, 1, 2, 4, 8, 16, 32, 64, -128,
6699        );
6700        let ones8 = _mm256_set1_epi8(1);
6701        let ones16 = _mm256_set1_epi16(1);
6702        let mut acc = 0f32;
6703        for gi in 0..gpr {
6704            let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
6705            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
6706            let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
6707            let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
6708            let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
6709            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6710            let sel = _mm256_and_si256(x, mask);
6711            // Σ of selected i8 lanes: maddubs(1u8, sel_i8) pairs → madd.
6712            let p16 = _mm256_maddubs_epi16(ones8, sel);
6713            let d32 = _mm256_madd_epi16(p16, ones16);
6714            let hi128 = _mm256_extracti128_si256::<1>(d32);
6715            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
6716            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
6717            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
6718            let msum = _mm_cvtsi128_si32(s32);
6719            // The and-select keeps x UN-negated (unlike ARM's −1-mask
6720            // sdot): d = Σ_set − Σ_unset = 2·Σ_set − Σ_all.
6721            let d = 2 * msum - gsum[gi];
6722            acc += d as f32 * s;
6723        }
6724        acc
6725    }
6726}
6727
6728/// VNNI twin of `dot_q1_row_avx2`: the masked-select sum goes through
6729/// one `vpdpbusd(1u8, sel)` (see `dpbusd_hsum` — bit-identical).
6730#[cfg(target_arch = "x86_64")]
6731#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
6732unsafe fn dot_q1_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
6733    // SAFETY: callers uphold the 6B-tile and xq/gsum length contracts.
6734    unsafe {
6735        use core::arch::x86_64::*;
6736        let expand = _mm256_setr_epi8(
6737            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,
6738            3, 3, 3,
6739        );
6740        let bitsel = _mm256_setr_epi8(
6741            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
6742            -128, 1, 2, 4, 8, 16, 32, 64, -128,
6743        );
6744        let ones8 = _mm256_set1_epi8(1);
6745        let mut acc = 0f32;
6746        for gi in 0..gpr {
6747            let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
6748            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
6749            let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
6750            let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
6751            let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
6752            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6753            let msum = dpbusd_hsum(ones8, _mm256_and_si256(x, mask));
6754            let d = 2 * msum - gsum[gi];
6755            acc += d as f32 * s;
6756        }
6757        acc
6758    }
6759}
6760
6761/// VNNI twin of `dot_q1_row_1x4_avx2` (see `dpbusd_hsum`).
6762#[cfg(target_arch = "x86_64")]
6763#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
6764unsafe fn dot_q1_row_1x4_vnni(
6765    bytes: &[u8],
6766    r: usize,
6767    gpr: usize,
6768    xs: [&[i8]; 4],
6769    gsums: [&[i32]; 4],
6770) -> [f32; 4] {
6771    // SAFETY: callers uphold the 6B-tile and xq/gsum length contracts.
6772    unsafe {
6773        use core::arch::x86_64::*;
6774        let expand = _mm256_setr_epi8(
6775            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,
6776            3, 3, 3,
6777        );
6778        let bitsel = _mm256_setr_epi8(
6779            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
6780            -128, 1, 2, 4, 8, 16, 32, 64, -128,
6781        );
6782        let ones8 = _mm256_set1_epi8(1);
6783        let mut acc = [0f32; 4];
6784        for gi in 0..gpr {
6785            let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
6786            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
6787            let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
6788            let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
6789            let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
6790            for (k, xq) in xs.iter().enumerate() {
6791                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6792                let msum = dpbusd_hsum(ones8, _mm256_and_si256(x, mask));
6793                let d = 2 * msum - gsums[k][gi];
6794                acc[k] += d as f32 * s;
6795            }
6796        }
6797        acc
6798    }
6799}
6800
6801/// The blocked 1×4 flavor: the expanded bit mask serves four activation
6802/// streams per group (mask build once, four select+reduce chains).
6803#[cfg(target_arch = "x86_64")]
6804#[target_feature(enable = "avx2")]
6805unsafe fn dot_q1_row_1x4_avx2(
6806    bytes: &[u8],
6807    r: usize,
6808    gpr: usize,
6809    xs: [&[i8]; 4],
6810    gsums: [&[i32]; 4],
6811) -> [f32; 4] {
6812    // SAFETY: callers uphold the 6B-tile and xq/gsum length contracts.
6813    unsafe {
6814        use core::arch::x86_64::*;
6815        let expand = _mm256_setr_epi8(
6816            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,
6817            3, 3, 3,
6818        );
6819        let bitsel = _mm256_setr_epi8(
6820            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
6821            -128, 1, 2, 4, 8, 16, 32, 64, -128,
6822        );
6823        let ones8 = _mm256_set1_epi8(1);
6824        let ones16 = _mm256_set1_epi16(1);
6825        let mut acc = [0f32; 4];
6826        for gi in 0..gpr {
6827            let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
6828            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
6829            let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
6830            let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
6831            let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
6832            for (k, xq) in xs.iter().enumerate() {
6833                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
6834                let sel = _mm256_and_si256(x, mask);
6835                let p16 = _mm256_maddubs_epi16(ones8, sel);
6836                let d32 = _mm256_madd_epi16(p16, ones16);
6837                let hi128 = _mm256_extracti128_si256::<1>(d32);
6838                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
6839                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
6840                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
6841                let msum = _mm_cvtsi128_si32(s32);
6842                let d = 2 * msum - gsums[k][gi];
6843                acc[k] += d as f32 * s;
6844            }
6845        }
6846        acc
6847    }
6848}
6849
6850#[allow(unreachable_code)]
6851fn dot_q1_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
6852    #[cfg(target_arch = "aarch64")]
6853    unsafe {
6854        return dot_q1_row_sdot(bytes, r, gpr, xq, gsum);
6855    }
6856    #[cfg(target_arch = "x86_64")]
6857    if avx2_enabled() {
6858        unsafe {
6859            if vnni_tiles_enabled() {
6860                return dot_q1_row_vnni(bytes, r, gpr, xq, gsum);
6861            }
6862            return dot_q1_row_avx2(bytes, r, gpr, xq, gsum);
6863        }
6864    }
6865    let _ = gsum;
6866    let mut acc = 0f32;
6867    for gi in 0..gpr {
6868        let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
6869        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
6870        let mut d = 0i32;
6871        for (j, &b) in tile[2..].iter().enumerate() {
6872            for k in 0..8 {
6873                let w = ((b >> k) & 1) as i32 * 2 - 1;
6874                d += w * xq[gi * GROUP_SIZE + j * 8 + k] as i32;
6875            }
6876        }
6877        acc += d as f32 * s;
6878    }
6879    acc
6880}
6881
6882/// SDOT q1 row via the ±1 identity: the vtst mask (0xFF where the bit
6883/// is set, i.e. −1 as i8) feeds `sdot` DIRECTLY — no expansion to ±1
6884/// lanes at all — and `dot = −(2·sdot(mask, x) + Σx_group)`, with the
6885/// per-group activation sums shared across every row of the matvec.
6886/// Four tiles (128 weights) per iteration: integer dots reduce through
6887/// a vpaddq tree into ONE i32x4 that meets its four scales in a single
6888/// fused f32 multiply-add. Integer math throughout — bit-identical to
6889/// the scalar ±1 reference.
6890#[cfg(target_arch = "aarch64")]
6891#[target_feature(enable = "neon,dotprod")]
6892unsafe fn dot_q1_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
6893    // SAFETY: callers uphold slice-length contracts (6B tile per group,
6894    // xq.len() == gpr·GROUP_SIZE, gsum.len() == gpr).
6895    unsafe {
6896        use core::arch::aarch64::*;
6897        use core::arch::asm;
6898        const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
6899        let m = vld1q_u8(MASKS.as_ptr());
6900        // One tile's −Σ_set(x) as an UNREDUCED i32x4 (two mask-sdots).
6901        macro_rules! tile_dot {
6902            ($t:expr, $x:expr) => {{
6903                let v0 = vcombine_u8(vdup_n_u8(*$t.add(2)), vdup_n_u8(*$t.add(3)));
6904                let v1 = vcombine_u8(vdup_n_u8(*$t.add(4)), vdup_n_u8(*$t.add(5)));
6905                let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
6906                let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
6907                let x0 = vld1q_s8($x);
6908                let x1 = vld1q_s8($x.add(16));
6909                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6910                asm!(
6911                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
6912                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
6913                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6914                    w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
6915                    options(pure, nomem, nostack),
6916                );
6917                vaddq_s32(a0, a1)
6918            }};
6919        }
6920        // TBL unpack over PAIR loads: one vld1q covers two 6B tiles
6921        // ([s s b b b b][s s b b b b] + 4B slack), TBL replicates each
6922        // bit-byte across 8 lanes for vtst, and the four scales gather
6923        // through tbl2 into one fcvtl — the 16 ld1r broadcast loads and
6924        // 4 branchy software f16 conversions per 128 weights (the
6925        // measured load-port wall of this kernel) become 2 vector
6926        // loads + 9 table lookups. Integer math order is unchanged —
6927        // bit-identical results (FCVTL is exact on every f16).
6928        const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
6929        const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
6930        const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
6931        const IW11: [u8; 16] = [
6932            10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
6933        ];
6934        const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
6935        let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
6936        let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
6937        let isc = vld1_u8(ISC.as_ptr());
6938        // One tile's −Σ_set(x) from a TBL-unpacked pair load.
6939        macro_rules! tile_dot_tbl {
6940            ($ld:expr, $i0:expr, $i1:expr, $x:expr) => {{
6941                let w0 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i0), m));
6942                let w1 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i1), m));
6943                let x0 = vld1q_s8($x);
6944                let x1 = vld1q_s8($x.add(16));
6945                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6946                asm!(
6947                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
6948                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
6949                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6950                    w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
6951                    options(pure, nomem, nostack),
6952                );
6953                vaddq_s32(a0, a1)
6954            }};
6955        }
6956        let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
6957        let row_base = r * gpr * Q1_TILE;
6958        let abs_end = bytes.len();
6959        let xp = xq.as_ptr();
6960        let gp = gsum.as_ptr();
6961        let mut accv = vdupq_n_f32(0.0);
6962        let mut gi = 0;
6963        // The second pair load reads 4B past tile gi+3 — stay inside
6964        // the payload slice (only the file's final tiles fall back).
6965        while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
6966            let t0 = base.add(gi * Q1_TILE);
6967            let ld_a = vld1q_u8(t0);
6968            let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
6969            let d0 = tile_dot_tbl!(ld_a, iw00, iw01, xp.add(gi * GROUP_SIZE));
6970            let d1 = tile_dot_tbl!(ld_a, iw10, iw11, xp.add((gi + 1) * GROUP_SIZE));
6971            let d2 = tile_dot_tbl!(ld_b, iw00, iw01, xp.add((gi + 2) * GROUP_SIZE));
6972            let d3 = tile_dot_tbl!(ld_b, iw10, iw11, xp.add((gi + 3) * GROUP_SIZE));
6973            // [−Σ0, −Σ1, −Σ2, −Σ3] → dots = −(2·Σset_neg + gsum)
6974            let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
6975            let g = vld1q_s32(gp.add(gi));
6976            let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
6977            let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
6978            let scf: float32x4_t;
6979            asm!(
6980                "fcvtl {o:v}.4s, {i:v}.4h",
6981                o = out(vreg) scf, i = in(vreg) sc16,
6982                options(pure, nomem, nostack),
6983            );
6984            accv = vfmaq_f32(accv, vcvtq_f32_s32(dots), scf);
6985            gi += 4;
6986        }
6987        let mut acc = vaddvq_f32(accv);
6988        while gi < gpr {
6989            let t = base.add(gi * Q1_TILE);
6990            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
6991            let d = vaddvq_s32(tile_dot!(t, xp.add(gi * GROUP_SIZE)));
6992            acc += (-(2 * d + *gp.add(gi))) as f32 * s;
6993            gi += 1;
6994        }
6995        acc
6996    }
6997}
6998
6999/// Blocked q1 1×4: one TBL unpack of the tile pair serves FOUR
7000/// activation streams (prefill amortization — the same idea as the
7001/// AVX2 twin; per stream the group order, fma order and tail match the
7002/// single-row kernel exactly, so batch == matvec bit-for-bit).
7003#[cfg(target_arch = "aarch64")]
7004#[target_feature(enable = "neon,dotprod")]
7005unsafe fn dot_q1_row_1x4_sdot(
7006    bytes: &[u8],
7007    r: usize,
7008    gpr: usize,
7009    xs: [&[i8]; 4],
7010    gs: [&[i32]; 4],
7011) -> [f32; 4] {
7012    // SAFETY: same slice-length contracts as `dot_q1_row_sdot`, ×4.
7013    unsafe {
7014        use core::arch::aarch64::*;
7015        use core::arch::asm;
7016        const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
7017        const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
7018        const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
7019        const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
7020        const IW11: [u8; 16] = [
7021            10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
7022        ];
7023        const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
7024        let m = vld1q_u8(MASKS.as_ptr());
7025        let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
7026        let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
7027        let isc = vld1_u8(ISC.as_ptr());
7028        macro_rules! sdot2 {
7029            ($w0:expr, $w1:expr, $x:expr) => {{
7030                let x0 = vld1q_s8($x);
7031                let x1 = vld1q_s8($x.add(16));
7032                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
7033                asm!(
7034                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
7035                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
7036                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
7037                    w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
7038                    options(pure, nomem, nostack),
7039                );
7040                vaddq_s32(a0, a1)
7041            }};
7042        }
7043        let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
7044        let row_base = r * gpr * Q1_TILE;
7045        let abs_end = bytes.len();
7046        let mut accv = [vdupq_n_f32(0.0); 4];
7047        let mut gi = 0;
7048        while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
7049            let t0 = base.add(gi * Q1_TILE);
7050            let ld_a = vld1q_u8(t0);
7051            let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
7052            // Unpack ONCE — eight ±mask vectors serve all four streams.
7053            let w00 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw00), m));
7054            let w01 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw01), m));
7055            let w10 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw10), m));
7056            let w11 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw11), m));
7057            let w20 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw00), m));
7058            let w21 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw01), m));
7059            let w30 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw10), m));
7060            let w31 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw11), m));
7061            let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
7062            let scf: float32x4_t;
7063            asm!(
7064                "fcvtl {o:v}.4s, {i:v}.4h",
7065                o = out(vreg) scf, i = in(vreg) sc16,
7066                options(pure, nomem, nostack),
7067            );
7068            for k in 0..4 {
7069                let xp = xs[k].as_ptr();
7070                let d0 = sdot2!(w00, w01, xp.add(gi * GROUP_SIZE));
7071                let d1 = sdot2!(w10, w11, xp.add((gi + 1) * GROUP_SIZE));
7072                let d2 = sdot2!(w20, w21, xp.add((gi + 2) * GROUP_SIZE));
7073                let d3 = sdot2!(w30, w31, xp.add((gi + 3) * GROUP_SIZE));
7074                let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
7075                let g = vld1q_s32(gs[k].as_ptr().add(gi));
7076                let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
7077                accv[k] = vfmaq_f32(accv[k], vcvtq_f32_s32(dots), scf);
7078            }
7079            gi += 4;
7080        }
7081        let mut acc = [
7082            vaddvq_f32(accv[0]),
7083            vaddvq_f32(accv[1]),
7084            vaddvq_f32(accv[2]),
7085            vaddvq_f32(accv[3]),
7086        ];
7087        while gi < gpr {
7088            let t = base.add(gi * Q1_TILE);
7089            let sc = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
7090            let v0 = vcombine_u8(vdup_n_u8(*t.add(2)), vdup_n_u8(*t.add(3)));
7091            let v1 = vcombine_u8(vdup_n_u8(*t.add(4)), vdup_n_u8(*t.add(5)));
7092            let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
7093            let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
7094            for k in 0..4 {
7095                let d = vaddvq_s32(sdot2!(w0, w1, xs[k].as_ptr().add(gi * GROUP_SIZE)));
7096                acc[k] += (-(2 * d + *gs[k].as_ptr().add(gi))) as f32 * sc;
7097            }
7098            gi += 1;
7099        }
7100        acc
7101    }
7102}
7103
7104/// (weight ±1, scale) of one q1 element — the exact outlier term.
7105#[inline]
7106fn q1_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
7107    let gi = j / GROUP_SIZE;
7108    let k = j % GROUP_SIZE;
7109    let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
7110    let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
7111    let bit = (tile[2 + k / 8] >> (k % 8)) & 1;
7112    ((bit as i32 * 2 - 1) as f32, s)
7113}
7114
7115/// Exact scalar q1 row (CMF_SDOT=0 contract).
7116#[inline]
7117fn q1_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
7118    let mut acc = 0f32;
7119    for gi in 0..gpr {
7120        let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
7121        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
7122        let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
7123        let mut ga = 0f32;
7124        for (j, &b) in tile[2..].iter().enumerate() {
7125            for k in 0..8 {
7126                ga += (((b >> k) & 1) as f32 * 2.0 - 1.0) * xg[j * 8 + k];
7127            }
7128        }
7129        acc += ga * s;
7130    }
7131    acc
7132}
7133
7134/// One q1 row range via A8W8 (the body of `q1_matvec`'s hot loop,
7135/// extracted so multi-matrix jobs drive the same kernel).
7136#[allow(clippy::too_many_arguments)]
7137fn q1_range_a8w8(
7138    bytes: &[u8],
7139    gpr: usize,
7140    act: &SplitAct,
7141    gsum: &[i32],
7142    out: SendMut,
7143    start: usize,
7144    end: usize,
7145) {
7146    for r in start..end {
7147        let mut acc = dot_q1_row_i8(bytes, r, gpr, &act.xq, gsum) * act.sx;
7148        for &(j, xv) in &act.outliers {
7149            let (w, s) = q1_outlier(bytes, r, gpr, j);
7150            acc += w * s * xv;
7151        }
7152        // SAFETY: disjoint row ranges per worker.
7153        unsafe { *out.at(r) = acc };
7154    }
7155}
7156
7157/// Exact-scalar q1 row range (CMF_SDOT=0 contract).
7158fn q1_range_f32(bytes: &[u8], gpr: usize, x: &[f32], out: SendMut, start: usize, end: usize) {
7159    for r in start..end {
7160        // SAFETY: disjoint row ranges per worker.
7161        unsafe { *out.at(r) = q1_row_exact(bytes, r, gpr, x) };
7162    }
7163}
7164
7165/// q1t per-row overlay locator. After the base (`base_len`) come
7166/// `[u32 row_ptr[rows+1]]` then `[(u16 col, f16 val)]` grouped by row (row
7167/// `r`'s entries are `[row_ptr[r], row_ptr[r+1])`). Returns
7168/// `(row_ptr offset, entries offset, present)`.
7169fn q1t_overlay(bytes: &[u8], base_len: usize, rows: usize) -> (usize, usize, bool) {
7170    let entries = base_len + (rows + 1) * 4;
7171    (base_len, entries, entries <= bytes.len())
7172}
7173
7174/// Read `row_ptr[r]` from the overlay's prefix-sum table.
7175#[inline]
7176fn q1t_rowptr(bytes: &[u8], rp_off: usize, r: usize) -> usize {
7177    let o = rp_off + r * 4;
7178    u32::from_le_bytes([bytes[o], bytes[o + 1], bytes[o + 2], bytes[o + 3]]) as usize
7179}
7180
7181/// Byte → the 5 ternary signs it packs `{−1,0,+1}` as f32, precomputed so
7182/// decoding a q1t code is a table load, not the base-3 divide/modulo per
7183/// weight (division is ~20–40× the cost of a load). Built at compile time.
7184const SIGN5: [[f32; 5]; 256] = {
7185    let mut lut = [[0.0f32; 5]; 256];
7186    let pow3 = [1u16, 3, 9, 27, 81];
7187    let mut byte = 0usize;
7188    while byte < 256 {
7189        let mut i = 0usize;
7190        while i < 5 {
7191            let code = (byte as u16 / pow3[i]) % 3;
7192            lut[byte][i] = if code == 1 {
7193                1.0
7194            } else if code == 2 {
7195                -1.0
7196            } else {
7197                0.0
7198            };
7199            i += 1;
7200        }
7201        byte += 1;
7202    }
7203    lut
7204};
7205
7206/// Same table, as i8 signs — the operand for the int8 SDOT base kernel.
7207const SIGN5_I8: [[i8; 5]; 256] = {
7208    let mut lut = [[0i8; 5]; 256];
7209    let pow3 = [1u16, 3, 9, 27, 81];
7210    let mut byte = 0usize;
7211    while byte < 256 {
7212        let mut i = 0usize;
7213        while i < 5 {
7214            let code = (byte as u16 / pow3[i]) % 3;
7215            lut[byte][i] = if code == 1 {
7216                1
7217            } else if code == 2 {
7218                -1
7219            } else {
7220                0
7221            };
7222            i += 1;
7223        }
7224        byte += 1;
7225    }
7226    lut
7227};
7228
7229/// The same 5 i8 signs packed into a u64 (`[s0 s1 s2 s3 s4 0 0 0]`, LE) so the
7230/// group unpack is 7 unaligned u64 stores at offsets 0,5,10,…,30 instead of
7231/// six 5-byte copies + LUT indexing — each store's trailing zeros are fixed by
7232/// the next store, and the last one runs 6 B past the 32nd weight (the unpack
7233/// buffer is padded to 40). This is the decode/prefill hot inner op.
7234const SIGN5_U64: [u64; 256] = {
7235    let mut lut = [0u64; 256];
7236    let pow3 = [1u16, 3, 9, 27, 81];
7237    let mut byte = 0usize;
7238    while byte < 256 {
7239        let mut v = 0u64;
7240        let mut i = 0usize;
7241        while i < 5 {
7242            let code = (byte as u16 / pow3[i]) % 3;
7243            let s: u8 = if code == 1 {
7244                1
7245            } else if code == 2 {
7246                0xFF
7247            } else {
7248                0
7249            };
7250            v |= (s as u64) << (i * 8);
7251            i += 1;
7252        }
7253        lut[byte] = v;
7254        byte += 1;
7255    }
7256    lut
7257};
7258
7259/// Ternary base weight at `(row r, col j)` = `sign(code)·s_group`. Used to add
7260/// back activation-outlier columns, whose `x` was zeroed for the int8 bulk dot
7261/// (`split_act`). At a weight-outlier position the code is 0, so this is 0 and
7262/// the overlay correction owns that column — no double counting.
7263#[inline]
7264fn q1t_base_weight(bytes: &[u8], r: usize, gpr: usize, j: usize) -> f32 {
7265    const TILE: usize = cortiq_core::quant::Q1T_TILE;
7266    let off = (r * gpr + j / GROUP_SIZE) * TILE;
7267    let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
7268    let within = j % GROUP_SIZE;
7269    SIGN5[bytes[off + 2 + within / 5] as usize][within % 5] * s
7270}
7271
7272/// One 32-group int8 dot via two SDOTs. Bit-exact vs the scalar i8 sum
7273/// (integer accumulation is order-independent).
7274#[cfg(target_arch = "aarch64")]
7275#[target_feature(enable = "neon,dotprod")]
7276#[inline]
7277unsafe fn sdot32_i8(w: *const i8, x: *const i8) -> i32 {
7278    // SAFETY: caller guarantees 32 readable i8 at each pointer.
7279    unsafe {
7280        use core::arch::aarch64::*;
7281        use core::arch::asm;
7282        let w0 = vld1q_s8(w);
7283        let w1 = vld1q_s8(w.add(16));
7284        let x0 = vld1q_s8(x);
7285        let x1 = vld1q_s8(x.add(16));
7286        let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
7287        asm!(
7288            "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
7289            "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
7290            a0 = inout(vreg) a0, a1 = inout(vreg) a1,
7291            w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
7292            options(pure, nomem, nostack),
7293        );
7294        vaddvq_s32(vaddq_s32(a0, a1))
7295    }
7296}
7297
7298/// One 32-group int8 dot via AVX2: signed·signed as `maddubs(|w|, sign(x,w))`
7299/// then `madd` and a horizontal reduce (the same idiom as `dot_q4t_row_avx2`).
7300#[cfg(target_arch = "x86_64")]
7301#[target_feature(enable = "avx2")]
7302#[inline]
7303unsafe fn i8dot32_avx2(w: *const i8, x: *const i8) -> i32 {
7304    // SAFETY: caller guarantees 32 readable i8 at each pointer.
7305    unsafe {
7306        use core::arch::x86_64::*;
7307        let wv = _mm256_loadu_si256(w as *const __m256i);
7308        let xv = _mm256_loadu_si256(x as *const __m256i);
7309        let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
7310        let d = _mm256_madd_epi16(p16, _mm256_set1_epi16(1));
7311        let hi128 = _mm256_extracti128_si256::<1>(d);
7312        let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
7313        let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7314        let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7315        _mm_cvtsi128_si32(s32)
7316    }
7317}
7318
7319/// Unpack one q1t group's base-3 codes into 32 i8 signs via 7 unaligned u64
7320/// stores (see `SIGN5_U64`). `dst` MUST have ≥ 40 bytes: the 7th store writes
7321/// `dst[30..38]`. Stores go in order so each one's trailing zeros are
7322/// overwritten by the next; the final 6 padding bytes are unused by the dot.
7323#[inline]
7324fn q1t_unpack_group_i8(codes: *const u8, dst: &mut [i8]) {
7325    debug_assert!(dst.len() >= 40);
7326    // SAFETY: codes points at 7 readable bytes; dst has ≥ 40 bytes so every
7327    // 8-byte store at offset bi*5 (bi ≤ 6 → ≤ 30) stays in bounds.
7328    unsafe {
7329        let p = dst.as_mut_ptr();
7330        for bi in 0..7 {
7331            core::ptr::write_unaligned(
7332                p.add(bi * 5) as *mut u64,
7333                SIGN5_U64[*codes.add(bi) as usize],
7334            );
7335        }
7336    }
7337}
7338
7339/// One 32-group int8 dot, arch-dispatched (the matmat inner loop, where the
7340/// row's signs are unpacked once and dotted against every batch input).
7341/// Callers are gated by `a8w8_enabled()`, so the target-feature arms are
7342/// reachable; the scalar arm is a non-SIMD-arch fallback.
7343#[inline]
7344fn q1t_i8dot32(w: *const i8, x: *const i8) -> i32 {
7345    #[cfg(target_arch = "aarch64")]
7346    unsafe {
7347        return sdot32_i8(w, x);
7348    }
7349    #[cfg(target_arch = "x86_64")]
7350    unsafe {
7351        return i8dot32_avx2(w, x);
7352    }
7353    #[allow(unreachable_code)]
7354    unsafe {
7355        let mut s = 0i32;
7356        for k in 0..GROUP_SIZE {
7357            s += *w.add(k) as i32 * *x.add(k) as i32;
7358        }
7359        s
7360    }
7361}
7362
7363#[inline]
7364unsafe fn q1t_unpack_reg_u64s(codes: *const u8) -> (u64, u64, u64, u64) {
7365    let (s0, s1, s2, s3, s4, s5, s6) = unsafe {
7366        (
7367            SIGN5_U64[*codes as usize],
7368            SIGN5_U64[*codes.add(1) as usize],
7369            SIGN5_U64[*codes.add(2) as usize],
7370            SIGN5_U64[*codes.add(3) as usize],
7371            SIGN5_U64[*codes.add(4) as usize],
7372            SIGN5_U64[*codes.add(5) as usize],
7373            SIGN5_U64[*codes.add(6) as usize],
7374        )
7375    };
7376
7377    let u0 = s0 | (s1 << 40);
7378    let u1 = (s1 >> 24) | (s2 << 16) | (s3 << 56);
7379    let u2 = (s3 >> 8) | (s4 << 32);
7380    let u3 = (s4 >> 32) | (s5 << 8) | (s6 << 48);
7381
7382    (u0, u1, u2, u3)
7383}
7384
7385/// One q1t row's int8 base dot: `Σ_group s·dot(signs, xq)` (before the shared
7386/// `sx`). Direct register unpacking (zero stack stores/loads, no STLF stalls).
7387/// ARM SDOT.
7388#[cfg(target_arch = "aarch64")]
7389#[target_feature(enable = "neon,dotprod")]
7390unsafe fn q1t_dot_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
7391    use core::arch::aarch64::*;
7392    use core::arch::asm;
7393    unsafe {
7394        const TILE: usize = cortiq_core::quant::Q1T_TILE;
7395        let mut acc = 0f32;
7396        let bytes_ptr = bytes.as_ptr();
7397        let xq_ptr = xq.as_ptr();
7398        let row_off = r * gpr * TILE;
7399
7400        let gpr2 = gpr & !1;
7401        let mut gi = 0;
7402        while gi < gpr2 {
7403            let off0 = row_off + gi * TILE;
7404            let off1 = off0 + TILE;
7405            let s0 = f16_to_f32(u16::from_le_bytes([
7406                *bytes_ptr.add(off0),
7407                *bytes_ptr.add(off0 + 1),
7408            ]));
7409            let s1 = f16_to_f32(u16::from_le_bytes([
7410                *bytes_ptr.add(off1),
7411                *bytes_ptr.add(off1 + 1),
7412            ]));
7413
7414            let (u0_0, u1_0, u2_0, u3_0) = q1t_unpack_reg_u64s(bytes_ptr.add(off0 + 2));
7415            let (u0_1, u1_1, u2_1, u3_1) = q1t_unpack_reg_u64s(bytes_ptr.add(off1 + 2));
7416
7417            let w0_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_0), vcreate_u64(u1_0)));
7418            let w1_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_0), vcreate_u64(u3_0)));
7419            let w0_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_1), vcreate_u64(u1_1)));
7420            let w1_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_1), vcreate_u64(u3_1)));
7421
7422            let x0_0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE));
7423            let x1_0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE + 16));
7424            let x0_1 = vld1q_s8(xq_ptr.add((gi + 1) * GROUP_SIZE));
7425            let x1_1 = vld1q_s8(xq_ptr.add((gi + 1) * GROUP_SIZE + 16));
7426
7427            let (mut a0_0, mut a1_0) = (vdupq_n_s32(0), vdupq_n_s32(0));
7428            let (mut a0_1, mut a1_1) = (vdupq_n_s32(0), vdupq_n_s32(0));
7429            asm!(
7430                "sdot {a0_0:v}.4s, {w0_0:v}.16b, {x0_0:v}.16b",
7431                "sdot {a1_0:v}.4s, {w1_0:v}.16b, {x1_0:v}.16b",
7432                "sdot {a0_1:v}.4s, {w0_1:v}.16b, {x0_1:v}.16b",
7433                "sdot {a1_1:v}.4s, {w1_1:v}.16b, {x1_1:v}.16b",
7434                a0_0 = inout(vreg) a0_0, a1_0 = inout(vreg) a1_0,
7435                a0_1 = inout(vreg) a0_1, a1_1 = inout(vreg) a1_1,
7436                w0_0 = in(vreg) w0_0, x0_0 = in(vreg) x0_0, w1_0 = in(vreg) w1_0, x1_0 = in(vreg) x1_0,
7437                w0_1 = in(vreg) w0_1, x0_1 = in(vreg) x0_1, w1_1 = in(vreg) w1_1, x1_1 = in(vreg) x1_1,
7438                options(pure, nomem, nostack),
7439            );
7440            let d0 = vaddvq_s32(vaddq_s32(a0_0, a1_0));
7441            let d1 = vaddvq_s32(vaddq_s32(a0_1, a1_1));
7442            acc += d0 as f32 * s0 + d1 as f32 * s1;
7443            gi += 2;
7444        }
7445
7446        if gi < gpr {
7447            let off = row_off + gi * TILE;
7448            let s = f16_to_f32(u16::from_le_bytes([
7449                *bytes_ptr.add(off),
7450                *bytes_ptr.add(off + 1),
7451            ]));
7452            let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
7453            let w0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0), vcreate_u64(u1)));
7454            let w1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2), vcreate_u64(u3)));
7455            let x0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE));
7456            let x1 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE + 16));
7457            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
7458            asm!(
7459                "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
7460                "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
7461                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
7462                w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
7463                options(pure, nomem, nostack),
7464            );
7465            let d = vaddvq_s32(vaddq_s32(a0, a1));
7466            acc += d as f32 * s;
7467        }
7468        acc
7469    }
7470}
7471
7472/// x86 AVX2 mirror of `q1t_dot_row_sdot` (maddubs int8 dot per group).
7473#[cfg(target_arch = "x86_64")]
7474#[target_feature(enable = "avx2")]
7475unsafe fn q1t_dot_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
7476    use core::arch::x86_64::*;
7477    unsafe {
7478        const TILE: usize = cortiq_core::quant::Q1T_TILE;
7479        let mut acc = 0f32;
7480        let bytes_ptr = bytes.as_ptr();
7481        let xq_ptr = xq.as_ptr();
7482        let row_off = r * gpr * TILE;
7483
7484        let ones = _mm256_set1_epi16(1);
7485        for gi in 0..gpr {
7486            let off = row_off + gi * TILE;
7487            let s = f16_to_f32(u16::from_le_bytes([
7488                *bytes_ptr.add(off),
7489                *bytes_ptr.add(off + 1),
7490            ]));
7491            let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
7492            let wv = _mm256_set_epi64x(u3 as i64, u2 as i64, u1 as i64, u0 as i64);
7493            let xv = _mm256_loadu_si256(xq_ptr.add(gi * GROUP_SIZE) as *const __m256i);
7494            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
7495            let d256 = _mm256_madd_epi16(p16, ones);
7496            let d128 = _mm_add_epi32(
7497                _mm256_castsi256_si128(d256),
7498                _mm256_extracti128_si256(d256, 1),
7499            );
7500            let d64 = _mm_add_epi32(d128, _mm_shuffle_epi32(d128, 0xee));
7501            let d32 = _mm_cvtsi128_si32(_mm_add_epi32(d64, _mm_shuffle_epi32(d64, 0x55)));
7502            acc += d32 as f32 * s;
7503        }
7504        acc
7505    }
7506}
7507
7508/// VNNI twin of `q1t_dot_row_avx2` (see `dpbusd_hsum`).
7509#[cfg(target_arch = "x86_64")]
7510#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
7511unsafe fn q1t_dot_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
7512    use core::arch::x86_64::*;
7513    // SAFETY: same tile/xq contracts as `q1t_dot_row_avx2`.
7514    unsafe {
7515        const TILE: usize = cortiq_core::quant::Q1T_TILE;
7516        let mut acc = 0f32;
7517        let bytes_ptr = bytes.as_ptr();
7518        let xq_ptr = xq.as_ptr();
7519        let row_off = r * gpr * TILE;
7520        for gi in 0..gpr {
7521            let off = row_off + gi * TILE;
7522            let s = f16_to_f32(u16::from_le_bytes([
7523                *bytes_ptr.add(off),
7524                *bytes_ptr.add(off + 1),
7525            ]));
7526            let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
7527            let wv = _mm256_set_epi64x(u3 as i64, u2 as i64, u1 as i64, u0 as i64);
7528            let xv = _mm256_loadu_si256(xq_ptr.add(gi * GROUP_SIZE) as *const __m256i);
7529            let d = dpbusd_hsum(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
7530            acc += d as f32 * s;
7531        }
7532        acc
7533    }
7534}
7535
7536/// Per-row int8 base dot, dispatched once per row (matvec decode hot path).
7537/// Callers are gated by `a8w8_enabled()`, so the target-feature kernels are
7538/// reachable.
7539#[inline]
7540fn q1t_dot_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
7541    #[cfg(target_arch = "aarch64")]
7542    unsafe {
7543        return q1t_dot_row_sdot(bytes, r, gpr, xq);
7544    }
7545    #[cfg(target_arch = "x86_64")]
7546    unsafe {
7547        if vnni_tiles_enabled() {
7548            return q1t_dot_row_vnni(bytes, r, gpr, xq);
7549        }
7550        return q1t_dot_row_avx2(bytes, r, gpr, xq);
7551    }
7552    #[allow(unreachable_code)]
7553    {
7554        const TILE: usize = cortiq_core::quant::Q1T_TILE;
7555        let mut acc = 0f32;
7556        let mut sg = [0i8; GROUP_SIZE + 8]; // +8 slack for the u64-store unpack
7557        for gi in 0..gpr {
7558            let off = (r * gpr + gi) * TILE;
7559            let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
7560            q1t_unpack_group_i8(bytes.as_ptr().wrapping_add(off + 2), &mut sg);
7561            let mut d = 0i32;
7562            for k in 0..GROUP_SIZE {
7563                d += sg[k] as i32 * xq[gi * GROUP_SIZE + k] as i32;
7564            }
7565            acc += d as f32 * s;
7566        }
7567        acc
7568    }
7569}
7570
7571/// Σ over a row's outliers of `value·x[col]` — the correction that adds the
7572/// overlay's exact weights on top of the base dot. INVARIANT: the encoder
7573/// writes ternary code 0 at every outlier position (`quantize_q1t`), so the
7574/// base contributes nothing there and this is a plain `value·x`, not
7575/// `(value − base)·x` — no scattered per-outlier scale read. Row `r`'s entries
7576/// are the contiguous slice `[row_ptr[r], row_ptr[r+1])`, so no binary search.
7577fn q1t_row_outlier_correction(
7578    bytes: &[u8],
7579    r: usize,
7580    rp_off: usize,
7581    entries_off: usize,
7582    has_ov: bool,
7583    x: &[f32],
7584) -> f32 {
7585    if !has_ov {
7586        return 0.0;
7587    }
7588    let (c0, c1) = (
7589        q1t_rowptr(bytes, rp_off, r),
7590        q1t_rowptr(bytes, rp_off, r + 1),
7591    );
7592    let mut corr = 0f32;
7593    for p in c0..c1 {
7594        let e = entries_off + p * 4;
7595        let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
7596        let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
7597        corr += val * x[col];
7598    }
7599    corr
7600}
7601
7602/// Dequantize one q1t row into `buf[..cols]` via the sign LUT (no division),
7603/// then apply the row's outliers (its `[row_ptr[r], row_ptr[r+1])` slice).
7604/// Used by the batched (prefill) path where the decode amortizes over the batch.
7605fn q1t_dequant_row(
7606    bytes: &[u8],
7607    r: usize,
7608    gpr: usize,
7609    rp_off: usize,
7610    entries_off: usize,
7611    has_ov: bool,
7612    buf: &mut [f32],
7613) {
7614    const TILE: usize = cortiq_core::quant::Q1T_TILE;
7615    for g in 0..gpr {
7616        let off = (r * gpr + g) * TILE;
7617        let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
7618        let codes = &bytes[off + 2..off + TILE];
7619        let bc = g * GROUP_SIZE;
7620        // 6 full bytes (30 codes) + a 7th byte holding the last 2.
7621        for bi in 0..6 {
7622            let lut = &SIGN5[codes[bi] as usize];
7623            let d = &mut buf[bc + bi * 5..bc + bi * 5 + 5];
7624            for i in 0..5 {
7625                d[i] = lut[i] * s;
7626            }
7627        }
7628        let lut = &SIGN5[codes[6] as usize];
7629        buf[bc + 30] = lut[0] * s;
7630        buf[bc + 31] = lut[1] * s;
7631    }
7632    if !has_ov {
7633        return;
7634    }
7635    let (c0, c1) = (
7636        q1t_rowptr(bytes, rp_off, r),
7637        q1t_rowptr(bytes, rp_off, r + 1),
7638    );
7639    for p in c0..c1 {
7640        let e = entries_off + p * 4;
7641        let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
7642        buf[col] = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
7643    }
7644}
7645
7646/// Add the sparse outlier overlay onto a base dot already in `out` (the GPU
7647/// computes the ternary base; the overlay stays on the CPU — its entries are
7648/// few and its per-row gather doesn't vectorize on the GPU). Row-parallel.
7649fn q1t_add_overlay(
7650    bytes: &[u8],
7651    x: &[f32],
7652    rows: usize,
7653    cols: usize,
7654    out: &mut [f32],
7655    pool: Option<&Pool>,
7656) {
7657    const TILE: usize = cortiq_core::quant::Q1T_TILE;
7658    let gpr = cols / GROUP_SIZE;
7659    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
7660    if !has_ov {
7661        return;
7662    }
7663    let out_addr = SendMut(out.as_mut_ptr());
7664    let run = move |start: usize, end: usize| {
7665        for r in start..end {
7666            let corr = q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
7667            // SAFETY: disjoint rows; add onto the base the GPU already wrote.
7668            unsafe { *out_addr.at(r) += corr };
7669        }
7670    };
7671    dispatch_rows(pool, rows, &run);
7672}
7673
7674/// Q1T row range via the A8W8 int8 path — shared activation split,
7675/// per-row: base SDOT dot + outlier correction + overlay.
7676#[allow(clippy::too_many_arguments)]
7677fn q1t_range_a8w8(
7678    bytes: &[u8],
7679    gpr: usize,
7680    rp_off: usize,
7681    ent_off: usize,
7682    has_ov: bool,
7683    act: &SplitAct,
7684    x: &[f32],
7685    out: SendMut,
7686    start: usize,
7687    end: usize,
7688) {
7689    for r in start..end {
7690        let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
7691        for &(j, xv) in &act.outliers {
7692            acc += q1t_base_weight(bytes, r, gpr, j) * xv;
7693        }
7694        acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
7695        // SAFETY: disjoint row ranges per worker.
7696        unsafe { *out.at(r) = acc };
7697    }
7698}
7699
7700/// Q1T row range via the f32 path (no SDOT) — for matvec_many batched
7701/// dispatch when a8w8 is unavailable.
7702#[allow(clippy::too_many_arguments)]
7703fn q1t_range_f32_batch(
7704    bytes: &[u8],
7705    gpr: usize,
7706    rp_off: usize,
7707    ent_off: usize,
7708    has_ov: bool,
7709    x: &[f32],
7710    out: SendMut,
7711    start: usize,
7712    end: usize,
7713) {
7714    const TILE: usize = cortiq_core::quant::Q1T_TILE;
7715    let mut sg = [0f32; GROUP_SIZE];
7716    for r in start..end {
7717        let mut acc = 0f32;
7718        for g in 0..gpr {
7719            let off = (r * gpr + g) * TILE;
7720            let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
7721            let codes = &bytes[off + 2..off + TILE];
7722            let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
7723            for bi in 0..6 {
7724                sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
7725            }
7726            let lut = &SIGN5[codes[6] as usize];
7727            sg[30] = lut[0];
7728            sg[31] = lut[1];
7729            let mut gsum = 0f32;
7730            for k in 0..GROUP_SIZE {
7731                gsum += sg[k] * xg[k];
7732            }
7733            acc += s * gsum;
7734        }
7735        acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
7736        // SAFETY: disjoint row ranges per worker.
7737        unsafe { *out.at(r) = acc };
7738    }
7739}
7740
7741/// Ternary (q1t) matvec — decode+dot straight from mmap, one group at a time:
7742/// no per-ROW buffer, no division (the sign LUT), and a tiny per-group sign
7743/// buffer so the 32-wide dot vectorizes. This is the decode hot path.
7744fn q1t_matvec(
7745    bytes: &[u8],
7746    x: &[f32],
7747    rows: usize,
7748    cols: usize,
7749    out: &mut [f32],
7750    pool: Option<&Pool>,
7751) {
7752    debug_assert_eq!(out.len(), rows);
7753    const TILE: usize = cortiq_core::quant::Q1T_TILE;
7754    let gpr = cols / GROUP_SIZE;
7755    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
7756    let out_addr = SendMut(out.as_mut_ptr());
7757    // int8 SDOT base dot (ARM dotprod): ~4× the f32 arithmetic. x → i8 once
7758    // (`split_act`), activation outliers added back exactly in f32, weight
7759    // overlay on top. ARM SDOT / x86 AVX2; CMF_SDOT=0 keeps the exact f32 path.
7760    if a8w8_enabled() {
7761        let act = split_act(x);
7762        let act = &act;
7763        let run = move |start: usize, end: usize| {
7764            for r in start..end {
7765                let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
7766                for &(j, xv) in &act.outliers {
7767                    acc += q1t_base_weight(bytes, r, gpr, j) * xv;
7768                }
7769                acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
7770                // SAFETY: disjoint row ranges per worker.
7771                unsafe { *out_addr.at(r) = acc };
7772            }
7773        };
7774        dispatch_rows(pool, rows, &run);
7775        return;
7776    }
7777    let run = move |start: usize, end: usize| {
7778        // Per-group signs, unpacked contiguously so the dot below is a clean
7779        // 32-wide reduction the autovectorizer turns into f32x4 FMAs — the
7780        // 5-values-per-byte base-3 layout won't SIMD in place.
7781        let mut sg = [0f32; GROUP_SIZE];
7782        for r in start..end {
7783            let mut acc = 0f32;
7784            for g in 0..gpr {
7785                let off = (r * gpr + g) * TILE;
7786                let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
7787                let codes = &bytes[off + 2..off + TILE];
7788                let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
7789                for bi in 0..6 {
7790                    sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
7791                }
7792                let lut = &SIGN5[codes[6] as usize];
7793                sg[30] = lut[0];
7794                sg[31] = lut[1];
7795                let mut gsum = 0f32;
7796                for k in 0..GROUP_SIZE {
7797                    gsum += sg[k] * xg[k];
7798                }
7799                acc += s * gsum;
7800            }
7801            acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
7802            unsafe { *out_addr.at(r) = acc };
7803        }
7804    };
7805    dispatch_rows(pool, rows, &run);
7806}
7807
7808/// Fused-pair twin of `q1t_dot_row_sdot`: ONE register unpack of the
7809/// ternary codes serves BOTH activation streams (the unpack chain is
7810/// the dominant per-row cost — MTP verify pairs paid it twice). Per
7811/// stream the group order and f32 accumulation match the single-row
7812/// kernel exactly, so pair == 2×matvec bit-for-bit.
7813#[cfg(target_arch = "aarch64")]
7814#[target_feature(enable = "neon,dotprod")]
7815unsafe fn q1t_dot_row_sdot2(bytes: &[u8], r: usize, gpr: usize, xa: &[i8], xb: &[i8]) -> [f32; 2] {
7816    use core::arch::aarch64::*;
7817    use core::arch::asm;
7818    // SAFETY: same slice-length contracts as `q1t_dot_row_sdot`, ×2.
7819    unsafe {
7820        const TILE: usize = cortiq_core::quant::Q1T_TILE;
7821        let bytes_ptr = bytes.as_ptr();
7822        let row_off = r * gpr * TILE;
7823        let xp = [xa.as_ptr(), xb.as_ptr()];
7824        let mut acc = [0f32; 2];
7825        macro_rules! sdot2 {
7826            ($w0:expr, $w1:expr, $x:expr) => {{
7827                let x0 = vld1q_s8($x);
7828                let x1 = vld1q_s8($x.add(16));
7829                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
7830                asm!(
7831                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
7832                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
7833                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
7834                    w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
7835                    options(pure, nomem, nostack),
7836                );
7837                vaddvq_s32(vaddq_s32(a0, a1))
7838            }};
7839        }
7840        let gpr2 = gpr & !1;
7841        let mut gi = 0;
7842        while gi < gpr2 {
7843            let off0 = row_off + gi * TILE;
7844            let off1 = off0 + TILE;
7845            let s0 = f16_to_f32(u16::from_le_bytes([
7846                *bytes_ptr.add(off0),
7847                *bytes_ptr.add(off0 + 1),
7848            ]));
7849            let s1 = f16_to_f32(u16::from_le_bytes([
7850                *bytes_ptr.add(off1),
7851                *bytes_ptr.add(off1 + 1),
7852            ]));
7853            let (u0_0, u1_0, u2_0, u3_0) = q1t_unpack_reg_u64s(bytes_ptr.add(off0 + 2));
7854            let (u0_1, u1_1, u2_1, u3_1) = q1t_unpack_reg_u64s(bytes_ptr.add(off1 + 2));
7855            let w0_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_0), vcreate_u64(u1_0)));
7856            let w1_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_0), vcreate_u64(u3_0)));
7857            let w0_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_1), vcreate_u64(u1_1)));
7858            let w1_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_1), vcreate_u64(u3_1)));
7859            for k in 0..2 {
7860                let d0 = sdot2!(w0_0, w1_0, xp[k].add(gi * GROUP_SIZE));
7861                let d1 = sdot2!(w0_1, w1_1, xp[k].add((gi + 1) * GROUP_SIZE));
7862                acc[k] += d0 as f32 * s0 + d1 as f32 * s1;
7863            }
7864            gi += 2;
7865        }
7866        if gi < gpr {
7867            let off = row_off + gi * TILE;
7868            let s = f16_to_f32(u16::from_le_bytes([
7869                *bytes_ptr.add(off),
7870                *bytes_ptr.add(off + 1),
7871            ]));
7872            let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
7873            let w0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0), vcreate_u64(u1)));
7874            let w1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2), vcreate_u64(u3)));
7875            for k in 0..2 {
7876                let d = sdot2!(w0, w1, xp[k].add(gi * GROUP_SIZE));
7877                acc[k] += d as f32 * s;
7878            }
7879        }
7880        acc
7881    }
7882}
7883
7884/// Fused Q1T pair matvec: ONE pass over the rows serves both
7885/// activation streams — on ARM the ternary register unpack happens
7886/// once per tile pair (`q1t_dot_row_sdot2`); elsewhere the second dot
7887/// rides the row's L1-warm tile bytes. Per stream the math matches
7888/// `q1t_matvec` exactly.
7889fn q1t_matvec2(
7890    bytes: &[u8],
7891    x1: &[f32],
7892    x2: &[f32],
7893    rows: usize,
7894    cols: usize,
7895    o1: &mut [f32],
7896    o2: &mut [f32],
7897    pool: Option<&Pool>,
7898) {
7899    debug_assert_eq!(o1.len(), rows);
7900    debug_assert_eq!(o2.len(), rows);
7901    const TILE: usize = cortiq_core::quant::Q1T_TILE;
7902    let gpr = cols / GROUP_SIZE;
7903    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
7904    let out1 = SendMut(o1.as_mut_ptr());
7905    let out2 = SendMut(o2.as_mut_ptr());
7906    if a8w8_enabled() {
7907        let a1 = split_act(x1);
7908        let a2 = split_act(x2);
7909        let (a1, a2) = (&a1, &a2);
7910        let run = move |start: usize, end: usize| {
7911            for r in start..end {
7912                #[cfg(target_arch = "aarch64")]
7913                // a8w8 on aarch64 ⇔ sdot_enabled(), so the kernel's
7914                // target features are present.
7915                let ds = unsafe { q1t_dot_row_sdot2(bytes, r, gpr, &a1.xq, &a2.xq) };
7916                #[cfg(not(target_arch = "aarch64"))]
7917                let ds = [
7918                    q1t_dot_row_i8(bytes, r, gpr, &a1.xq),
7919                    q1t_dot_row_i8(bytes, r, gpr, &a2.xq),
7920                ];
7921                let mut acc1 = ds[0] * a1.sx;
7922                for &(j, xv) in &a1.outliers {
7923                    acc1 += q1t_base_weight(bytes, r, gpr, j) * xv;
7924                }
7925                acc1 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x1);
7926                let mut acc2 = ds[1] * a2.sx;
7927                for &(j, xv) in &a2.outliers {
7928                    acc2 += q1t_base_weight(bytes, r, gpr, j) * xv;
7929                }
7930                acc2 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x2);
7931                // SAFETY: disjoint row ranges per worker.
7932                unsafe {
7933                    *out1.at(r) = acc1;
7934                    *out2.at(r) = acc2;
7935                }
7936            }
7937        };
7938        dispatch_rows(pool, rows, &run);
7939        return;
7940    }
7941    let run = move |start: usize, end: usize| {
7942        // Exact path (CMF_SDOT=0): unpack the sign LUT once per group,
7943        // dot both streams — same op order per stream as `q1t_matvec`.
7944        let mut sg = [0f32; GROUP_SIZE];
7945        for r in start..end {
7946            let mut acc1 = 0f32;
7947            let mut acc2 = 0f32;
7948            for g in 0..gpr {
7949                let off = (r * gpr + g) * TILE;
7950                let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
7951                let codes = &bytes[off + 2..off + TILE];
7952                for bi in 0..6 {
7953                    sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
7954                }
7955                let lut = &SIGN5[codes[6] as usize];
7956                sg[30] = lut[0];
7957                sg[31] = lut[1];
7958                let xg1 = &x1[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
7959                let xg2 = &x2[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
7960                let mut gsum1 = 0f32;
7961                for k in 0..GROUP_SIZE {
7962                    gsum1 += sg[k] * xg1[k];
7963                }
7964                acc1 += s * gsum1;
7965                let mut gsum2 = 0f32;
7966                for k in 0..GROUP_SIZE {
7967                    gsum2 += sg[k] * xg2[k];
7968                }
7969                acc2 += s * gsum2;
7970            }
7971            acc1 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x1);
7972            acc2 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x2);
7973            // SAFETY: disjoint row ranges per worker.
7974            unsafe {
7975                *out1.at(r) = acc1;
7976                *out2.at(r) = acc2;
7977            }
7978        }
7979    };
7980    dispatch_rows(pool, rows, &run);
7981}
7982
7983/// Ternary (q1t) matmat (prefill) — dequant each row once, dot the whole
7984/// batch against it (amortizes the per-row decode).
7985fn q1t_matmat(
7986    bytes: &[u8],
7987    xs: &[f32],
7988    b: usize,
7989    rows: usize,
7990    cols: usize,
7991    out: &mut [f32],
7992    pool: Option<&Pool>,
7993) {
7994    debug_assert_eq!(out.len(), b * rows);
7995    const TILE: usize = cortiq_core::quant::Q1T_TILE;
7996    let gpr = cols / GROUP_SIZE;
7997    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
7998    let out_addr = SendMut(out.as_mut_ptr());
7999    // int8 prefill (ARM SDOT / x86 AVX2): quantize the B inputs once, unpack
8000    // each weight row's signs to i8 ONCE, then int8-dot against every input —
8001    // the row sign-decode amortizes over the whole batch. CMF_SDOT=0 → f32.
8002    if a8w8_enabled() {
8003        let acts: Vec<SplitAct> = (0..b)
8004            .map(|bi| split_act(&xs[bi * cols..(bi + 1) * cols]))
8005            .collect();
8006        let acts = &acts;
8007        let run = move |start: usize, end: usize| {
8008            let mut sg = vec![0i8; cols + 8]; // row signs, i8 (+8 unpack slack)
8009            let mut sc = vec![0f32; gpr]; // per-group scales
8010            let mut accs = vec![0f32; b]; // per-batch accumulators, reused per row
8011            for r in start..end {
8012                for g in 0..gpr {
8013                    let off = (r * gpr + g) * TILE;
8014                    sc[g] = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
8015                    q1t_unpack_group_i8(
8016                        bytes.as_ptr().wrapping_add(off + 2),
8017                        &mut sg[g * GROUP_SIZE..],
8018                    );
8019                }
8020                for bi in 0..b {
8021                    let act = &acts[bi];
8022                    let mut isum = 0f32;
8023                    for g in 0..gpr {
8024                        let d = q1t_i8dot32(
8025                            sg.as_ptr().wrapping_add(g * GROUP_SIZE),
8026                            act.xq.as_ptr().wrapping_add(g * GROUP_SIZE),
8027                        );
8028                        isum += d as f32 * sc[g];
8029                    }
8030                    let mut acc = isum * act.sx;
8031                    for &(j, xv) in &act.outliers {
8032                        acc += q1t_base_weight(bytes, r, gpr, j) * xv;
8033                    }
8034                    accs[bi] = acc;
8035                }
8036                // Overlay ONCE per row for the whole batch: read each (col, val)
8037                // from mmap a single time (was b× — the re-read dominated prefill)
8038                // and fan it out over the batch via the cached inputs.
8039                if has_ov {
8040                    let (c0, c1) = (
8041                        q1t_rowptr(bytes, rp_off, r),
8042                        q1t_rowptr(bytes, rp_off, r + 1),
8043                    );
8044                    for p in c0..c1 {
8045                        let e = ent_off + p * 4;
8046                        let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
8047                        let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
8048                        for bi in 0..b {
8049                            accs[bi] += val * xs[bi * cols + col];
8050                        }
8051                    }
8052                }
8053                for bi in 0..b {
8054                    unsafe { *out_addr.at(bi * rows + r) = accs[bi] };
8055                }
8056            }
8057        };
8058        dispatch_rows(pool, rows, &run);
8059        return;
8060    }
8061    let run = move |start: usize, end: usize| {
8062        let mut buf = vec![0f32; cols];
8063        for r in start..end {
8064            q1t_dequant_row(bytes, r, gpr, rp_off, ent_off, has_ov, &mut buf);
8065            for bi in 0..b {
8066                let xr = &xs[bi * cols..(bi + 1) * cols];
8067                let mut acc = 0f32;
8068                for j in 0..cols {
8069                    acc += buf[j] * xr[j];
8070                }
8071                unsafe { *out_addr.at(bi * rows + r) = acc };
8072            }
8073        }
8074    };
8075    dispatch_rows(pool, rows, &run);
8076}
8077
8078fn q1_matvec(
8079    bytes: &[u8],
8080    x: &[f32],
8081    rows: usize,
8082    cols: usize,
8083    out: &mut [f32],
8084    pool: Option<&Pool>,
8085) {
8086    debug_assert_eq!(out.len(), rows);
8087    let gpr = cols / GROUP_SIZE;
8088    let out_addr = SendMut(out.as_mut_ptr());
8089    if a8w8_enabled() {
8090        let act = split_act(x);
8091        let gsum = q1_group_sums(&act.xq, gpr);
8092        let (act, gsum) = (&act, &gsum);
8093        let run = move |start: usize, end: usize| {
8094            q1_range_a8w8(bytes, gpr, act, gsum, out_addr, start, end)
8095        };
8096        dispatch_rows(pool, rows, &run);
8097        return;
8098    }
8099    let run = move |start: usize, end: usize| q1_range_f32(bytes, gpr, x, out_addr, start, end);
8100    dispatch_rows(pool, rows, &run);
8101}
8102
8103/// Fused two-input q1 matvec (weights read once per pair).
8104#[allow(clippy::too_many_arguments)]
8105fn q1_matvec2(
8106    bytes: &[u8],
8107    x1: &[f32],
8108    x2: &[f32],
8109    rows: usize,
8110    cols: usize,
8111    o1: &mut [f32],
8112    o2: &mut [f32],
8113    pool: Option<&Pool>,
8114) {
8115    let gpr = cols / GROUP_SIZE;
8116    let p1 = SendMut(o1.as_mut_ptr());
8117    let p2 = SendMut(o2.as_mut_ptr());
8118    if a8w8_enabled() {
8119        let a1 = split_act(x1);
8120        let a2 = split_act(x2);
8121        let g1 = q1_group_sums(&a1.xq, gpr);
8122        let g2 = q1_group_sums(&a2.xq, gpr);
8123        let (a1, a2, g1, g2) = (&a1, &a2, &g1, &g2);
8124        let run = move |start: usize, end: usize| {
8125            for r in start..end {
8126                let mut v1 = dot_q1_row_i8(bytes, r, gpr, &a1.xq, g1) * a1.sx;
8127                let mut v2 = dot_q1_row_i8(bytes, r, gpr, &a2.xq, g2) * a2.sx;
8128                for &(j, xv) in &a1.outliers {
8129                    let (w, s) = q1_outlier(bytes, r, gpr, j);
8130                    v1 += w * s * xv;
8131                }
8132                for &(j, xv) in &a2.outliers {
8133                    let (w, s) = q1_outlier(bytes, r, gpr, j);
8134                    v2 += w * s * xv;
8135                }
8136                // SAFETY: disjoint row ranges per worker.
8137                unsafe {
8138                    *p1.at(r) = v1;
8139                    *p2.at(r) = v2;
8140                }
8141            }
8142        };
8143        dispatch_rows(pool, rows, &run);
8144        return;
8145    }
8146    let run = move |start: usize, end: usize| {
8147        for r in start..end {
8148            // SAFETY: disjoint row ranges per worker.
8149            unsafe {
8150                *p1.at(r) = q1_row_exact(bytes, r, gpr, x1);
8151                *p2.at(r) = q1_row_exact(bytes, r, gpr, x2);
8152            }
8153        }
8154    };
8155    dispatch_rows(pool, rows, &run);
8156}
8157
8158/// Batched q1 matmat: each row's tiles stream once per microbatch.
8159#[allow(clippy::too_many_arguments)]
8160fn q1_matmat(
8161    bytes: &[u8],
8162    xs_all: &[f32],
8163    b: usize,
8164    rows: usize,
8165    cols: usize,
8166    out: &mut [f32],
8167    pool: Option<&Pool>,
8168) {
8169    debug_assert_eq!(out.len(), b * rows);
8170    let gpr = cols / GROUP_SIZE;
8171    let out_addr = SendMut(out.as_mut_ptr());
8172    if a8w8_enabled() {
8173        let acts: Vec<(SplitAct, Vec<i32>)> = (0..b)
8174            .map(|bi| {
8175                let act = split_act(&xs_all[bi * cols..(bi + 1) * cols]);
8176                let gsum = q1_group_sums(&act.xq, gpr);
8177                (act, gsum)
8178            })
8179            .collect();
8180        let acts = &acts;
8181        #[cfg(target_arch = "x86_64")]
8182        let blocked_ok = avx2_enabled() && blocked_enabled();
8183        #[cfg(target_arch = "aarch64")]
8184        let blocked_ok = sdot_enabled() && blocked_enabled();
8185        let run = move |start: usize, end: usize| {
8186            for r in start..end {
8187                let mut bi = 0usize;
8188                // Blocked 1×4: the unpacked bit mask serves four
8189                // activation streams per group.
8190                #[cfg(target_arch = "aarch64")]
8191                if blocked_ok {
8192                    while bi + 4 <= acts.len() {
8193                        let xs = [
8194                            acts[bi].0.xq.as_slice(),
8195                            acts[bi + 1].0.xq.as_slice(),
8196                            acts[bi + 2].0.xq.as_slice(),
8197                            acts[bi + 3].0.xq.as_slice(),
8198                        ];
8199                        let gs = [
8200                            acts[bi].1.as_slice(),
8201                            acts[bi + 1].1.as_slice(),
8202                            acts[bi + 2].1.as_slice(),
8203                            acts[bi + 3].1.as_slice(),
8204                        ];
8205                        let d = unsafe { dot_q1_row_1x4_sdot(bytes, r, gpr, xs, gs) };
8206                        for k in 0..4 {
8207                            let (act, _) = &acts[bi + k];
8208                            let mut acc = d[k] * act.sx;
8209                            for &(j, xv) in &act.outliers {
8210                                let (w, sc) = q1_outlier(bytes, r, gpr, j);
8211                                acc += w * sc * xv;
8212                            }
8213                            // SAFETY: disjoint (bi, r) cells per worker.
8214                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
8215                        }
8216                        bi += 4;
8217                    }
8218                }
8219                #[cfg(target_arch = "x86_64")]
8220                if blocked_ok {
8221                    while bi + 4 <= acts.len() {
8222                        let xs = [
8223                            acts[bi].0.xq.as_slice(),
8224                            acts[bi + 1].0.xq.as_slice(),
8225                            acts[bi + 2].0.xq.as_slice(),
8226                            acts[bi + 3].0.xq.as_slice(),
8227                        ];
8228                        let gs = [
8229                            acts[bi].1.as_slice(),
8230                            acts[bi + 1].1.as_slice(),
8231                            acts[bi + 2].1.as_slice(),
8232                            acts[bi + 3].1.as_slice(),
8233                        ];
8234                        let d = unsafe {
8235                            if vnni_tiles_enabled() {
8236                                dot_q1_row_1x4_vnni(bytes, r, gpr, xs, gs)
8237                            } else {
8238                                dot_q1_row_1x4_avx2(bytes, r, gpr, xs, gs)
8239                            }
8240                        };
8241                        for k in 0..4 {
8242                            let (act, _) = &acts[bi + k];
8243                            let mut acc = d[k] * act.sx;
8244                            for &(j, xv) in &act.outliers {
8245                                let (w, sc) = q1_outlier(bytes, r, gpr, j);
8246                                acc += w * sc * xv;
8247                            }
8248                            // SAFETY: disjoint (bi, r) cells per worker.
8249                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
8250                        }
8251                        bi += 4;
8252                    }
8253                }
8254                while bi < acts.len() {
8255                    let (act, gsum) = &acts[bi];
8256                    let mut acc = dot_q1_row_i8(bytes, r, gpr, &act.xq, gsum) * act.sx;
8257                    for &(j, xv) in &act.outliers {
8258                        let (w, s) = q1_outlier(bytes, r, gpr, j);
8259                        acc += w * s * xv;
8260                    }
8261                    // SAFETY: disjoint (bi, r) cells per worker range.
8262                    unsafe { *out_addr.at(bi * rows + r) = acc };
8263                    bi += 1;
8264                }
8265            }
8266        };
8267        dispatch_rows(pool, rows, &run);
8268        return;
8269    }
8270    let run = move |start: usize, end: usize| {
8271        for r in start..end {
8272            for bi in 0..b {
8273                let x = &xs_all[bi * cols..(bi + 1) * cols];
8274                // SAFETY: disjoint (bi, r) cells per worker range.
8275                unsafe { *out_addr.at(bi * rows + r) = q1_row_exact(bytes, r, gpr, x) };
8276            }
8277        }
8278    };
8279    dispatch_rows(pool, rows, &run);
8280}
8281
8282/// Fused q4_block matvec straight from the mapped bytes. SDOT path when
8283/// dotprod is available (port of vmfcore `dot_q4_block_sdot`, measured
8284/// +23% on q4 decode): nibbles → centered i8, int8×int8 `sdot` per
8285/// 32-group, exact outlier correction — the same A8W8 contract as q8.
8286/// `CMF_SDOT=0` keeps the exact scalar path.
8287fn q4matvec(
8288    bytes: &[u8],
8289    x: &[f32],
8290    rows: usize,
8291    cols: usize,
8292    out: &mut [f32],
8293    pool: Option<&Pool>,
8294) {
8295    debug_assert_eq!(out.len(), rows);
8296    let (packed, scales) = q4_split(bytes, rows, cols);
8297    let gpr = cols / GROUP_SIZE;
8298    let out_addr = SendMut(out.as_mut_ptr());
8299
8300    if a8w8_enabled() {
8301        let act = split_act(x);
8302        let run = move |start: usize, end: usize| {
8303            q4_range_a8w8(packed, scales, gpr, cols, &act, out_addr, start, end)
8304        };
8305        dispatch_rows(pool, rows, &run);
8306        return;
8307    }
8308
8309    let run =
8310        move |start: usize, end: usize| q4_range_f32(packed, scales, gpr, x, out_addr, start, end);
8311    dispatch_rows(pool, rows, &run);
8312}
8313
8314/// One q4 row via the A8W8 int8 path — SDOT on ARM, AVX2 maddubs on
8315/// x86 (scalar fallback is unreachable: callers gate on a8w8_enabled).
8316#[inline]
8317#[allow(unreachable_code)]
8318/// One UNPACKED q4 row (centered i8 in `buf`) against four activation
8319/// streams: the 32-byte weight chunk and its abs() load once per group,
8320/// the per-group f16 scale decodes once — four maddubs+reduce chains
8321/// instead of four full (load, abs, dot) rounds.
8322#[cfg(target_arch = "x86_64")]
8323#[target_feature(enable = "avx2")]
8324unsafe fn dot_q4b_row_1x4_avx2(
8325    buf: &[u8],
8326    scales: &[u8],
8327    g0: usize,
8328    gpr: usize,
8329    xs: [&[i8]; 4],
8330) -> [f32; 4] {
8331    // SAFETY: callers uphold buffer contracts (buf.len() == gpr·32).
8332    unsafe {
8333        use core::arch::x86_64::*;
8334        let ones = _mm256_set1_epi16(1);
8335        let mut acc = [0f32; 4];
8336        for gi in 0..gpr {
8337            let s = f16_to_f32(u16::from_le_bytes([
8338                scales[(g0 + gi) * 2],
8339                scales[(g0 + gi) * 2 + 1],
8340            ]));
8341            let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8342            let aw = _mm256_abs_epi8(w);
8343            for (k, xq) in xs.iter().enumerate() {
8344                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8345                let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
8346                let d = _mm256_madd_epi16(p16, ones);
8347                let hi128 = _mm256_extracti128_si256::<1>(d);
8348                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
8349                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8350                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8351                acc[k] += _mm_cvtsi128_si32(s32) as f32 * s;
8352            }
8353        }
8354        acc
8355    }
8356}
8357
8358/// VNNI twin of `dot_q4b_row_1x4_avx2` (see `dpbusd_hsum`).
8359#[cfg(target_arch = "x86_64")]
8360#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
8361unsafe fn dot_q4b_row_1x4_vnni(
8362    buf: &[u8],
8363    scales: &[u8],
8364    g0: usize,
8365    gpr: usize,
8366    xs: [&[i8]; 4],
8367) -> [f32; 4] {
8368    // SAFETY: callers uphold buffer contracts (buf.len() == gpr·32).
8369    unsafe {
8370        use core::arch::x86_64::*;
8371        let mut acc = [0f32; 4];
8372        for gi in 0..gpr {
8373            let s = f16_to_f32(u16::from_le_bytes([
8374                scales[(g0 + gi) * 2],
8375                scales[(g0 + gi) * 2 + 1],
8376            ]));
8377            let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8378            let aw = _mm256_abs_epi8(w);
8379            for (k, xq) in xs.iter().enumerate() {
8380                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8381                let d = dpbusd_hsum(aw, _mm256_sign_epi8(x, w));
8382                acc[k] += d as f32 * s;
8383            }
8384        }
8385        acc
8386    }
8387}
8388
8389/// The vbit flavor of the blocked 1×4: the per-activation A8W8 scale
8390/// folds in PER GROUP as `(d·sx)·s` — bit-matching the single-matvec
8391/// accumulation order (the q4_block flavor applies sx once at the end,
8392/// matching ITS single path; the two conventions are historical and
8393/// each blocked leg must mirror its own).
8394#[cfg(target_arch = "x86_64")]
8395#[target_feature(enable = "avx2")]
8396unsafe fn dot_q4b_row_1x4_sx_avx2(
8397    buf: &[u8],
8398    scales: &[u8],
8399    g0: usize,
8400    gpr: usize,
8401    xs: [&[i8]; 4],
8402    sxs: [f32; 4],
8403) -> [f32; 4] {
8404    // SAFETY: callers uphold buffer contracts (buf.len() == gpr·32).
8405    unsafe {
8406        use core::arch::x86_64::*;
8407        let ones = _mm256_set1_epi16(1);
8408        let mut acc = [0f32; 4];
8409        for gi in 0..gpr {
8410            let s = f16_to_f32(u16::from_le_bytes([
8411                scales[(g0 + gi) * 2],
8412                scales[(g0 + gi) * 2 + 1],
8413            ]));
8414            let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8415            let aw = _mm256_abs_epi8(w);
8416            for (k, xq) in xs.iter().enumerate() {
8417                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8418                let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
8419                let d = _mm256_madd_epi16(p16, ones);
8420                let hi128 = _mm256_extracti128_si256::<1>(d);
8421                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
8422                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8423                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8424                acc[k] += (_mm_cvtsi128_si32(s32) as f32 * sxs[k]) * s;
8425            }
8426        }
8427        acc
8428    }
8429}
8430
8431/// VNNI twin of `dot_q4b_row_1x4_sx_avx2` (see `dpbusd_hsum`; the
8432/// per-group `(d·sx)·s` fold mirrors the vbit single path).
8433#[cfg(target_arch = "x86_64")]
8434#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
8435unsafe fn dot_q4b_row_1x4_sx_vnni(
8436    buf: &[u8],
8437    scales: &[u8],
8438    g0: usize,
8439    gpr: usize,
8440    xs: [&[i8]; 4],
8441    sxs: [f32; 4],
8442) -> [f32; 4] {
8443    // SAFETY: callers uphold buffer contracts (buf.len() == gpr·32).
8444    unsafe {
8445        use core::arch::x86_64::*;
8446        let mut acc = [0f32; 4];
8447        for gi in 0..gpr {
8448            let s = f16_to_f32(u16::from_le_bytes([
8449                scales[(g0 + gi) * 2],
8450                scales[(g0 + gi) * 2 + 1],
8451            ]));
8452            let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8453            let aw = _mm256_abs_epi8(w);
8454            for (k, xq) in xs.iter().enumerate() {
8455                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8456                let d = dpbusd_hsum(aw, _mm256_sign_epi8(x, w));
8457                acc[k] += (d as f32 * sxs[k]) * s;
8458            }
8459        }
8460        acc
8461    }
8462}
8463
8464#[allow(unreachable_code)]
8465fn dot_q4_row_i8(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
8466    #[cfg(target_arch = "aarch64")]
8467    unsafe {
8468        return dot_q4_row_sdot(packed, scales, g0, gpr, xq);
8469    }
8470    #[cfg(target_arch = "x86_64")]
8471    unsafe {
8472        return dot_q4_row_avx2(packed, scales, g0, gpr, xq);
8473    }
8474    let mut acc = 0f32;
8475    for gi in 0..gpr {
8476        let g = g0 + gi;
8477        let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
8478        let mut d = 0i32;
8479        for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
8480            d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
8481                + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
8482        }
8483        acc += d as f32 * s;
8484    }
8485    acc
8486}
8487
8488/// Two-activation q4 row via the A8W8 int8 path (see `dot_q4_row_i8`).
8489#[inline]
8490#[allow(unreachable_code)]
8491fn dot_q4_row_i8_2(
8492    packed: &[u8],
8493    scales: &[u8],
8494    g0: usize,
8495    gpr: usize,
8496    xq1: &[i8],
8497    xq2: &[i8],
8498) -> (f32, f32) {
8499    #[cfg(target_arch = "aarch64")]
8500    unsafe {
8501        return dot_q4_row_sdot2(packed, scales, g0, gpr, xq1, xq2);
8502    }
8503    #[cfg(target_arch = "x86_64")]
8504    unsafe {
8505        return dot_q4_row_avx2_2(packed, scales, g0, gpr, xq1, xq2);
8506    }
8507    (
8508        dot_q4_row_i8(packed, scales, g0, gpr, xq1),
8509        dot_q4_row_i8(packed, scales, g0, gpr, xq2),
8510    )
8511}
8512
8513/// One q4 row range via SDOT (kernel body of `q4matvec`, extracted so
8514/// multi-matrix jobs can drive it for several tensors in one dispatch).
8515#[allow(clippy::too_many_arguments)]
8516fn q4_range_a8w8(
8517    packed: &[u8],
8518    scales: &[u8],
8519    gpr: usize,
8520    cols: usize,
8521    act: &SplitAct,
8522    out: SendMut,
8523    start: usize,
8524    end: usize,
8525) {
8526    for r in start..end {
8527        let mut acc = dot_q4_row_i8(packed, scales, r * gpr, gpr, &act.xq) * act.sx;
8528        // xq is zeroed at outlier slots — add the exact terms.
8529        for &(j, xv) in &act.outliers {
8530            let flat = r * cols + j;
8531            let byte = packed[flat / 2];
8532            let nib = if flat & 1 == 0 {
8533                byte & 0x0F
8534            } else {
8535                byte >> 4
8536            };
8537            let s = f16_to_f32(u16::from_le_bytes([
8538                scales[(flat / GROUP_SIZE) * 2],
8539                scales[(flat / GROUP_SIZE) * 2 + 1],
8540            ]));
8541            acc += ((nib as i32 - 8) as f32) * s * xv;
8542        }
8543        // SAFETY: disjoint row ranges per worker.
8544        unsafe { *out.at(r) = acc };
8545    }
8546}
8547
8548/// Two-input q4 row range via the A8W8 int8 path — kernel body of
8549/// `q4matvec2`, extracted for pair multi-matrix jobs.
8550#[allow(clippy::too_many_arguments)]
8551fn q4_range2_a8w8(
8552    packed: &[u8],
8553    scales: &[u8],
8554    gpr: usize,
8555    cols: usize,
8556    a1: &SplitAct,
8557    a2: &SplitAct,
8558    p1: SendMut,
8559    p2: SendMut,
8560    start: usize,
8561    end: usize,
8562) {
8563    for r in start..end {
8564        let (s1, s2) = dot_q4_row_i8_2(packed, scales, r * gpr, gpr, &a1.xq, &a2.xq);
8565        let mut acc1 = s1 * a1.sx;
8566        let mut acc2 = s2 * a2.sx;
8567        // xq is zeroed at outlier slots — add the exact terms.
8568        let fix = |outliers: &[(usize, f32)], acc: &mut f32| {
8569            for &(j, xv) in outliers {
8570                let flat = r * cols + j;
8571                let byte = packed[flat / 2];
8572                let nib = if flat & 1 == 0 {
8573                    byte & 0x0F
8574                } else {
8575                    byte >> 4
8576                };
8577                let s = f16_to_f32(u16::from_le_bytes([
8578                    scales[(flat / GROUP_SIZE) * 2],
8579                    scales[(flat / GROUP_SIZE) * 2 + 1],
8580                ]));
8581                *acc += ((nib as i32 - 8) as f32) * s * xv;
8582            }
8583        };
8584        fix(&a1.outliers, &mut acc1);
8585        fix(&a2.outliers, &mut acc2);
8586        // SAFETY: disjoint row ranges per worker.
8587        unsafe {
8588            *p1.at(r) = acc1;
8589            *p2.at(r) = acc2;
8590        }
8591    }
8592}
8593
8594/// Exact scalar q4 row range (same extraction, non-SDOT path).
8595fn q4_range_f32(
8596    packed: &[u8],
8597    scales: &[u8],
8598    gpr: usize,
8599    x: &[f32],
8600    out: SendMut,
8601    start: usize,
8602    end: usize,
8603) {
8604    for r in start..end {
8605        let mut acc = 0f32;
8606        for gi in 0..gpr {
8607            let g = r * gpr + gi;
8608            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
8609            let pk = &packed[g * 16..(g + 1) * 16];
8610            let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
8611            let mut ga = 0f32;
8612            for (k, &b) in pk.iter().enumerate() {
8613                ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
8614                    + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
8615            }
8616            acc += ga * s;
8617        }
8618        // SAFETY: disjoint row ranges per worker.
8619        unsafe { *out.at(r) = acc };
8620    }
8621}
8622
8623/// Fused two-input q4 matvec: nibbles are unpacked ONCE per group and
8624/// dotted against both activations (was: two full matvecs — double
8625/// weight traffic). Per-lane math matches `q4matvec` exactly.
8626#[allow(clippy::too_many_arguments)]
8627fn q4matvec2(
8628    bytes: &[u8],
8629    x1: &[f32],
8630    x2: &[f32],
8631    rows: usize,
8632    cols: usize,
8633    o1: &mut [f32],
8634    o2: &mut [f32],
8635    pool: Option<&Pool>,
8636) {
8637    debug_assert_eq!(o1.len(), rows);
8638    debug_assert_eq!(o2.len(), rows);
8639    let (packed, scales) = q4_split(bytes, rows, cols);
8640    let gpr = cols / GROUP_SIZE;
8641
8642    if a8w8_enabled() {
8643        let a1 = split_act(x1);
8644        let a2 = split_act(x2);
8645        let p1 = SendMut(o1.as_mut_ptr());
8646        let p2 = SendMut(o2.as_mut_ptr());
8647        let run = move |start: usize, end: usize| {
8648            q4_range2_a8w8(packed, scales, gpr, cols, &a1, &a2, p1, p2, start, end)
8649        };
8650        dispatch_rows(pool, rows, &run);
8651        return;
8652    }
8653
8654    let p1 = SendMut(o1.as_mut_ptr());
8655    let p2 = SendMut(o2.as_mut_ptr());
8656    let run = move |start: usize, end: usize| {
8657        q4_range2_f32(packed, scales, gpr, x1, x2, p1, p2, start, end)
8658    };
8659    dispatch_rows(pool, rows, &run);
8660}
8661
8662/// Two-input exact scalar q4 row range (same extraction).
8663#[allow(clippy::too_many_arguments)]
8664fn q4_range2_f32(
8665    packed: &[u8],
8666    scales: &[u8],
8667    gpr: usize,
8668    x1: &[f32],
8669    x2: &[f32],
8670    p1: SendMut,
8671    p2: SendMut,
8672    start: usize,
8673    end: usize,
8674) {
8675    for r in start..end {
8676        let (mut acc1, mut acc2) = (0f32, 0f32);
8677        for gi in 0..gpr {
8678            let g = r * gpr + gi;
8679            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
8680            let pk = &packed[g * 16..(g + 1) * 16];
8681            let x1g = &x1[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
8682            let x2g = &x2[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
8683            let (mut g1, mut g2) = (0f32, 0f32);
8684            for (k, &b) in pk.iter().enumerate() {
8685                let wl = (b & 0x0F) as f32 - 8.0;
8686                let wh = ((b >> 4) & 0x0F) as f32 - 8.0;
8687                g1 += wl * x1g[k * 2] + wh * x1g[k * 2 + 1];
8688                g2 += wl * x2g[k * 2] + wh * x2g[k * 2 + 1];
8689            }
8690            acc1 += g1 * s;
8691            acc2 += g2 * s;
8692        }
8693        // SAFETY: disjoint row ranges per worker.
8694        unsafe {
8695            *p1.at(r) = acc1;
8696            *p2.at(r) = acc2;
8697        }
8698    }
8699}
8700
8701thread_local! {
8702    /// Per-worker decoded-row scratch for the batched q4/vbit kernels
8703    /// (centered i8 for SDOT, f32 for the exact/scalar paths).
8704    static ROW_I8: std::cell::RefCell<Vec<u8>> = const { std::cell::RefCell::new(Vec::new()) };
8705    static ROW_F32: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
8706}
8707
8708/// Batched q4 matmat: each weight row is unpacked from the mmap ONCE
8709/// and dotted against ALL b activations (prefill used to fall back to b
8710/// full matvecs — b× weight traffic and b× nibble decode). Per-position
8711/// math matches `q4matvec` exactly: same group order, same accumulation.
8712/// `out` is row-major [b, rows] like `qmatmat`.
8713#[allow(clippy::too_many_arguments)]
8714fn q4matmat(
8715    bytes: &[u8],
8716    xs_all: &[f32],
8717    b: usize,
8718    rows: usize,
8719    cols: usize,
8720    out: &mut [f32],
8721    pool: Option<&Pool>,
8722) {
8723    debug_assert_eq!(xs_all.len(), b * cols);
8724    debug_assert_eq!(out.len(), b * rows);
8725    let (packed, scales) = q4_split(bytes, rows, cols);
8726    let gpr = cols / GROUP_SIZE;
8727    let gscale = |g: usize| f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
8728
8729    if a8w8_enabled() {
8730        let acts: Vec<SplitAct> = (0..b)
8731            .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
8732            .collect();
8733        let acts = &acts;
8734        let out_addr = SendMut(out.as_mut_ptr());
8735        let run = move |start: usize, end: usize| {
8736            ROW_I8.with(|rb| {
8737                let mut buf = rb.borrow_mut();
8738                buf.resize(cols, 0);
8739                for r in start..end {
8740                    // Unpack the row's nibbles to centered i8 once
8741                    // (element 2k = low nibble, 2k+1 = high — flat order,
8742                    // same as dot_q4_row_sdot's zip).
8743                    for gi in 0..gpr {
8744                        let g = r * gpr + gi;
8745                        for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
8746                            buf[gi * GROUP_SIZE + k * 2] = ((bt & 0x0F) as i32 - 8) as i8 as u8;
8747                            buf[gi * GROUP_SIZE + k * 2 + 1] =
8748                                (((bt >> 4) & 0x0F) as i32 - 8) as i8 as u8;
8749                        }
8750                    }
8751                    let mut bi = 0usize;
8752                    #[cfg(target_arch = "x86_64")]
8753                    if avx2_enabled() && blocked_enabled() {
8754                        while bi + 4 <= acts.len() {
8755                            let xs = [
8756                                acts[bi].xq.as_slice(),
8757                                acts[bi + 1].xq.as_slice(),
8758                                acts[bi + 2].xq.as_slice(),
8759                                acts[bi + 3].xq.as_slice(),
8760                            ];
8761                            let d = unsafe {
8762                                if vnni_tiles_enabled() {
8763                                    dot_q4b_row_1x4_vnni(&buf, scales, r * gpr, gpr, xs)
8764                                } else {
8765                                    dot_q4b_row_1x4_avx2(&buf, scales, r * gpr, gpr, xs)
8766                                }
8767                            };
8768                            for k in 0..4 {
8769                                let act = &acts[bi + k];
8770                                let mut acc = d[k] * act.sx;
8771                                for &(j, xv) in &act.outliers {
8772                                    acc += (buf[j] as i8) as f32
8773                                        * gscale((r * cols + j) / GROUP_SIZE)
8774                                        * xv;
8775                                }
8776                                // SAFETY: disjoint (bi, r) cells per worker.
8777                                unsafe { *out_addr.at((bi + k) * rows + r) = acc };
8778                            }
8779                            bi += 4;
8780                        }
8781                    }
8782                    while bi < acts.len() {
8783                        let act = &acts[bi];
8784                        let mut acc = 0f32;
8785                        for gi in 0..gpr {
8786                            let d = dot_i8_i8(
8787                                &buf[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
8788                                &act.xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
8789                            );
8790                            acc += d as f32 * gscale(r * gpr + gi);
8791                        }
8792                        acc *= act.sx;
8793                        // xq is zeroed at outlier slots — exact terms.
8794                        for &(j, xv) in &act.outliers {
8795                            acc += (buf[j] as i8) as f32 * gscale((r * cols + j) / GROUP_SIZE) * xv;
8796                        }
8797                        // SAFETY: disjoint (bi, r) cells per worker row range.
8798                        unsafe { *out_addr.at(bi * rows + r) = acc };
8799                        bi += 1;
8800                    }
8801                }
8802            })
8803        };
8804        dispatch_rows(pool, rows, &run);
8805        return;
8806    }
8807
8808    let out_addr = SendMut(out.as_mut_ptr());
8809    let run = move |start: usize, end: usize| {
8810        ROW_F32.with(|rb| {
8811            let mut buf = rb.borrow_mut();
8812            buf.resize(cols, 0.0);
8813            for r in start..end {
8814                // Decode raw (nib − 8) values once; scales stay per-group
8815                // so the accumulation order matches q4matvec bit-for-bit.
8816                for gi in 0..gpr {
8817                    let g = r * gpr + gi;
8818                    for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
8819                        buf[gi * GROUP_SIZE + k * 2] = (bt & 0x0F) as f32 - 8.0;
8820                        buf[gi * GROUP_SIZE + k * 2 + 1] = ((bt >> 4) & 0x0F) as f32 - 8.0;
8821                    }
8822                }
8823                for bi in 0..b {
8824                    let x = &xs_all[bi * cols..(bi + 1) * cols];
8825                    let mut acc = 0f32;
8826                    for gi in 0..gpr {
8827                        let mut ga = 0f32;
8828                        // Pairwise (lo + hi) addition, matching
8829                        // q4matvec's `ga += lo·x + hi·x` shape exactly —
8830                        // a flat one-per-element loop rounds differently
8831                        // and broke bit-parity on the scalar (x86) path.
8832                        for k in 0..GROUP_SIZE / 2 {
8833                            let e = gi * GROUP_SIZE + k * 2;
8834                            ga += buf[e] * x[e] + buf[e + 1] * x[e + 1];
8835                        }
8836                        acc += ga * gscale(r * gpr + gi);
8837                    }
8838                    // SAFETY: disjoint (bi, r) cells per worker row range.
8839                    unsafe { *out_addr.at(bi * rows + r) = acc };
8840                }
8841            }
8842        })
8843    };
8844    dispatch_rows(pool, rows, &run);
8845}
8846
8847/// Batched vbit matmat: each variable-bit row is decoded from the mmap
8848/// ONCE for the whole microbatch. Same per-position math as
8849/// `vbitmatvec` (SDOT A8W8 with exact outliers / exact f32 for b=8 rows
8850/// and the scalar path).
8851#[allow(clippy::too_many_arguments)]
8852fn vbitmatmat(
8853    bytes: &[u8],
8854    offsets: &[usize],
8855    xs_all: &[f32],
8856    b: usize,
8857    rows: usize,
8858    cols: usize,
8859    out: &mut [f32],
8860    pool: Option<&Pool>,
8861) {
8862    debug_assert_eq!(xs_all.len(), b * cols);
8863    debug_assert_eq!(out.len(), b * rows);
8864    debug_assert_eq!(offsets.len(), rows + 1);
8865    let ng = cols / GROUP_SIZE;
8866    let bits = &bytes[..rows];
8867    let sc_off = rows;
8868    let gscale = |r: usize, g: usize| {
8869        let so = (r * ng + g) * 2;
8870        f16_to_f32(u16::from_le_bytes([
8871            bytes[sc_off + so],
8872            bytes[sc_off + so + 1],
8873        ]))
8874    };
8875
8876    // Decode row r's raw (u − L) values into `dst` (f32, unscaled).
8877    let decode_f32 = |r: usize, dst: &mut [f32]| {
8878        let bw = bits[r] as usize;
8879        let l = ((1i32 << (bw - 1)) - 1) as f32;
8880        let data = &bytes[offsets[r]..offsets[r + 1]];
8881        let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
8882        for d in dst.iter_mut() {
8883            while nbits < bw {
8884                acc = (acc << 8) | data[idx] as u64;
8885                idx += 1;
8886                nbits += 8;
8887            }
8888            let u = ((acc >> (nbits - bw)) & ((1u64 << bw) - 1)) as f32;
8889            nbits -= bw;
8890            *d = u - l;
8891        }
8892    };
8893
8894    if a8w8_enabled() {
8895        let acts: Vec<SplitAct> = (0..b)
8896            .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
8897            .collect();
8898        let acts = &acts;
8899        let out_addr = SendMut(out.as_mut_ptr());
8900        let run = move |start: usize, end: usize| {
8901            for r in start..end {
8902                let bw = bits[r] as usize;
8903                if bw == 8 {
8904                    // u−L reaches 128 → no i8 path; decode once, exact
8905                    // f32 dots for every position (same as vbitmatvec).
8906                    ROW_F32.with(|rb| {
8907                        let mut buf = rb.borrow_mut();
8908                        buf.resize(cols, 0.0);
8909                        decode_f32(r, &mut buf);
8910                        for bi in 0..b {
8911                            let x = &xs_all[bi * cols..(bi + 1) * cols];
8912                            let mut dot = 0f32;
8913                            for g in 0..ng {
8914                                let mut gd = 0f32;
8915                                for k in 0..GROUP_SIZE {
8916                                    gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
8917                                }
8918                                dot += gd * gscale(r, g);
8919                            }
8920                            // SAFETY: disjoint (bi, r) cells per worker range.
8921                            unsafe { *out_addr.at(bi * rows + r) = dot };
8922                        }
8923                    });
8924                    continue;
8925                }
8926                let l = (1i32 << (bw - 1)) - 1;
8927                let data = &bytes[offsets[r]..offsets[r + 1]];
8928                ROW_I8.with(|rb| {
8929                    let mut buf = rb.borrow_mut();
8930                    buf.resize(cols, 0);
8931                    #[inline(always)]
8932                    fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
8933                        for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
8934                            let u = unpack8::<B>(&data[blk * B..]);
8935                            for k in 0..8 {
8936                                chunk[k] = (u[k] - l) as i8 as u8;
8937                            }
8938                        }
8939                    }
8940                    match bw {
8941                        3 => fill::<3>(data, l, &mut buf),
8942                        4 => vbit_fill4(data, &mut buf),
8943                        5 => fill::<5>(data, l, &mut buf),
8944                        6 => fill::<6>(data, l, &mut buf),
8945                        _ => unreachable!("vbit bit-width {bw} (validated at load)"),
8946                    }
8947                    let mut bi = 0usize;
8948                    // The vbit scale table shares q4_block's layout
8949                    // (contiguous f16 per (row·ng + g)), so the same
8950                    // blocked 1×4 kernel serves the decoded row.
8951                    #[cfg(target_arch = "x86_64")]
8952                    if avx2_enabled() && blocked_enabled() {
8953                        while bi + 4 <= acts.len() {
8954                            let xs = [
8955                                acts[bi].xq.as_slice(),
8956                                acts[bi + 1].xq.as_slice(),
8957                                acts[bi + 2].xq.as_slice(),
8958                                acts[bi + 3].xq.as_slice(),
8959                            ];
8960                            let sxs = [
8961                                acts[bi].sx,
8962                                acts[bi + 1].sx,
8963                                acts[bi + 2].sx,
8964                                acts[bi + 3].sx,
8965                            ];
8966                            let d = unsafe {
8967                                if vnni_tiles_enabled() {
8968                                    dot_q4b_row_1x4_sx_vnni(
8969                                        &buf,
8970                                        &bytes[sc_off..],
8971                                        r * ng,
8972                                        ng,
8973                                        xs,
8974                                        sxs,
8975                                    )
8976                                } else {
8977                                    dot_q4b_row_1x4_sx_avx2(
8978                                        &buf,
8979                                        &bytes[sc_off..],
8980                                        r * ng,
8981                                        ng,
8982                                        xs,
8983                                        sxs,
8984                                    )
8985                                }
8986                            };
8987                            for k in 0..4 {
8988                                let act = &acts[bi + k];
8989                                let mut dot = d[k];
8990                                for &(j, xv) in &act.outliers {
8991                                    dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
8992                                }
8993                                // SAFETY: disjoint (bi, r) cells per worker.
8994                                unsafe { *out_addr.at((bi + k) * rows + r) = dot };
8995                            }
8996                            bi += 4;
8997                        }
8998                    }
8999                    while bi < acts.len() {
9000                        let act = &acts[bi];
9001                        let mut dot = 0f32;
9002                        for g in 0..ng {
9003                            let d = dot_i8_i8(
9004                                &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
9005                                &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
9006                            ) as f32
9007                                * act.sx;
9008                            dot += d * gscale(r, g);
9009                        }
9010                        for &(j, xv) in &act.outliers {
9011                            dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
9012                        }
9013                        // SAFETY: disjoint (bi, r) cells per worker range.
9014                        unsafe { *out_addr.at(bi * rows + r) = dot };
9015                        bi += 1;
9016                    }
9017                });
9018            }
9019        };
9020        dispatch_rows(pool, rows, &run);
9021        return;
9022    }
9023
9024    let out_addr = SendMut(out.as_mut_ptr());
9025    let run = move |start: usize, end: usize| {
9026        ROW_F32.with(|rb| {
9027            let mut buf = rb.borrow_mut();
9028            buf.resize(cols, 0.0);
9029            for r in start..end {
9030                decode_f32(r, &mut buf);
9031                for bi in 0..b {
9032                    let x = &xs_all[bi * cols..(bi + 1) * cols];
9033                    let mut dot = 0f32;
9034                    for g in 0..ng {
9035                        let mut gd = 0f32;
9036                        for k in 0..GROUP_SIZE {
9037                            gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
9038                        }
9039                        dot += gd * gscale(r, g);
9040                    }
9041                    // SAFETY: disjoint (bi, r) cells per worker range.
9042                    unsafe { *out_addr.at(bi * rows + r) = dot };
9043                }
9044            }
9045        })
9046    };
9047    dispatch_rows(pool, rows, &run);
9048}
9049
9050/// Build a GPU batch job for a q8-family mapped tensor (primary
9051/// shard): prescaled input + directory coordinates. None → not
9052/// GPU-eligible, caller stays on the CPU.
9053pub(crate) fn gpu_batch_job<'a>(
9054    t: &'a QTensor,
9055    x: &[f32],
9056) -> Option<(std::sync::Arc<CmfModel>, crate::gpu::BatchJob<'a>)> {
9057    match t {
9058        QTensor::Mapped {
9059            model,
9060            idx,
9061            dtype: dt @ (TensorDtype::Q8Row | TensorDtype::Q8_2f),
9062            rows,
9063            cols,
9064            row_scale,
9065            col_field,
9066            ..
9067        } => Some((
9068            model.clone(),
9069            crate::gpu::BatchJob {
9070                idx: *idx,
9071                rows: *rows,
9072                cols: *cols,
9073                row_scale,
9074                xs: prescale(x, col_field, *dt).into_owned(),
9075                layout: crate::gpu::BatchLayout::Q8,
9076            },
9077        )),
9078        // q1: raw f32 activations, tile-embedded scales.
9079        QTensor::Mapped {
9080            model,
9081            idx,
9082            dtype: TensorDtype::Q1,
9083            rows,
9084            cols,
9085            ..
9086        } => Some((
9087            model.clone(),
9088            crate::gpu::BatchJob {
9089                idx: *idx,
9090                rows: *rows,
9091                cols: *cols,
9092                row_scale: &[],
9093                xs: x.to_vec(),
9094                layout: crate::gpu::BatchLayout::Q1,
9095            },
9096        )),
9097        // q4_tiled / q4tp: raw f32 activations; the scales live in the
9098        // payload (inline tiles / row ladder), so row_scale stays empty.
9099        // The GDN projection batch already runs these layouts on Metal —
9100        // this arm lets the attention QKV batch reach the same kernels.
9101        QTensor::Mapped {
9102            model,
9103            idx,
9104            dtype: dt @ (TensorDtype::Q4Tiled | TensorDtype::Q4TiledP),
9105            rows,
9106            cols,
9107            ..
9108        } => Some((
9109            model.clone(),
9110            crate::gpu::BatchJob {
9111                idx: *idx,
9112                rows: *rows,
9113                cols: *cols,
9114                row_scale: &[],
9115                xs: x.to_vec(),
9116                layout: if *dt == TensorDtype::Q4Tiled {
9117                    crate::gpu::BatchLayout::Q4t
9118                } else {
9119                    crate::gpu::BatchLayout::Q4tp
9120                },
9121            },
9122        )),
9123        _ => None,
9124    }
9125}
9126
9127thread_local! {
9128    static PRESCALE_BUF1: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
9129    static PRESCALE_BUF2: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
9130}
9131
9132pub(crate) fn prescale<'a>(
9133    x: &'a [f32],
9134    col_field: &[f32],
9135    dtype: TensorDtype,
9136) -> std::borrow::Cow<'a, [f32]> {
9137    if dtype == TensorDtype::Q8_2f {
9138        x.iter().zip(col_field).map(|(a, c)| a * c).collect()
9139    } else {
9140        std::borrow::Cow::Borrowed(x)
9141    }
9142}
9143
9144/// θ col-field fold for q8_2f activations. Borrowed pass-through for
9145/// every other dtype, using thread-local buffers to eliminate per-matvec allocations.
9146pub(crate) fn prescale_with<R, F: FnOnce(&[f32]) -> R>(
9147    x: &[f32],
9148    col_field: &[f32],
9149    dtype: TensorDtype,
9150    buf_id: u8,
9151    f: F,
9152) -> R {
9153    if dtype == TensorDtype::Q8_2f {
9154        if buf_id == 1 {
9155            PRESCALE_BUF1.with(|b| {
9156                let mut buf = b.borrow_mut();
9157                buf.clear();
9158                buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
9159                f(&buf)
9160            })
9161        } else {
9162            PRESCALE_BUF2.with(|b| {
9163                let mut buf = b.borrow_mut();
9164                buf.clear();
9165                buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
9166                f(&buf)
9167            })
9168        }
9169    } else {
9170        f(x)
9171    }
9172}
9173
9174// ───────────────────── x86-64 AVX2 kernels (roadmap этап 2) ─────────────────────
9175
9176/// AVX2+FMA available? Default ON when the CPU supports both;
9177/// `CMF_AVX2=0` disables (falls back to the autovectorized loops).
9178#[cfg(target_arch = "x86_64")]
9179pub(crate) fn avx2_enabled() -> bool {
9180    use std::sync::OnceLock;
9181    static ON: OnceLock<bool> = OnceLock::new();
9182    *ON.get_or_init(|| {
9183        std::env::var("CMF_AVX2").map(|v| v != "0").unwrap_or(true)
9184            && std::arch::is_x86_feature_detected!("avx2")
9185            && std::arch::is_x86_feature_detected!("fma")
9186    })
9187}
9188
9189/// AVX2 A8W8 allowed? The quantized-activation contract is switched by
9190/// the SAME env as the ARM SDOT path: `CMF_SDOT=0` keeps exact kernels
9191/// (the golden-parity exact gate relies on it) — AVX2 f32 kernels stay
9192/// active either way, they are exact (regrouped sums only).
9193#[cfg(target_arch = "x86_64")]
9194fn avx2_a8w8_enabled() -> bool {
9195    use std::sync::OnceLock;
9196    static ON: OnceLock<bool> = OnceLock::new();
9197    *ON.get_or_init(|| {
9198        avx2_enabled() && std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true)
9199    })
9200}
9201
9202/// A8W8 quantized-activation path available on THIS machine? One
9203/// switch across architectures: ARM dotprod (CMF_SDOT) or x86 AVX2
9204/// (CMF_AVX2 + the same CMF_SDOT exact-contract override).
9205#[inline]
9206pub(crate) fn a8w8_enabled() -> bool {
9207    #[cfg(target_arch = "aarch64")]
9208    {
9209        sdot_enabled()
9210    }
9211    #[cfg(target_arch = "x86_64")]
9212    {
9213        avx2_a8w8_enabled()
9214    }
9215    #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
9216    {
9217        false
9218    }
9219}
9220
9221/// int8·int8 dot dispatch: SDOT on ARM; AVX-512 VNNI (vpdpbusd) or AVX2
9222/// maddubs on x86. Callers are gated by `a8w8_enabled()`.
9223#[inline]
9224#[allow(unreachable_code)]
9225fn dot_i8_i8(w: &[u8], xq: &[i8]) -> i32 {
9226    #[cfg(target_arch = "aarch64")]
9227    unsafe {
9228        return dot_i8_sdot(w, xq);
9229    }
9230    #[cfg(target_arch = "x86_64")]
9231    unsafe {
9232        if avx512vnni_enabled() {
9233            return dot_i8_i8_vnni(w, xq);
9234        }
9235        return dot_i8_i8_avx2(w, xq);
9236    }
9237    w.iter()
9238        .zip(xq)
9239        .map(|(&a, &b)| (a as i8) as i32 * b as i32)
9240        .sum()
9241}
9242
9243/// AVX-512 VNNI available? (F+BW+VL+VNNI; `CMF_AVX512=0` falls back to
9244/// AVX2.) VL matters: short 32-byte groups (q4/vbit) ride the 256-bit
9245/// `vpdpbusd` encoding.
9246#[cfg(target_arch = "x86_64")]
9247fn avx512vnni_enabled() -> bool {
9248    use std::sync::OnceLock;
9249    static ON: OnceLock<bool> = OnceLock::new();
9250    *ON.get_or_init(|| {
9251        std::env::var("CMF_AVX512")
9252            .map(|v| v != "0")
9253            .unwrap_or(true)
9254            && std::arch::is_x86_feature_detected!("avx512f")
9255            && std::arch::is_x86_feature_detected!("avx512bw")
9256            && std::arch::is_x86_feature_detected!("avx512vl")
9257            && std::arch::is_x86_feature_detected!("avx512vnni")
9258    })
9259}
9260
9261/// Grouped-codec VNNI arms (the q4t/q4b/q1/q1t tile kernels): default
9262/// ON where AVX-512 VNNI exists (`CMF_VNNI_TILES=0` opt-out). Measured
9263/// on Ryzen 7950X (Zen4, 3 alternating process pairs, blocked GEMM
9264/// 4864×896 b=256): q4t 63→68 GF/s (+8%), q1 53→56 (+6%), q4b 72→75
9265/// (+4%) — consistent, no leg regressed. The tile kernels keep a
9266/// horizontal reduce per 32-weight group, so the `vpdpbusd` saving is
9267/// smaller than the long-dot q8 win (+13%), but it is real and free.
9268#[cfg(target_arch = "x86_64")]
9269fn vnni_tiles_enabled() -> bool {
9270    use std::sync::OnceLock;
9271    static ON: OnceLock<bool> = OnceLock::new();
9272    *ON.get_or_init(|| {
9273        std::env::var("CMF_VNNI_TILES")
9274            .map(|v| v != "0")
9275            .unwrap_or(true)
9276            && avx512vnni_enabled()
9277    })
9278}
9279
9280/// One 256-bit u8×i8 dot → i32 via `vpdpbusd` into a fresh accumulator
9281/// plus the same horizontal reduce the AVX2 kernels use. Products are
9282/// bounded (|w| ≤ 8 or ≤ 1), so maddubs never saturated — the i32 sum
9283/// is bit-identical to the maddubs+madd pair it replaces.
9284#[cfg(target_arch = "x86_64")]
9285#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
9286#[inline]
9287unsafe fn dpbusd_hsum(aw: core::arch::x86_64::__m256i, xs: core::arch::x86_64::__m256i) -> i32 {
9288    // SAFETY: pure register math.
9289    unsafe {
9290        use core::arch::x86_64::*;
9291        let d = _mm256_dpbusd_epi32(_mm256_setzero_si256(), aw, xs);
9292        let hi128 = _mm256_extracti128_si256::<1>(d);
9293        let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
9294        let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
9295        let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
9296        _mm_cvtsi128_si32(s32)
9297    }
9298}
9299
9300/// int8·int8 via AVX-512 VNNI: `vpdpbusd` fuses the maddubs+madd+add
9301/// triple into one u8×i8 dot-accumulate. AVX-512 has no vpsignb, so the
9302/// |w|·sign(x,w) trick becomes |w| × (x negated where w<0) via a mask
9303/// subtract — w==0 lanes contribute 0 through |w|=0 either way.
9304#[cfg(target_arch = "x86_64")]
9305#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
9306unsafe fn dot_i8_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
9307    // SAFETY: callers uphold slice-length contracts (see call sites).
9308    unsafe {
9309        use core::arch::x86_64::*;
9310        let n = w.len();
9311        let mut j = 0usize;
9312        let mut total: i32;
9313        // 4 independent accumulators: vpdpbusd is its own loop-carried
9314        // dependency (~5-cycle latency) — a single-acc loop runs
9315        // latency-bound and LOSES to the AVX2 maddubs kernel, measured
9316        // on Granite Rapids.
9317        {
9318            #[inline(always)]
9319            unsafe fn step(
9320                w: *const u8,
9321                x: *const i8,
9322                acc: core::arch::x86_64::__m512i,
9323            ) -> core::arch::x86_64::__m512i {
9324                unsafe {
9325                    use core::arch::x86_64::*;
9326                    let wv = _mm512_loadu_si512(w as *const _);
9327                    let xv = _mm512_loadu_si512(x as *const _);
9328                    let aw = _mm512_abs_epi8(wv);
9329                    let neg = _mm512_movepi8_mask(wv);
9330                    let sx = _mm512_mask_sub_epi8(xv, neg, _mm512_setzero_si512(), xv);
9331                    _mm512_dpbusd_epi32(acc, aw, sx)
9332                }
9333            }
9334            let (mut a0, mut a1, mut a2, mut a3) = (
9335                _mm512_setzero_si512(),
9336                _mm512_setzero_si512(),
9337                _mm512_setzero_si512(),
9338                _mm512_setzero_si512(),
9339            );
9340            while j + 256 <= n {
9341                a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
9342                a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), a1);
9343                a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), a2);
9344                a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), a3);
9345                j += 256;
9346            }
9347            while j + 64 <= n {
9348                a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
9349                j += 64;
9350            }
9351            let s01 = _mm512_add_epi32(a0, a1);
9352            let s23 = _mm512_add_epi32(a2, a3);
9353            total = _mm512_reduce_add_epi32(_mm512_add_epi32(s01, s23));
9354        }
9355        // 32-wide (q4/vbit groups are exactly 32 bytes).
9356        if j + 32 <= n {
9357            let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
9358            let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
9359            let d = _mm256_dpbusd_epi32(
9360                _mm256_setzero_si256(),
9361                _mm256_abs_epi8(wv),
9362                _mm256_sign_epi8(xv, wv),
9363            );
9364            let hi128 = _mm256_extracti128_si256::<1>(d);
9365            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
9366            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
9367            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
9368            total += _mm_cvtsi128_si32(s32);
9369            j += 32;
9370        }
9371        while j < n {
9372            total += (w[j] as i8) as i32 * xq[j] as i32;
9373            j += 1;
9374        }
9375        total
9376    }
9377}
9378
9379/// i8 row · f32 x via AVX2/FMA (x86 mirror of `dot_i8_f32_neon`).
9380#[cfg(target_arch = "x86_64")]
9381#[target_feature(enable = "avx2,fma")]
9382unsafe fn dot_i8_f32_avx2(w: &[u8], x: &[f32]) -> f32 {
9383    // SAFETY: callers uphold slice-length contracts (see call sites).
9384    unsafe {
9385        use core::arch::x86_64::*;
9386        let n = x.len();
9387        let wp = w.as_ptr();
9388        let xp = x.as_ptr();
9389        let (mut a0, mut a1) = (_mm256_setzero_ps(), _mm256_setzero_ps());
9390        let mut j = 0usize;
9391        while j + 16 <= n {
9392            let wb = _mm_loadu_si128(wp.add(j) as *const __m128i);
9393            let lo = _mm256_cvtepi8_epi32(wb);
9394            let hi = _mm256_cvtepi8_epi32(_mm_srli_si128::<8>(wb));
9395            a0 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(lo), _mm256_loadu_ps(xp.add(j)), a0);
9396            a1 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(hi), _mm256_loadu_ps(xp.add(j + 8)), a1);
9397            j += 16;
9398        }
9399        let acc = _mm256_add_ps(a0, a1);
9400        let hi128 = _mm256_extractf128_ps::<1>(acc);
9401        let s128 = _mm_add_ps(_mm256_castps256_ps128(acc), hi128);
9402        let s64 = _mm_add_ps(s128, _mm_movehl_ps(s128, s128));
9403        let s32 = _mm_add_ss(s64, _mm_shuffle_ps::<1>(s64, s64));
9404        let mut sum = _mm_cvtss_f32(s32);
9405        while j < n {
9406            sum += (*wp.add(j) as i8) as f32 * *xp.add(j);
9407            j += 1;
9408        }
9409        sum
9410    }
9411}
9412
9413/// int8(weight)·int8(activation) → i32 via AVX2 maddubs — the x86
9414/// analogue of the SDOT A8W8 path. `maddubs` takes u8×i8, so the
9415/// standard sign trick applies: |w| × sign(x, w) ≡ w × x per lane.
9416/// Pair saturation is safe: |w|≤128, |x|≤127 → 2·128·127 < 32767.
9417#[cfg(target_arch = "x86_64")]
9418#[target_feature(enable = "avx2")]
9419unsafe fn dot_i8_i8_avx2(w: &[u8], xq: &[i8]) -> i32 {
9420    // SAFETY: callers uphold slice-length contracts (see call sites).
9421    unsafe {
9422        use core::arch::x86_64::*;
9423        let n = w.len();
9424        let ones = _mm256_set1_epi16(1);
9425        let mut acc = _mm256_setzero_si256();
9426        let mut j = 0usize;
9427        while j + 32 <= n {
9428            let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
9429            let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
9430            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
9431            acc = _mm256_add_epi32(acc, _mm256_madd_epi16(p16, ones));
9432            j += 32;
9433        }
9434        let hi128 = _mm256_extracti128_si256::<1>(acc);
9435        let s128 = _mm_add_epi32(_mm256_castsi256_si128(acc), hi128);
9436        let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
9437        let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
9438        let mut s = _mm_cvtsi128_si32(s32);
9439        while j < n {
9440            s += (w[j] as i8) as i32 * xq[j] as i32;
9441            j += 1;
9442        }
9443        s
9444    }
9445}
9446
9447/// smmla 2×4: one instruction covers a 2-row × 2-activation × 8-deep
9448/// tile (32 MACs vs sdot's 16) — the weight pair loads once per 8-k
9449/// slice as a combined 2×8 register and meets two activation pairs.
9450#[cfg(target_arch = "aarch64")]
9451#[target_feature(enable = "neon,i8mm")]
9452unsafe fn dot_i8_smmla_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
9453    // SAFETY: callers uphold slice-length contracts.
9454    unsafe {
9455        use core::arch::aarch64::*;
9456        use core::arch::asm;
9457        let n = w0.len();
9458        let w0p = w0.as_ptr() as *const i8;
9459        let w1p = w1.as_ptr() as *const i8;
9460        // acc01 holds [c(r0,x0) c(r0,x1) c(r1,x0) c(r1,x1)]; acc23 the
9461        // same for x2/x3.
9462        let mut acc01 = vdupq_n_s32(0);
9463        let mut acc23 = vdupq_n_s32(0);
9464        let mut i = 0usize;
9465        while i + 8 <= n {
9466            let wa = vcombine_s8(vld1_s8(w0p.add(i)), vld1_s8(w1p.add(i)));
9467            let xb01 = vcombine_s8(
9468                vld1_s8(xs[0].as_ptr().add(i)),
9469                vld1_s8(xs[1].as_ptr().add(i)),
9470            );
9471            let xb23 = vcombine_s8(
9472                vld1_s8(xs[2].as_ptr().add(i)),
9473                vld1_s8(xs[3].as_ptr().add(i)),
9474            );
9475            asm!(
9476                "smmla {a01:v}.4s, {w:v}.16b, {x01:v}.16b",
9477                "smmla {a23:v}.4s, {w:v}.16b, {x23:v}.16b",
9478                a01 = inout(vreg) acc01, a23 = inout(vreg) acc23,
9479                w = in(vreg) wa, x01 = in(vreg) xb01, x23 = in(vreg) xb23,
9480                options(pure, nomem, nostack),
9481            );
9482            i += 8;
9483        }
9484        let mut out = [[0i32; 4]; 2];
9485        let a01: [i32; 4] = core::mem::transmute(acc01);
9486        let a23: [i32; 4] = core::mem::transmute(acc23);
9487        out[0][0] = a01[0];
9488        out[0][1] = a01[1];
9489        out[1][0] = a01[2];
9490        out[1][1] = a01[3];
9491        out[0][2] = a23[0];
9492        out[0][3] = a23[1];
9493        out[1][2] = a23[2];
9494        out[1][3] = a23[3];
9495        if i < n {
9496            for (k, x) in xs.iter().enumerate() {
9497                for j in i..n {
9498                    out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
9499                    out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
9500                }
9501            }
9502        }
9503        out
9504    }
9505}
9506
9507/// ARM twin of the x86 blocked prefill GEMM: two weight rows stay in
9508/// registers across four activation streams, eight sdot accumulators.
9509/// (The per-row form re-read each W row once per activation.)
9510#[cfg(target_arch = "aarch64")]
9511#[target_feature(enable = "neon,dotprod")]
9512unsafe fn dot_i8_sdot_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
9513    // SAFETY: callers uphold slice-length contracts.
9514    unsafe {
9515        use core::arch::aarch64::*;
9516        use core::arch::asm;
9517        let n = w0.len();
9518        let w0p = w0.as_ptr() as *const i8;
9519        let w1p = w1.as_ptr() as *const i8;
9520        let mut acc = [[vdupq_n_s32(0); 4]; 2];
9521        let mut i = 0usize;
9522        while i + 16 <= n {
9523            let wv0 = vld1q_s8(w0p.add(i));
9524            let wv1 = vld1q_s8(w1p.add(i));
9525            for (k, x) in xs.iter().enumerate() {
9526                let xv = vld1q_s8(x.as_ptr().add(i));
9527                let (mut a0, mut a1) = (acc[0][k], acc[1][k]);
9528                asm!(
9529                    "sdot {a0:v}.4s, {w0:v}.16b, {x:v}.16b",
9530                    "sdot {a1:v}.4s, {w1:v}.16b, {x:v}.16b",
9531                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
9532                    w0 = in(vreg) wv0, w1 = in(vreg) wv1, x = in(vreg) xv,
9533                    options(pure, nomem, nostack),
9534                );
9535                acc[0][k] = a0;
9536                acc[1][k] = a1;
9537            }
9538            i += 16;
9539        }
9540        let mut out = [[0i32; 4]; 2];
9541        for r in 0..2 {
9542            for k in 0..4 {
9543                out[r][k] = vaddvq_s32(acc[r][k]);
9544            }
9545        }
9546        if i < n {
9547            for (k, x) in xs.iter().enumerate() {
9548                for j in i..n {
9549                    out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
9550                    out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
9551                }
9552            }
9553        }
9554        out
9555    }
9556}
9557
9558/// Blocked 2 weight rows × 4 activations for the prefill GEMM
9559/// (roadmap P0: packed panels + multi-row accumulators). The two rows'
9560/// abs() live in registers across all four activation streams; the
9561/// sign-fixup is recomputed per pair (the price of the maddubs trick).
9562/// Returns raw i8·i8 dots; the caller applies scales and outliers.
9563#[cfg(target_arch = "x86_64")]
9564#[target_feature(enable = "avx2")]
9565unsafe fn dot_i8_i8_avx2_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
9566    // SAFETY: callers uphold slice-length contracts.
9567    unsafe {
9568        use core::arch::x86_64::*;
9569        let n = w0.len();
9570        let ones = _mm256_set1_epi16(1);
9571        let mut acc = [[_mm256_setzero_si256(); 4]; 2];
9572        let mut j = 0usize;
9573        while j + 32 <= n {
9574            let wv0 = _mm256_loadu_si256(w0.as_ptr().add(j) as *const __m256i);
9575            let wv1 = _mm256_loadu_si256(w1.as_ptr().add(j) as *const __m256i);
9576            let aw0 = _mm256_abs_epi8(wv0);
9577            let aw1 = _mm256_abs_epi8(wv1);
9578            for (k, x) in xs.iter().enumerate() {
9579                let xv = _mm256_loadu_si256(x.as_ptr().add(j) as *const __m256i);
9580                let p0 = _mm256_maddubs_epi16(aw0, _mm256_sign_epi8(xv, wv0));
9581                acc[0][k] = _mm256_add_epi32(acc[0][k], _mm256_madd_epi16(p0, ones));
9582                let p1 = _mm256_maddubs_epi16(aw1, _mm256_sign_epi8(xv, wv1));
9583                acc[1][k] = _mm256_add_epi32(acc[1][k], _mm256_madd_epi16(p1, ones));
9584            }
9585            j += 32;
9586        }
9587        let mut out = [[0i32; 4]; 2];
9588        for r in 0..2 {
9589            for k in 0..4 {
9590                let a = acc[r][k];
9591                let hi128 = _mm256_extracti128_si256::<1>(a);
9592                let s128 = _mm_add_epi32(_mm256_castsi256_si128(a), hi128);
9593                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
9594                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
9595                out[r][k] = _mm_cvtsi128_si32(s32);
9596            }
9597        }
9598        if j < n {
9599            for (k, x) in xs.iter().enumerate() {
9600                for i in j..n {
9601                    out[0][k] += (w0[i] as i8) as i32 * x[i] as i32;
9602                    out[1][k] += (w1[i] as i8) as i32 * x[i] as i32;
9603                }
9604            }
9605        }
9606        out
9607    }
9608}
9609
9610/// AVX2/VNNI q8 row dot with exact outlier correction (x86 mirror of
9611/// `row_dot_sdot` — same A8W8 contract). With AVX-512 VNNI the row goes
9612/// through the bias trick: Σ(w+128)·x via pure `vpdpbusd` (no per-lane
9613/// sign fixups), corrected by −128·Σx with Σx precomputed per split.
9614#[cfg(target_arch = "x86_64")]
9615#[inline]
9616fn row_dot_avx2(row: &[u8], act: &SplitAct) -> f32 {
9617    let dot = if avx512vnni_enabled() && row.len() >= 64 {
9618        (unsafe { dot_u8p128_i8_vnni(row, &act.xq) }) - 128 * act.xsum
9619    } else {
9620        unsafe { dot_i8_i8_avx2(row, &act.xq) }
9621    };
9622    let mut acc = dot as f32 * act.sx;
9623    for &(j, xv) in &act.outliers {
9624        acc += (row[j] as i8) as f32 * xv;
9625    }
9626    acc
9627}
9628
9629/// Σ (w[i]+128)·x[i] via pure `vpdpbusd` — the caller subtracts
9630/// 128·Σx. Four independent accumulators (dpbusd is ~5-cycle latency;
9631/// a single-acc loop runs latency-bound, measured on Granite Rapids).
9632#[cfg(target_arch = "x86_64")]
9633#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
9634unsafe fn dot_u8p128_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
9635    // SAFETY: callers uphold slice-length contracts (see call sites).
9636    unsafe {
9637        use core::arch::x86_64::*;
9638        let n = w.len();
9639        let flip = _mm512_set1_epi8(-128); // XOR 0x80: i8 w → u8 (w+128)
9640        #[inline(always)]
9641        unsafe fn step(
9642            w: *const u8,
9643            x: *const i8,
9644            flip: core::arch::x86_64::__m512i,
9645            acc: core::arch::x86_64::__m512i,
9646        ) -> core::arch::x86_64::__m512i {
9647            unsafe {
9648                use core::arch::x86_64::*;
9649                let wv = _mm512_xor_si512(_mm512_loadu_si512(w as *const _), flip);
9650                _mm512_dpbusd_epi32(acc, wv, _mm512_loadu_si512(x as *const _))
9651            }
9652        }
9653        let (mut a0, mut a1, mut a2, mut a3) = (
9654            _mm512_setzero_si512(),
9655            _mm512_setzero_si512(),
9656            _mm512_setzero_si512(),
9657            _mm512_setzero_si512(),
9658        );
9659        let mut j = 0usize;
9660        while j + 256 <= n {
9661            a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
9662            a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), flip, a1);
9663            a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), flip, a2);
9664            a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), flip, a3);
9665            j += 256;
9666        }
9667        while j + 64 <= n {
9668            a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
9669            j += 64;
9670        }
9671        let mut total = _mm512_reduce_add_epi32(_mm512_add_epi32(
9672            _mm512_add_epi32(a0, a1),
9673            _mm512_add_epi32(a2, a3),
9674        ));
9675        // Scalar tail: (w as i8) + 128 ≡ (w as u8) ^ 0x80.
9676        while j < n {
9677            total += ((w[j] ^ 0x80) as i32) * xq[j] as i32;
9678            j += 1;
9679        }
9680        total
9681    }
9682}
9683
9684/// One q4 row via AVX2: nibbles → centered i8 (unpacklo/hi restores the
9685/// writer's flat order, same as the NEON vzip pair), maddubs against
9686/// the pre-quantized activation group, × the group's f16 scale. Pair
9687/// saturation safe: |w|≤8, |x|≤127 → 2·8·127 ≪ 32767. Mirror of
9688/// `dot_q4_row_sdot`.
9689#[cfg(target_arch = "x86_64")]
9690#[target_feature(enable = "avx2")]
9691unsafe fn dot_q4_row_avx2(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
9692    // SAFETY: callers uphold slice-length contracts (16 packed bytes and
9693    // 2 scale bytes per group; xq.len() == gpr·GROUP_SIZE).
9694    unsafe {
9695        use core::arch::x86_64::*;
9696        let lomask = _mm_set1_epi8(0x0F);
9697        let eight = _mm256_set1_epi8(8);
9698        let ones = _mm256_set1_epi16(1);
9699        let mut acc = 0f32;
9700        for gi in 0..gpr {
9701            let g = g0 + gi;
9702            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
9703            let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
9704            let lo = _mm_and_si128(b, lomask);
9705            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
9706            let w = _mm256_sub_epi8(
9707                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
9708                eight,
9709            );
9710            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9711            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
9712            let d = _mm256_madd_epi16(p16, ones);
9713            let hi128 = _mm256_extracti128_si256::<1>(d);
9714            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
9715            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
9716            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
9717            acc += _mm_cvtsi128_si32(s32) as f32 * s;
9718        }
9719        acc
9720    }
9721}
9722
9723/// Two-activation q4 row via AVX2: nibbles unpacked ONCE per group,
9724/// both activations dotted against the same centered i8 register.
9725#[cfg(target_arch = "x86_64")]
9726#[target_feature(enable = "avx2")]
9727unsafe fn dot_q4_row_avx2_2(
9728    packed: &[u8],
9729    scales: &[u8],
9730    g0: usize,
9731    gpr: usize,
9732    xq1: &[i8],
9733    xq2: &[i8],
9734) -> (f32, f32) {
9735    // SAFETY: callers uphold slice-length contracts (see dot_q4_row_avx2).
9736    unsafe {
9737        use core::arch::x86_64::*;
9738        let lomask = _mm_set1_epi8(0x0F);
9739        let eight = _mm256_set1_epi8(8);
9740        let ones = _mm256_set1_epi16(1);
9741        let (mut acc1, mut acc2) = (0f32, 0f32);
9742        #[inline(always)]
9743        unsafe fn hsum(d: core::arch::x86_64::__m256i) -> i32 {
9744            unsafe {
9745                use core::arch::x86_64::*;
9746                let hi128 = _mm256_extracti128_si256::<1>(d);
9747                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
9748                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
9749                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
9750                _mm_cvtsi128_si32(s32)
9751            }
9752        }
9753        for gi in 0..gpr {
9754            let g = g0 + gi;
9755            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
9756            let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
9757            let lo = _mm_and_si128(b, lomask);
9758            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
9759            let w = _mm256_sub_epi8(
9760                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
9761                eight,
9762            );
9763            let aw = _mm256_abs_epi8(w);
9764            let x1 = _mm256_loadu_si256(xq1.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9765            let x2 = _mm256_loadu_si256(xq2.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
9766            let d1 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x1, w)), ones);
9767            let d2 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x2, w)), ones);
9768            acc1 += hsum(d1) as f32 * s;
9769            acc2 += hsum(d2) as f32 * s;
9770        }
9771        (acc1, acc2)
9772    }
9773}
9774
9775/// One q8 row range via AVX2 (x86 mirror of `q8_range_sdot`).
9776#[cfg(target_arch = "x86_64")]
9777fn q8_range_avx2(
9778    q: &[u8],
9779    row_scale: &[f32],
9780    act: &SplitAct,
9781    cols: usize,
9782    out_addr: SendMut,
9783    start: usize,
9784    end: usize,
9785) {
9786    for o in start..end {
9787        let v = row_dot_avx2(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
9788        // SAFETY: disjoint row ranges per worker.
9789        unsafe { *out_addr.at(o) = v };
9790    }
9791}
9792
9793/// Two-input q8 row range via AVX2 (x86 mirror of `q8_range2_sdot`).
9794#[cfg(target_arch = "x86_64")]
9795#[allow(clippy::too_many_arguments)]
9796fn q8_range2_avx2(
9797    q: &[u8],
9798    row_scale: &[f32],
9799    a1: &SplitAct,
9800    a2: &SplitAct,
9801    cols: usize,
9802    p1: SendMut,
9803    p2: SendMut,
9804    start: usize,
9805    end: usize,
9806) {
9807    for o in start..end {
9808        let row = &q[o * cols..(o + 1) * cols];
9809        // SAFETY: disjoint row ranges per worker.
9810        unsafe {
9811            *p1.at(o) = row_dot_avx2(row, a1) * row_scale[o];
9812            *p2.at(o) = row_dot_avx2(row, a2) * row_scale[o];
9813        }
9814    }
9815}
9816
9817// ───────────────────── A8W8 SDOT path (port of vmfcore, ×1.78 decode) ─────────────────────
9818
9819/// ARMv8.6 i8mm (smmla): 32 int8 MACs per instruction vs sdot's 16 —
9820/// yet MEASURED 2.4× SLOWER than the blocked sdot on Apple silicon
9821/// (108 vs 264 GF/s): the on-the-fly vcombine packing and the two-
9822/// accumulator dependency chain swamp the MAC advantage, and Apple's
9823/// four SIMD pipes already keep sdot fed. OPT-IN (CMF_I8MM=1) for
9824/// field trials on Cortex-A710/X-class parts with two pipes, where the
9825/// balance may differ; a pre-interleaved weight layout (repack infra)
9826/// is the known path if it ever earns its keep.
9827#[cfg(target_arch = "aarch64")]
9828fn i8mm_enabled() -> bool {
9829    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
9830    *ON.get_or_init(|| {
9831        std::env::var("CMF_I8MM").map(|v| v == "1").unwrap_or(false)
9832            && std::arch::is_aarch64_feature_detected!("i8mm")
9833    })
9834}
9835
9836/// SDOT enabled? Default ON when the CPU has ARMv8.2 dotprod;
9837/// `CMF_SDOT=0` disables (falls back to i8×f32 NEON).
9838/// (On non-ARM release builds only the test tolerance switch calls it.)
9839#[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
9840fn sdot_enabled() -> bool {
9841    use std::sync::OnceLock;
9842    static ON: OnceLock<bool> = OnceLock::new();
9843    *ON.get_or_init(|| {
9844        let want = std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true);
9845        if !want {
9846            return false;
9847        }
9848
9849        #[cfg(target_arch = "aarch64")]
9850        {
9851            if std::arch::is_aarch64_feature_detected!("dotprod") {
9852                return true;
9853            }
9854            #[cfg(target_os = "android")]
9855            {
9856                if let Ok(cpuinfo) = std::fs::read_to_string("/proc/cpuinfo") {
9857                    if cpuinfo.lines().any(|l| {
9858                        (l.starts_with("Features") || l.starts_with("features"))
9859                            && l.contains("asimddp")
9860                    }) {
9861                        return true;
9862                    }
9863                }
9864            }
9865            false
9866        }
9867        #[cfg(not(target_arch = "aarch64"))]
9868        {
9869            false
9870        }
9871    })
9872}
9873
9874/// Two-field activation split (≡ vmfcore `q8_split_prep`): outlier
9875/// channels (>8·rms) are computed exactly in f32; the bulk (outliers
9876/// zeroed → clean absmax) goes through int8 SDOT. Computed ONCE per
9877/// matvec, shared by all rows/workers.
9878struct SplitAct {
9879    xq: Vec<i8>,
9880    sx: f32,
9881    outliers: Vec<(usize, f32)>,
9882    /// Σ xq — the VNNI bias-trick correction (`(w+128)·x` sums need
9883    /// `−128·Σx`); one i32 per split, computed once per matvec.
9884    #[cfg_attr(not(target_arch = "x86_64"), allow(dead_code))]
9885    xsum: i32,
9886}
9887
9888thread_local! {
9889    /// Recycled xq buffers: split_act runs for every matvec (~200/token)
9890    /// and its hidden-size allocation was steady-state heap churn.
9891    static XQ_FREE: std::cell::RefCell<Vec<Vec<i8>>> =
9892        const { std::cell::RefCell::new(Vec::new()) };
9893}
9894
9895impl Drop for SplitAct {
9896    fn drop(&mut self) {
9897        let buf = std::mem::take(&mut self.xq);
9898        if buf.capacity() > 0 {
9899            XQ_FREE.with(|f| {
9900                let mut f = f.borrow_mut();
9901                if f.len() < 16 {
9902                    f.push(buf);
9903                }
9904            });
9905        }
9906    }
9907}
9908
9909thread_local! {
9910    /// One scratch row per WORKER, kept for the life of the thread.
9911    ///
9912    /// The kernels take a row of group scales per dispatch, and a fresh
9913    /// `vec![0f32; gpr]` inside the closure is one allocation per worker per
9914    /// dispatch — on the release checkpoint about six thousand a token, a
9915    /// quarter of everything the benchmark counts.
9916    static KROW: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
9917}
9918
9919/// Borrow `n` floats of the calling worker's scratch. Nothing inside a
9920/// kernel body borrows it again, which is what keeps the RefCell honest.
9921#[inline]
9922fn with_krow<R>(n: usize, f: impl FnOnce(&mut [f32]) -> R) -> R {
9923    KROW.with(|s| {
9924        let mut b = s.borrow_mut();
9925        if b.len() < n {
9926            b.resize(n, 0.0);
9927        }
9928        f(&mut b[..n])
9929    })
9930}
9931
9932/// `t.round().clamp(-127.0, 127.0) as i8`, bit for bit, without the libm
9933/// call. On baseline x86-64 (no SSE4.1 `roundps`) `f32::round` is a
9934/// function call per element, and split_act runs it over every hidden
9935/// state before every matvec: measured 27 us a call on a 2048-wide
9936/// activation on an EPYC 7763 — 5.4 ms of a 55 ms decode token, all of
9937/// it on the caller's thread while thirty workers wait. Clamping first is
9938/// equivalent (round is monotonic and ±127 are integers), and after the
9939/// clamp `t - trunc(t)` is exact, so the half-away-from-zero decision is
9940/// the one `round` makes. NaN clamps to NaN and converts to 0, as before.
9941/// The loop vectorizes (cvttps2dq + compare/select).
9942#[inline(always)]
9943fn q8_round(t: f32) -> i8 {
9944    let t = t.clamp(-127.0, 127.0);
9945    let i = t as i32;
9946    let f = t - i as f32;
9947    let r = if f >= 0.5 {
9948        i + 1
9949    } else if f <= -0.5 {
9950        i - 1
9951    } else {
9952        i
9953    };
9954    r as i8
9955}
9956
9957fn split_act(x: &[f32]) -> SplitAct {
9958    let _prof = crate::cpuprof::time(crate::cpuprof::Slot::SplitAct);
9959    let n = x.len();
9960    let rms = (x.iter().map(|&v| (v * v) as f64).sum::<f64>() / n.max(1) as f64).sqrt() as f32;
9961    let thr = 8.0 * rms;
9962    // One pass: collect outliers and the bulk absmax (outliers excluded —
9963    // identical to the old zero-then-fold over a copied buffer, minus the
9964    // full-vector copy).
9965    let mut outliers: Vec<(usize, f32)> = Vec::new();
9966    let mut amax = 0f32;
9967    for (j, &v) in x.iter().enumerate() {
9968        let a = v.abs();
9969        if a > thr {
9970            outliers.push((j, v));
9971        } else if a > amax {
9972            amax = a;
9973        }
9974    }
9975    let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
9976    let inv = 1.0 / sx;
9977    let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
9978    xq.clear();
9979    xq.reserve(n);
9980    if outliers.is_empty() {
9981        xq.extend(
9982            x.iter()
9983                .map(|&v| q8_round(v * inv)),
9984        );
9985    } else {
9986        // Outlier slots quantize to 0 (their exact term is added later).
9987        xq.extend(x.iter().map(|&v| {
9988            if v.abs() > thr {
9989                0
9990            } else {
9991                q8_round(v * inv)
9992            }
9993        }));
9994    }
9995    let xsum = xq.iter().map(|&v| v as i32).sum();
9996    SplitAct {
9997        xq,
9998        sx,
9999        outliers,
10000        xsum,
10001    }
10002}
10003
10004fn split_act_q8_2f(x: &[f32], col: &[f32]) -> SplitAct {
10005    let _prof = crate::cpuprof::time(crate::cpuprof::Slot::SplitAct);
10006    let n = x.len();
10007    let rms = (x
10008        .iter()
10009        .zip(col)
10010        .map(|(&a, &c)| {
10011            let v = a * c;
10012            (v * v) as f64
10013        })
10014        .sum::<f64>()
10015        / n.max(1) as f64)
10016        .sqrt() as f32;
10017    let thr = 8.0 * rms;
10018
10019    let mut outliers = Vec::new();
10020    let mut amax = 0f32;
10021    for (j, (&a, &c)) in x.iter().zip(col).enumerate() {
10022        let v = a * c;
10023        let s = v.abs();
10024        if s > thr {
10025            outliers.push((j, v));
10026        } else if s > amax {
10027            amax = s;
10028        }
10029    }
10030
10031    let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
10032    let inv = 1.0 / sx;
10033    let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
10034    xq.clear();
10035    xq.reserve(n);
10036    if outliers.is_empty() {
10037        xq.extend(
10038            x.iter()
10039                .zip(col)
10040                .map(|(&a, &c)| q8_round((a * c) * inv)),
10041        );
10042    } else {
10043        xq.extend(x.iter().zip(col).map(|(&a, &c)| {
10044            let v = a * c;
10045            if v.abs() > thr {
10046                0
10047            } else {
10048                q8_round(v * inv)
10049            }
10050        }));
10051    }
10052    let xsum = xq.iter().map(|&v| v as i32).sum();
10053    SplitAct {
10054        xq,
10055        sx,
10056        outliers,
10057        xsum,
10058    }
10059}
10060
10061/// int8(weight)·int8(activation) → i32 via `sdot` (inline asm — the
10062/// vdotq intrinsic is unstable; port of vmfcore `dot_i8_sdot`).
10063#[cfg(target_arch = "aarch64")]
10064#[target_feature(enable = "neon,dotprod")]
10065unsafe fn dot_i8_sdot(w: &[u8], xq: &[i8]) -> i32 {
10066    // SAFETY: callers uphold slice-length contracts (see call sites).
10067    unsafe {
10068        use core::arch::aarch64::*;
10069        use core::arch::asm;
10070        let wp = w.as_ptr() as *const i8;
10071        let n = w.len();
10072        let (mut a0, mut a1, mut a2, mut a3) = (
10073            vdupq_n_s32(0),
10074            vdupq_n_s32(0),
10075            vdupq_n_s32(0),
10076            vdupq_n_s32(0),
10077        );
10078        let mut i = 0;
10079        while i + 64 <= n {
10080            let (w0, x0) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
10081            let (w1, x1) = (vld1q_s8(wp.add(i + 16)), vld1q_s8(xq.as_ptr().add(i + 16)));
10082            let (w2, x2) = (vld1q_s8(wp.add(i + 32)), vld1q_s8(xq.as_ptr().add(i + 32)));
10083            let (w3, x3) = (vld1q_s8(wp.add(i + 48)), vld1q_s8(xq.as_ptr().add(i + 48)));
10084            asm!(
10085                "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
10086                "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
10087                "sdot {a2:v}.4s, {w2:v}.16b, {x2:v}.16b",
10088                "sdot {a3:v}.4s, {w3:v}.16b, {x3:v}.16b",
10089                a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
10090                w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
10091                w2 = in(vreg) w2, x2 = in(vreg) x2, w3 = in(vreg) w3, x3 = in(vreg) x3,
10092                options(pure, nomem, nostack),
10093            );
10094            i += 64;
10095        }
10096        while i + 16 <= n {
10097            let (wv, xv) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
10098            asm!("sdot {a:v}.4s, {w:v}.16b, {x:v}.16b",
10099                 a = inout(vreg) a0, w = in(vreg) wv, x = in(vreg) xv, options(pure, nomem, nostack));
10100            i += 16;
10101        }
10102        let mut s = vaddvq_s32(vaddq_s32(vaddq_s32(a0, a1), vaddq_s32(a2, a3)));
10103        while i < n {
10104            s += (*wp.add(i)) as i32 * xq[i] as i32;
10105            i += 1;
10106        }
10107        s
10108    }
10109}
10110
10111/// Row-blocked SDOT: 4 output rows per pass — the activation chunk is
10112/// loaded once and reused, 4 independent accumulators hide sdot latency
10113/// (port of vmfcore `dot_i8_sdot_4rows`).
10114#[cfg(target_arch = "aarch64")]
10115#[target_feature(enable = "neon,dotprod")]
10116unsafe fn dot_i8_sdot_4rows(w0: &[u8], w1: &[u8], w2: &[u8], w3: &[u8], xq: &[i8]) -> [i32; 4] {
10117    // SAFETY: callers uphold slice-length contracts (see call sites).
10118    unsafe {
10119        use core::arch::aarch64::*;
10120        use core::arch::asm;
10121        let n = xq.len();
10122        let px = xq.as_ptr();
10123        let (p0, p1, p2, p3) = (
10124            w0.as_ptr() as *const i8,
10125            w1.as_ptr() as *const i8,
10126            w2.as_ptr() as *const i8,
10127            w3.as_ptr() as *const i8,
10128        );
10129        let (mut a0, mut a1, mut a2, mut a3) = (
10130            vdupq_n_s32(0),
10131            vdupq_n_s32(0),
10132            vdupq_n_s32(0),
10133            vdupq_n_s32(0),
10134        );
10135        let mut i = 0;
10136        while i + 16 <= n {
10137            let x = vld1q_s8(px.add(i));
10138            let v0 = vld1q_s8(p0.add(i));
10139            let v1 = vld1q_s8(p1.add(i));
10140            let v2 = vld1q_s8(p2.add(i));
10141            let v3 = vld1q_s8(p3.add(i));
10142            asm!(
10143                "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
10144                "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
10145                "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
10146                "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
10147                a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
10148                v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
10149                options(pure, nomem, nostack),
10150            );
10151            i += 16;
10152        }
10153        let mut r = [
10154            vaddvq_s32(a0),
10155            vaddvq_s32(a1),
10156            vaddvq_s32(a2),
10157            vaddvq_s32(a3),
10158        ];
10159        while i < n {
10160            let xi = *px.add(i) as i32;
10161            r[0] += (*p0.add(i)) as i32 * xi;
10162            r[1] += (*p1.add(i)) as i32 * xi;
10163            r[2] += (*p2.add(i)) as i32 * xi;
10164            r[3] += (*p3.add(i)) as i32 * xi;
10165            i += 1;
10166        }
10167        r
10168    }
10169}
10170
10171/// 4 interleaved rows in one pass: the repacked group is [r0[c], r1[c],
10172/// r2[c], r3[c]] per 16-byte chunk, so each iteration reads ONE 64-byte
10173/// line plus the shared activation chunk — a single sequential weight
10174/// stream per worker. Per-row accumulation is the same one-accumulator
10175/// scheme as `dot_i8_sdot_4rows`; integer sums are exact, so outputs
10176/// are bit-identical to the mmap-layout kernel.
10177#[cfg(target_arch = "aarch64")]
10178#[target_feature(enable = "neon,dotprod")]
10179unsafe fn dot_i8_sdot_4rows_il(g: &[u8], xq: &[i8]) -> [i32; 4] {
10180    // SAFETY: callers uphold slice-length contracts (g.len() == 4·n,
10181    // n % 16 == 0 — guaranteed by the repack gate).
10182    unsafe {
10183        use core::arch::aarch64::*;
10184        use core::arch::asm;
10185        let n = xq.len();
10186        let px = xq.as_ptr();
10187        let pg = g.as_ptr() as *const i8;
10188        let (mut a0, mut a1, mut a2, mut a3) = (
10189            vdupq_n_s32(0),
10190            vdupq_n_s32(0),
10191            vdupq_n_s32(0),
10192            vdupq_n_s32(0),
10193        );
10194        let mut i = 0;
10195        while i + 16 <= n {
10196            let x = vld1q_s8(px.add(i));
10197            let base = pg.add(4 * i);
10198            let v0 = vld1q_s8(base);
10199            let v1 = vld1q_s8(base.add(16));
10200            let v2 = vld1q_s8(base.add(32));
10201            let v3 = vld1q_s8(base.add(48));
10202            asm!(
10203                "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
10204                "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
10205                "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
10206                "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
10207                a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
10208                v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
10209                options(pure, nomem, nostack),
10210            );
10211            i += 16;
10212        }
10213        [
10214            vaddvq_s32(a0),
10215            vaddvq_s32(a1),
10216            vaddvq_s32(a2),
10217            vaddvq_s32(a3),
10218        ]
10219    }
10220}
10221
10222/// One q8 row range via SDOT (4-row blocks + tail) — the body of
10223/// `qmatvec`'s hot loop, extracted so multi-matrix jobs can drive the
10224/// SAME kernel for several tensors under one pool dispatch. `rep` — the
10225/// load-time interleaved repack (empty = mmap layout only); rows outside
10226/// full 4-row groups always come from the mmap layout.
10227#[cfg(target_arch = "aarch64")]
10228fn q8_range_sdot(
10229    q: &[u8],
10230    rep: &[u8],
10231    row_scale: &[f32],
10232    act: &SplitAct,
10233    cols: usize,
10234    out_addr: SendMut,
10235    start: usize,
10236    end: usize,
10237) {
10238    let mut o = start;
10239    // Leading rows to the group boundary (repack path only): the pool
10240    // splits row ranges arbitrarily, groups are absolute.
10241    if !rep.is_empty() {
10242        while o < end && o % 4 != 0 {
10243            let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
10244            unsafe { *out_addr.at(o) = v };
10245            o += 1;
10246        }
10247    }
10248    while o + 4 <= end {
10249        let r = if rep.is_empty() {
10250            unsafe {
10251                dot_i8_sdot_4rows(
10252                    &q[o * cols..(o + 1) * cols],
10253                    &q[(o + 1) * cols..(o + 2) * cols],
10254                    &q[(o + 2) * cols..(o + 3) * cols],
10255                    &q[(o + 3) * cols..(o + 4) * cols],
10256                    &act.xq,
10257                )
10258            }
10259        } else {
10260            unsafe { dot_i8_sdot_4rows_il(&rep[o * cols..(o + 4) * cols], &act.xq) }
10261        };
10262        for k in 0..4 {
10263            let mut acc = r[k] as f32 * act.sx;
10264            for &(j, xv) in &act.outliers {
10265                acc += (q[(o + k) * cols + j] as i8) as f32 * xv;
10266            }
10267            // SAFETY: disjoint row ranges per worker.
10268            unsafe { *out_addr.at(o + k) = acc * row_scale[o + k] };
10269        }
10270        o += 4;
10271    }
10272    while o < end {
10273        let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
10274        unsafe { *out_addr.at(o) = v };
10275        o += 1;
10276    }
10277}
10278
10279/// Two-input q8 row range via SDOT — `qmatvec2`'s hot loop, extracted
10280/// for the fused pair multi-matrix job (`matvec2_many`).
10281#[cfg(target_arch = "aarch64")]
10282#[allow(clippy::too_many_arguments)]
10283fn q8_range2_sdot(
10284    q: &[u8],
10285    row_scale: &[f32],
10286    a1: &SplitAct,
10287    a2: &SplitAct,
10288    cols: usize,
10289    p1: SendMut,
10290    p2: SendMut,
10291    start: usize,
10292    end: usize,
10293) {
10294    for o in start..end {
10295        let row = &q[o * cols..(o + 1) * cols];
10296        // SAFETY: disjoint row ranges per worker.
10297        unsafe {
10298            *p1.at(o) = row_dot_sdot(row, a1) * row_scale[o];
10299            *p2.at(o) = row_dot_sdot(row, a2) * row_scale[o];
10300        }
10301    }
10302}
10303
10304/// Two-input q8 row range, f32 kernel (non-SDOT) — same extraction.
10305#[allow(clippy::too_many_arguments)]
10306fn q8_range2_f32(
10307    q: &[u8],
10308    row_scale: &[f32],
10309    x1: &[f32],
10310    x2: &[f32],
10311    cols: usize,
10312    p1: SendMut,
10313    p2: SendMut,
10314    start: usize,
10315    end: usize,
10316) {
10317    for o in start..end {
10318        let row = &q[o * cols..(o + 1) * cols];
10319        // SAFETY: disjoint row ranges per worker.
10320        unsafe {
10321            *p1.at(o) = dot_i8_f32(row, x1) * row_scale[o];
10322            *p2.at(o) = dot_i8_f32(row, x2) * row_scale[o];
10323        }
10324    }
10325}
10326
10327/// Scalar/NEON-f32 q8 row range (non-SDOT platforms) — same extraction.
10328fn q8_range_f32(
10329    q: &[u8],
10330    row_scale: &[f32],
10331    xs: &[f32],
10332    cols: usize,
10333    out_addr: SendMut,
10334    start: usize,
10335    end: usize,
10336) {
10337    for o in start..end {
10338        let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
10339        // SAFETY: disjoint row ranges per worker.
10340        unsafe { *out_addr.at(o) = v };
10341    }
10342}
10343
10344/// One q8 row against a split activation, portable: the per-arch fast
10345/// dots where they exist, the exact scalar loop elsewhere. The scalar
10346/// arm is also the test oracle for both fast arms.
10347#[inline]
10348fn q8_row_dot(row: &[u8], act: &SplitAct) -> f32 {
10349    #[cfg(target_arch = "aarch64")]
10350    return row_dot_sdot(row, act);
10351    #[cfg(target_arch = "x86_64")]
10352    return row_dot_avx2(row, act);
10353    #[allow(unreachable_code)]
10354    q8_row_dot_scalar(row, act)
10355}
10356
10357#[allow(dead_code)]
10358fn q8_row_dot_scalar(row: &[u8], act: &SplitAct) -> f32 {
10359    let mut acc = 0i32;
10360    for (k, &b) in row.iter().enumerate() {
10361        acc += (b as i8) as i32 * act.xq[k] as i32;
10362    }
10363    let mut acc = acc as f32 * act.sx;
10364    for &(j, xv) in &act.outliers {
10365        acc += (row[j] as i8) as f32 * xv;
10366    }
10367    acc
10368}
10369
10370/// SDOT row dot with exact outlier correction:
10371/// `dot = sdot(w, xq)·sx + Σ_outl w[j]·x[j]` (then × row_scale by caller).
10372#[cfg(target_arch = "aarch64")]
10373#[inline]
10374fn row_dot_sdot(row: &[u8], act: &SplitAct) -> f32 {
10375    let mut acc = unsafe { dot_i8_sdot(row, &act.xq) } as f32 * act.sx;
10376    for &(j, xv) in &act.outliers {
10377        acc += (row[j] as i8) as f32 * xv;
10378    }
10379    acc
10380}
10381
10382/// One q4 row via SDOT: each 32-group's nibbles unpack to centered i8
10383/// (nib−8 ∈ [−8,7]), int8×int8 `sdot` against the pre-quantized
10384/// activation group, × the group's f16 scale. Returns Σ_g dot_g·s_g;
10385/// the caller multiplies by the activation scale and adds the exact
10386/// outlier terms (port of vmfcore `dot_q4_block_sdot`, +23% measured).
10387/// Nibble order matches the writer: element 2k = low nibble, 2k+1 = high
10388/// → zip(lo,hi) restores flat order.
10389#[cfg(target_arch = "aarch64")]
10390#[target_feature(enable = "neon,dotprod")]
10391unsafe fn dot_q4_row_sdot(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
10392    // SAFETY: callers uphold slice-length contracts (16 packed bytes and
10393    // 2 scale bytes per group; xq.len() == gpr·GROUP_SIZE).
10394    unsafe {
10395        use core::arch::aarch64::*;
10396        use core::arch::asm;
10397        let lomask = vdupq_n_u8(0x0F);
10398        let eight = vdupq_n_s8(8);
10399        let mut acc = 0f32;
10400        for gi in 0..gpr {
10401            let g = g0 + gi;
10402            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
10403            let b = vld1q_u8(packed.as_ptr().add(g * 16));
10404            let lo = vandq_u8(b, lomask);
10405            let hi = vshrq_n_u8::<4>(b);
10406            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
10407            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
10408            let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
10409            let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
10410            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
10411            asm!(
10412                "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
10413                "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
10414                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
10415                e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
10416                options(pure, nomem, nostack),
10417            );
10418            acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
10419        }
10420        acc
10421    }
10422}
10423
10424/// Two-activation q4 row via SDOT: the nibble unpack (the expensive
10425/// part) happens ONCE per group; both pre-quantized activations are
10426/// dotted against the same centered i8 registers. Per-lane math matches
10427/// `dot_q4_row_sdot` exactly.
10428#[cfg(target_arch = "aarch64")]
10429#[target_feature(enable = "neon,dotprod")]
10430unsafe fn dot_q4_row_sdot2(
10431    packed: &[u8],
10432    scales: &[u8],
10433    g0: usize,
10434    gpr: usize,
10435    xq1: &[i8],
10436    xq2: &[i8],
10437) -> (f32, f32) {
10438    // SAFETY: callers uphold slice-length contracts (16 packed bytes and
10439    // 2 scale bytes per group; xq*.len() == gpr·GROUP_SIZE).
10440    unsafe {
10441        use core::arch::aarch64::*;
10442        use core::arch::asm;
10443        let lomask = vdupq_n_u8(0x0F);
10444        let eight = vdupq_n_s8(8);
10445        let (mut acc1, mut acc2) = (0f32, 0f32);
10446        for gi in 0..gpr {
10447            let g = g0 + gi;
10448            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
10449            let b = vld1q_u8(packed.as_ptr().add(g * 16));
10450            let lo = vandq_u8(b, lomask);
10451            let hi = vshrq_n_u8::<4>(b);
10452            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
10453            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
10454            let x10 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE));
10455            let x11 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE + 16));
10456            let x20 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE));
10457            let x21 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE + 16));
10458            let (mut a0, mut a1, mut b0, mut b1) = (
10459                vdupq_n_s32(0),
10460                vdupq_n_s32(0),
10461                vdupq_n_s32(0),
10462                vdupq_n_s32(0),
10463            );
10464            asm!(
10465                "sdot {a0:v}.4s, {e0:v}.16b, {x10:v}.16b",
10466                "sdot {a1:v}.4s, {e1:v}.16b, {x11:v}.16b",
10467                "sdot {b0:v}.4s, {e0:v}.16b, {x20:v}.16b",
10468                "sdot {b1:v}.4s, {e1:v}.16b, {x21:v}.16b",
10469                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
10470                b0 = inout(vreg) b0, b1 = inout(vreg) b1,
10471                e0 = in(vreg) e0, e1 = in(vreg) e1,
10472                x10 = in(vreg) x10, x11 = in(vreg) x11,
10473                x20 = in(vreg) x20, x21 = in(vreg) x21,
10474                options(pure, nomem, nostack),
10475            );
10476            acc1 += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
10477            acc2 += vaddvq_s32(vaddq_s32(b0, b1)) as f32 * s;
10478        }
10479        (acc1, acc2)
10480    }
10481}
10482
10483// ───────────────────── fused int8 kernels ─────────────────────
10484
10485/// `acc += w · row` where the row is centered i8 — NEON widen+fma on
10486/// aarch64, scalar elsewhere. The KV-cache q8 value path rides on this.
10487#[inline]
10488pub(crate) fn axpy_i8_f32(acc: &mut [f32], row: &[i8], w: f32) {
10489    #[cfg(target_arch = "aarch64")]
10490    unsafe {
10491        return axpy_i8_f32_neon(acc, row, w);
10492    }
10493    #[cfg(target_arch = "x86_64")]
10494    if avx2_enabled() {
10495        return unsafe { axpy_i8_f32_avx2(acc, row, w) };
10496    }
10497    #[allow(unreachable_code)]
10498    {
10499        for (a, &b) in acc.iter_mut().zip(row) {
10500            *a += w * b as f32;
10501        }
10502    }
10503}
10504
10505/// i8→f32 axpy via AVX2/FMA (x86 mirror of `axpy_i8_f32_neon`).
10506#[cfg(target_arch = "x86_64")]
10507#[target_feature(enable = "avx2,fma")]
10508unsafe fn axpy_i8_f32_avx2(acc: &mut [f32], row: &[i8], w: f32) {
10509    // SAFETY: callers uphold slice-length contracts (see call sites).
10510    unsafe {
10511        use core::arch::x86_64::*;
10512        let n = acc.len().min(row.len());
10513        let ap = acc.as_mut_ptr();
10514        let rp = row.as_ptr();
10515        let wv = _mm256_set1_ps(w);
10516        let mut j = 0usize;
10517        while j + 16 <= n {
10518            let rb = _mm_loadu_si128(rp.add(j) as *const __m128i);
10519            let lo = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(rb));
10520            let hi = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(rb)));
10521            let v0 = _mm256_fmadd_ps(wv, lo, _mm256_loadu_ps(ap.add(j)));
10522            let v1 = _mm256_fmadd_ps(wv, hi, _mm256_loadu_ps(ap.add(j + 8)));
10523            _mm256_storeu_ps(ap.add(j), v0);
10524            _mm256_storeu_ps(ap.add(j + 8), v1);
10525            j += 16;
10526        }
10527        while j < n {
10528            *ap.add(j) += w * (*rp.add(j)) as f32;
10529            j += 1;
10530        }
10531    }
10532}
10533
10534#[cfg(target_arch = "aarch64")]
10535#[target_feature(enable = "neon")]
10536unsafe fn axpy_i8_f32_neon(acc: &mut [f32], row: &[i8], w: f32) {
10537    // SAFETY: callers uphold slice-length contracts (see call sites).
10538    unsafe {
10539        use core::arch::aarch64::*;
10540        let n = acc.len().min(row.len());
10541        let ap = acc.as_mut_ptr();
10542        let rp = row.as_ptr();
10543        let wv = vdupq_n_f32(w);
10544        let mut j = 0usize;
10545        while j + 16 <= n {
10546            let rb = vld1q_s8(rp.add(j));
10547            let lo = vmovl_s8(vget_low_s8(rb));
10548            let hi = vmovl_s8(vget_high_s8(rb));
10549            for (off, half) in [(0, lo), (8, hi)] {
10550                let f0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(half)));
10551                let f1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(half)));
10552                let o = j + off;
10553                vst1q_f32(ap.add(o), vfmaq_f32(vld1q_f32(ap.add(o)), wv, f0));
10554                vst1q_f32(ap.add(o + 4), vfmaq_f32(vld1q_f32(ap.add(o + 4)), wv, f1));
10555            }
10556            j += 16;
10557        }
10558        while j < n {
10559            *ap.add(j) += w * (*rp.add(j)) as f32;
10560            j += 1;
10561        }
10562    }
10563}
10564
10565/// i8 row · f32 x. NEON on aarch64 (ported from vmfcore `dot_i8_f32_neon`,
10566/// ≈9× scalar), scalar elsewhere.
10567#[inline]
10568pub(crate) fn dot_i8_f32(w: &[u8], x: &[f32]) -> f32 {
10569    #[cfg(target_arch = "aarch64")]
10570    unsafe {
10571        return dot_i8_f32_neon(w, x);
10572    }
10573    #[cfg(target_arch = "x86_64")]
10574    if avx2_enabled() {
10575        return unsafe { dot_i8_f32_avx2(w, x) };
10576    }
10577    #[allow(unreachable_code)]
10578    {
10579        let mut sum = 0.0f32;
10580        for (j, &b) in w.iter().enumerate() {
10581            sum += (b as i8) as f32 * x[j];
10582        }
10583        sum
10584    }
10585}
10586
10587/// i8 row · (x ⊙ col_field) — the q8_2f row dot with the θ col-field
10588/// folded into the product (no prescaled copy of x). NEON on aarch64,
10589/// scalar elsewhere. Used by the active-neuron path `row_dot`.
10590#[inline]
10591fn dot_i8_col_f32(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
10592    #[cfg(target_arch = "aarch64")]
10593    unsafe {
10594        return dot_i8_col_f32_neon(w, x, col);
10595    }
10596    #[allow(unreachable_code)]
10597    {
10598        let mut sum = 0.0f32;
10599        for (j, &b) in w.iter().enumerate() {
10600            sum += (b as i8) as f32 * x[j] * col[j];
10601        }
10602        sum
10603    }
10604}
10605
10606#[cfg(target_arch = "aarch64")]
10607#[target_feature(enable = "neon")]
10608unsafe fn dot_i8_col_f32_neon(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
10609    // SAFETY: callers uphold slice-length contracts (see call sites).
10610    unsafe {
10611        use core::arch::aarch64::*;
10612        let n = x.len();
10613        let wp = w.as_ptr() as *const i8;
10614        let xp = x.as_ptr();
10615        let cp = col.as_ptr();
10616        let (mut a0, mut a1, mut a2, mut a3) = (
10617            vdupq_n_f32(0.0),
10618            vdupq_n_f32(0.0),
10619            vdupq_n_f32(0.0),
10620            vdupq_n_f32(0.0),
10621        );
10622        let mut j = 0usize;
10623        while j + 16 <= n {
10624            let wb = vld1q_s8(wp.add(j));
10625            let lo = vmovl_s8(vget_low_s8(wb));
10626            let hi = vmovl_s8(vget_high_s8(wb));
10627            let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
10628            let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
10629            let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
10630            let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
10631            a0 = vfmaq_f32(
10632                a0,
10633                w0,
10634                vmulq_f32(vld1q_f32(xp.add(j)), vld1q_f32(cp.add(j))),
10635            );
10636            a1 = vfmaq_f32(
10637                a1,
10638                w1,
10639                vmulq_f32(vld1q_f32(xp.add(j + 4)), vld1q_f32(cp.add(j + 4))),
10640            );
10641            a2 = vfmaq_f32(
10642                a2,
10643                w2,
10644                vmulq_f32(vld1q_f32(xp.add(j + 8)), vld1q_f32(cp.add(j + 8))),
10645            );
10646            a3 = vfmaq_f32(
10647                a3,
10648                w3,
10649                vmulq_f32(vld1q_f32(xp.add(j + 12)), vld1q_f32(cp.add(j + 12))),
10650            );
10651            j += 16;
10652        }
10653        let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
10654        while j < n {
10655            sum += (*wp.add(j)) as f32 * *xp.add(j) * *cp.add(j);
10656            j += 1;
10657        }
10658        sum
10659    }
10660}
10661
10662#[cfg(target_arch = "aarch64")]
10663#[target_feature(enable = "neon")]
10664unsafe fn dot_i8_f32_neon(w: &[u8], x: &[f32]) -> f32 {
10665    // SAFETY: callers uphold slice-length contracts (see call sites).
10666    unsafe {
10667        use core::arch::aarch64::*;
10668        let n = x.len();
10669        let wp = w.as_ptr() as *const i8;
10670        let xp = x.as_ptr();
10671        let (mut a0, mut a1, mut a2, mut a3) = (
10672            vdupq_n_f32(0.0),
10673            vdupq_n_f32(0.0),
10674            vdupq_n_f32(0.0),
10675            vdupq_n_f32(0.0),
10676        );
10677        let mut j = 0usize;
10678        while j + 16 <= n {
10679            let wb = vld1q_s8(wp.add(j));
10680            let lo = vmovl_s8(vget_low_s8(wb));
10681            let hi = vmovl_s8(vget_high_s8(wb));
10682            let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
10683            let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
10684            let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
10685            let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
10686            a0 = vfmaq_f32(a0, w0, vld1q_f32(xp.add(j)));
10687            a1 = vfmaq_f32(a1, w1, vld1q_f32(xp.add(j + 4)));
10688            a2 = vfmaq_f32(a2, w2, vld1q_f32(xp.add(j + 8)));
10689            a3 = vfmaq_f32(a3, w3, vld1q_f32(xp.add(j + 12)));
10690            j += 16;
10691        }
10692        let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
10693        while j < n {
10694            sum += (*wp.add(j)) as f32 * *xp.add(j);
10695            j += 1;
10696        }
10697        sum
10698    }
10699}
10700
10701#[allow(clippy::too_many_arguments)]
10702fn qmatvec(
10703    q: &[u8],
10704    rep: &[u8],
10705    row_scale: &[f32],
10706    x: &[f32],
10707    col_field: &[f32],
10708    dtype: TensorDtype,
10709    rows: usize,
10710    cols: usize,
10711    out: &mut [f32],
10712    pool: Option<&Pool>,
10713) {
10714    debug_assert_eq!(out.len(), rows);
10715    #[cfg(not(target_arch = "aarch64"))]
10716    let _ = rep;
10717
10718    #[cfg(target_arch = "aarch64")]
10719    if sdot_enabled() {
10720        let act = if dtype == TensorDtype::Q8_2f {
10721            split_act_q8_2f(x, col_field)
10722        } else {
10723            split_act(x)
10724        };
10725        let out_addr = SendMut(out.as_mut_ptr());
10726        let run_range = |start: usize, end: usize| {
10727            q8_range_sdot(q, rep, row_scale, &act, cols, out_addr, start, end)
10728        };
10729        match pool {
10730            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
10731            _ => run_range(0, rows),
10732        }
10733        return;
10734    }
10735    // x86 A8W8 via AVX2 maddubs — same quantized-activation contract as
10736    // the SDOT path (CMF_AVX2=0 keeps the exact i8×f32 loop).
10737    #[cfg(target_arch = "x86_64")]
10738    if avx2_a8w8_enabled() {
10739        let act = if dtype == TensorDtype::Q8_2f {
10740            split_act_q8_2f(x, col_field)
10741        } else {
10742            split_act(x)
10743        };
10744        let out_addr = SendMut(out.as_mut_ptr());
10745        let run_range = |start: usize, end: usize| {
10746            q8_range_avx2(q, row_scale, &act, cols, out_addr, start, end)
10747        };
10748        match pool {
10749            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
10750            _ => run_range(0, rows),
10751        }
10752        return;
10753    }
10754
10755    prescale_with(x, col_field, dtype, 1, |xs| {
10756        let out_addr = SendMut(out.as_mut_ptr());
10757        let run_range = move |start: usize, end: usize| {
10758            for o in start..end {
10759                let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
10760                // SAFETY: disjoint row ranges per worker.
10761                unsafe { *out_addr.at(o) = v };
10762            }
10763        };
10764        match pool {
10765            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
10766            _ => run_range(0, rows),
10767        }
10768    });
10769}
10770
10771#[allow(clippy::too_many_arguments)]
10772fn qmatvec2(
10773    q: &[u8],
10774    row_scale: &[f32],
10775    x1: &[f32],
10776    x2: &[f32],
10777    col_field: &[f32],
10778    dtype: TensorDtype,
10779    rows: usize,
10780    cols: usize,
10781    o1: &mut [f32],
10782    o2: &mut [f32],
10783    pool: Option<&Pool>,
10784) {
10785    #[cfg(target_arch = "aarch64")]
10786    if sdot_enabled() {
10787        let a1s = if dtype == TensorDtype::Q8_2f {
10788            split_act_q8_2f(x1, col_field)
10789        } else {
10790            split_act(x1)
10791        };
10792        let a2s = if dtype == TensorDtype::Q8_2f {
10793            split_act_q8_2f(x2, col_field)
10794        } else {
10795            split_act(x2)
10796        };
10797        let p1 = SendMut(o1.as_mut_ptr());
10798        let p2 = SendMut(o2.as_mut_ptr());
10799        let run_range = |start: usize, end: usize| {
10800            q8_range2_sdot(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
10801        };
10802        match pool {
10803            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
10804            _ => run_range(0, rows),
10805        }
10806        return;
10807    }
10808    #[cfg(target_arch = "x86_64")]
10809    if avx2_a8w8_enabled() {
10810        let a1s = if dtype == TensorDtype::Q8_2f {
10811            split_act_q8_2f(x1, col_field)
10812        } else {
10813            split_act(x1)
10814        };
10815        let a2s = if dtype == TensorDtype::Q8_2f {
10816            split_act_q8_2f(x2, col_field)
10817        } else {
10818            split_act(x2)
10819        };
10820        let p1 = SendMut(o1.as_mut_ptr());
10821        let p2 = SendMut(o2.as_mut_ptr());
10822        let run_range = |start: usize, end: usize| {
10823            q8_range2_avx2(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
10824        };
10825        match pool {
10826            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
10827            _ => run_range(0, rows),
10828        }
10829        return;
10830    }
10831
10832    prescale_with(x1, col_field, dtype, 1, |x1s| {
10833        prescale_with(x2, col_field, dtype, 2, |x2s| {
10834            let p1 = SendMut(o1.as_mut_ptr());
10835            let p2 = SendMut(o2.as_mut_ptr());
10836            let run_range = move |start: usize, end: usize| {
10837                for o in start..end {
10838                    let row = &q[o * cols..(o + 1) * cols];
10839                    let s1 = dot_i8_f32(row, x1s) * row_scale[o];
10840                    let s2 = dot_i8_f32(row, x2s) * row_scale[o];
10841                    // SAFETY: disjoint row ranges per worker.
10842                    unsafe {
10843                        *p1.at(o) = s1;
10844                        *p2.at(o) = s2;
10845                    }
10846                }
10847            };
10848            match pool {
10849                Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
10850                _ => run_range(0, rows),
10851            }
10852        });
10853    });
10854}
10855
10856#[derive(Clone, Copy)]
10857struct SendMut(*mut f32);
10858unsafe impl Send for SendMut {}
10859unsafe impl Sync for SendMut {}
10860
10861impl SendMut {
10862    #[inline]
10863    fn at(self, i: usize) -> *mut f32 {
10864        unsafe { self.0.add(i) }
10865    }
10866}
10867
10868#[cfg(test)]
10869mod tests {
10870    /// `q8_round` must be `round().clamp(±127) as i8` bit for bit: every
10871    /// half-integer, their neighbours one ulp either side, the clamp
10872    /// boundary, huge values, infinities and NaN, plus a dense sweep.
10873    #[test]
10874    fn q8_round_is_round_clamp() {
10875        let reference = |t: f32| t.round().clamp(-127.0, 127.0) as i8;
10876        let mut probe = vec![
10877            0.0f32,
10878            -0.0,
10879            f32::NAN,
10880            f32::INFINITY,
10881            f32::NEG_INFINITY,
10882            f32::MAX,
10883            f32::MIN,
10884            1e30,
10885            -1e30,
10886            f32::MIN_POSITIVE,
10887            -f32::MIN_POSITIVE,
10888        ];
10889        for k in -300i32..=300 {
10890            let h = k as f32 * 0.5;
10891            let up = f32::from_bits(h.to_bits() + 1);
10892            let down = f32::from_bits(h.to_bits().wrapping_sub(1));
10893            for t in [h, up, down] {
10894                probe.push(t);
10895                probe.push(-t);
10896            }
10897        }
10898        let mut t = -140.0f32;
10899        while t < 140.0 {
10900            probe.push(t);
10901            t += 0.000_731;
10902        }
10903        for t in probe {
10904            assert_eq!(q8_round(t), reference(t), "t = {t:e} ({:#x})", t.to_bits());
10905        }
10906    }
10907
10908    use super::*;
10909
10910    #[test]
10911    fn q2tp_i8_dot_matches_exact_on_grid() {
10912        // On-grid activations (±1 → sx=1/127, xq=±127 dequantizes
10913        // exactly, no outliers) must make the integer path agree with
10914        // the exact scalar walk to f32 rounding.
10915        let (rows, cols) = (5, 64);
10916        let gpr = cols / GROUP_SIZE;
10917        // Synthetic codes plane + a flat ladder: scales_into is not under
10918        // test here, so drive dot_q2tp_row_i8 / q2tp_row_exact directly
10919        // with hand-made scales.
10920        let chunks: Vec<u8> = (0..rows * gpr * Q2TP_CHUNK)
10921            .map(|i| (i as u32).wrapping_mul(2654435761) as u8)
10922            .collect();
10923        let scales: Vec<f32> = (0..gpr).map(|g| 0.5 + g as f32 * 0.25).collect();
10924        let x: Vec<f32> = (0..cols)
10925            .map(|i| if i % 3 == 0 { -1.0 } else { 1.0 })
10926            .collect();
10927        let act = split_act(&x);
10928        assert!(
10929            act.outliers.is_empty(),
10930            "on-grid input must have no outliers"
10931        );
10932        let gsum = q1_group_sums(&act.xq, gpr);
10933        for r in 0..rows {
10934            let exact = q2tp_row_exact(&chunks, r, gpr, &x, &scales);
10935            let fast = dot_q2tp_row_i8(&chunks, r, gpr, &act.xq, &gsum, &scales) * act.sx;
10936            assert!(
10937                (exact - fast).abs() <= exact.abs() * 1e-5 + 1e-5,
10938                "row {r}: exact {exact} vs i8 {fast}"
10939            );
10940        }
10941    }
10942
10943    #[test]
10944    fn q2tp_affine_fuses_half_scale_correction_without_changing_raw_decode() {
10945        let (rows, cols) = (1usize, GROUP_SIZE);
10946        let mut bytes = vec![0u8; Q2TP_CHUNK + 4 + 1];
10947        // Repeating symbols 0,1,2,0 at unit scale.  q2tp's raw B is
10948        // (c-1.5), while the affine Prism operator is (c-1.0).
10949        bytes[..Q2TP_CHUNK].fill(0x24); // codes 0,1,2,0 in LSB-first order
10950        bytes[Q2TP_CHUNK..Q2TP_CHUNK + 2].copy_from_slice(&0u16.to_le_bytes());
10951        bytes[Q2TP_CHUNK + 2..Q2TP_CHUNK + 4].copy_from_slice(&0u16.to_le_bytes());
10952        bytes[Q2TP_CHUNK + 4] = 1; // dtype16 rung 1 = 1.0
10953        let x = vec![1.0f32; cols];
10954        let mut raw = vec![0.0f32; rows];
10955        let mut affine = vec![0.0f32; rows];
10956        q2tp_matvec_for_test(&bytes, &x, rows, cols, &mut raw);
10957        q2tp_affine_matvec_for_test(&bytes, &x, rows, cols, &mut affine);
10958        assert_eq!(raw, vec![-24.0]);
10959        assert_eq!(affine, vec![-8.0]);
10960        assert!((affine[0] - (raw[0] + 0.5 * cols as f32)).abs() < 1e-6);
10961    }
10962
10963    #[cfg(target_arch = "x86_64")]
10964    #[test]
10965    fn q2tp_avx2_dot_matches_scalar_for_random_patterns() {
10966        // Compare the release AVX2 integer dot against the scalar oracle over
10967        // arbitrary packed bytes/activation signs.  This guards the exact
10968        // table-load path used after rejecting a faster-looking decoder whose
10969        // full-checkpoint greedy output drifted.
10970        if !std::arch::is_x86_feature_detected!("avx2") {
10971            return;
10972        }
10973        let mut seed = 0x9e3779b9u32;
10974        let mut next = || {
10975            seed = seed.wrapping_mul(1664525).wrapping_add(1013904223);
10976            seed
10977        };
10978        for _ in 0..20_000 {
10979            let mut ch = [0u8; Q2TP_CHUNK];
10980            let mut x = [0i8; GROUP_SIZE];
10981            for b in &mut ch {
10982                *b = next() as u8;
10983            }
10984            for v in &mut x {
10985                *v = (next() >> 24) as i8;
10986            }
10987            let mut reference = 0i32;
10988            for (k, &b) in ch.iter().enumerate() {
10989                reference += (b & 3) as i32 * x[k * 4] as i32;
10990                reference += ((b >> 2) & 3) as i32 * x[k * 4 + 1] as i32;
10991                reference += ((b >> 4) & 3) as i32 * x[k * 4 + 2] as i32;
10992                reference += ((b >> 6) & 3) as i32 * x[k * 4 + 3] as i32;
10993            }
10994            // SAFETY: guarded by the runtime AVX2 feature check and fixed
10995            // 8-byte/32-byte slice lengths above.
10996            let got = unsafe { q2tp_code_dot_avx2(&ch, &x) };
10997            assert_eq!(got, reference, "packed q2 lane mismatch");
10998        }
10999    }
11000
11001    #[test]
11002    fn q8_row_dot_fast_matches_scalar() {
11003        // The per-arch fast dot must agree with the exact scalar oracle
11004        // (same contract the fused q8 FFN arm rides on).
11005        let cols = 96;
11006        let row: Vec<u8> = (0..cols)
11007            .map(|i| ((i * 37 % 251) - 125) as i8 as u8)
11008            .collect();
11009        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.13).sin()).collect();
11010        let act = split_act(&x);
11011        let fast = q8_row_dot(&row, &act);
11012        let scalar = q8_row_dot_scalar(&row, &act);
11013        assert!(
11014            (fast - scalar).abs() <= scalar.abs() * 1e-5 + 1e-5,
11015            "fast {fast} vs scalar {scalar}"
11016        );
11017    }
11018
11019    #[test]
11020    fn f32_matvec_matches_matvec_rows_bitexact() {
11021        let (rows, cols) = (300, 40);
11022        let w: Vec<f32> = (0..rows * cols).map(|i| (i as f32 * 0.017).sin()).collect();
11023        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.05).cos()).collect();
11024        let qt = QTensor::from_f32(w.clone(), rows, cols);
11025
11026        let mut a = vec![0.0f32; rows];
11027        matvec_rows(None, &w, &x, &mut a);
11028        let mut b = vec![0.0f32; rows];
11029        qt.matvec(&x, &mut b, None);
11030        assert_eq!(a, b);
11031    }
11032
11033    #[test]
11034    fn sdot_kernel_exact_on_grid() {
11035        // Activations already on the i8 grid (±1 with amax=1 → sx=1/127,
11036        // xq=±127 dequantizes EXACTLY) → the SDOT path must match the
11037        // exact f32 dot to float rounding. This isolates kernel
11038        // correctness from quantization noise.
11039        eprintln!("sdot_enabled = {}", sdot_enabled());
11040        let (rows, cols) = (9, 80); // odd rows → exercises 4-row + tail
11041        let w: Vec<u8> = (0..rows * cols)
11042            .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
11043            .collect();
11044        let scales: Vec<f32> = (0..rows).map(|o| 0.005 + o as f32 * 0.001).collect();
11045        let x: Vec<f32> = (0..cols)
11046            .map(|i| match i % 3 {
11047                0 => 1.0,
11048                1 => -1.0,
11049                _ => 0.0,
11050            })
11051            .collect();
11052        let mut a = vec![0.0f32; rows];
11053        qmatvec(
11054            &w,
11055            &[],
11056            &scales,
11057            &x,
11058            &[],
11059            TensorDtype::Q8Row,
11060            rows,
11061            cols,
11062            &mut a,
11063            None,
11064        );
11065        for o in 0..rows {
11066            let mut acc = 0.0f32;
11067            for j in 0..cols {
11068                acc += (w[o * cols + j] as i8) as f32 * x[j];
11069            }
11070            let expect = acc * scales[o];
11071            assert!(
11072                (a[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
11073                "row {o}: {} vs {expect}",
11074                a[o]
11075            );
11076        }
11077    }
11078
11079    #[test]
11080    fn q1_tbl_fast_path_matches_reference() {
11081        // gpr = 8 exercises the TBL pair-load fast loop, and the LAST
11082        // row's final 4-tile window trips the 4B-overread guard (the
11083        // payload ends exactly at the last tile) — both paths must
11084        // agree with the dequant reference.
11085        let (rows, cols) = (5, 256);
11086        let gpr = cols / GROUP_SIZE;
11087        let mut bytes = Vec::new();
11088        for t in 0..rows * gpr {
11089            let s = 0.007 + (t % 11) as f32 * 0.004;
11090            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
11091            for j in 0..4 {
11092                bytes.push(((t * 53 + j * 89 + 7) % 249) as u8);
11093            }
11094        }
11095        let x: Vec<f32> = (0..cols)
11096            .map(|i| if (i * 5) % 7 < 3 { 1.0 } else { -1.0 })
11097            .collect();
11098        let mut w = vec![0.0f32; rows * cols];
11099        cortiq_core::quant::dequant_q1(&bytes, &mut w);
11100        let mut got = vec![0.0f32; rows];
11101        q1_matvec(&bytes, &x, rows, cols, &mut got, None);
11102        for o in 0..rows {
11103            let expect: f32 = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
11104            assert!(
11105                (got[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
11106                "row {o}: {} vs {expect}",
11107                got[o]
11108            );
11109        }
11110        // Blocked 1×4 batch (b=5: one quad + remainder) must equal the
11111        // single-matvec path bit-for-bit.
11112        let b = 5usize;
11113        let mut xs_all = Vec::new();
11114        for bi in 0..b {
11115            xs_all.extend(x.iter().map(|v| if bi % 2 == 0 { *v } else { -*v }));
11116        }
11117        let mut mm = vec![0.0f32; b * rows];
11118        q1_matmat(&bytes, &xs_all, b, rows, cols, &mut mm, None);
11119        for bi in 0..b {
11120            let mut single = vec![0.0f32; rows];
11121            q1_matvec(
11122                &bytes,
11123                &xs_all[bi * cols..(bi + 1) * cols],
11124                rows,
11125                cols,
11126                &mut single,
11127                None,
11128            );
11129            assert_eq!(&mm[bi * rows..(bi + 1) * rows], &single[..], "stream {bi}");
11130        }
11131    }
11132
11133    #[test]
11134    fn q1_kernels_match_exact_reference() {
11135        // Synthetic q1 payload: 6-byte tiles [f16 scale][4B bits].
11136        let (rows, cols) = (7, 96);
11137        let gpr = cols / GROUP_SIZE;
11138        let mut bytes = Vec::new();
11139        for t in 0..rows * gpr {
11140            let s = 0.01 + (t % 13) as f32 * 0.003;
11141            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
11142            for j in 0..4 {
11143                bytes.push(((t * 31 + j * 97) % 251) as u8);
11144            }
11145        }
11146        // On-grid activations (±1, amax 1) → the SDOT path is exact.
11147        let x: Vec<f32> = (0..cols)
11148            .map(|i| if i % 3 == 0 { 1.0 } else { -1.0 })
11149            .collect();
11150        // Reference through the core dequant.
11151        let mut w = vec![0.0f32; rows * cols];
11152        cortiq_core::quant::dequant_q1(&bytes, &mut w);
11153        let mut expect = vec![0.0f32; rows];
11154        for o in 0..rows {
11155            expect[o] = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
11156        }
11157        let mut got = vec![0.0f32; rows];
11158        q1_matvec(&bytes, &x, rows, cols, &mut got, None);
11159        for o in 0..rows {
11160            assert!(
11161                (got[o] - expect[o]).abs() < 1e-3 * expect[o].abs().max(1e-3),
11162                "row {o}: {} vs {}",
11163                got[o],
11164                expect[o]
11165            );
11166        }
11167        // Pair and batch paths agree with the single path.
11168        let x2: Vec<f32> = x.iter().map(|v| -v).collect();
11169        let (mut a1, mut a2) = (vec![0.0f32; rows], vec![0.0f32; rows]);
11170        q1_matvec2(&bytes, &x, &x2, rows, cols, &mut a1, &mut a2, None);
11171        assert_eq!(a1, got);
11172        let mut xs = x.clone();
11173        xs.extend_from_slice(&x2);
11174        let mut mm = vec![0.0f32; 2 * rows];
11175        q1_matmat(&bytes, &xs, 2, rows, cols, &mut mm, None);
11176        assert_eq!(&mm[..rows], got.as_slice());
11177        assert_eq!(&mm[rows..], a2.as_slice());
11178    }
11179
11180    #[test]
11181    fn repack_is_bit_identical() {
11182        // The interleaved-repack kernel must produce EXACTLY the same
11183        // bits as the mmap-layout kernel: integer accumulation is order-
11184        // exact, the f32 epilogue is identical. Odd rows exercise the
11185        // tail; direct range calls exercise unaligned pool splits.
11186        let (rows, cols) = (267, 96); // 66 groups + 3 tail rows, cols % 16 == 0
11187        let w: Vec<u8> = (0..rows * cols)
11188            .map(|i| (((i * 89) % 253) as i32 - 126) as i8 as u8)
11189            .collect();
11190        let scales: Vec<f32> = (0..rows).map(|o| 0.003 + o as f32 * 0.0007).collect();
11191        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.37).sin() * 2.0).collect();
11192        let rep = q8_repack_layout(&w, rows, cols);
11193        // Group interleave round-trips.
11194        for g in 0..rows / 4 {
11195            for c in 0..cols / 16 {
11196                for lane in 0..4 {
11197                    assert_eq!(
11198                        &rep[g * 4 * cols + c * 64 + lane * 16
11199                            ..g * 4 * cols + c * 64 + lane * 16 + 16],
11200                        &w[(g * 4 + lane) * cols + c * 16..(g * 4 + lane) * cols + c * 16 + 16],
11201                    );
11202                }
11203            }
11204        }
11205        let mut a = vec![0.0f32; rows];
11206        qmatvec(
11207            &w,
11208            &[],
11209            &scales,
11210            &x,
11211            &[],
11212            TensorDtype::Q8Row,
11213            rows,
11214            cols,
11215            &mut a,
11216            None,
11217        );
11218        let mut b = vec![0.0f32; rows];
11219        qmatvec(
11220            &w,
11221            &rep,
11222            &scales,
11223            &x,
11224            &[],
11225            TensorDtype::Q8Row,
11226            rows,
11227            cols,
11228            &mut b,
11229            None,
11230        );
11231        assert_eq!(a, b, "full-range repack output diverged");
11232
11233        #[cfg(target_arch = "aarch64")]
11234        if sdot_enabled() {
11235            // Unaligned range split (pool workers get arbitrary bounds).
11236            let act = split_act(&x);
11237            let mut c1 = vec![0.0f32; rows];
11238            let mut c2 = vec![0.0f32; rows];
11239            q8_range_sdot(
11240                &w,
11241                &[],
11242                &scales,
11243                &act,
11244                cols,
11245                SendMut(c1.as_mut_ptr()),
11246                3,
11247                rows - 2,
11248            );
11249            q8_range_sdot(
11250                &w,
11251                &rep,
11252                &scales,
11253                &act,
11254                cols,
11255                SendMut(c2.as_mut_ptr()),
11256                3,
11257                rows - 2,
11258            );
11259            assert_eq!(c1, c2, "unaligned-range repack output diverged");
11260        }
11261    }
11262
11263    #[test]
11264    fn sdot_a8w8_noise_is_bounded() {
11265        // Off-grid activations: A8 quantization noise must stay small in
11266        // relative L2 over the whole output (realistic accuracy contract;
11267        // vmfcore measured argmax-identical decode on real models).
11268        let (rows, cols) = (16, 512);
11269        let w: Vec<u8> = (0..rows * cols)
11270            .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
11271            .collect();
11272        let scales = vec![0.01f32; rows];
11273        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.21).sin()).collect();
11274        let mut a = vec![0.0f32; rows];
11275        qmatvec(
11276            &w,
11277            &[],
11278            &scales,
11279            &x,
11280            &[],
11281            TensorDtype::Q8Row,
11282            rows,
11283            cols,
11284            &mut a,
11285            None,
11286        );
11287        let (mut num, mut den) = (0f64, 0f64);
11288        for o in 0..rows {
11289            let mut acc = 0.0f32;
11290            for j in 0..cols {
11291                acc += (w[o * cols + j] as i8) as f32 * x[j];
11292            }
11293            let expect = acc * scales[o];
11294            num += ((a[o] - expect) as f64).powi(2);
11295            den += (expect as f64).powi(2);
11296        }
11297        let rel = (num / den.max(1e-12)).sqrt();
11298        assert!(rel < 0.05, "A8W8 relative L2 error too high: {rel}");
11299    }
11300
11301    #[test]
11302    fn i8_dot_neon_matches_scalar() {
11303        let n = 100;
11304        let w: Vec<u8> = (0..n).map(|i| ((i * 37 + 11) % 251) as u8).collect();
11305        let x: Vec<f32> = (0..n).map(|i| (i as f32 * 0.13).sin()).collect();
11306        let mut scalar = 0.0f32;
11307        for j in 0..n {
11308            scalar += (w[j] as i8) as f32 * x[j];
11309        }
11310        let fast = dot_i8_f32(&w, &x);
11311        assert!((scalar - fast).abs() < 1e-3 * scalar.abs().max(1.0));
11312    }
11313
11314    /// Fused vbit matvec must match full dequant_vbit + dense matvec.
11315    #[test]
11316    fn vbitmatvec_matches_full_dequant() {
11317        let (rows, cols) = (6, 64);
11318        let ng = cols / GROUP_SIZE;
11319        // Hand-craft: bits per row, f16 scales, packed rows.
11320        let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
11321        let mut bytes = bits.clone();
11322        for g in 0..rows * ng {
11323            let s = 0.02 + 0.001 * g as f32;
11324            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
11325        }
11326        for r in 0..rows {
11327            let b = bits[r] as usize;
11328            let (mut acc, mut nb) = (0u64, 0usize);
11329            let mut rowbytes = Vec::new();
11330            for i in 0..cols {
11331                let v = ((i * 7 + r * 13) % (1 << b)) as u64;
11332                acc = (acc << b) | v;
11333                nb += b;
11334                while nb >= 8 {
11335                    nb -= 8;
11336                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
11337                }
11338            }
11339            if nb > 0 {
11340                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
11341            }
11342            bytes.extend_from_slice(&rowbytes);
11343        }
11344        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
11345
11346        let mut reference = vec![0f32; rows * cols];
11347        cortiq_core::quant::dequant_vbit(&bytes, rows, cols, &mut reference).unwrap();
11348        let mut expect = vec![0f32; rows];
11349        for r in 0..rows {
11350            expect[r] = reference[r * cols..(r + 1) * cols]
11351                .iter()
11352                .zip(&x)
11353                .map(|(w, xv)| w * xv)
11354                .sum();
11355        }
11356        let mut got = vec![0f32; rows];
11357        let offsets = vbit_row_offsets(&bytes, rows, cols);
11358        vbitmatvec(&bytes, &offsets, &x, rows, cols, &mut got, None);
11359        // SDOT path quantizes activations to i8 (A8W8): bounded noise,
11360        // same contract as q8 (exact path is pinned by CMF_SDOT=0 in
11361        // the golden-parity gate).
11362        let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
11363        let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
11364        for r in 0..rows {
11365            assert!(
11366                (got[r] - expect[r]).abs() < tol * scale,
11367                "row {r}: {} vs {}",
11368                got[r],
11369                expect[r]
11370            );
11371        }
11372    }
11373
11374    /// Fused q4 matvec must match the reference full-dequant + dense
11375    /// matvec bit-for-bit in structure (same f32 math, group order).
11376    /// vbit matmat: the blocked 1×4 leg must match the per-row path
11377    /// (paired env toggle; larger shape so both code paths engage).
11378    #[test]
11379    #[cfg(target_arch = "x86_64")]
11380    fn vbit_matmat_blocked_matches_per_row() {
11381        let (rows, cols, b) = (64usize, 128usize, 9usize);
11382        let ng = cols / GROUP_SIZE;
11383        let bits: Vec<u8> = (0..rows).map(|r| [3u8, 4, 5, 6][r % 4]).collect();
11384        let mut bytes = bits.clone();
11385        for g in 0..rows * ng {
11386            let sc = 0.02 + 0.0005 * g as f32;
11387            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
11388        }
11389        for r in 0..rows {
11390            let bw = bits[r] as usize;
11391            let (mut acc, mut nb) = (0u64, 0usize);
11392            let mut rowbytes = Vec::new();
11393            for i in 0..cols {
11394                let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
11395                acc = (acc << bw) | v;
11396                nb += bw;
11397                while nb >= 8 {
11398                    nb -= 8;
11399                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
11400                }
11401            }
11402            if nb > 0 {
11403                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
11404            }
11405            bytes.extend_from_slice(&rowbytes);
11406        }
11407        let x: Vec<f32> = (0..b * cols)
11408            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
11409            .collect();
11410        let offsets = vbit_row_offsets(&bytes, rows, cols);
11411        let mut y_a = vec![0f32; b * rows];
11412        let mut y_b = vec![0f32; b * rows];
11413        unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
11414        vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_a, None);
11415        unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
11416        vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_b, None);
11417        unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
11418        let max_d = y_a
11419            .iter()
11420            .zip(&y_b)
11421            .map(|(p, q)| (p - q).abs())
11422            .fold(0.0f32, f32::max);
11423        assert!(max_d < 1e-4, "vbit blocked ≠ per-row: max|Δ| = {max_d}");
11424    }
11425
11426    /// q4t blocked 1×4 (SDOT on ARM, AVX2 on x86) must equal the
11427    /// per-row path exactly: same nibble unpack, same group order,
11428    /// same f32 accumulation — batch == matvec bit-for-bit. b=9 covers
11429    /// two full 1×4 blocks plus a remainder through the single-row
11430    /// kernel. (Both paths produce identical output, so the shared
11431    /// CMF_X86_BLOCKED env var racing with other tests cannot flip
11432    /// the verdict — worst case both sides take the same path.)
11433    #[test]
11434    fn q4t_matmat_blocked_matches_per_row() {
11435        let (rows, cols, b) = (16usize, 64usize, 9usize);
11436        let gpr = cols / GROUP_SIZE;
11437        let mut bytes = vec![0u8; rows * gpr * Q4_TILE];
11438        for r in 0..rows {
11439            for g in 0..gpr {
11440                let t = (r * gpr + g) * Q4_TILE;
11441                let sc = 0.02 + 0.001 * (r * gpr + g) as f32;
11442                bytes[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
11443                for k in 0..16 {
11444                    bytes[t + 2 + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
11445                }
11446            }
11447        }
11448        let x: Vec<f32> = (0..b * cols)
11449            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
11450            .collect();
11451        let mut y_blk = vec![0f32; b * rows];
11452        let mut y_row = vec![0f32; b * rows];
11453        unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
11454        q4t_matmat(&bytes, &x, b, rows, cols, &mut y_blk, None);
11455        unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
11456        q4t_matmat(&bytes, &x, b, rows, cols, &mut y_row, None);
11457        unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
11458        assert_eq!(y_blk, y_row, "q4t blocked 1x4 ≠ per-row");
11459    }
11460
11461    /// The wide-batch Accelerate arm of q4t_matmat vs a brute-force
11462    /// f32 dequant matmul: both are f32 GEMMs, so only reduction
11463    /// order differs — tight tolerance.
11464    /// A synthetic q4tp payload: random nibbles plus a per-row ladder whose
11465    /// span varies row to row, so the codes actually exercise the full 0..31
11466    /// range rather than clustering on one rung.
11467    fn synth_q4tp(rows: usize, cols: usize) -> Vec<u8> {
11468        use cortiq_core::quant::{f32_to_f16, q4tp_code_stride, q4tp_put_code};
11469        let gpr = cols / GROUP_SIZE;
11470        let stride = q4tp_code_stride(gpr);
11471        let (params_off, codes_off, _) = q4tp_sections(rows, cols);
11472        let mut b = vec![0u8; codes_off + rows * stride];
11473        for r in 0..rows {
11474            for g in 0..gpr {
11475                let t = (r * gpr + g) * Q4TP_NIB;
11476                for k in 0..16 {
11477                    b[t + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
11478                }
11479            }
11480            let lo = -6.0 - 0.03 * (r % 17) as f32;
11481            let step = 0.01 + 0.004 * (r % 11) as f32;
11482            let p = params_off + r * 4;
11483            b[p..p + 2].copy_from_slice(&f32_to_f16(lo).to_le_bytes());
11484            b[p + 2..p + 4].copy_from_slice(&f32_to_f16(step).to_le_bytes());
11485            let crow = &mut b[codes_off + r * stride..codes_off + (r + 1) * stride];
11486            for g in 0..gpr {
11487                q4tp_put_code(crow, g, (r * 5 + g * 3) % 32);
11488            }
11489        }
11490        b
11491    }
11492
11493    /// The same weights re-expressed as q4_tiled, so the proven kernel can
11494    /// be the reference: each tile stores the ladder scale its code selects.
11495    /// Only the f16 rounding of that scale separates the two payloads.
11496    fn q4tp_as_q4t(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
11497        let gpr = cols / GROUP_SIZE;
11498        let v = Q4tpView::new(bytes, rows, cols);
11499        let mut out = vec![0u8; rows * gpr * Q4_TILE];
11500        let mut sc = vec![0f32; gpr];
11501        for r in 0..rows {
11502            v.scales_into(r, gpr, &mut sc);
11503            for g in 0..gpr {
11504                let t = (r * gpr + g) * Q4_TILE;
11505                let s = sc[g];
11506                out[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
11507                let src = (r * gpr + g) * Q4TP_NIB;
11508                out[t + 2..t + Q4_TILE].copy_from_slice(&v.nib[src..src + Q4TP_NIB]);
11509            }
11510        }
11511        out
11512    }
11513
11514    /// The exact (`CMF_SDOT=0`) path must reproduce `dequant_q4tp` to f32
11515    /// rounding — that scalar routine is the format's definition, and the
11516    /// kernels re-derive the scale from the ladder independently. Call the
11517    /// row kernel directly: `matmat` picks the int8 arm when a8w8 is on,
11518    /// so routing through it would test the other path by accident.
11519    #[test]
11520    fn q4tp_exact_path_matches_dequant_reference() {
11521        let (rows, cols) = (256usize, 512usize);
11522        let gpr = cols / GROUP_SIZE;
11523        let bytes = synth_q4tp(rows, cols);
11524        let mut w = vec![0f32; rows * cols];
11525        cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
11526
11527        let x: Vec<f32> = (0..cols)
11528            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
11529            .collect();
11530        let v = Q4tpView::new(&bytes, rows, cols);
11531        let mut sc = vec![0f32; gpr];
11532        for r in 0..rows {
11533            v.scales_into(r, gpr, &mut sc);
11534            let got = q4tp_row_exact(v.nib, r, gpr, &x, &sc);
11535            let want: f32 = (0..cols).map(|c| w[r * cols + c] * x[c]).sum();
11536            // These dot products cancel down to ~1e-3 from terms of ~5e-2, so
11537            // the meaningful yardstick is the summed magnitude, not the result:
11538            // against the result any reordering of a 512-term f32 sum "fails".
11539            let mag: f32 = (0..cols).map(|c| (w[r * cols + c] * x[c]).abs()).sum();
11540            assert!(
11541                (got - want).abs() <= 1e-5 * mag,
11542                "row {r}: kernel {got} vs dequant {want}"
11543            );
11544        }
11545    }
11546
11547    /// The int8 (a8w8) path can't be checked against an f32 reference — the
11548    /// activation quantization dominates. Check it against the q4t kernel it
11549    /// was ported from instead, on payloads holding the same weights: that
11550    /// isolates exactly what the port could break (16 B stride, ladder
11551    /// lookup, nibble unpack) from what it deliberately shares.
11552    #[test]
11553    fn q4tp_matvec_matches_the_q4t_kernel_it_was_ported_from() {
11554        let (rows, cols) = (256usize, 512usize);
11555        let bytes = synth_q4tp(rows, cols);
11556        let twin = q4tp_as_q4t(&bytes, rows, cols);
11557        let x: Vec<f32> = (0..cols)
11558            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
11559            .collect();
11560
11561        let mut got = vec![0f32; rows];
11562        q4tp_matvec(&bytes, &x, rows, cols, &mut got, None);
11563        let mut want = vec![0f32; rows];
11564        q4t_matvec(&twin, &x, rows, cols, &mut want, None);
11565
11566        // Scale is f16 in the twin and f32 here, so allow that rounding on
11567        // top of the summed magnitude (same cancellation argument as above).
11568        let mut w = vec![0f32; rows * cols];
11569        cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
11570        for r in 0..rows {
11571            let mag: f32 = (0..cols).map(|c| (w[r * cols + c] * x[c]).abs()).sum();
11572            assert!(
11573                (got[r] - want[r]).abs() <= 1e-3 * mag,
11574                "row {r}: q4tp {} vs q4t {}",
11575                got[r],
11576                want[r]
11577            );
11578        }
11579    }
11580
11581    /// `matmat` carries three arms (Accelerate, blocked int8 1x4, scalar).
11582    /// Batch 5 crosses the blocked kernel's stride, so this exercises the
11583    /// 1x4 path AND its scalar tail in one run — the blocked kernel is new
11584    /// code and its four accumulators are exactly what tends to go wrong.
11585    #[test]
11586    fn q4tp_matmat_matches_the_q4t_kernel_it_was_ported_from() {
11587        let (rows, cols, b) = (256usize, 512usize, 5usize);
11588        let bytes = synth_q4tp(rows, cols);
11589        let twin = q4tp_as_q4t(&bytes, rows, cols);
11590        let xs: Vec<f32> = (0..b * cols)
11591            .map(|i| ((i * 29 + 11) % 89) as f32 / 89.0 - 0.5)
11592            .collect();
11593
11594        let mut got = vec![0f32; b * rows];
11595        q4tp_matmat(&bytes, &xs, b, rows, cols, &mut got, None);
11596        let mut want = vec![0f32; b * rows];
11597        q4t_matmat(&twin, &xs, b, rows, cols, &mut want, None);
11598
11599        let mut w = vec![0f32; rows * cols];
11600        cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
11601        for t in 0..b {
11602            for r in 0..rows {
11603                let mag: f32 = (0..cols)
11604                    .map(|c| (w[r * cols + c] * xs[t * cols + c]).abs())
11605                    .sum();
11606                let (g, wa) = (got[t * rows + r], want[t * rows + r]);
11607                assert!(
11608                    (g - wa).abs() <= 1e-3 * mag,
11609                    "batch {t} row {r}: q4tp {g} vs q4t {wa}"
11610                );
11611            }
11612        }
11613    }
11614
11615    #[test]
11616    fn q4tp_matvec2_matches_the_single_stream_kernel() {
11617        let (rows, cols) = (128usize, 256usize);
11618        let gpr = cols / GROUP_SIZE;
11619        let bytes = synth_q4tp(rows, cols);
11620        let xs: Vec<f32> = (0..2 * cols)
11621            .map(|i| ((i * 29 + 11) % 89) as f32 / 89.0 - 0.5)
11622            .collect();
11623
11624        let (mut o1, mut o2) = (vec![0f32; rows], vec![0f32; rows]);
11625        q4tp_matvec2(
11626            &bytes,
11627            &xs[..cols],
11628            &xs[cols..],
11629            rows,
11630            cols,
11631            &mut o1,
11632            &mut o2,
11633            None,
11634        );
11635
11636        // matvec2 takes the exact path for both streams, so the single-row
11637        // kernel is an exact reference — no tolerance for path differences.
11638        let v = Q4tpView::new(&bytes, rows, cols);
11639        let mut sc = vec![0f32; gpr];
11640        for r in 0..rows {
11641            v.scales_into(r, gpr, &mut sc);
11642            assert_eq!(o1[r], q4tp_row_exact(v.nib, r, gpr, &xs[..cols], &sc));
11643            assert_eq!(o2[r], q4tp_row_exact(v.nib, r, gpr, &xs[cols..], &sc));
11644        }
11645    }
11646
11647    /// q4tp must not COST speed — it exists to save bytes, and a format that
11648    /// trades 7% of a file for a slower model is a bad trade. This guard is
11649    /// here because correctness tests happily passed while `q4tp_matmat` was
11650    /// missing its int8 and Accelerate arms and the model ran 5x slower.
11651    /// Measured on M-series: 0.97-1.04x, i.e. parity (16 B tiles are better
11652    /// aligned than q4t's 18 B, which pays for the scale indirection).
11653    #[test]
11654    fn q4tp_matvec_keeps_pace_with_q4t() {
11655        let (rows, cols) = (4096usize, 3072usize);
11656        let bytes = synth_q4tp(rows, cols);
11657        let twin = q4tp_as_q4t(&bytes, rows, cols);
11658        let x: Vec<f32> = (0..cols).map(|i| (i % 97) as f32 / 97.0 - 0.5).collect();
11659        let mut o = vec![0f32; rows];
11660        let n = 12;
11661        let mut best = (f64::MAX, f64::MAX);
11662        // Interleaved A/B, minimum statistic: this machine throttles, and a
11663        // mean over a thermal ramp reliably indicts whichever ran second.
11664        for _ in 0..3 {
11665            let t0 = std::time::Instant::now();
11666            for _ in 0..n {
11667                q4t_matvec(&twin, &x, rows, cols, &mut o, None);
11668            }
11669            best.0 = best.0.min(t0.elapsed().as_secs_f64());
11670            let t0 = std::time::Instant::now();
11671            for _ in 0..n {
11672                q4tp_matvec(&bytes, &x, rows, cols, &mut o, None);
11673            }
11674            best.1 = best.1.min(t0.elapsed().as_secs_f64());
11675        }
11676        let ratio = best.1 / best.0;
11677        println!(
11678            "q4t {:.3} ms | q4tp {:.3} ms | {ratio:.2}x",
11679            best.0 * 1e3 / n as f64,
11680            best.1 * 1e3 / n as f64
11681        );
11682        assert!(ratio < 2.0, "q4tp matvec {ratio:.2}x slower than q4t");
11683    }
11684
11685    #[cfg(target_os = "macos")]
11686    #[test]
11687    fn q4t_matmat_accel_matches_dequant_reference() {
11688        if !accel_gemm_enabled() {
11689            return; // CMF_ACCEL=0
11690        }
11691        let (rows, cols, b) = (512usize, 1024usize, 8usize); // ≥500K → accel arm
11692        let gpr = cols / GROUP_SIZE;
11693        let mut bytes = vec![0u8; rows * gpr * Q4_TILE];
11694        for r in 0..rows {
11695            for g in 0..gpr {
11696                let t = (r * gpr + g) * Q4_TILE;
11697                let sc = 0.02 + 0.0005 * ((r * gpr + g) % 64) as f32;
11698                bytes[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
11699                for k in 0..16 {
11700                    bytes[t + 2 + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
11701                }
11702            }
11703        }
11704        let x: Vec<f32> = (0..b * cols)
11705            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
11706            .collect();
11707        let mut got = vec![0f32; b * rows];
11708        q4t_matmat(&bytes, &x, b, rows, cols, &mut got, None);
11709        // Brute-force reference off the same tiles.
11710        let mut w = vec![0f32; rows * cols];
11711        for r in 0..rows {
11712            for g in 0..gpr {
11713                let t = (r * gpr + g) * Q4_TILE;
11714                let s = f16_to_f32(u16::from_le_bytes([bytes[t], bytes[t + 1]]));
11715                for (k, &bb) in bytes[t + 2..t + Q4_TILE].iter().enumerate() {
11716                    w[r * cols + g * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s;
11717                    w[r * cols + g * GROUP_SIZE + k * 2 + 1] =
11718                        (((bb >> 4) & 0x0F) as f32 - 8.0) * s;
11719                }
11720            }
11721        }
11722        for bi in 0..b {
11723            for r in 0..rows {
11724                let want: f32 = (0..cols).map(|j| x[bi * cols + j] * w[r * cols + j]).sum();
11725                let d = (got[bi * rows + r] - want).abs();
11726                assert!(
11727                    d <= want.abs().max(1.0) * 1e-4,
11728                    "accel q4t GEMM diverged at ({bi},{r}): {} vs {want}",
11729                    got[bi * rows + r]
11730                );
11731            }
11732        }
11733    }
11734
11735    #[test]
11736    fn q4matvec_matches_full_dequant() {
11737        let (rows, cols) = (8, 64);
11738        let groups = rows * cols / GROUP_SIZE;
11739        // Hand-craft a q4_block blob: nibbles then f16 scales.
11740        let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
11741        for i in 0..groups * 16 {
11742            bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
11743        }
11744        for g in 0..groups {
11745            let s = 0.01 + 0.003 * g as f32;
11746            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
11747        }
11748        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
11749
11750        let mut reference = vec![0.0f32; rows * cols];
11751        cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
11752        let mut expect = vec![0.0f32; rows];
11753        for r in 0..rows {
11754            expect[r] = reference[r * cols..(r + 1) * cols]
11755                .iter()
11756                .zip(&x)
11757                .map(|(w, xv)| w * xv)
11758                .sum();
11759        }
11760
11761        let mut got = vec![0.0f32; rows];
11762        q4matvec(&bytes, &x, rows, cols, &mut got, None);
11763        // SDOT path quantizes activations to i8 (A8W8): bounded noise,
11764        // same contract as q8/vbit (exact path is pinned by CMF_SDOT=0
11765        // in the golden-parity gate).
11766        let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
11767        let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
11768        for r in 0..rows {
11769            assert!(
11770                (got[r] - expect[r]).abs() < tol * scale,
11771                "row {r}: {} vs {}",
11772                got[r],
11773                expect[r]
11774            );
11775        }
11776    }
11777
11778    /// Fused two-input vbit matvec must equal two single matvecs exactly
11779    /// (same per-lane accumulation order on both scalar and SDOT paths).
11780    #[test]
11781    fn vbitmatvec2_equals_two_singles() {
11782        let (rows, cols) = (6, 64);
11783        let ng = cols / GROUP_SIZE;
11784        let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
11785        let mut bytes = bits.clone();
11786        for g in 0..rows * ng {
11787            let s = 0.02 + 0.001 * g as f32;
11788            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
11789        }
11790        for r in 0..rows {
11791            let b = bits[r] as usize;
11792            let (mut acc, mut nb) = (0u64, 0usize);
11793            let mut rowbytes = Vec::new();
11794            for i in 0..cols {
11795                let v = ((i * 7 + r * 13) % (1 << b)) as u64;
11796                acc = (acc << b) | v;
11797                nb += b;
11798                while nb >= 8 {
11799                    nb -= 8;
11800                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
11801                }
11802            }
11803            if nb > 0 {
11804                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
11805            }
11806            bytes.extend_from_slice(&rowbytes);
11807        }
11808        let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
11809        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).cos()).collect();
11810        let offsets = vbit_row_offsets(&bytes, rows, cols);
11811
11812        let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
11813        vbitmatvec(&bytes, &offsets, &x1, rows, cols, &mut a1, None);
11814        vbitmatvec(&bytes, &offsets, &x2, rows, cols, &mut a2, None);
11815        let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
11816        vbitmatvec2(
11817            &bytes, &offsets, &x1, &x2, rows, cols, &mut b1, &mut b2, None,
11818        );
11819        assert_eq!(a1, b1, "fused vbit lane 1 must be bit-identical");
11820        assert_eq!(a2, b2, "fused vbit lane 2 must be bit-identical");
11821    }
11822
11823    /// Fused two-input q4 matvec must equal two single matvecs exactly.
11824    #[test]
11825    fn q4matvec2_equals_two_singles() {
11826        let (rows, cols) = (8, 128);
11827        let groups = rows * cols / GROUP_SIZE;
11828        let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
11829        for i in 0..groups * 16 {
11830            bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
11831        }
11832        for g in 0..groups {
11833            let s = 0.01 + 0.003 * g as f32;
11834            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
11835        }
11836        // Include an outlier channel so the SDOT correction path is
11837        // exercised in the pair kernel too.
11838        let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
11839        x1[9] = 250.0;
11840        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
11841
11842        let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
11843        q4matvec(&bytes, &x1, rows, cols, &mut a1, None);
11844        q4matvec(&bytes, &x2, rows, cols, &mut a2, None);
11845        let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
11846        q4matvec2(&bytes, &x1, &x2, rows, cols, &mut b1, &mut b2, None);
11847        assert_eq!(a1, b1, "fused q4 lane 1 must be bit-identical");
11848        assert_eq!(a2, b2, "fused q4 lane 2 must be bit-identical");
11849    }
11850
11851    /// Multi-matrix job must equal separate matvecs exactly — same
11852    /// kernels, only the dispatch is fused.
11853    #[test]
11854    fn matvec_many_equals_separate_matvecs() {
11855        use crate::pool::Pool;
11856        let (r1, r2, cols) = (300, 200, 64);
11857        let mk = |salt: usize, rows: usize| {
11858            QTensor::from_f32(
11859                (0..rows * cols)
11860                    .map(|i| ((i * 7 + salt) % 97) as f32 / 97.0 - 0.5)
11861                    .collect(),
11862                rows,
11863                cols,
11864            )
11865        };
11866        let (a, b) = (mk(1, r1), mk(5, r2));
11867        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).sin()).collect();
11868        let pool = Pool::new(3);
11869
11870        let (mut ea, mut eb) = (vec![0f32; r1], vec![0f32; r2]);
11871        a.matvec(&x, &mut ea, Some(&pool));
11872        b.matvec(&x, &mut eb, Some(&pool));
11873        let (mut ga, mut gb) = (vec![0f32; r1], vec![0f32; r2]);
11874        QTensor::matvec_many([&a, &b], &x, [&mut ga, &mut gb], Some(&pool));
11875        assert_eq!(ea, ga, "fused multi-matrix lane 1 must be bit-identical");
11876        assert_eq!(eb, gb, "fused multi-matrix lane 2 must be bit-identical");
11877    }
11878
11879    /// The public Q4TP operator must take the real mapped matvec_many arm,
11880    /// rather than the F32 fallback above.  Build a tiny valid CMF so both
11881    /// handles retain their mmap payloads, then compare the fused dispatch
11882    /// with two ordinary mapped matvec calls bit-for-bit.
11883    #[test]
11884    fn q4tp_matvec_many_equals_separate_matvecs() {
11885        use crate::pool::Pool;
11886        use cortiq_core::{CMF_VERSION, CmfHeader, CmfModel, QuantType, TensorSpec};
11887
11888        let (r1, r2, cols) = (300usize, 200usize, 64usize);
11889        let arch: cortiq_core::ModelArch = serde_json::from_value(serde_json::json!({
11890            "arch_name": "tiny-q4tp",
11891            "hidden_size": cols,
11892            "intermediate_size": cols * 2,
11893            "num_layers": 1,
11894            "num_attention_heads": 2,
11895            "num_kv_heads": 1,
11896            "head_dim": 32,
11897            "vocab_size": r1,
11898            "layer_types": ["FullAttention"],
11899            "rms_norm_eps": 1e-6,
11900            "max_position_embeddings": 8,
11901            "linear_conv_kernel_dim": 0,
11902            "linear_num_key_heads": 0,
11903            "linear_num_value_heads": 0
11904        }))
11905        .unwrap();
11906        let header = CmfHeader {
11907            format: "cmf".into(),
11908            version: CMF_VERSION,
11909            arch,
11910            quant_type: QuantType::Q4Block,
11911            provenance: None,
11912            tokenizer_config: None,
11913            section_hashes: None,
11914            skills: Vec::new(),
11915            shard: None,
11916            calibration: None,
11917            routing: None,
11918        };
11919        let specs = [
11920            TensorSpec {
11921                name: "q".into(),
11922                dtype: TensorDtype::Q4TiledP,
11923                shape: vec![r1, cols],
11924                data: synth_q4tp(r1, cols),
11925            },
11926            TensorSpec {
11927                name: "kv".into(),
11928                dtype: TensorDtype::Q4TiledP,
11929                shape: vec![r2, cols],
11930                data: synth_q4tp(r2, cols),
11931            },
11932        ];
11933        let dir = std::env::temp_dir().join(format!("cmf-q4tp-many-{}", std::process::id()));
11934        std::fs::create_dir_all(&dir).unwrap();
11935        let path = dir.join("m.cmf");
11936        CmfModel::write(&path, &header, &specs, None, None).unwrap();
11937        let model = Arc::new(CmfModel::open(&path).unwrap());
11938        let (a, b) = (
11939            QTensor::from_model(&model, "q").unwrap(),
11940            QTensor::from_model(&model, "kv").unwrap(),
11941        );
11942        assert_eq!(a.model_dtype(), Some(TensorDtype::Q4TiledP));
11943        assert_eq!(b.model_dtype(), Some(TensorDtype::Q4TiledP));
11944        let x: Vec<f32> = (0..cols)
11945            .map(|i| ((i * 17 + 3) % 97) as f32 / 97.0 - 0.5)
11946            .collect();
11947        let pool = Pool::new(3);
11948        let (mut ea, mut eb) = (vec![0.0f32; r1], vec![0.0f32; r2]);
11949        a.matvec(&x, &mut ea, Some(&pool));
11950        b.matvec(&x, &mut eb, Some(&pool));
11951        let (mut ga, mut gb) = (vec![0.0f32; r1], vec![0.0f32; r2]);
11952        QTensor::matvec_many([&a, &b], &x, [&mut ga, &mut gb], Some(&pool));
11953        assert_eq!(ea, ga, "Q4TP fused lane 1 must be bit-identical");
11954        assert_eq!(eb, gb, "Q4TP fused lane 2 must be bit-identical");
11955        let _ = std::fs::remove_dir_all(&dir);
11956    }
11957
11958    /// Batched q4/vbit matmat must equal per-position matvec calls
11959    /// exactly (the fallback it replaced) — same kernels, same order.
11960    #[test]
11961    fn batched_matmat_equals_per_position_matvec() {
11962        let (rows, cols, b) = (8, 64, 5);
11963        // q4 blob.
11964        let groups = rows * cols / GROUP_SIZE;
11965        let mut q4 = Vec::new();
11966        for i in 0..groups * 16 {
11967            q4.push((((i * 7 + 3) % 256) & 0xFF) as u8);
11968        }
11969        for g in 0..groups {
11970            q4.extend_from_slice(
11971                &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
11972            );
11973        }
11974        // vbit blob (mixed widths incl. 8).
11975        let ng = cols / GROUP_SIZE;
11976        let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4, 5, 3];
11977        let mut vb = bits.clone();
11978        for g in 0..rows * ng {
11979            vb.extend_from_slice(
11980                &cortiq_core::quant::f32_to_f16(0.02 + 0.001 * g as f32).to_le_bytes(),
11981            );
11982        }
11983        for r in 0..rows {
11984            let bw = bits[r] as usize;
11985            let (mut acc, mut nb) = (0u64, 0usize);
11986            let mut rowbytes = Vec::new();
11987            for i in 0..cols {
11988                let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
11989                acc = (acc << bw) | v;
11990                nb += bw;
11991                while nb >= 8 {
11992                    nb -= 8;
11993                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
11994                }
11995            }
11996            if nb > 0 {
11997                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
11998            }
11999            vb.extend_from_slice(&rowbytes);
12000        }
12001        let offsets = vbit_row_offsets(&vb, rows, cols);
12002
12003        let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
12004
12005        // q4: batch vs singles.
12006        let mut got = vec![0f32; b * rows];
12007        q4matmat(&q4, &xs, b, rows, cols, &mut got, None);
12008        for bi in 0..b {
12009            let mut expect = vec![0f32; rows];
12010            q4matvec(
12011                &q4,
12012                &xs[bi * cols..(bi + 1) * cols],
12013                rows,
12014                cols,
12015                &mut expect,
12016                None,
12017            );
12018            assert_eq!(
12019                &got[bi * rows..(bi + 1) * rows],
12020                &expect[..],
12021                "q4 batch pos {bi}"
12022            );
12023        }
12024
12025        // vbit: batch vs singles.
12026        let mut got = vec![0f32; b * rows];
12027        vbitmatmat(&vb, &offsets, &xs, b, rows, cols, &mut got, None);
12028        for bi in 0..b {
12029            let mut expect = vec![0f32; rows];
12030            vbitmatvec(
12031                &vb,
12032                &offsets,
12033                &xs[bi * cols..(bi + 1) * cols],
12034                rows,
12035                cols,
12036                &mut expect,
12037                None,
12038            );
12039            assert_eq!(
12040                &got[bi * rows..(bi + 1) * rows],
12041                &expect[..],
12042                "vbit batch pos {bi}"
12043            );
12044        }
12045    }
12046
12047    /// q4_tiled kernels must produce BIT-identical outputs to the q4
12048    /// split kernels on the same values (same ints, same order — only
12049    /// the byte placement differs).
12050    #[test]
12051    fn q4_tiled_matches_q4_block_bitexact() {
12052        let (rows, cols, b) = (8usize, 128usize, 3usize);
12053        let groups = rows * cols / GROUP_SIZE;
12054        let mut split = Vec::with_capacity(groups * 18);
12055        for i in 0..groups * 16 {
12056            split.push((((i * 7 + 3) % 256) & 0xFF) as u8);
12057        }
12058        for g in 0..groups {
12059            split.extend_from_slice(
12060                &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
12061            );
12062        }
12063        // Re-tile: [scale][nibbles] per group.
12064        let (packed, scales) = split.split_at(groups * 16);
12065        let mut tiled = Vec::with_capacity(groups * Q4_TILE);
12066        for g in 0..groups {
12067            tiled.extend_from_slice(&scales[g * 2..g * 2 + 2]);
12068            tiled.extend_from_slice(&packed[g * 16..(g + 1) * 16]);
12069        }
12070
12071        let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
12072        x1[9] = 250.0; // exercise the outlier path
12073        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
12074
12075        let (mut a, mut t) = (vec![0f32; rows], vec![0f32; rows]);
12076        q4matvec(&split, &x1, rows, cols, &mut a, None);
12077        q4t_matvec(&tiled, &x1, rows, cols, &mut t, None);
12078        assert_eq!(a, t, "q4t matvec must match q4 bit-for-bit");
12079
12080        let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
12081        let (mut t1, mut t2) = (vec![0f32; rows], vec![0f32; rows]);
12082        q4matvec2(&split, &x1, &x2, rows, cols, &mut a1, &mut a2, None);
12083        q4t_matvec2(&tiled, &x1, &x2, rows, cols, &mut t1, &mut t2, None);
12084        assert_eq!(a1, t1);
12085        assert_eq!(a2, t2);
12086
12087        let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
12088        let (mut am, mut tm) = (vec![0f32; b * rows], vec![0f32; b * rows]);
12089        q4matmat(&split, &xs, b, rows, cols, &mut am, None);
12090        q4t_matmat(&tiled, &xs, b, rows, cols, &mut tm, None);
12091        assert_eq!(am, tm, "q4t matmat must match q4 bit-for-bit");
12092    }
12093
12094    /// q4 SDOT outlier correction: a single huge activation channel
12095    /// (>8·rms → outlier, zeroed in xq) must still contribute its EXACT
12096    /// term. On-grid bulk (±1/0 → xq dequantizes exactly) isolates the
12097    /// correction from A8W8 noise. cols must exceed 64: at n=64 the
12098    /// 8·rms threshold equals sqrt(v²+rest) ≥ v, so a single outlier
12099    /// can never qualify (8² = n).
12100    #[test]
12101    fn q4matvec_sdot_outlier_exact() {
12102        let (rows, cols) = (4, 128);
12103        let groups = rows * cols / GROUP_SIZE;
12104        let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
12105        for i in 0..groups * 16 {
12106            bytes.push(((i * 11 + 5) % 256) as u8);
12107        }
12108        for g in 0..groups {
12109            let s = 0.02 + 0.002 * g as f32;
12110            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
12111        }
12112        let mut x: Vec<f32> = (0..cols)
12113            .map(|i| match i % 3 {
12114                0 => 1.0,
12115                1 => -1.0,
12116                _ => 0.0,
12117            })
12118            .collect();
12119        x[17] = 300.0; // ≫ 8·rms → outlier channel
12120
12121        let mut reference = vec![0.0f32; rows * cols];
12122        cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
12123        let mut expect = vec![0.0f32; rows];
12124        for r in 0..rows {
12125            expect[r] = reference[r * cols..(r + 1) * cols]
12126                .iter()
12127                .zip(&x)
12128                .map(|(w, xv)| w * xv)
12129                .sum();
12130        }
12131        let mut got = vec![0.0f32; rows];
12132        q4matvec(&bytes, &x, rows, cols, &mut got, None);
12133        let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
12134        for r in 0..rows {
12135            assert!(
12136                (got[r] - expect[r]).abs() < 2e-3 * scale,
12137                "row {r}: {} vs {} (outlier term must be exact)",
12138                got[r],
12139                expect[r]
12140            );
12141        }
12142    }
12143
12144    /// The fused q1t matvec must equal the reference (dequant_q1t → dot),
12145    /// including the ternary zero level and the binary-searched outlier
12146    /// overlay. Guards the mmap kernel that makes a 12B q1t runnable.
12147    #[test]
12148    fn q1t_matvec_matches_reference() {
12149        use cortiq_core::quant::{dequant_q1t, f32_to_f16};
12150        let (rows, cols) = (3usize, 64usize); // gpr = 2
12151        let gpr = cols / GROUP_SIZE;
12152        let scales = [0.5f32, 0.3, 0.7, 0.2, 0.6, 0.15];
12153        // Overlay (must be sorted by flat index): a few spikes across rows.
12154        let outliers: [(u32, f32); 3] = [(5, 9.0), (70, -4.5), (150, 3.25)];
12155        let is_out = |flat: usize| outliers.iter().any(|&(i, _)| i as usize == flat);
12156        let mut bytes = Vec::new();
12157        for r in 0..rows {
12158            for g in 0..gpr {
12159                bytes.extend_from_slice(&f32_to_f16(scales[r * gpr + g]).to_le_bytes());
12160                let mut c = [0u8; 7];
12161                for k in 0..GROUP_SIZE {
12162                    // Encoder invariant: code 0 at outlier positions.
12163                    let code = if is_out(r * cols + g * GROUP_SIZE + k) {
12164                        0
12165                    } else {
12166                        ((k + r * 3 + g) % 3) as u8 // 0,1,2
12167                    };
12168                    cortiq_core::quant::q1t_pack(&mut c, k, code);
12169                }
12170                bytes.extend_from_slice(&c);
12171            }
12172        }
12173        // Per-row overlay: [u32 row_ptr[rows+1]] then [(u16 col, f16 val)] by
12174        // row (outliers are sorted by flat index → already grouped by row).
12175        let mut row_ptr = vec![0u32; rows + 1];
12176        for &(idx, _) in &outliers {
12177            row_ptr[idx as usize / cols + 1] += 1;
12178        }
12179        for r in 0..rows {
12180            row_ptr[r + 1] += row_ptr[r];
12181        }
12182        for &p in &row_ptr {
12183            bytes.extend_from_slice(&p.to_le_bytes());
12184        }
12185        for &(idx, v) in &outliers {
12186            bytes.extend_from_slice(&((idx as usize % cols) as u16).to_le_bytes());
12187            bytes.extend_from_slice(&f32_to_f16(v).to_le_bytes());
12188        }
12189
12190        let mut refw = vec![0f32; rows * cols];
12191        dequant_q1t(&bytes, rows, cols, &mut refw);
12192        // On-grid activations (±1, amax 1) so the int8 SDOT path reconstructs
12193        // x exactly and matches the f32 reference (same trick as the q1 test).
12194        let x: Vec<f32> = (0..cols)
12195            .map(|j| if j % 3 == 0 { 1.0 } else { -1.0 })
12196            .collect();
12197        let mut expect = vec![0f32; rows];
12198        for r in 0..rows {
12199            let mut a = 0.0f32;
12200            for j in 0..cols {
12201                a += refw[r * cols + j] * x[j];
12202            }
12203            expect[r] = a;
12204        }
12205        let tol = |e: f32| 1e-3 * e.abs().max(1e-3);
12206        let mut got = vec![0f32; rows];
12207        q1t_matvec(&bytes, &x, rows, cols, &mut got, None);
12208        for r in 0..rows {
12209            assert!(
12210                (got[r] - expect[r]).abs() < tol(expect[r]),
12211                "row {r}: {} vs {}",
12212                got[r],
12213                expect[r]
12214            );
12215        }
12216        // matmat (b=2, f32 decode path) must agree too.
12217        let x2: Vec<f32> = x.iter().chain(x.iter()).copied().collect();
12218        let mut gm = vec![0f32; 2 * rows];
12219        q1t_matmat(&bytes, &x2, 2, rows, cols, &mut gm, None);
12220        for r in 0..rows {
12221            assert!((gm[r] - expect[r]).abs() < tol(expect[r]));
12222            assert!((gm[rows + r] - expect[r]).abs() < tol(expect[r]));
12223        }
12224        // Fused pair (q1t_matvec2) must equal two single matvecs
12225        // bit-for-bit: same unpack, same group order, same f32
12226        // accumulation per stream. Distinct x2 exercises both lanes.
12227        let xb: Vec<f32> = (0..cols)
12228            .map(|j| if j % 5 == 0 { -1.0 } else { 1.0 })
12229            .collect();
12230        let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
12231        q1t_matvec(&bytes, &x, rows, cols, &mut s1, None);
12232        q1t_matvec(&bytes, &xb, rows, cols, &mut s2, None);
12233        let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
12234        q1t_matvec2(&bytes, &x, &xb, rows, cols, &mut p1, &mut p2, None);
12235        assert_eq!(p1, s1, "q1t pair lane 1 ≠ single matvec");
12236        assert_eq!(p2, s2, "q1t pair lane 2 ≠ single matvec");
12237    }
12238
12239    /// Pair == 2×matvec with an ODD group count (the kernel's tail
12240    /// group) and no overlay section.
12241    #[test]
12242    fn q1t_matvec2_odd_gpr_matches_singles() {
12243        use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_pack};
12244        let (rows, cols) = (5usize, 96usize); // gpr = 3 → paired + tail
12245        let gpr = cols / GROUP_SIZE;
12246        let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE);
12247        for r in 0..rows {
12248            for g in 0..gpr {
12249                bytes.extend_from_slice(&f32_to_f16(0.1 + 0.05 * (r + g) as f32).to_le_bytes());
12250                let mut c = [0u8; 7];
12251                for k in 0..GROUP_SIZE {
12252                    q1t_pack(&mut c, k, ((k * 7 + r * 5 + g * 3) % 3) as u8);
12253                }
12254                bytes.extend_from_slice(&c);
12255            }
12256        }
12257        let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.31).sin()).collect();
12258        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).cos()).collect();
12259        let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
12260        q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
12261        q1t_matvec(&bytes, &x2, rows, cols, &mut s2, None);
12262        let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
12263        q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
12264        assert_eq!(p1, s1, "odd-gpr pair lane 1 ≠ single");
12265        assert_eq!(p2, s2, "odd-gpr pair lane 2 ≠ single");
12266    }
12267
12268    // Speed A/B: fused pair (one unpack, two streams) vs two single
12269    // matvecs. Single-threaded, FFN-sized, min-of paired in-process.
12270    //   cargo test -p cortiq-engine --release q1t_matvec2_speed -- --ignored --nocapture
12271    #[test]
12272    #[ignore]
12273    fn q1t_matvec2_speed() {
12274        use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_pack};
12275        use std::time::Instant;
12276        let (rows, cols) = (8192usize, 4096usize);
12277        let gpr = cols / GROUP_SIZE;
12278        let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE);
12279        for r in 0..rows {
12280            for g in 0..gpr {
12281                let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
12282                bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
12283                let mut c = [0u8; 7];
12284                for k in 0..GROUP_SIZE {
12285                    q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
12286                }
12287                bytes.extend_from_slice(&c);
12288            }
12289        }
12290        let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.31).sin()).collect();
12291        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).cos()).collect();
12292        let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
12293        let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
12294        // Warm both paths once.
12295        q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
12296        q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
12297        let (mut t_pair, mut t_two) = (f64::MAX, f64::MAX);
12298        for _ in 0..8 {
12299            let t0 = Instant::now();
12300            q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
12301            t_pair = t_pair.min(t0.elapsed().as_secs_f64() * 1000.0);
12302            let t1 = Instant::now();
12303            q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
12304            q1t_matvec(&bytes, &x2, rows, cols, &mut s2, None);
12305            t_two = t_two.min(t1.elapsed().as_secs_f64() * 1000.0);
12306        }
12307        assert_eq!(p1, s1);
12308        assert_eq!(p2, s2);
12309        println!("q1t pair {rows}x{cols}: fused {t_pair:.2} ms | two singles {t_two:.2} ms");
12310    }
12311
12312    // Speed A/B: the base-3-division decode (what the packing commit left in
12313    // place) vs the fused sign-LUT matvec. Both single-threaded, same bytes.
12314    //   cargo test -p cortiq-engine q1t_matvec_speed -- --ignored --nocapture
12315    #[test]
12316    #[ignore]
12317    fn q1t_matvec_speed() {
12318        use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_code, q1t_pack};
12319        use std::time::Instant;
12320        let (rows, cols) = (8192usize, 4096usize); // FFN-sized
12321        let gpr = cols / GROUP_SIZE;
12322        let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE + 16);
12323        for r in 0..rows {
12324            for g in 0..gpr {
12325                let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
12326                bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
12327                let mut c = [0u8; 7];
12328                for k in 0..GROUP_SIZE {
12329                    q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
12330                }
12331                bytes.extend_from_slice(&c);
12332            }
12333        }
12334        let (n, stride) = (rows * cols, 40usize); // ~2.5% outliers, per-row overlay
12335        let mut row_ptr = vec![0u32; rows + 1];
12336        let mut idx = 0usize;
12337        while idx < n {
12338            row_ptr[idx / cols + 1] += 1;
12339            idx += stride;
12340        }
12341        for r in 0..rows {
12342            row_ptr[r + 1] += row_ptr[r];
12343        }
12344        for &p in &row_ptr {
12345            bytes.extend_from_slice(&p.to_le_bytes());
12346        }
12347        let mut idx = 0usize;
12348        while idx < n {
12349            bytes.extend_from_slice(&((idx % cols) as u16).to_le_bytes());
12350            bytes.extend_from_slice(&f32_to_f16((idx % 13) as f32 * 0.1 - 0.6).to_le_bytes());
12351            idx += stride;
12352        }
12353        // On-grid ±1 so the fast path's int8 SDOT is exact vs the f32 "slow"
12354        // reference (the A/B is a timing check; values must still agree).
12355        let x: Vec<f32> = (0..cols)
12356            .map(|j| if j % 3 == 0 { 1.0 } else { -1.0 })
12357            .collect();
12358        let (rp_off, ent_off, has_ov) = q1t_overlay(&bytes, rows * gpr * Q1T_TILE, rows);
12359
12360        // "before": base-3 division decode into a buffer, then dot.
12361        let slow = |out: &mut [f32]| {
12362            let mut buf = vec![0f32; cols];
12363            for r in 0..rows {
12364                for g in 0..gpr {
12365                    let off = (r * gpr + g) * Q1T_TILE;
12366                    let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
12367                    let codes = &bytes[off + 2..off + Q1T_TILE];
12368                    for k in 0..GROUP_SIZE {
12369                        buf[g * GROUP_SIZE + k] = match q1t_code(codes, k) {
12370                            1 => s,
12371                            2 => -s,
12372                            _ => 0.0,
12373                        };
12374                    }
12375                }
12376                out[r] = q1t_row_outlier_correction(&bytes, r, rp_off, ent_off, has_ov, &x)
12377                    + (0..cols).map(|j| buf[j] * x[j]).sum::<f32>();
12378            }
12379        };
12380        let iters = 5;
12381        let mut a = vec![0f32; rows];
12382        slow(&mut a); // warm
12383        let t = Instant::now();
12384        for _ in 0..iters {
12385            slow(&mut a);
12386        }
12387        let slow_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
12388
12389        let mut b = vec![0f32; rows];
12390        q1t_matvec(&bytes, &x, rows, cols, &mut b, None); // warm
12391        let t = Instant::now();
12392        for _ in 0..iters {
12393            q1t_matvec(&bytes, &x, rows, cols, &mut b, None);
12394        }
12395        let fast_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
12396
12397        for r in 0..rows {
12398            assert!((a[r] - b[r]).abs() < 1e-2, "mismatch row {r}");
12399        }
12400        println!(
12401            "q1t matvec {rows}x{cols} (1 thread): div-decode {slow_ms:.2} ms  fused-LUT {fast_ms:.2} ms  => {:.2}x",
12402            slow_ms / fast_ms
12403        );
12404    }
12405}
12406
12407#[cfg(test)]
12408mod gemm_bench {
12409    /// `cargo test -p cortiq-engine --release q4tp_matmat_throughput -- --ignored --nocapture`
12410    /// Times the batched q4tp GEMM at the shapes the image DiT runs
12411    /// (b=296 tokens, 2304 -> 9216), on synthetic bytes: no model, no
12412    /// mmap, no thermal drift over minutes — a kernel change shows up
12413    /// here in seconds where a full render hides it in noise.
12414    ///
12415    /// On macOS add `CMF_ACCEL=0`: this shape is over the 500k-cell mark
12416    /// where the matmat hands off to Accelerate's dequant sgemm, and
12417    /// without the opt-out both rows below measure the AMX, not the
12418    /// kernel under test.
12419    #[test]
12420    #[ignore]
12421    fn q4tp_matmat_throughput() {
12422        // 296 is a prompt-encode batch; the image DiT runs 2085 at
12423        // 512x512, where the activation panel stops fitting L2 and the
12424        // loop's shape starts to matter more than its instructions.
12425        let b: usize = std::env::var("CMF_BENCH_B")
12426            .ok()
12427            .and_then(|v| v.parse().ok())
12428            .unwrap_or(296);
12429        let (rows, cols) = (9216usize, 2304usize);
12430        let (_, _, _) = (rows, cols, b);
12431        let total =
12432            cortiq_core::quant::expected_nbytes(cortiq_core::TensorDtype::Q4TiledP, &[rows, cols])
12433                .unwrap();
12434        // Random nibbles are fine, but the row params are f16 (lo, step)
12435        // of a geometric ladder: garbage there gives exp2 of a huge
12436        // exponent, the scales come back inf, and the whole bench times
12437        // NaN arithmetic instead of the kernel.
12438        let (params_off, codes_off, _) = cortiq_core::quant::q4tp_sections(rows, cols);
12439        let mut bytes: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
12440        let lo = cortiq_core::quant::f32_to_f16(-4.0);
12441        let step = cortiq_core::quant::f32_to_f16(0.1);
12442        for r in 0..rows {
12443            let o = params_off + r * 4;
12444            bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
12445            bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
12446        }
12447        let _ = codes_off;
12448        let xs: Vec<f32> = (0..b * cols)
12449            .map(|i| ((i % 97) as f32 - 48.0) / 48.0)
12450            .collect();
12451        let mut out = vec![0f32; b * rows];
12452        let pool = crate::pool::Pool::from_env();
12453        // A shared 48-core stand drifts ±25% run to run, which is wider
12454        // than any kernel change worth making. So: alternate the two
12455        // kernels inside one process and keep the BEST time for
12456        // each. Interleaving makes both see the same interference, and a
12457        // minimum is the one statistic another tenant cannot inflate.
12458        super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
12459        let reps: usize = std::env::var("CMF_BENCH_REPS")
12460            .ok()
12461            .and_then(|v| v.parse().ok())
12462            .unwrap_or(10);
12463        let mut best = [f64::MAX; 2];
12464        let mut sums = [0f32; 2];
12465        for _ in 0..reps {
12466            for (k, w) in [(0usize, 1u8), (1usize, 2u8)] {
12467                super::Q4TP_ALT.store(w, std::sync::atomic::Ordering::Relaxed);
12468                let t = std::time::Instant::now();
12469                super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
12470                best[k] = best[k].min(t.elapsed().as_secs_f64());
12471                sums[k] = out.iter().take(64).sum::<f32>();
12472            }
12473        }
12474        let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
12475        for (k, name) in ["previous", "tuned   "].iter().enumerate() {
12476            println!(
12477                "q4tp matmat {rows}x{cols} b={b} {name}: {:.1} ms  {:.1} GFLOP/s  (checksum {:.3})",
12478                best[k] * 1e3,
12479                flops / best[k] / 1e9,
12480                sums[k]
12481            );
12482        }
12483        assert!(
12484            (sums[0] - sums[1]).abs() < 1e-2,
12485            "the tuned kernel changed the result: {} vs {}",
12486            sums[0],
12487            sums[1]
12488        );
12489    }
12490
12491    /// The blocked kernel must agree with the per-column path exactly —
12492    /// same weights, same activation split, only a different instruction
12493    /// mix. Shapes are chosen to hit the awkward cases: a column count
12494    /// that leaves an odd group (the 512-bit kernel does two at a time),
12495    /// and a batch that does not divide by four.
12496    #[test]
12497    fn q4tp_matmat_blocked_matches_scalar() {
12498        use std::sync::atomic::Ordering::Relaxed;
12499        // The last shape carries the image DiT's column count — 2304, so
12500        // 72 groups of accumulation, which is where a reordered sum can
12501        // actually drift — and runs through the thread pool, since the
12502        // blocked path splits rows across workers. Its row count stays
12503        // under 500k cells on purpose: above that, macOS diverts the whole
12504        // matmat to the Accelerate/AMX dequant sgemm and neither kernel
12505        // here would run.
12506        for &(rows, cols, b) in &[
12507            (64usize, 128usize, 7usize),
12508            (33, 96, 4),
12509            (16, 256, 9),
12510            (192, 2304, 37),
12511        ] {
12512            let total = cortiq_core::quant::expected_nbytes(
12513                cortiq_core::TensorDtype::Q4TiledP,
12514                &[rows, cols],
12515            )
12516            .unwrap();
12517            let (params_off, _, _) = cortiq_core::quant::q4tp_sections(rows, cols);
12518            let mut bytes: Vec<u8> = (0..total).map(|i| (i * 61 % 251) as u8).collect();
12519            let lo = cortiq_core::quant::f32_to_f16(-4.0);
12520            let step = cortiq_core::quant::f32_to_f16(0.1);
12521            for r in 0..rows {
12522                let o = params_off + r * 4;
12523                bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
12524                bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
12525            }
12526            let xs: Vec<f32> = (0..b * cols)
12527                .map(|i| ((i % 89) as f32 - 44.0) / 44.0)
12528                .collect();
12529            let mut got = vec![0f32; b * rows];
12530            let mut want = vec![0f32; b * rows];
12531            let gpr = cols / 32;
12532            let view = super::Q4tpView::new(&bytes, rows, cols);
12533            let pool = crate::pool::Pool::from_env();
12534            super::Q4TP_ALT.store(2, Relaxed);
12535            super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut got, pool.as_deref());
12536            super::Q4TP_ALT.store(1, Relaxed);
12537            super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut want, pool.as_deref());
12538            super::Q4TP_ALT.store(0, Relaxed);
12539            // Measured against the output's scale, not cell by cell: a
12540            // dot product of 2304 terms lands near zero wherever the row
12541            // and the activation nearly cancel, and there a per-cell
12542            // ratio reports 1e-3 for an absolute error of 5e-6 — f32's
12543            // own rounding, reordered. What must stay small is the error
12544            // relative to what the layer actually outputs.
12545            let scale = want.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
12546            let (mut worst, mut at) = (0f32, 0usize);
12547            for (i, (g, w)) in got.iter().zip(&want).enumerate() {
12548                if (g - w).abs() > worst {
12549                    worst = (g - w).abs();
12550                    at = i;
12551                }
12552            }
12553            assert!(
12554                worst <= 1e-4 * scale,
12555                "{rows}x{cols} b={b}: blocked and scalar disagree by {worst:.3e} \
12556                 (scale {scale:.3e}) at cell {at}: {} vs {}",
12557                got[at],
12558                want[at]
12559            );
12560
12561            // "Same speed, no quality loss" is a claim about which answer
12562            // is RIGHT, not about which two agree. Both paths sum the same
12563            // 2304 products in different orders, so f64 decides: the
12564            // blocked kernel keeps sixteen partial sums and folds them at
12565            // the end, which is a shallower addition tree than the
12566            // per-column path's running scalar, and it must not be worse.
12567            let (mut e_blocked, mut e_scalar) = (0f64, 0f64);
12568            for bi in 0..b {
12569                let act = super::split_act(&xs[bi * cols..(bi + 1) * cols]);
12570                for r in 0..rows {
12571                    let mut sc = vec![0f32; gpr];
12572                    view.scales_into(r, gpr, &mut sc);
12573                    let mut exact = 0f64;
12574                    for j in 0..cols {
12575                        let (w, sq) = super::q4tp_outlier(view.nib, r, gpr, j, &sc);
12576                        exact += w as f64 * sq as f64 * act.xq[j] as f64;
12577                    }
12578                    exact *= act.sx as f64;
12579                    for &(j, xv) in &act.outliers {
12580                        let (w, sq) = super::q4tp_outlier(view.nib, r, gpr, j, &sc);
12581                        exact += w as f64 * sq as f64 * xv as f64;
12582                    }
12583                    let i = bi * rows + r;
12584                    e_blocked = e_blocked.max((got[i] as f64 - exact).abs());
12585                    e_scalar = e_scalar.max((want[i] as f64 - exact).abs());
12586                }
12587            }
12588            println!(
12589                "{rows}x{cols} b={b}: worst error vs f64 — blocked {e_blocked:.3e}, \
12590                 per-column {e_scalar:.3e}"
12591            );
12592            // An absolute bar, not a race between the two: at these
12593            // magnitudes both sit in f32's last bits, and on a small shape
12594            // whichever one happens to round the unluckiest cell "wins" by
12595            // a factor the next seed reverses.
12596            assert!(
12597                e_blocked <= 1e-5 * scale as f64 && e_scalar <= 1e-5 * scale as f64,
12598                "{rows}x{cols} b={b}: error against f64 too large — blocked \
12599                 {e_blocked:.3e}, per-column {e_scalar:.3e}, scale {scale:.3e}"
12600            );
12601        }
12602    }
12603
12604    /// The q4t twin of the throughput bench, same shape and rules, so the
12605    /// two quantisations' batch kernels can be read against each other.
12606    /// `cargo test -p cortiq-engine --release q4t_matmat_throughput -- --ignored --nocapture`
12607    #[test]
12608    #[ignore]
12609    fn q4t_matmat_throughput() {
12610        let (rows, cols, b) = (9216usize, 2304usize, 296usize);
12611        let total =
12612            cortiq_core::quant::expected_nbytes(cortiq_core::TensorDtype::Q4Tiled, &[rows, cols])
12613                .unwrap();
12614        // q4t carries a per-group f16 scale in the tile's first two bytes;
12615        // random bytes there decode to inf and the bench would time NaNs.
12616        let mut bytes: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
12617        let sc = cortiq_core::quant::f32_to_f16(0.02);
12618        for t in bytes.chunks_mut(super::Q4_TILE) {
12619            t[..2].copy_from_slice(&sc.to_le_bytes());
12620        }
12621        let xs: Vec<f32> = (0..b * cols)
12622            .map(|i| ((i % 97) as f32 - 48.0) / 48.0)
12623            .collect();
12624        let mut out = vec![0f32; b * rows];
12625        let pool = crate::pool::Pool::from_env();
12626        super::q4t_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
12627        let reps: usize = std::env::var("CMF_BENCH_REPS")
12628            .ok()
12629            .and_then(|v| v.parse().ok())
12630            .unwrap_or(10);
12631        let mut best = f64::MAX;
12632        for _ in 0..reps {
12633            let t = std::time::Instant::now();
12634            super::q4t_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
12635            best = best.min(t.elapsed().as_secs_f64());
12636        }
12637        let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
12638        println!(
12639            "q4t matmat {rows}x{cols} b={b}: {:.1} ms  {:.1} GFLOP/s  (checksum {:.3})",
12640            best * 1e3,
12641            flops / best / 1e9,
12642            out.iter().take(64).sum::<f32>()
12643        );
12644    }
12645}