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").map(|v| v != "0").unwrap_or(true)
130            })
131        }
132    }
133}
134
135static BLOCKED_OVERRIDE: std::sync::atomic::AtomicU8 = std::sync::atomic::AtomicU8::new(0);
136
137/// Force the blocked GEMM on or off, ignoring the environment; `None`
138/// restores it. For tests that need to run BOTH paths and compare them:
139/// `blocked_enabled` caches its answer for the life of the process, which
140/// is right when the environment is the only input, but leaves a test that
141/// flips `CMF_X86_BLOCKED` between two calls comparing a path against
142/// itself — or against whatever a test running in parallel latched first.
143pub fn set_blocked_override(on: Option<bool>) {
144    let v = match on {
145        None => 0,
146        Some(false) => 1,
147        Some(true) => 2,
148    };
149    BLOCKED_OVERRIDE.store(v, std::sync::atomic::Ordering::Relaxed);
150}
151
152fn gpu_lmhead_enabled() -> bool {
153    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
154    *ON.get_or_init(|| std::env::var("CMF_GPU_LMHEAD").map(|v| v != "0").unwrap_or(true))
155}
156
157fn gpu_split_frac() -> f32 {
158    static F: std::sync::OnceLock<f32> = std::sync::OnceLock::new();
159    *F.get_or_init(|| {
160        std::env::var("CMF_GPU_SPLIT")
161            .ok()
162            .and_then(|v| v.parse::<f32>().ok())
163            .unwrap_or(0.5)
164            .clamp(0.0, 1.0)
165    })
166}
167
168impl QTensor {
169    pub fn from_f32(data: Vec<f32>, rows: usize, cols: usize) -> Self {
170        debug_assert_eq!(data.len(), rows * cols);
171        Self::F32 { data, rows, cols }
172    }
173
174    /// Wrap a directory tensor without dequantizing the payload.
175    /// Falls back to dequantized f32 for dtypes without a fused kernel.
176    pub fn from_model(model: &Arc<CmfModel>, name: &str) -> Result<Self, String> {
177        // Indexed lookup: the linear directory scan made pipeline build
178        // O(N²) on MoE/skills files with thousands of tensors.
179        let idx = model
180            .tensor_index(name)
181            .ok_or_else(|| format!("tensor '{name}' not found in CMF directory"))?;
182        let entry = &model.tensors[idx];
183        if entry.shape.len() != 2 {
184            return Err(format!("QTensor::from_model needs 2-D, got '{name}'"));
185        }
186        let (rows, cols) = (entry.shape[0], entry.shape[1]);
187        let bytes = model.entry_bytes(entry);
188
189        match entry.dtype {
190            TensorDtype::Q8Row | TensorDtype::Q8_2f => {
191                let n = rows * cols;
192                let scales_off = n;
193                let row_scale: Vec<f32> = (0..rows)
194                    .map(|o| {
195                        f16_to_f32(u16::from_le_bytes([
196                            bytes[scales_off + o * 2],
197                            bytes[scales_off + o * 2 + 1],
198                        ]))
199                    })
200                    .collect();
201                let col_field: Vec<f32> = if entry.dtype == TensorDtype::Q8_2f {
202                    let col_off = n + rows * 2;
203                    (0..cols)
204                        .map(|i| {
205                            f16_to_f32(u16::from_le_bytes([
206                                bytes[col_off + i * 2],
207                                bytes[col_off + i * 2 + 1],
208                            ]))
209                        })
210                        .collect()
211                } else {
212                    Vec::new()
213                };
214                Ok(Self::Mapped {
215                    model: model.clone(),
216                    idx,
217                    dtype: entry.dtype,
218                    rows,
219                    cols,
220                    row_scale,
221                    col_field,
222                    vbit_offsets: Vec::new(),
223                    repack: q8_repack(bytes, rows, cols),
224                })
225            }
226            // vbit: fused kernel unpacks variable-bit rows from mmap.
227            TensorDtype::Vbit if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
228                model: model.clone(),
229                idx,
230                dtype: entry.dtype,
231                rows,
232                cols,
233                row_scale: Vec::new(),
234                col_field: Vec::new(),
235                vbit_offsets: vbit_row_offsets(bytes, rows, cols),
236                repack: Vec::new(),
237            }),
238            // vbit_ro (§4.2): the offset table comes straight from the
239            // file — no load-time prefix scan; kernels are shared with
240            // legacy vbit (they consume absolute offsets either way).
241            TensorDtype::VbitRo if cols % GROUP_SIZE == 0 => {
242                let (_, off_off, packed_off) = cortiq_core::quant::vbit_ro_sections(rows, cols);
243                let offsets: Vec<usize> = (0..=rows)
244                    .map(|r| packed_off + cortiq_core::quant::vbit_ro_offset(bytes, off_off, r))
245                    .collect();
246                Ok(Self::Mapped {
247                    model: model.clone(),
248                    idx,
249                    dtype: entry.dtype,
250                    rows,
251                    cols,
252                    row_scale: Vec::new(),
253                    col_field: Vec::new(),
254                    vbit_offsets: offsets,
255                    repack: Vec::new(),
256                })
257            }
258            // q4_block: fused kernel reads nibbles straight from mmap —
259            // a 14B q4 file no longer explodes into ×8 f32 RAM.
260            // q4_tiled (§4.3): interleaved [scale][nibbles] tiles — one
261            // sequential memory stream (measured ×1.66 ARM / ×1.13 AVX2
262            // at kernel level over the split layout).
263            TensorDtype::Q4Tiled if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
264                model: model.clone(),
265                idx,
266                dtype: entry.dtype,
267                rows,
268                cols,
269                row_scale: Vec::new(),
270                col_field: Vec::new(),
271                vbit_offsets: Vec::new(),
272                repack: Vec::new(),
273            }),
274            // q4tp (§4.10): nibbles from mmap, scale from the row ladder —
275            // 7.3% less file than q4t at the same 4-bit grid.
276            TensorDtype::Q4TiledP if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
277                model: model.clone(),
278                idx,
279                dtype: entry.dtype,
280                rows,
281                cols,
282                row_scale: Vec::new(),
283                col_field: Vec::new(),
284                vbit_offsets: Vec::new(),
285                repack: Vec::new(),
286            }),
287            // q2tp: 2-bit chunks from mmap, scale from the same row ladder.
288            TensorDtype::Q2TiledP if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
289                model: model.clone(),
290                idx,
291                dtype: entry.dtype,
292                rows,
293                cols,
294                row_scale: Vec::new(),
295                col_field: Vec::new(),
296                vbit_offsets: Vec::new(),
297                repack: Vec::new(),
298            }),
299            TensorDtype::Q4Block if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
300                model: model.clone(),
301                idx,
302                dtype: entry.dtype,
303                rows,
304                cols,
305                row_scale: Vec::new(),
306                col_field: Vec::new(),
307                vbit_offsets: Vec::new(),
308                repack: Vec::new(),
309            }),
310            // q1: binary sign-bit tiles from mmap (1-bit-trained models).
311            TensorDtype::Q1 if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
312                model: model.clone(),
313                idx,
314                dtype: entry.dtype,
315                rows,
316                cols,
317                row_scale: Vec::new(),
318                col_field: Vec::new(),
319                vbit_offsets: Vec::new(),
320                repack: Vec::new(),
321            }),
322            // q1t (ternary + outlier overlay): fused per-row dequant kernel
323            // reads straight from mmap — a 12B q1t stays ~its file size in
324            // RAM instead of dequantizing to ~48 GB of f32.
325            TensorDtype::Q1T if cols % GROUP_SIZE == 0 => Ok(Self::Mapped {
326                model: model.clone(),
327                idx,
328                dtype: entry.dtype,
329                rows,
330                cols,
331                row_scale: Vec::new(),
332                col_field: Vec::new(),
333                vbit_offsets: Vec::new(),
334                repack: Vec::new(),
335            }),
336            // No fused kernel yet → dequantize once (correct, more RAM).
337            _ => {
338                let mut data = vec![0.0f32; rows * cols];
339                cortiq_core::quant::dequant_tensor(entry, bytes, &mut data)?;
340                Ok(Self::from_f32(data, rows, cols))
341            }
342        }
343    }
344
345    /// q1-mapped tensor? (GPU gates: the q1 CPU kernel is
346    /// compute-bound, so offload pays at much smaller shapes than q8.)
347    pub(crate) fn is_q1(&self) -> bool {
348        matches!(
349            self,
350            Self::Mapped {
351                dtype: TensorDtype::Q1,
352                ..
353            }
354        )
355    }
356
357    /// Owned-f32 view (data, rows, cols) — the GDN a/b gate projections
358    /// arrive dequantized (force-f16 in the converter → F32 in RAM).
359    pub(crate) fn f32_parts(&self) -> Option<(&[f32], usize, usize)> {
360        match self {
361            Self::F32 { data, rows, cols } => Some((data, *rows, *cols)),
362            _ => None,
363        }
364    }
365
366    /// (directory idx, rows, cols) of a q1-mapped tensor — the
367    /// whole-block GPU path resolves offsets itself.
368    /// (idx, rows, cols) of a mapped tensor the whole-token GPU graph can drive
369    /// — Q1, Q1T or Q4-block (it resolves the offset and picks the kernel by
370    /// dtype). Q4-block lets a precise down_proj/lm_head stay on-device.
371    /// Named `q1_parts` for historical reasons.
372    pub(crate) fn q1_parts(&self) -> Option<(usize, usize, usize)> {
373        match self {
374            #[cfg(target_os = "macos")]
375            Self::Mapped {
376                dtype: TensorDtype::Q1T,
377                ..
378            } if !crate::gpu::metal_q1t_enabled() => None,
379            Self::Mapped {
380                idx,
381                dtype:
382                    TensorDtype::Q1
383                    | TensorDtype::Q1T
384                    | TensorDtype::Q4Block
385                    | TensorDtype::Q4Tiled
386                    // Q2TiledP deliberately absent: the Metal graph has no
387                    // q2tp kernel, and advertising it here made the block
388                    // plan truncate mid-run at the first q2tp layer.
389                    | TensorDtype::Q4TiledP
390                    | TensorDtype::Q8Row
391                    | TensorDtype::Q8_2f,
392                rows,
393                cols,
394                ..
395            } => Some((*idx, *rows, *cols)),
396            _ => None,
397        }
398    }
399
400    /// (directory idx, rows, cols) of a q4_tiled mapped tensor. The
401    /// chunk-prefill graph takes it in the same 4-tuple slot as
402    /// `q8_row_parts` with an EMPTY row_scale — q4t carries its scales
403    /// inside the 18-byte tiles, and the empty slice is what tells the
404    /// encoder to reach for the q4t kernels.
405    pub(crate) fn q4t_parts(&self) -> Option<(usize, usize, usize)> {
406        match self {
407            Self::Mapped {
408                idx,
409                dtype: TensorDtype::Q4Tiled,
410                rows,
411                cols,
412                ..
413            } => Some((*idx, *rows, *cols)),
414            _ => None,
415        }
416    }
417
418    /// (directory idx, rows, cols) of a q4tp mapped tensor. Same empty-scale
419    /// slot as `q4t_parts` in the chunk graph — the encoder tells the two
420    /// apart by the tensor's dtype, not by the slot.
421    pub(crate) fn q4tp_parts(&self) -> Option<(usize, usize, usize)> {
422        match self {
423            Self::Mapped {
424                idx,
425                dtype: TensorDtype::Q4TiledP,
426                rows,
427                cols,
428                ..
429            } => Some((*idx, *rows, *cols)),
430            _ => None,
431        }
432    }
433
434    /// (directory idx, rows, cols, row_scale) of a plain q8_row mapped
435    /// tensor — the chunk-prefill GPU graph resolves offsets itself.
436    /// q8_2f is excluded on purpose: its column field would need a
437    /// prescale stage on the device.
438    pub(crate) fn q8_row_parts(&self) -> Option<(usize, usize, usize, &[f32])> {
439        match self {
440            Self::Mapped {
441                idx,
442                dtype: TensorDtype::Q8Row,
443                rows,
444                cols,
445                row_scale,
446                col_field,
447                ..
448            } if col_field.is_empty() => Some((*idx, *rows, *cols, row_scale)),
449            _ => None,
450        }
451    }
452
453    /// The layout this tensor is stored in, when it is mapped from a model.
454    /// The frames branch on it — a q2tp gate against a q4tp down is a real
455    /// combination in the 2-bit profile and needs a different kernel.
456    pub fn model_dtype(&self) -> Option<cortiq_core::TensorDtype> {
457        match self {
458            Self::Mapped { dtype, .. } => Some(*dtype),
459            _ => None,
460        }
461    }
462
463    /// The tensor's index in the model directory, when it is mapped from one.
464    /// The GPU frames bind by index rather than by name — a name lookup per
465    /// layer per token is not free, and the index is what the device cache is
466    /// keyed on anyway.
467    pub fn model_idx(&self) -> Option<usize> {
468        match self {
469            Self::Mapped { idx, .. } => Some(*idx),
470            _ => None,
471        }
472    }
473
474    /// The model this tensor is mapped from, when it is mapped at all. The
475    /// GPU frames need the container to reach the bytes; a QTensor already
476    /// holds it, and threading a second handle down every call site to say
477    /// the same thing invites the two to disagree.
478    pub fn model_arc(&self) -> Option<std::sync::Arc<cortiq_core::CmfModel>> {
479        match self {
480            Self::Mapped { model, .. } => Some(model.clone()),
481            _ => None,
482        }
483    }
484
485    pub fn rows(&self) -> usize {
486        match self {
487            Self::F32 { rows, .. } | Self::Mapped { rows, .. } => *rows,
488        }
489    }
490
491    /// Mapped q4t handle (model + directory index) — the fused GPU FFN
492    /// needs the raw file coordinates of its three projections.
493    pub(crate) fn mapped_q4t(&self) -> Option<(&Arc<CmfModel>, usize)> {
494        match self {
495            Self::Mapped {
496                model,
497                idx,
498                dtype: TensorDtype::Q4Tiled,
499                ..
500            } => Some((model, *idx)),
501            _ => None,
502        }
503    }
504
505    /// Same slot as `mapped_q4t` for a q4tp tensor — the fused DiT FFN picks
506    /// its kernels by which of the two answers.
507    pub fn mapped_q4tp(&self) -> Option<(&Arc<CmfModel>, usize)> {
508        match self {
509            Self::Mapped {
510                model,
511                idx,
512                dtype: TensorDtype::Q4TiledP,
513                ..
514            } => Some((model, *idx)),
515            _ => None,
516        }
517    }
518
519    /// (model, tensor idx) for a q2tp mapped weight — the 2-bit twin of
520    /// `mapped_q4tp`, used by the mixed MoE profile.
521    pub fn mapped_q2tp(&self) -> Option<(&Arc<CmfModel>, usize)> {
522        match self {
523            Self::Mapped {
524                model,
525                idx,
526                dtype: TensorDtype::Q2TiledP,
527                ..
528            } => Some((model, *idx)),
529            _ => None,
530        }
531    }
532
533    pub fn cols(&self) -> usize {
534        match self {
535            Self::F32 { cols, .. } | Self::Mapped { cols, .. } => *cols,
536        }
537    }
538
539    /// (model, tensor idx) for a q1 mapped weight — the wgpu token graph
540    /// keys its resident VRAM cache by idx. None for any other dtype/kind.
541    pub fn mapped_q1(&self) -> Option<(&std::sync::Arc<CmfModel>, usize)> {
542        match self {
543            Self::Mapped {
544                model,
545                idx,
546                dtype: TensorDtype::Q1,
547                ..
548            } => Some((model, *idx)),
549            _ => None,
550        }
551    }
552
553    /// (model, idx, kind, row_scale) for a graph-capable mapped weight. kind:
554    /// 0=q8_row (per-row scales), 1=q1, 2=q4_tiled, 3=q1t (tile-embedded, no
555    /// rs). None for dtypes the token graph does not handle (q8_2f/q4_block/vbit).
556    pub fn graph_weight(&self) -> Option<(&std::sync::Arc<CmfModel>, usize, u8, &[f32])> {
557        match self {
558            Self::Mapped {
559                model,
560                idx,
561                dtype: TensorDtype::Q8Row,
562                row_scale,
563                ..
564            } => Some((model, *idx, 0, row_scale.as_slice())),
565            Self::Mapped {
566                model,
567                idx,
568                dtype: TensorDtype::Q1,
569                ..
570            } => Some((model, *idx, 1, &[])),
571            // Q4Tiled is kind 5, NOT 2: both carried 2 historically, and
572            // the wgpu token graph fed 18B interleaved tiles to the
573            // split-layout q4b kernel — garbage output on q4t models
574            // (caught by an end-to-end answer check on real Vulkan).
575            Self::Mapped {
576                model,
577                idx,
578                dtype: TensorDtype::Q4Tiled,
579                ..
580            } => Some((model, *idx, 5, &[])),
581            // Kind 6, not 5: q4tp's nibble stride and scale planes differ,
582            // and feeding them to the q4t kernel is exactly the mistake that
583            // produced garbage when Q4Tiled shared kind 2 with Q4Block.
584            Self::Mapped {
585                model,
586                idx,
587                dtype: TensorDtype::Q4TiledP,
588                ..
589            } => Some((model, *idx, 6, &[])),
590            Self::Mapped {
591                model,
592                idx,
593                dtype: TensorDtype::Q4Block,
594                ..
595            } => Some((model, *idx, 2, &[])),
596            Self::Mapped {
597                model,
598                idx,
599                dtype: TensorDtype::Q1T,
600                ..
601            } => Some((model, *idx, 3, &[])),
602            _ => None,
603        }
604    }
605
606    /// Dense f32 view — only for owned tensors. Masked/sparse execution
607    /// paths require it; quantized weights don't support masks yet.
608    pub fn as_f32(&self) -> Option<&[f32]> {
609        match self {
610            Self::F32 { data, .. } => Some(data),
611            Self::Mapped { .. } => None,
612        }
613    }
614
615    fn quant_bytes(&self) -> &[u8] {
616        match self {
617            Self::Mapped { model, idx, .. } => model.entry_bytes(&model.tensors[*idx]),
618            Self::F32 { .. } => unreachable!("quant_bytes on F32"),
619        }
620    }
621
622    /// Dequantize one row into `dst` (embedding lookup).
623    pub fn row_f32(&self, r: usize, dst: &mut [f32]) {
624        let cols = self.cols();
625        debug_assert_eq!(dst.len(), cols);
626        match self {
627            Self::F32 { data, .. } => dst.copy_from_slice(&data[r * cols..(r + 1) * cols]),
628            Self::Mapped {
629                dtype,
630                row_scale,
631                col_field,
632                vbit_offsets,
633                ..
634            } => {
635                if *dtype == TensorDtype::Q4Tiled {
636                    let bytes = self.quant_bytes();
637                    let gpr = cols / GROUP_SIZE;
638                    for gi in 0..gpr {
639                        let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
640                        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
641                        for (k, &b) in tile[2..].iter().enumerate() {
642                            dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
643                            dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
644                        }
645                    }
646                    return;
647                }
648                if *dtype == TensorDtype::Q4TiledP {
649                    let bytes = self.quant_bytes();
650                    let gpr = cols / GROUP_SIZE;
651                    let v = Q4tpView::new(bytes, self.rows(), cols);
652                    let mut sc = vec![0f32; gpr];
653                    v.scales_into(r, gpr, &mut sc);
654                    for gi in 0..gpr {
655                        let tile = &v.nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
656                        let s = sc[gi];
657                        for (k, &b) in tile.iter().enumerate() {
658                            dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
659                            dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
660                        }
661                    }
662                    return;
663                }
664                if *dtype == TensorDtype::Q2TiledP {
665                    let bytes = self.quant_bytes();
666                    let gpr = cols / GROUP_SIZE;
667                    let v = Q4tpView::new_q2(bytes, self.rows(), cols);
668                    let mut sc = vec![0f32; gpr];
669                    v.scales_into(r, gpr, &mut sc);
670                    for gi in 0..gpr {
671                        let ch =
672                            &v.nib[(r * gpr + gi) * Q2TP_CHUNK..(r * gpr + gi + 1) * Q2TP_CHUNK];
673                        let s = sc[gi];
674                        for (k, &b) in ch.iter().enumerate() {
675                            for j in 0..4 {
676                                dst[gi * GROUP_SIZE + k * 4 + j] =
677                                    (((b >> (2 * j)) & 3) as f32 - 1.5) * s;
678                            }
679                        }
680                    }
681                    return;
682                }
683                if *dtype == TensorDtype::Q4Block {
684                    let (packed, scales) = q4_split(self.quant_bytes(), self.rows(), cols);
685                    let gpr = cols / GROUP_SIZE;
686                    for gi in 0..gpr {
687                        let g = r * gpr + gi;
688                        let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
689                        for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
690                            dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
691                            dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
692                        }
693                    }
694                    return;
695                }
696                if *dtype == TensorDtype::Q1 {
697                    let bytes = self.quant_bytes();
698                    let gpr = cols / GROUP_SIZE;
699                    for gi in 0..gpr {
700                        let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
701                        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
702                        for (j, &b) in tile[2..].iter().enumerate() {
703                            for k in 0..8 {
704                                dst[gi * GROUP_SIZE + j * 8 + k] =
705                                    (((b >> k) & 1) as f32 * 2.0 - 1.0) * s;
706                            }
707                        }
708                    }
709                    return;
710                }
711                if *dtype == TensorDtype::Q1T {
712                    let bytes = self.quant_bytes();
713                    let gpr = cols / GROUP_SIZE;
714                    let base_len = self.rows() * gpr * cortiq_core::quant::Q1T_TILE;
715                    for gi in 0..gpr {
716                        let off = (r * gpr + gi) * cortiq_core::quant::Q1T_TILE;
717                        let s = cortiq_core::quant::f16_to_f32(u16::from_le_bytes([
718                            bytes[off],
719                            bytes[off + 1],
720                        ]));
721                        let codes = &bytes[off + 2..off + cortiq_core::quant::Q1T_TILE];
722                        for k in 0..GROUP_SIZE {
723                            dst[gi * GROUP_SIZE + k] = match cortiq_core::quant::q1t_code(codes, k)
724                            {
725                                1 => s,
726                                2 => -s,
727                                _ => 0.0,
728                            };
729                        }
730                    }
731                    // Overlay
732                    let rows = self.rows();
733                    let entries = base_len + (rows + 1) * 4;
734                    if entries <= bytes.len() {
735                        let ptrs = &bytes[base_len..base_len + (rows + 1) * 4];
736                        let r0 = u32::from_le_bytes([
737                            ptrs[r * 4],
738                            ptrs[r * 4 + 1],
739                            ptrs[r * 4 + 2],
740                            ptrs[r * 4 + 3],
741                        ]) as usize;
742                        let r1 = u32::from_le_bytes([
743                            ptrs[(r + 1) * 4],
744                            ptrs[(r + 1) * 4 + 1],
745                            ptrs[(r + 1) * 4 + 2],
746                            ptrs[(r + 1) * 4 + 3],
747                        ]) as usize;
748                        let off = entries + r0 * 4;
749                        for i in 0..r1 - r0 {
750                            let item = &bytes[off + i * 4..off + i * 4 + 4];
751                            let c = u16::from_le_bytes([item[0], item[1]]) as usize;
752                            let v = cortiq_core::quant::f16_to_f32(u16::from_le_bytes([
753                                item[2], item[3],
754                            ]));
755                            if c < cols {
756                                dst[c] = v;
757                            }
758                        }
759                    }
760                    return;
761                }
762                if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
763                    let bytes = self.quant_bytes();
764                    let rows = self.rows();
765                    let ng = cols / GROUP_SIZE;
766                    let bits = &bytes[..rows];
767                    let sc_off = rows;
768                    // Precomputed at load — embedding lookup used to scan
769                    // the bit-widths of every preceding row (O(token_id)).
770                    let off = vbit_offsets[r];
771                    let b = bits[r] as usize;
772                    let l = ((1usize << (b - 1)) - 1) as f32;
773                    let data = &bytes[off..];
774                    let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
775                    for (i, d) in dst.iter_mut().enumerate() {
776                        while nbits < b {
777                            acc = (acc << 8) | data[idx] as u64;
778                            idx += 1;
779                            nbits += 8;
780                        }
781                        let u = ((acc >> (nbits - b)) & ((1u64 << b) - 1)) as f32;
782                        nbits -= b;
783                        let so = (r * ng + i / GROUP_SIZE) * 2;
784                        let sv = f16_to_f32(u16::from_le_bytes([
785                            bytes[sc_off + so],
786                            bytes[sc_off + so + 1],
787                        ]));
788                        *d = (u - l) * sv;
789                    }
790                    return;
791                }
792                let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
793                let s = row_scale[r];
794                match dtype {
795                    TensorDtype::Q8Row => {
796                        for (d, &b) in dst.iter_mut().zip(q) {
797                            *d = (b as i8) as f32 * s;
798                        }
799                    }
800                    TensorDtype::Q8_2f => {
801                        for (i, (d, &b)) in dst.iter_mut().zip(q).enumerate() {
802                            *d = (b as i8) as f32 * s * col_field[i];
803                        }
804                    }
805                    _ => unreachable!(),
806                }
807            }
808        }
809    }
810
811    /// Can this tensor's columns be read cheaply (for sparse down_proj)?
812    /// True for F32/Q8Row/Q8_2f (per-row scale, direct strided access);
813    /// false for group-packed q4/vbit (column access would unpack whole
814    /// groups — sparse execution falls back to f32 for those).
815    pub fn sparse_col_ok(&self) -> bool {
816        match self {
817            Self::F32 { .. } => true,
818            Self::Mapped { dtype, .. } => {
819                matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
820            }
821        }
822    }
823
824    /// down_proj [hidden, inter]: accumulate `w · col(c)` into `out`
825    /// [hidden] — reads ONLY column `c` (one neuron) from the mmap,
826    /// no full-matrix dequant. `out[k] += w · down[k, c]`.
827    pub fn add_col_scaled(&self, c: usize, w: f32, out: &mut [f32]) {
828        let inter = self.cols();
829        let hidden = self.rows();
830        debug_assert_eq!(out.len(), hidden);
831        match self {
832            Self::F32 { data, .. } => {
833                for (k, o) in out.iter_mut().enumerate() {
834                    *o += w * data[k * inter + c];
835                }
836            }
837            Self::Mapped {
838                dtype,
839                row_scale,
840                col_field,
841                ..
842            } => {
843                let q = self.quant_bytes();
844                let colf = if *dtype == TensorDtype::Q8_2f {
845                    col_field[c]
846                } else {
847                    1.0
848                };
849                let wc = w * colf;
850                for (k, o) in out.iter_mut().enumerate() {
851                    let b = q[k * inter + c] as i8 as f32;
852                    *o += wc * b * row_scale[k];
853                }
854            }
855        }
856    }
857
858    /// Dot of row `r` with `x` (gate/up active-neuron path). Reads only
859    /// row `r` from the mmap — no full dequant. q4/vbit dequant the row
860    /// into `scratch` first (rare for active-FFN weights).
861    pub fn row_dot(&self, r: usize, x: &[f32], scratch: &mut [f32]) -> f32 {
862        let cols = self.cols();
863        match self {
864            Self::F32 { data, .. } => {
865                let row = &data[r * cols..(r + 1) * cols];
866                row.iter().zip(x).map(|(w, v)| w * v).sum()
867            }
868            Self::Mapped {
869                dtype,
870                row_scale,
871                col_field,
872                ..
873            } => match dtype {
874                TensorDtype::Q8Row => {
875                    let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
876                    dot_i8_f32(q, x) * row_scale[r]
877                }
878                TensorDtype::Q8_2f => {
879                    let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
880                    dot_i8_col_f32(q, x, col_field) * row_scale[r]
881                }
882                _ => {
883                    self.row_f32(r, scratch);
884                    scratch.iter().zip(x).map(|(w, v)| w * v).sum()
885                }
886            },
887        }
888    }
889
890    /// `out = W · x` (row-major). F32 delegates to the historical
891    /// bit-exact path; Mapped runs the fused int8 kernel.
892    pub fn matvec(&self, x: &[f32], out: &mut [f32], pool: Option<&Pool>) {
893        match self {
894            // NOTE: `out.len()` DRIVES this arm — it computes that many rows,
895            // and `x.len()` is the stride. A short `out` is legitimate here,
896            // which is why the check below lives in the Mapped arm only.
897            Self::F32 { data, .. } => matvec_rows(pool, data, x, out),
898            Self::Mapped {
899                model,
900                idx,
901                dtype,
902                rows,
903                cols,
904                row_scale,
905                col_field,
906                vbit_offsets,
907                repack,
908            } => {
909                let _ = (model, idx);
910                // Every kernel below writes `rows` entries through a raw
911                // pointer, so a short `out` is an out-of-bounds WRITE, not a
912                // wrong answer: it scribbles on the allocator's metadata and
913                // the process aborts much later, somewhere innocent
914                // (`double free or corruption`, `corrupted double-linked
915                // list`). The debug_assert two of the kernels carried is
916                // compiled out of the release — exactly the build where it
917                // matters. Fail here instead, while the caller is still on
918                // the stack to be named.
919                assert!(
920                    out.len() >= *rows && x.len() >= *cols,
921                    "matvec {rows}x{cols}: out {} (need {rows}), x {} (need {cols})",
922                    out.len(),
923                    x.len(),
924                );
925                if *dtype == TensorDtype::Q4Block {
926                    // GPU route (wgpu q4b kernel) for large q4_block matvecs —
927                    // gives NVIDIA/AMD/Intel q4 models a GPU path. Probe keeps
928                    // the winner; Metal returns false → the CPU kernel below.
929                    if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
930                        let t0 = std::time::Instant::now();
931                        match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
932                            crate::gpu::ProbeArm::Gpu => {
933                                if crate::gpu::q4b_matvec(model, *idx, x, *rows, *cols, out) {
934                                    crate::gpu::probe_record(
935                                        crate::gpu::OpClass::Matvec,
936                                        true,
937                                        t0.elapsed(),
938                                    );
939                                    return;
940                                }
941                            }
942                            crate::gpu::ProbeArm::CpuTimed => {
943                                q4matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
944                                crate::gpu::probe_record(
945                                    crate::gpu::OpClass::Matvec,
946                                    false,
947                                    t0.elapsed(),
948                                );
949                                return;
950                            }
951                            crate::gpu::ProbeArm::Cpu => {}
952                        }
953                    }
954                    q4matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
955                    return;
956                }
957                if *dtype == TensorDtype::Q4Tiled {
958                    // GPU route for large q4t matvecs — the lm_head class,
959                    // same shape as the q4tp arm below. The probe keeps the
960                    // winner; a backend without the kernel refuses and the
961                    // CPU path stays.
962                    if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
963                        let t0 = std::time::Instant::now();
964                        let cls = crate::gpu::matvec_class(*rows, *cols);
965                        match crate::gpu::probe_arm(cls) {
966                            crate::gpu::ProbeArm::Gpu => {
967                                if crate::gpu::q4t_matvec(model, *idx, x, *rows, *cols, out) {
968                                    crate::gpu::probe_record(cls, true, t0.elapsed());
969                                    return;
970                                }
971                            }
972                            crate::gpu::ProbeArm::CpuTimed => {
973                                q4t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
974                                crate::gpu::probe_record(cls, false, t0.elapsed());
975                                return;
976                            }
977                            crate::gpu::ProbeArm::Cpu => {}
978                        }
979                    }
980                    q4t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
981                    return;
982                }
983                if *dtype == TensorDtype::Q4TiledP {
984                    // GPU route for large q4tp matvecs — the lm_head class.
985                    // On a q4tp checkpoint the head is the biggest single
986                    // host matvec left in the decode step, and the batched
987                    // kernel at b=1 already exists on both backends. Probe
988                    // keeps the winner, same as q4_block above.
989                    if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
990                        let t0 = std::time::Instant::now();
991                        let cls = crate::gpu::matvec_class(*rows, *cols);
992                        match crate::gpu::probe_arm(cls) {
993                            crate::gpu::ProbeArm::Gpu => {
994                                if crate::gpu::q4tp_matvec(model, *idx, x, *rows, *cols, out) {
995                                    crate::gpu::probe_record(cls, true, t0.elapsed());
996                                    return;
997                                }
998                            }
999                            crate::gpu::ProbeArm::CpuTimed => {
1000                                q4tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1001                                crate::gpu::probe_record(cls, false, t0.elapsed());
1002                                return;
1003                            }
1004                            crate::gpu::ProbeArm::Cpu => {}
1005                        }
1006                    }
1007                    q4tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1008                    return;
1009                }
1010                if *dtype == TensorDtype::Q2TiledP {
1011                    q2tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1012                    return;
1013                }
1014                if *dtype == TensorDtype::Q1 {
1015                    // GPU route for large q1 matvecs (out_proj / lm_head
1016                    // class): the CPU q1 kernel is load-port-bound at
1017                    // ~4 GB/s/core, the GPU one is bandwidth-bound — the
1018                    // probe measures both arms and keeps the winner.
1019                    if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1020                        let t0 = std::time::Instant::now();
1021                        let arm = if crate::gpu::q1_force() {
1022                            crate::gpu::ProbeArm::Gpu
1023                        } else {
1024                            crate::gpu::probe_arm(crate::gpu::OpClass::Matvec)
1025                        };
1026                        match arm {
1027                            crate::gpu::ProbeArm::Gpu => {
1028                                if crate::gpu::q1_matvec(model, *idx, x, *rows, *cols, out) {
1029                                    crate::gpu::probe_record(
1030                                        crate::gpu::OpClass::Matvec,
1031                                        true,
1032                                        t0.elapsed(),
1033                                    );
1034                                    return;
1035                                }
1036                            }
1037                            crate::gpu::ProbeArm::CpuTimed => {
1038                                q1_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1039                                crate::gpu::probe_record(
1040                                    crate::gpu::OpClass::Matvec,
1041                                    false,
1042                                    t0.elapsed(),
1043                                );
1044                                return;
1045                            }
1046                            crate::gpu::ProbeArm::Cpu => {}
1047                        }
1048                    }
1049                    q1_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1050                    return;
1051                }
1052                if *dtype == TensorDtype::Q1T {
1053                    // GPU route for large q1t matvecs: the ternary BASE dot runs
1054                    // on the GPU (load-port-bound on CPU, like q1), then the
1055                    // sparse overlay is added on the CPU. Probe keeps the winner.
1056                    if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1057                        let t0 = std::time::Instant::now();
1058                        match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
1059                            crate::gpu::ProbeArm::Gpu => {
1060                                if crate::gpu::q1t_matvec(model, *idx, x, *rows, *cols, out) {
1061                                    q1t_add_overlay(self.quant_bytes(), x, *rows, *cols, out, pool);
1062                                    crate::gpu::probe_record(
1063                                        crate::gpu::OpClass::Matvec,
1064                                        true,
1065                                        t0.elapsed(),
1066                                    );
1067                                    return;
1068                                }
1069                            }
1070                            crate::gpu::ProbeArm::CpuTimed => {
1071                                q1t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1072                                crate::gpu::probe_record(
1073                                    crate::gpu::OpClass::Matvec,
1074                                    false,
1075                                    t0.elapsed(),
1076                                );
1077                                return;
1078                            }
1079                            crate::gpu::ProbeArm::Cpu => {}
1080                        }
1081                    }
1082                    q1t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1083                    return;
1084                }
1085                if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1086                    vbitmatvec(self.quant_bytes(), vbit_offsets, x, *rows, *cols, out, pool);
1087                    return;
1088                }
1089                let xs = prescale(x, col_field, *dtype);
1090                // D5: large q8 matrices (lm_head-class) — hybrid
1091                // CPU∥GPU: split the rows, both sides compute
1092                // SIMULTANEOUSLY (same math, shared prescale).
1093                // GPU share: CMF_GPU_SPLIT (0..1, default 0.5).
1094                if *rows >= crate::gpu::min_rows()
1095                    && matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
1096                    && gpu_lmhead_enabled()
1097                    && crate::gpu::enabled_here()
1098                {
1099                    // Runtime probe: alternate the hybrid against the
1100                    // pure-CPU matvec, keep whichever is faster HERE.
1101                    let t0 = std::time::Instant::now();
1102                    match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
1103                        crate::gpu::ProbeArm::Gpu => {}
1104                        crate::gpu::ProbeArm::CpuTimed => {
1105                            qmatvec(
1106                                self.quant_bytes(),
1107                                repack,
1108                                row_scale,
1109                                x,
1110                                col_field,
1111                                *dtype,
1112                                *rows,
1113                                *cols,
1114                                out,
1115                                pool,
1116                            );
1117                            crate::gpu::probe_record(
1118                                crate::gpu::OpClass::Matvec,
1119                                false,
1120                                t0.elapsed(),
1121                            );
1122                            return;
1123                        }
1124                        crate::gpu::ProbeArm::Cpu => {
1125                            qmatvec(
1126                                self.quant_bytes(),
1127                                repack,
1128                                row_scale,
1129                                x,
1130                                col_field,
1131                                *dtype,
1132                                *rows,
1133                                *cols,
1134                                out,
1135                                pool,
1136                            );
1137                            return;
1138                        }
1139                    }
1140                    let frac = gpu_split_frac();
1141                    let cpu_rows = ((*rows as f32) * (1.0 - frac)) as usize;
1142                    let (out_cpu, out_gpu) = out.split_at_mut(cpu_rows);
1143                    let bytes = self.quant_bytes();
1144                    let ok = std::thread::scope(|sc| {
1145                        let g = sc.spawn(|| {
1146                            crate::gpu::q8_matvec_range(
1147                                model,
1148                                *idx,
1149                                cpu_rows,
1150                                &row_scale[cpu_rows..],
1151                                &xs,
1152                                *rows - cpu_rows,
1153                                *cols,
1154                                out_gpu,
1155                            )
1156                        });
1157                        if cpu_rows > 0 {
1158                            // Repack prefix covers the full groups of the
1159                            // CPU half (the split starts at row 0).
1160                            let rep_cpu = if repack.is_empty() {
1161                                &[][..]
1162                            } else {
1163                                &repack[..(cpu_rows / 4) * 4 * *cols]
1164                            };
1165                            qmatvec(
1166                                &bytes[..cpu_rows * *cols],
1167                                rep_cpu,
1168                                &row_scale[..cpu_rows],
1169                                x,
1170                                col_field,
1171                                *dtype,
1172                                cpu_rows,
1173                                *cols,
1174                                out_cpu,
1175                                pool,
1176                            );
1177                        }
1178                        g.join().unwrap_or(false)
1179                    });
1180                    if ok {
1181                        crate::gpu::probe_record(crate::gpu::OpClass::Matvec, true, t0.elapsed());
1182                        return;
1183                    }
1184                    // GPU failed — CPU finishes its half (rows rebased —
1185                    // group offsets don't line up, mmap layout only).
1186                    qmatvec(
1187                        &bytes[cpu_rows * *cols..(*rows) * *cols],
1188                        &[],
1189                        &row_scale[cpu_rows..],
1190                        x,
1191                        col_field,
1192                        *dtype,
1193                        *rows - cpu_rows,
1194                        *cols,
1195                        out_gpu,
1196                        pool,
1197                    );
1198                    return;
1199                }
1200                qmatvec(
1201                    self.quant_bytes(),
1202                    repack,
1203                    row_scale,
1204                    x,
1205                    col_field,
1206                    *dtype,
1207                    *rows,
1208                    *cols,
1209                    out,
1210                    pool,
1211                );
1212            }
1213        }
1214    }
1215
1216    /// Fused two-input matvec (MTP verify pair): weights streamed once.
1217    pub fn matvec2(
1218        &self,
1219        x1: &[f32],
1220        x2: &[f32],
1221        o1: &mut [f32],
1222        o2: &mut [f32],
1223        pool: Option<&Pool>,
1224    ) {
1225        match self {
1226            Self::F32 { data, .. } => matvec_rows2(pool, data, x1, x2, o1, o2),
1227            Self::Mapped {
1228                dtype,
1229                rows,
1230                cols,
1231                row_scale,
1232                col_field,
1233                vbit_offsets,
1234                ..
1235            } => {
1236                if *dtype == TensorDtype::Q4Block {
1237                    q4matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1238                    return;
1239                }
1240                if *dtype == TensorDtype::Q4Tiled {
1241                    q4t_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1242                    return;
1243                }
1244                if *dtype == TensorDtype::Q4TiledP {
1245                    q4tp_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1246                    return;
1247                }
1248                if *dtype == TensorDtype::Q2TiledP {
1249                    q2tp_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1250                    return;
1251                }
1252                if *dtype == TensorDtype::Q1 {
1253                    q1_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1254                    return;
1255                }
1256                if *dtype == TensorDtype::Q1T {
1257                    // Fused ternary pair: one row pass, the register
1258                    // unpack shared across both streams on ARM. (Q1T
1259                    // lacks a row_scale array — scales live inline in
1260                    // the tiles — so it must not fall through to the
1261                    // q8 qmatvec2 below.)
1262                    q1t_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1263                    return;
1264                }
1265                if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1266                    vbitmatvec2(
1267                        self.quant_bytes(),
1268                        vbit_offsets,
1269                        x1,
1270                        x2,
1271                        *rows,
1272                        *cols,
1273                        o1,
1274                        o2,
1275                        pool,
1276                    );
1277                    return;
1278                }
1279                qmatvec2(
1280                    self.quant_bytes(),
1281                    row_scale,
1282                    x1,
1283                    x2,
1284                    col_field,
1285                    *dtype,
1286                    *rows,
1287                    *cols,
1288                    o1,
1289                    o2,
1290                    pool,
1291                );
1292            }
1293        }
1294    }
1295}
1296
1297impl QTensor {
1298    /// Batched matvec (prefill-GEMM): xs — row-major [b, cols],
1299    /// out — row-major [b, rows]. Element-wise semantics are IDENTICAL
1300    /// to b matvec calls (same dot kernels in the same order); the win —
1301    /// the weight row streams from DRAM once per batch, not b times.
1302    pub fn matmat(&self, xs_all: &[f32], b: usize, out: &mut [f32], pool: Option<&Pool>) {
1303        let cols = self.cols();
1304        let rows = self.rows();
1305        debug_assert_eq!(xs_all.len(), b * cols);
1306        debug_assert_eq!(out.len(), b * rows);
1307        // GPTQ calibration: fold this layer's inputs into its Hessian. Only
1308        // Mapped tensors carry a directory name; the check is a relaxed
1309        // atomic load, free when not calibrating.
1310        if crate::gptq_capture::capturing() {
1311            if let Self::Mapped { model, idx, .. } = self {
1312                crate::gptq_capture::accumulate(&model.tensors[*idx].name, xs_all, b, cols);
1313            }
1314        }
1315        match self {
1316            Self::F32 { data, .. } => {
1317                let out_addr = SendMut(out.as_mut_ptr());
1318                let run = |start: usize, end: usize| {
1319                    for o in start..end {
1320                        let row = &data[o * cols..(o + 1) * cols];
1321                        for bi in 0..b {
1322                            let x = &xs_all[bi * cols..(bi + 1) * cols];
1323                            let mut acc = 0f32;
1324                            for j in 0..cols {
1325                                acc += row[j] * x[j];
1326                            }
1327                            unsafe { *out_addr.at(bi * rows + o) = acc };
1328                        }
1329                    }
1330                };
1331                dispatch_rows(pool, rows, &run);
1332            }
1333            Self::Mapped {
1334                dtype,
1335                row_scale,
1336                col_field,
1337                vbit_offsets,
1338                ..
1339            } => {
1340                if *dtype == TensorDtype::Q4Block {
1341                    q4matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1342                    return;
1343                }
1344                if *dtype == TensorDtype::Q4TiledP {
1345                    // GPU batched q4tp GEMM (dequant + f32nt mul_mm on the
1346                    // device); the probe keeps whichever beats the CPU arm.
1347                    // Narrow (prompt-encode) and wide (DiT) batches probe
1348                    // as separate classes — the regimes have opposite
1349                    // winners and one shared verdict locked the wrong arm.
1350                    // Kill switch (gpu::mm_kill): one grossly slow GPU op
1351                    // (a fair-condition op is ≤~100 ms even at 1024px)
1352                    // means the device is contended by another process
1353                    // (e.g. a simulator) — verdicts are per-process, so
1354                    // without the bail the whole render crawls behind
1355                    // someone else's queue.
1356                    if b >= 32
1357                        && b * rows * cols >= 128_000_000
1358                        && cols % 32 == 0
1359                        && !crate::gpu::mm_killed()
1360                        && crate::gpu::enabled_here()
1361                    {
1362                        let class = if b >= 128 {
1363                            crate::gpu::OpClass::MatmatWide
1364                        } else {
1365                            crate::gpu::OpClass::Matmat
1366                        };
1367                        if let Self::Mapped { model, idx, .. } = self {
1368                            let t0 = std::time::Instant::now();
1369                            // A cold call takes the device arm: its sample
1370                            // is discarded either way, and the upload is
1371                            // what the next step needs.
1372                            let resident = crate::gpu::weight_is_resident(model, *idx);
1373                            match crate::gpu::probe_arm_cold_prefers_gpu(class, resident) {
1374                                crate::gpu::ProbeArm::Gpu => {
1375                                    if crate::gpu::q4tp_matmat(
1376                                        model, *idx, xs_all, b, rows, cols, out,
1377                                    ) {
1378                                        let el = t0.elapsed();
1379                                        // Work-proportional budget: ~8× the
1380                                        // fair-device estimate (+20 ms slack).
1381                                        // An absolute cap missed the worst
1382                                        // case — contended ops sit at
1383                                        // 100–240 ms each and still bury a
1384                                        // render whose fair op is 3–9 ms.
1385                                        // Cold ops (first PSO build, buffer
1386                                        // alloc) are exempt: a one-off
1387                                        // ~50 ms compile is not contention.
1388                                        let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
1389                                        let budget = std::time::Duration::from_secs_f64(
1390                                            flops / 1.5e12 * 8.0 + 0.020,
1391                                        );
1392                                        if el > budget && !crate::gpu::probe_was_cold() {
1393                                            tracing::warn!(
1394                                                "gpu q4tp matmat took {el:?} (budget {budget:?}) — \
1395                                                 device contended, CPU for the rest of the process"
1396                                            );
1397                                            crate::gpu::mm_kill();
1398                                        }
1399                                        crate::gpu::probe_record(class, true, el);
1400                                        return;
1401                                    }
1402                                }
1403                                crate::gpu::ProbeArm::CpuTimed => {
1404                                    q4tp_matmat(
1405                                        self.quant_bytes(),
1406                                        xs_all,
1407                                        b,
1408                                        rows,
1409                                        cols,
1410                                        out,
1411                                        pool,
1412                                    );
1413                                    crate::gpu::probe_record(class, false, t0.elapsed());
1414                                    return;
1415                                }
1416                                crate::gpu::ProbeArm::Cpu => {}
1417                            }
1418                        }
1419                    }
1420                    q4tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1421                    return;
1422                }
1423                if *dtype == TensorDtype::Q2TiledP {
1424                    // Same device arm as q4tp, behind the same probe:
1425                    // the planes differ, the dispatch does not. Without
1426                    // this a q2tp file ran its widest projections on the
1427                    // host while the 4-bit one had the card, which is a
1428                    // codec paying for its size twice.
1429                    if b >= 32
1430                        && b * rows * cols >= 128_000_000
1431                        && cols % 32 == 0
1432                        && !crate::gpu::mm_killed()
1433                        && crate::gpu::enabled_here()
1434                    {
1435                        let class = if b >= 128 {
1436                            crate::gpu::OpClass::MatmatWide
1437                        } else {
1438                            crate::gpu::OpClass::Matmat
1439                        };
1440                        if let Self::Mapped { model, idx, .. } = self {
1441                            let t0 = std::time::Instant::now();
1442                            match crate::gpu::probe_arm(class) {
1443                                crate::gpu::ProbeArm::Gpu => {
1444                                    if crate::gpu::q2tp_matmat(
1445                                        model, *idx, xs_all, b, rows, cols, out,
1446                                    ) {
1447                                        crate::gpu::probe_record(class, true, t0.elapsed());
1448                                        return;
1449                                    }
1450                                }
1451                                crate::gpu::ProbeArm::CpuTimed => {
1452                                    q2tp_matmat(
1453                                        self.quant_bytes(),
1454                                        xs_all,
1455                                        b,
1456                                        rows,
1457                                        cols,
1458                                        out,
1459                                        pool,
1460                                    );
1461                                    crate::gpu::probe_record(class, false, t0.elapsed());
1462                                    return;
1463                                }
1464                                crate::gpu::ProbeArm::Cpu => {}
1465                            }
1466                        }
1467                    }
1468                    // Without a host arm a q2tp tensor falls through to
1469                    // the q8 fallback, which reads it at one BYTE per
1470                    // weight — a 2x overrun that killed pool workers
1471                    // mid-prefill while the dispatcher waited forever.
1472                    q2tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1473                    return;
1474                }
1475                if *dtype == TensorDtype::Q4Tiled {
1476                    // GPU batched q4t GEMM (dequant + f32nt mul_mm on the
1477                    // device); the probe keeps whichever beats the CPU arm.
1478                    // Narrow (prompt-encode) and wide (DiT) batches probe
1479                    // as separate classes — the regimes have opposite
1480                    // winners and one shared verdict locked the wrong arm.
1481                    // Kill switch (gpu::mm_kill): one grossly slow GPU op
1482                    // (a fair-condition op is ≤~100 ms even at 1024px)
1483                    // means the device is contended by another process
1484                    // (e.g. a simulator) — verdicts are per-process, so
1485                    // without the bail the whole render crawls behind
1486                    // someone else's queue.
1487                    if b >= 32
1488                        && b * rows * cols >= 128_000_000
1489                        && cols % 32 == 0
1490                        && !crate::gpu::mm_killed()
1491                        && crate::gpu::enabled_here()
1492                    {
1493                        let class = if b >= 128 {
1494                            crate::gpu::OpClass::MatmatWide
1495                        } else {
1496                            crate::gpu::OpClass::Matmat
1497                        };
1498                        if let Self::Mapped { model, idx, .. } = self {
1499                            let t0 = std::time::Instant::now();
1500                            match crate::gpu::probe_arm(class) {
1501                                crate::gpu::ProbeArm::Gpu => {
1502                                    if crate::gpu::q4t_matmat(
1503                                        model, *idx, xs_all, b, rows, cols, out,
1504                                    ) {
1505                                        let el = t0.elapsed();
1506                                        // Work-proportional budget: ~8× the
1507                                        // fair-device estimate (+20 ms slack).
1508                                        // An absolute cap missed the worst
1509                                        // case — contended ops sit at
1510                                        // 100–240 ms each and still bury a
1511                                        // render whose fair op is 3–9 ms.
1512                                        // Cold ops (first PSO build, buffer
1513                                        // alloc) are exempt: a one-off
1514                                        // ~50 ms compile is not contention.
1515                                        let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
1516                                        let budget = std::time::Duration::from_secs_f64(
1517                                            flops / 1.5e12 * 8.0 + 0.020,
1518                                        );
1519                                        if el > budget && !crate::gpu::probe_was_cold() {
1520                                            tracing::warn!(
1521                                                "gpu q4t matmat took {el:?} (budget {budget:?}) — \
1522                                                 device contended, CPU for the rest of the process"
1523                                            );
1524                                            crate::gpu::mm_kill();
1525                                        }
1526                                        crate::gpu::probe_record(class, true, el);
1527                                        return;
1528                                    }
1529                                }
1530                                crate::gpu::ProbeArm::CpuTimed => {
1531                                    q4t_matmat(
1532                                        self.quant_bytes(),
1533                                        xs_all,
1534                                        b,
1535                                        rows,
1536                                        cols,
1537                                        out,
1538                                        pool,
1539                                    );
1540                                    crate::gpu::probe_record(class, false, t0.elapsed());
1541                                    return;
1542                                }
1543                                crate::gpu::ProbeArm::Cpu => {}
1544                            }
1545                        }
1546                    }
1547                    q4t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1548                    return;
1549                }
1550                if *dtype == TensorDtype::Q1 {
1551                    // GPU batched q1 GEMM for wide prefill (q1_mul_mm on the
1552                    // device); the probe keeps whichever beats the CPU matmat.
1553                    if b >= 32
1554                        && b * rows * cols >= 128_000_000
1555                        && cols % 64 == 0
1556                        && crate::gpu::enabled_here()
1557                    {
1558                        if let Self::Mapped { model, idx, .. } = self {
1559                            let t0 = std::time::Instant::now();
1560                            match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
1561                                crate::gpu::ProbeArm::Gpu => {
1562                                    if crate::gpu::q1_matmat(
1563                                        model, *idx, xs_all, b, rows, cols, out,
1564                                    ) {
1565                                        crate::gpu::probe_record(
1566                                            crate::gpu::OpClass::Matmat,
1567                                            true,
1568                                            t0.elapsed(),
1569                                        );
1570                                        return;
1571                                    }
1572                                }
1573                                crate::gpu::ProbeArm::CpuTimed => {
1574                                    q1_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1575                                    crate::gpu::probe_record(
1576                                        crate::gpu::OpClass::Matmat,
1577                                        false,
1578                                        t0.elapsed(),
1579                                    );
1580                                    return;
1581                                }
1582                                crate::gpu::ProbeArm::Cpu => {}
1583                            }
1584                        }
1585                    }
1586                    q1_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1587                    return;
1588                }
1589                if *dtype == TensorDtype::Q1T {
1590                    // GPU batched GEMM for wide prefill (base + overlay on the
1591                    // device); probe keeps the winner vs the CPU matmat.
1592                    if b >= 32 && b * rows * cols >= 128_000_000 && crate::gpu::enabled_here() {
1593                        if let Self::Mapped { model, idx, .. } = self {
1594                            let t0 = std::time::Instant::now();
1595                            match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
1596                                crate::gpu::ProbeArm::Gpu => {
1597                                    if crate::gpu::q1t_matmat(
1598                                        model, *idx, xs_all, b, rows, cols, out,
1599                                    ) {
1600                                        crate::gpu::probe_record(
1601                                            crate::gpu::OpClass::Matmat,
1602                                            true,
1603                                            t0.elapsed(),
1604                                        );
1605                                        return;
1606                                    }
1607                                }
1608                                crate::gpu::ProbeArm::CpuTimed => {
1609                                    q1t_matmat(
1610                                        self.quant_bytes(),
1611                                        xs_all,
1612                                        b,
1613                                        rows,
1614                                        cols,
1615                                        out,
1616                                        pool,
1617                                    );
1618                                    crate::gpu::probe_record(
1619                                        crate::gpu::OpClass::Matmat,
1620                                        false,
1621                                        t0.elapsed(),
1622                                    );
1623                                    return;
1624                                }
1625                                crate::gpu::ProbeArm::Cpu => {}
1626                            }
1627                        }
1628                    }
1629                    q1t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1630                    return;
1631                }
1632                if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1633                    vbitmatmat(
1634                        self.quant_bytes(),
1635                        vbit_offsets,
1636                        xs_all,
1637                        b,
1638                        rows,
1639                        cols,
1640                        out,
1641                        pool,
1642                    );
1643                    return;
1644                }
1645                let pre: Vec<std::borrow::Cow<'_, [f32]>> = (0..b)
1646                    .map(|bi| prescale(&xs_all[bi * cols..(bi + 1) * cols], col_field, *dtype))
1647                    .collect();
1648                // D5: large prefill-batch GEMMs — on the GPU (threshold by
1649                // work volume: submission carries b×rows×cols MACs).
1650                // Runtime probe: the naive GEMM shader + sync readback
1651                // lose to the CPU GEMM on slow driver stacks — alternate
1652                // both arms and keep the winner.
1653                if b >= 8 && b * rows * cols >= 128_000_000 && crate::gpu::enabled_here() {
1654                    if let Self::Mapped { model, idx, .. } = self {
1655                        let t0 = std::time::Instant::now();
1656                        match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
1657                            crate::gpu::ProbeArm::Gpu
1658                                if crate::gpu::probe_deciding(crate::gpu::OpClass::Matmat)
1659                                    && !crate::gpu::q8_resident_or_upload(model, *idx) =>
1660                            {
1661                                // Cold weights during probing: the upload
1662                                // has started, the count runs on the CPU —
1663                                // the GPU arm samples on the next touch.
1664                                let q = self.quant_bytes();
1665                                qmatmat(q, row_scale, &pre, rows, cols, out, pool);
1666                                return;
1667                            }
1668                            crate::gpu::ProbeArm::Gpu => {
1669                                let flat: Vec<f32> =
1670                                    pre.iter().flat_map(|v| v.iter().copied()).collect();
1671                                if crate::gpu::q8_matmat(
1672                                    model, *idx, row_scale, &flat, b, rows, cols, out,
1673                                ) {
1674                                    crate::gpu::probe_record(
1675                                        crate::gpu::OpClass::Matmat,
1676                                        true,
1677                                        t0.elapsed(),
1678                                    );
1679                                    return;
1680                                }
1681                            }
1682                            crate::gpu::ProbeArm::CpuTimed => {
1683                                let q = self.quant_bytes();
1684                                qmatmat(q, row_scale, &pre, rows, cols, out, pool);
1685                                crate::gpu::probe_record(
1686                                    crate::gpu::OpClass::Matmat,
1687                                    false,
1688                                    t0.elapsed(),
1689                                );
1690                                return;
1691                            }
1692                            crate::gpu::ProbeArm::Cpu => {}
1693                        }
1694                    }
1695                }
1696                let q = self.quant_bytes();
1697                qmatmat(q, row_scale, &pre, rows, cols, out, pool);
1698            }
1699        }
1700    }
1701}
1702
1703impl QTensor {
1704    /// Multi-matrix job (roadmap §3 P0): N tensors sharing one input
1705    /// run under a SINGLE pool dispatch — QKV or gate+up cost one
1706    /// barrier instead of N. Per-row math is the exact same kernel as
1707    /// `matvec` (bit-identical outputs); only the dispatch is fused.
1708    /// Falls back to N sequential matvecs when the set is not a uniform
1709    /// q8-family/F32 group or there is no pool.
1710    pub fn matvec_many<const N: usize>(
1711        ts: [&QTensor; N],
1712        x: &[f32],
1713        mut outs: [&mut [f32]; N],
1714        pool: Option<&Pool>,
1715    ) {
1716        let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
1717        let uniform_q8 = ts.iter().all(|t| {
1718            matches!(
1719                t,
1720                Self::Mapped {
1721                    dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f,
1722                    ..
1723                }
1724            )
1725        });
1726        let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
1727        let uniform_q4 = ts.iter().all(|t| {
1728            matches!(
1729                t,
1730                Self::Mapped {
1731                    dtype: TensorDtype::Q4Block,
1732                    ..
1733                }
1734            )
1735        });
1736        let uniform_vbit = ts.iter().all(|t| {
1737            matches!(
1738                t,
1739                Self::Mapped {
1740                    dtype: TensorDtype::Vbit | TensorDtype::VbitRo,
1741                    ..
1742                }
1743            )
1744        });
1745        let uniform_q1 = ts.iter().all(|t| {
1746            matches!(
1747                t,
1748                Self::Mapped {
1749                    dtype: TensorDtype::Q1,
1750                    ..
1751                }
1752            )
1753        });
1754        let uniform_q1t = ts.iter().all(|t| {
1755            matches!(
1756                t,
1757                Self::Mapped {
1758                    dtype: TensorDtype::Q1T,
1759                    ..
1760                }
1761            )
1762        });
1763        // q4tp is the skeleton dtype of the big MoE files, and without an arm
1764        // here every projection that shares an input paid its own pool
1765        // barrier: DeepSeek-V4's attention step alone hands this function
1766        // wq_a, wkv and both compressors' pairs off the same hidden state.
1767        let uniform_q4tp = ts.iter().all(|t| {
1768            matches!(
1769                t,
1770                Self::Mapped {
1771                    dtype: TensorDtype::Q4TiledP,
1772                    ..
1773                }
1774            )
1775        }) && ts.iter().all(|t| t.cols() == ts[0].cols() && t.cols() % GROUP_SIZE == 0);
1776        let Some(pool) = pool else {
1777            for (t, o) in ts.iter().zip(outs.iter_mut()) {
1778                t.matvec(x, o, None);
1779            }
1780            return;
1781        };
1782        if total_rows < 256
1783            || !(uniform_q8
1784                || uniform_f32
1785                || uniform_q4
1786                || uniform_vbit
1787                || uniform_q1
1788                || uniform_q1t
1789                || uniform_q4tp)
1790        {
1791            for (t, o) in ts.iter().zip(outs.iter_mut()) {
1792                t.matvec(x, o, Some(pool));
1793            }
1794            return;
1795        }
1796
1797        if uniform_q4tp {
1798            // Every tensor's rows laid end to end in one virtual row space,
1799            // so the whole set is ONE dispatch. The per-row body is the
1800            // `q4tp_matvec` arm verbatim — same activation split, same
1801            // accumulation order — so the outputs are bit-identical to the
1802            // sequential calls this replaces.
1803            let cols = ts[0].cols();
1804            let gpr = cols / GROUP_SIZE;
1805            let views: Vec<Q4tpView> = ts
1806                .iter()
1807                .map(|t| Q4tpView::new(t.quant_bytes(), t.rows(), cols))
1808                .collect();
1809            let rows_of: Vec<usize> = ts.iter().map(|t| t.rows()).collect();
1810            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1811            // flat index -> (which tensor, which of its rows)
1812            let locate = |flat: usize| -> (usize, usize) {
1813                let mut acc = 0;
1814                for (i, &r) in rows_of.iter().enumerate() {
1815                    if flat < acc + r {
1816                        return (i, flat - acc);
1817                    }
1818                    acc += r;
1819                }
1820                (rows_of.len() - 1, 0)
1821            };
1822            let (views, outs_addr) = (&views, &outs_addr);
1823            if a8w8_enabled() {
1824                let act = split_act(x);
1825                let act = &act;
1826                let run = |start: usize, end: usize| {
1827                    let mut sc = vec![0f32; gpr];
1828                    for flat in start..end {
1829                        let (t, r) = locate(flat);
1830                        let v = &views[t];
1831                        v.scales_into(r, gpr, &mut sc);
1832                        let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, &sc) * act.sx;
1833                        for &(j, xv) in &act.outliers {
1834                            let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
1835                            acc += w * s * xv;
1836                        }
1837                        // SAFETY: one worker owns each (tensor, row) pair.
1838                        unsafe { *outs_addr[t].at(r) = acc };
1839                    }
1840                };
1841                pool.run_rows(total_rows, &run);
1842            } else {
1843                let run = |start: usize, end: usize| {
1844                    let mut sc = vec![0f32; gpr];
1845                    for flat in start..end {
1846                        let (t, r) = locate(flat);
1847                        let v = &views[t];
1848                        v.scales_into(r, gpr, &mut sc);
1849                        // SAFETY: one worker owns each (tensor, row) pair.
1850                        unsafe { *outs_addr[t].at(r) = q4tp_row_exact(v.nib, r, gpr, x, &sc) };
1851                    }
1852                };
1853                pool.run_rows(total_rows, &run);
1854            }
1855            return;
1856        }
1857
1858        if uniform_q1 {
1859            // One shared activation split + group sums (q1 has no col
1860            // field; the same input feeds every tensor).
1861            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1862            if a8w8_enabled() {
1863                let act = split_act(x);
1864                let gsum = q1_group_sums(&act.xq, ts[0].cols() / GROUP_SIZE);
1865                let (act, gsum) = (&act, &gsum);
1866                let closures: [_; N] = std::array::from_fn(|i| {
1867                    let (bytes, gpr, out) =
1868                        (ts[i].quant_bytes(), ts[i].cols() / GROUP_SIZE, outs_addr[i]);
1869                    move |s: usize, e: usize| q1_range_a8w8(bytes, gpr, act, gsum, out, s, e)
1870                });
1871                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1872                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1873                pool.run_many(&parts);
1874            } else {
1875                let closures: [_; N] = std::array::from_fn(|i| {
1876                    let (bytes, gpr, out) =
1877                        (ts[i].quant_bytes(), ts[i].cols() / GROUP_SIZE, outs_addr[i]);
1878                    move |s: usize, e: usize| q1_range_f32(bytes, gpr, x, out, s, e)
1879                });
1880                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1881                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1882                pool.run_many(&parts);
1883            }
1884            return;
1885        }
1886
1887        if uniform_q1t {
1888            // Q1T batched: one shared activation split + overlay decode,
1889            // all tensors' rows in ONE pool dispatch (saves N−1 dispatches
1890            // and N−1 redundant split_act calls per layer).
1891            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1892            const TILE: usize = cortiq_core::quant::Q1T_TILE;
1893            if a8w8_enabled() {
1894                let act = split_act(x);
1895                let act = &act;
1896                let x_ref = x;
1897                let closures: [_; N] = std::array::from_fn(|i| {
1898                    let bytes = ts[i].quant_bytes();
1899                    let (rows, cols) = (ts[i].rows(), ts[i].cols());
1900                    let gpr = cols / GROUP_SIZE;
1901                    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
1902                    let out = outs_addr[i];
1903                    move |s: usize, e: usize| {
1904                        q1t_range_a8w8(bytes, gpr, rp_off, ent_off, has_ov, act, x_ref, out, s, e)
1905                    }
1906                });
1907                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1908                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1909                pool.run_many(&parts);
1910            } else {
1911                let x_ref = x;
1912                let closures: [_; N] = std::array::from_fn(|i| {
1913                    let bytes = ts[i].quant_bytes();
1914                    let (rows, cols) = (ts[i].rows(), ts[i].cols());
1915                    let gpr = cols / GROUP_SIZE;
1916                    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
1917                    let out = outs_addr[i];
1918                    move |s: usize, e: usize| {
1919                        q1t_range_f32_batch(bytes, gpr, rp_off, ent_off, has_ov, x_ref, out, s, e)
1920                    }
1921                });
1922                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1923                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1924                pool.run_many(&parts);
1925            }
1926            return;
1927        }
1928
1929        if uniform_q4 || uniform_vbit {
1930            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1931            // q4/vbit share one activation split — no per-tensor col field.
1932            if a8w8_enabled() {
1933                let act = split_act(x);
1934                let act = &act;
1935                if uniform_q4 {
1936                    let closures: [_; N] = std::array::from_fn(|i| {
1937                        let (packed, scales) =
1938                            q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
1939                        let (gpr, cols, out) =
1940                            (ts[i].cols() / GROUP_SIZE, ts[i].cols(), outs_addr[i]);
1941                        move |s: usize, e: usize| {
1942                            q4_range_a8w8(packed, scales, gpr, cols, act, out, s, e)
1943                        }
1944                    });
1945                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1946                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1947                    pool.run_many(&parts);
1948                } else {
1949                    let closures: [_; N] = std::array::from_fn(|i| {
1950                        let Self::Mapped { vbit_offsets, .. } = ts[i] else {
1951                            unreachable!()
1952                        };
1953                        let (bytes, rows, cols, out) = (
1954                            ts[i].quant_bytes(),
1955                            ts[i].rows(),
1956                            ts[i].cols(),
1957                            outs_addr[i],
1958                        );
1959                        move |s: usize, e: usize| {
1960                            vbit_range_a8w8(bytes, vbit_offsets, x, act, rows, cols, out, s, e)
1961                        }
1962                    });
1963                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1964                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1965                    pool.run_many(&parts);
1966                }
1967                return;
1968            }
1969            if uniform_q4 {
1970                let closures: [_; N] = std::array::from_fn(|i| {
1971                    let (packed, scales) =
1972                        q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
1973                    let (gpr, out) = (ts[i].cols() / GROUP_SIZE, outs_addr[i]);
1974                    move |s: usize, e: usize| q4_range_f32(packed, scales, gpr, x, out, s, e)
1975                });
1976                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1977                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1978                pool.run_many(&parts);
1979            } else {
1980                let closures: [_; N] = std::array::from_fn(|i| {
1981                    let Self::Mapped { vbit_offsets, .. } = ts[i] else {
1982                        unreachable!()
1983                    };
1984                    let (bytes, rows, cols, out) = (
1985                        ts[i].quant_bytes(),
1986                        ts[i].rows(),
1987                        ts[i].cols(),
1988                        outs_addr[i],
1989                    );
1990                    move |s: usize, e: usize| {
1991                        vbit_range_f32(bytes, vbit_offsets, x, rows, cols, out, s, e)
1992                    }
1993                });
1994                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1995                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1996                pool.run_many(&parts);
1997            }
1998            return;
1999        }
2000
2001        if uniform_f32 {
2002            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2003            let closures: [_; N] = std::array::from_fn(|i| {
2004                let Self::F32 { data, cols, .. } = ts[i] else {
2005                    unreachable!()
2006                };
2007                let out = outs_addr[i];
2008                move |start: usize, end: usize| {
2009                    for o in start..end {
2010                        let row = &data[o * cols..(o + 1) * cols];
2011                        let mut sum = 0.0f32;
2012                        for j in 0..*cols {
2013                            sum += row[j] * x[j];
2014                        }
2015                        // SAFETY: disjoint (tensor, row) cells per worker.
2016                        unsafe { *out.at(o) = sum };
2017                    }
2018                }
2019            });
2020            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2021                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2022            pool.run_many(&parts);
2023            return;
2024        }
2025
2026        // Uniform q8-family: per-tensor prescale (q8_2f col fields
2027        // differ per tensor) + the shared range kernels.
2028        struct Ctx<'a> {
2029            bytes: &'a [u8],
2030            #[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
2031            rep: &'a [u8],
2032            row_scale: &'a [f32],
2033            cols: usize,
2034            xs: std::borrow::Cow<'a, [f32]>,
2035        }
2036        let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
2037            let Self::Mapped {
2038                dtype,
2039                cols,
2040                row_scale,
2041                col_field,
2042                repack,
2043                ..
2044            } = ts[i]
2045            else {
2046                unreachable!()
2047            };
2048            Ctx {
2049                bytes: ts[i].quant_bytes(),
2050                rep: repack,
2051                row_scale,
2052                cols: *cols,
2053                xs: prescale(x, col_field, *dtype),
2054            }
2055        });
2056        let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2057        #[cfg(target_arch = "aarch64")]
2058        if sdot_enabled() {
2059            let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
2060            let closures: [_; N] = std::array::from_fn(|i| {
2061                let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
2062                move |start: usize, end: usize| {
2063                    q8_range_sdot(c.bytes, c.rep, c.row_scale, act, c.cols, out, start, end)
2064                }
2065            });
2066            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2067                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2068            pool.run_many(&parts);
2069            return;
2070        }
2071        #[cfg(target_arch = "x86_64")]
2072        if avx2_a8w8_enabled() {
2073            let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
2074            let closures: [_; N] = std::array::from_fn(|i| {
2075                let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
2076                move |start: usize, end: usize| {
2077                    q8_range_avx2(c.bytes, c.row_scale, act, c.cols, out, start, end)
2078                }
2079            });
2080            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2081                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2082            pool.run_many(&parts);
2083            return;
2084        }
2085        let closures: [_; N] = std::array::from_fn(|i| {
2086            let (c, out) = (&ctxs[i], outs_addr[i]);
2087            move |start: usize, end: usize| {
2088                q8_range_f32(c.bytes, c.row_scale, &c.xs, c.cols, out, start, end)
2089            }
2090        });
2091        let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2092            std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2093        pool.run_many(&parts);
2094    }
2095}
2096
2097impl QTensor {
2098    /// Pair-input multi-matrix job: N tensors × 2 shared inputs under a
2099    /// single pool dispatch — the MTP/pair decode path publishes one job
2100    /// for Q/K/V (and one for gate+up) instead of one per tensor.
2101    /// Per-row math is exactly `matvec2`'s kernels; bit-identical.
2102    #[allow(clippy::needless_range_loop)]
2103    pub fn matvec2_many<const N: usize>(
2104        ts: [&QTensor; N],
2105        x1: &[f32],
2106        x2: &[f32],
2107        mut o1s: [&mut [f32]; N],
2108        mut o2s: [&mut [f32]; N],
2109        pool: Option<&Pool>,
2110    ) {
2111        let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
2112        let uniform_q8 = ts.iter().all(|t| {
2113            matches!(
2114                t,
2115                Self::Mapped {
2116                    dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f,
2117                    ..
2118                }
2119            )
2120        });
2121        let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
2122        let uniform_q4 = ts.iter().all(|t| {
2123            matches!(
2124                t,
2125                Self::Mapped {
2126                    dtype: TensorDtype::Q4Block,
2127                    ..
2128                }
2129            )
2130        });
2131        let uniform_vbit = ts.iter().all(|t| {
2132            matches!(
2133                t,
2134                Self::Mapped {
2135                    dtype: TensorDtype::Vbit | TensorDtype::VbitRo,
2136                    ..
2137                }
2138            )
2139        });
2140        let fusable = pool.is_some()
2141            && total_rows >= 256
2142            && (uniform_q8 || uniform_f32 || uniform_q4 || uniform_vbit);
2143        if !fusable {
2144            for i in 0..N {
2145                ts[i].matvec2(x1, x2, o1s[i], o2s[i], pool);
2146            }
2147            return;
2148        }
2149        let pool = pool.unwrap();
2150
2151        if uniform_q4 || uniform_vbit {
2152            let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
2153            let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
2154            // q4/vbit share activation splits — no per-tensor col field.
2155            if a8w8_enabled() {
2156                let a1 = split_act(x1);
2157                let a2 = split_act(x2);
2158                let (a1, a2) = (&a1, &a2);
2159                if uniform_q4 {
2160                    let closures: [_; N] = std::array::from_fn(|i| {
2161                        let (packed, scales) =
2162                            q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
2163                        let (gpr, cols, o1, o2) =
2164                            (ts[i].cols() / GROUP_SIZE, ts[i].cols(), p1[i], p2[i]);
2165                        move |s: usize, e: usize| {
2166                            q4_range2_a8w8(packed, scales, gpr, cols, a1, a2, o1, o2, s, e)
2167                        }
2168                    });
2169                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2170                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2171                    pool.run_many(&parts);
2172                } else {
2173                    let closures: [_; N] = std::array::from_fn(|i| {
2174                        let Self::Mapped { vbit_offsets, .. } = ts[i] else {
2175                            unreachable!()
2176                        };
2177                        let (bytes, rows, cols, o1, o2) = (
2178                            ts[i].quant_bytes(),
2179                            ts[i].rows(),
2180                            ts[i].cols(),
2181                            p1[i],
2182                            p2[i],
2183                        );
2184                        move |s: usize, e: usize| {
2185                            vbit_range2_a8w8(
2186                                bytes,
2187                                vbit_offsets,
2188                                x1,
2189                                x2,
2190                                a1,
2191                                a2,
2192                                rows,
2193                                cols,
2194                                o1,
2195                                o2,
2196                                s,
2197                                e,
2198                            )
2199                        }
2200                    });
2201                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2202                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2203                    pool.run_many(&parts);
2204                }
2205                return;
2206            }
2207            if uniform_q4 {
2208                let closures: [_; N] = std::array::from_fn(|i| {
2209                    let (packed, scales) =
2210                        q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
2211                    let (gpr, o1, o2) = (ts[i].cols() / GROUP_SIZE, p1[i], p2[i]);
2212                    move |s: usize, e: usize| {
2213                        q4_range2_f32(packed, scales, gpr, x1, x2, o1, o2, s, e)
2214                    }
2215                });
2216                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2217                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2218                pool.run_many(&parts);
2219            } else {
2220                let closures: [_; N] = std::array::from_fn(|i| {
2221                    let Self::Mapped { vbit_offsets, .. } = ts[i] else {
2222                        unreachable!()
2223                    };
2224                    let (bytes, rows, cols, o1, o2) = (
2225                        ts[i].quant_bytes(),
2226                        ts[i].rows(),
2227                        ts[i].cols(),
2228                        p1[i],
2229                        p2[i],
2230                    );
2231                    move |s: usize, e: usize| {
2232                        vbit_range2_f32(bytes, vbit_offsets, x1, x2, rows, cols, o1, o2, s, e)
2233                    }
2234                });
2235                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2236                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2237                pool.run_many(&parts);
2238            }
2239            return;
2240        }
2241
2242        if uniform_f32 {
2243            let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
2244            let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
2245            let closures: [_; N] = std::array::from_fn(|i| {
2246                let Self::F32 { data, cols, .. } = ts[i] else {
2247                    unreachable!()
2248                };
2249                let (o1, o2) = (p1[i], p2[i]);
2250                move |start: usize, end: usize| {
2251                    for o in start..end {
2252                        let row = &data[o * cols..(o + 1) * cols];
2253                        let (mut s1, mut s2) = (0.0f32, 0.0f32);
2254                        for j in 0..*cols {
2255                            s1 += row[j] * x1[j];
2256                            s2 += row[j] * x2[j];
2257                        }
2258                        // SAFETY: disjoint (tensor, row) cells per worker.
2259                        unsafe {
2260                            *o1.at(o) = s1;
2261                            *o2.at(o) = s2;
2262                        }
2263                    }
2264                }
2265            });
2266            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2267                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2268            pool.run_many(&parts);
2269            return;
2270        }
2271
2272        struct Ctx<'a> {
2273            bytes: &'a [u8],
2274            row_scale: &'a [f32],
2275            cols: usize,
2276            xs1: std::borrow::Cow<'a, [f32]>,
2277            xs2: std::borrow::Cow<'a, [f32]>,
2278        }
2279        let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
2280            let Self::Mapped {
2281                dtype,
2282                cols,
2283                row_scale,
2284                col_field,
2285                ..
2286            } = ts[i]
2287            else {
2288                unreachable!()
2289            };
2290            Ctx {
2291                bytes: ts[i].quant_bytes(),
2292                row_scale,
2293                cols: *cols,
2294                xs1: prescale(x1, col_field, *dtype),
2295                xs2: prescale(x2, col_field, *dtype),
2296            }
2297        });
2298        let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
2299        let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
2300        #[cfg(target_arch = "aarch64")]
2301        if sdot_enabled() {
2302            let acts: [(SplitAct, SplitAct); N] =
2303                std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
2304            let closures: [_; N] = std::array::from_fn(|i| {
2305                let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
2306                move |start: usize, end: usize| {
2307                    q8_range2_sdot(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
2308                }
2309            });
2310            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2311                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2312            pool.run_many(&parts);
2313            return;
2314        }
2315        #[cfg(target_arch = "x86_64")]
2316        if avx2_a8w8_enabled() {
2317            let acts: [(SplitAct, SplitAct); N] =
2318                std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
2319            let closures: [_; N] = std::array::from_fn(|i| {
2320                let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
2321                move |start: usize, end: usize| {
2322                    q8_range2_avx2(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
2323                }
2324            });
2325            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2326                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2327            pool.run_many(&parts);
2328            return;
2329        }
2330        let closures: [_; N] = std::array::from_fn(|i| {
2331            let (c, o1, o2) = (&ctxs[i], p1[i], p2[i]);
2332            move |start: usize, end: usize| {
2333                q8_range2_f32(
2334                    c.bytes,
2335                    c.row_scale,
2336                    &c.xs1,
2337                    &c.xs2,
2338                    c.cols,
2339                    o1,
2340                    o2,
2341                    start,
2342                    end,
2343                )
2344            }
2345        });
2346        let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2347            std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2348        pool.run_many(&parts);
2349    }
2350
2351    /// Fused gate+up matvec with SiLU·mul: for each row r, computes
2352    /// `silu(gate·x) * (up·x)` and writes to `out[r]`. ONE pool dispatch,
2353    /// no intermediate g/u buffers, no separate silu pass. Falls back
2354    /// (returns false) for unsupported dtype combos.
2355    pub fn matvec_silu_mul(
2356        gate: &QTensor,
2357        up: &QTensor,
2358        x: &[f32],
2359        out: &mut [f32],
2360        pool: Option<&Pool>,
2361    ) -> bool {
2362        let inter = gate.rows();
2363        debug_assert_eq!(up.rows(), inter);
2364        debug_assert_eq!(out.len(), inter);
2365        debug_assert_eq!(gate.cols(), up.cols());
2366        if !a8w8_enabled() {
2367            return false;
2368        }
2369        let act = split_act(x);
2370        let act = &act;
2371        let x_ref = x;
2372        let out_addr = SendMut(out.as_mut_ptr());
2373
2374        match (gate, up) {
2375            // Q4Block gate + Q4Block up (most common mobile q4 models)
2376            (
2377                Self::Mapped {
2378                    dtype: TensorDtype::Q4Block,
2379                    ..
2380                },
2381                Self::Mapped {
2382                    dtype: TensorDtype::Q4Block,
2383                    ..
2384                },
2385            ) => {
2386                let (gp, gs) = q4_split(gate.quant_bytes(), gate.rows(), gate.cols());
2387                let (up_p, up_s) = q4_split(up.quant_bytes(), up.rows(), up.cols());
2388                let gpr = gate.cols() / GROUP_SIZE;
2389                let cols = gate.cols();
2390                let run = move |start: usize, end: usize| {
2391                    for r in start..end {
2392                        let mut gv = dot_q4_row_i8(gp, gs, r * gpr, gpr, &act.xq) * act.sx;
2393                        let mut uv = dot_q4_row_i8(up_p, up_s, r * gpr, gpr, &act.xq) * act.sx;
2394                        for &(j, xv) in &act.outliers {
2395                            let flat = r * cols + j;
2396                            let gb = gp[flat / 2];
2397                            let gn = if flat & 1 == 0 { gb & 0x0F } else { gb >> 4 };
2398                            let gsc = f16_to_f32(u16::from_le_bytes([
2399                                gs[(flat / GROUP_SIZE) * 2],
2400                                gs[(flat / GROUP_SIZE) * 2 + 1],
2401                            ]));
2402                            gv += ((gn as i32 - 8) as f32) * gsc * xv;
2403                            let ub = up_p[flat / 2];
2404                            let un = if flat & 1 == 0 { ub & 0x0F } else { ub >> 4 };
2405                            let usc = f16_to_f32(u16::from_le_bytes([
2406                                up_s[(flat / GROUP_SIZE) * 2],
2407                                up_s[(flat / GROUP_SIZE) * 2 + 1],
2408                            ]));
2409                            uv += ((un as i32 - 8) as f32) * usc * xv;
2410                        }
2411                        let silu_g = gv / (1.0 + (-gv).exp());
2412                        // SAFETY: disjoint row ranges per worker.
2413                        unsafe { *out_addr.at(r) = silu_g * uv };
2414                    }
2415                };
2416                dispatch_rows(pool, inter, &run);
2417                true
2418            }
2419            // Q4Tiled gate + Q4Tiled up — one row pass, both tile
2420            // streams sequential, silu·mul fused (same per-row math as
2421            // `q4t_matvec`).
2422            (
2423                Self::Mapped {
2424                    dtype: TensorDtype::Q4Tiled,
2425                    ..
2426                },
2427                Self::Mapped {
2428                    dtype: TensorDtype::Q4Tiled,
2429                    ..
2430                },
2431            ) => {
2432                let g_bytes = gate.quant_bytes();
2433                let u_bytes = up.quant_bytes();
2434                let gpr = gate.cols() / GROUP_SIZE;
2435                let run = move |start: usize, end: usize| {
2436                    for r in start..end {
2437                        let mut gv = dot_q4t_row_i8(g_bytes, r, gpr, &act.xq) * act.sx;
2438                        let mut uv = dot_q4t_row_i8(u_bytes, r, gpr, &act.xq) * act.sx;
2439                        for &(j, xv) in &act.outliers {
2440                            let (w, s) = q4t_outlier(g_bytes, r, gpr, j);
2441                            gv += w * s * xv;
2442                            let (w, s) = q4t_outlier(u_bytes, r, gpr, j);
2443                            uv += w * s * xv;
2444                        }
2445                        let silu_g = gv / (1.0 + (-gv).exp());
2446                        // SAFETY: disjoint row ranges per worker.
2447                        unsafe { *out_addr.at(r) = silu_g * uv };
2448                    }
2449                };
2450                dispatch_rows(pool, inter, &run);
2451                true
2452            }
2453            // Q4TiledP gate + Q4TiledP up — the same fused row pass, with
2454            // each row's two ladders built once and spent on both streams.
2455            (
2456                Self::Mapped {
2457                    dtype: TensorDtype::Q4TiledP,
2458                    ..
2459                },
2460                Self::Mapped {
2461                    dtype: TensorDtype::Q4TiledP,
2462                    ..
2463                },
2464            ) => {
2465                let cols = gate.cols();
2466                let gpr = cols / GROUP_SIZE;
2467                let gv_view = Q4tpView::new(gate.quant_bytes(), inter, cols);
2468                let uv_view = Q4tpView::new(up.quant_bytes(), inter, cols);
2469                let run = |start: usize, end: usize| {
2470                    let (mut gsc, mut usc) = (vec![0f32; gpr], vec![0f32; gpr]);
2471                    for r in start..end {
2472                        gv_view.scales_into(r, gpr, &mut gsc);
2473                        uv_view.scales_into(r, gpr, &mut usc);
2474                        let mut gv = dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, &gsc) * act.sx;
2475                        let mut uv = dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, &usc) * act.sx;
2476                        for &(j, xv) in &act.outliers {
2477                            let (w, s) = q4tp_outlier(gv_view.nib, r, gpr, j, &gsc);
2478                            gv += w * s * xv;
2479                            let (w, s) = q4tp_outlier(uv_view.nib, r, gpr, j, &usc);
2480                            uv += w * s * xv;
2481                        }
2482                        let silu_g = gv / (1.0 + (-gv).exp());
2483                        // SAFETY: disjoint row ranges per worker.
2484                        unsafe { *out_addr.at(r) = silu_g * uv };
2485                    }
2486                };
2487                dispatch_rows(pool, inter, &run);
2488                true
2489            }
2490            // Q1 gate + Q1 up — one row pass over both sign streams,
2491            // silu·mul fused (the per-row math of `q1_range_a8w8`); the
2492            // activation group sums are shared by both streams. Without
2493            // this arm a q1 dense FFN paid two dispatches + a combine
2494            // loop — the exact barrier this function exists to remove.
2495            (
2496                Self::Mapped {
2497                    dtype: TensorDtype::Q1,
2498                    ..
2499                },
2500                Self::Mapped {
2501                    dtype: TensorDtype::Q1,
2502                    ..
2503                },
2504            ) => {
2505                let g_bytes = gate.quant_bytes();
2506                let u_bytes = up.quant_bytes();
2507                let gpr = gate.cols() / GROUP_SIZE;
2508                let gsum = q1_group_sums(&act.xq, gpr);
2509                let gsum = &gsum;
2510                let run = move |start: usize, end: usize| {
2511                    for r in start..end {
2512                        let mut gv = dot_q1_row_i8(g_bytes, r, gpr, &act.xq, gsum) * act.sx;
2513                        let mut uv = dot_q1_row_i8(u_bytes, r, gpr, &act.xq, gsum) * act.sx;
2514                        for &(j, xv) in &act.outliers {
2515                            let (w, s) = q1_outlier(g_bytes, r, gpr, j);
2516                            gv += w * s * xv;
2517                            let (w, s) = q1_outlier(u_bytes, r, gpr, j);
2518                            uv += w * s * xv;
2519                        }
2520                        let silu_g = gv / (1.0 + (-gv).exp());
2521                        // SAFETY: disjoint row ranges per worker.
2522                        unsafe { *out_addr.at(r) = silu_g * uv };
2523                    }
2524                };
2525                dispatch_rows(pool, inter, &run);
2526                true
2527            }
2528            // Q2TiledP gate + Q2TiledP up — the 2-bit expert pair (MoE
2529            // FFNs of the W2 class): one row pass, both ladders built
2530            // once, integer code dots with shared group sums.
2531            (
2532                Self::Mapped {
2533                    dtype: TensorDtype::Q2TiledP,
2534                    ..
2535                },
2536                Self::Mapped {
2537                    dtype: TensorDtype::Q2TiledP,
2538                    ..
2539                },
2540            ) => {
2541                let cols = gate.cols();
2542                let gpr = cols / GROUP_SIZE;
2543                let gv_view = Q4tpView::new_q2(gate.quant_bytes(), inter, cols);
2544                let uv_view = Q4tpView::new_q2(up.quant_bytes(), inter, cols);
2545                let gsum = q1_group_sums(&act.xq, gpr);
2546                let gsum = &gsum;
2547                let run = move |start: usize, end: usize| {
2548                    let (mut gsc, mut usc) = (vec![0f32; gpr], vec![0f32; gpr]);
2549                    for r in start..end {
2550                        gv_view.scales_into(r, gpr, &mut gsc);
2551                        uv_view.scales_into(r, gpr, &mut usc);
2552                        let mut gv =
2553                            dot_q2tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsum, &gsc) * act.sx;
2554                        let mut uv =
2555                            dot_q2tp_row_i8(uv_view.nib, r, gpr, &act.xq, gsum, &usc) * act.sx;
2556                        for &(j, xv) in &act.outliers {
2557                            let (w, s) = q2tp_outlier(gv_view.nib, r, gpr, j, &gsc);
2558                            gv += w * s * xv;
2559                            let (w, s) = q2tp_outlier(uv_view.nib, r, gpr, j, &usc);
2560                            uv += w * s * xv;
2561                        }
2562                        let silu_g = gv / (1.0 + (-gv).exp());
2563                        // SAFETY: disjoint row ranges per worker.
2564                        unsafe { *out_addr.at(r) = silu_g * uv };
2565                    }
2566                };
2567                dispatch_rows(pool, inter, &run);
2568                true
2569            }
2570            // Q8Row gate + Q8Row up — one row pass over both i8 streams.
2571            // Q8_2f stays out on purpose: its column field prescales the
2572            // activations PER TENSOR, which breaks this fn's shared
2573            // split_act contract — it keeps the two-dispatch path.
2574            (
2575                Self::Mapped {
2576                    dtype: TensorDtype::Q8Row,
2577                    row_scale: g_rs,
2578                    ..
2579                },
2580                Self::Mapped {
2581                    dtype: TensorDtype::Q8Row,
2582                    row_scale: u_rs,
2583                    ..
2584                },
2585            ) => {
2586                let g_bytes = gate.quant_bytes();
2587                let u_bytes = up.quant_bytes();
2588                let cols = gate.cols();
2589                let run = move |start: usize, end: usize| {
2590                    for r in start..end {
2591                        let gv = q8_row_dot(&g_bytes[r * cols..(r + 1) * cols], act) * g_rs[r];
2592                        let uv = q8_row_dot(&u_bytes[r * cols..(r + 1) * cols], act) * u_rs[r];
2593                        let silu_g = gv / (1.0 + (-gv).exp());
2594                        // SAFETY: disjoint row ranges per worker.
2595                        unsafe { *out_addr.at(r) = silu_g * uv };
2596                    }
2597                };
2598                dispatch_rows(pool, inter, &run);
2599                true
2600            }
2601            // Q1T gate + Q1T up
2602            (
2603                Self::Mapped {
2604                    dtype: TensorDtype::Q1T,
2605                    ..
2606                },
2607                Self::Mapped {
2608                    dtype: TensorDtype::Q1T,
2609                    ..
2610                },
2611            ) => {
2612                const TILE: usize = cortiq_core::quant::Q1T_TILE;
2613                let g_bytes = gate.quant_bytes();
2614                let u_bytes = up.quant_bytes();
2615                let gpr = gate.cols() / GROUP_SIZE;
2616                let (g_rp, g_ent, g_ov) = q1t_overlay(g_bytes, inter * gpr * TILE, inter);
2617                let (u_rp, u_ent, u_ov) = q1t_overlay(u_bytes, inter * gpr * TILE, inter);
2618                let run = move |start: usize, end: usize| {
2619                    for r in start..end {
2620                        let mut gv = q1t_dot_row_i8(g_bytes, r, gpr, &act.xq) * act.sx;
2621                        let mut uv = q1t_dot_row_i8(u_bytes, r, gpr, &act.xq) * act.sx;
2622                        for &(j, xv) in &act.outliers {
2623                            gv += q1t_base_weight(g_bytes, r, gpr, j) * xv;
2624                            uv += q1t_base_weight(u_bytes, r, gpr, j) * xv;
2625                        }
2626                        gv += q1t_row_outlier_correction(g_bytes, r, g_rp, g_ent, g_ov, x_ref);
2627                        uv += q1t_row_outlier_correction(u_bytes, r, u_rp, u_ent, u_ov, x_ref);
2628                        let silu_g = gv / (1.0 + (-gv).exp());
2629                        // SAFETY: disjoint row ranges per worker.
2630                        unsafe { *out_addr.at(r) = silu_g * uv };
2631                    }
2632                };
2633                dispatch_rows(pool, inter, &run);
2634                true
2635            }
2636            _ => false,
2637        }
2638    }
2639
2640    /// Every routed expert's fused gate/up/SiLU under ONE pool dispatch.
2641    ///
2642    /// The per-expert path pays a pool barrier per expert per stage: at 9
2643    /// experts over 40 layers that is ~720 barriers a token, and a decode
2644    /// profile of Qwen3.6-35B-A3B showed the pool parked in
2645    /// `psynch_cvwait` about twice as long as it spent computing. Laying
2646    /// every expert's rows end-to-end in one virtual row space collapses
2647    /// the stage to a single dispatch. The per-row body is the
2648    /// single-expert q4tp arm verbatim, so outputs are bit-identical.
2649    ///
2650    /// `false` = something is outside the fused q4tp kernel (dtype, shape,
2651    /// or the `CMF_SDOT=0` exact contract); the caller walks the ordinary
2652    /// per-expert path.
2653    pub fn moe_gate_up_many(
2654        pairs: &[(&QTensor, &QTensor)],
2655        x: &[f32],
2656        outs: &mut [Vec<f32>],
2657        pool: Option<&Pool>,
2658    ) -> bool {
2659        if pairs.is_empty() || pairs.len() != outs.len() || !a8w8_enabled() {
2660            return false;
2661        }
2662        let inter = pairs[0].0.rows();
2663        let cols = pairs[0].0.cols();
2664        if cols % GROUP_SIZE != 0 {
2665            return false;
2666        }
2667        let gpr = cols / GROUP_SIZE;
2668        // Uniform layout across every routed pair: q4tp, or the 2-bit
2669        // profile's q2tp gate/up (the W2 class). Mixed sets refuse.
2670        let q2 = matches!(
2671            pairs[0].0,
2672            Self::Mapped {
2673                dtype: TensorDtype::Q2TiledP,
2674                ..
2675            }
2676        );
2677        let want = if q2 {
2678            TensorDtype::Q2TiledP
2679        } else {
2680            TensorDtype::Q4TiledP
2681        };
2682        let mut views = Vec::with_capacity(pairs.len() * 2);
2683        for ((g, u), o) in pairs.iter().zip(outs.iter()) {
2684            let both = matches!(g, Self::Mapped { dtype, .. } if *dtype == want)
2685                && matches!(u, Self::Mapped { dtype, .. } if *dtype == want);
2686            if !both
2687                || g.rows() != inter
2688                || u.rows() != inter
2689                || g.cols() != cols
2690                || u.cols() != cols
2691                || o.len() != inter
2692            {
2693                return false;
2694            }
2695            let mk = if q2 { Q4tpView::new_q2 } else { Q4tpView::new };
2696            views.push(mk(g.quant_bytes(), inter, cols));
2697            views.push(mk(u.quant_bytes(), inter, cols));
2698        }
2699        let act = split_act(x);
2700        let gsum = if q2 {
2701            q1_group_sums(&act.xq, gpr)
2702        } else {
2703            Vec::new()
2704        };
2705        let (act, gsum) = (&act, &gsum);
2706        let ptrs: Vec<SendMut> = outs.iter_mut().map(|o| SendMut(o.as_mut_ptr())).collect();
2707        let (views, ptrs) = (&views, &ptrs);
2708        let run = |start: usize, end: usize| {
2709            let (mut gsc, mut usc) = (vec![0f32; gpr], vec![0f32; gpr]);
2710            for flat in start..end {
2711                let (e, r) = (flat / inter, flat % inter);
2712                let gv_view = &views[e * 2];
2713                let uv_view = &views[e * 2 + 1];
2714                gv_view.scales_into(r, gpr, &mut gsc);
2715                uv_view.scales_into(r, gpr, &mut usc);
2716                let (mut gv, mut uv) = if q2 {
2717                    (
2718                        dot_q2tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsum, &gsc) * act.sx,
2719                        dot_q2tp_row_i8(uv_view.nib, r, gpr, &act.xq, gsum, &usc) * act.sx,
2720                    )
2721                } else {
2722                    (
2723                        dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, &gsc) * act.sx,
2724                        dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, &usc) * act.sx,
2725                    )
2726                };
2727                for &(j, xv) in &act.outliers {
2728                    let (og, ou) = if q2 {
2729                        (
2730                            q2tp_outlier(gv_view.nib, r, gpr, j, &gsc),
2731                            q2tp_outlier(uv_view.nib, r, gpr, j, &usc),
2732                        )
2733                    } else {
2734                        (
2735                            q4tp_outlier(gv_view.nib, r, gpr, j, &gsc),
2736                            q4tp_outlier(uv_view.nib, r, gpr, j, &usc),
2737                        )
2738                    };
2739                    gv += og.0 * og.1 * xv;
2740                    uv += ou.0 * ou.1 * xv;
2741                }
2742                let silu_g = gv / (1.0 + (-gv).exp());
2743                // SAFETY: one worker owns each (expert, row) pair.
2744                unsafe { *ptrs[e].at(r) = silu_g * uv };
2745            }
2746        };
2747        dispatch_rows(pool, pairs.len() * inter, &run);
2748        true
2749    }
2750
2751    /// Every routed expert's down projection, weighted and summed into
2752    /// `out`, under ONE pool dispatch.
2753    ///
2754    /// Partitioned by OUTPUT row rather than by expert: each row is owned
2755    /// by a single worker, so the experts are summed in the caller's order
2756    /// — the same sequence of f32 adds the serial `out[i] += w·eo[i]` loop
2757    /// performs, hence bit-identical. Partitioning by expert instead would
2758    /// race on the shared accumulator.
2759    pub fn moe_down_many(
2760        downs: &[&QTensor],
2761        gs: &[Vec<f32>],
2762        weights: &[f32],
2763        out: &mut [f32],
2764        pool: Option<&Pool>,
2765    ) -> bool {
2766        if downs.is_empty()
2767            || downs.len() != gs.len()
2768            || downs.len() != weights.len()
2769            || !a8w8_enabled()
2770        {
2771            return false;
2772        }
2773        let rows = out.len();
2774        let cols = downs[0].cols();
2775        if cols % GROUP_SIZE != 0 {
2776            return false;
2777        }
2778        let gpr = cols / GROUP_SIZE;
2779        let mut views = Vec::with_capacity(downs.len());
2780        for (d, g) in downs.iter().zip(gs.iter()) {
2781            if !matches!(
2782                d,
2783                Self::Mapped {
2784                    dtype: TensorDtype::Q4TiledP,
2785                    ..
2786                }
2787            ) || d.rows() != rows
2788                || d.cols() != cols
2789                || g.len() != cols
2790            {
2791                return false;
2792            }
2793            views.push(Q4tpView::new(d.quant_bytes(), rows, cols));
2794        }
2795        // One int8 split per expert — the activation vectors differ.
2796        let acts: Vec<SplitAct> = gs.iter().map(|g| split_act(g)).collect();
2797        // Partitioned by OUTPUT row, with the experts folded inside: each
2798        // row is owned by one worker, so they are summed in the caller's
2799        // order — the same f32 sequence the serial `out[i] += w·eo[i]`
2800        // loop produces. Partitioning by expert instead would either race
2801        // on the accumulator or need a scratch plane and a second pass;
2802        // measured, that variant was a wash, so this keeps the simpler
2803        // shape.
2804        let out_addr = SendMut(out.as_mut_ptr());
2805        let (views, acts, weights) = (&views, &acts, &weights);
2806        let run = |start: usize, end: usize| {
2807            let mut sc = vec![0f32; gpr];
2808            for r in start..end {
2809                let mut acc = 0f32;
2810                for (e, v) in views.iter().enumerate() {
2811                    v.scales_into(r, gpr, &mut sc);
2812                    let a = &acts[e];
2813                    let mut d = dot_q4tp_row_i8(v.nib, r, gpr, &a.xq, &sc) * a.sx;
2814                    for &(j, xv) in &a.outliers {
2815                        let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
2816                        d += w * s * xv;
2817                    }
2818                    acc += weights[e] * d;
2819                }
2820                // SAFETY: disjoint row ranges per worker.
2821                unsafe { *out_addr.at(r) = acc };
2822            }
2823        };
2824        dispatch_rows(pool, rows, &run);
2825        true
2826    }
2827}
2828
2829/// Batched q8 kernel: same math as qmatvec, the row makes a single
2830/// pass from memory for the whole batch.
2831/// Accelerate CBLAS — the Apple AMX matrix units, the same engine
2832/// llama.cpp's `-ngl 0` prefill rides via ggml-blas.
2833#[cfg(target_os = "macos")]
2834mod accel_blas {
2835    #[link(name = "Accelerate", kind = "framework")]
2836    unsafe extern "C" {
2837        pub fn cblas_sgemm(
2838            order: i32,
2839            trans_a: i32,
2840            trans_b: i32,
2841            m: i32,
2842            n: i32,
2843            k: i32,
2844            alpha: f32,
2845            a: *const f32,
2846            lda: i32,
2847            b: *const f32,
2848            ldb: i32,
2849            beta: f32,
2850            c: *mut f32,
2851            ldc: i32,
2852        );
2853    }
2854}
2855
2856#[cfg(target_os = "macos")]
2857pub(crate) fn accel_gemm_enabled() -> bool {
2858    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2859    *ON.get_or_init(|| std::env::var("CMF_ACCEL").map(|v| v != "0").unwrap_or(true))
2860}
2861
2862/// Off macOS the "accel" GEMM is the portable NEON micro-kernel below —
2863/// same entry point, so the batched-attention path opens on mobile.
2864#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
2865pub(crate) fn accel_gemm_enabled() -> bool {
2866    true
2867}
2868
2869/// Portable NEON f32 GEMM (row-major, optional Bᵀ): a 4×8 fmla
2870/// micro-kernel with A broadcast against B panels — the mobile stand-in
2871/// for Accelerate in the batched causal attention (QKᵀ and P·V). Not a
2872/// BLAS: shapes here are the attention panels (m ≤ heads·chunk,
2873/// k = head_dim or context), and the goal is removing the per-position
2874/// quadratic wall, not peak GEMM.
2875#[cfg(target_arch = "aarch64")]
2876#[allow(clippy::too_many_arguments)]
2877pub(crate) fn neon_gemm_rm(
2878    m: usize,
2879    n: usize,
2880    k: usize,
2881    alpha: f32,
2882    a: &[f32],
2883    lda: usize,
2884    b_mat: &[f32],
2885    ldb: usize,
2886    b_rows_are_n: bool,
2887    c: &mut [f32],
2888    ldc: usize,
2889) {
2890    debug_assert!(a.len() >= (m - 1) * lda + k);
2891    debug_assert!(c.len() >= (m - 1) * ldc + n);
2892    // SAFETY: bounds asserted above; NEON is baseline on aarch64.
2893    unsafe {
2894        use core::arch::aarch64::*;
2895        let mut i = 0usize;
2896        while i < m {
2897            let mi = (m - i).min(4);
2898            let mut j = 0usize;
2899            while j < n {
2900                let nj = (n - j).min(8);
2901                if mi == 4 && nj == 8 {
2902                    let (mut c0a, mut c0b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
2903                    let (mut c1a, mut c1b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
2904                    let (mut c2a, mut c2b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
2905                    let (mut c3a, mut c3b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
2906                    for p in 0..k {
2907                        let (b0, b1) = if b_rows_are_n {
2908                            // B is [n, k]: column p of Bᵀ = element p of
2909                            // eight consecutive B rows — gathered.
2910                            let base = b_mat.as_ptr().add(j * ldb + p);
2911                            let g = |o: usize| *base.add(o * ldb);
2912                            ([g(0), g(1), g(2), g(3)], [g(4), g(5), g(6), g(7)])
2913                        } else {
2914                            let base = b_mat.as_ptr().add(p * ldb + j);
2915                            (
2916                                [*base, *base.add(1), *base.add(2), *base.add(3)],
2917                                [*base.add(4), *base.add(5), *base.add(6), *base.add(7)],
2918                            )
2919                        };
2920                        let bv0 = vld1q_f32(b0.as_ptr());
2921                        let bv1 = vld1q_f32(b1.as_ptr());
2922                        let a0 = vdupq_n_f32(*a.as_ptr().add(i * lda + p));
2923                        let a1 = vdupq_n_f32(*a.as_ptr().add((i + 1) * lda + p));
2924                        let a2 = vdupq_n_f32(*a.as_ptr().add((i + 2) * lda + p));
2925                        let a3 = vdupq_n_f32(*a.as_ptr().add((i + 3) * lda + p));
2926                        c0a = vfmaq_f32(c0a, a0, bv0);
2927                        c0b = vfmaq_f32(c0b, a0, bv1);
2928                        c1a = vfmaq_f32(c1a, a1, bv0);
2929                        c1b = vfmaq_f32(c1b, a1, bv1);
2930                        c2a = vfmaq_f32(c2a, a2, bv0);
2931                        c2b = vfmaq_f32(c2b, a2, bv1);
2932                        c3a = vfmaq_f32(c3a, a3, bv0);
2933                        c3b = vfmaq_f32(c3b, a3, bv1);
2934                    }
2935                    let al = vdupq_n_f32(alpha);
2936                    for (r, (ca, cb)) in [(c0a, c0b), (c1a, c1b), (c2a, c2b), (c3a, c3b)]
2937                        .iter()
2938                        .enumerate()
2939                    {
2940                        let dst = c.as_mut_ptr().add((i + r) * ldc + j);
2941                        vst1q_f32(dst, vmulq_f32(*ca, al));
2942                        vst1q_f32(dst.add(4), vmulq_f32(*cb, al));
2943                    }
2944                } else {
2945                    for r in 0..mi {
2946                        for q in 0..nj {
2947                            let mut acc = 0f32;
2948                            for p in 0..k {
2949                                let bv = if b_rows_are_n {
2950                                    b_mat[(j + q) * ldb + p]
2951                                } else {
2952                                    b_mat[p * ldb + j + q]
2953                                };
2954                                acc += a[(i + r) * lda + p] * bv;
2955                            }
2956                            c[(i + r) * ldc + j + q] = acc * alpha;
2957                        }
2958                    }
2959                }
2960                j += nj;
2961            }
2962            i += mi;
2963        }
2964    }
2965}
2966
2967/// Off-macOS aarch64: the batched attention rides the NEON micro-GEMM.
2968#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
2969#[allow(clippy::too_many_arguments)]
2970pub(crate) fn sgemm_rm(
2971    m: usize,
2972    n: usize,
2973    k: usize,
2974    alpha: f32,
2975    a: &[f32],
2976    lda: usize,
2977    b_mat: &[f32],
2978    ldb: usize,
2979    b_rows_are_n: bool,
2980    c: &mut [f32],
2981    ldc: usize,
2982) {
2983    neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
2984}
2985
2986/// Row-major f32 GEMM, exposed for offline tools (the AWNP pass builds a
2987/// per-layer projection and applies it to every expert; a naive triple loop
2988/// would turn a two-minute job into half an hour).
2989#[allow(clippy::too_many_arguments)]
2990pub fn sgemm_public(
2991    m: usize,
2992    n: usize,
2993    k: usize,
2994    alpha: f32,
2995    a: &[f32],
2996    lda: usize,
2997    b_mat: &[f32],
2998    ldb: usize,
2999    b_rows_are_n: bool,
3000    c: &mut [f32],
3001    ldc: usize,
3002) {
3003    #[cfg(any(target_os = "macos", target_arch = "aarch64"))]
3004    {
3005        sgemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
3006    }
3007    // x86 without Accelerate has no sgemm_rm: the specialized paths there are
3008    // quantized kernels, not an f32 GEMM. Only the offline AWNP pass reaches
3009    // this, so correctness matters and throughput does not — a triple loop is
3010    // the honest fallback rather than a reason to make the tool macOS-only.
3011    #[cfg(not(any(target_os = "macos", target_arch = "aarch64")))]
3012    {
3013        for i in 0..m {
3014            for j in 0..n {
3015                let mut acc = 0f32;
3016                for p in 0..k {
3017                    let bv = if b_rows_are_n {
3018                        b_mat[j * ldb + p]
3019                    } else {
3020                        b_mat[p * ldb + j]
3021                    };
3022                    acc += a[i * lda + p] * bv;
3023                }
3024                c[i * ldc + j] = alpha * acc;
3025            }
3026        }
3027    }
3028}
3029
3030/// Row-major f32 GEMM on Accelerate: C[m,n] = alpha·A[m,k] × B(ᵀ).
3031/// `b_rows_are_n` = true multiplies by Bᵀ where B is stored [n, k].
3032#[cfg(target_os = "macos")]
3033#[allow(clippy::too_many_arguments)]
3034pub(crate) fn sgemm_rm(
3035    m: usize,
3036    n: usize,
3037    k: usize,
3038    alpha: f32,
3039    a: &[f32],
3040    lda: usize,
3041    b_mat: &[f32],
3042    ldb: usize,
3043    b_rows_are_n: bool,
3044    c: &mut [f32],
3045    ldc: usize,
3046) {
3047    debug_assert!(a.len() >= (m - 1) * lda + k);
3048    debug_assert!(c.len() >= (m - 1) * ldc + n);
3049    // Test hook: route the attention GEMMs through the portable NEON
3050    // micro-kernel ON APPLE SILICON — how the mobile batched attend is
3051    // measured without a phone in the loop. (Intel macOS has no NEON —
3052    // the hook is a no-op there, Accelerate continues below.)
3053    #[cfg(target_arch = "aarch64")]
3054    if std::env::var("CMF_FORCE_NEON_GEMM")
3055        .map(|v| v == "1")
3056        .unwrap_or(false)
3057    {
3058        return neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
3059    }
3060    unsafe {
3061        accel_blas::cblas_sgemm(
3062            101, // RowMajor
3063            111, // NoTrans A
3064            if b_rows_are_n { 112 } else { 111 },
3065            m as i32,
3066            n as i32,
3067            k as i32,
3068            alpha,
3069            a.as_ptr(),
3070            lda as i32,
3071            b_mat.as_ptr(),
3072            ldb as i32,
3073            0.0,
3074            c.as_mut_ptr(),
3075            ldc as i32,
3076        );
3077    }
3078}
3079
3080/// Prefill GEMM through Accelerate (macOS): dequantize q8 rows into
3081/// f32 tiles (scale folded in, pool-parallel) and multiply each tile
3082/// on the AMX with one row-major sgemm. Tiles live in cache, weights
3083/// stream once. Numerics are f32-GEMM (not the int8 dot): prefill
3084/// logits shift within f32 rounding — tolerance-class, like every
3085/// reduction-order change; decode (M=1) never takes this path.
3086#[cfg(target_os = "macos")]
3087fn qmatmat_accel(
3088    q: &[u8],
3089    row_scale: &[f32],
3090    pre: &[std::borrow::Cow<'_, [f32]>],
3091    rows: usize,
3092    cols: usize,
3093    out: &mut [f32],
3094    pool: Option<&Pool>,
3095) {
3096    // NOTE: double-buffering the dequant against the sgemm (a scoped
3097    // thread driving the pool on tile k+1 while the caller multiplies
3098    // tile k) was tried and LOST ~6%: Accelerate's sgemm is itself
3099    // multithreaded, and the dequant workers just steal its cores.
3100    const TR: usize = 2048;
3101    let b = pre.len();
3102    thread_local! {
3103        static XPANEL: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
3104        static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
3105    }
3106    XPANEL.with(|xp| {
3107        WTILE.with(|wt| {
3108            let mut xpanel = xp.borrow_mut();
3109            xpanel.clear();
3110            for x in pre {
3111                xpanel.extend_from_slice(x);
3112            }
3113            let mut wtile = wt.borrow_mut();
3114            wtile.resize(TR * cols, 0.0);
3115            let mut r0 = 0usize;
3116            while r0 < rows {
3117                let tr = TR.min(rows - r0);
3118                // Dequant the tile (scale folded) — pool-parallel.
3119                let wt_addr = SendMut(wtile.as_mut_ptr());
3120                let run = |start: usize, end: usize| {
3121                    for r in start..end {
3122                        let row = &q[(r0 + r) * cols..(r0 + r + 1) * cols];
3123                        let s = row_scale[r0 + r];
3124                        // SAFETY: workers cover disjoint r ranges.
3125                        let dst =
3126                            unsafe { std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols) };
3127                        for (d, &v) in dst.iter_mut().zip(row) {
3128                            *d = (v as i8) as f32 * s;
3129                        }
3130                    }
3131                };
3132                dispatch_rows(pool, tr, &run);
3133                // C[b, tr] (at column r0 of out[b, rows]) = X · Wtileᵀ
3134                unsafe {
3135                    accel_blas::cblas_sgemm(
3136                        101, // RowMajor
3137                        111, // NoTrans A
3138                        112, // Trans B
3139                        b as i32,
3140                        tr as i32,
3141                        cols as i32,
3142                        1.0,
3143                        xpanel.as_ptr(),
3144                        cols as i32,
3145                        wtile.as_ptr(),
3146                        cols as i32,
3147                        0.0,
3148                        out.as_mut_ptr().add(r0),
3149                        rows as i32,
3150                    );
3151                }
3152                r0 += tr;
3153            }
3154        })
3155    });
3156}
3157
3158fn qmatmat(
3159    q: &[u8],
3160    row_scale: &[f32],
3161    pre: &[std::borrow::Cow<'_, [f32]>],
3162    rows: usize,
3163    cols: usize,
3164    out: &mut [f32],
3165    pool: Option<&Pool>,
3166) {
3167    let b = pre.len();
3168    debug_assert_eq!(out.len(), b * rows);
3169    // Big prefill batches ride the AMX (roadmap PR3): the row×batch
3170    // SDOT loop below peaks near the CPU's dot throughput, an order
3171    // below the matrix units. Small tensors and tiny test models stay
3172    // on the exact integer path.
3173    #[cfg(target_os = "macos")]
3174    if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
3175        qmatmat_accel(q, row_scale, pre, rows, cols, out, pool);
3176        return;
3177    }
3178    #[cfg(target_arch = "aarch64")]
3179    if sdot_enabled() {
3180        let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
3181        let out_addr = SendMut(out.as_mut_ptr());
3182        // Blocked 2×4 (mobile prefill: no AMX to fall back on — this
3183        // path IS the ARM prefill GEMM off Apple silicon).
3184        let blocked_ok = blocked_enabled();
3185        let use_i8mm = i8mm_enabled();
3186        if blocked_ok {
3187            let run = |start: usize, end: usize| {
3188                let mut o = start;
3189                while o < end {
3190                    if o + 2 <= end {
3191                        let r0 = &q[o * cols..(o + 1) * cols];
3192                        let r1 = &q[(o + 1) * cols..(o + 2) * cols];
3193                        let mut bi = 0usize;
3194                        while bi + 4 <= acts.len() {
3195                            let xs = [
3196                                acts[bi].xq.as_slice(),
3197                                acts[bi + 1].xq.as_slice(),
3198                                acts[bi + 2].xq.as_slice(),
3199                                acts[bi + 3].xq.as_slice(),
3200                            ];
3201                            let d = if use_i8mm {
3202                                unsafe { dot_i8_smmla_2x4(r0, r1, xs) }
3203                            } else {
3204                                unsafe { dot_i8_sdot_2x4(r0, r1, xs) }
3205                            };
3206                            for (r, row) in [r0, r1].into_iter().enumerate() {
3207                                for k in 0..4 {
3208                                    let act = &acts[bi + k];
3209                                    let mut v = d[r][k] as f32 * act.sx;
3210                                    for &(j, xv) in &act.outliers {
3211                                        v += (row[j] as i8) as f32 * xv;
3212                                    }
3213                                    unsafe {
3214                                        *out_addr.at((bi + k) * rows + o + r) = v * row_scale[o + r]
3215                                    };
3216                                }
3217                            }
3218                            bi += 4;
3219                        }
3220                        while bi < acts.len() {
3221                            for (r, row) in [r0, r1].into_iter().enumerate() {
3222                                let v = row_dot_sdot(row, &acts[bi]) * row_scale[o + r];
3223                                unsafe { *out_addr.at(bi * rows + o + r) = v };
3224                            }
3225                            bi += 1;
3226                        }
3227                        o += 2;
3228                    } else {
3229                        let row = &q[o * cols..(o + 1) * cols];
3230                        for (bi, act) in acts.iter().enumerate() {
3231                            let v = row_dot_sdot(row, act) * row_scale[o];
3232                            unsafe { *out_addr.at(bi * rows + o) = v };
3233                        }
3234                        o += 1;
3235                    }
3236                }
3237            };
3238            dispatch_rows(pool, rows, &run);
3239            return;
3240        }
3241        let run = |start: usize, end: usize| {
3242            for o in start..end {
3243                let row = &q[o * cols..(o + 1) * cols];
3244                for (bi, act) in acts.iter().enumerate() {
3245                    let v = row_dot_sdot(row, act) * row_scale[o];
3246                    unsafe { *out_addr.at(bi * rows + o) = v };
3247                }
3248            }
3249        };
3250        dispatch_rows(pool, rows, &run);
3251        return;
3252    }
3253    // x86 A8W8 batch. Non-VNNI parts take the BLOCKED 2×4 kernel
3254    // (roadmap P0: two weight rows' abs() stay in registers across four
3255    // activation streams); VNNI machines keep the per-row bias-trick
3256    // dot, which is already throughput-bound there.
3257    #[cfg(target_arch = "x86_64")]
3258    if avx2_a8w8_enabled() {
3259        let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
3260        let out_addr = SendMut(out.as_mut_ptr());
3261        // CMF_X86_BLOCKED=0 forces the per-row path (paired in-process
3262        // A/B on noisy shared-vCPU hosts).
3263        let blocked_ok = blocked_enabled();
3264        if !avx512vnni_enabled() && blocked_ok {
3265            let run = |start: usize, end: usize| {
3266                let mut o = start;
3267                while o < end {
3268                    if o + 2 <= end {
3269                        let r0 = &q[o * cols..(o + 1) * cols];
3270                        let r1 = &q[(o + 1) * cols..(o + 2) * cols];
3271                        let mut bi = 0usize;
3272                        while bi + 4 <= acts.len() {
3273                            let xs = [
3274                                acts[bi].xq.as_slice(),
3275                                acts[bi + 1].xq.as_slice(),
3276                                acts[bi + 2].xq.as_slice(),
3277                                acts[bi + 3].xq.as_slice(),
3278                            ];
3279                            let d = unsafe { dot_i8_i8_avx2_2x4(r0, r1, xs) };
3280                            for (r, row) in [r0, r1].into_iter().enumerate() {
3281                                for k in 0..4 {
3282                                    let act = &acts[bi + k];
3283                                    let mut v = d[r][k] as f32 * act.sx;
3284                                    for &(j, xv) in &act.outliers {
3285                                        v += (row[j] as i8) as f32 * xv;
3286                                    }
3287                                    unsafe {
3288                                        *out_addr.at((bi + k) * rows + o + r) = v * row_scale[o + r]
3289                                    };
3290                                }
3291                            }
3292                            bi += 4;
3293                        }
3294                        while bi < acts.len() {
3295                            for (r, row) in [r0, r1].into_iter().enumerate() {
3296                                let v = row_dot_avx2(row, &acts[bi]) * row_scale[o + r];
3297                                unsafe { *out_addr.at(bi * rows + o + r) = v };
3298                            }
3299                            bi += 1;
3300                        }
3301                        o += 2;
3302                    } else {
3303                        let row = &q[o * cols..(o + 1) * cols];
3304                        for (bi, act) in acts.iter().enumerate() {
3305                            let v = row_dot_avx2(row, act) * row_scale[o];
3306                            unsafe { *out_addr.at(bi * rows + o) = v };
3307                        }
3308                        o += 1;
3309                    }
3310                }
3311            };
3312            dispatch_rows(pool, rows, &run);
3313            return;
3314        }
3315        let run = |start: usize, end: usize| {
3316            for o in start..end {
3317                let row = &q[o * cols..(o + 1) * cols];
3318                for (bi, act) in acts.iter().enumerate() {
3319                    let v = row_dot_avx2(row, act) * row_scale[o];
3320                    unsafe { *out_addr.at(bi * rows + o) = v };
3321                }
3322            }
3323        };
3324        dispatch_rows(pool, rows, &run);
3325        return;
3326    }
3327    let out_addr = SendMut(out.as_mut_ptr());
3328    let run = |start: usize, end: usize| {
3329        for o in start..end {
3330            let row = &q[o * cols..(o + 1) * cols];
3331            for (bi, x) in pre.iter().enumerate() {
3332                let mut acc = 0f32;
3333                for j in 0..cols {
3334                    acc += (row[j] as i8) as f32 * x[j];
3335                }
3336                unsafe { *out_addr.at(bi * rows + o) = acc * row_scale[o] };
3337            }
3338        }
3339    };
3340    dispatch_rows(pool, rows, &run);
3341}
3342
3343/// Split rows across pool workers (shared qmatvec pattern). Self-balancing
3344/// — see `Pool::run_rows` for why a static 1/n split is wrong here.
3345fn dispatch_rows(pool: Option<&Pool>, rows: usize, run: &(dyn Fn(usize, usize) + Sync)) {
3346    match pool {
3347        Some(pool) if rows >= 256 => pool.run_rows(rows, run),
3348        _ => run(0, rows),
3349    }
3350}
3351
3352/// Split a q4_block blob into (packed nibbles, f16 group scales).
3353fn q4_split(bytes: &[u8], rows: usize, cols: usize) -> (&[u8], &[u8]) {
3354    let groups = rows * cols / GROUP_SIZE;
3355    bytes.split_at(groups * 16)
3356}
3357
3358/// SIMD unpack for the dominant vbit width B=4 (94% of rows on the
3359/// log2-shape calibration): 16 packed bytes -> 32 centered i8 values.
3360/// vbit packs MSB-first, so the HIGH nibble is the even element
3361/// (opposite of q4_block's lo-first interleave). Centering is u-7.
3362#[inline]
3363fn vbit_fill4(data: &[u8], buf: &mut [u8]) {
3364    #[cfg(target_arch = "aarch64")]
3365    unsafe {
3366        return vbit_fill4_neon(data, buf);
3367    }
3368    #[cfg(target_arch = "x86_64")]
3369    if avx2_enabled() {
3370        return unsafe { vbit_fill4_avx2(data, buf) };
3371    }
3372    #[allow(unreachable_code)]
3373    for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
3374        let u = unpack8::<4>(&data[blk * 4..]);
3375        for k in 0..8 {
3376            chunk[k] = (u[k] - 7) as i8 as u8;
3377        }
3378    }
3379}
3380
3381#[cfg(target_arch = "aarch64")]
3382#[target_feature(enable = "neon")]
3383unsafe fn vbit_fill4_neon(data: &[u8], buf: &mut [u8]) {
3384    // SAFETY: buf.len() is a multiple of GROUP_SIZE=32; data holds
3385    // buf.len()/2 packed bytes (validated at load).
3386    unsafe {
3387        use core::arch::aarch64::*;
3388        let n = buf.len();
3389        let mask = vdupq_n_u8(0x0F);
3390        let seven = vdupq_n_s8(7);
3391        let mut g = 0usize;
3392        while g * 32 + 32 <= n {
3393            let b = vld1q_u8(data.as_ptr().add(g * 16));
3394            let hi = vshrq_n_u8::<4>(b);
3395            let lo = vandq_u8(b, mask);
3396            let z0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(hi, lo)), seven);
3397            let z1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(hi, lo)), seven);
3398            vst1q_u8(buf.as_mut_ptr().add(g * 32), vreinterpretq_u8_s8(z0));
3399            vst1q_u8(buf.as_mut_ptr().add(g * 32 + 16), vreinterpretq_u8_s8(z1));
3400            g += 1;
3401        }
3402    }
3403}
3404
3405#[cfg(target_arch = "x86_64")]
3406#[target_feature(enable = "avx2")]
3407unsafe fn vbit_fill4_avx2(data: &[u8], buf: &mut [u8]) {
3408    // SAFETY: see vbit_fill4_neon.
3409    unsafe {
3410        use core::arch::x86_64::*;
3411        let n = buf.len();
3412        let mask = _mm_set1_epi8(0x0F);
3413        let seven = _mm256_set1_epi8(7);
3414        let mut g = 0usize;
3415        while g * 32 + 32 <= n {
3416            let b = _mm_loadu_si128(data.as_ptr().add(g * 16) as *const __m128i);
3417            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), mask);
3418            let lo = _mm_and_si128(b, mask);
3419            let z = _mm256_sub_epi8(
3420                _mm256_set_m128i(_mm_unpackhi_epi8(hi, lo), _mm_unpacklo_epi8(hi, lo)),
3421                seven,
3422            );
3423            _mm256_storeu_si256(buf.as_mut_ptr().add(g * 32) as *mut __m256i, z);
3424            g += 1;
3425        }
3426    }
3427}
3428
3429/// Unpack 8 MSB-first B-bit values from exactly B bytes (fixed shifts —
3430/// no serial bit-buffer, auto-vectorizable). Every 32-value group starts
3431/// byte-aligned (32·B/8 is integral for B∈3..8), so groups decompose
3432/// into 4 such blocks.
3433#[inline(always)]
3434fn unpack8<const B: usize>(data: &[u8]) -> [i32; 8] {
3435    let mut acc = 0u64;
3436    for i in 0..B {
3437        acc = (acc << 8) | data[i] as u64;
3438    }
3439    let mask = (1u64 << B) - 1;
3440    let mut out = [0i32; 8];
3441    for (k, o) in out.iter_mut().enumerate() {
3442        *o = ((acc >> ((7 - k) * B)) & mask) as i32;
3443    }
3444    out
3445}
3446
3447/// Fused vbit matvec straight from the mapped bytes (spec §3, P13
3448/// FIG.3): [u8 bits: rows][f16 scales: rows·cols/32][bit-packed rows,
3449/// MSB-first, byte-padded]. Row data offsets are precomputed at load
3450/// (`vbit_row_offsets`) — the per-call prefix scan was O(rows) pure
3451/// overhead on every matvec.
3452#[allow(clippy::too_many_arguments)]
3453fn vbitmatvec(
3454    bytes: &[u8],
3455    offsets: &[usize],
3456    x: &[f32],
3457    rows: usize,
3458    cols: usize,
3459    out: &mut [f32],
3460    pool: Option<&Pool>,
3461) {
3462    debug_assert_eq!(out.len(), rows);
3463    debug_assert_eq!(offsets.len(), rows + 1);
3464
3465    // SDOT path: unpack the row to centered i8 once, then per-group
3466    // int8 dot against the quantized activations — same A8W8 contract
3467    // as q8 (bounded noise; CMF_SDOT=0 keeps the exact scalar path).
3468    if a8w8_enabled() {
3469        let act = split_act(x);
3470        let out_addr = SendMut(out.as_mut_ptr());
3471        let run = move |start: usize, end: usize| {
3472            vbit_range_a8w8(bytes, offsets, x, &act, rows, cols, out_addr, start, end)
3473        };
3474        dispatch_rows(pool, rows, &run);
3475        return;
3476    }
3477
3478    let out_addr = SendMut(out.as_mut_ptr());
3479    let run = move |start: usize, end: usize| {
3480        vbit_range_f32(bytes, offsets, x, rows, cols, out_addr, start, end)
3481    };
3482    dispatch_rows(pool, rows, &run);
3483}
3484
3485/// One vbit row range via the A8W8 int8 path — kernel body of
3486/// `vbitmatvec`, extracted so multi-matrix jobs can drive it for
3487/// several tensors in one dispatch (b=8 rows go exact f32).
3488#[allow(clippy::too_many_arguments)]
3489fn vbit_range_a8w8(
3490    bytes: &[u8],
3491    offsets: &[usize],
3492    x: &[f32],
3493    act: &SplitAct,
3494    rows: usize,
3495    cols: usize,
3496    out: SendMut,
3497    start: usize,
3498    end: usize,
3499) {
3500    let ng = cols / GROUP_SIZE;
3501    let bits = &bytes[..rows];
3502    let sc_off = rows;
3503    let row_dot = |r: usize| -> f32 {
3504        let b = bits[r] as usize;
3505        let l = (1i32 << (b - 1)) - 1;
3506        let mask = (1u64 << b) - 1;
3507        let data = &bytes[offsets[r]..offsets[r + 1]];
3508        if b == 8 {
3509            // u−L reaches 128 → does not fit i8; exact f32 path.
3510            let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
3511            let mut dot = 0f32;
3512            for g in 0..ng {
3513                let so = (r * ng + g) * 2;
3514                let sgf = f16_to_f32(u16::from_le_bytes([
3515                    bytes[sc_off + so],
3516                    bytes[sc_off + so + 1],
3517                ]));
3518                let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
3519                let mut gd = 0f32;
3520                for &xv in xg.iter() {
3521                    if nbits < 8 {
3522                        acc = (acc << 8) | data[idx] as u64;
3523                        idx += 1;
3524                        nbits += 8;
3525                    }
3526                    let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
3527                    nbits -= 8;
3528                    gd += (u - l) as f32 * xv;
3529                }
3530                dot += gd * sgf;
3531            }
3532            return dot;
3533        }
3534        // Per-worker scratch: this closure runs for every row of the
3535        // tensor (lm_head ≈ 150k rows/token) — a heap allocation per
3536        // row was measurable pure overhead.
3537        thread_local! {
3538            static VBIT_SCRATCH: std::cell::RefCell<Vec<u8>> =
3539                const { std::cell::RefCell::new(Vec::new()) };
3540        }
3541        #[inline(always)]
3542        fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
3543            for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
3544                let u = unpack8::<B>(&data[blk * B..]);
3545                for k in 0..8 {
3546                    chunk[k] = (u[k] - l) as i8 as u8;
3547                }
3548            }
3549        }
3550        let _ = mask;
3551        VBIT_SCRATCH.with(|scratch| {
3552            let mut buf = scratch.borrow_mut();
3553            buf.resize(cols, 0);
3554            match b {
3555                3 => fill::<3>(data, l, &mut buf),
3556                4 => vbit_fill4(data, &mut buf),
3557                5 => fill::<5>(data, l, &mut buf),
3558                6 => fill::<6>(data, l, &mut buf),
3559                _ => unreachable!(),
3560            }
3561            let mut dot = 0f32;
3562            for g in 0..ng {
3563                let so = (r * ng + g) * 2;
3564                let s = f16_to_f32(u16::from_le_bytes([
3565                    bytes[sc_off + so],
3566                    bytes[sc_off + so + 1],
3567                ]));
3568                let d = dot_i8_i8(
3569                    &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
3570                    &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
3571                ) as f32
3572                    * act.sx;
3573                dot += d * s;
3574            }
3575            for &(j, xv) in &act.outliers {
3576                let so = (r * ng + j / GROUP_SIZE) * 2;
3577                let s = f16_to_f32(u16::from_le_bytes([
3578                    bytes[sc_off + so],
3579                    bytes[sc_off + so + 1],
3580                ]));
3581                // xq is zeroed at outlier slots — add the exact term.
3582                dot += (buf[j] as i8) as f32 * s * xv;
3583            }
3584            dot
3585        })
3586    };
3587    for r in start..end {
3588        // SAFETY: disjoint row ranges per worker.
3589        unsafe { *out.at(r) = row_dot(r) };
3590    }
3591}
3592
3593/// Exact scalar vbit row range (same extraction, non-SDOT path).
3594#[allow(clippy::too_many_arguments)]
3595fn vbit_range_f32(
3596    bytes: &[u8],
3597    offsets: &[usize],
3598    x: &[f32],
3599    rows: usize,
3600    cols: usize,
3601    out: SendMut,
3602    start: usize,
3603    end: usize,
3604) {
3605    let ng = cols / GROUP_SIZE;
3606    let bits = &bytes[..rows];
3607    let sc_off = rows;
3608    // Per-bit-width specialized inner loops: the compiler unrolls the
3609    // constant shifts (the generic bit-buffer loop was branch-bound —
3610    // 5.6 vs 13.2 tok/s q4 on the 0.8B).
3611    #[inline(always)]
3612    fn dot_row<const B: usize>(
3613        data: &[u8],
3614        bytes: &[u8],
3615        sc_off: usize,
3616        r: usize,
3617        ng: usize,
3618        x: &[f32],
3619    ) -> f32 {
3620        let l = ((1i32 << (B - 1)) - 1) as f32;
3621        let gbytes = GROUP_SIZE * B / 8;
3622        let mut dot = 0f32;
3623        for g in 0..ng {
3624            let so = (r * ng + g) * 2;
3625            let s = f16_to_f32(u16::from_le_bytes([
3626                bytes[sc_off + so],
3627                bytes[sc_off + so + 1],
3628            ]));
3629            let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
3630            let gd0 = &data[g * gbytes..(g + 1) * gbytes];
3631            let mut gd = 0f32;
3632            for blk in 0..GROUP_SIZE / 8 {
3633                let u = unpack8::<B>(&gd0[blk * B..]);
3634                let xb = &xg[blk * 8..blk * 8 + 8];
3635                for k in 0..8 {
3636                    gd += (u[k] as f32 - l) * xb[k];
3637                }
3638            }
3639            dot += gd * s;
3640        }
3641        dot
3642    }
3643    for r in start..end {
3644        let data = &bytes[offsets[r]..offsets[r + 1]];
3645        let v = match bits[r] {
3646            3 => dot_row::<3>(data, bytes, sc_off, r, ng, x),
3647            4 => dot_row::<4>(data, bytes, sc_off, r, ng, x),
3648            5 => dot_row::<5>(data, bytes, sc_off, r, ng, x),
3649            6 => dot_row::<6>(data, bytes, sc_off, r, ng, x),
3650            8 => dot_row::<8>(data, bytes, sc_off, r, ng, x),
3651            b => unreachable!("vbit bit-width {b} (validated at load)"),
3652        };
3653        // SAFETY: disjoint row ranges per worker.
3654        unsafe { *out.at(r) = v };
3655    }
3656}
3657
3658/// Fused two-input vbit matvec: each row is unpacked from the mmap ONCE
3659/// and dotted against BOTH activations (MTP verify / pair prefill used
3660/// to run two full matvecs — double weight traffic and double unpack).
3661/// Per-input math is identical to `vbitmatvec` → same accuracy contract.
3662#[allow(clippy::too_many_arguments)]
3663fn vbitmatvec2(
3664    bytes: &[u8],
3665    offsets: &[usize],
3666    x1: &[f32],
3667    x2: &[f32],
3668    rows: usize,
3669    cols: usize,
3670    o1: &mut [f32],
3671    o2: &mut [f32],
3672    pool: Option<&Pool>,
3673) {
3674    debug_assert_eq!(o1.len(), rows);
3675    debug_assert_eq!(o2.len(), rows);
3676
3677    if a8w8_enabled() {
3678        let a1 = split_act(x1);
3679        let a2 = split_act(x2);
3680        let p1 = SendMut(o1.as_mut_ptr());
3681        let p2 = SendMut(o2.as_mut_ptr());
3682        let run = move |start: usize, end: usize| {
3683            vbit_range2_a8w8(
3684                bytes, offsets, x1, x2, &a1, &a2, rows, cols, p1, p2, start, end,
3685            )
3686        };
3687        dispatch_rows(pool, rows, &run);
3688        return;
3689    }
3690
3691    let p1 = SendMut(o1.as_mut_ptr());
3692    let p2 = SendMut(o2.as_mut_ptr());
3693    let run = move |start: usize, end: usize| {
3694        vbit_range2_f32(bytes, offsets, x1, x2, rows, cols, p1, p2, start, end)
3695    };
3696    dispatch_rows(pool, rows, &run);
3697}
3698
3699/// Two-input vbit row range via the A8W8 int8 path — kernel body of
3700/// `vbitmatvec2`, extracted for pair multi-matrix jobs (b=8 rows go
3701/// exact f32 for both lanes, bits streamed once).
3702#[allow(clippy::too_many_arguments)]
3703fn vbit_range2_a8w8(
3704    bytes: &[u8],
3705    offsets: &[usize],
3706    x1: &[f32],
3707    x2: &[f32],
3708    a1: &SplitAct,
3709    a2: &SplitAct,
3710    rows: usize,
3711    cols: usize,
3712    p1: SendMut,
3713    p2: SendMut,
3714    start: usize,
3715    end: usize,
3716) {
3717    let ng = cols / GROUP_SIZE;
3718    let bits = &bytes[..rows];
3719    let sc_off = rows;
3720    let row_dots = |r: usize| -> (f32, f32) {
3721        let b = bits[r] as usize;
3722        let l = (1i32 << (b - 1)) - 1;
3723        let data = &bytes[offsets[r]..offsets[r + 1]];
3724        if b == 8 {
3725            // u−L reaches 128 → does not fit i8; exact f32 path,
3726            // bits still streamed once for both lanes.
3727            let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
3728            let (mut d1, mut d2) = (0f32, 0f32);
3729            for g in 0..ng {
3730                let so = (r * ng + g) * 2;
3731                let sgf = f16_to_f32(u16::from_le_bytes([
3732                    bytes[sc_off + so],
3733                    bytes[sc_off + so + 1],
3734                ]));
3735                let (mut g1, mut g2) = (0f32, 0f32);
3736                for k in 0..GROUP_SIZE {
3737                    if nbits < 8 {
3738                        acc = (acc << 8) | data[idx] as u64;
3739                        idx += 1;
3740                        nbits += 8;
3741                    }
3742                    let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
3743                    nbits -= 8;
3744                    let w = (u - l) as f32;
3745                    g1 += w * x1[g * GROUP_SIZE + k];
3746                    g2 += w * x2[g * GROUP_SIZE + k];
3747                }
3748                d1 += g1 * sgf;
3749                d2 += g2 * sgf;
3750            }
3751            return (d1, d2);
3752        }
3753        thread_local! {
3754            static VBIT_SCRATCH2: std::cell::RefCell<Vec<u8>> =
3755                const { std::cell::RefCell::new(Vec::new()) };
3756        }
3757        #[inline(always)]
3758        fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
3759            for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
3760                let u = unpack8::<B>(&data[blk * B..]);
3761                for k in 0..8 {
3762                    chunk[k] = (u[k] - l) as i8 as u8;
3763                }
3764            }
3765        }
3766        VBIT_SCRATCH2.with(|scratch| {
3767            let mut buf = scratch.borrow_mut();
3768            buf.resize(cols, 0);
3769            match b {
3770                3 => fill::<3>(data, l, &mut buf),
3771                4 => vbit_fill4(data, &mut buf),
3772                5 => fill::<5>(data, l, &mut buf),
3773                6 => fill::<6>(data, l, &mut buf),
3774                _ => unreachable!(),
3775            }
3776            let (mut d1, mut d2) = (0f32, 0f32);
3777            for g in 0..ng {
3778                let so = (r * ng + g) * 2;
3779                let s = f16_to_f32(u16::from_le_bytes([
3780                    bytes[sc_off + so],
3781                    bytes[sc_off + so + 1],
3782                ]));
3783                let wg = &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
3784                let v1 = dot_i8_i8(wg, &a1.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a1.sx;
3785                let v2 = dot_i8_i8(wg, &a2.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a2.sx;
3786                d1 += v1 * s;
3787                d2 += v2 * s;
3788            }
3789            for &(j, xv) in &a1.outliers {
3790                let so = (r * ng + j / GROUP_SIZE) * 2;
3791                let s = f16_to_f32(u16::from_le_bytes([
3792                    bytes[sc_off + so],
3793                    bytes[sc_off + so + 1],
3794                ]));
3795                d1 += (buf[j] as i8) as f32 * s * xv;
3796            }
3797            for &(j, xv) in &a2.outliers {
3798                let so = (r * ng + j / GROUP_SIZE) * 2;
3799                let s = f16_to_f32(u16::from_le_bytes([
3800                    bytes[sc_off + so],
3801                    bytes[sc_off + so + 1],
3802                ]));
3803                d2 += (buf[j] as i8) as f32 * s * xv;
3804            }
3805            (d1, d2)
3806        })
3807    };
3808    for r in start..end {
3809        let (v1, v2) = row_dots(r);
3810        // SAFETY: disjoint row ranges per worker.
3811        unsafe {
3812            *p1.at(r) = v1;
3813            *p2.at(r) = v2;
3814        }
3815    }
3816}
3817
3818/// Two-input exact scalar vbit row range (same extraction) —
3819/// per-bit-width specialized, two accumulators per row; per-lane
3820/// accumulation order matches `vbitmatvec` exactly.
3821#[allow(clippy::too_many_arguments)]
3822fn vbit_range2_f32(
3823    bytes: &[u8],
3824    offsets: &[usize],
3825    x1: &[f32],
3826    x2: &[f32],
3827    rows: usize,
3828    cols: usize,
3829    p1: SendMut,
3830    p2: SendMut,
3831    start: usize,
3832    end: usize,
3833) {
3834    let ng = cols / GROUP_SIZE;
3835    let bits = &bytes[..rows];
3836    let sc_off = rows;
3837    #[inline(always)]
3838    #[allow(clippy::too_many_arguments)]
3839    fn dot_row2<const B: usize>(
3840        data: &[u8],
3841        bytes: &[u8],
3842        sc_off: usize,
3843        r: usize,
3844        ng: usize,
3845        x1: &[f32],
3846        x2: &[f32],
3847    ) -> (f32, f32) {
3848        let l = ((1i32 << (B - 1)) - 1) as f32;
3849        let gbytes = GROUP_SIZE * B / 8;
3850        let (mut d1, mut d2) = (0f32, 0f32);
3851        for g in 0..ng {
3852            let so = (r * ng + g) * 2;
3853            let s = f16_to_f32(u16::from_le_bytes([
3854                bytes[sc_off + so],
3855                bytes[sc_off + so + 1],
3856            ]));
3857            let x1g = &x1[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
3858            let x2g = &x2[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
3859            let gd0 = &data[g * gbytes..(g + 1) * gbytes];
3860            let (mut g1, mut g2) = (0f32, 0f32);
3861            for blk in 0..GROUP_SIZE / 8 {
3862                let u = unpack8::<B>(&gd0[blk * B..]);
3863                for k in 0..8 {
3864                    let w = u[k] as f32 - l;
3865                    g1 += w * x1g[blk * 8 + k];
3866                    g2 += w * x2g[blk * 8 + k];
3867                }
3868            }
3869            d1 += g1 * s;
3870            d2 += g2 * s;
3871        }
3872        (d1, d2)
3873    }
3874    for r in start..end {
3875        let data = &bytes[offsets[r]..offsets[r + 1]];
3876        let (v1, v2) = match bits[r] {
3877            3 => dot_row2::<3>(data, bytes, sc_off, r, ng, x1, x2),
3878            4 => dot_row2::<4>(data, bytes, sc_off, r, ng, x1, x2),
3879            5 => dot_row2::<5>(data, bytes, sc_off, r, ng, x1, x2),
3880            6 => dot_row2::<6>(data, bytes, sc_off, r, ng, x1, x2),
3881            8 => dot_row2::<8>(data, bytes, sc_off, r, ng, x1, x2),
3882            b => unreachable!("vbit bit-width {b} (validated at load)"),
3883        };
3884        // SAFETY: disjoint row ranges per worker.
3885        unsafe {
3886            *p1.at(r) = v1;
3887            *p2.at(r) = v2;
3888        }
3889    }
3890}
3891
3892// ───────────────────── q4_tiled kernels (§4.3) ─────────────────────
3893
3894/// One q4_tiled row dot on the A8W8 int8 path: per 32-group the tile
3895/// is ONE sequential read — [f16 scale][16B nibbles] — versus the two
3896/// distant streams of the split layout. Values/order identical to the
3897/// split kernels.
3898#[inline]
3899#[allow(unreachable_code)]
3900fn dot_q4t_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
3901    #[cfg(target_arch = "aarch64")]
3902    unsafe {
3903        return dot_q4t_row_sdot(bytes, r, gpr, xq);
3904    }
3905    #[cfg(target_arch = "x86_64")]
3906    unsafe {
3907        if vnni_tiles_enabled() {
3908            return dot_q4t_row_vnni(bytes, r, gpr, xq);
3909        }
3910        return dot_q4t_row_avx2(bytes, r, gpr, xq);
3911    }
3912    let mut acc = 0f32;
3913    for gi in 0..gpr {
3914        let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
3915        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
3916        let mut d = 0i32;
3917        for (k, &b) in tile[2..].iter().enumerate() {
3918            d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
3919                + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
3920        }
3921        acc += d as f32 * s;
3922    }
3923    acc
3924}
3925
3926#[cfg(target_arch = "aarch64")]
3927#[target_feature(enable = "neon,dotprod")]
3928unsafe fn dot_q4t_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
3929    // SAFETY: callers uphold slice-length contracts (18B tile per group,
3930    // xq.len() == gpr·GROUP_SIZE).
3931    unsafe {
3932        use core::arch::aarch64::*;
3933        use core::arch::asm;
3934        let lomask = vdupq_n_u8(0x0F);
3935        let eight = vdupq_n_s8(8);
3936        let mut acc = 0f32;
3937        for gi in 0..gpr {
3938            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
3939            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
3940            let b = vld1q_u8(t.add(2));
3941            let lo = vandq_u8(b, lomask);
3942            let hi = vshrq_n_u8::<4>(b);
3943            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
3944            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
3945            let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
3946            let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
3947            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
3948            asm!(
3949                "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
3950                "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
3951                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
3952                e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
3953                options(pure, nomem, nostack),
3954            );
3955            acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
3956        }
3957        acc
3958    }
3959}
3960
3961#[cfg(target_arch = "x86_64")]
3962#[target_feature(enable = "avx2")]
3963unsafe fn dot_q4t_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
3964    // SAFETY: see dot_q4t_row_sdot.
3965    unsafe {
3966        use core::arch::x86_64::*;
3967        let lomask = _mm_set1_epi8(0x0F);
3968        let eight = _mm256_set1_epi8(8);
3969        let ones = _mm256_set1_epi16(1);
3970        let mut acc = 0f32;
3971        for gi in 0..gpr {
3972            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
3973            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
3974            let b = _mm_loadu_si128(t.add(2) as *const __m128i);
3975            let lo = _mm_and_si128(b, lomask);
3976            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
3977            let w = _mm256_sub_epi8(
3978                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
3979                eight,
3980            );
3981            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
3982            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
3983            let d = _mm256_madd_epi16(p16, ones);
3984            let hi128 = _mm256_extracti128_si256::<1>(d);
3985            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
3986            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
3987            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
3988            acc += _mm_cvtsi128_si32(s32) as f32 * s;
3989        }
3990        acc
3991    }
3992}
3993
3994/// VNNI twin of `dot_q4t_row_avx2`: same unpack, `vpdpbusd` replaces
3995/// the maddubs+madd pair (see `dpbusd_hsum` — sums are bit-identical).
3996/// 256-bit VL encoding, so the VEX `vpsignb` stays usable.
3997#[cfg(target_arch = "x86_64")]
3998#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
3999unsafe fn dot_q4t_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
4000    // SAFETY: see dot_q4t_row_sdot.
4001    unsafe {
4002        use core::arch::x86_64::*;
4003        let lomask = _mm_set1_epi8(0x0F);
4004        let eight = _mm256_set1_epi8(8);
4005        let mut acc = 0f32;
4006        for gi in 0..gpr {
4007            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4008            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4009            let b = _mm_loadu_si128(t.add(2) as *const __m128i);
4010            let lo = _mm_and_si128(b, lomask);
4011            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
4012            let w = _mm256_sub_epi8(
4013                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4014                eight,
4015            );
4016            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4017            let d = dpbusd_hsum(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
4018            acc += d as f32 * s;
4019        }
4020        acc
4021    }
4022}
4023
4024/// One q4_tiled row against FOUR activation streams: the nibble unpack
4025/// and abs() happen once per group instead of once per (group,
4026/// activation) — the unpack is the dominant per-element cost of the
4027/// tiled format (roadmap P0 portable blocking, q4t leg).
4028#[cfg(target_arch = "x86_64")]
4029// `fma` is NOT implied by `avx2`: without it LLVM lowers _mm256_fmadd_ps
4030// to a libm call per lane — measured 2x slower than the reduction this
4031// kernel replaces. The runtime gate (`avx2_enabled`) already requires
4032// both features, so declaring it here is safe.
4033#[target_feature(enable = "avx2,fma")]
4034unsafe fn dot_q4t_row_1x4_avx2(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
4035    // SAFETY: callers uphold the 18B-tile and xq-length contracts.
4036    unsafe {
4037        use core::arch::x86_64::*;
4038        let lomask = _mm_set1_epi8(0x0F);
4039        let eight = _mm256_set1_epi8(8);
4040        let ones = _mm256_set1_epi16(1);
4041        // One f32 accumulator VECTOR per activation, reduced once at the
4042        // end. Folding each group's i32 lanes to a scalar inside the loop
4043        // costs an extracti128 + three shift/add + a movd — a cross-lane
4044        // dependency chain per (group, activation), 288 of them per row at
4045        // cols=2304. The per-group scale is what forces a float
4046        // accumulator; it does not force a horizontal sum.
4047        //
4048        // The four accumulators are NAMED, not an array: as `[__m256; 4]`
4049        // indexed by a loop variable LLVM keeps them in memory and every
4050        // group pays four 32-byte loads and stores. That alone made this
4051        // kernel 2x SLOWER than the per-group reduction it replaces
4052        // (measured on the EPYC box: 150 s vs 71 s for two 256² steps).
4053        let mut f0 = _mm256_setzero_ps();
4054        let mut f1 = _mm256_setzero_ps();
4055        let mut f2 = _mm256_setzero_ps();
4056        let mut f3 = _mm256_setzero_ps();
4057        for gi in 0..gpr {
4058            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4059            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4060            let sv = _mm256_set1_ps(s);
4061            let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
4062            let lo = _mm_and_si128(bb, lomask);
4063            let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
4064            let w = _mm256_sub_epi8(
4065                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4066                eight,
4067            );
4068            let aw = _mm256_abs_epi8(w);
4069            let off = gi * GROUP_SIZE;
4070            let dot = |xq: &[i8]| {
4071                let x = _mm256_loadu_si256(xq.as_ptr().add(off) as *const __m256i);
4072                let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
4073                _mm256_cvtepi32_ps(_mm256_madd_epi16(p16, ones))
4074            };
4075            f0 = _mm256_fmadd_ps(dot(xs[0]), sv, f0);
4076            f1 = _mm256_fmadd_ps(dot(xs[1]), sv, f1);
4077            f2 = _mm256_fmadd_ps(dot(xs[2]), sv, f2);
4078            f3 = _mm256_fmadd_ps(dot(xs[3]), sv, f3);
4079        }
4080        [
4081            hsum256_ps(f0),
4082            hsum256_ps(f1),
4083            hsum256_ps(f2),
4084            hsum256_ps(f3),
4085        ]
4086    }
4087}
4088
4089/// Horizontal sum of eight f32 lanes — the one cross-lane reduction the
4090/// blocked kernels pay, once per row instead of once per group.
4091#[cfg(target_arch = "x86_64")]
4092#[target_feature(enable = "avx2")]
4093#[inline]
4094unsafe fn hsum256_ps(v: core::arch::x86_64::__m256) -> f32 {
4095    // SAFETY: pure register arithmetic on the caller's vector.
4096    unsafe {
4097        use core::arch::x86_64::*;
4098        let hi = _mm256_extractf128_ps::<1>(v);
4099        let s = _mm_add_ps(_mm256_castps256_ps128(v), hi);
4100        let s = _mm_add_ps(s, _mm_movehl_ps(s, s));
4101        let s = _mm_add_ss(s, _mm_shuffle_ps::<0x55>(s, s));
4102        _mm_cvtss_f32(s)
4103    }
4104}
4105
4106/// VNNI twin of `dot_q4t_row_1x4_avx2` (see `dpbusd_hsum`).
4107#[cfg(target_arch = "x86_64")]
4108#[target_feature(enable = "avx2,fma,avx512f,avx512bw,avx512vl,avx512vnni")]
4109unsafe fn dot_q4t_row_1x4_vnni(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
4110    // SAFETY: callers uphold the 18B-tile and xq-length contracts.
4111    unsafe {
4112        use core::arch::x86_64::*;
4113        let lomask = _mm_set1_epi8(0x0F);
4114        let eight = _mm256_set1_epi8(8);
4115        // Same shape as the AVX2 twin: accumulate in f32 vectors and pay
4116        // one cross-lane reduction per row, not per (group, activation).
4117        let mut f0 = _mm256_setzero_ps();
4118        let mut f1 = _mm256_setzero_ps();
4119        let mut f2 = _mm256_setzero_ps();
4120        let mut f3 = _mm256_setzero_ps();
4121        for gi in 0..gpr {
4122            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4123            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4124            let sv = _mm256_set1_ps(s);
4125            let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
4126            let lo = _mm_and_si128(bb, lomask);
4127            let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
4128            let w = _mm256_sub_epi8(
4129                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4130                eight,
4131            );
4132            let aw = _mm256_abs_epi8(w);
4133            let off = gi * GROUP_SIZE;
4134            let dot = |xq: &[i8]| {
4135                let x = _mm256_loadu_si256(xq.as_ptr().add(off) as *const __m256i);
4136                _mm256_cvtepi32_ps(_mm256_dpbusd_epi32(
4137                    _mm256_setzero_si256(),
4138                    aw,
4139                    _mm256_sign_epi8(x, w),
4140                ))
4141            };
4142            f0 = _mm256_fmadd_ps(dot(xs[0]), sv, f0);
4143            f1 = _mm256_fmadd_ps(dot(xs[1]), sv, f1);
4144            f2 = _mm256_fmadd_ps(dot(xs[2]), sv, f2);
4145            f3 = _mm256_fmadd_ps(dot(xs[3]), sv, f3);
4146        }
4147        let acc = [
4148            hsum256_ps(f0),
4149            hsum256_ps(f1),
4150            hsum256_ps(f2),
4151            hsum256_ps(f3),
4152        ];
4153        acc
4154    }
4155}
4156
4157/// ARM twin of `dot_q4t_row_1x4_avx2`: one nibble unpack per group
4158/// serves FOUR activation streams. Per stream the group order and f32
4159/// accumulation match `dot_q4t_row_sdot` exactly — batch == matvec
4160/// bit-for-bit.
4161#[cfg(target_arch = "aarch64")]
4162#[target_feature(enable = "neon,dotprod")]
4163unsafe fn dot_q4t_row_1x4_sdot(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
4164    // SAFETY: callers uphold the 18B-tile and xq-length contracts.
4165    unsafe {
4166        use core::arch::aarch64::*;
4167        use core::arch::asm;
4168        let lomask = vdupq_n_u8(0x0F);
4169        let eight = vdupq_n_s8(8);
4170        let mut acc = [0f32; 4];
4171        for gi in 0..gpr {
4172            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4173            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4174            let b = vld1q_u8(t.add(2));
4175            let lo = vandq_u8(b, lomask);
4176            let hi = vshrq_n_u8::<4>(b);
4177            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
4178            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
4179            for (k, xq) in xs.iter().enumerate() {
4180                let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
4181                let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
4182                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
4183                asm!(
4184                    "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
4185                    "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
4186                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4187                    e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
4188                    options(pure, nomem, nostack),
4189                );
4190                acc[k] += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
4191            }
4192        }
4193        acc
4194    }
4195}
4196
4197/// Exact-term correction for A8W8 outliers on a tiled row.
4198#[inline]
4199fn q4t_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
4200    let gi = j / GROUP_SIZE;
4201    let k = j % GROUP_SIZE;
4202    let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
4203    let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
4204    let byte = tile[2 + k / 2];
4205    let nib = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
4206    ((nib as i32 - 8) as f32, s)
4207}
4208
4209/// Exact scalar q4_tiled row (CMF_SDOT=0 contract) — same pairwise
4210/// accumulation shape as `q4_range_f32`.
4211#[inline]
4212fn q4t_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
4213    let mut acc = 0f32;
4214    for gi in 0..gpr {
4215        let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
4216        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
4217        let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
4218        let mut ga = 0f32;
4219        for (k, &b) in tile[2..].iter().enumerate() {
4220            ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
4221                + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
4222        }
4223        acc += ga * s;
4224    }
4225    acc
4226}
4227
4228/// Split view of a `q4tp` payload. The three planes are resolved once per
4229/// matvec instead of per row — `q4tp_sections` is cheap, but doing it inside
4230/// the row loop would put a division on the hot path for nothing.
4231struct Q4tpView<'a> {
4232    nib: &'a [u8],
4233    params: &'a [u8],
4234    codes: &'a [u8],
4235    stride: usize,
4236    /// q2tp reads the ladder with rung 0 = exact zero.
4237    zero_rung: bool,
4238}
4239
4240impl<'a> Q4tpView<'a> {
4241    fn new(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
4242        let (params_off, codes_off, stride) = q4tp_sections(rows, cols);
4243        Self {
4244            nib: &bytes[..params_off],
4245            params: &bytes[params_off..codes_off],
4246            codes: &bytes[codes_off..],
4247            stride,
4248            zero_rung: false,
4249        }
4250    }
4251
4252    /// The q2tp view: identical params/codes planes, 8 B weight chunks.
4253    fn new_q2(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
4254        let (params_off, codes_off, stride) = q2tp_sections(rows, cols);
4255        Self {
4256            nib: &bytes[..params_off],
4257            params: &bytes[params_off..codes_off],
4258            codes: &bytes[codes_off..],
4259            stride,
4260            zero_rung: true,
4261        }
4262    }
4263
4264    /// Expand row `r`'s per-tile scales into `out` (length `gpr`).
4265    ///
4266    /// Doing this once per row — rather than decoding a 5-bit code inside the
4267    /// tile loop — is what makes the format free at runtime. Random access to
4268    /// a packed 5-bit field costs a division, two bounds checks and a branch;
4269    /// the tile's actual work is two `sdot`s, so per-tile decoding dominated
4270    /// the kernel and cost 5x (measured: 1.4 vs 6.9 tok/s on Nanbeige-3B).
4271    /// Walking the plane sequentially with a bit accumulator is ~3 ops.
4272    /// Eight 5-bit codes are exactly five bytes, so a whole group of
4273    /// eight decodes from one little-endian word at fixed shifts. The
4274    /// bit-accumulator this replaces carried a data-dependent `while
4275    /// have < 5` refill whose branch sat in the innermost loop of every
4276    /// q4tp row; a decode profile put this function above the dot
4277    /// products it feeds. Same bitstream, same codes — just no branch
4278    /// and eight independent extractions.
4279    #[inline]
4280    fn scales_into(&self, r: usize, gpr: usize, out: &mut [f32]) {
4281        let tab = if self.zero_rung {
4282            q2tp_ladder(self.params, r)
4283        } else {
4284            q4tp_ladder(self.params, r)
4285        };
4286        let codes = &self.codes[r * self.stride..(r + 1) * self.stride];
4287        let out = &mut out[..gpr];
4288        let mut chunks = out.chunks_exact_mut(8);
4289        let mut ci = 0usize;
4290        for c in &mut chunks {
4291            let w = u64::from(codes[ci])
4292                | u64::from(codes[ci + 1]) << 8
4293                | u64::from(codes[ci + 2]) << 16
4294                | u64::from(codes[ci + 3]) << 24
4295                | u64::from(codes[ci + 4]) << 32;
4296            for (k, o) in c.iter_mut().enumerate() {
4297                *o = tab[((w >> (5 * k)) & 31) as usize];
4298            }
4299            ci += 5;
4300        }
4301        // Fewer than eight codes left: the shared total accessor, which
4302        // tolerates a 5-bit field whose spill byte is past the stride.
4303        let tail = &codes[ci..];
4304        for (k, o) in chunks.into_remainder().iter_mut().enumerate() {
4305            *o = tab[q4tp_code(tail, k)];
4306        }
4307    }
4308}
4309
4310#[inline]
4311fn dot_q4tp_row_i8(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
4312    #[cfg(target_arch = "aarch64")]
4313    unsafe {
4314        return dot_q4tp_row_sdot(nib, r, gpr, xq, scales);
4315    }
4316    #[cfg(target_arch = "x86_64")]
4317    unsafe {
4318        if vnni_tiles_enabled() {
4319            return dot_q4tp_row_vnni(nib, r, gpr, xq, scales);
4320        }
4321        return dot_q4tp_row_avx2(nib, r, gpr, xq, scales);
4322    }
4323    #[allow(unreachable_code)]
4324    {
4325        let mut acc = 0f32;
4326        for gi in 0..gpr {
4327            let tile = &nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
4328            let s = scales[gi];
4329            let mut d = 0i32;
4330            for (k, &b) in tile.iter().enumerate() {
4331                d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
4332                    + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
4333            }
4334            acc += d as f32 * s;
4335        }
4336        acc
4337    }
4338}
4339
4340/// q4tp twin of `dot_q4t_row_sdot`: identical nibble math, but the tile
4341/// stride is 16 B (no inline scale) and the scale is a ladder lookup.
4342#[cfg(target_arch = "aarch64")]
4343#[target_feature(enable = "neon,dotprod")]
4344unsafe fn dot_q4tp_row_sdot(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
4345    // SAFETY: callers uphold slice-length contracts (16B tile per group,
4346    // xq.len() == gpr·GROUP_SIZE, codes covering gpr 5-bit fields).
4347    unsafe {
4348        use core::arch::aarch64::*;
4349        use core::arch::asm;
4350        let lomask = vdupq_n_u8(0x0F);
4351        let eight = vdupq_n_s8(8);
4352        let mut acc = 0f32;
4353        for gi in 0..gpr {
4354            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
4355            let s = *scales.get_unchecked(gi);
4356            let b = vld1q_u8(t);
4357            let lo = vandq_u8(b, lomask);
4358            let hi = vshrq_n_u8::<4>(b);
4359            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
4360            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
4361            let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
4362            let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
4363            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
4364            asm!(
4365                "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
4366                "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
4367                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4368                e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
4369                options(pure, nomem, nostack),
4370            );
4371            acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
4372        }
4373        acc
4374    }
4375}
4376
4377#[cfg(target_arch = "x86_64")]
4378#[target_feature(enable = "avx2")]
4379unsafe fn dot_q4tp_row_avx2(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
4380    // SAFETY: see dot_q4tp_row_sdot.
4381    unsafe {
4382        use core::arch::x86_64::*;
4383        let lomask = _mm_set1_epi8(0x0F);
4384        let eight = _mm256_set1_epi8(8);
4385        let ones = _mm256_set1_epi16(1);
4386        let mut acc = 0f32;
4387        for gi in 0..gpr {
4388            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
4389            let s = *scales.get_unchecked(gi);
4390            let b = _mm_loadu_si128(t as *const __m128i);
4391            let lo = _mm_and_si128(b, lomask);
4392            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
4393            let w = _mm256_sub_epi8(
4394                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4395                eight,
4396            );
4397            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4398            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
4399            let d = _mm256_madd_epi16(p16, ones);
4400            let hi128 = _mm256_extracti128_si256::<1>(d);
4401            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
4402            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
4403            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
4404            acc += _mm_cvtsi128_si32(s32) as f32 * s;
4405        }
4406        acc
4407    }
4408}
4409
4410
4411/// VNNI twin of `dot_q4tp_row_avx2` (see `dot_q4t_row_vnni` for why the
4412/// 256-bit VL encoding is the one to use here).
4413#[cfg(target_arch = "x86_64")]
4414#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
4415unsafe fn dot_q4tp_row_vnni(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
4416    // SAFETY: see dot_q4tp_row_sdot.
4417    unsafe {
4418        use core::arch::x86_64::*;
4419        let lomask = _mm_set1_epi8(0x0F);
4420        let eight = _mm256_set1_epi8(8);
4421        let mut acc = 0f32;
4422        for gi in 0..gpr {
4423            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
4424            let s = *scales.get_unchecked(gi);
4425            let b = _mm_loadu_si128(t as *const __m128i);
4426            let lo = _mm_and_si128(b, lomask);
4427            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
4428            let w = _mm256_sub_epi8(
4429                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4430                eight,
4431            );
4432            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4433            acc += dpbusd_hsum(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w)) as f32 * s;
4434        }
4435        acc
4436    }
4437}
4438
4439/// Exact scalar q4tp row — the `CMF_SDOT=0` contract, same pairwise
4440/// accumulation shape as `q4t_row_exact`.
4441#[inline]
4442fn q4tp_row_exact(nib: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
4443    let mut acc = 0f32;
4444    for gi in 0..gpr {
4445        let tile = &nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
4446        let s = scales[gi];
4447        let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
4448        let mut ga = 0f32;
4449        for (k, &b) in tile.iter().enumerate() {
4450            ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
4451                + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
4452        }
4453        acc += ga * s;
4454    }
4455    acc
4456}
4457
4458/// Single weight of a q4tp tensor — the a8w8 outlier path, which restores
4459/// activation outliers at full precision after the int8 pass.
4460#[inline]
4461fn q4tp_outlier(nib: &[u8], r: usize, gpr: usize, j: usize, scales: &[f32]) -> (f32, f32) {
4462    let (gi, k) = (j / GROUP_SIZE, j % GROUP_SIZE);
4463    let byte = nib[(r * gpr + gi) * Q4TP_NIB + k / 2];
4464    let n = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
4465    ((n as i32 - 8) as f32, scales[gi])
4466}
4467
4468/// Fused q4tp matvec (dispatch mirrors `q4t_matvec`).
4469fn q4tp_matvec(
4470    bytes: &[u8],
4471    x: &[f32],
4472    rows: usize,
4473    cols: usize,
4474    out: &mut [f32],
4475    pool: Option<&Pool>,
4476) {
4477    debug_assert_eq!(out.len(), rows);
4478    let gpr = cols / GROUP_SIZE;
4479    let v = Q4tpView::new(bytes, rows, cols);
4480    let out_addr = SendMut(out.as_mut_ptr());
4481    if a8w8_enabled() {
4482        let act = split_act(x);
4483        let run = |start: usize, end: usize| {
4484            // One scratch row of scales per worker — borrowed, not minted.
4485            with_krow(gpr, |sc| {
4486                for r in start..end {
4487                    v.scales_into(r, gpr, sc);
4488                    let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, sc) * act.sx;
4489                    for &(j, xv) in &act.outliers {
4490                        let (w, s) = q4tp_outlier(v.nib, r, gpr, j, sc);
4491                        acc += w * s * xv;
4492                    }
4493                    // SAFETY: disjoint row ranges per worker.
4494                    unsafe { *out_addr.at(r) = acc };
4495                }
4496            })
4497        };
4498        dispatch_rows(pool, rows, &run);
4499        return;
4500    }
4501    let run = |start: usize, end: usize| {
4502        with_krow(gpr, |sc| {
4503            for r in start..end {
4504                v.scales_into(r, gpr, sc);
4505                // SAFETY: disjoint row ranges per worker.
4506                unsafe { *out_addr.at(r) = q4tp_row_exact(v.nib, r, gpr, x, sc) };
4507            }
4508        })
4509    };
4510    dispatch_rows(pool, rows, &run);
4511}
4512
4513/// Fused two-input q4tp matvec — the SwiGLU gate/up pair. Weights and the
4514/// row ladder are read once and spent on both activation streams.
4515#[allow(clippy::too_many_arguments)]
4516fn q4tp_matvec2(
4517    bytes: &[u8],
4518    x1: &[f32],
4519    x2: &[f32],
4520    rows: usize,
4521    cols: usize,
4522    o1: &mut [f32],
4523    o2: &mut [f32],
4524    pool: Option<&Pool>,
4525) {
4526    let gpr = cols / GROUP_SIZE;
4527    let v = Q4tpView::new(bytes, rows, cols);
4528    let (p1, p2) = (SendMut(o1.as_mut_ptr()), SendMut(o2.as_mut_ptr()));
4529    let run = |start: usize, end: usize| {
4530        let mut sc = vec![0f32; gpr];
4531        for r in start..end {
4532            v.scales_into(r, gpr, &mut sc);
4533            // SAFETY: disjoint row ranges per worker.
4534            unsafe {
4535                *p1.at(r) = q4tp_row_exact(v.nib, r, gpr, x1, &sc);
4536                *p2.at(r) = q4tp_row_exact(v.nib, r, gpr, x2, &sc);
4537            }
4538        }
4539    };
4540    dispatch_rows(pool, rows, &run);
4541}
4542
4543/// One q2tp outlier weight at column `j` of row `r`: the 2-bit code and
4544/// its group scale, mirrored on `q4tp_outlier`.
4545#[inline]
4546fn q2tp_outlier(chunks: &[u8], r: usize, gpr: usize, j: usize, scales: &[f32]) -> (f32, f32) {
4547    let (gi, k) = (j / GROUP_SIZE, j % GROUP_SIZE);
4548    let byte = chunks[(r * gpr + gi) * Q2TP_CHUNK + k / 4];
4549    let c = (byte >> (2 * (k % 4))) & 3;
4550    (c as f32 - 1.5, scales[gi])
4551}
4552
4553/// Integer dot of one q2tp row against pre-quantized activations:
4554/// Σ_g s_g · (Σ c·xq − 1.5·Σ xq). The half-integer grid (c − 1.5)
4555/// becomes exact integer math through the group sums — the same trick
4556/// every a8w8 kernel in this file rides. The codes decode into a
4557/// 32-byte scratch in natural order and the dot itself is the shared
4558/// SDOT primitive; elsewhere a scalar integer loop.
4559#[inline]
4560fn dot_q2tp_row_i8(
4561    chunks: &[u8],
4562    r: usize,
4563    gpr: usize,
4564    xq: &[i8],
4565    gsum: &[i32],
4566    scales: &[f32],
4567) -> f32 {
4568    let mut acc = 0f32;
4569    let base = r * gpr * Q2TP_CHUNK;
4570    #[cfg(not(target_arch = "aarch64"))]
4571    let mut codes = [0i8; GROUP_SIZE];
4572    for gi in 0..gpr {
4573        let ch = &chunks[base + gi * Q2TP_CHUNK..base + (gi + 1) * Q2TP_CHUNK];
4574        let xg = &xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
4575        #[cfg(target_arch = "aarch64")]
4576        // NEON: the byte's four 2-bit fields land in four lane vectors
4577        // (shift+mask), vld4 de-interleaves xq to match (xj[k] =
4578        // xq[4k+j]), widening MACs accumulate exactly in i32. A scalar
4579        // decode here cost as much as the dot it fed — the profile put
4580        // it at the top of the whole W2 decode.
4581        let dot = unsafe {
4582            use core::arch::aarch64::*;
4583            let b = vld1_u8(ch.as_ptr());
4584            let three = vdup_n_u8(3);
4585            let c0 = vreinterpret_s8_u8(vand_u8(b, three));
4586            let c1 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 2), three));
4587            let c2 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 4), three));
4588            let c3 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 6), three));
4589            let x4 = vld4_s8(xg.as_ptr());
4590            let mut acc4 = vdupq_n_s32(0);
4591            acc4 = vpadalq_s16(acc4, vmull_s8(c0, x4.0));
4592            acc4 = vpadalq_s16(acc4, vmull_s8(c1, x4.1));
4593            acc4 = vpadalq_s16(acc4, vmull_s8(c2, x4.2));
4594            acc4 = vpadalq_s16(acc4, vmull_s8(c3, x4.3));
4595            vaddvq_s32(acc4)
4596        };
4597        #[cfg(not(target_arch = "aarch64"))]
4598        let dot: i32 = {
4599            for (k, &b) in ch.iter().enumerate() {
4600                codes[k * 4] = (b & 3) as i8;
4601                codes[k * 4 + 1] = ((b >> 2) & 3) as i8;
4602                codes[k * 4 + 2] = ((b >> 4) & 3) as i8;
4603                codes[k * 4 + 3] = ((b >> 6) & 3) as i8;
4604            }
4605            codes
4606                .iter()
4607                .zip(xg)
4608                .map(|(&c, &x)| c as i32 * x as i32)
4609                .sum()
4610        };
4611        acc += scales[gi] * (dot as f32 - 1.5 * gsum[gi] as f32);
4612    }
4613    acc
4614}
4615
4616/// Exact f32 dot of one q2tp row: 2-bit fields LSB-first, (c − 1.5)·s.
4617/// Scalar on purpose — the 2-bit class targets the GPU graph; the CPU
4618/// path exists for parity gates and small-machine fallback.
4619fn q2tp_row_exact(chunks: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
4620    let mut acc = 0f32;
4621    for gi in 0..gpr {
4622        let ch = &chunks[(r * gpr + gi) * Q2TP_CHUNK..(r * gpr + gi + 1) * Q2TP_CHUNK];
4623        let s = scales[gi];
4624        let xb = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
4625        let mut g = 0f32;
4626        for (k, &b) in ch.iter().enumerate() {
4627            g += ((b & 3) as f32 - 1.5) * xb[k * 4]
4628                + (((b >> 2) & 3) as f32 - 1.5) * xb[k * 4 + 1]
4629                + (((b >> 4) & 3) as f32 - 1.5) * xb[k * 4 + 2]
4630                + (((b >> 6) & 3) as f32 - 1.5) * xb[k * 4 + 3];
4631        }
4632        acc += s * g;
4633    }
4634    acc
4635}
4636
4637fn q2tp_matvec(
4638    bytes: &[u8],
4639    x: &[f32],
4640    rows: usize,
4641    cols: usize,
4642    out: &mut [f32],
4643    pool: Option<&Pool>,
4644) {
4645    debug_assert_eq!(out.len(), rows);
4646    let gpr = cols / GROUP_SIZE;
4647    let v = Q4tpView::new_q2(bytes, rows, cols);
4648    let out_addr = SendMut(out.as_mut_ptr());
4649    // a8w8 fast path (CMF_SDOT=0 keeps the exact scalar walk): integer
4650    // code dots + group sums, exact outlier correction — the same
4651    // contract as every sibling kernel; measured 2-bit rows were the
4652    // only scalar holdout in the family.
4653    if a8w8_enabled() {
4654        let act = split_act(x);
4655        let gsum = q1_group_sums(&act.xq, gpr);
4656        let (act, gsum) = (&act, &gsum);
4657        let run = move |start: usize, end: usize| {
4658            with_krow(gpr, |sc| {
4659                for r in start..end {
4660                    v.scales_into(r, gpr, sc);
4661                    let mut acc = dot_q2tp_row_i8(v.nib, r, gpr, &act.xq, gsum, sc) * act.sx;
4662                    for &(j, xv) in &act.outliers {
4663                        let (w, s) = q2tp_outlier(v.nib, r, gpr, j, sc);
4664                        acc += w * s * xv;
4665                    }
4666                    // SAFETY: disjoint row ranges per worker.
4667                    unsafe { *out_addr.at(r) = acc };
4668                }
4669            })
4670        };
4671        dispatch_rows(pool, rows, &run);
4672        return;
4673    }
4674    let run = |start: usize, end: usize| {
4675        with_krow(gpr, |sc| {
4676            for r in start..end {
4677                v.scales_into(r, gpr, sc);
4678                // SAFETY: disjoint row ranges per worker.
4679                unsafe { *out_addr.at(r) = q2tp_row_exact(v.nib, r, gpr, x, sc) };
4680            }
4681        })
4682    };
4683    dispatch_rows(pool, rows, &run);
4684}
4685
4686/// Fused two-input q2tp matvec — the SwiGLU gate/up pair.
4687#[allow(clippy::too_many_arguments)]
4688fn q2tp_matvec2(
4689    bytes: &[u8],
4690    x1: &[f32],
4691    x2: &[f32],
4692    rows: usize,
4693    cols: usize,
4694    o1: &mut [f32],
4695    o2: &mut [f32],
4696    pool: Option<&Pool>,
4697) {
4698    let gpr = cols / GROUP_SIZE;
4699    let v = Q4tpView::new_q2(bytes, rows, cols);
4700    let (p1, p2) = (SendMut(o1.as_mut_ptr()), SendMut(o2.as_mut_ptr()));
4701    let run = |start: usize, end: usize| {
4702        let mut sc = vec![0f32; gpr];
4703        for r in start..end {
4704            v.scales_into(r, gpr, &mut sc);
4705            // SAFETY: disjoint row ranges per worker.
4706            unsafe {
4707                *p1.at(r) = q2tp_row_exact(v.nib, r, gpr, x1, &sc);
4708                *p2.at(r) = q2tp_row_exact(v.nib, r, gpr, x2, &sc);
4709            }
4710        }
4711    };
4712    dispatch_rows(pool, rows, &run);
4713}
4714
4715/// Batched q2tp matmat: scalar row kernel over every batch column. CPU
4716/// prefill only — decode rides the graph, so plain and correct beats
4717/// clever here.
4718/// Test doors into the host 2-bit kernels: the stand's heap corruption
4719/// pointed at down-shaped tensors, and the private fns need a way to be
4720/// held to a reference without a model file around them.
4721pub fn q2tp_matvec_for_test(bytes: &[u8], x: &[f32], rows: usize, cols: usize, out: &mut [f32]) {
4722    // The facade IS the reference: encoder oracles hold requant output
4723    // to the exact scalar walk. The production dispatch may take the i8
4724    // fast path, whose error scale is the ACTIVATIONS' — a different
4725    // claim than the encoder correctness these tests pin.
4726    let gpr = cols / GROUP_SIZE;
4727    let v = Q4tpView::new_q2(bytes, rows, cols);
4728    with_krow(gpr, |sc| {
4729        for r in 0..rows {
4730            v.scales_into(r, gpr, sc);
4731            out[r] = q2tp_row_exact(v.nib, r, gpr, x, sc);
4732        }
4733    });
4734}
4735
4736pub fn q2tp_matmat_for_test(
4737    bytes: &[u8],
4738    xs_all: &[f32],
4739    b: usize,
4740    rows: usize,
4741    cols: usize,
4742    out: &mut [f32],
4743) {
4744    q2tp_matmat(bytes, xs_all, b, rows, cols, out, None);
4745}
4746
4747fn q2tp_matmat(
4748    bytes: &[u8],
4749    xs_all: &[f32],
4750    b: usize,
4751    rows: usize,
4752    cols: usize,
4753    out: &mut [f32],
4754    pool: Option<&Pool>,
4755) {
4756    debug_assert_eq!(out.len(), b * rows);
4757    let gpr = cols / GROUP_SIZE;
4758    let v = Q4tpView::new_q2(bytes, rows, cols);
4759    let out_addr = SendMut(out.as_mut_ptr());
4760    let run = |start: usize, end: usize| {
4761        let mut sc = vec![0f32; gpr];
4762        for r in start..end {
4763            v.scales_into(r, gpr, &mut sc);
4764            for bi in 0..b {
4765                let x = &xs_all[bi * cols..(bi + 1) * cols];
4766                // SAFETY: disjoint row ranges per worker.
4767                unsafe { *out_addr.at(bi * rows + r) = q2tp_row_exact(v.nib, r, gpr, x, &sc) };
4768            }
4769        }
4770    };
4771    dispatch_rows(pool, rows, &run);
4772}
4773
4774/// The pre-vectorised shape, kept for A/B (`CMF_Q4TP_V1=1`): the
4775/// horizontal add lands once per group per column instead of once per
4776/// row. Same weights, same activations — only the reduction differs.
4777#[cfg(target_arch = "aarch64")]
4778#[target_feature(enable = "neon,dotprod")]
4779unsafe fn dot_q4tp_row_1x4_sdot_v1(
4780    nib: &[u8],
4781    r: usize,
4782    gpr: usize,
4783    xs: [&[i8]; 4],
4784    scales: &[f32],
4785) -> [f32; 4] {
4786    unsafe {
4787        use core::arch::aarch64::*;
4788        use core::arch::asm;
4789        let lomask = vdupq_n_u8(0x0F);
4790        let eight = vdupq_n_s8(8);
4791        let (mut f0, mut f1, mut f2, mut f3) = (0f32, 0f32, 0f32, 0f32);
4792        for gi in 0..gpr {
4793            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
4794            let s = *scales.get_unchecked(gi);
4795            let bb = vld1q_u8(t);
4796            let lo = vandq_u8(bb, lomask);
4797            let hi = vshrq_n_u8::<4>(bb);
4798            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
4799            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
4800            let mut d = [0f32; 4];
4801            for (k, dk) in d.iter_mut().enumerate() {
4802                let x0 = vld1q_s8(xs[k].as_ptr().add(gi * GROUP_SIZE));
4803                let x1 = vld1q_s8(xs[k].as_ptr().add(gi * GROUP_SIZE + 16));
4804                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
4805                asm!(
4806                    "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
4807                    "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
4808                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4809                    e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
4810                    options(pure, nomem, nostack),
4811                );
4812                *dk = vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
4813            }
4814            f0 += d[0];
4815            f1 += d[1];
4816            f2 += d[2];
4817            f3 += d[3];
4818        }
4819        [f0, f1, f2, f3]
4820    }
4821}
4822
4823/// Which q4tp batch kernel to run: 1 = the previous one, 2 = the tuned
4824/// one, 0 = decide from the CPU. An atomic rather than a `OnceLock` so a
4825/// benchmark can alternate the two inside one process, where the machine's
4826/// mood — a shared box drifts ±25% between runs — is the same for both.
4827/// What the two mean is per-architecture: on x86 the blocked AVX-512 path
4828/// against the per-column one, on ARM the two reduction shapes.
4829#[allow(dead_code)]
4830static Q4TP_ALT: std::sync::atomic::AtomicU8 = std::sync::atomic::AtomicU8::new(0);
4831
4832/// Blocking pays on x86 only with 512-bit VNNI. With AVX2 alone, four
4833/// columns sharing an unpack still measured slower than the per-column
4834/// path (23.2 ms against 19.4 on a 48-thread EPYC), because that path
4835/// already dequantizes the row once — so the blocked kernel bought a
4836/// second unpack-free pass at the price of half the vector width.
4837#[cfg(target_arch = "x86_64")]
4838fn q4tp_blocked_x86() -> bool {
4839    match Q4TP_ALT.load(std::sync::atomic::Ordering::Relaxed) {
4840        1 => false,
4841        // A forced ON still asks the CPU. The switch exists so a bench can
4842        // pick a kernel, not so it can promise instructions the machine
4843        // does not have — CI caught that as a SIGILL on a runner without
4844        // AVX-512, where the parity test had turned the path on by hand.
4845        2 => avx512vnni_enabled(),
4846        // Deliberately not cached back into the switch: both gates below
4847        // hold their own `OnceLock`, and latching their answer here would
4848        // make a test's override outlive the test that set it.
4849        _ => blocked_enabled() && avx512vnni_enabled(),
4850    }
4851}
4852
4853/// `CMF_Q4TP_V1=1` picks the old reduction shape (A/B only).
4854#[cfg(target_arch = "aarch64")]
4855#[allow(dead_code)]
4856fn q4tp_v1() -> bool {
4857    match Q4TP_ALT.load(std::sync::atomic::Ordering::Relaxed) {
4858        1 => true,
4859        2 => false,
4860        _ => {
4861            static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4862            *ON.get_or_init(|| std::env::var("CMF_Q4TP_V1").is_ok_and(|v| v != "0"))
4863        }
4864    }
4865}
4866
4867/// Two weight rows against eight columns. The activation load is the
4868/// same for both rows, so it is paid once for twice the arithmetic, and
4869/// sixteen accumulator chains run where eight did — which is what a kernel
4870/// retiring 0.29 instructions a cycle is short of. Register pressure is
4871/// the limit: sixteen `zmm` accumulators, two weight tiles, one
4872/// activation, of thirty-two.
4873///
4874/// Four rows by four columns spends the same sixteen accumulators the
4875/// other way and measured worse — 1488 GFLOP/s against 1644 — so the
4876/// unpack, which four rows pay twice as often, costs more than the extra
4877/// sharing of one activation load buys.
4878#[cfg(target_arch = "x86_64")]
4879#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
4880unsafe fn dot_q4tp_2x8_avx512(
4881    nib: &[u8],
4882    r0: usize,
4883    gpr: usize,
4884    xs: [&[i8]; 8],
4885    sc0: &[f32],
4886    sc1: &[f32],
4887) -> [[f32; 8]; 2] {
4888    // SAFETY: as dot_q4tp_row_1x8_avx512, two adjacent rows at once; the
4889    // caller guarantees r0 + 1 < rows and the ISA.
4890    unsafe {
4891        use core::arch::x86_64::*;
4892        let lomask = _mm256_set1_epi8(0x0F);
4893        let eight = _mm256_set1_epi8(8);
4894        let zero = _mm512_setzero_si512();
4895        let mut v0 = [_mm512_setzero_ps(); 8];
4896        let mut v1 = [_mm512_setzero_ps(); 8];
4897        let pairs = gpr / 2;
4898        let unpack = |r: usize, gi: usize| -> (__m512i, __mmask64) {
4899            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
4900            let bb = _mm256_loadu_si256(t as *const __m256i);
4901            let lo = _mm256_and_si256(bb, lomask);
4902            let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
4903            let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
4904            let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
4905            let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
4906            let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
4907            (_mm512_abs_epi8(w), _mm512_movepi8_mask(w))
4908        };
4909        for gp in 0..pairs {
4910            let gi = gp * 2;
4911            let (wa0, neg0) = unpack(r0, gi);
4912            let (wa1, neg1) = unpack(r0 + 1, gi);
4913            let off = gi * GROUP_SIZE;
4914            let sv = |sc: &[f32]| {
4915                _mm512_insertf32x8::<1>(
4916                    _mm512_castps256_ps512(_mm256_set1_ps(*sc.get_unchecked(gi))),
4917                    _mm256_set1_ps(*sc.get_unchecked(gi + 1)),
4918                )
4919            };
4920            let s0 = sv(sc0);
4921            let s1 = sv(sc1);
4922            for k in 0..8 {
4923                let xv = _mm512_loadu_si512(xs[k].as_ptr().add(off) as *const __m512i);
4924                let d0 = _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(
4925                    zero,
4926                    wa0,
4927                    _mm512_mask_sub_epi8(xv, neg0, zero, xv),
4928                ));
4929                let d1 = _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(
4930                    zero,
4931                    wa1,
4932                    _mm512_mask_sub_epi8(xv, neg1, zero, xv),
4933                ));
4934                v0[k] = _mm512_fmadd_ps(d0, s0, v0[k]);
4935                v1[k] = _mm512_fmadd_ps(d1, s1, v1[k]);
4936            }
4937        }
4938        let mut acc = [[0f32; 8]; 2];
4939        for k in 0..8 {
4940            acc[0][k] = _mm512_reduce_add_ps(v0[k]);
4941            acc[1][k] = _mm512_reduce_add_ps(v1[k]);
4942        }
4943        if gpr % 2 == 1 {
4944            let off = (gpr - 1) * GROUP_SIZE;
4945            for j in off..off + GROUP_SIZE {
4946                let (w0, sa) = q4tp_outlier(nib, r0, gpr, j, sc0);
4947                let (w1, sb) = q4tp_outlier(nib, r0 + 1, gpr, j, sc1);
4948                for k in 0..8 {
4949                    let x = *xs[k].get_unchecked(j) as f32;
4950                    acc[0][k] += w0 * sa * x;
4951                    acc[1][k] += w1 * sb * x;
4952                }
4953            }
4954        }
4955        acc
4956    }
4957}
4958
4959/// The same, eight columns at a time. One unpack then feeds twice as many
4960/// activation streams, so a wide batch reads the weight tile half as
4961/// often; the price is eight accumulators live at once. Measured 9.0 ->
4962/// 8.3 ms at 9216x2304, b=296 on a 48-thread EPYC 9B45.
4963#[cfg(target_arch = "x86_64")]
4964#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
4965unsafe fn dot_q4tp_row_1x8_avx512(
4966    nib: &[u8],
4967    r: usize,
4968    gpr: usize,
4969    xs: [&[i8]; 8],
4970    scales: &[f32],
4971) -> [f32; 8] {
4972    // SAFETY: as dot_q4tp_row_1x4_avx2; caller guarantees the ISA.
4973    unsafe {
4974        use core::arch::x86_64::*;
4975        let lomask = _mm256_set1_epi8(0x0F);
4976        let eight = _mm256_set1_epi8(8);
4977        let zero = _mm512_setzero_si512();
4978        let (mut v0, mut v1, mut v2, mut v3) = (
4979            _mm512_setzero_ps(),
4980            _mm512_setzero_ps(),
4981            _mm512_setzero_ps(),
4982            _mm512_setzero_ps(),
4983        );
4984        let (mut v4, mut v5, mut v6, mut v7) = (
4985            _mm512_setzero_ps(),
4986            _mm512_setzero_ps(),
4987            _mm512_setzero_ps(),
4988            _mm512_setzero_ps(),
4989        );
4990        let pairs = gpr / 2;
4991        for gp in 0..pairs {
4992            let gi = gp * 2;
4993            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
4994            let bb = _mm256_loadu_si256(t as *const __m256i);
4995            let lo = _mm256_and_si256(bb, lomask);
4996            let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
4997            // `unpack` works per 128-bit lane, so the halves come out as
4998            // [A.lo, B.lo] and [A.hi, B.hi]; the shuffle reorders the four
4999            // 128-bit lanes into the weights' natural order, which is what
5000            // the straight activation load expects.
5001            let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
5002            let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
5003            let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
5004            let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
5005            let wabs = _mm512_abs_epi8(w);
5006            let neg = _mm512_movepi8_mask(w);
5007            let off = gi * GROUP_SIZE;
5008            let sv = _mm512_insertf32x8::<1>(
5009                _mm512_castps256_ps512(_mm256_set1_ps(*scales.get_unchecked(gi))),
5010                _mm256_set1_ps(*scales.get_unchecked(gi + 1)),
5011            );
5012            let dot = |x: &[i8]| -> __m512 {
5013                let xv = _mm512_loadu_si512(x.as_ptr().add(off) as *const __m512i);
5014                let sx = _mm512_mask_sub_epi8(xv, neg, zero, xv);
5015                _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(zero, wabs, sx))
5016            };
5017            v0 = _mm512_fmadd_ps(dot(xs[0]), sv, v0);
5018            v1 = _mm512_fmadd_ps(dot(xs[1]), sv, v1);
5019            v2 = _mm512_fmadd_ps(dot(xs[2]), sv, v2);
5020            v3 = _mm512_fmadd_ps(dot(xs[3]), sv, v3);
5021            v4 = _mm512_fmadd_ps(dot(xs[4]), sv, v4);
5022            v5 = _mm512_fmadd_ps(dot(xs[5]), sv, v5);
5023            v6 = _mm512_fmadd_ps(dot(xs[6]), sv, v6);
5024            v7 = _mm512_fmadd_ps(dot(xs[7]), sv, v7);
5025        }
5026        let mut acc = [
5027            _mm512_reduce_add_ps(v0),
5028            _mm512_reduce_add_ps(v1),
5029            _mm512_reduce_add_ps(v2),
5030            _mm512_reduce_add_ps(v3),
5031            _mm512_reduce_add_ps(v4),
5032            _mm512_reduce_add_ps(v5),
5033            _mm512_reduce_add_ps(v6),
5034            _mm512_reduce_add_ps(v7),
5035        ];
5036        // An odd group count leaves one group over; the narrow kernel
5037        // finishes it rather than the tail being a special case here.
5038        if gpr % 2 == 1 {
5039            let off = (gpr - 1) * GROUP_SIZE;
5040            for j in off..off + GROUP_SIZE {
5041                let (w, s) = q4tp_outlier(nib, r, gpr, j, scales);
5042                let ws = w * s;
5043                for k in 0..8 {
5044                    acc[k] += ws * *xs[k].get_unchecked(j) as f32;
5045                }
5046            }
5047        }
5048        acc
5049    }
5050}
5051
5052/// The same four columns, 512 bits wide. Two groups (64 weights) ride one
5053/// unpack and one `vpdpbusd`, where AVX2 needs two unpacks and four
5054/// `maddubs`/`madd` pairs — about 2.3x fewer instructions for the same
5055/// arithmetic. The two groups carry different scales, so the fma takes a
5056/// vector whose halves hold each group's scale rather than a broadcast.
5057///
5058/// There is no 512-bit `vpsignb`, so the activation's sign is applied by
5059/// negating under a mask taken from the weight's sign bits. That mask is
5060/// per-tile, so it is hoisted out of the column loop and the per-column
5061/// cost stays exactly one instruction, as with `sign_epi8`. Weights of
5062/// zero are not zeroed by the mask trick and do not need to be: their
5063/// magnitude is zero, so the product is.
5064#[cfg(target_arch = "x86_64")]
5065#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
5066unsafe fn dot_q4tp_row_1x4_avx512(
5067    nib: &[u8],
5068    r: usize,
5069    gpr: usize,
5070    xs: [&[i8]; 4],
5071    scales: &[f32],
5072) -> [f32; 4] {
5073    // SAFETY: as dot_q4tp_row_1x4_avx2; caller guarantees the ISA.
5074    unsafe {
5075        use core::arch::x86_64::*;
5076        let lomask = _mm256_set1_epi8(0x0F);
5077        let eight = _mm256_set1_epi8(8);
5078        let zero = _mm512_setzero_si512();
5079        let (mut v0, mut v1, mut v2, mut v3) = (
5080            _mm512_setzero_ps(),
5081            _mm512_setzero_ps(),
5082            _mm512_setzero_ps(),
5083            _mm512_setzero_ps(),
5084        );
5085        let pairs = gpr / 2;
5086        for gp in 0..pairs {
5087            let gi = gp * 2;
5088            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5089            let bb = _mm256_loadu_si256(t as *const __m256i);
5090            let lo = _mm256_and_si256(bb, lomask);
5091            let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
5092            // `unpack` works per 128-bit lane, so the halves come out as
5093            // [A.lo, B.lo] and [A.hi, B.hi]; the shuffle reorders the four
5094            // 128-bit lanes into the weights' natural order, which is what
5095            // the straight activation load expects.
5096            let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
5097            let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
5098            let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
5099            let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
5100            let wabs = _mm512_abs_epi8(w);
5101            let neg = _mm512_movepi8_mask(w);
5102            let off = gi * GROUP_SIZE;
5103            let sv = _mm512_insertf32x8::<1>(
5104                _mm512_castps256_ps512(_mm256_set1_ps(*scales.get_unchecked(gi))),
5105                _mm256_set1_ps(*scales.get_unchecked(gi + 1)),
5106            );
5107            let dot = |x: &[i8]| -> __m512 {
5108                let xv = _mm512_loadu_si512(x.as_ptr().add(off) as *const __m512i);
5109                let sx = _mm512_mask_sub_epi8(xv, neg, zero, xv);
5110                _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(zero, wabs, sx))
5111            };
5112            v0 = _mm512_fmadd_ps(dot(xs[0]), sv, v0);
5113            v1 = _mm512_fmadd_ps(dot(xs[1]), sv, v1);
5114            v2 = _mm512_fmadd_ps(dot(xs[2]), sv, v2);
5115            v3 = _mm512_fmadd_ps(dot(xs[3]), sv, v3);
5116        }
5117        let mut acc = [
5118            _mm512_reduce_add_ps(v0),
5119            _mm512_reduce_add_ps(v1),
5120            _mm512_reduce_add_ps(v2),
5121            _mm512_reduce_add_ps(v3),
5122        ];
5123        // An odd group count leaves one group over; the narrow kernel
5124        // finishes it rather than the tail being a special case here.
5125        if gpr % 2 == 1 {
5126            let off = (gpr - 1) * GROUP_SIZE;
5127            for j in off..off + GROUP_SIZE {
5128                let (w, s) = q4tp_outlier(nib, r, gpr, j, scales);
5129                let ws = w * s;
5130                for k in 0..4 {
5131                    acc[k] += ws * *xs[k].get_unchecked(j) as f32;
5132                }
5133            }
5134        }
5135        acc
5136    }
5137}
5138
5139/// Four batch columns against one q4tp row: the tile is unpacked ONCE and
5140/// spent on four activation streams, which is where a prefill batch stops
5141/// being weight-bandwidth-bound. Twin of `dot_q4t_row_1x4_sdot`.
5142#[cfg(target_arch = "aarch64")]
5143#[target_feature(enable = "neon,dotprod")]
5144unsafe fn dot_q4tp_row_1x4_sdot(
5145    nib: &[u8],
5146    r: usize,
5147    gpr: usize,
5148    xs: [&[i8]; 4],
5149    scales: &[f32],
5150) -> [f32; 4] {
5151    // SAFETY: see dot_q4tp_row_sdot; every xs[k] is gpr·GROUP_SIZE long.
5152    unsafe {
5153        use core::arch::aarch64::*;
5154        use core::arch::asm;
5155        let lomask = vdupq_n_u8(0x0F);
5156        let eight = vdupq_n_s8(8);
5157        // Named accumulators, NOT an array indexed by a loop variable: the
5158        // latter does not stay in registers (the same defect cost 2x in the
5159        // AVX2 q4t kernel and again in WGSL).
5160        //
5161        // They are VECTORS, and the horizontal add happens once at the end
5162        // instead of once per group per column. `vaddvq` is a cross-lane
5163        // reduction — with 72 groups and four columns the old shape paid
5164        // 288 of them per row, each one a dependency stall the pipeline
5165        // cannot hide, to save four float adds. The group's scale now
5166        // rides an fma into the lane accumulators, so the arithmetic per
5167        // group is one convert and one fma. Summation order changes (the
5168        // lanes carry independent partial sums), which is the same
5169        // round-off class the SDOT path already lives in — the strict
5170        // kernel (`CMF_SDOT=0`, what `cortiq ppl` runs) is unchanged and
5171        // stays the reference.
5172        let (mut v0, mut v1, mut v2, mut v3) = (
5173            vdupq_n_f32(0.0),
5174            vdupq_n_f32(0.0),
5175            vdupq_n_f32(0.0),
5176            vdupq_n_f32(0.0),
5177        );
5178        for gi in 0..gpr {
5179            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5180            let s = *scales.get_unchecked(gi);
5181            let bb = vld1q_u8(t);
5182            let lo = vandq_u8(bb, lomask);
5183            let hi = vshrq_n_u8::<4>(bb);
5184            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
5185            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
5186            let off = gi * GROUP_SIZE;
5187            let dot4 = |x: &[i8]| -> int32x4_t {
5188                let x0 = vld1q_s8(x.as_ptr().add(off));
5189                let x1 = vld1q_s8(x.as_ptr().add(off + 16));
5190                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5191                asm!(
5192                    "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
5193                    "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
5194                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5195                    e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
5196                    options(pure, nomem, nostack),
5197                );
5198                vaddq_s32(a0, a1)
5199            };
5200            v0 = vfmaq_n_f32(v0, vcvtq_f32_s32(dot4(xs[0])), s);
5201            v1 = vfmaq_n_f32(v1, vcvtq_f32_s32(dot4(xs[1])), s);
5202            v2 = vfmaq_n_f32(v2, vcvtq_f32_s32(dot4(xs[2])), s);
5203            v3 = vfmaq_n_f32(v3, vcvtq_f32_s32(dot4(xs[3])), s);
5204        }
5205        [
5206            vaddvq_f32(v0),
5207            vaddvq_f32(v1),
5208            vaddvq_f32(v2),
5209            vaddvq_f32(v3),
5210        ]
5211    }
5212}
5213
5214/// Fused q4tp matmat — the same three arms `q4t_matmat` has. Shipping only
5215/// the scalar one made Nanbeige-3B decode at 1.2 tok/s against q4t's 5.9:
5216/// the format was fine, the missing arms were the whole regression.
5217fn q4tp_matmat(
5218    bytes: &[u8],
5219    xs_all: &[f32],
5220    b: usize,
5221    rows: usize,
5222    cols: usize,
5223    out: &mut [f32],
5224    pool: Option<&Pool>,
5225) {
5226    debug_assert_eq!(out.len(), b * rows);
5227    let gpr = cols / GROUP_SIZE;
5228    let v = Q4tpView::new(bytes, rows, cols);
5229
5230    // Wide batches ride the AMX through a dequant-tile sgemm, as in q4t.
5231    #[cfg(target_os = "macos")]
5232    if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
5233        dequant_matmat_accel(
5234            &|r, dst| {
5235                let mut sc = [0f32; 32];
5236                let mut scv;
5237                let s: &[f32] = if gpr <= 32 {
5238                    v.scales_into(r, gpr, &mut sc);
5239                    &sc[..gpr]
5240                } else {
5241                    scv = vec![0f32; gpr];
5242                    v.scales_into(r, gpr, &mut scv);
5243                    &scv
5244                };
5245                for gi in 0..gpr {
5246                    let tile = &v.nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
5247                    for (k, &bb) in tile.iter().enumerate() {
5248                        dst[gi * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s[gi];
5249                        dst[gi * GROUP_SIZE + k * 2 + 1] =
5250                            (((bb >> 4) & 0x0F) as f32 - 8.0) * s[gi];
5251                    }
5252                }
5253            },
5254            xs_all,
5255            b,
5256            rows,
5257            cols,
5258            out,
5259            pool,
5260        );
5261        return;
5262    }
5263
5264    let out_addr = SendMut(out.as_mut_ptr());
5265    if a8w8_enabled() {
5266        let acts: Vec<SplitAct> = (0..b)
5267            .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
5268            .collect();
5269        let acts = &acts;
5270        #[cfg(target_arch = "aarch64")]
5271        let blocked_ok = sdot_enabled() && blocked_enabled();
5272        // x86 gets the same blocking: one tile unpack spent on four
5273        // columns. Without it every column re-decoded the row, which is
5274        // why a 48-core EPYC measured a sixth of an M4's per-core rate.
5275        // The gate is `avx2_enabled`, as in q4t — `sdot_enabled` answers
5276        // for ARM's dotprod and is hard-wired false everywhere else, so
5277        // asking it here left the whole blocked path unreachable on x86.
5278        #[cfg(target_arch = "x86_64")]
5279        let blocked_ok = q4tp_blocked_x86();
5280        #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
5281        let blocked_ok = false;
5282        // Columns are swept in panels that fit L2. Without this a
5283        // row-pair walks every activation in the batch — 4.8 MB at
5284        // 512x512 — and does it again for the next pair, so the whole
5285        // batch streams out of the shared cache once per row. Measured
5286        // 800 GB/s of it, flat across batch sizes, which is the signature
5287        // of a loop bound by traffic rather than by arithmetic. A panel of
5288        // 256 columns is 590 KB beside 221 KB of this worker's weights:
5289        // both stay resident and the batch crosses L3 once instead of
5290        // once per row.
5291        let panel_cols: usize = std::env::var("CMF_Q4TP_PANEL")
5292            .ok()
5293            .and_then(|v| v.parse().ok())
5294            .filter(|v| *v > 0)
5295            .unwrap_or(256);
5296        let run = |start: usize, end: usize| {
5297            for abase in (0..acts.len()).step_by(panel_cols) {
5298                let alen = (acts.len() - abase).min(panel_cols);
5299                let mut sc = vec![0f32; gpr];
5300                #[cfg(target_arch = "x86_64")]
5301                let mut r_lo = start;
5302                #[cfg(target_arch = "x86_64")]
5303                if blocked_ok && alen >= 8 {
5304                    let mut sc1 = vec![0f32; gpr];
5305                    while r_lo + 2 <= end {
5306                        v.scales_into(r_lo, gpr, &mut sc);
5307                        v.scales_into(r_lo + 1, gpr, &mut sc1);
5308                        let mut bi = 0usize;
5309                        while bi + 8 <= alen {
5310                            let xs = [
5311                                acts[abase + bi].xq.as_slice(),
5312                                acts[abase + bi + 1].xq.as_slice(),
5313                                acts[abase + bi + 2].xq.as_slice(),
5314                                acts[abase + bi + 3].xq.as_slice(),
5315                                acts[abase + bi + 4].xq.as_slice(),
5316                                acts[abase + bi + 5].xq.as_slice(),
5317                                acts[abase + bi + 6].xq.as_slice(),
5318                                acts[abase + bi + 7].xq.as_slice(),
5319                            ];
5320                            let d =
5321                                unsafe { dot_q4tp_2x8_avx512(v.nib, r_lo, gpr, xs, &sc, &sc1) };
5322                            for (row, dr, scr) in
5323                                [(r_lo, &d[0], &sc), (r_lo + 1, &d[1], &sc1)]
5324                            {
5325                                for k in 0..8 {
5326                                    let act = &acts[abase + bi + k];
5327                                    let mut acc = dr[k] * act.sx;
5328                                    for &(j, xv) in &act.outliers {
5329                                        let (w, s) = q4tp_outlier(v.nib, row, gpr, j, scr);
5330                                        acc += w * s * xv;
5331                                    }
5332                                    // SAFETY: disjoint (bi, r) cells per worker.
5333                                    unsafe { *out_addr.at((abase + bi + k) * rows + row) = acc };
5334                                }
5335                            }
5336                            bi += 8;
5337                        }
5338                        // Columns past the last group of eight, both rows —
5339                        // the same single-row kernel the tail below uses.
5340                        for row in [r_lo, r_lo + 1] {
5341                            let scr: &[f32] = if row == r_lo { &sc } else { &sc1 };
5342                            for b2 in bi..alen {
5343                                let act = &acts[abase + b2];
5344                                let xs4 = [
5345                                    act.xq.as_slice(),
5346                                    act.xq.as_slice(),
5347                                    act.xq.as_slice(),
5348                                    act.xq.as_slice(),
5349                                ];
5350                                let d =
5351                                    unsafe { dot_q4tp_row_1x4_avx512(v.nib, row, gpr, xs4, scr) };
5352                                let mut acc = d[0] * act.sx;
5353                                for &(j, xv) in &act.outliers {
5354                                    let (w, s) = q4tp_outlier(v.nib, row, gpr, j, scr);
5355                                    acc += w * s * xv;
5356                                }
5357                                // SAFETY: disjoint (bi, r) cells per worker.
5358                                unsafe { *out_addr.at((abase + b2) * rows + row) = acc };
5359                            }
5360                        }
5361                        r_lo += 2;
5362                    }
5363                }
5364                #[cfg(target_arch = "x86_64")]
5365                let row_start = r_lo;
5366                #[cfg(not(target_arch = "x86_64"))]
5367                let row_start = start;
5368                for r in row_start..end {
5369                    v.scales_into(r, gpr, &mut sc);
5370                    let mut bi = 0usize;
5371                    #[cfg(target_arch = "x86_64")]
5372                    if blocked_ok {
5373                        while bi + 8 <= alen {
5374                            let xs = [
5375                                acts[abase + bi].xq.as_slice(),
5376                                acts[abase + bi + 1].xq.as_slice(),
5377                                acts[abase + bi + 2].xq.as_slice(),
5378                                acts[abase + bi + 3].xq.as_slice(),
5379                                acts[abase + bi + 4].xq.as_slice(),
5380                                acts[abase + bi + 5].xq.as_slice(),
5381                                acts[abase + bi + 6].xq.as_slice(),
5382                                acts[abase + bi + 7].xq.as_slice(),
5383                            ];
5384                            let d = unsafe { dot_q4tp_row_1x8_avx512(v.nib, r, gpr, xs, &sc) };
5385                            for k in 0..8 {
5386                                let act = &acts[abase + bi + k];
5387                                let mut acc = d[k] * act.sx;
5388                                for &(j, xv) in &act.outliers {
5389                                    let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
5390                                    acc += w * s * xv;
5391                                }
5392                                // SAFETY: disjoint (bi, r) cells per worker.
5393                                unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
5394                            }
5395                            bi += 8;
5396                        }
5397                        while bi + 4 <= alen {
5398                            let xs = [
5399                                acts[abase + bi].xq.as_slice(),
5400                                acts[abase + bi + 1].xq.as_slice(),
5401                                acts[abase + bi + 2].xq.as_slice(),
5402                                acts[abase + bi + 3].xq.as_slice(),
5403                            ];
5404                            let d = unsafe { dot_q4tp_row_1x4_avx512(v.nib, r, gpr, xs, &sc) };
5405                            for k in 0..4 {
5406                                let act = &acts[abase + bi + k];
5407                                let mut acc = d[k] * act.sx;
5408                                for &(j, xv) in &act.outliers {
5409                                    let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
5410                                    acc += w * s * xv;
5411                                }
5412                                // SAFETY: disjoint (bi, r) cells per worker.
5413                                unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
5414                            }
5415                            bi += 4;
5416                        }
5417                    }
5418                    #[cfg(target_arch = "aarch64")]
5419                    if blocked_ok {
5420                        while bi + 4 <= alen {
5421                            let xs = [
5422                                acts[abase + bi].xq.as_slice(),
5423                                acts[abase + bi + 1].xq.as_slice(),
5424                                acts[abase + bi + 2].xq.as_slice(),
5425                                acts[abase + bi + 3].xq.as_slice(),
5426                            ];
5427                            let d = unsafe {
5428                                if q4tp_v1() {
5429                                    dot_q4tp_row_1x4_sdot_v1(v.nib, r, gpr, xs, &sc)
5430                                } else {
5431                                    dot_q4tp_row_1x4_sdot(v.nib, r, gpr, xs, &sc)
5432                                }
5433                            };
5434                            for k in 0..4 {
5435                                let act = &acts[abase + bi + k];
5436                                let mut acc = d[k] * act.sx;
5437                                for &(j, xv) in &act.outliers {
5438                                    let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
5439                                    acc += w * s * xv;
5440                                }
5441                                // SAFETY: disjoint (bi, r) cells per worker.
5442                                unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
5443                            }
5444                            bi += 4;
5445                        }
5446                    }
5447                    let _ = blocked_ok;
5448                    while bi < alen {
5449                        let act = &acts[abase + bi];
5450                        let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, &sc) * act.sx;
5451                        for &(j, xv) in &act.outliers {
5452                            let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
5453                            acc += w * s * xv;
5454                        }
5455                        // SAFETY: disjoint (bi, r) cells per worker range.
5456                        unsafe { *out_addr.at((abase + bi) * rows + r) = acc };
5457                        bi += 1;
5458                    }
5459                }
5460        
5461            }
5462        };
5463        dispatch_rows(pool, rows, &run);
5464        return;
5465    }
5466
5467    let run = |start: usize, end: usize| {
5468        let mut sc = vec![0f32; gpr];
5469        for r in start..end {
5470            v.scales_into(r, gpr, &mut sc);
5471            for bi in 0..b {
5472                let x = &xs_all[bi * cols..(bi + 1) * cols];
5473                // SAFETY: disjoint (bi, r) cells per worker range.
5474                unsafe { *out_addr.at(bi * rows + r) = q4tp_row_exact(v.nib, r, gpr, x, &sc) };
5475            }
5476        }
5477    };
5478    dispatch_rows(pool, rows, &run);
5479}
5480
5481/// Fused q4_tiled matvec (dispatch mirrors `q4matvec`).
5482fn q4t_matvec(
5483    bytes: &[u8],
5484    x: &[f32],
5485    rows: usize,
5486    cols: usize,
5487    out: &mut [f32],
5488    pool: Option<&Pool>,
5489) {
5490    debug_assert_eq!(out.len(), rows);
5491    let gpr = cols / GROUP_SIZE;
5492    let out_addr = SendMut(out.as_mut_ptr());
5493    if a8w8_enabled() {
5494        let act = split_act(x);
5495        let run = move |start: usize, end: usize| {
5496            for r in start..end {
5497                let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
5498                for &(j, xv) in &act.outliers {
5499                    let (w, s) = q4t_outlier(bytes, r, gpr, j);
5500                    acc += w * s * xv;
5501                }
5502                // SAFETY: disjoint row ranges per worker.
5503                unsafe { *out_addr.at(r) = acc };
5504            }
5505        };
5506        dispatch_rows(pool, rows, &run);
5507        return;
5508    }
5509    let run = move |start: usize, end: usize| {
5510        for r in start..end {
5511            // SAFETY: disjoint row ranges per worker.
5512            unsafe { *out_addr.at(r) = q4t_row_exact(bytes, r, gpr, x) };
5513        }
5514    };
5515    dispatch_rows(pool, rows, &run);
5516}
5517
5518/// Fused two-input q4_tiled matvec (weights read once per pair).
5519#[allow(clippy::too_many_arguments)]
5520fn q4t_matvec2(
5521    bytes: &[u8],
5522    x1: &[f32],
5523    x2: &[f32],
5524    rows: usize,
5525    cols: usize,
5526    o1: &mut [f32],
5527    o2: &mut [f32],
5528    pool: Option<&Pool>,
5529) {
5530    let gpr = cols / GROUP_SIZE;
5531    let p1 = SendMut(o1.as_mut_ptr());
5532    let p2 = SendMut(o2.as_mut_ptr());
5533    if a8w8_enabled() {
5534        let a1 = split_act(x1);
5535        let a2 = split_act(x2);
5536        let run = move |start: usize, end: usize| {
5537            for r in start..end {
5538                let mut v1 = dot_q4t_row_i8(bytes, r, gpr, &a1.xq) * a1.sx;
5539                let mut v2 = dot_q4t_row_i8(bytes, r, gpr, &a2.xq) * a2.sx;
5540                for &(j, xv) in &a1.outliers {
5541                    let (w, s) = q4t_outlier(bytes, r, gpr, j);
5542                    v1 += w * s * xv;
5543                }
5544                for &(j, xv) in &a2.outliers {
5545                    let (w, s) = q4t_outlier(bytes, r, gpr, j);
5546                    v2 += w * s * xv;
5547                }
5548                // SAFETY: disjoint row ranges per worker.
5549                unsafe {
5550                    *p1.at(r) = v1;
5551                    *p2.at(r) = v2;
5552                }
5553            }
5554        };
5555        dispatch_rows(pool, rows, &run);
5556        return;
5557    }
5558    let run = move |start: usize, end: usize| {
5559        for r in start..end {
5560            // SAFETY: disjoint row ranges per worker.
5561            unsafe {
5562                *p1.at(r) = q4t_row_exact(bytes, r, gpr, x1);
5563                *p2.at(r) = q4t_row_exact(bytes, r, gpr, x2);
5564            }
5565        }
5566    };
5567    dispatch_rows(pool, rows, &run);
5568}
5569
5570/// Batched q4_tiled matmat: each row's tiles stream once per microbatch.
5571#[allow(clippy::too_many_arguments)]
5572/// Prefill GEMM through Accelerate for group-quantized codecs: a
5573/// caller-supplied row dequantizer fills f32 tiles (pool-parallel) and
5574/// each tile rides the AMX with one sgemm — the generic sibling of
5575/// `qmatmat_accel` (q8). Numerics are f32-GEMM (tolerance class);
5576/// decode (b=1) never takes this path.
5577#[cfg(target_os = "macos")]
5578fn dequant_matmat_accel(
5579    dequant_row: &(dyn Fn(usize, &mut [f32]) + Sync),
5580    xs_all: &[f32],
5581    b: usize,
5582    rows: usize,
5583    cols: usize,
5584    out: &mut [f32],
5585    pool: Option<&Pool>,
5586) {
5587    const TR: usize = 2048;
5588    thread_local! {
5589        static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
5590    }
5591    WTILE.with(|wt| {
5592        let mut wtile = wt.borrow_mut();
5593        wtile.resize(TR * cols, 0.0);
5594        let mut r0 = 0usize;
5595        while r0 < rows {
5596            let tr = TR.min(rows - r0);
5597            let wt_addr = SendMut(wtile.as_mut_ptr());
5598            let run = |start: usize, end: usize| {
5599                for r in start..end {
5600                    // SAFETY: workers cover disjoint r ranges.
5601                    let dst = unsafe { std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols) };
5602                    dequant_row(r0 + r, dst);
5603                }
5604            };
5605            dispatch_rows(pool, tr, &run);
5606            unsafe {
5607                accel_blas::cblas_sgemm(
5608                    101, // RowMajor
5609                    111, // NoTrans A
5610                    112, // Trans B
5611                    b as i32,
5612                    tr as i32,
5613                    cols as i32,
5614                    1.0,
5615                    xs_all.as_ptr(),
5616                    cols as i32,
5617                    wtile.as_ptr(),
5618                    cols as i32,
5619                    0.0,
5620                    out.as_mut_ptr().add(r0),
5621                    rows as i32,
5622                );
5623            }
5624            r0 += tr;
5625        }
5626    });
5627}
5628
5629fn q4t_matmat(
5630    bytes: &[u8],
5631    xs_all: &[f32],
5632    b: usize,
5633    rows: usize,
5634    cols: usize,
5635    out: &mut [f32],
5636    pool: Option<&Pool>,
5637) {
5638    debug_assert_eq!(out.len(), b * rows);
5639    let gpr = cols / GROUP_SIZE;
5640    // Wide batches ride the AMX like q8's qmatmat: on Apple silicon
5641    // the dequant-tile sgemm is an order above the SDOT row loop for
5642    // prefill shapes (imagegen DiT forwards are exactly this).
5643    #[cfg(target_os = "macos")]
5644    if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
5645        dequant_matmat_accel(
5646            &|r, dst| {
5647                for gi in 0..gpr {
5648                    let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
5649                    let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
5650                    for (k, &bb) in tile[2..].iter().enumerate() {
5651                        dst[gi * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s;
5652                        dst[gi * GROUP_SIZE + k * 2 + 1] = (((bb >> 4) & 0x0F) as f32 - 8.0) * s;
5653                    }
5654                }
5655            },
5656            xs_all,
5657            b,
5658            rows,
5659            cols,
5660            out,
5661            pool,
5662        );
5663        return;
5664    }
5665    let out_addr = SendMut(out.as_mut_ptr());
5666    if a8w8_enabled() {
5667        let acts: Vec<SplitAct> = (0..b)
5668            .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
5669            .collect();
5670        let acts = &acts;
5671        #[cfg(target_arch = "x86_64")]
5672        let blocked_ok = avx2_enabled()
5673            && blocked_enabled();
5674        #[cfg(target_arch = "aarch64")]
5675        let blocked_ok = sdot_enabled()
5676            && blocked_enabled();
5677        #[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))]
5678        let blocked_ok = false;
5679        let run = move |start: usize, end: usize| {
5680            for r in start..end {
5681                let mut bi = 0usize;
5682                #[cfg(target_arch = "aarch64")]
5683                if blocked_ok {
5684                    while bi + 4 <= acts.len() {
5685                        let xs = [
5686                            acts[bi].xq.as_slice(),
5687                            acts[bi + 1].xq.as_slice(),
5688                            acts[bi + 2].xq.as_slice(),
5689                            acts[bi + 3].xq.as_slice(),
5690                        ];
5691                        let d = unsafe { dot_q4t_row_1x4_sdot(bytes, r, gpr, xs) };
5692                        for k in 0..4 {
5693                            let act = &acts[bi + k];
5694                            let mut acc = d[k] * act.sx;
5695                            for &(j, xv) in &act.outliers {
5696                                let (w, sc) = q4t_outlier(bytes, r, gpr, j);
5697                                acc += w * sc * xv;
5698                            }
5699                            // SAFETY: disjoint (bi, r) cells per worker.
5700                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
5701                        }
5702                        bi += 4;
5703                    }
5704                }
5705                #[cfg(target_arch = "x86_64")]
5706                if blocked_ok {
5707                    while bi + 4 <= acts.len() {
5708                        let xs = [
5709                            acts[bi].xq.as_slice(),
5710                            acts[bi + 1].xq.as_slice(),
5711                            acts[bi + 2].xq.as_slice(),
5712                            acts[bi + 3].xq.as_slice(),
5713                        ];
5714                        let d = unsafe {
5715                            if vnni_tiles_enabled() {
5716                                dot_q4t_row_1x4_vnni(bytes, r, gpr, xs)
5717                            } else {
5718                                dot_q4t_row_1x4_avx2(bytes, r, gpr, xs)
5719                            }
5720                        };
5721                        for k in 0..4 {
5722                            let act = &acts[bi + k];
5723                            let mut acc = d[k] * act.sx;
5724                            for &(j, xv) in &act.outliers {
5725                                let (w, sc) = q4t_outlier(bytes, r, gpr, j);
5726                                acc += w * sc * xv;
5727                            }
5728                            // SAFETY: disjoint (bi, r) cells per worker.
5729                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
5730                        }
5731                        bi += 4;
5732                    }
5733                }
5734                let _ = blocked_ok;
5735                while bi < acts.len() {
5736                    let act = &acts[bi];
5737                    let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
5738                    for &(j, xv) in &act.outliers {
5739                        let (w, s) = q4t_outlier(bytes, r, gpr, j);
5740                        acc += w * s * xv;
5741                    }
5742                    // SAFETY: disjoint (bi, r) cells per worker range.
5743                    unsafe { *out_addr.at(bi * rows + r) = acc };
5744                    bi += 1;
5745                }
5746            }
5747        };
5748        dispatch_rows(pool, rows, &run);
5749        return;
5750    }
5751    let run = move |start: usize, end: usize| {
5752        for r in start..end {
5753            for bi in 0..b {
5754                let x = &xs_all[bi * cols..(bi + 1) * cols];
5755                // SAFETY: disjoint (bi, r) cells per worker range.
5756                unsafe { *out_addr.at(bi * rows + r) = q4t_row_exact(bytes, r, gpr, x) };
5757            }
5758        }
5759    };
5760    dispatch_rows(pool, rows, &run);
5761}
5762
5763// ── q1 (dtype 12): binary weights, [f16 scale][4B sign bits] per
5764// 32-group tile. The kernel family mirrors q4_tiled: one sequential
5765// stream of 6-byte tiles, per-tile integer dot × scale, exact outlier
5766// correction (A8W8 contract), exact scalar path under CMF_SDOT=0. ──
5767
5768/// Per-32-group sums of the quantized activation — the ±1 identity's
5769/// shared half: `dot = −2·sdot(mask, x) − gsum[g]`, computed ONCE per
5770/// matvec and reused by every row.
5771fn q1_group_sums(xq: &[i8], gpr: usize) -> Vec<i32> {
5772    (0..gpr)
5773        .map(|gi| {
5774            xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE]
5775                .iter()
5776                .map(|&v| v as i32)
5777                .sum()
5778        })
5779        .collect()
5780}
5781
5782/// One q1 row via the A8W8 int8 path — mask-SDOT on ARM (no ±1
5783/// expansion at all), scalar bit loop elsewhere (AVX2 queued with the
5784/// x86 pass).
5785#[inline]
5786#[allow(unreachable_code)]
5787/// AVX2 q1 row via the same ±1 identity as the ARM sdot kernel: the
5788/// sign bits expand to a {0, −1} byte mask through shuffle+cmpeq, the
5789/// masked activation sums through maddubs(1, x&mask), and
5790/// `dot = −(2·masked_sum + Σx_group)` — bit-identical integer math.
5791#[cfg(target_arch = "x86_64")]
5792#[target_feature(enable = "avx2")]
5793unsafe fn dot_q1_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
5794    // SAFETY: callers uphold the 6B-tile and xq/gsum length contracts.
5795    unsafe {
5796        use core::arch::x86_64::*;
5797        // Byte j of the mask must replicate bits-byte j/8.
5798        let expand = _mm256_setr_epi8(
5799            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,
5800            3, 3, 3,
5801        );
5802        let bitsel = _mm256_setr_epi8(
5803            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
5804            -128, 1, 2, 4, 8, 16, 32, 64, -128,
5805        );
5806        let ones8 = _mm256_set1_epi8(1);
5807        let ones16 = _mm256_set1_epi16(1);
5808        let mut acc = 0f32;
5809        for gi in 0..gpr {
5810            let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
5811            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5812            let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
5813            let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
5814            let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
5815            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5816            let sel = _mm256_and_si256(x, mask);
5817            // Σ of selected i8 lanes: maddubs(1u8, sel_i8) pairs → madd.
5818            let p16 = _mm256_maddubs_epi16(ones8, sel);
5819            let d32 = _mm256_madd_epi16(p16, ones16);
5820            let hi128 = _mm256_extracti128_si256::<1>(d32);
5821            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
5822            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
5823            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
5824            let msum = _mm_cvtsi128_si32(s32);
5825            // The and-select keeps x UN-negated (unlike ARM's −1-mask
5826            // sdot): d = Σ_set − Σ_unset = 2·Σ_set − Σ_all.
5827            let d = 2 * msum - gsum[gi];
5828            acc += d as f32 * s;
5829        }
5830        acc
5831    }
5832}
5833
5834/// VNNI twin of `dot_q1_row_avx2`: the masked-select sum goes through
5835/// one `vpdpbusd(1u8, sel)` (see `dpbusd_hsum` — bit-identical).
5836#[cfg(target_arch = "x86_64")]
5837#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
5838unsafe fn dot_q1_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
5839    // SAFETY: callers uphold the 6B-tile and xq/gsum length contracts.
5840    unsafe {
5841        use core::arch::x86_64::*;
5842        let expand = _mm256_setr_epi8(
5843            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,
5844            3, 3, 3,
5845        );
5846        let bitsel = _mm256_setr_epi8(
5847            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
5848            -128, 1, 2, 4, 8, 16, 32, 64, -128,
5849        );
5850        let ones8 = _mm256_set1_epi8(1);
5851        let mut acc = 0f32;
5852        for gi in 0..gpr {
5853            let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
5854            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5855            let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
5856            let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
5857            let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
5858            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5859            let msum = dpbusd_hsum(ones8, _mm256_and_si256(x, mask));
5860            let d = 2 * msum - gsum[gi];
5861            acc += d as f32 * s;
5862        }
5863        acc
5864    }
5865}
5866
5867/// VNNI twin of `dot_q1_row_1x4_avx2` (see `dpbusd_hsum`).
5868#[cfg(target_arch = "x86_64")]
5869#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
5870unsafe fn dot_q1_row_1x4_vnni(
5871    bytes: &[u8],
5872    r: usize,
5873    gpr: usize,
5874    xs: [&[i8]; 4],
5875    gsums: [&[i32]; 4],
5876) -> [f32; 4] {
5877    // SAFETY: callers uphold the 6B-tile and xq/gsum length contracts.
5878    unsafe {
5879        use core::arch::x86_64::*;
5880        let expand = _mm256_setr_epi8(
5881            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,
5882            3, 3, 3,
5883        );
5884        let bitsel = _mm256_setr_epi8(
5885            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
5886            -128, 1, 2, 4, 8, 16, 32, 64, -128,
5887        );
5888        let ones8 = _mm256_set1_epi8(1);
5889        let mut acc = [0f32; 4];
5890        for gi in 0..gpr {
5891            let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
5892            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5893            let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
5894            let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
5895            let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
5896            for (k, xq) in xs.iter().enumerate() {
5897                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5898                let msum = dpbusd_hsum(ones8, _mm256_and_si256(x, mask));
5899                let d = 2 * msum - gsums[k][gi];
5900                acc[k] += d as f32 * s;
5901            }
5902        }
5903        acc
5904    }
5905}
5906
5907/// The blocked 1×4 flavor: the expanded bit mask serves four activation
5908/// streams per group (mask build once, four select+reduce chains).
5909#[cfg(target_arch = "x86_64")]
5910#[target_feature(enable = "avx2")]
5911unsafe fn dot_q1_row_1x4_avx2(
5912    bytes: &[u8],
5913    r: usize,
5914    gpr: usize,
5915    xs: [&[i8]; 4],
5916    gsums: [&[i32]; 4],
5917) -> [f32; 4] {
5918    // SAFETY: callers uphold the 6B-tile and xq/gsum length contracts.
5919    unsafe {
5920        use core::arch::x86_64::*;
5921        let expand = _mm256_setr_epi8(
5922            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,
5923            3, 3, 3,
5924        );
5925        let bitsel = _mm256_setr_epi8(
5926            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
5927            -128, 1, 2, 4, 8, 16, 32, 64, -128,
5928        );
5929        let ones8 = _mm256_set1_epi8(1);
5930        let ones16 = _mm256_set1_epi16(1);
5931        let mut acc = [0f32; 4];
5932        for gi in 0..gpr {
5933            let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
5934            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5935            let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
5936            let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
5937            let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
5938            for (k, xq) in xs.iter().enumerate() {
5939                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5940                let sel = _mm256_and_si256(x, mask);
5941                let p16 = _mm256_maddubs_epi16(ones8, sel);
5942                let d32 = _mm256_madd_epi16(p16, ones16);
5943                let hi128 = _mm256_extracti128_si256::<1>(d32);
5944                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
5945                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
5946                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
5947                let msum = _mm_cvtsi128_si32(s32);
5948                let d = 2 * msum - gsums[k][gi];
5949                acc[k] += d as f32 * s;
5950            }
5951        }
5952        acc
5953    }
5954}
5955
5956#[allow(unreachable_code)]
5957fn dot_q1_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
5958    #[cfg(target_arch = "aarch64")]
5959    unsafe {
5960        return dot_q1_row_sdot(bytes, r, gpr, xq, gsum);
5961    }
5962    #[cfg(target_arch = "x86_64")]
5963    if avx2_enabled() {
5964        unsafe {
5965            if vnni_tiles_enabled() {
5966                return dot_q1_row_vnni(bytes, r, gpr, xq, gsum);
5967            }
5968            return dot_q1_row_avx2(bytes, r, gpr, xq, gsum);
5969        }
5970    }
5971    let _ = gsum;
5972    let mut acc = 0f32;
5973    for gi in 0..gpr {
5974        let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
5975        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
5976        let mut d = 0i32;
5977        for (j, &b) in tile[2..].iter().enumerate() {
5978            for k in 0..8 {
5979                let w = ((b >> k) & 1) as i32 * 2 - 1;
5980                d += w * xq[gi * GROUP_SIZE + j * 8 + k] as i32;
5981            }
5982        }
5983        acc += d as f32 * s;
5984    }
5985    acc
5986}
5987
5988/// SDOT q1 row via the ±1 identity: the vtst mask (0xFF where the bit
5989/// is set, i.e. −1 as i8) feeds `sdot` DIRECTLY — no expansion to ±1
5990/// lanes at all — and `dot = −(2·sdot(mask, x) + Σx_group)`, with the
5991/// per-group activation sums shared across every row of the matvec.
5992/// Four tiles (128 weights) per iteration: integer dots reduce through
5993/// a vpaddq tree into ONE i32x4 that meets its four scales in a single
5994/// fused f32 multiply-add. Integer math throughout — bit-identical to
5995/// the scalar ±1 reference.
5996#[cfg(target_arch = "aarch64")]
5997#[target_feature(enable = "neon,dotprod")]
5998unsafe fn dot_q1_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
5999    // SAFETY: callers uphold slice-length contracts (6B tile per group,
6000    // xq.len() == gpr·GROUP_SIZE, gsum.len() == gpr).
6001    unsafe {
6002        use core::arch::aarch64::*;
6003        use core::arch::asm;
6004        const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
6005        let m = vld1q_u8(MASKS.as_ptr());
6006        // One tile's −Σ_set(x) as an UNREDUCED i32x4 (two mask-sdots).
6007        macro_rules! tile_dot {
6008            ($t:expr, $x:expr) => {{
6009                let v0 = vcombine_u8(vdup_n_u8(*$t.add(2)), vdup_n_u8(*$t.add(3)));
6010                let v1 = vcombine_u8(vdup_n_u8(*$t.add(4)), vdup_n_u8(*$t.add(5)));
6011                let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
6012                let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
6013                let x0 = vld1q_s8($x);
6014                let x1 = vld1q_s8($x.add(16));
6015                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6016                asm!(
6017                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
6018                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
6019                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6020                    w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
6021                    options(pure, nomem, nostack),
6022                );
6023                vaddq_s32(a0, a1)
6024            }};
6025        }
6026        // TBL unpack over PAIR loads: one vld1q covers two 6B tiles
6027        // ([s s b b b b][s s b b b b] + 4B slack), TBL replicates each
6028        // bit-byte across 8 lanes for vtst, and the four scales gather
6029        // through tbl2 into one fcvtl — the 16 ld1r broadcast loads and
6030        // 4 branchy software f16 conversions per 128 weights (the
6031        // measured load-port wall of this kernel) become 2 vector
6032        // loads + 9 table lookups. Integer math order is unchanged —
6033        // bit-identical results (FCVTL is exact on every f16).
6034        const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
6035        const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
6036        const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
6037        const IW11: [u8; 16] = [
6038            10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
6039        ];
6040        const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
6041        let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
6042        let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
6043        let isc = vld1_u8(ISC.as_ptr());
6044        // One tile's −Σ_set(x) from a TBL-unpacked pair load.
6045        macro_rules! tile_dot_tbl {
6046            ($ld:expr, $i0:expr, $i1:expr, $x:expr) => {{
6047                let w0 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i0), m));
6048                let w1 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i1), m));
6049                let x0 = vld1q_s8($x);
6050                let x1 = vld1q_s8($x.add(16));
6051                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6052                asm!(
6053                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
6054                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
6055                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6056                    w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
6057                    options(pure, nomem, nostack),
6058                );
6059                vaddq_s32(a0, a1)
6060            }};
6061        }
6062        let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
6063        let row_base = r * gpr * Q1_TILE;
6064        let abs_end = bytes.len();
6065        let xp = xq.as_ptr();
6066        let gp = gsum.as_ptr();
6067        let mut accv = vdupq_n_f32(0.0);
6068        let mut gi = 0;
6069        // The second pair load reads 4B past tile gi+3 — stay inside
6070        // the payload slice (only the file's final tiles fall back).
6071        while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
6072            let t0 = base.add(gi * Q1_TILE);
6073            let ld_a = vld1q_u8(t0);
6074            let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
6075            let d0 = tile_dot_tbl!(ld_a, iw00, iw01, xp.add(gi * GROUP_SIZE));
6076            let d1 = tile_dot_tbl!(ld_a, iw10, iw11, xp.add((gi + 1) * GROUP_SIZE));
6077            let d2 = tile_dot_tbl!(ld_b, iw00, iw01, xp.add((gi + 2) * GROUP_SIZE));
6078            let d3 = tile_dot_tbl!(ld_b, iw10, iw11, xp.add((gi + 3) * GROUP_SIZE));
6079            // [−Σ0, −Σ1, −Σ2, −Σ3] → dots = −(2·Σset_neg + gsum)
6080            let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
6081            let g = vld1q_s32(gp.add(gi));
6082            let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
6083            let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
6084            let scf: float32x4_t;
6085            asm!(
6086                "fcvtl {o:v}.4s, {i:v}.4h",
6087                o = out(vreg) scf, i = in(vreg) sc16,
6088                options(pure, nomem, nostack),
6089            );
6090            accv = vfmaq_f32(accv, vcvtq_f32_s32(dots), scf);
6091            gi += 4;
6092        }
6093        let mut acc = vaddvq_f32(accv);
6094        while gi < gpr {
6095            let t = base.add(gi * Q1_TILE);
6096            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
6097            let d = vaddvq_s32(tile_dot!(t, xp.add(gi * GROUP_SIZE)));
6098            acc += (-(2 * d + *gp.add(gi))) as f32 * s;
6099            gi += 1;
6100        }
6101        acc
6102    }
6103}
6104
6105/// Blocked q1 1×4: one TBL unpack of the tile pair serves FOUR
6106/// activation streams (prefill amortization — the same idea as the
6107/// AVX2 twin; per stream the group order, fma order and tail match the
6108/// single-row kernel exactly, so batch == matvec bit-for-bit).
6109#[cfg(target_arch = "aarch64")]
6110#[target_feature(enable = "neon,dotprod")]
6111unsafe fn dot_q1_row_1x4_sdot(
6112    bytes: &[u8],
6113    r: usize,
6114    gpr: usize,
6115    xs: [&[i8]; 4],
6116    gs: [&[i32]; 4],
6117) -> [f32; 4] {
6118    // SAFETY: same slice-length contracts as `dot_q1_row_sdot`, ×4.
6119    unsafe {
6120        use core::arch::aarch64::*;
6121        use core::arch::asm;
6122        const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
6123        const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
6124        const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
6125        const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
6126        const IW11: [u8; 16] = [
6127            10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
6128        ];
6129        const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
6130        let m = vld1q_u8(MASKS.as_ptr());
6131        let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
6132        let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
6133        let isc = vld1_u8(ISC.as_ptr());
6134        macro_rules! sdot2 {
6135            ($w0:expr, $w1:expr, $x:expr) => {{
6136                let x0 = vld1q_s8($x);
6137                let x1 = vld1q_s8($x.add(16));
6138                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6139                asm!(
6140                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
6141                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
6142                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6143                    w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
6144                    options(pure, nomem, nostack),
6145                );
6146                vaddq_s32(a0, a1)
6147            }};
6148        }
6149        let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
6150        let row_base = r * gpr * Q1_TILE;
6151        let abs_end = bytes.len();
6152        let mut accv = [vdupq_n_f32(0.0); 4];
6153        let mut gi = 0;
6154        while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
6155            let t0 = base.add(gi * Q1_TILE);
6156            let ld_a = vld1q_u8(t0);
6157            let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
6158            // Unpack ONCE — eight ±mask vectors serve all four streams.
6159            let w00 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw00), m));
6160            let w01 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw01), m));
6161            let w10 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw10), m));
6162            let w11 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw11), m));
6163            let w20 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw00), m));
6164            let w21 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw01), m));
6165            let w30 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw10), m));
6166            let w31 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw11), m));
6167            let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
6168            let scf: float32x4_t;
6169            asm!(
6170                "fcvtl {o:v}.4s, {i:v}.4h",
6171                o = out(vreg) scf, i = in(vreg) sc16,
6172                options(pure, nomem, nostack),
6173            );
6174            for k in 0..4 {
6175                let xp = xs[k].as_ptr();
6176                let d0 = sdot2!(w00, w01, xp.add(gi * GROUP_SIZE));
6177                let d1 = sdot2!(w10, w11, xp.add((gi + 1) * GROUP_SIZE));
6178                let d2 = sdot2!(w20, w21, xp.add((gi + 2) * GROUP_SIZE));
6179                let d3 = sdot2!(w30, w31, xp.add((gi + 3) * GROUP_SIZE));
6180                let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
6181                let g = vld1q_s32(gs[k].as_ptr().add(gi));
6182                let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
6183                accv[k] = vfmaq_f32(accv[k], vcvtq_f32_s32(dots), scf);
6184            }
6185            gi += 4;
6186        }
6187        let mut acc = [
6188            vaddvq_f32(accv[0]),
6189            vaddvq_f32(accv[1]),
6190            vaddvq_f32(accv[2]),
6191            vaddvq_f32(accv[3]),
6192        ];
6193        while gi < gpr {
6194            let t = base.add(gi * Q1_TILE);
6195            let sc = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
6196            let v0 = vcombine_u8(vdup_n_u8(*t.add(2)), vdup_n_u8(*t.add(3)));
6197            let v1 = vcombine_u8(vdup_n_u8(*t.add(4)), vdup_n_u8(*t.add(5)));
6198            let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
6199            let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
6200            for k in 0..4 {
6201                let d = vaddvq_s32(sdot2!(w0, w1, xs[k].as_ptr().add(gi * GROUP_SIZE)));
6202                acc[k] += (-(2 * d + *gs[k].as_ptr().add(gi))) as f32 * sc;
6203            }
6204            gi += 1;
6205        }
6206        acc
6207    }
6208}
6209
6210/// (weight ±1, scale) of one q1 element — the exact outlier term.
6211#[inline]
6212fn q1_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
6213    let gi = j / GROUP_SIZE;
6214    let k = j % GROUP_SIZE;
6215    let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
6216    let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
6217    let bit = (tile[2 + k / 8] >> (k % 8)) & 1;
6218    ((bit as i32 * 2 - 1) as f32, s)
6219}
6220
6221/// Exact scalar q1 row (CMF_SDOT=0 contract).
6222#[inline]
6223fn q1_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
6224    let mut acc = 0f32;
6225    for gi in 0..gpr {
6226        let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
6227        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
6228        let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
6229        let mut ga = 0f32;
6230        for (j, &b) in tile[2..].iter().enumerate() {
6231            for k in 0..8 {
6232                ga += (((b >> k) & 1) as f32 * 2.0 - 1.0) * xg[j * 8 + k];
6233            }
6234        }
6235        acc += ga * s;
6236    }
6237    acc
6238}
6239
6240/// One q1 row range via A8W8 (the body of `q1_matvec`'s hot loop,
6241/// extracted so multi-matrix jobs drive the same kernel).
6242#[allow(clippy::too_many_arguments)]
6243fn q1_range_a8w8(
6244    bytes: &[u8],
6245    gpr: usize,
6246    act: &SplitAct,
6247    gsum: &[i32],
6248    out: SendMut,
6249    start: usize,
6250    end: usize,
6251) {
6252    for r in start..end {
6253        let mut acc = dot_q1_row_i8(bytes, r, gpr, &act.xq, gsum) * act.sx;
6254        for &(j, xv) in &act.outliers {
6255            let (w, s) = q1_outlier(bytes, r, gpr, j);
6256            acc += w * s * xv;
6257        }
6258        // SAFETY: disjoint row ranges per worker.
6259        unsafe { *out.at(r) = acc };
6260    }
6261}
6262
6263/// Exact-scalar q1 row range (CMF_SDOT=0 contract).
6264fn q1_range_f32(bytes: &[u8], gpr: usize, x: &[f32], out: SendMut, start: usize, end: usize) {
6265    for r in start..end {
6266        // SAFETY: disjoint row ranges per worker.
6267        unsafe { *out.at(r) = q1_row_exact(bytes, r, gpr, x) };
6268    }
6269}
6270
6271/// q1t per-row overlay locator. After the base (`base_len`) come
6272/// `[u32 row_ptr[rows+1]]` then `[(u16 col, f16 val)]` grouped by row (row
6273/// `r`'s entries are `[row_ptr[r], row_ptr[r+1])`). Returns
6274/// `(row_ptr offset, entries offset, present)`.
6275fn q1t_overlay(bytes: &[u8], base_len: usize, rows: usize) -> (usize, usize, bool) {
6276    let entries = base_len + (rows + 1) * 4;
6277    (base_len, entries, entries <= bytes.len())
6278}
6279
6280/// Read `row_ptr[r]` from the overlay's prefix-sum table.
6281#[inline]
6282fn q1t_rowptr(bytes: &[u8], rp_off: usize, r: usize) -> usize {
6283    let o = rp_off + r * 4;
6284    u32::from_le_bytes([bytes[o], bytes[o + 1], bytes[o + 2], bytes[o + 3]]) as usize
6285}
6286
6287/// Byte → the 5 ternary signs it packs `{−1,0,+1}` as f32, precomputed so
6288/// decoding a q1t code is a table load, not the base-3 divide/modulo per
6289/// weight (division is ~20–40× the cost of a load). Built at compile time.
6290const SIGN5: [[f32; 5]; 256] = {
6291    let mut lut = [[0.0f32; 5]; 256];
6292    let pow3 = [1u16, 3, 9, 27, 81];
6293    let mut byte = 0usize;
6294    while byte < 256 {
6295        let mut i = 0usize;
6296        while i < 5 {
6297            let code = (byte as u16 / pow3[i]) % 3;
6298            lut[byte][i] = if code == 1 {
6299                1.0
6300            } else if code == 2 {
6301                -1.0
6302            } else {
6303                0.0
6304            };
6305            i += 1;
6306        }
6307        byte += 1;
6308    }
6309    lut
6310};
6311
6312/// Same table, as i8 signs — the operand for the int8 SDOT base kernel.
6313const SIGN5_I8: [[i8; 5]; 256] = {
6314    let mut lut = [[0i8; 5]; 256];
6315    let pow3 = [1u16, 3, 9, 27, 81];
6316    let mut byte = 0usize;
6317    while byte < 256 {
6318        let mut i = 0usize;
6319        while i < 5 {
6320            let code = (byte as u16 / pow3[i]) % 3;
6321            lut[byte][i] = if code == 1 {
6322                1
6323            } else if code == 2 {
6324                -1
6325            } else {
6326                0
6327            };
6328            i += 1;
6329        }
6330        byte += 1;
6331    }
6332    lut
6333};
6334
6335/// The same 5 i8 signs packed into a u64 (`[s0 s1 s2 s3 s4 0 0 0]`, LE) so the
6336/// group unpack is 7 unaligned u64 stores at offsets 0,5,10,…,30 instead of
6337/// six 5-byte copies + LUT indexing — each store's trailing zeros are fixed by
6338/// the next store, and the last one runs 6 B past the 32nd weight (the unpack
6339/// buffer is padded to 40). This is the decode/prefill hot inner op.
6340const SIGN5_U64: [u64; 256] = {
6341    let mut lut = [0u64; 256];
6342    let pow3 = [1u16, 3, 9, 27, 81];
6343    let mut byte = 0usize;
6344    while byte < 256 {
6345        let mut v = 0u64;
6346        let mut i = 0usize;
6347        while i < 5 {
6348            let code = (byte as u16 / pow3[i]) % 3;
6349            let s: u8 = if code == 1 {
6350                1
6351            } else if code == 2 {
6352                0xFF
6353            } else {
6354                0
6355            };
6356            v |= (s as u64) << (i * 8);
6357            i += 1;
6358        }
6359        lut[byte] = v;
6360        byte += 1;
6361    }
6362    lut
6363};
6364
6365/// Ternary base weight at `(row r, col j)` = `sign(code)·s_group`. Used to add
6366/// back activation-outlier columns, whose `x` was zeroed for the int8 bulk dot
6367/// (`split_act`). At a weight-outlier position the code is 0, so this is 0 and
6368/// the overlay correction owns that column — no double counting.
6369#[inline]
6370fn q1t_base_weight(bytes: &[u8], r: usize, gpr: usize, j: usize) -> f32 {
6371    const TILE: usize = cortiq_core::quant::Q1T_TILE;
6372    let off = (r * gpr + j / GROUP_SIZE) * TILE;
6373    let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
6374    let within = j % GROUP_SIZE;
6375    SIGN5[bytes[off + 2 + within / 5] as usize][within % 5] * s
6376}
6377
6378/// One 32-group int8 dot via two SDOTs. Bit-exact vs the scalar i8 sum
6379/// (integer accumulation is order-independent).
6380#[cfg(target_arch = "aarch64")]
6381#[target_feature(enable = "neon,dotprod")]
6382#[inline]
6383unsafe fn sdot32_i8(w: *const i8, x: *const i8) -> i32 {
6384    // SAFETY: caller guarantees 32 readable i8 at each pointer.
6385    unsafe {
6386        use core::arch::aarch64::*;
6387        use core::arch::asm;
6388        let w0 = vld1q_s8(w);
6389        let w1 = vld1q_s8(w.add(16));
6390        let x0 = vld1q_s8(x);
6391        let x1 = vld1q_s8(x.add(16));
6392        let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6393        asm!(
6394            "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
6395            "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
6396            a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6397            w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
6398            options(pure, nomem, nostack),
6399        );
6400        vaddvq_s32(vaddq_s32(a0, a1))
6401    }
6402}
6403
6404/// One 32-group int8 dot via AVX2: signed·signed as `maddubs(|w|, sign(x,w))`
6405/// then `madd` and a horizontal reduce (the same idiom as `dot_q4t_row_avx2`).
6406#[cfg(target_arch = "x86_64")]
6407#[target_feature(enable = "avx2")]
6408#[inline]
6409unsafe fn i8dot32_avx2(w: *const i8, x: *const i8) -> i32 {
6410    // SAFETY: caller guarantees 32 readable i8 at each pointer.
6411    unsafe {
6412        use core::arch::x86_64::*;
6413        let wv = _mm256_loadu_si256(w as *const __m256i);
6414        let xv = _mm256_loadu_si256(x as *const __m256i);
6415        let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
6416        let d = _mm256_madd_epi16(p16, _mm256_set1_epi16(1));
6417        let hi128 = _mm256_extracti128_si256::<1>(d);
6418        let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
6419        let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
6420        let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
6421        _mm_cvtsi128_si32(s32)
6422    }
6423}
6424
6425/// Unpack one q1t group's base-3 codes into 32 i8 signs via 7 unaligned u64
6426/// stores (see `SIGN5_U64`). `dst` MUST have ≥ 40 bytes: the 7th store writes
6427/// `dst[30..38]`. Stores go in order so each one's trailing zeros are
6428/// overwritten by the next; the final 6 padding bytes are unused by the dot.
6429#[inline]
6430fn q1t_unpack_group_i8(codes: *const u8, dst: &mut [i8]) {
6431    debug_assert!(dst.len() >= 40);
6432    // SAFETY: codes points at 7 readable bytes; dst has ≥ 40 bytes so every
6433    // 8-byte store at offset bi*5 (bi ≤ 6 → ≤ 30) stays in bounds.
6434    unsafe {
6435        let p = dst.as_mut_ptr();
6436        for bi in 0..7 {
6437            core::ptr::write_unaligned(
6438                p.add(bi * 5) as *mut u64,
6439                SIGN5_U64[*codes.add(bi) as usize],
6440            );
6441        }
6442    }
6443}
6444
6445/// One 32-group int8 dot, arch-dispatched (the matmat inner loop, where the
6446/// row's signs are unpacked once and dotted against every batch input).
6447/// Callers are gated by `a8w8_enabled()`, so the target-feature arms are
6448/// reachable; the scalar arm is a non-SIMD-arch fallback.
6449#[inline]
6450fn q1t_i8dot32(w: *const i8, x: *const i8) -> i32 {
6451    #[cfg(target_arch = "aarch64")]
6452    unsafe {
6453        return sdot32_i8(w, x);
6454    }
6455    #[cfg(target_arch = "x86_64")]
6456    unsafe {
6457        return i8dot32_avx2(w, x);
6458    }
6459    #[allow(unreachable_code)]
6460    unsafe {
6461        let mut s = 0i32;
6462        for k in 0..GROUP_SIZE {
6463            s += *w.add(k) as i32 * *x.add(k) as i32;
6464        }
6465        s
6466    }
6467}
6468
6469#[inline]
6470unsafe fn q1t_unpack_reg_u64s(codes: *const u8) -> (u64, u64, u64, u64) {
6471    let (s0, s1, s2, s3, s4, s5, s6) = unsafe {
6472        (
6473            SIGN5_U64[*codes as usize],
6474            SIGN5_U64[*codes.add(1) as usize],
6475            SIGN5_U64[*codes.add(2) as usize],
6476            SIGN5_U64[*codes.add(3) as usize],
6477            SIGN5_U64[*codes.add(4) as usize],
6478            SIGN5_U64[*codes.add(5) as usize],
6479            SIGN5_U64[*codes.add(6) as usize],
6480        )
6481    };
6482
6483    let u0 = s0 | (s1 << 40);
6484    let u1 = (s1 >> 24) | (s2 << 16) | (s3 << 56);
6485    let u2 = (s3 >> 8) | (s4 << 32);
6486    let u3 = (s4 >> 32) | (s5 << 8) | (s6 << 48);
6487
6488    (u0, u1, u2, u3)
6489}
6490
6491/// One q1t row's int8 base dot: `Σ_group s·dot(signs, xq)` (before the shared
6492/// `sx`). Direct register unpacking (zero stack stores/loads, no STLF stalls).
6493/// ARM SDOT.
6494#[cfg(target_arch = "aarch64")]
6495#[target_feature(enable = "neon,dotprod")]
6496unsafe fn q1t_dot_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
6497    use core::arch::aarch64::*;
6498    use core::arch::asm;
6499    unsafe {
6500        const TILE: usize = cortiq_core::quant::Q1T_TILE;
6501        let mut acc = 0f32;
6502        let bytes_ptr = bytes.as_ptr();
6503        let xq_ptr = xq.as_ptr();
6504        let row_off = r * gpr * TILE;
6505
6506        let gpr2 = gpr & !1;
6507        let mut gi = 0;
6508        while gi < gpr2 {
6509            let off0 = row_off + gi * TILE;
6510            let off1 = off0 + TILE;
6511            let s0 = f16_to_f32(u16::from_le_bytes([
6512                *bytes_ptr.add(off0),
6513                *bytes_ptr.add(off0 + 1),
6514            ]));
6515            let s1 = f16_to_f32(u16::from_le_bytes([
6516                *bytes_ptr.add(off1),
6517                *bytes_ptr.add(off1 + 1),
6518            ]));
6519
6520            let (u0_0, u1_0, u2_0, u3_0) = q1t_unpack_reg_u64s(bytes_ptr.add(off0 + 2));
6521            let (u0_1, u1_1, u2_1, u3_1) = q1t_unpack_reg_u64s(bytes_ptr.add(off1 + 2));
6522
6523            let w0_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_0), vcreate_u64(u1_0)));
6524            let w1_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_0), vcreate_u64(u3_0)));
6525            let w0_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_1), vcreate_u64(u1_1)));
6526            let w1_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_1), vcreate_u64(u3_1)));
6527
6528            let x0_0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE));
6529            let x1_0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE + 16));
6530            let x0_1 = vld1q_s8(xq_ptr.add((gi + 1) * GROUP_SIZE));
6531            let x1_1 = vld1q_s8(xq_ptr.add((gi + 1) * GROUP_SIZE + 16));
6532
6533            let (mut a0_0, mut a1_0) = (vdupq_n_s32(0), vdupq_n_s32(0));
6534            let (mut a0_1, mut a1_1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6535            asm!(
6536                "sdot {a0_0:v}.4s, {w0_0:v}.16b, {x0_0:v}.16b",
6537                "sdot {a1_0:v}.4s, {w1_0:v}.16b, {x1_0:v}.16b",
6538                "sdot {a0_1:v}.4s, {w0_1:v}.16b, {x0_1:v}.16b",
6539                "sdot {a1_1:v}.4s, {w1_1:v}.16b, {x1_1:v}.16b",
6540                a0_0 = inout(vreg) a0_0, a1_0 = inout(vreg) a1_0,
6541                a0_1 = inout(vreg) a0_1, a1_1 = inout(vreg) a1_1,
6542                w0_0 = in(vreg) w0_0, x0_0 = in(vreg) x0_0, w1_0 = in(vreg) w1_0, x1_0 = in(vreg) x1_0,
6543                w0_1 = in(vreg) w0_1, x0_1 = in(vreg) x0_1, w1_1 = in(vreg) w1_1, x1_1 = in(vreg) x1_1,
6544                options(pure, nomem, nostack),
6545            );
6546            let d0 = vaddvq_s32(vaddq_s32(a0_0, a1_0));
6547            let d1 = vaddvq_s32(vaddq_s32(a0_1, a1_1));
6548            acc += d0 as f32 * s0 + d1 as f32 * s1;
6549            gi += 2;
6550        }
6551
6552        if gi < gpr {
6553            let off = row_off + gi * TILE;
6554            let s = f16_to_f32(u16::from_le_bytes([
6555                *bytes_ptr.add(off),
6556                *bytes_ptr.add(off + 1),
6557            ]));
6558            let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
6559            let w0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0), vcreate_u64(u1)));
6560            let w1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2), vcreate_u64(u3)));
6561            let x0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE));
6562            let x1 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE + 16));
6563            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6564            asm!(
6565                "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
6566                "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
6567                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6568                w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
6569                options(pure, nomem, nostack),
6570            );
6571            let d = vaddvq_s32(vaddq_s32(a0, a1));
6572            acc += d as f32 * s;
6573        }
6574        acc
6575    }
6576}
6577
6578/// x86 AVX2 mirror of `q1t_dot_row_sdot` (maddubs int8 dot per group).
6579#[cfg(target_arch = "x86_64")]
6580#[target_feature(enable = "avx2")]
6581unsafe fn q1t_dot_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
6582    use core::arch::x86_64::*;
6583    unsafe {
6584        const TILE: usize = cortiq_core::quant::Q1T_TILE;
6585        let mut acc = 0f32;
6586        let bytes_ptr = bytes.as_ptr();
6587        let xq_ptr = xq.as_ptr();
6588        let row_off = r * gpr * TILE;
6589
6590        let ones = _mm256_set1_epi16(1);
6591        for gi in 0..gpr {
6592            let off = row_off + gi * TILE;
6593            let s = f16_to_f32(u16::from_le_bytes([
6594                *bytes_ptr.add(off),
6595                *bytes_ptr.add(off + 1),
6596            ]));
6597            let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
6598            let wv = _mm256_set_epi64x(u3 as i64, u2 as i64, u1 as i64, u0 as i64);
6599            let xv = _mm256_loadu_si256(xq_ptr.add(gi * GROUP_SIZE) as *const __m256i);
6600            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
6601            let d256 = _mm256_madd_epi16(p16, ones);
6602            let d128 = _mm_add_epi32(
6603                _mm256_castsi256_si128(d256),
6604                _mm256_extracti128_si256(d256, 1),
6605            );
6606            let d64 = _mm_add_epi32(d128, _mm_shuffle_epi32(d128, 0xee));
6607            let d32 = _mm_cvtsi128_si32(_mm_add_epi32(d64, _mm_shuffle_epi32(d64, 0x55)));
6608            acc += d32 as f32 * s;
6609        }
6610        acc
6611    }
6612}
6613
6614/// VNNI twin of `q1t_dot_row_avx2` (see `dpbusd_hsum`).
6615#[cfg(target_arch = "x86_64")]
6616#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
6617unsafe fn q1t_dot_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
6618    use core::arch::x86_64::*;
6619    // SAFETY: same tile/xq contracts as `q1t_dot_row_avx2`.
6620    unsafe {
6621        const TILE: usize = cortiq_core::quant::Q1T_TILE;
6622        let mut acc = 0f32;
6623        let bytes_ptr = bytes.as_ptr();
6624        let xq_ptr = xq.as_ptr();
6625        let row_off = r * gpr * TILE;
6626        for gi in 0..gpr {
6627            let off = row_off + gi * TILE;
6628            let s = f16_to_f32(u16::from_le_bytes([
6629                *bytes_ptr.add(off),
6630                *bytes_ptr.add(off + 1),
6631            ]));
6632            let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
6633            let wv = _mm256_set_epi64x(u3 as i64, u2 as i64, u1 as i64, u0 as i64);
6634            let xv = _mm256_loadu_si256(xq_ptr.add(gi * GROUP_SIZE) as *const __m256i);
6635            let d = dpbusd_hsum(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
6636            acc += d as f32 * s;
6637        }
6638        acc
6639    }
6640}
6641
6642/// Per-row int8 base dot, dispatched once per row (matvec decode hot path).
6643/// Callers are gated by `a8w8_enabled()`, so the target-feature kernels are
6644/// reachable.
6645#[inline]
6646fn q1t_dot_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
6647    #[cfg(target_arch = "aarch64")]
6648    unsafe {
6649        return q1t_dot_row_sdot(bytes, r, gpr, xq);
6650    }
6651    #[cfg(target_arch = "x86_64")]
6652    unsafe {
6653        if vnni_tiles_enabled() {
6654            return q1t_dot_row_vnni(bytes, r, gpr, xq);
6655        }
6656        return q1t_dot_row_avx2(bytes, r, gpr, xq);
6657    }
6658    #[allow(unreachable_code)]
6659    {
6660        const TILE: usize = cortiq_core::quant::Q1T_TILE;
6661        let mut acc = 0f32;
6662        let mut sg = [0i8; GROUP_SIZE + 8]; // +8 slack for the u64-store unpack
6663        for gi in 0..gpr {
6664            let off = (r * gpr + gi) * TILE;
6665            let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
6666            q1t_unpack_group_i8(bytes.as_ptr().wrapping_add(off + 2), &mut sg);
6667            let mut d = 0i32;
6668            for k in 0..GROUP_SIZE {
6669                d += sg[k] as i32 * xq[gi * GROUP_SIZE + k] as i32;
6670            }
6671            acc += d as f32 * s;
6672        }
6673        acc
6674    }
6675}
6676
6677/// Σ over a row's outliers of `value·x[col]` — the correction that adds the
6678/// overlay's exact weights on top of the base dot. INVARIANT: the encoder
6679/// writes ternary code 0 at every outlier position (`quantize_q1t`), so the
6680/// base contributes nothing there and this is a plain `value·x`, not
6681/// `(value − base)·x` — no scattered per-outlier scale read. Row `r`'s entries
6682/// are the contiguous slice `[row_ptr[r], row_ptr[r+1])`, so no binary search.
6683fn q1t_row_outlier_correction(
6684    bytes: &[u8],
6685    r: usize,
6686    rp_off: usize,
6687    entries_off: usize,
6688    has_ov: bool,
6689    x: &[f32],
6690) -> f32 {
6691    if !has_ov {
6692        return 0.0;
6693    }
6694    let (c0, c1) = (
6695        q1t_rowptr(bytes, rp_off, r),
6696        q1t_rowptr(bytes, rp_off, r + 1),
6697    );
6698    let mut corr = 0f32;
6699    for p in c0..c1 {
6700        let e = entries_off + p * 4;
6701        let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
6702        let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
6703        corr += val * x[col];
6704    }
6705    corr
6706}
6707
6708/// Dequantize one q1t row into `buf[..cols]` via the sign LUT (no division),
6709/// then apply the row's outliers (its `[row_ptr[r], row_ptr[r+1])` slice).
6710/// Used by the batched (prefill) path where the decode amortizes over the batch.
6711fn q1t_dequant_row(
6712    bytes: &[u8],
6713    r: usize,
6714    gpr: usize,
6715    rp_off: usize,
6716    entries_off: usize,
6717    has_ov: bool,
6718    buf: &mut [f32],
6719) {
6720    const TILE: usize = cortiq_core::quant::Q1T_TILE;
6721    for g in 0..gpr {
6722        let off = (r * gpr + g) * TILE;
6723        let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
6724        let codes = &bytes[off + 2..off + TILE];
6725        let bc = g * GROUP_SIZE;
6726        // 6 full bytes (30 codes) + a 7th byte holding the last 2.
6727        for bi in 0..6 {
6728            let lut = &SIGN5[codes[bi] as usize];
6729            let d = &mut buf[bc + bi * 5..bc + bi * 5 + 5];
6730            for i in 0..5 {
6731                d[i] = lut[i] * s;
6732            }
6733        }
6734        let lut = &SIGN5[codes[6] as usize];
6735        buf[bc + 30] = lut[0] * s;
6736        buf[bc + 31] = lut[1] * s;
6737    }
6738    if !has_ov {
6739        return;
6740    }
6741    let (c0, c1) = (
6742        q1t_rowptr(bytes, rp_off, r),
6743        q1t_rowptr(bytes, rp_off, r + 1),
6744    );
6745    for p in c0..c1 {
6746        let e = entries_off + p * 4;
6747        let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
6748        buf[col] = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
6749    }
6750}
6751
6752/// Add the sparse outlier overlay onto a base dot already in `out` (the GPU
6753/// computes the ternary base; the overlay stays on the CPU — its entries are
6754/// few and its per-row gather doesn't vectorize on the GPU). Row-parallel.
6755fn q1t_add_overlay(
6756    bytes: &[u8],
6757    x: &[f32],
6758    rows: usize,
6759    cols: usize,
6760    out: &mut [f32],
6761    pool: Option<&Pool>,
6762) {
6763    const TILE: usize = cortiq_core::quant::Q1T_TILE;
6764    let gpr = cols / GROUP_SIZE;
6765    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
6766    if !has_ov {
6767        return;
6768    }
6769    let out_addr = SendMut(out.as_mut_ptr());
6770    let run = move |start: usize, end: usize| {
6771        for r in start..end {
6772            let corr = q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
6773            // SAFETY: disjoint rows; add onto the base the GPU already wrote.
6774            unsafe { *out_addr.at(r) += corr };
6775        }
6776    };
6777    dispatch_rows(pool, rows, &run);
6778}
6779
6780/// Q1T row range via the A8W8 int8 path — shared activation split,
6781/// per-row: base SDOT dot + outlier correction + overlay.
6782#[allow(clippy::too_many_arguments)]
6783fn q1t_range_a8w8(
6784    bytes: &[u8],
6785    gpr: usize,
6786    rp_off: usize,
6787    ent_off: usize,
6788    has_ov: bool,
6789    act: &SplitAct,
6790    x: &[f32],
6791    out: SendMut,
6792    start: usize,
6793    end: usize,
6794) {
6795    for r in start..end {
6796        let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
6797        for &(j, xv) in &act.outliers {
6798            acc += q1t_base_weight(bytes, r, gpr, j) * xv;
6799        }
6800        acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
6801        // SAFETY: disjoint row ranges per worker.
6802        unsafe { *out.at(r) = acc };
6803    }
6804}
6805
6806/// Q1T row range via the f32 path (no SDOT) — for matvec_many batched
6807/// dispatch when a8w8 is unavailable.
6808#[allow(clippy::too_many_arguments)]
6809fn q1t_range_f32_batch(
6810    bytes: &[u8],
6811    gpr: usize,
6812    rp_off: usize,
6813    ent_off: usize,
6814    has_ov: bool,
6815    x: &[f32],
6816    out: SendMut,
6817    start: usize,
6818    end: usize,
6819) {
6820    const TILE: usize = cortiq_core::quant::Q1T_TILE;
6821    let mut sg = [0f32; GROUP_SIZE];
6822    for r in start..end {
6823        let mut acc = 0f32;
6824        for g in 0..gpr {
6825            let off = (r * gpr + g) * TILE;
6826            let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
6827            let codes = &bytes[off + 2..off + TILE];
6828            let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
6829            for bi in 0..6 {
6830                sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
6831            }
6832            let lut = &SIGN5[codes[6] as usize];
6833            sg[30] = lut[0];
6834            sg[31] = lut[1];
6835            let mut gsum = 0f32;
6836            for k in 0..GROUP_SIZE {
6837                gsum += sg[k] * xg[k];
6838            }
6839            acc += s * gsum;
6840        }
6841        acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
6842        // SAFETY: disjoint row ranges per worker.
6843        unsafe { *out.at(r) = acc };
6844    }
6845}
6846
6847/// Ternary (q1t) matvec — decode+dot straight from mmap, one group at a time:
6848/// no per-ROW buffer, no division (the sign LUT), and a tiny per-group sign
6849/// buffer so the 32-wide dot vectorizes. This is the decode hot path.
6850fn q1t_matvec(
6851    bytes: &[u8],
6852    x: &[f32],
6853    rows: usize,
6854    cols: usize,
6855    out: &mut [f32],
6856    pool: Option<&Pool>,
6857) {
6858    debug_assert_eq!(out.len(), rows);
6859    const TILE: usize = cortiq_core::quant::Q1T_TILE;
6860    let gpr = cols / GROUP_SIZE;
6861    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
6862    let out_addr = SendMut(out.as_mut_ptr());
6863    // int8 SDOT base dot (ARM dotprod): ~4× the f32 arithmetic. x → i8 once
6864    // (`split_act`), activation outliers added back exactly in f32, weight
6865    // overlay on top. ARM SDOT / x86 AVX2; CMF_SDOT=0 keeps the exact f32 path.
6866    if a8w8_enabled() {
6867        let act = split_act(x);
6868        let act = &act;
6869        let run = move |start: usize, end: usize| {
6870            for r in start..end {
6871                let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
6872                for &(j, xv) in &act.outliers {
6873                    acc += q1t_base_weight(bytes, r, gpr, j) * xv;
6874                }
6875                acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
6876                // SAFETY: disjoint row ranges per worker.
6877                unsafe { *out_addr.at(r) = acc };
6878            }
6879        };
6880        dispatch_rows(pool, rows, &run);
6881        return;
6882    }
6883    let run = move |start: usize, end: usize| {
6884        // Per-group signs, unpacked contiguously so the dot below is a clean
6885        // 32-wide reduction the autovectorizer turns into f32x4 FMAs — the
6886        // 5-values-per-byte base-3 layout won't SIMD in place.
6887        let mut sg = [0f32; GROUP_SIZE];
6888        for r in start..end {
6889            let mut acc = 0f32;
6890            for g in 0..gpr {
6891                let off = (r * gpr + g) * TILE;
6892                let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
6893                let codes = &bytes[off + 2..off + TILE];
6894                let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
6895                for bi in 0..6 {
6896                    sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
6897                }
6898                let lut = &SIGN5[codes[6] as usize];
6899                sg[30] = lut[0];
6900                sg[31] = lut[1];
6901                let mut gsum = 0f32;
6902                for k in 0..GROUP_SIZE {
6903                    gsum += sg[k] * xg[k];
6904                }
6905                acc += s * gsum;
6906            }
6907            acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
6908            unsafe { *out_addr.at(r) = acc };
6909        }
6910    };
6911    dispatch_rows(pool, rows, &run);
6912}
6913
6914/// Fused-pair twin of `q1t_dot_row_sdot`: ONE register unpack of the
6915/// ternary codes serves BOTH activation streams (the unpack chain is
6916/// the dominant per-row cost — MTP verify pairs paid it twice). Per
6917/// stream the group order and f32 accumulation match the single-row
6918/// kernel exactly, so pair == 2×matvec bit-for-bit.
6919#[cfg(target_arch = "aarch64")]
6920#[target_feature(enable = "neon,dotprod")]
6921unsafe fn q1t_dot_row_sdot2(bytes: &[u8], r: usize, gpr: usize, xa: &[i8], xb: &[i8]) -> [f32; 2] {
6922    use core::arch::aarch64::*;
6923    use core::arch::asm;
6924    // SAFETY: same slice-length contracts as `q1t_dot_row_sdot`, ×2.
6925    unsafe {
6926        const TILE: usize = cortiq_core::quant::Q1T_TILE;
6927        let bytes_ptr = bytes.as_ptr();
6928        let row_off = r * gpr * TILE;
6929        let xp = [xa.as_ptr(), xb.as_ptr()];
6930        let mut acc = [0f32; 2];
6931        macro_rules! sdot2 {
6932            ($w0:expr, $w1:expr, $x:expr) => {{
6933                let x0 = vld1q_s8($x);
6934                let x1 = vld1q_s8($x.add(16));
6935                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6936                asm!(
6937                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
6938                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
6939                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6940                    w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
6941                    options(pure, nomem, nostack),
6942                );
6943                vaddvq_s32(vaddq_s32(a0, a1))
6944            }};
6945        }
6946        let gpr2 = gpr & !1;
6947        let mut gi = 0;
6948        while gi < gpr2 {
6949            let off0 = row_off + gi * TILE;
6950            let off1 = off0 + TILE;
6951            let s0 = f16_to_f32(u16::from_le_bytes([
6952                *bytes_ptr.add(off0),
6953                *bytes_ptr.add(off0 + 1),
6954            ]));
6955            let s1 = f16_to_f32(u16::from_le_bytes([
6956                *bytes_ptr.add(off1),
6957                *bytes_ptr.add(off1 + 1),
6958            ]));
6959            let (u0_0, u1_0, u2_0, u3_0) = q1t_unpack_reg_u64s(bytes_ptr.add(off0 + 2));
6960            let (u0_1, u1_1, u2_1, u3_1) = q1t_unpack_reg_u64s(bytes_ptr.add(off1 + 2));
6961            let w0_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_0), vcreate_u64(u1_0)));
6962            let w1_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_0), vcreate_u64(u3_0)));
6963            let w0_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_1), vcreate_u64(u1_1)));
6964            let w1_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_1), vcreate_u64(u3_1)));
6965            for k in 0..2 {
6966                let d0 = sdot2!(w0_0, w1_0, xp[k].add(gi * GROUP_SIZE));
6967                let d1 = sdot2!(w0_1, w1_1, xp[k].add((gi + 1) * GROUP_SIZE));
6968                acc[k] += d0 as f32 * s0 + d1 as f32 * s1;
6969            }
6970            gi += 2;
6971        }
6972        if gi < gpr {
6973            let off = row_off + gi * TILE;
6974            let s = f16_to_f32(u16::from_le_bytes([
6975                *bytes_ptr.add(off),
6976                *bytes_ptr.add(off + 1),
6977            ]));
6978            let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
6979            let w0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0), vcreate_u64(u1)));
6980            let w1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2), vcreate_u64(u3)));
6981            for k in 0..2 {
6982                let d = sdot2!(w0, w1, xp[k].add(gi * GROUP_SIZE));
6983                acc[k] += d as f32 * s;
6984            }
6985        }
6986        acc
6987    }
6988}
6989
6990/// Fused Q1T pair matvec: ONE pass over the rows serves both
6991/// activation streams — on ARM the ternary register unpack happens
6992/// once per tile pair (`q1t_dot_row_sdot2`); elsewhere the second dot
6993/// rides the row's L1-warm tile bytes. Per stream the math matches
6994/// `q1t_matvec` exactly.
6995fn q1t_matvec2(
6996    bytes: &[u8],
6997    x1: &[f32],
6998    x2: &[f32],
6999    rows: usize,
7000    cols: usize,
7001    o1: &mut [f32],
7002    o2: &mut [f32],
7003    pool: Option<&Pool>,
7004) {
7005    debug_assert_eq!(o1.len(), rows);
7006    debug_assert_eq!(o2.len(), rows);
7007    const TILE: usize = cortiq_core::quant::Q1T_TILE;
7008    let gpr = cols / GROUP_SIZE;
7009    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
7010    let out1 = SendMut(o1.as_mut_ptr());
7011    let out2 = SendMut(o2.as_mut_ptr());
7012    if a8w8_enabled() {
7013        let a1 = split_act(x1);
7014        let a2 = split_act(x2);
7015        let (a1, a2) = (&a1, &a2);
7016        let run = move |start: usize, end: usize| {
7017            for r in start..end {
7018                #[cfg(target_arch = "aarch64")]
7019                // a8w8 on aarch64 ⇔ sdot_enabled(), so the kernel's
7020                // target features are present.
7021                let ds = unsafe { q1t_dot_row_sdot2(bytes, r, gpr, &a1.xq, &a2.xq) };
7022                #[cfg(not(target_arch = "aarch64"))]
7023                let ds = [
7024                    q1t_dot_row_i8(bytes, r, gpr, &a1.xq),
7025                    q1t_dot_row_i8(bytes, r, gpr, &a2.xq),
7026                ];
7027                let mut acc1 = ds[0] * a1.sx;
7028                for &(j, xv) in &a1.outliers {
7029                    acc1 += q1t_base_weight(bytes, r, gpr, j) * xv;
7030                }
7031                acc1 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x1);
7032                let mut acc2 = ds[1] * a2.sx;
7033                for &(j, xv) in &a2.outliers {
7034                    acc2 += q1t_base_weight(bytes, r, gpr, j) * xv;
7035                }
7036                acc2 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x2);
7037                // SAFETY: disjoint row ranges per worker.
7038                unsafe {
7039                    *out1.at(r) = acc1;
7040                    *out2.at(r) = acc2;
7041                }
7042            }
7043        };
7044        dispatch_rows(pool, rows, &run);
7045        return;
7046    }
7047    let run = move |start: usize, end: usize| {
7048        // Exact path (CMF_SDOT=0): unpack the sign LUT once per group,
7049        // dot both streams — same op order per stream as `q1t_matvec`.
7050        let mut sg = [0f32; GROUP_SIZE];
7051        for r in start..end {
7052            let mut acc1 = 0f32;
7053            let mut acc2 = 0f32;
7054            for g in 0..gpr {
7055                let off = (r * gpr + g) * TILE;
7056                let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
7057                let codes = &bytes[off + 2..off + TILE];
7058                for bi in 0..6 {
7059                    sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
7060                }
7061                let lut = &SIGN5[codes[6] as usize];
7062                sg[30] = lut[0];
7063                sg[31] = lut[1];
7064                let xg1 = &x1[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
7065                let xg2 = &x2[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
7066                let mut gsum1 = 0f32;
7067                for k in 0..GROUP_SIZE {
7068                    gsum1 += sg[k] * xg1[k];
7069                }
7070                acc1 += s * gsum1;
7071                let mut gsum2 = 0f32;
7072                for k in 0..GROUP_SIZE {
7073                    gsum2 += sg[k] * xg2[k];
7074                }
7075                acc2 += s * gsum2;
7076            }
7077            acc1 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x1);
7078            acc2 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x2);
7079            // SAFETY: disjoint row ranges per worker.
7080            unsafe {
7081                *out1.at(r) = acc1;
7082                *out2.at(r) = acc2;
7083            }
7084        }
7085    };
7086    dispatch_rows(pool, rows, &run);
7087}
7088
7089/// Ternary (q1t) matmat (prefill) — dequant each row once, dot the whole
7090/// batch against it (amortizes the per-row decode).
7091fn q1t_matmat(
7092    bytes: &[u8],
7093    xs: &[f32],
7094    b: usize,
7095    rows: usize,
7096    cols: usize,
7097    out: &mut [f32],
7098    pool: Option<&Pool>,
7099) {
7100    debug_assert_eq!(out.len(), b * rows);
7101    const TILE: usize = cortiq_core::quant::Q1T_TILE;
7102    let gpr = cols / GROUP_SIZE;
7103    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
7104    let out_addr = SendMut(out.as_mut_ptr());
7105    // int8 prefill (ARM SDOT / x86 AVX2): quantize the B inputs once, unpack
7106    // each weight row's signs to i8 ONCE, then int8-dot against every input —
7107    // the row sign-decode amortizes over the whole batch. CMF_SDOT=0 → f32.
7108    if a8w8_enabled() {
7109        let acts: Vec<SplitAct> = (0..b)
7110            .map(|bi| split_act(&xs[bi * cols..(bi + 1) * cols]))
7111            .collect();
7112        let acts = &acts;
7113        let run = move |start: usize, end: usize| {
7114            let mut sg = vec![0i8; cols + 8]; // row signs, i8 (+8 unpack slack)
7115            let mut sc = vec![0f32; gpr]; // per-group scales
7116            let mut accs = vec![0f32; b]; // per-batch accumulators, reused per row
7117            for r in start..end {
7118                for g in 0..gpr {
7119                    let off = (r * gpr + g) * TILE;
7120                    sc[g] = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
7121                    q1t_unpack_group_i8(
7122                        bytes.as_ptr().wrapping_add(off + 2),
7123                        &mut sg[g * GROUP_SIZE..],
7124                    );
7125                }
7126                for bi in 0..b {
7127                    let act = &acts[bi];
7128                    let mut isum = 0f32;
7129                    for g in 0..gpr {
7130                        let d = q1t_i8dot32(
7131                            sg.as_ptr().wrapping_add(g * GROUP_SIZE),
7132                            act.xq.as_ptr().wrapping_add(g * GROUP_SIZE),
7133                        );
7134                        isum += d as f32 * sc[g];
7135                    }
7136                    let mut acc = isum * act.sx;
7137                    for &(j, xv) in &act.outliers {
7138                        acc += q1t_base_weight(bytes, r, gpr, j) * xv;
7139                    }
7140                    accs[bi] = acc;
7141                }
7142                // Overlay ONCE per row for the whole batch: read each (col, val)
7143                // from mmap a single time (was b× — the re-read dominated prefill)
7144                // and fan it out over the batch via the cached inputs.
7145                if has_ov {
7146                    let (c0, c1) = (
7147                        q1t_rowptr(bytes, rp_off, r),
7148                        q1t_rowptr(bytes, rp_off, r + 1),
7149                    );
7150                    for p in c0..c1 {
7151                        let e = ent_off + p * 4;
7152                        let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
7153                        let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
7154                        for bi in 0..b {
7155                            accs[bi] += val * xs[bi * cols + col];
7156                        }
7157                    }
7158                }
7159                for bi in 0..b {
7160                    unsafe { *out_addr.at(bi * rows + r) = accs[bi] };
7161                }
7162            }
7163        };
7164        dispatch_rows(pool, rows, &run);
7165        return;
7166    }
7167    let run = move |start: usize, end: usize| {
7168        let mut buf = vec![0f32; cols];
7169        for r in start..end {
7170            q1t_dequant_row(bytes, r, gpr, rp_off, ent_off, has_ov, &mut buf);
7171            for bi in 0..b {
7172                let xr = &xs[bi * cols..(bi + 1) * cols];
7173                let mut acc = 0f32;
7174                for j in 0..cols {
7175                    acc += buf[j] * xr[j];
7176                }
7177                unsafe { *out_addr.at(bi * rows + r) = acc };
7178            }
7179        }
7180    };
7181    dispatch_rows(pool, rows, &run);
7182}
7183
7184fn q1_matvec(
7185    bytes: &[u8],
7186    x: &[f32],
7187    rows: usize,
7188    cols: usize,
7189    out: &mut [f32],
7190    pool: Option<&Pool>,
7191) {
7192    debug_assert_eq!(out.len(), rows);
7193    let gpr = cols / GROUP_SIZE;
7194    let out_addr = SendMut(out.as_mut_ptr());
7195    if a8w8_enabled() {
7196        let act = split_act(x);
7197        let gsum = q1_group_sums(&act.xq, gpr);
7198        let (act, gsum) = (&act, &gsum);
7199        let run = move |start: usize, end: usize| {
7200            q1_range_a8w8(bytes, gpr, act, gsum, out_addr, start, end)
7201        };
7202        dispatch_rows(pool, rows, &run);
7203        return;
7204    }
7205    let run = move |start: usize, end: usize| q1_range_f32(bytes, gpr, x, out_addr, start, end);
7206    dispatch_rows(pool, rows, &run);
7207}
7208
7209/// Fused two-input q1 matvec (weights read once per pair).
7210#[allow(clippy::too_many_arguments)]
7211fn q1_matvec2(
7212    bytes: &[u8],
7213    x1: &[f32],
7214    x2: &[f32],
7215    rows: usize,
7216    cols: usize,
7217    o1: &mut [f32],
7218    o2: &mut [f32],
7219    pool: Option<&Pool>,
7220) {
7221    let gpr = cols / GROUP_SIZE;
7222    let p1 = SendMut(o1.as_mut_ptr());
7223    let p2 = SendMut(o2.as_mut_ptr());
7224    if a8w8_enabled() {
7225        let a1 = split_act(x1);
7226        let a2 = split_act(x2);
7227        let g1 = q1_group_sums(&a1.xq, gpr);
7228        let g2 = q1_group_sums(&a2.xq, gpr);
7229        let (a1, a2, g1, g2) = (&a1, &a2, &g1, &g2);
7230        let run = move |start: usize, end: usize| {
7231            for r in start..end {
7232                let mut v1 = dot_q1_row_i8(bytes, r, gpr, &a1.xq, g1) * a1.sx;
7233                let mut v2 = dot_q1_row_i8(bytes, r, gpr, &a2.xq, g2) * a2.sx;
7234                for &(j, xv) in &a1.outliers {
7235                    let (w, s) = q1_outlier(bytes, r, gpr, j);
7236                    v1 += w * s * xv;
7237                }
7238                for &(j, xv) in &a2.outliers {
7239                    let (w, s) = q1_outlier(bytes, r, gpr, j);
7240                    v2 += w * s * xv;
7241                }
7242                // SAFETY: disjoint row ranges per worker.
7243                unsafe {
7244                    *p1.at(r) = v1;
7245                    *p2.at(r) = v2;
7246                }
7247            }
7248        };
7249        dispatch_rows(pool, rows, &run);
7250        return;
7251    }
7252    let run = move |start: usize, end: usize| {
7253        for r in start..end {
7254            // SAFETY: disjoint row ranges per worker.
7255            unsafe {
7256                *p1.at(r) = q1_row_exact(bytes, r, gpr, x1);
7257                *p2.at(r) = q1_row_exact(bytes, r, gpr, x2);
7258            }
7259        }
7260    };
7261    dispatch_rows(pool, rows, &run);
7262}
7263
7264/// Batched q1 matmat: each row's tiles stream once per microbatch.
7265#[allow(clippy::too_many_arguments)]
7266fn q1_matmat(
7267    bytes: &[u8],
7268    xs_all: &[f32],
7269    b: usize,
7270    rows: usize,
7271    cols: usize,
7272    out: &mut [f32],
7273    pool: Option<&Pool>,
7274) {
7275    debug_assert_eq!(out.len(), b * rows);
7276    let gpr = cols / GROUP_SIZE;
7277    let out_addr = SendMut(out.as_mut_ptr());
7278    if a8w8_enabled() {
7279        let acts: Vec<(SplitAct, Vec<i32>)> = (0..b)
7280            .map(|bi| {
7281                let act = split_act(&xs_all[bi * cols..(bi + 1) * cols]);
7282                let gsum = q1_group_sums(&act.xq, gpr);
7283                (act, gsum)
7284            })
7285            .collect();
7286        let acts = &acts;
7287        #[cfg(target_arch = "x86_64")]
7288        let blocked_ok = avx2_enabled()
7289            && blocked_enabled();
7290        #[cfg(target_arch = "aarch64")]
7291        let blocked_ok = sdot_enabled()
7292            && blocked_enabled();
7293        let run = move |start: usize, end: usize| {
7294            for r in start..end {
7295                let mut bi = 0usize;
7296                // Blocked 1×4: the unpacked bit mask serves four
7297                // activation streams per group.
7298                #[cfg(target_arch = "aarch64")]
7299                if blocked_ok {
7300                    while bi + 4 <= acts.len() {
7301                        let xs = [
7302                            acts[bi].0.xq.as_slice(),
7303                            acts[bi + 1].0.xq.as_slice(),
7304                            acts[bi + 2].0.xq.as_slice(),
7305                            acts[bi + 3].0.xq.as_slice(),
7306                        ];
7307                        let gs = [
7308                            acts[bi].1.as_slice(),
7309                            acts[bi + 1].1.as_slice(),
7310                            acts[bi + 2].1.as_slice(),
7311                            acts[bi + 3].1.as_slice(),
7312                        ];
7313                        let d = unsafe { dot_q1_row_1x4_sdot(bytes, r, gpr, xs, gs) };
7314                        for k in 0..4 {
7315                            let (act, _) = &acts[bi + k];
7316                            let mut acc = d[k] * act.sx;
7317                            for &(j, xv) in &act.outliers {
7318                                let (w, sc) = q1_outlier(bytes, r, gpr, j);
7319                                acc += w * sc * xv;
7320                            }
7321                            // SAFETY: disjoint (bi, r) cells per worker.
7322                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
7323                        }
7324                        bi += 4;
7325                    }
7326                }
7327                #[cfg(target_arch = "x86_64")]
7328                if blocked_ok {
7329                    while bi + 4 <= acts.len() {
7330                        let xs = [
7331                            acts[bi].0.xq.as_slice(),
7332                            acts[bi + 1].0.xq.as_slice(),
7333                            acts[bi + 2].0.xq.as_slice(),
7334                            acts[bi + 3].0.xq.as_slice(),
7335                        ];
7336                        let gs = [
7337                            acts[bi].1.as_slice(),
7338                            acts[bi + 1].1.as_slice(),
7339                            acts[bi + 2].1.as_slice(),
7340                            acts[bi + 3].1.as_slice(),
7341                        ];
7342                        let d = unsafe {
7343                            if vnni_tiles_enabled() {
7344                                dot_q1_row_1x4_vnni(bytes, r, gpr, xs, gs)
7345                            } else {
7346                                dot_q1_row_1x4_avx2(bytes, r, gpr, xs, gs)
7347                            }
7348                        };
7349                        for k in 0..4 {
7350                            let (act, _) = &acts[bi + k];
7351                            let mut acc = d[k] * act.sx;
7352                            for &(j, xv) in &act.outliers {
7353                                let (w, sc) = q1_outlier(bytes, r, gpr, j);
7354                                acc += w * sc * xv;
7355                            }
7356                            // SAFETY: disjoint (bi, r) cells per worker.
7357                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
7358                        }
7359                        bi += 4;
7360                    }
7361                }
7362                while bi < acts.len() {
7363                    let (act, gsum) = &acts[bi];
7364                    let mut acc = dot_q1_row_i8(bytes, r, gpr, &act.xq, gsum) * act.sx;
7365                    for &(j, xv) in &act.outliers {
7366                        let (w, s) = q1_outlier(bytes, r, gpr, j);
7367                        acc += w * s * xv;
7368                    }
7369                    // SAFETY: disjoint (bi, r) cells per worker range.
7370                    unsafe { *out_addr.at(bi * rows + r) = acc };
7371                    bi += 1;
7372                }
7373            }
7374        };
7375        dispatch_rows(pool, rows, &run);
7376        return;
7377    }
7378    let run = move |start: usize, end: usize| {
7379        for r in start..end {
7380            for bi in 0..b {
7381                let x = &xs_all[bi * cols..(bi + 1) * cols];
7382                // SAFETY: disjoint (bi, r) cells per worker range.
7383                unsafe { *out_addr.at(bi * rows + r) = q1_row_exact(bytes, r, gpr, x) };
7384            }
7385        }
7386    };
7387    dispatch_rows(pool, rows, &run);
7388}
7389
7390/// Fused q4_block matvec straight from the mapped bytes. SDOT path when
7391/// dotprod is available (port of vmfcore `dot_q4_block_sdot`, measured
7392/// +23% on q4 decode): nibbles → centered i8, int8×int8 `sdot` per
7393/// 32-group, exact outlier correction — the same A8W8 contract as q8.
7394/// `CMF_SDOT=0` keeps the exact scalar path.
7395fn q4matvec(
7396    bytes: &[u8],
7397    x: &[f32],
7398    rows: usize,
7399    cols: usize,
7400    out: &mut [f32],
7401    pool: Option<&Pool>,
7402) {
7403    debug_assert_eq!(out.len(), rows);
7404    let (packed, scales) = q4_split(bytes, rows, cols);
7405    let gpr = cols / GROUP_SIZE;
7406    let out_addr = SendMut(out.as_mut_ptr());
7407
7408    if a8w8_enabled() {
7409        let act = split_act(x);
7410        let run = move |start: usize, end: usize| {
7411            q4_range_a8w8(packed, scales, gpr, cols, &act, out_addr, start, end)
7412        };
7413        dispatch_rows(pool, rows, &run);
7414        return;
7415    }
7416
7417    let run =
7418        move |start: usize, end: usize| q4_range_f32(packed, scales, gpr, x, out_addr, start, end);
7419    dispatch_rows(pool, rows, &run);
7420}
7421
7422/// One q4 row via the A8W8 int8 path — SDOT on ARM, AVX2 maddubs on
7423/// x86 (scalar fallback is unreachable: callers gate on a8w8_enabled).
7424#[inline]
7425#[allow(unreachable_code)]
7426/// One UNPACKED q4 row (centered i8 in `buf`) against four activation
7427/// streams: the 32-byte weight chunk and its abs() load once per group,
7428/// the per-group f16 scale decodes once — four maddubs+reduce chains
7429/// instead of four full (load, abs, dot) rounds.
7430#[cfg(target_arch = "x86_64")]
7431#[target_feature(enable = "avx2")]
7432unsafe fn dot_q4b_row_1x4_avx2(
7433    buf: &[u8],
7434    scales: &[u8],
7435    g0: usize,
7436    gpr: usize,
7437    xs: [&[i8]; 4],
7438) -> [f32; 4] {
7439    // SAFETY: callers uphold buffer contracts (buf.len() == gpr·32).
7440    unsafe {
7441        use core::arch::x86_64::*;
7442        let ones = _mm256_set1_epi16(1);
7443        let mut acc = [0f32; 4];
7444        for gi in 0..gpr {
7445            let s = f16_to_f32(u16::from_le_bytes([
7446                scales[(g0 + gi) * 2],
7447                scales[(g0 + gi) * 2 + 1],
7448            ]));
7449            let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7450            let aw = _mm256_abs_epi8(w);
7451            for (k, xq) in xs.iter().enumerate() {
7452                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7453                let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
7454                let d = _mm256_madd_epi16(p16, ones);
7455                let hi128 = _mm256_extracti128_si256::<1>(d);
7456                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
7457                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7458                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7459                acc[k] += _mm_cvtsi128_si32(s32) as f32 * s;
7460            }
7461        }
7462        acc
7463    }
7464}
7465
7466/// VNNI twin of `dot_q4b_row_1x4_avx2` (see `dpbusd_hsum`).
7467#[cfg(target_arch = "x86_64")]
7468#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
7469unsafe fn dot_q4b_row_1x4_vnni(
7470    buf: &[u8],
7471    scales: &[u8],
7472    g0: usize,
7473    gpr: usize,
7474    xs: [&[i8]; 4],
7475) -> [f32; 4] {
7476    // SAFETY: callers uphold buffer contracts (buf.len() == gpr·32).
7477    unsafe {
7478        use core::arch::x86_64::*;
7479        let mut acc = [0f32; 4];
7480        for gi in 0..gpr {
7481            let s = f16_to_f32(u16::from_le_bytes([
7482                scales[(g0 + gi) * 2],
7483                scales[(g0 + gi) * 2 + 1],
7484            ]));
7485            let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7486            let aw = _mm256_abs_epi8(w);
7487            for (k, xq) in xs.iter().enumerate() {
7488                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7489                let d = dpbusd_hsum(aw, _mm256_sign_epi8(x, w));
7490                acc[k] += d as f32 * s;
7491            }
7492        }
7493        acc
7494    }
7495}
7496
7497/// The vbit flavor of the blocked 1×4: the per-activation A8W8 scale
7498/// folds in PER GROUP as `(d·sx)·s` — bit-matching the single-matvec
7499/// accumulation order (the q4_block flavor applies sx once at the end,
7500/// matching ITS single path; the two conventions are historical and
7501/// each blocked leg must mirror its own).
7502#[cfg(target_arch = "x86_64")]
7503#[target_feature(enable = "avx2")]
7504unsafe fn dot_q4b_row_1x4_sx_avx2(
7505    buf: &[u8],
7506    scales: &[u8],
7507    g0: usize,
7508    gpr: usize,
7509    xs: [&[i8]; 4],
7510    sxs: [f32; 4],
7511) -> [f32; 4] {
7512    // SAFETY: callers uphold buffer contracts (buf.len() == gpr·32).
7513    unsafe {
7514        use core::arch::x86_64::*;
7515        let ones = _mm256_set1_epi16(1);
7516        let mut acc = [0f32; 4];
7517        for gi in 0..gpr {
7518            let s = f16_to_f32(u16::from_le_bytes([
7519                scales[(g0 + gi) * 2],
7520                scales[(g0 + gi) * 2 + 1],
7521            ]));
7522            let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7523            let aw = _mm256_abs_epi8(w);
7524            for (k, xq) in xs.iter().enumerate() {
7525                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7526                let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
7527                let d = _mm256_madd_epi16(p16, ones);
7528                let hi128 = _mm256_extracti128_si256::<1>(d);
7529                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
7530                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7531                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7532                acc[k] += (_mm_cvtsi128_si32(s32) as f32 * sxs[k]) * s;
7533            }
7534        }
7535        acc
7536    }
7537}
7538
7539/// VNNI twin of `dot_q4b_row_1x4_sx_avx2` (see `dpbusd_hsum`; the
7540/// per-group `(d·sx)·s` fold mirrors the vbit single path).
7541#[cfg(target_arch = "x86_64")]
7542#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
7543unsafe fn dot_q4b_row_1x4_sx_vnni(
7544    buf: &[u8],
7545    scales: &[u8],
7546    g0: usize,
7547    gpr: usize,
7548    xs: [&[i8]; 4],
7549    sxs: [f32; 4],
7550) -> [f32; 4] {
7551    // SAFETY: callers uphold buffer contracts (buf.len() == gpr·32).
7552    unsafe {
7553        use core::arch::x86_64::*;
7554        let mut acc = [0f32; 4];
7555        for gi in 0..gpr {
7556            let s = f16_to_f32(u16::from_le_bytes([
7557                scales[(g0 + gi) * 2],
7558                scales[(g0 + gi) * 2 + 1],
7559            ]));
7560            let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7561            let aw = _mm256_abs_epi8(w);
7562            for (k, xq) in xs.iter().enumerate() {
7563                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7564                let d = dpbusd_hsum(aw, _mm256_sign_epi8(x, w));
7565                acc[k] += (d as f32 * sxs[k]) * s;
7566            }
7567        }
7568        acc
7569    }
7570}
7571
7572#[allow(unreachable_code)]
7573fn dot_q4_row_i8(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
7574    #[cfg(target_arch = "aarch64")]
7575    unsafe {
7576        return dot_q4_row_sdot(packed, scales, g0, gpr, xq);
7577    }
7578    #[cfg(target_arch = "x86_64")]
7579    unsafe {
7580        return dot_q4_row_avx2(packed, scales, g0, gpr, xq);
7581    }
7582    let mut acc = 0f32;
7583    for gi in 0..gpr {
7584        let g = g0 + gi;
7585        let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
7586        let mut d = 0i32;
7587        for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
7588            d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
7589                + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
7590        }
7591        acc += d as f32 * s;
7592    }
7593    acc
7594}
7595
7596/// Two-activation q4 row via the A8W8 int8 path (see `dot_q4_row_i8`).
7597#[inline]
7598#[allow(unreachable_code)]
7599fn dot_q4_row_i8_2(
7600    packed: &[u8],
7601    scales: &[u8],
7602    g0: usize,
7603    gpr: usize,
7604    xq1: &[i8],
7605    xq2: &[i8],
7606) -> (f32, f32) {
7607    #[cfg(target_arch = "aarch64")]
7608    unsafe {
7609        return dot_q4_row_sdot2(packed, scales, g0, gpr, xq1, xq2);
7610    }
7611    #[cfg(target_arch = "x86_64")]
7612    unsafe {
7613        return dot_q4_row_avx2_2(packed, scales, g0, gpr, xq1, xq2);
7614    }
7615    (
7616        dot_q4_row_i8(packed, scales, g0, gpr, xq1),
7617        dot_q4_row_i8(packed, scales, g0, gpr, xq2),
7618    )
7619}
7620
7621/// One q4 row range via SDOT (kernel body of `q4matvec`, extracted so
7622/// multi-matrix jobs can drive it for several tensors in one dispatch).
7623#[allow(clippy::too_many_arguments)]
7624fn q4_range_a8w8(
7625    packed: &[u8],
7626    scales: &[u8],
7627    gpr: usize,
7628    cols: usize,
7629    act: &SplitAct,
7630    out: SendMut,
7631    start: usize,
7632    end: usize,
7633) {
7634    for r in start..end {
7635        let mut acc = dot_q4_row_i8(packed, scales, r * gpr, gpr, &act.xq) * act.sx;
7636        // xq is zeroed at outlier slots — add the exact terms.
7637        for &(j, xv) in &act.outliers {
7638            let flat = r * cols + j;
7639            let byte = packed[flat / 2];
7640            let nib = if flat & 1 == 0 {
7641                byte & 0x0F
7642            } else {
7643                byte >> 4
7644            };
7645            let s = f16_to_f32(u16::from_le_bytes([
7646                scales[(flat / GROUP_SIZE) * 2],
7647                scales[(flat / GROUP_SIZE) * 2 + 1],
7648            ]));
7649            acc += ((nib as i32 - 8) as f32) * s * xv;
7650        }
7651        // SAFETY: disjoint row ranges per worker.
7652        unsafe { *out.at(r) = acc };
7653    }
7654}
7655
7656/// Two-input q4 row range via the A8W8 int8 path — kernel body of
7657/// `q4matvec2`, extracted for pair multi-matrix jobs.
7658#[allow(clippy::too_many_arguments)]
7659fn q4_range2_a8w8(
7660    packed: &[u8],
7661    scales: &[u8],
7662    gpr: usize,
7663    cols: usize,
7664    a1: &SplitAct,
7665    a2: &SplitAct,
7666    p1: SendMut,
7667    p2: SendMut,
7668    start: usize,
7669    end: usize,
7670) {
7671    for r in start..end {
7672        let (s1, s2) = dot_q4_row_i8_2(packed, scales, r * gpr, gpr, &a1.xq, &a2.xq);
7673        let mut acc1 = s1 * a1.sx;
7674        let mut acc2 = s2 * a2.sx;
7675        // xq is zeroed at outlier slots — add the exact terms.
7676        let fix = |outliers: &[(usize, f32)], acc: &mut f32| {
7677            for &(j, xv) in outliers {
7678                let flat = r * cols + j;
7679                let byte = packed[flat / 2];
7680                let nib = if flat & 1 == 0 {
7681                    byte & 0x0F
7682                } else {
7683                    byte >> 4
7684                };
7685                let s = f16_to_f32(u16::from_le_bytes([
7686                    scales[(flat / GROUP_SIZE) * 2],
7687                    scales[(flat / GROUP_SIZE) * 2 + 1],
7688                ]));
7689                *acc += ((nib as i32 - 8) as f32) * s * xv;
7690            }
7691        };
7692        fix(&a1.outliers, &mut acc1);
7693        fix(&a2.outliers, &mut acc2);
7694        // SAFETY: disjoint row ranges per worker.
7695        unsafe {
7696            *p1.at(r) = acc1;
7697            *p2.at(r) = acc2;
7698        }
7699    }
7700}
7701
7702/// Exact scalar q4 row range (same extraction, non-SDOT path).
7703fn q4_range_f32(
7704    packed: &[u8],
7705    scales: &[u8],
7706    gpr: usize,
7707    x: &[f32],
7708    out: SendMut,
7709    start: usize,
7710    end: usize,
7711) {
7712    for r in start..end {
7713        let mut acc = 0f32;
7714        for gi in 0..gpr {
7715            let g = r * gpr + gi;
7716            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
7717            let pk = &packed[g * 16..(g + 1) * 16];
7718            let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
7719            let mut ga = 0f32;
7720            for (k, &b) in pk.iter().enumerate() {
7721                ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
7722                    + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
7723            }
7724            acc += ga * s;
7725        }
7726        // SAFETY: disjoint row ranges per worker.
7727        unsafe { *out.at(r) = acc };
7728    }
7729}
7730
7731/// Fused two-input q4 matvec: nibbles are unpacked ONCE per group and
7732/// dotted against both activations (was: two full matvecs — double
7733/// weight traffic). Per-lane math matches `q4matvec` exactly.
7734#[allow(clippy::too_many_arguments)]
7735fn q4matvec2(
7736    bytes: &[u8],
7737    x1: &[f32],
7738    x2: &[f32],
7739    rows: usize,
7740    cols: usize,
7741    o1: &mut [f32],
7742    o2: &mut [f32],
7743    pool: Option<&Pool>,
7744) {
7745    debug_assert_eq!(o1.len(), rows);
7746    debug_assert_eq!(o2.len(), rows);
7747    let (packed, scales) = q4_split(bytes, rows, cols);
7748    let gpr = cols / GROUP_SIZE;
7749
7750    if a8w8_enabled() {
7751        let a1 = split_act(x1);
7752        let a2 = split_act(x2);
7753        let p1 = SendMut(o1.as_mut_ptr());
7754        let p2 = SendMut(o2.as_mut_ptr());
7755        let run = move |start: usize, end: usize| {
7756            q4_range2_a8w8(packed, scales, gpr, cols, &a1, &a2, p1, p2, start, end)
7757        };
7758        dispatch_rows(pool, rows, &run);
7759        return;
7760    }
7761
7762    let p1 = SendMut(o1.as_mut_ptr());
7763    let p2 = SendMut(o2.as_mut_ptr());
7764    let run = move |start: usize, end: usize| {
7765        q4_range2_f32(packed, scales, gpr, x1, x2, p1, p2, start, end)
7766    };
7767    dispatch_rows(pool, rows, &run);
7768}
7769
7770/// Two-input exact scalar q4 row range (same extraction).
7771#[allow(clippy::too_many_arguments)]
7772fn q4_range2_f32(
7773    packed: &[u8],
7774    scales: &[u8],
7775    gpr: usize,
7776    x1: &[f32],
7777    x2: &[f32],
7778    p1: SendMut,
7779    p2: SendMut,
7780    start: usize,
7781    end: usize,
7782) {
7783    for r in start..end {
7784        let (mut acc1, mut acc2) = (0f32, 0f32);
7785        for gi in 0..gpr {
7786            let g = r * gpr + gi;
7787            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
7788            let pk = &packed[g * 16..(g + 1) * 16];
7789            let x1g = &x1[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
7790            let x2g = &x2[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
7791            let (mut g1, mut g2) = (0f32, 0f32);
7792            for (k, &b) in pk.iter().enumerate() {
7793                let wl = (b & 0x0F) as f32 - 8.0;
7794                let wh = ((b >> 4) & 0x0F) as f32 - 8.0;
7795                g1 += wl * x1g[k * 2] + wh * x1g[k * 2 + 1];
7796                g2 += wl * x2g[k * 2] + wh * x2g[k * 2 + 1];
7797            }
7798            acc1 += g1 * s;
7799            acc2 += g2 * s;
7800        }
7801        // SAFETY: disjoint row ranges per worker.
7802        unsafe {
7803            *p1.at(r) = acc1;
7804            *p2.at(r) = acc2;
7805        }
7806    }
7807}
7808
7809thread_local! {
7810    /// Per-worker decoded-row scratch for the batched q4/vbit kernels
7811    /// (centered i8 for SDOT, f32 for the exact/scalar paths).
7812    static ROW_I8: std::cell::RefCell<Vec<u8>> = const { std::cell::RefCell::new(Vec::new()) };
7813    static ROW_F32: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
7814}
7815
7816/// Batched q4 matmat: each weight row is unpacked from the mmap ONCE
7817/// and dotted against ALL b activations (prefill used to fall back to b
7818/// full matvecs — b× weight traffic and b× nibble decode). Per-position
7819/// math matches `q4matvec` exactly: same group order, same accumulation.
7820/// `out` is row-major [b, rows] like `qmatmat`.
7821#[allow(clippy::too_many_arguments)]
7822fn q4matmat(
7823    bytes: &[u8],
7824    xs_all: &[f32],
7825    b: usize,
7826    rows: usize,
7827    cols: usize,
7828    out: &mut [f32],
7829    pool: Option<&Pool>,
7830) {
7831    debug_assert_eq!(xs_all.len(), b * cols);
7832    debug_assert_eq!(out.len(), b * rows);
7833    let (packed, scales) = q4_split(bytes, rows, cols);
7834    let gpr = cols / GROUP_SIZE;
7835    let gscale = |g: usize| f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
7836
7837    if a8w8_enabled() {
7838        let acts: Vec<SplitAct> = (0..b)
7839            .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
7840            .collect();
7841        let acts = &acts;
7842        let out_addr = SendMut(out.as_mut_ptr());
7843        let run = move |start: usize, end: usize| {
7844            ROW_I8.with(|rb| {
7845                let mut buf = rb.borrow_mut();
7846                buf.resize(cols, 0);
7847                for r in start..end {
7848                    // Unpack the row's nibbles to centered i8 once
7849                    // (element 2k = low nibble, 2k+1 = high — flat order,
7850                    // same as dot_q4_row_sdot's zip).
7851                    for gi in 0..gpr {
7852                        let g = r * gpr + gi;
7853                        for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
7854                            buf[gi * GROUP_SIZE + k * 2] = ((bt & 0x0F) as i32 - 8) as i8 as u8;
7855                            buf[gi * GROUP_SIZE + k * 2 + 1] =
7856                                (((bt >> 4) & 0x0F) as i32 - 8) as i8 as u8;
7857                        }
7858                    }
7859                    let mut bi = 0usize;
7860                    #[cfg(target_arch = "x86_64")]
7861                    if avx2_enabled()
7862                        && blocked_enabled()
7863                    {
7864                        while bi + 4 <= acts.len() {
7865                            let xs = [
7866                                acts[bi].xq.as_slice(),
7867                                acts[bi + 1].xq.as_slice(),
7868                                acts[bi + 2].xq.as_slice(),
7869                                acts[bi + 3].xq.as_slice(),
7870                            ];
7871                            let d = unsafe {
7872                                if vnni_tiles_enabled() {
7873                                    dot_q4b_row_1x4_vnni(&buf, scales, r * gpr, gpr, xs)
7874                                } else {
7875                                    dot_q4b_row_1x4_avx2(&buf, scales, r * gpr, gpr, xs)
7876                                }
7877                            };
7878                            for k in 0..4 {
7879                                let act = &acts[bi + k];
7880                                let mut acc = d[k] * act.sx;
7881                                for &(j, xv) in &act.outliers {
7882                                    acc += (buf[j] as i8) as f32
7883                                        * gscale((r * cols + j) / GROUP_SIZE)
7884                                        * xv;
7885                                }
7886                                // SAFETY: disjoint (bi, r) cells per worker.
7887                                unsafe { *out_addr.at((bi + k) * rows + r) = acc };
7888                            }
7889                            bi += 4;
7890                        }
7891                    }
7892                    while bi < acts.len() {
7893                        let act = &acts[bi];
7894                        let mut acc = 0f32;
7895                        for gi in 0..gpr {
7896                            let d = dot_i8_i8(
7897                                &buf[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
7898                                &act.xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
7899                            );
7900                            acc += d as f32 * gscale(r * gpr + gi);
7901                        }
7902                        acc *= act.sx;
7903                        // xq is zeroed at outlier slots — exact terms.
7904                        for &(j, xv) in &act.outliers {
7905                            acc += (buf[j] as i8) as f32 * gscale((r * cols + j) / GROUP_SIZE) * xv;
7906                        }
7907                        // SAFETY: disjoint (bi, r) cells per worker row range.
7908                        unsafe { *out_addr.at(bi * rows + r) = acc };
7909                        bi += 1;
7910                    }
7911                }
7912            })
7913        };
7914        dispatch_rows(pool, rows, &run);
7915        return;
7916    }
7917
7918    let out_addr = SendMut(out.as_mut_ptr());
7919    let run = move |start: usize, end: usize| {
7920        ROW_F32.with(|rb| {
7921            let mut buf = rb.borrow_mut();
7922            buf.resize(cols, 0.0);
7923            for r in start..end {
7924                // Decode raw (nib − 8) values once; scales stay per-group
7925                // so the accumulation order matches q4matvec bit-for-bit.
7926                for gi in 0..gpr {
7927                    let g = r * gpr + gi;
7928                    for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
7929                        buf[gi * GROUP_SIZE + k * 2] = (bt & 0x0F) as f32 - 8.0;
7930                        buf[gi * GROUP_SIZE + k * 2 + 1] = ((bt >> 4) & 0x0F) as f32 - 8.0;
7931                    }
7932                }
7933                for bi in 0..b {
7934                    let x = &xs_all[bi * cols..(bi + 1) * cols];
7935                    let mut acc = 0f32;
7936                    for gi in 0..gpr {
7937                        let mut ga = 0f32;
7938                        // Pairwise (lo + hi) addition, matching
7939                        // q4matvec's `ga += lo·x + hi·x` shape exactly —
7940                        // a flat one-per-element loop rounds differently
7941                        // and broke bit-parity on the scalar (x86) path.
7942                        for k in 0..GROUP_SIZE / 2 {
7943                            let e = gi * GROUP_SIZE + k * 2;
7944                            ga += buf[e] * x[e] + buf[e + 1] * x[e + 1];
7945                        }
7946                        acc += ga * gscale(r * gpr + gi);
7947                    }
7948                    // SAFETY: disjoint (bi, r) cells per worker row range.
7949                    unsafe { *out_addr.at(bi * rows + r) = acc };
7950                }
7951            }
7952        })
7953    };
7954    dispatch_rows(pool, rows, &run);
7955}
7956
7957/// Batched vbit matmat: each variable-bit row is decoded from the mmap
7958/// ONCE for the whole microbatch. Same per-position math as
7959/// `vbitmatvec` (SDOT A8W8 with exact outliers / exact f32 for b=8 rows
7960/// and the scalar path).
7961#[allow(clippy::too_many_arguments)]
7962fn vbitmatmat(
7963    bytes: &[u8],
7964    offsets: &[usize],
7965    xs_all: &[f32],
7966    b: usize,
7967    rows: usize,
7968    cols: usize,
7969    out: &mut [f32],
7970    pool: Option<&Pool>,
7971) {
7972    debug_assert_eq!(xs_all.len(), b * cols);
7973    debug_assert_eq!(out.len(), b * rows);
7974    debug_assert_eq!(offsets.len(), rows + 1);
7975    let ng = cols / GROUP_SIZE;
7976    let bits = &bytes[..rows];
7977    let sc_off = rows;
7978    let gscale = |r: usize, g: usize| {
7979        let so = (r * ng + g) * 2;
7980        f16_to_f32(u16::from_le_bytes([
7981            bytes[sc_off + so],
7982            bytes[sc_off + so + 1],
7983        ]))
7984    };
7985
7986    // Decode row r's raw (u − L) values into `dst` (f32, unscaled).
7987    let decode_f32 = |r: usize, dst: &mut [f32]| {
7988        let bw = bits[r] as usize;
7989        let l = ((1i32 << (bw - 1)) - 1) as f32;
7990        let data = &bytes[offsets[r]..offsets[r + 1]];
7991        let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
7992        for d in dst.iter_mut() {
7993            while nbits < bw {
7994                acc = (acc << 8) | data[idx] as u64;
7995                idx += 1;
7996                nbits += 8;
7997            }
7998            let u = ((acc >> (nbits - bw)) & ((1u64 << bw) - 1)) as f32;
7999            nbits -= bw;
8000            *d = u - l;
8001        }
8002    };
8003
8004    if a8w8_enabled() {
8005        let acts: Vec<SplitAct> = (0..b)
8006            .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
8007            .collect();
8008        let acts = &acts;
8009        let out_addr = SendMut(out.as_mut_ptr());
8010        let run = move |start: usize, end: usize| {
8011            for r in start..end {
8012                let bw = bits[r] as usize;
8013                if bw == 8 {
8014                    // u−L reaches 128 → no i8 path; decode once, exact
8015                    // f32 dots for every position (same as vbitmatvec).
8016                    ROW_F32.with(|rb| {
8017                        let mut buf = rb.borrow_mut();
8018                        buf.resize(cols, 0.0);
8019                        decode_f32(r, &mut buf);
8020                        for bi in 0..b {
8021                            let x = &xs_all[bi * cols..(bi + 1) * cols];
8022                            let mut dot = 0f32;
8023                            for g in 0..ng {
8024                                let mut gd = 0f32;
8025                                for k in 0..GROUP_SIZE {
8026                                    gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
8027                                }
8028                                dot += gd * gscale(r, g);
8029                            }
8030                            // SAFETY: disjoint (bi, r) cells per worker range.
8031                            unsafe { *out_addr.at(bi * rows + r) = dot };
8032                        }
8033                    });
8034                    continue;
8035                }
8036                let l = (1i32 << (bw - 1)) - 1;
8037                let data = &bytes[offsets[r]..offsets[r + 1]];
8038                ROW_I8.with(|rb| {
8039                    let mut buf = rb.borrow_mut();
8040                    buf.resize(cols, 0);
8041                    #[inline(always)]
8042                    fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
8043                        for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
8044                            let u = unpack8::<B>(&data[blk * B..]);
8045                            for k in 0..8 {
8046                                chunk[k] = (u[k] - l) as i8 as u8;
8047                            }
8048                        }
8049                    }
8050                    match bw {
8051                        3 => fill::<3>(data, l, &mut buf),
8052                        4 => vbit_fill4(data, &mut buf),
8053                        5 => fill::<5>(data, l, &mut buf),
8054                        6 => fill::<6>(data, l, &mut buf),
8055                        _ => unreachable!("vbit bit-width {bw} (validated at load)"),
8056                    }
8057                    let mut bi = 0usize;
8058                    // The vbit scale table shares q4_block's layout
8059                    // (contiguous f16 per (row·ng + g)), so the same
8060                    // blocked 1×4 kernel serves the decoded row.
8061                    #[cfg(target_arch = "x86_64")]
8062                    if avx2_enabled()
8063                        && blocked_enabled()
8064                    {
8065                        while bi + 4 <= acts.len() {
8066                            let xs = [
8067                                acts[bi].xq.as_slice(),
8068                                acts[bi + 1].xq.as_slice(),
8069                                acts[bi + 2].xq.as_slice(),
8070                                acts[bi + 3].xq.as_slice(),
8071                            ];
8072                            let sxs = [
8073                                acts[bi].sx,
8074                                acts[bi + 1].sx,
8075                                acts[bi + 2].sx,
8076                                acts[bi + 3].sx,
8077                            ];
8078                            let d = unsafe {
8079                                if vnni_tiles_enabled() {
8080                                    dot_q4b_row_1x4_sx_vnni(
8081                                        &buf,
8082                                        &bytes[sc_off..],
8083                                        r * ng,
8084                                        ng,
8085                                        xs,
8086                                        sxs,
8087                                    )
8088                                } else {
8089                                    dot_q4b_row_1x4_sx_avx2(
8090                                        &buf,
8091                                        &bytes[sc_off..],
8092                                        r * ng,
8093                                        ng,
8094                                        xs,
8095                                        sxs,
8096                                    )
8097                                }
8098                            };
8099                            for k in 0..4 {
8100                                let act = &acts[bi + k];
8101                                let mut dot = d[k];
8102                                for &(j, xv) in &act.outliers {
8103                                    dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
8104                                }
8105                                // SAFETY: disjoint (bi, r) cells per worker.
8106                                unsafe { *out_addr.at((bi + k) * rows + r) = dot };
8107                            }
8108                            bi += 4;
8109                        }
8110                    }
8111                    while bi < acts.len() {
8112                        let act = &acts[bi];
8113                        let mut dot = 0f32;
8114                        for g in 0..ng {
8115                            let d = dot_i8_i8(
8116                                &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
8117                                &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
8118                            ) as f32
8119                                * act.sx;
8120                            dot += d * gscale(r, g);
8121                        }
8122                        for &(j, xv) in &act.outliers {
8123                            dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
8124                        }
8125                        // SAFETY: disjoint (bi, r) cells per worker range.
8126                        unsafe { *out_addr.at(bi * rows + r) = dot };
8127                        bi += 1;
8128                    }
8129                });
8130            }
8131        };
8132        dispatch_rows(pool, rows, &run);
8133        return;
8134    }
8135
8136    let out_addr = SendMut(out.as_mut_ptr());
8137    let run = move |start: usize, end: usize| {
8138        ROW_F32.with(|rb| {
8139            let mut buf = rb.borrow_mut();
8140            buf.resize(cols, 0.0);
8141            for r in start..end {
8142                decode_f32(r, &mut buf);
8143                for bi in 0..b {
8144                    let x = &xs_all[bi * cols..(bi + 1) * cols];
8145                    let mut dot = 0f32;
8146                    for g in 0..ng {
8147                        let mut gd = 0f32;
8148                        for k in 0..GROUP_SIZE {
8149                            gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
8150                        }
8151                        dot += gd * gscale(r, g);
8152                    }
8153                    // SAFETY: disjoint (bi, r) cells per worker range.
8154                    unsafe { *out_addr.at(bi * rows + r) = dot };
8155                }
8156            }
8157        })
8158    };
8159    dispatch_rows(pool, rows, &run);
8160}
8161
8162/// Build a GPU batch job for a q8-family mapped tensor (primary
8163/// shard): prescaled input + directory coordinates. None → not
8164/// GPU-eligible, caller stays on the CPU.
8165pub(crate) fn gpu_batch_job<'a>(
8166    t: &'a QTensor,
8167    x: &[f32],
8168) -> Option<(std::sync::Arc<CmfModel>, crate::gpu::BatchJob<'a>)> {
8169    match t {
8170        QTensor::Mapped {
8171            model,
8172            idx,
8173            dtype: dt @ (TensorDtype::Q8Row | TensorDtype::Q8_2f),
8174            rows,
8175            cols,
8176            row_scale,
8177            col_field,
8178            ..
8179        } => Some((
8180            model.clone(),
8181            crate::gpu::BatchJob {
8182                idx: *idx,
8183                rows: *rows,
8184                cols: *cols,
8185                row_scale,
8186                xs: prescale(x, col_field, *dt).into_owned(),
8187                layout: crate::gpu::BatchLayout::Q8,
8188            },
8189        )),
8190        // q1: raw f32 activations, tile-embedded scales.
8191        QTensor::Mapped {
8192            model,
8193            idx,
8194            dtype: TensorDtype::Q1,
8195            rows,
8196            cols,
8197            ..
8198        } => Some((
8199            model.clone(),
8200            crate::gpu::BatchJob {
8201                idx: *idx,
8202                rows: *rows,
8203                cols: *cols,
8204                row_scale: &[],
8205                xs: x.to_vec(),
8206                layout: crate::gpu::BatchLayout::Q1,
8207            },
8208        )),
8209        // q4_tiled / q4tp: raw f32 activations; the scales live in the
8210        // payload (inline tiles / row ladder), so row_scale stays empty.
8211        // The GDN projection batch already runs these layouts on Metal —
8212        // this arm lets the attention QKV batch reach the same kernels.
8213        QTensor::Mapped {
8214            model,
8215            idx,
8216            dtype: dt @ (TensorDtype::Q4Tiled | TensorDtype::Q4TiledP),
8217            rows,
8218            cols,
8219            ..
8220        } => Some((
8221            model.clone(),
8222            crate::gpu::BatchJob {
8223                idx: *idx,
8224                rows: *rows,
8225                cols: *cols,
8226                row_scale: &[],
8227                xs: x.to_vec(),
8228                layout: if *dt == TensorDtype::Q4Tiled {
8229                    crate::gpu::BatchLayout::Q4t
8230                } else {
8231                    crate::gpu::BatchLayout::Q4tp
8232                },
8233            },
8234        )),
8235        _ => None,
8236    }
8237}
8238
8239thread_local! {
8240    static PRESCALE_BUF1: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
8241    static PRESCALE_BUF2: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
8242}
8243
8244pub(crate) fn prescale<'a>(
8245    x: &'a [f32],
8246    col_field: &[f32],
8247    dtype: TensorDtype,
8248) -> std::borrow::Cow<'a, [f32]> {
8249    if dtype == TensorDtype::Q8_2f {
8250        x.iter().zip(col_field).map(|(a, c)| a * c).collect()
8251    } else {
8252        std::borrow::Cow::Borrowed(x)
8253    }
8254}
8255
8256/// θ col-field fold for q8_2f activations. Borrowed pass-through for
8257/// every other dtype, using thread-local buffers to eliminate per-matvec allocations.
8258pub(crate) fn prescale_with<R, F: FnOnce(&[f32]) -> R>(
8259    x: &[f32],
8260    col_field: &[f32],
8261    dtype: TensorDtype,
8262    buf_id: u8,
8263    f: F,
8264) -> R {
8265    if dtype == TensorDtype::Q8_2f {
8266        if buf_id == 1 {
8267            PRESCALE_BUF1.with(|b| {
8268                let mut buf = b.borrow_mut();
8269                buf.clear();
8270                buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
8271                f(&buf)
8272            })
8273        } else {
8274            PRESCALE_BUF2.with(|b| {
8275                let mut buf = b.borrow_mut();
8276                buf.clear();
8277                buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
8278                f(&buf)
8279            })
8280        }
8281    } else {
8282        f(x)
8283    }
8284}
8285
8286// ───────────────────── x86-64 AVX2 kernels (roadmap этап 2) ─────────────────────
8287
8288/// AVX2+FMA available? Default ON when the CPU supports both;
8289/// `CMF_AVX2=0` disables (falls back to the autovectorized loops).
8290#[cfg(target_arch = "x86_64")]
8291pub(crate) fn avx2_enabled() -> bool {
8292    use std::sync::OnceLock;
8293    static ON: OnceLock<bool> = OnceLock::new();
8294    *ON.get_or_init(|| {
8295        std::env::var("CMF_AVX2").map(|v| v != "0").unwrap_or(true)
8296            && std::arch::is_x86_feature_detected!("avx2")
8297            && std::arch::is_x86_feature_detected!("fma")
8298    })
8299}
8300
8301/// AVX2 A8W8 allowed? The quantized-activation contract is switched by
8302/// the SAME env as the ARM SDOT path: `CMF_SDOT=0` keeps exact kernels
8303/// (the golden-parity exact gate relies on it) — AVX2 f32 kernels stay
8304/// active either way, they are exact (regrouped sums only).
8305#[cfg(target_arch = "x86_64")]
8306fn avx2_a8w8_enabled() -> bool {
8307    use std::sync::OnceLock;
8308    static ON: OnceLock<bool> = OnceLock::new();
8309    *ON.get_or_init(|| {
8310        avx2_enabled() && std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true)
8311    })
8312}
8313
8314/// A8W8 quantized-activation path available on THIS machine? One
8315/// switch across architectures: ARM dotprod (CMF_SDOT) or x86 AVX2
8316/// (CMF_AVX2 + the same CMF_SDOT exact-contract override).
8317#[inline]
8318pub(crate) fn a8w8_enabled() -> bool {
8319    #[cfg(target_arch = "aarch64")]
8320    {
8321        sdot_enabled()
8322    }
8323    #[cfg(target_arch = "x86_64")]
8324    {
8325        avx2_a8w8_enabled()
8326    }
8327    #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
8328    {
8329        false
8330    }
8331}
8332
8333/// int8·int8 dot dispatch: SDOT on ARM; AVX-512 VNNI (vpdpbusd) or AVX2
8334/// maddubs on x86. Callers are gated by `a8w8_enabled()`.
8335#[inline]
8336#[allow(unreachable_code)]
8337fn dot_i8_i8(w: &[u8], xq: &[i8]) -> i32 {
8338    #[cfg(target_arch = "aarch64")]
8339    unsafe {
8340        return dot_i8_sdot(w, xq);
8341    }
8342    #[cfg(target_arch = "x86_64")]
8343    unsafe {
8344        if avx512vnni_enabled() {
8345            return dot_i8_i8_vnni(w, xq);
8346        }
8347        return dot_i8_i8_avx2(w, xq);
8348    }
8349    w.iter()
8350        .zip(xq)
8351        .map(|(&a, &b)| (a as i8) as i32 * b as i32)
8352        .sum()
8353}
8354
8355/// AVX-512 VNNI available? (F+BW+VL+VNNI; `CMF_AVX512=0` falls back to
8356/// AVX2.) VL matters: short 32-byte groups (q4/vbit) ride the 256-bit
8357/// `vpdpbusd` encoding.
8358#[cfg(target_arch = "x86_64")]
8359fn avx512vnni_enabled() -> bool {
8360    use std::sync::OnceLock;
8361    static ON: OnceLock<bool> = OnceLock::new();
8362    *ON.get_or_init(|| {
8363        std::env::var("CMF_AVX512")
8364            .map(|v| v != "0")
8365            .unwrap_or(true)
8366            && std::arch::is_x86_feature_detected!("avx512f")
8367            && std::arch::is_x86_feature_detected!("avx512bw")
8368            && std::arch::is_x86_feature_detected!("avx512vl")
8369            && std::arch::is_x86_feature_detected!("avx512vnni")
8370    })
8371}
8372
8373/// Grouped-codec VNNI arms (the q4t/q4b/q1/q1t tile kernels): default
8374/// ON where AVX-512 VNNI exists (`CMF_VNNI_TILES=0` opt-out). Measured
8375/// on Ryzen 7950X (Zen4, 3 alternating process pairs, blocked GEMM
8376/// 4864×896 b=256): q4t 63→68 GF/s (+8%), q1 53→56 (+6%), q4b 72→75
8377/// (+4%) — consistent, no leg regressed. The tile kernels keep a
8378/// horizontal reduce per 32-weight group, so the `vpdpbusd` saving is
8379/// smaller than the long-dot q8 win (+13%), but it is real and free.
8380#[cfg(target_arch = "x86_64")]
8381fn vnni_tiles_enabled() -> bool {
8382    use std::sync::OnceLock;
8383    static ON: OnceLock<bool> = OnceLock::new();
8384    *ON.get_or_init(|| {
8385        std::env::var("CMF_VNNI_TILES")
8386            .map(|v| v != "0")
8387            .unwrap_or(true)
8388            && avx512vnni_enabled()
8389    })
8390}
8391
8392/// One 256-bit u8×i8 dot → i32 via `vpdpbusd` into a fresh accumulator
8393/// plus the same horizontal reduce the AVX2 kernels use. Products are
8394/// bounded (|w| ≤ 8 or ≤ 1), so maddubs never saturated — the i32 sum
8395/// is bit-identical to the maddubs+madd pair it replaces.
8396#[cfg(target_arch = "x86_64")]
8397#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
8398#[inline]
8399unsafe fn dpbusd_hsum(aw: core::arch::x86_64::__m256i, xs: core::arch::x86_64::__m256i) -> i32 {
8400    // SAFETY: pure register math.
8401    unsafe {
8402        use core::arch::x86_64::*;
8403        let d = _mm256_dpbusd_epi32(_mm256_setzero_si256(), aw, xs);
8404        let hi128 = _mm256_extracti128_si256::<1>(d);
8405        let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
8406        let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8407        let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8408        _mm_cvtsi128_si32(s32)
8409    }
8410}
8411
8412/// int8·int8 via AVX-512 VNNI: `vpdpbusd` fuses the maddubs+madd+add
8413/// triple into one u8×i8 dot-accumulate. AVX-512 has no vpsignb, so the
8414/// |w|·sign(x,w) trick becomes |w| × (x negated where w<0) via a mask
8415/// subtract — w==0 lanes contribute 0 through |w|=0 either way.
8416#[cfg(target_arch = "x86_64")]
8417#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
8418unsafe fn dot_i8_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
8419    // SAFETY: callers uphold slice-length contracts (see call sites).
8420    unsafe {
8421        use core::arch::x86_64::*;
8422        let n = w.len();
8423        let mut j = 0usize;
8424        let mut total: i32;
8425        // 4 independent accumulators: vpdpbusd is its own loop-carried
8426        // dependency (~5-cycle latency) — a single-acc loop runs
8427        // latency-bound and LOSES to the AVX2 maddubs kernel, measured
8428        // on Granite Rapids.
8429        {
8430            #[inline(always)]
8431            unsafe fn step(
8432                w: *const u8,
8433                x: *const i8,
8434                acc: core::arch::x86_64::__m512i,
8435            ) -> core::arch::x86_64::__m512i {
8436                unsafe {
8437                    use core::arch::x86_64::*;
8438                    let wv = _mm512_loadu_si512(w as *const _);
8439                    let xv = _mm512_loadu_si512(x as *const _);
8440                    let aw = _mm512_abs_epi8(wv);
8441                    let neg = _mm512_movepi8_mask(wv);
8442                    let sx = _mm512_mask_sub_epi8(xv, neg, _mm512_setzero_si512(), xv);
8443                    _mm512_dpbusd_epi32(acc, aw, sx)
8444                }
8445            }
8446            let (mut a0, mut a1, mut a2, mut a3) = (
8447                _mm512_setzero_si512(),
8448                _mm512_setzero_si512(),
8449                _mm512_setzero_si512(),
8450                _mm512_setzero_si512(),
8451            );
8452            while j + 256 <= n {
8453                a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
8454                a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), a1);
8455                a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), a2);
8456                a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), a3);
8457                j += 256;
8458            }
8459            while j + 64 <= n {
8460                a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
8461                j += 64;
8462            }
8463            let s01 = _mm512_add_epi32(a0, a1);
8464            let s23 = _mm512_add_epi32(a2, a3);
8465            total = _mm512_reduce_add_epi32(_mm512_add_epi32(s01, s23));
8466        }
8467        // 32-wide (q4/vbit groups are exactly 32 bytes).
8468        if j + 32 <= n {
8469            let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
8470            let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
8471            let d = _mm256_dpbusd_epi32(
8472                _mm256_setzero_si256(),
8473                _mm256_abs_epi8(wv),
8474                _mm256_sign_epi8(xv, wv),
8475            );
8476            let hi128 = _mm256_extracti128_si256::<1>(d);
8477            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
8478            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8479            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8480            total += _mm_cvtsi128_si32(s32);
8481            j += 32;
8482        }
8483        while j < n {
8484            total += (w[j] as i8) as i32 * xq[j] as i32;
8485            j += 1;
8486        }
8487        total
8488    }
8489}
8490
8491/// i8 row · f32 x via AVX2/FMA (x86 mirror of `dot_i8_f32_neon`).
8492#[cfg(target_arch = "x86_64")]
8493#[target_feature(enable = "avx2,fma")]
8494unsafe fn dot_i8_f32_avx2(w: &[u8], x: &[f32]) -> f32 {
8495    // SAFETY: callers uphold slice-length contracts (see call sites).
8496    unsafe {
8497        use core::arch::x86_64::*;
8498        let n = x.len();
8499        let wp = w.as_ptr();
8500        let xp = x.as_ptr();
8501        let (mut a0, mut a1) = (_mm256_setzero_ps(), _mm256_setzero_ps());
8502        let mut j = 0usize;
8503        while j + 16 <= n {
8504            let wb = _mm_loadu_si128(wp.add(j) as *const __m128i);
8505            let lo = _mm256_cvtepi8_epi32(wb);
8506            let hi = _mm256_cvtepi8_epi32(_mm_srli_si128::<8>(wb));
8507            a0 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(lo), _mm256_loadu_ps(xp.add(j)), a0);
8508            a1 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(hi), _mm256_loadu_ps(xp.add(j + 8)), a1);
8509            j += 16;
8510        }
8511        let acc = _mm256_add_ps(a0, a1);
8512        let hi128 = _mm256_extractf128_ps::<1>(acc);
8513        let s128 = _mm_add_ps(_mm256_castps256_ps128(acc), hi128);
8514        let s64 = _mm_add_ps(s128, _mm_movehl_ps(s128, s128));
8515        let s32 = _mm_add_ss(s64, _mm_shuffle_ps::<1>(s64, s64));
8516        let mut sum = _mm_cvtss_f32(s32);
8517        while j < n {
8518            sum += (*wp.add(j) as i8) as f32 * *xp.add(j);
8519            j += 1;
8520        }
8521        sum
8522    }
8523}
8524
8525/// int8(weight)·int8(activation) → i32 via AVX2 maddubs — the x86
8526/// analogue of the SDOT A8W8 path. `maddubs` takes u8×i8, so the
8527/// standard sign trick applies: |w| × sign(x, w) ≡ w × x per lane.
8528/// Pair saturation is safe: |w|≤128, |x|≤127 → 2·128·127 < 32767.
8529#[cfg(target_arch = "x86_64")]
8530#[target_feature(enable = "avx2")]
8531unsafe fn dot_i8_i8_avx2(w: &[u8], xq: &[i8]) -> i32 {
8532    // SAFETY: callers uphold slice-length contracts (see call sites).
8533    unsafe {
8534        use core::arch::x86_64::*;
8535        let n = w.len();
8536        let ones = _mm256_set1_epi16(1);
8537        let mut acc = _mm256_setzero_si256();
8538        let mut j = 0usize;
8539        while j + 32 <= n {
8540            let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
8541            let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
8542            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
8543            acc = _mm256_add_epi32(acc, _mm256_madd_epi16(p16, ones));
8544            j += 32;
8545        }
8546        let hi128 = _mm256_extracti128_si256::<1>(acc);
8547        let s128 = _mm_add_epi32(_mm256_castsi256_si128(acc), hi128);
8548        let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8549        let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8550        let mut s = _mm_cvtsi128_si32(s32);
8551        while j < n {
8552            s += (w[j] as i8) as i32 * xq[j] as i32;
8553            j += 1;
8554        }
8555        s
8556    }
8557}
8558
8559/// smmla 2×4: one instruction covers a 2-row × 2-activation × 8-deep
8560/// tile (32 MACs vs sdot's 16) — the weight pair loads once per 8-k
8561/// slice as a combined 2×8 register and meets two activation pairs.
8562#[cfg(target_arch = "aarch64")]
8563#[target_feature(enable = "neon,i8mm")]
8564unsafe fn dot_i8_smmla_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
8565    // SAFETY: callers uphold slice-length contracts.
8566    unsafe {
8567        use core::arch::aarch64::*;
8568        use core::arch::asm;
8569        let n = w0.len();
8570        let w0p = w0.as_ptr() as *const i8;
8571        let w1p = w1.as_ptr() as *const i8;
8572        // acc01 holds [c(r0,x0) c(r0,x1) c(r1,x0) c(r1,x1)]; acc23 the
8573        // same for x2/x3.
8574        let mut acc01 = vdupq_n_s32(0);
8575        let mut acc23 = vdupq_n_s32(0);
8576        let mut i = 0usize;
8577        while i + 8 <= n {
8578            let wa = vcombine_s8(vld1_s8(w0p.add(i)), vld1_s8(w1p.add(i)));
8579            let xb01 = vcombine_s8(
8580                vld1_s8(xs[0].as_ptr().add(i)),
8581                vld1_s8(xs[1].as_ptr().add(i)),
8582            );
8583            let xb23 = vcombine_s8(
8584                vld1_s8(xs[2].as_ptr().add(i)),
8585                vld1_s8(xs[3].as_ptr().add(i)),
8586            );
8587            asm!(
8588                "smmla {a01:v}.4s, {w:v}.16b, {x01:v}.16b",
8589                "smmla {a23:v}.4s, {w:v}.16b, {x23:v}.16b",
8590                a01 = inout(vreg) acc01, a23 = inout(vreg) acc23,
8591                w = in(vreg) wa, x01 = in(vreg) xb01, x23 = in(vreg) xb23,
8592                options(pure, nomem, nostack),
8593            );
8594            i += 8;
8595        }
8596        let mut out = [[0i32; 4]; 2];
8597        let a01: [i32; 4] = core::mem::transmute(acc01);
8598        let a23: [i32; 4] = core::mem::transmute(acc23);
8599        out[0][0] = a01[0];
8600        out[0][1] = a01[1];
8601        out[1][0] = a01[2];
8602        out[1][1] = a01[3];
8603        out[0][2] = a23[0];
8604        out[0][3] = a23[1];
8605        out[1][2] = a23[2];
8606        out[1][3] = a23[3];
8607        if i < n {
8608            for (k, x) in xs.iter().enumerate() {
8609                for j in i..n {
8610                    out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
8611                    out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
8612                }
8613            }
8614        }
8615        out
8616    }
8617}
8618
8619/// ARM twin of the x86 blocked prefill GEMM: two weight rows stay in
8620/// registers across four activation streams, eight sdot accumulators.
8621/// (The per-row form re-read each W row once per activation.)
8622#[cfg(target_arch = "aarch64")]
8623#[target_feature(enable = "neon,dotprod")]
8624unsafe fn dot_i8_sdot_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
8625    // SAFETY: callers uphold slice-length contracts.
8626    unsafe {
8627        use core::arch::aarch64::*;
8628        use core::arch::asm;
8629        let n = w0.len();
8630        let w0p = w0.as_ptr() as *const i8;
8631        let w1p = w1.as_ptr() as *const i8;
8632        let mut acc = [[vdupq_n_s32(0); 4]; 2];
8633        let mut i = 0usize;
8634        while i + 16 <= n {
8635            let wv0 = vld1q_s8(w0p.add(i));
8636            let wv1 = vld1q_s8(w1p.add(i));
8637            for (k, x) in xs.iter().enumerate() {
8638                let xv = vld1q_s8(x.as_ptr().add(i));
8639                let (mut a0, mut a1) = (acc[0][k], acc[1][k]);
8640                asm!(
8641                    "sdot {a0:v}.4s, {w0:v}.16b, {x:v}.16b",
8642                    "sdot {a1:v}.4s, {w1:v}.16b, {x:v}.16b",
8643                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
8644                    w0 = in(vreg) wv0, w1 = in(vreg) wv1, x = in(vreg) xv,
8645                    options(pure, nomem, nostack),
8646                );
8647                acc[0][k] = a0;
8648                acc[1][k] = a1;
8649            }
8650            i += 16;
8651        }
8652        let mut out = [[0i32; 4]; 2];
8653        for r in 0..2 {
8654            for k in 0..4 {
8655                out[r][k] = vaddvq_s32(acc[r][k]);
8656            }
8657        }
8658        if i < n {
8659            for (k, x) in xs.iter().enumerate() {
8660                for j in i..n {
8661                    out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
8662                    out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
8663                }
8664            }
8665        }
8666        out
8667    }
8668}
8669
8670/// Blocked 2 weight rows × 4 activations for the prefill GEMM
8671/// (roadmap P0: packed panels + multi-row accumulators). The two rows'
8672/// abs() live in registers across all four activation streams; the
8673/// sign-fixup is recomputed per pair (the price of the maddubs trick).
8674/// Returns raw i8·i8 dots; the caller applies scales and outliers.
8675#[cfg(target_arch = "x86_64")]
8676#[target_feature(enable = "avx2")]
8677unsafe fn dot_i8_i8_avx2_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
8678    // SAFETY: callers uphold slice-length contracts.
8679    unsafe {
8680        use core::arch::x86_64::*;
8681        let n = w0.len();
8682        let ones = _mm256_set1_epi16(1);
8683        let mut acc = [[_mm256_setzero_si256(); 4]; 2];
8684        let mut j = 0usize;
8685        while j + 32 <= n {
8686            let wv0 = _mm256_loadu_si256(w0.as_ptr().add(j) as *const __m256i);
8687            let wv1 = _mm256_loadu_si256(w1.as_ptr().add(j) as *const __m256i);
8688            let aw0 = _mm256_abs_epi8(wv0);
8689            let aw1 = _mm256_abs_epi8(wv1);
8690            for (k, x) in xs.iter().enumerate() {
8691                let xv = _mm256_loadu_si256(x.as_ptr().add(j) as *const __m256i);
8692                let p0 = _mm256_maddubs_epi16(aw0, _mm256_sign_epi8(xv, wv0));
8693                acc[0][k] = _mm256_add_epi32(acc[0][k], _mm256_madd_epi16(p0, ones));
8694                let p1 = _mm256_maddubs_epi16(aw1, _mm256_sign_epi8(xv, wv1));
8695                acc[1][k] = _mm256_add_epi32(acc[1][k], _mm256_madd_epi16(p1, ones));
8696            }
8697            j += 32;
8698        }
8699        let mut out = [[0i32; 4]; 2];
8700        for r in 0..2 {
8701            for k in 0..4 {
8702                let a = acc[r][k];
8703                let hi128 = _mm256_extracti128_si256::<1>(a);
8704                let s128 = _mm_add_epi32(_mm256_castsi256_si128(a), hi128);
8705                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8706                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8707                out[r][k] = _mm_cvtsi128_si32(s32);
8708            }
8709        }
8710        if j < n {
8711            for (k, x) in xs.iter().enumerate() {
8712                for i in j..n {
8713                    out[0][k] += (w0[i] as i8) as i32 * x[i] as i32;
8714                    out[1][k] += (w1[i] as i8) as i32 * x[i] as i32;
8715                }
8716            }
8717        }
8718        out
8719    }
8720}
8721
8722/// AVX2/VNNI q8 row dot with exact outlier correction (x86 mirror of
8723/// `row_dot_sdot` — same A8W8 contract). With AVX-512 VNNI the row goes
8724/// through the bias trick: Σ(w+128)·x via pure `vpdpbusd` (no per-lane
8725/// sign fixups), corrected by −128·Σx with Σx precomputed per split.
8726#[cfg(target_arch = "x86_64")]
8727#[inline]
8728fn row_dot_avx2(row: &[u8], act: &SplitAct) -> f32 {
8729    let dot = if avx512vnni_enabled() && row.len() >= 64 {
8730        (unsafe { dot_u8p128_i8_vnni(row, &act.xq) }) - 128 * act.xsum
8731    } else {
8732        unsafe { dot_i8_i8_avx2(row, &act.xq) }
8733    };
8734    let mut acc = dot as f32 * act.sx;
8735    for &(j, xv) in &act.outliers {
8736        acc += (row[j] as i8) as f32 * xv;
8737    }
8738    acc
8739}
8740
8741/// Σ (w[i]+128)·x[i] via pure `vpdpbusd` — the caller subtracts
8742/// 128·Σx. Four independent accumulators (dpbusd is ~5-cycle latency;
8743/// a single-acc loop runs latency-bound, measured on Granite Rapids).
8744#[cfg(target_arch = "x86_64")]
8745#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
8746unsafe fn dot_u8p128_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
8747    // SAFETY: callers uphold slice-length contracts (see call sites).
8748    unsafe {
8749        use core::arch::x86_64::*;
8750        let n = w.len();
8751        let flip = _mm512_set1_epi8(-128); // XOR 0x80: i8 w → u8 (w+128)
8752        #[inline(always)]
8753        unsafe fn step(
8754            w: *const u8,
8755            x: *const i8,
8756            flip: core::arch::x86_64::__m512i,
8757            acc: core::arch::x86_64::__m512i,
8758        ) -> core::arch::x86_64::__m512i {
8759            unsafe {
8760                use core::arch::x86_64::*;
8761                let wv = _mm512_xor_si512(_mm512_loadu_si512(w as *const _), flip);
8762                _mm512_dpbusd_epi32(acc, wv, _mm512_loadu_si512(x as *const _))
8763            }
8764        }
8765        let (mut a0, mut a1, mut a2, mut a3) = (
8766            _mm512_setzero_si512(),
8767            _mm512_setzero_si512(),
8768            _mm512_setzero_si512(),
8769            _mm512_setzero_si512(),
8770        );
8771        let mut j = 0usize;
8772        while j + 256 <= n {
8773            a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
8774            a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), flip, a1);
8775            a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), flip, a2);
8776            a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), flip, a3);
8777            j += 256;
8778        }
8779        while j + 64 <= n {
8780            a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
8781            j += 64;
8782        }
8783        let mut total = _mm512_reduce_add_epi32(_mm512_add_epi32(
8784            _mm512_add_epi32(a0, a1),
8785            _mm512_add_epi32(a2, a3),
8786        ));
8787        // Scalar tail: (w as i8) + 128 ≡ (w as u8) ^ 0x80.
8788        while j < n {
8789            total += ((w[j] ^ 0x80) as i32) * xq[j] as i32;
8790            j += 1;
8791        }
8792        total
8793    }
8794}
8795
8796/// One q4 row via AVX2: nibbles → centered i8 (unpacklo/hi restores the
8797/// writer's flat order, same as the NEON vzip pair), maddubs against
8798/// the pre-quantized activation group, × the group's f16 scale. Pair
8799/// saturation safe: |w|≤8, |x|≤127 → 2·8·127 ≪ 32767. Mirror of
8800/// `dot_q4_row_sdot`.
8801#[cfg(target_arch = "x86_64")]
8802#[target_feature(enable = "avx2")]
8803unsafe fn dot_q4_row_avx2(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
8804    // SAFETY: callers uphold slice-length contracts (16 packed bytes and
8805    // 2 scale bytes per group; xq.len() == gpr·GROUP_SIZE).
8806    unsafe {
8807        use core::arch::x86_64::*;
8808        let lomask = _mm_set1_epi8(0x0F);
8809        let eight = _mm256_set1_epi8(8);
8810        let ones = _mm256_set1_epi16(1);
8811        let mut acc = 0f32;
8812        for gi in 0..gpr {
8813            let g = g0 + gi;
8814            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
8815            let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
8816            let lo = _mm_and_si128(b, lomask);
8817            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
8818            let w = _mm256_sub_epi8(
8819                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
8820                eight,
8821            );
8822            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8823            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
8824            let d = _mm256_madd_epi16(p16, ones);
8825            let hi128 = _mm256_extracti128_si256::<1>(d);
8826            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
8827            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8828            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8829            acc += _mm_cvtsi128_si32(s32) as f32 * s;
8830        }
8831        acc
8832    }
8833}
8834
8835/// Two-activation q4 row via AVX2: nibbles unpacked ONCE per group,
8836/// both activations dotted against the same centered i8 register.
8837#[cfg(target_arch = "x86_64")]
8838#[target_feature(enable = "avx2")]
8839unsafe fn dot_q4_row_avx2_2(
8840    packed: &[u8],
8841    scales: &[u8],
8842    g0: usize,
8843    gpr: usize,
8844    xq1: &[i8],
8845    xq2: &[i8],
8846) -> (f32, f32) {
8847    // SAFETY: callers uphold slice-length contracts (see dot_q4_row_avx2).
8848    unsafe {
8849        use core::arch::x86_64::*;
8850        let lomask = _mm_set1_epi8(0x0F);
8851        let eight = _mm256_set1_epi8(8);
8852        let ones = _mm256_set1_epi16(1);
8853        let (mut acc1, mut acc2) = (0f32, 0f32);
8854        #[inline(always)]
8855        unsafe fn hsum(d: core::arch::x86_64::__m256i) -> i32 {
8856            unsafe {
8857                use core::arch::x86_64::*;
8858                let hi128 = _mm256_extracti128_si256::<1>(d);
8859                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
8860                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8861                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8862                _mm_cvtsi128_si32(s32)
8863            }
8864        }
8865        for gi in 0..gpr {
8866            let g = g0 + gi;
8867            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
8868            let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
8869            let lo = _mm_and_si128(b, lomask);
8870            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
8871            let w = _mm256_sub_epi8(
8872                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
8873                eight,
8874            );
8875            let aw = _mm256_abs_epi8(w);
8876            let x1 = _mm256_loadu_si256(xq1.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8877            let x2 = _mm256_loadu_si256(xq2.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8878            let d1 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x1, w)), ones);
8879            let d2 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x2, w)), ones);
8880            acc1 += hsum(d1) as f32 * s;
8881            acc2 += hsum(d2) as f32 * s;
8882        }
8883        (acc1, acc2)
8884    }
8885}
8886
8887/// One q8 row range via AVX2 (x86 mirror of `q8_range_sdot`).
8888#[cfg(target_arch = "x86_64")]
8889fn q8_range_avx2(
8890    q: &[u8],
8891    row_scale: &[f32],
8892    act: &SplitAct,
8893    cols: usize,
8894    out_addr: SendMut,
8895    start: usize,
8896    end: usize,
8897) {
8898    for o in start..end {
8899        let v = row_dot_avx2(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
8900        // SAFETY: disjoint row ranges per worker.
8901        unsafe { *out_addr.at(o) = v };
8902    }
8903}
8904
8905/// Two-input q8 row range via AVX2 (x86 mirror of `q8_range2_sdot`).
8906#[cfg(target_arch = "x86_64")]
8907#[allow(clippy::too_many_arguments)]
8908fn q8_range2_avx2(
8909    q: &[u8],
8910    row_scale: &[f32],
8911    a1: &SplitAct,
8912    a2: &SplitAct,
8913    cols: usize,
8914    p1: SendMut,
8915    p2: SendMut,
8916    start: usize,
8917    end: usize,
8918) {
8919    for o in start..end {
8920        let row = &q[o * cols..(o + 1) * cols];
8921        // SAFETY: disjoint row ranges per worker.
8922        unsafe {
8923            *p1.at(o) = row_dot_avx2(row, a1) * row_scale[o];
8924            *p2.at(o) = row_dot_avx2(row, a2) * row_scale[o];
8925        }
8926    }
8927}
8928
8929// ───────────────────── A8W8 SDOT path (port of vmfcore, ×1.78 decode) ─────────────────────
8930
8931/// ARMv8.6 i8mm (smmla): 32 int8 MACs per instruction vs sdot's 16 —
8932/// yet MEASURED 2.4× SLOWER than the blocked sdot on Apple silicon
8933/// (108 vs 264 GF/s): the on-the-fly vcombine packing and the two-
8934/// accumulator dependency chain swamp the MAC advantage, and Apple's
8935/// four SIMD pipes already keep sdot fed. OPT-IN (CMF_I8MM=1) for
8936/// field trials on Cortex-A710/X-class parts with two pipes, where the
8937/// balance may differ; a pre-interleaved weight layout (repack infra)
8938/// is the known path if it ever earns its keep.
8939#[cfg(target_arch = "aarch64")]
8940fn i8mm_enabled() -> bool {
8941    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8942    *ON.get_or_init(|| {
8943        std::env::var("CMF_I8MM").map(|v| v == "1").unwrap_or(false)
8944            && std::arch::is_aarch64_feature_detected!("i8mm")
8945    })
8946}
8947
8948/// SDOT enabled? Default ON when the CPU has ARMv8.2 dotprod;
8949/// `CMF_SDOT=0` disables (falls back to i8×f32 NEON).
8950/// (On non-ARM release builds only the test tolerance switch calls it.)
8951#[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
8952fn sdot_enabled() -> bool {
8953    use std::sync::OnceLock;
8954    static ON: OnceLock<bool> = OnceLock::new();
8955    *ON.get_or_init(|| {
8956        let want = std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true);
8957        if !want {
8958            return false;
8959        }
8960
8961        #[cfg(target_arch = "aarch64")]
8962        {
8963            if std::arch::is_aarch64_feature_detected!("dotprod") {
8964                return true;
8965            }
8966            #[cfg(target_os = "android")]
8967            {
8968                if let Ok(cpuinfo) = std::fs::read_to_string("/proc/cpuinfo") {
8969                    if cpuinfo.lines().any(|l| {
8970                        (l.starts_with("Features") || l.starts_with("features"))
8971                            && l.contains("asimddp")
8972                    }) {
8973                        return true;
8974                    }
8975                }
8976            }
8977            false
8978        }
8979        #[cfg(not(target_arch = "aarch64"))]
8980        {
8981            false
8982        }
8983    })
8984}
8985
8986/// Two-field activation split (≡ vmfcore `q8_split_prep`): outlier
8987/// channels (>8·rms) are computed exactly in f32; the bulk (outliers
8988/// zeroed → clean absmax) goes through int8 SDOT. Computed ONCE per
8989/// matvec, shared by all rows/workers.
8990struct SplitAct {
8991    xq: Vec<i8>,
8992    sx: f32,
8993    outliers: Vec<(usize, f32)>,
8994    /// Σ xq — the VNNI bias-trick correction (`(w+128)·x` sums need
8995    /// `−128·Σx`); one i32 per split, computed once per matvec.
8996    #[cfg_attr(not(target_arch = "x86_64"), allow(dead_code))]
8997    xsum: i32,
8998}
8999
9000thread_local! {
9001    /// Recycled xq buffers: split_act runs for every matvec (~200/token)
9002    /// and its hidden-size allocation was steady-state heap churn.
9003    static XQ_FREE: std::cell::RefCell<Vec<Vec<i8>>> =
9004        const { std::cell::RefCell::new(Vec::new()) };
9005}
9006
9007impl Drop for SplitAct {
9008    fn drop(&mut self) {
9009        let buf = std::mem::take(&mut self.xq);
9010        if buf.capacity() > 0 {
9011            XQ_FREE.with(|f| {
9012                let mut f = f.borrow_mut();
9013                if f.len() < 16 {
9014                    f.push(buf);
9015                }
9016            });
9017        }
9018    }
9019}
9020
9021thread_local! {
9022    /// One scratch row per WORKER, kept for the life of the thread.
9023    ///
9024    /// The kernels take a row of group scales per dispatch, and a fresh
9025    /// `vec![0f32; gpr]` inside the closure is one allocation per worker per
9026    /// dispatch — on the release checkpoint about six thousand a token, a
9027    /// quarter of everything the benchmark counts.
9028    static KROW: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
9029}
9030
9031/// Borrow `n` floats of the calling worker's scratch. Nothing inside a
9032/// kernel body borrows it again, which is what keeps the RefCell honest.
9033#[inline]
9034fn with_krow<R>(n: usize, f: impl FnOnce(&mut [f32]) -> R) -> R {
9035    KROW.with(|s| {
9036        let mut b = s.borrow_mut();
9037        if b.len() < n {
9038            b.resize(n, 0.0);
9039        }
9040        f(&mut b[..n])
9041    })
9042}
9043
9044fn split_act(x: &[f32]) -> SplitAct {
9045    let n = x.len();
9046    let rms = (x.iter().map(|&v| (v * v) as f64).sum::<f64>() / n.max(1) as f64).sqrt() as f32;
9047    let thr = 8.0 * rms;
9048    // One pass: collect outliers and the bulk absmax (outliers excluded —
9049    // identical to the old zero-then-fold over a copied buffer, minus the
9050    // full-vector copy).
9051    let mut outliers: Vec<(usize, f32)> = Vec::new();
9052    let mut amax = 0f32;
9053    for (j, &v) in x.iter().enumerate() {
9054        let a = v.abs();
9055        if a > thr {
9056            outliers.push((j, v));
9057        } else if a > amax {
9058            amax = a;
9059        }
9060    }
9061    let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
9062    let inv = 1.0 / sx;
9063    let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
9064    xq.clear();
9065    xq.reserve(n);
9066    if outliers.is_empty() {
9067        xq.extend(
9068            x.iter()
9069                .map(|&v| (v * inv).round().clamp(-127.0, 127.0) as i8),
9070        );
9071    } else {
9072        // Outlier slots quantize to 0 (their exact term is added later).
9073        xq.extend(x.iter().map(|&v| {
9074            if v.abs() > thr {
9075                0
9076            } else {
9077                (v * inv).round().clamp(-127.0, 127.0) as i8
9078            }
9079        }));
9080    }
9081    let xsum = xq.iter().map(|&v| v as i32).sum();
9082    SplitAct {
9083        xq,
9084        sx,
9085        outliers,
9086        xsum,
9087    }
9088}
9089
9090fn split_act_q8_2f(x: &[f32], col: &[f32]) -> SplitAct {
9091    let n = x.len();
9092    let rms = (x
9093        .iter()
9094        .zip(col)
9095        .map(|(&a, &c)| {
9096            let v = a * c;
9097            (v * v) as f64
9098        })
9099        .sum::<f64>()
9100        / n.max(1) as f64)
9101        .sqrt() as f32;
9102    let thr = 8.0 * rms;
9103
9104    let mut outliers = Vec::new();
9105    let mut amax = 0f32;
9106    for (j, (&a, &c)) in x.iter().zip(col).enumerate() {
9107        let v = a * c;
9108        let s = v.abs();
9109        if s > thr {
9110            outliers.push((j, v));
9111        } else if s > amax {
9112            amax = s;
9113        }
9114    }
9115
9116    let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
9117    let inv = 1.0 / sx;
9118    let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
9119    xq.clear();
9120    xq.reserve(n);
9121    if outliers.is_empty() {
9122        xq.extend(
9123            x.iter()
9124                .zip(col)
9125                .map(|(&a, &c)| ((a * c) * inv).round().clamp(-127.0, 127.0) as i8),
9126        );
9127    } else {
9128        xq.extend(x.iter().zip(col).map(|(&a, &c)| {
9129            let v = a * c;
9130            if v.abs() > thr {
9131                0
9132            } else {
9133                (v * inv).round().clamp(-127.0, 127.0) as i8
9134            }
9135        }));
9136    }
9137    let xsum = xq.iter().map(|&v| v as i32).sum();
9138    SplitAct {
9139        xq,
9140        sx,
9141        outliers,
9142        xsum,
9143    }
9144}
9145
9146/// int8(weight)·int8(activation) → i32 via `sdot` (inline asm — the
9147/// vdotq intrinsic is unstable; port of vmfcore `dot_i8_sdot`).
9148#[cfg(target_arch = "aarch64")]
9149#[target_feature(enable = "neon,dotprod")]
9150unsafe fn dot_i8_sdot(w: &[u8], xq: &[i8]) -> i32 {
9151    // SAFETY: callers uphold slice-length contracts (see call sites).
9152    unsafe {
9153        use core::arch::aarch64::*;
9154        use core::arch::asm;
9155        let wp = w.as_ptr() as *const i8;
9156        let n = w.len();
9157        let (mut a0, mut a1, mut a2, mut a3) = (
9158            vdupq_n_s32(0),
9159            vdupq_n_s32(0),
9160            vdupq_n_s32(0),
9161            vdupq_n_s32(0),
9162        );
9163        let mut i = 0;
9164        while i + 64 <= n {
9165            let (w0, x0) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
9166            let (w1, x1) = (vld1q_s8(wp.add(i + 16)), vld1q_s8(xq.as_ptr().add(i + 16)));
9167            let (w2, x2) = (vld1q_s8(wp.add(i + 32)), vld1q_s8(xq.as_ptr().add(i + 32)));
9168            let (w3, x3) = (vld1q_s8(wp.add(i + 48)), vld1q_s8(xq.as_ptr().add(i + 48)));
9169            asm!(
9170                "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
9171                "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
9172                "sdot {a2:v}.4s, {w2:v}.16b, {x2:v}.16b",
9173                "sdot {a3:v}.4s, {w3:v}.16b, {x3:v}.16b",
9174                a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
9175                w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
9176                w2 = in(vreg) w2, x2 = in(vreg) x2, w3 = in(vreg) w3, x3 = in(vreg) x3,
9177                options(pure, nomem, nostack),
9178            );
9179            i += 64;
9180        }
9181        while i + 16 <= n {
9182            let (wv, xv) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
9183            asm!("sdot {a:v}.4s, {w:v}.16b, {x:v}.16b",
9184                 a = inout(vreg) a0, w = in(vreg) wv, x = in(vreg) xv, options(pure, nomem, nostack));
9185            i += 16;
9186        }
9187        let mut s = vaddvq_s32(vaddq_s32(vaddq_s32(a0, a1), vaddq_s32(a2, a3)));
9188        while i < n {
9189            s += (*wp.add(i)) as i32 * xq[i] as i32;
9190            i += 1;
9191        }
9192        s
9193    }
9194}
9195
9196/// Row-blocked SDOT: 4 output rows per pass — the activation chunk is
9197/// loaded once and reused, 4 independent accumulators hide sdot latency
9198/// (port of vmfcore `dot_i8_sdot_4rows`).
9199#[cfg(target_arch = "aarch64")]
9200#[target_feature(enable = "neon,dotprod")]
9201unsafe fn dot_i8_sdot_4rows(w0: &[u8], w1: &[u8], w2: &[u8], w3: &[u8], xq: &[i8]) -> [i32; 4] {
9202    // SAFETY: callers uphold slice-length contracts (see call sites).
9203    unsafe {
9204        use core::arch::aarch64::*;
9205        use core::arch::asm;
9206        let n = xq.len();
9207        let px = xq.as_ptr();
9208        let (p0, p1, p2, p3) = (
9209            w0.as_ptr() as *const i8,
9210            w1.as_ptr() as *const i8,
9211            w2.as_ptr() as *const i8,
9212            w3.as_ptr() as *const i8,
9213        );
9214        let (mut a0, mut a1, mut a2, mut a3) = (
9215            vdupq_n_s32(0),
9216            vdupq_n_s32(0),
9217            vdupq_n_s32(0),
9218            vdupq_n_s32(0),
9219        );
9220        let mut i = 0;
9221        while i + 16 <= n {
9222            let x = vld1q_s8(px.add(i));
9223            let v0 = vld1q_s8(p0.add(i));
9224            let v1 = vld1q_s8(p1.add(i));
9225            let v2 = vld1q_s8(p2.add(i));
9226            let v3 = vld1q_s8(p3.add(i));
9227            asm!(
9228                "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
9229                "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
9230                "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
9231                "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
9232                a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
9233                v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
9234                options(pure, nomem, nostack),
9235            );
9236            i += 16;
9237        }
9238        let mut r = [
9239            vaddvq_s32(a0),
9240            vaddvq_s32(a1),
9241            vaddvq_s32(a2),
9242            vaddvq_s32(a3),
9243        ];
9244        while i < n {
9245            let xi = *px.add(i) as i32;
9246            r[0] += (*p0.add(i)) as i32 * xi;
9247            r[1] += (*p1.add(i)) as i32 * xi;
9248            r[2] += (*p2.add(i)) as i32 * xi;
9249            r[3] += (*p3.add(i)) as i32 * xi;
9250            i += 1;
9251        }
9252        r
9253    }
9254}
9255
9256/// 4 interleaved rows in one pass: the repacked group is [r0[c], r1[c],
9257/// r2[c], r3[c]] per 16-byte chunk, so each iteration reads ONE 64-byte
9258/// line plus the shared activation chunk — a single sequential weight
9259/// stream per worker. Per-row accumulation is the same one-accumulator
9260/// scheme as `dot_i8_sdot_4rows`; integer sums are exact, so outputs
9261/// are bit-identical to the mmap-layout kernel.
9262#[cfg(target_arch = "aarch64")]
9263#[target_feature(enable = "neon,dotprod")]
9264unsafe fn dot_i8_sdot_4rows_il(g: &[u8], xq: &[i8]) -> [i32; 4] {
9265    // SAFETY: callers uphold slice-length contracts (g.len() == 4·n,
9266    // n % 16 == 0 — guaranteed by the repack gate).
9267    unsafe {
9268        use core::arch::aarch64::*;
9269        use core::arch::asm;
9270        let n = xq.len();
9271        let px = xq.as_ptr();
9272        let pg = g.as_ptr() as *const i8;
9273        let (mut a0, mut a1, mut a2, mut a3) = (
9274            vdupq_n_s32(0),
9275            vdupq_n_s32(0),
9276            vdupq_n_s32(0),
9277            vdupq_n_s32(0),
9278        );
9279        let mut i = 0;
9280        while i + 16 <= n {
9281            let x = vld1q_s8(px.add(i));
9282            let base = pg.add(4 * i);
9283            let v0 = vld1q_s8(base);
9284            let v1 = vld1q_s8(base.add(16));
9285            let v2 = vld1q_s8(base.add(32));
9286            let v3 = vld1q_s8(base.add(48));
9287            asm!(
9288                "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
9289                "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
9290                "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
9291                "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
9292                a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
9293                v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
9294                options(pure, nomem, nostack),
9295            );
9296            i += 16;
9297        }
9298        [
9299            vaddvq_s32(a0),
9300            vaddvq_s32(a1),
9301            vaddvq_s32(a2),
9302            vaddvq_s32(a3),
9303        ]
9304    }
9305}
9306
9307/// One q8 row range via SDOT (4-row blocks + tail) — the body of
9308/// `qmatvec`'s hot loop, extracted so multi-matrix jobs can drive the
9309/// SAME kernel for several tensors under one pool dispatch. `rep` — the
9310/// load-time interleaved repack (empty = mmap layout only); rows outside
9311/// full 4-row groups always come from the mmap layout.
9312#[cfg(target_arch = "aarch64")]
9313fn q8_range_sdot(
9314    q: &[u8],
9315    rep: &[u8],
9316    row_scale: &[f32],
9317    act: &SplitAct,
9318    cols: usize,
9319    out_addr: SendMut,
9320    start: usize,
9321    end: usize,
9322) {
9323    let mut o = start;
9324    // Leading rows to the group boundary (repack path only): the pool
9325    // splits row ranges arbitrarily, groups are absolute.
9326    if !rep.is_empty() {
9327        while o < end && o % 4 != 0 {
9328            let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
9329            unsafe { *out_addr.at(o) = v };
9330            o += 1;
9331        }
9332    }
9333    while o + 4 <= end {
9334        let r = if rep.is_empty() {
9335            unsafe {
9336                dot_i8_sdot_4rows(
9337                    &q[o * cols..(o + 1) * cols],
9338                    &q[(o + 1) * cols..(o + 2) * cols],
9339                    &q[(o + 2) * cols..(o + 3) * cols],
9340                    &q[(o + 3) * cols..(o + 4) * cols],
9341                    &act.xq,
9342                )
9343            }
9344        } else {
9345            unsafe { dot_i8_sdot_4rows_il(&rep[o * cols..(o + 4) * cols], &act.xq) }
9346        };
9347        for k in 0..4 {
9348            let mut acc = r[k] as f32 * act.sx;
9349            for &(j, xv) in &act.outliers {
9350                acc += (q[(o + k) * cols + j] as i8) as f32 * xv;
9351            }
9352            // SAFETY: disjoint row ranges per worker.
9353            unsafe { *out_addr.at(o + k) = acc * row_scale[o + k] };
9354        }
9355        o += 4;
9356    }
9357    while o < end {
9358        let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
9359        unsafe { *out_addr.at(o) = v };
9360        o += 1;
9361    }
9362}
9363
9364/// Two-input q8 row range via SDOT — `qmatvec2`'s hot loop, extracted
9365/// for the fused pair multi-matrix job (`matvec2_many`).
9366#[cfg(target_arch = "aarch64")]
9367#[allow(clippy::too_many_arguments)]
9368fn q8_range2_sdot(
9369    q: &[u8],
9370    row_scale: &[f32],
9371    a1: &SplitAct,
9372    a2: &SplitAct,
9373    cols: usize,
9374    p1: SendMut,
9375    p2: SendMut,
9376    start: usize,
9377    end: usize,
9378) {
9379    for o in start..end {
9380        let row = &q[o * cols..(o + 1) * cols];
9381        // SAFETY: disjoint row ranges per worker.
9382        unsafe {
9383            *p1.at(o) = row_dot_sdot(row, a1) * row_scale[o];
9384            *p2.at(o) = row_dot_sdot(row, a2) * row_scale[o];
9385        }
9386    }
9387}
9388
9389/// Two-input q8 row range, f32 kernel (non-SDOT) — same extraction.
9390#[allow(clippy::too_many_arguments)]
9391fn q8_range2_f32(
9392    q: &[u8],
9393    row_scale: &[f32],
9394    x1: &[f32],
9395    x2: &[f32],
9396    cols: usize,
9397    p1: SendMut,
9398    p2: SendMut,
9399    start: usize,
9400    end: usize,
9401) {
9402    for o in start..end {
9403        let row = &q[o * cols..(o + 1) * cols];
9404        // SAFETY: disjoint row ranges per worker.
9405        unsafe {
9406            *p1.at(o) = dot_i8_f32(row, x1) * row_scale[o];
9407            *p2.at(o) = dot_i8_f32(row, x2) * row_scale[o];
9408        }
9409    }
9410}
9411
9412/// Scalar/NEON-f32 q8 row range (non-SDOT platforms) — same extraction.
9413fn q8_range_f32(
9414    q: &[u8],
9415    row_scale: &[f32],
9416    xs: &[f32],
9417    cols: usize,
9418    out_addr: SendMut,
9419    start: usize,
9420    end: usize,
9421) {
9422    for o in start..end {
9423        let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
9424        // SAFETY: disjoint row ranges per worker.
9425        unsafe { *out_addr.at(o) = v };
9426    }
9427}
9428
9429/// One q8 row against a split activation, portable: the per-arch fast
9430/// dots where they exist, the exact scalar loop elsewhere. The scalar
9431/// arm is also the test oracle for both fast arms.
9432#[inline]
9433fn q8_row_dot(row: &[u8], act: &SplitAct) -> f32 {
9434    #[cfg(target_arch = "aarch64")]
9435    return row_dot_sdot(row, act);
9436    #[cfg(target_arch = "x86_64")]
9437    return row_dot_avx2(row, act);
9438    #[allow(unreachable_code)]
9439    q8_row_dot_scalar(row, act)
9440}
9441
9442#[allow(dead_code)]
9443fn q8_row_dot_scalar(row: &[u8], act: &SplitAct) -> f32 {
9444    let mut acc = 0i32;
9445    for (k, &b) in row.iter().enumerate() {
9446        acc += (b as i8) as i32 * act.xq[k] as i32;
9447    }
9448    let mut acc = acc as f32 * act.sx;
9449    for &(j, xv) in &act.outliers {
9450        acc += (row[j] as i8) as f32 * xv;
9451    }
9452    acc
9453}
9454
9455/// SDOT row dot with exact outlier correction:
9456/// `dot = sdot(w, xq)·sx + Σ_outl w[j]·x[j]` (then × row_scale by caller).
9457#[cfg(target_arch = "aarch64")]
9458#[inline]
9459fn row_dot_sdot(row: &[u8], act: &SplitAct) -> f32 {
9460    let mut acc = unsafe { dot_i8_sdot(row, &act.xq) } as f32 * act.sx;
9461    for &(j, xv) in &act.outliers {
9462        acc += (row[j] as i8) as f32 * xv;
9463    }
9464    acc
9465}
9466
9467/// One q4 row via SDOT: each 32-group's nibbles unpack to centered i8
9468/// (nib−8 ∈ [−8,7]), int8×int8 `sdot` against the pre-quantized
9469/// activation group, × the group's f16 scale. Returns Σ_g dot_g·s_g;
9470/// the caller multiplies by the activation scale and adds the exact
9471/// outlier terms (port of vmfcore `dot_q4_block_sdot`, +23% measured).
9472/// Nibble order matches the writer: element 2k = low nibble, 2k+1 = high
9473/// → zip(lo,hi) restores flat order.
9474#[cfg(target_arch = "aarch64")]
9475#[target_feature(enable = "neon,dotprod")]
9476unsafe fn dot_q4_row_sdot(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
9477    // SAFETY: callers uphold slice-length contracts (16 packed bytes and
9478    // 2 scale bytes per group; xq.len() == gpr·GROUP_SIZE).
9479    unsafe {
9480        use core::arch::aarch64::*;
9481        use core::arch::asm;
9482        let lomask = vdupq_n_u8(0x0F);
9483        let eight = vdupq_n_s8(8);
9484        let mut acc = 0f32;
9485        for gi in 0..gpr {
9486            let g = g0 + gi;
9487            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
9488            let b = vld1q_u8(packed.as_ptr().add(g * 16));
9489            let lo = vandq_u8(b, lomask);
9490            let hi = vshrq_n_u8::<4>(b);
9491            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
9492            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
9493            let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
9494            let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
9495            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
9496            asm!(
9497                "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
9498                "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
9499                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
9500                e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
9501                options(pure, nomem, nostack),
9502            );
9503            acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
9504        }
9505        acc
9506    }
9507}
9508
9509/// Two-activation q4 row via SDOT: the nibble unpack (the expensive
9510/// part) happens ONCE per group; both pre-quantized activations are
9511/// dotted against the same centered i8 registers. Per-lane math matches
9512/// `dot_q4_row_sdot` exactly.
9513#[cfg(target_arch = "aarch64")]
9514#[target_feature(enable = "neon,dotprod")]
9515unsafe fn dot_q4_row_sdot2(
9516    packed: &[u8],
9517    scales: &[u8],
9518    g0: usize,
9519    gpr: usize,
9520    xq1: &[i8],
9521    xq2: &[i8],
9522) -> (f32, f32) {
9523    // SAFETY: callers uphold slice-length contracts (16 packed bytes and
9524    // 2 scale bytes per group; xq*.len() == gpr·GROUP_SIZE).
9525    unsafe {
9526        use core::arch::aarch64::*;
9527        use core::arch::asm;
9528        let lomask = vdupq_n_u8(0x0F);
9529        let eight = vdupq_n_s8(8);
9530        let (mut acc1, mut acc2) = (0f32, 0f32);
9531        for gi in 0..gpr {
9532            let g = g0 + gi;
9533            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
9534            let b = vld1q_u8(packed.as_ptr().add(g * 16));
9535            let lo = vandq_u8(b, lomask);
9536            let hi = vshrq_n_u8::<4>(b);
9537            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
9538            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
9539            let x10 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE));
9540            let x11 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE + 16));
9541            let x20 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE));
9542            let x21 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE + 16));
9543            let (mut a0, mut a1, mut b0, mut b1) = (
9544                vdupq_n_s32(0),
9545                vdupq_n_s32(0),
9546                vdupq_n_s32(0),
9547                vdupq_n_s32(0),
9548            );
9549            asm!(
9550                "sdot {a0:v}.4s, {e0:v}.16b, {x10:v}.16b",
9551                "sdot {a1:v}.4s, {e1:v}.16b, {x11:v}.16b",
9552                "sdot {b0:v}.4s, {e0:v}.16b, {x20:v}.16b",
9553                "sdot {b1:v}.4s, {e1:v}.16b, {x21:v}.16b",
9554                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
9555                b0 = inout(vreg) b0, b1 = inout(vreg) b1,
9556                e0 = in(vreg) e0, e1 = in(vreg) e1,
9557                x10 = in(vreg) x10, x11 = in(vreg) x11,
9558                x20 = in(vreg) x20, x21 = in(vreg) x21,
9559                options(pure, nomem, nostack),
9560            );
9561            acc1 += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
9562            acc2 += vaddvq_s32(vaddq_s32(b0, b1)) as f32 * s;
9563        }
9564        (acc1, acc2)
9565    }
9566}
9567
9568// ───────────────────── fused int8 kernels ─────────────────────
9569
9570/// `acc += w · row` where the row is centered i8 — NEON widen+fma on
9571/// aarch64, scalar elsewhere. The KV-cache q8 value path rides on this.
9572#[inline]
9573pub(crate) fn axpy_i8_f32(acc: &mut [f32], row: &[i8], w: f32) {
9574    #[cfg(target_arch = "aarch64")]
9575    unsafe {
9576        return axpy_i8_f32_neon(acc, row, w);
9577    }
9578    #[cfg(target_arch = "x86_64")]
9579    if avx2_enabled() {
9580        return unsafe { axpy_i8_f32_avx2(acc, row, w) };
9581    }
9582    #[allow(unreachable_code)]
9583    {
9584        for (a, &b) in acc.iter_mut().zip(row) {
9585            *a += w * b as f32;
9586        }
9587    }
9588}
9589
9590/// i8→f32 axpy via AVX2/FMA (x86 mirror of `axpy_i8_f32_neon`).
9591#[cfg(target_arch = "x86_64")]
9592#[target_feature(enable = "avx2,fma")]
9593unsafe fn axpy_i8_f32_avx2(acc: &mut [f32], row: &[i8], w: f32) {
9594    // SAFETY: callers uphold slice-length contracts (see call sites).
9595    unsafe {
9596        use core::arch::x86_64::*;
9597        let n = acc.len().min(row.len());
9598        let ap = acc.as_mut_ptr();
9599        let rp = row.as_ptr();
9600        let wv = _mm256_set1_ps(w);
9601        let mut j = 0usize;
9602        while j + 16 <= n {
9603            let rb = _mm_loadu_si128(rp.add(j) as *const __m128i);
9604            let lo = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(rb));
9605            let hi = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(rb)));
9606            let v0 = _mm256_fmadd_ps(wv, lo, _mm256_loadu_ps(ap.add(j)));
9607            let v1 = _mm256_fmadd_ps(wv, hi, _mm256_loadu_ps(ap.add(j + 8)));
9608            _mm256_storeu_ps(ap.add(j), v0);
9609            _mm256_storeu_ps(ap.add(j + 8), v1);
9610            j += 16;
9611        }
9612        while j < n {
9613            *ap.add(j) += w * (*rp.add(j)) as f32;
9614            j += 1;
9615        }
9616    }
9617}
9618
9619#[cfg(target_arch = "aarch64")]
9620#[target_feature(enable = "neon")]
9621unsafe fn axpy_i8_f32_neon(acc: &mut [f32], row: &[i8], w: f32) {
9622    // SAFETY: callers uphold slice-length contracts (see call sites).
9623    unsafe {
9624        use core::arch::aarch64::*;
9625        let n = acc.len().min(row.len());
9626        let ap = acc.as_mut_ptr();
9627        let rp = row.as_ptr();
9628        let wv = vdupq_n_f32(w);
9629        let mut j = 0usize;
9630        while j + 16 <= n {
9631            let rb = vld1q_s8(rp.add(j));
9632            let lo = vmovl_s8(vget_low_s8(rb));
9633            let hi = vmovl_s8(vget_high_s8(rb));
9634            for (off, half) in [(0, lo), (8, hi)] {
9635                let f0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(half)));
9636                let f1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(half)));
9637                let o = j + off;
9638                vst1q_f32(ap.add(o), vfmaq_f32(vld1q_f32(ap.add(o)), wv, f0));
9639                vst1q_f32(ap.add(o + 4), vfmaq_f32(vld1q_f32(ap.add(o + 4)), wv, f1));
9640            }
9641            j += 16;
9642        }
9643        while j < n {
9644            *ap.add(j) += w * (*rp.add(j)) as f32;
9645            j += 1;
9646        }
9647    }
9648}
9649
9650/// i8 row · f32 x. NEON on aarch64 (ported from vmfcore `dot_i8_f32_neon`,
9651/// ≈9× scalar), scalar elsewhere.
9652#[inline]
9653pub(crate) fn dot_i8_f32(w: &[u8], x: &[f32]) -> f32 {
9654    #[cfg(target_arch = "aarch64")]
9655    unsafe {
9656        return dot_i8_f32_neon(w, x);
9657    }
9658    #[cfg(target_arch = "x86_64")]
9659    if avx2_enabled() {
9660        return unsafe { dot_i8_f32_avx2(w, x) };
9661    }
9662    #[allow(unreachable_code)]
9663    {
9664        let mut sum = 0.0f32;
9665        for (j, &b) in w.iter().enumerate() {
9666            sum += (b as i8) as f32 * x[j];
9667        }
9668        sum
9669    }
9670}
9671
9672/// i8 row · (x ⊙ col_field) — the q8_2f row dot with the θ col-field
9673/// folded into the product (no prescaled copy of x). NEON on aarch64,
9674/// scalar elsewhere. Used by the active-neuron path `row_dot`.
9675#[inline]
9676fn dot_i8_col_f32(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
9677    #[cfg(target_arch = "aarch64")]
9678    unsafe {
9679        return dot_i8_col_f32_neon(w, x, col);
9680    }
9681    #[allow(unreachable_code)]
9682    {
9683        let mut sum = 0.0f32;
9684        for (j, &b) in w.iter().enumerate() {
9685            sum += (b as i8) as f32 * x[j] * col[j];
9686        }
9687        sum
9688    }
9689}
9690
9691#[cfg(target_arch = "aarch64")]
9692#[target_feature(enable = "neon")]
9693unsafe fn dot_i8_col_f32_neon(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
9694    // SAFETY: callers uphold slice-length contracts (see call sites).
9695    unsafe {
9696        use core::arch::aarch64::*;
9697        let n = x.len();
9698        let wp = w.as_ptr() as *const i8;
9699        let xp = x.as_ptr();
9700        let cp = col.as_ptr();
9701        let (mut a0, mut a1, mut a2, mut a3) = (
9702            vdupq_n_f32(0.0),
9703            vdupq_n_f32(0.0),
9704            vdupq_n_f32(0.0),
9705            vdupq_n_f32(0.0),
9706        );
9707        let mut j = 0usize;
9708        while j + 16 <= n {
9709            let wb = vld1q_s8(wp.add(j));
9710            let lo = vmovl_s8(vget_low_s8(wb));
9711            let hi = vmovl_s8(vget_high_s8(wb));
9712            let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
9713            let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
9714            let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
9715            let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
9716            a0 = vfmaq_f32(
9717                a0,
9718                w0,
9719                vmulq_f32(vld1q_f32(xp.add(j)), vld1q_f32(cp.add(j))),
9720            );
9721            a1 = vfmaq_f32(
9722                a1,
9723                w1,
9724                vmulq_f32(vld1q_f32(xp.add(j + 4)), vld1q_f32(cp.add(j + 4))),
9725            );
9726            a2 = vfmaq_f32(
9727                a2,
9728                w2,
9729                vmulq_f32(vld1q_f32(xp.add(j + 8)), vld1q_f32(cp.add(j + 8))),
9730            );
9731            a3 = vfmaq_f32(
9732                a3,
9733                w3,
9734                vmulq_f32(vld1q_f32(xp.add(j + 12)), vld1q_f32(cp.add(j + 12))),
9735            );
9736            j += 16;
9737        }
9738        let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
9739        while j < n {
9740            sum += (*wp.add(j)) as f32 * *xp.add(j) * *cp.add(j);
9741            j += 1;
9742        }
9743        sum
9744    }
9745}
9746
9747#[cfg(target_arch = "aarch64")]
9748#[target_feature(enable = "neon")]
9749unsafe fn dot_i8_f32_neon(w: &[u8], x: &[f32]) -> f32 {
9750    // SAFETY: callers uphold slice-length contracts (see call sites).
9751    unsafe {
9752        use core::arch::aarch64::*;
9753        let n = x.len();
9754        let wp = w.as_ptr() as *const i8;
9755        let xp = x.as_ptr();
9756        let (mut a0, mut a1, mut a2, mut a3) = (
9757            vdupq_n_f32(0.0),
9758            vdupq_n_f32(0.0),
9759            vdupq_n_f32(0.0),
9760            vdupq_n_f32(0.0),
9761        );
9762        let mut j = 0usize;
9763        while j + 16 <= n {
9764            let wb = vld1q_s8(wp.add(j));
9765            let lo = vmovl_s8(vget_low_s8(wb));
9766            let hi = vmovl_s8(vget_high_s8(wb));
9767            let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
9768            let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
9769            let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
9770            let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
9771            a0 = vfmaq_f32(a0, w0, vld1q_f32(xp.add(j)));
9772            a1 = vfmaq_f32(a1, w1, vld1q_f32(xp.add(j + 4)));
9773            a2 = vfmaq_f32(a2, w2, vld1q_f32(xp.add(j + 8)));
9774            a3 = vfmaq_f32(a3, w3, vld1q_f32(xp.add(j + 12)));
9775            j += 16;
9776        }
9777        let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
9778        while j < n {
9779            sum += (*wp.add(j)) as f32 * *xp.add(j);
9780            j += 1;
9781        }
9782        sum
9783    }
9784}
9785
9786#[allow(clippy::too_many_arguments)]
9787fn qmatvec(
9788    q: &[u8],
9789    rep: &[u8],
9790    row_scale: &[f32],
9791    x: &[f32],
9792    col_field: &[f32],
9793    dtype: TensorDtype,
9794    rows: usize,
9795    cols: usize,
9796    out: &mut [f32],
9797    pool: Option<&Pool>,
9798) {
9799    debug_assert_eq!(out.len(), rows);
9800    #[cfg(not(target_arch = "aarch64"))]
9801    let _ = rep;
9802
9803    #[cfg(target_arch = "aarch64")]
9804    if sdot_enabled() {
9805        let act = if dtype == TensorDtype::Q8_2f {
9806            split_act_q8_2f(x, col_field)
9807        } else {
9808            split_act(x)
9809        };
9810        let out_addr = SendMut(out.as_mut_ptr());
9811        let run_range = |start: usize, end: usize| {
9812            q8_range_sdot(q, rep, row_scale, &act, cols, out_addr, start, end)
9813        };
9814        match pool {
9815            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
9816            _ => run_range(0, rows),
9817        }
9818        return;
9819    }
9820    // x86 A8W8 via AVX2 maddubs — same quantized-activation contract as
9821    // the SDOT path (CMF_AVX2=0 keeps the exact i8×f32 loop).
9822    #[cfg(target_arch = "x86_64")]
9823    if avx2_a8w8_enabled() {
9824        let act = if dtype == TensorDtype::Q8_2f {
9825            split_act_q8_2f(x, col_field)
9826        } else {
9827            split_act(x)
9828        };
9829        let out_addr = SendMut(out.as_mut_ptr());
9830        let run_range = |start: usize, end: usize| {
9831            q8_range_avx2(q, row_scale, &act, cols, out_addr, start, end)
9832        };
9833        match pool {
9834            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
9835            _ => run_range(0, rows),
9836        }
9837        return;
9838    }
9839
9840    prescale_with(x, col_field, dtype, 1, |xs| {
9841        let out_addr = SendMut(out.as_mut_ptr());
9842        let run_range = move |start: usize, end: usize| {
9843            for o in start..end {
9844                let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
9845                // SAFETY: disjoint row ranges per worker.
9846                unsafe { *out_addr.at(o) = v };
9847            }
9848        };
9849        match pool {
9850            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
9851            _ => run_range(0, rows),
9852        }
9853    });
9854}
9855
9856#[allow(clippy::too_many_arguments)]
9857fn qmatvec2(
9858    q: &[u8],
9859    row_scale: &[f32],
9860    x1: &[f32],
9861    x2: &[f32],
9862    col_field: &[f32],
9863    dtype: TensorDtype,
9864    rows: usize,
9865    cols: usize,
9866    o1: &mut [f32],
9867    o2: &mut [f32],
9868    pool: Option<&Pool>,
9869) {
9870    #[cfg(target_arch = "aarch64")]
9871    if sdot_enabled() {
9872        let a1s = if dtype == TensorDtype::Q8_2f {
9873            split_act_q8_2f(x1, col_field)
9874        } else {
9875            split_act(x1)
9876        };
9877        let a2s = if dtype == TensorDtype::Q8_2f {
9878            split_act_q8_2f(x2, col_field)
9879        } else {
9880            split_act(x2)
9881        };
9882        let p1 = SendMut(o1.as_mut_ptr());
9883        let p2 = SendMut(o2.as_mut_ptr());
9884        let run_range = |start: usize, end: usize| {
9885            q8_range2_sdot(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
9886        };
9887        match pool {
9888            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
9889            _ => run_range(0, rows),
9890        }
9891        return;
9892    }
9893    #[cfg(target_arch = "x86_64")]
9894    if avx2_a8w8_enabled() {
9895        let a1s = if dtype == TensorDtype::Q8_2f {
9896            split_act_q8_2f(x1, col_field)
9897        } else {
9898            split_act(x1)
9899        };
9900        let a2s = if dtype == TensorDtype::Q8_2f {
9901            split_act_q8_2f(x2, col_field)
9902        } else {
9903            split_act(x2)
9904        };
9905        let p1 = SendMut(o1.as_mut_ptr());
9906        let p2 = SendMut(o2.as_mut_ptr());
9907        let run_range = |start: usize, end: usize| {
9908            q8_range2_avx2(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
9909        };
9910        match pool {
9911            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
9912            _ => run_range(0, rows),
9913        }
9914        return;
9915    }
9916
9917    prescale_with(x1, col_field, dtype, 1, |x1s| {
9918        prescale_with(x2, col_field, dtype, 2, |x2s| {
9919            let p1 = SendMut(o1.as_mut_ptr());
9920            let p2 = SendMut(o2.as_mut_ptr());
9921            let run_range = move |start: usize, end: usize| {
9922                for o in start..end {
9923                    let row = &q[o * cols..(o + 1) * cols];
9924                    let s1 = dot_i8_f32(row, x1s) * row_scale[o];
9925                    let s2 = dot_i8_f32(row, x2s) * row_scale[o];
9926                    // SAFETY: disjoint row ranges per worker.
9927                    unsafe {
9928                        *p1.at(o) = s1;
9929                        *p2.at(o) = s2;
9930                    }
9931                }
9932            };
9933            match pool {
9934                Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
9935                _ => run_range(0, rows),
9936            }
9937        });
9938    });
9939}
9940
9941#[derive(Clone, Copy)]
9942struct SendMut(*mut f32);
9943unsafe impl Send for SendMut {}
9944unsafe impl Sync for SendMut {}
9945
9946impl SendMut {
9947    #[inline]
9948    fn at(self, i: usize) -> *mut f32 {
9949        unsafe { self.0.add(i) }
9950    }
9951}
9952
9953#[cfg(test)]
9954mod tests {
9955    use super::*;
9956
9957    #[test]
9958    fn q2tp_i8_dot_matches_exact_on_grid() {
9959        // On-grid activations (±1 → sx=1/127, xq=±127 dequantizes
9960        // exactly, no outliers) must make the integer path agree with
9961        // the exact scalar walk to f32 rounding.
9962        let (rows, cols) = (5, 64);
9963        let gpr = cols / GROUP_SIZE;
9964        // Synthetic codes plane + a flat ladder: scales_into is not under
9965        // test here, so drive dot_q2tp_row_i8 / q2tp_row_exact directly
9966        // with hand-made scales.
9967        let chunks: Vec<u8> = (0..rows * gpr * Q2TP_CHUNK)
9968            .map(|i| (i as u32).wrapping_mul(2654435761) as u8)
9969            .collect();
9970        let scales: Vec<f32> = (0..gpr).map(|g| 0.5 + g as f32 * 0.25).collect();
9971        let x: Vec<f32> = (0..cols).map(|i| if i % 3 == 0 { -1.0 } else { 1.0 }).collect();
9972        let act = split_act(&x);
9973        assert!(act.outliers.is_empty(), "on-grid input must have no outliers");
9974        let gsum = q1_group_sums(&act.xq, gpr);
9975        for r in 0..rows {
9976            let exact = q2tp_row_exact(&chunks, r, gpr, &x, &scales);
9977            let fast = dot_q2tp_row_i8(&chunks, r, gpr, &act.xq, &gsum, &scales) * act.sx;
9978            assert!(
9979                (exact - fast).abs() <= exact.abs() * 1e-5 + 1e-5,
9980                "row {r}: exact {exact} vs i8 {fast}"
9981            );
9982        }
9983    }
9984
9985    #[test]
9986    fn q8_row_dot_fast_matches_scalar() {
9987        // The per-arch fast dot must agree with the exact scalar oracle
9988        // (same contract the fused q8 FFN arm rides on).
9989        let cols = 96;
9990        let row: Vec<u8> = (0..cols)
9991            .map(|i| ((i as i32 * 37 % 251) - 125) as i8 as u8)
9992            .collect();
9993        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.13).sin()).collect();
9994        let act = split_act(&x);
9995        let fast = q8_row_dot(&row, &act);
9996        let scalar = q8_row_dot_scalar(&row, &act);
9997        assert!(
9998            (fast - scalar).abs() <= scalar.abs() * 1e-5 + 1e-5,
9999            "fast {fast} vs scalar {scalar}"
10000        );
10001    }
10002
10003    #[test]
10004    fn f32_matvec_matches_matvec_rows_bitexact() {
10005        let (rows, cols) = (300, 40);
10006        let w: Vec<f32> = (0..rows * cols).map(|i| (i as f32 * 0.017).sin()).collect();
10007        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.05).cos()).collect();
10008        let qt = QTensor::from_f32(w.clone(), rows, cols);
10009
10010        let mut a = vec![0.0f32; rows];
10011        matvec_rows(None, &w, &x, &mut a);
10012        let mut b = vec![0.0f32; rows];
10013        qt.matvec(&x, &mut b, None);
10014        assert_eq!(a, b);
10015    }
10016
10017    #[test]
10018    fn sdot_kernel_exact_on_grid() {
10019        // Activations already on the i8 grid (±1 with amax=1 → sx=1/127,
10020        // xq=±127 dequantizes EXACTLY) → the SDOT path must match the
10021        // exact f32 dot to float rounding. This isolates kernel
10022        // correctness from quantization noise.
10023        eprintln!("sdot_enabled = {}", sdot_enabled());
10024        let (rows, cols) = (9, 80); // odd rows → exercises 4-row + tail
10025        let w: Vec<u8> = (0..rows * cols)
10026            .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
10027            .collect();
10028        let scales: Vec<f32> = (0..rows).map(|o| 0.005 + o as f32 * 0.001).collect();
10029        let x: Vec<f32> = (0..cols)
10030            .map(|i| match i % 3 {
10031                0 => 1.0,
10032                1 => -1.0,
10033                _ => 0.0,
10034            })
10035            .collect();
10036        let mut a = vec![0.0f32; rows];
10037        qmatvec(
10038            &w,
10039            &[],
10040            &scales,
10041            &x,
10042            &[],
10043            TensorDtype::Q8Row,
10044            rows,
10045            cols,
10046            &mut a,
10047            None,
10048        );
10049        for o in 0..rows {
10050            let mut acc = 0.0f32;
10051            for j in 0..cols {
10052                acc += (w[o * cols + j] as i8) as f32 * x[j];
10053            }
10054            let expect = acc * scales[o];
10055            assert!(
10056                (a[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
10057                "row {o}: {} vs {expect}",
10058                a[o]
10059            );
10060        }
10061    }
10062
10063    #[test]
10064    fn q1_tbl_fast_path_matches_reference() {
10065        // gpr = 8 exercises the TBL pair-load fast loop, and the LAST
10066        // row's final 4-tile window trips the 4B-overread guard (the
10067        // payload ends exactly at the last tile) — both paths must
10068        // agree with the dequant reference.
10069        let (rows, cols) = (5, 256);
10070        let gpr = cols / GROUP_SIZE;
10071        let mut bytes = Vec::new();
10072        for t in 0..rows * gpr {
10073            let s = 0.007 + (t % 11) as f32 * 0.004;
10074            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
10075            for j in 0..4 {
10076                bytes.push(((t * 53 + j * 89 + 7) % 249) as u8);
10077            }
10078        }
10079        let x: Vec<f32> = (0..cols)
10080            .map(|i| if (i * 5) % 7 < 3 { 1.0 } else { -1.0 })
10081            .collect();
10082        let mut w = vec![0.0f32; rows * cols];
10083        cortiq_core::quant::dequant_q1(&bytes, &mut w);
10084        let mut got = vec![0.0f32; rows];
10085        q1_matvec(&bytes, &x, rows, cols, &mut got, None);
10086        for o in 0..rows {
10087            let expect: f32 = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
10088            assert!(
10089                (got[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
10090                "row {o}: {} vs {expect}",
10091                got[o]
10092            );
10093        }
10094        // Blocked 1×4 batch (b=5: one quad + remainder) must equal the
10095        // single-matvec path bit-for-bit.
10096        let b = 5usize;
10097        let mut xs_all = Vec::new();
10098        for bi in 0..b {
10099            xs_all.extend(x.iter().map(|v| if bi % 2 == 0 { *v } else { -*v }));
10100        }
10101        let mut mm = vec![0.0f32; b * rows];
10102        q1_matmat(&bytes, &xs_all, b, rows, cols, &mut mm, None);
10103        for bi in 0..b {
10104            let mut single = vec![0.0f32; rows];
10105            q1_matvec(
10106                &bytes,
10107                &xs_all[bi * cols..(bi + 1) * cols],
10108                rows,
10109                cols,
10110                &mut single,
10111                None,
10112            );
10113            assert_eq!(&mm[bi * rows..(bi + 1) * rows], &single[..], "stream {bi}");
10114        }
10115    }
10116
10117    #[test]
10118    fn q1_kernels_match_exact_reference() {
10119        // Synthetic q1 payload: 6-byte tiles [f16 scale][4B bits].
10120        let (rows, cols) = (7, 96);
10121        let gpr = cols / GROUP_SIZE;
10122        let mut bytes = Vec::new();
10123        for t in 0..rows * gpr {
10124            let s = 0.01 + (t % 13) as f32 * 0.003;
10125            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
10126            for j in 0..4 {
10127                bytes.push(((t * 31 + j * 97) % 251) as u8);
10128            }
10129        }
10130        // On-grid activations (±1, amax 1) → the SDOT path is exact.
10131        let x: Vec<f32> = (0..cols)
10132            .map(|i| if i % 3 == 0 { 1.0 } else { -1.0 })
10133            .collect();
10134        // Reference through the core dequant.
10135        let mut w = vec![0.0f32; rows * cols];
10136        cortiq_core::quant::dequant_q1(&bytes, &mut w);
10137        let mut expect = vec![0.0f32; rows];
10138        for o in 0..rows {
10139            expect[o] = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
10140        }
10141        let mut got = vec![0.0f32; rows];
10142        q1_matvec(&bytes, &x, rows, cols, &mut got, None);
10143        for o in 0..rows {
10144            assert!(
10145                (got[o] - expect[o]).abs() < 1e-3 * expect[o].abs().max(1e-3),
10146                "row {o}: {} vs {}",
10147                got[o],
10148                expect[o]
10149            );
10150        }
10151        // Pair and batch paths agree with the single path.
10152        let x2: Vec<f32> = x.iter().map(|v| -v).collect();
10153        let (mut a1, mut a2) = (vec![0.0f32; rows], vec![0.0f32; rows]);
10154        q1_matvec2(&bytes, &x, &x2, rows, cols, &mut a1, &mut a2, None);
10155        assert_eq!(a1, got);
10156        let mut xs = x.clone();
10157        xs.extend_from_slice(&x2);
10158        let mut mm = vec![0.0f32; 2 * rows];
10159        q1_matmat(&bytes, &xs, 2, rows, cols, &mut mm, None);
10160        assert_eq!(&mm[..rows], got.as_slice());
10161        assert_eq!(&mm[rows..], a2.as_slice());
10162    }
10163
10164    #[test]
10165    fn repack_is_bit_identical() {
10166        // The interleaved-repack kernel must produce EXACTLY the same
10167        // bits as the mmap-layout kernel: integer accumulation is order-
10168        // exact, the f32 epilogue is identical. Odd rows exercise the
10169        // tail; direct range calls exercise unaligned pool splits.
10170        let (rows, cols) = (267, 96); // 66 groups + 3 tail rows, cols % 16 == 0
10171        let w: Vec<u8> = (0..rows * cols)
10172            .map(|i| (((i * 89) % 253) as i32 - 126) as i8 as u8)
10173            .collect();
10174        let scales: Vec<f32> = (0..rows).map(|o| 0.003 + o as f32 * 0.0007).collect();
10175        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.37).sin() * 2.0).collect();
10176        let rep = q8_repack_layout(&w, rows, cols);
10177        // Group interleave round-trips.
10178        for g in 0..rows / 4 {
10179            for c in 0..cols / 16 {
10180                for lane in 0..4 {
10181                    assert_eq!(
10182                        &rep[g * 4 * cols + c * 64 + lane * 16
10183                            ..g * 4 * cols + c * 64 + lane * 16 + 16],
10184                        &w[(g * 4 + lane) * cols + c * 16..(g * 4 + lane) * cols + c * 16 + 16],
10185                    );
10186                }
10187            }
10188        }
10189        let mut a = vec![0.0f32; rows];
10190        qmatvec(
10191            &w,
10192            &[],
10193            &scales,
10194            &x,
10195            &[],
10196            TensorDtype::Q8Row,
10197            rows,
10198            cols,
10199            &mut a,
10200            None,
10201        );
10202        let mut b = vec![0.0f32; rows];
10203        qmatvec(
10204            &w,
10205            &rep,
10206            &scales,
10207            &x,
10208            &[],
10209            TensorDtype::Q8Row,
10210            rows,
10211            cols,
10212            &mut b,
10213            None,
10214        );
10215        assert_eq!(a, b, "full-range repack output diverged");
10216
10217        #[cfg(target_arch = "aarch64")]
10218        if sdot_enabled() {
10219            // Unaligned range split (pool workers get arbitrary bounds).
10220            let act = split_act(&x);
10221            let mut c1 = vec![0.0f32; rows];
10222            let mut c2 = vec![0.0f32; rows];
10223            q8_range_sdot(
10224                &w,
10225                &[],
10226                &scales,
10227                &act,
10228                cols,
10229                SendMut(c1.as_mut_ptr()),
10230                3,
10231                rows - 2,
10232            );
10233            q8_range_sdot(
10234                &w,
10235                &rep,
10236                &scales,
10237                &act,
10238                cols,
10239                SendMut(c2.as_mut_ptr()),
10240                3,
10241                rows - 2,
10242            );
10243            assert_eq!(c1, c2, "unaligned-range repack output diverged");
10244        }
10245    }
10246
10247    #[test]
10248    fn sdot_a8w8_noise_is_bounded() {
10249        // Off-grid activations: A8 quantization noise must stay small in
10250        // relative L2 over the whole output (realistic accuracy contract;
10251        // vmfcore measured argmax-identical decode on real models).
10252        let (rows, cols) = (16, 512);
10253        let w: Vec<u8> = (0..rows * cols)
10254            .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
10255            .collect();
10256        let scales = vec![0.01f32; rows];
10257        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.21).sin()).collect();
10258        let mut a = vec![0.0f32; rows];
10259        qmatvec(
10260            &w,
10261            &[],
10262            &scales,
10263            &x,
10264            &[],
10265            TensorDtype::Q8Row,
10266            rows,
10267            cols,
10268            &mut a,
10269            None,
10270        );
10271        let (mut num, mut den) = (0f64, 0f64);
10272        for o in 0..rows {
10273            let mut acc = 0.0f32;
10274            for j in 0..cols {
10275                acc += (w[o * cols + j] as i8) as f32 * x[j];
10276            }
10277            let expect = acc * scales[o];
10278            num += ((a[o] - expect) as f64).powi(2);
10279            den += (expect as f64).powi(2);
10280        }
10281        let rel = (num / den.max(1e-12)).sqrt();
10282        assert!(rel < 0.05, "A8W8 relative L2 error too high: {rel}");
10283    }
10284
10285    #[test]
10286    fn i8_dot_neon_matches_scalar() {
10287        let n = 100;
10288        let w: Vec<u8> = (0..n).map(|i| ((i * 37 + 11) % 251) as u8).collect();
10289        let x: Vec<f32> = (0..n).map(|i| (i as f32 * 0.13).sin()).collect();
10290        let mut scalar = 0.0f32;
10291        for j in 0..n {
10292            scalar += (w[j] as i8) as f32 * x[j];
10293        }
10294        let fast = dot_i8_f32(&w, &x);
10295        assert!((scalar - fast).abs() < 1e-3 * scalar.abs().max(1.0));
10296    }
10297
10298    /// Fused vbit matvec must match full dequant_vbit + dense matvec.
10299    #[test]
10300    fn vbitmatvec_matches_full_dequant() {
10301        let (rows, cols) = (6, 64);
10302        let ng = cols / GROUP_SIZE;
10303        // Hand-craft: bits per row, f16 scales, packed rows.
10304        let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
10305        let mut bytes = bits.clone();
10306        for g in 0..rows * ng {
10307            let s = 0.02 + 0.001 * g as f32;
10308            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
10309        }
10310        for r in 0..rows {
10311            let b = bits[r] as usize;
10312            let (mut acc, mut nb) = (0u64, 0usize);
10313            let mut rowbytes = Vec::new();
10314            for i in 0..cols {
10315                let v = ((i * 7 + r * 13) % (1 << b)) as u64;
10316                acc = (acc << b) | v;
10317                nb += b;
10318                while nb >= 8 {
10319                    nb -= 8;
10320                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
10321                }
10322            }
10323            if nb > 0 {
10324                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
10325            }
10326            bytes.extend_from_slice(&rowbytes);
10327        }
10328        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
10329
10330        let mut reference = vec![0f32; rows * cols];
10331        cortiq_core::quant::dequant_vbit(&bytes, rows, cols, &mut reference).unwrap();
10332        let mut expect = vec![0f32; rows];
10333        for r in 0..rows {
10334            expect[r] = reference[r * cols..(r + 1) * cols]
10335                .iter()
10336                .zip(&x)
10337                .map(|(w, xv)| w * xv)
10338                .sum();
10339        }
10340        let mut got = vec![0f32; rows];
10341        let offsets = vbit_row_offsets(&bytes, rows, cols);
10342        vbitmatvec(&bytes, &offsets, &x, rows, cols, &mut got, None);
10343        // SDOT path quantizes activations to i8 (A8W8): bounded noise,
10344        // same contract as q8 (exact path is pinned by CMF_SDOT=0 in
10345        // the golden-parity gate).
10346        let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
10347        let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
10348        for r in 0..rows {
10349            assert!(
10350                (got[r] - expect[r]).abs() < tol * scale,
10351                "row {r}: {} vs {}",
10352                got[r],
10353                expect[r]
10354            );
10355        }
10356    }
10357
10358    /// Fused q4 matvec must match the reference full-dequant + dense
10359    /// matvec bit-for-bit in structure (same f32 math, group order).
10360    /// vbit matmat: the blocked 1×4 leg must match the per-row path
10361    /// (paired env toggle; larger shape so both code paths engage).
10362    #[test]
10363    #[cfg(target_arch = "x86_64")]
10364    fn vbit_matmat_blocked_matches_per_row() {
10365        let (rows, cols, b) = (64usize, 128usize, 9usize);
10366        let ng = cols / GROUP_SIZE;
10367        let bits: Vec<u8> = (0..rows).map(|r| [3u8, 4, 5, 6][r % 4]).collect();
10368        let mut bytes = bits.clone();
10369        for g in 0..rows * ng {
10370            let sc = 0.02 + 0.0005 * g as f32;
10371            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
10372        }
10373        for r in 0..rows {
10374            let bw = bits[r] as usize;
10375            let (mut acc, mut nb) = (0u64, 0usize);
10376            let mut rowbytes = Vec::new();
10377            for i in 0..cols {
10378                let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
10379                acc = (acc << bw) | v;
10380                nb += bw;
10381                while nb >= 8 {
10382                    nb -= 8;
10383                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
10384                }
10385            }
10386            if nb > 0 {
10387                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
10388            }
10389            bytes.extend_from_slice(&rowbytes);
10390        }
10391        let x: Vec<f32> = (0..b * cols)
10392            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
10393            .collect();
10394        let offsets = vbit_row_offsets(&bytes, rows, cols);
10395        let mut y_a = vec![0f32; b * rows];
10396        let mut y_b = vec![0f32; b * rows];
10397        unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
10398        vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_a, None);
10399        unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
10400        vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_b, None);
10401        unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
10402        let max_d = y_a
10403            .iter()
10404            .zip(&y_b)
10405            .map(|(p, q)| (p - q).abs())
10406            .fold(0.0f32, f32::max);
10407        assert!(max_d < 1e-4, "vbit blocked ≠ per-row: max|Δ| = {max_d}");
10408    }
10409
10410    /// q4t blocked 1×4 (SDOT on ARM, AVX2 on x86) must equal the
10411    /// per-row path exactly: same nibble unpack, same group order,
10412    /// same f32 accumulation — batch == matvec bit-for-bit. b=9 covers
10413    /// two full 1×4 blocks plus a remainder through the single-row
10414    /// kernel. (Both paths produce identical output, so the shared
10415    /// CMF_X86_BLOCKED env var racing with other tests cannot flip
10416    /// the verdict — worst case both sides take the same path.)
10417    #[test]
10418    fn q4t_matmat_blocked_matches_per_row() {
10419        let (rows, cols, b) = (16usize, 64usize, 9usize);
10420        let gpr = cols / GROUP_SIZE;
10421        let mut bytes = vec![0u8; rows * gpr * Q4_TILE];
10422        for r in 0..rows {
10423            for g in 0..gpr {
10424                let t = (r * gpr + g) * Q4_TILE;
10425                let sc = 0.02 + 0.001 * (r * gpr + g) as f32;
10426                bytes[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
10427                for k in 0..16 {
10428                    bytes[t + 2 + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
10429                }
10430            }
10431        }
10432        let x: Vec<f32> = (0..b * cols)
10433            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
10434            .collect();
10435        let mut y_blk = vec![0f32; b * rows];
10436        let mut y_row = vec![0f32; b * rows];
10437        unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
10438        q4t_matmat(&bytes, &x, b, rows, cols, &mut y_blk, None);
10439        unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
10440        q4t_matmat(&bytes, &x, b, rows, cols, &mut y_row, None);
10441        unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
10442        assert_eq!(y_blk, y_row, "q4t blocked 1x4 ≠ per-row");
10443    }
10444
10445    /// The wide-batch Accelerate arm of q4t_matmat vs a brute-force
10446    /// f32 dequant matmul: both are f32 GEMMs, so only reduction
10447    /// order differs — tight tolerance.
10448    /// A synthetic q4tp payload: random nibbles plus a per-row ladder whose
10449    /// span varies row to row, so the codes actually exercise the full 0..31
10450    /// range rather than clustering on one rung.
10451    fn synth_q4tp(rows: usize, cols: usize) -> Vec<u8> {
10452        use cortiq_core::quant::{f32_to_f16, q4tp_code_stride, q4tp_put_code};
10453        let gpr = cols / GROUP_SIZE;
10454        let stride = q4tp_code_stride(gpr);
10455        let (params_off, codes_off, _) = q4tp_sections(rows, cols);
10456        let mut b = vec![0u8; codes_off + rows * stride];
10457        for r in 0..rows {
10458            for g in 0..gpr {
10459                let t = (r * gpr + g) * Q4TP_NIB;
10460                for k in 0..16 {
10461                    b[t + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
10462                }
10463            }
10464            let lo = -6.0 - 0.03 * (r % 17) as f32;
10465            let step = 0.01 + 0.004 * (r % 11) as f32;
10466            let p = params_off + r * 4;
10467            b[p..p + 2].copy_from_slice(&f32_to_f16(lo).to_le_bytes());
10468            b[p + 2..p + 4].copy_from_slice(&f32_to_f16(step).to_le_bytes());
10469            let crow = &mut b[codes_off + r * stride..codes_off + (r + 1) * stride];
10470            for g in 0..gpr {
10471                q4tp_put_code(crow, g, (r * 5 + g * 3) % 32);
10472            }
10473        }
10474        b
10475    }
10476
10477    /// The same weights re-expressed as q4_tiled, so the proven kernel can
10478    /// be the reference: each tile stores the ladder scale its code selects.
10479    /// Only the f16 rounding of that scale separates the two payloads.
10480    fn q4tp_as_q4t(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
10481        let gpr = cols / GROUP_SIZE;
10482        let v = Q4tpView::new(bytes, rows, cols);
10483        let mut out = vec![0u8; rows * gpr * Q4_TILE];
10484        let mut sc = vec![0f32; gpr];
10485        for r in 0..rows {
10486            v.scales_into(r, gpr, &mut sc);
10487            for g in 0..gpr {
10488                let t = (r * gpr + g) * Q4_TILE;
10489                let s = sc[g];
10490                out[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
10491                let src = (r * gpr + g) * Q4TP_NIB;
10492                out[t + 2..t + Q4_TILE].copy_from_slice(&v.nib[src..src + Q4TP_NIB]);
10493            }
10494        }
10495        out
10496    }
10497
10498    /// The exact (`CMF_SDOT=0`) path must reproduce `dequant_q4tp` to f32
10499    /// rounding — that scalar routine is the format's definition, and the
10500    /// kernels re-derive the scale from the ladder independently. Call the
10501    /// row kernel directly: `matmat` picks the int8 arm when a8w8 is on,
10502    /// so routing through it would test the other path by accident.
10503    #[test]
10504    fn q4tp_exact_path_matches_dequant_reference() {
10505        let (rows, cols) = (256usize, 512usize);
10506        let gpr = cols / GROUP_SIZE;
10507        let bytes = synth_q4tp(rows, cols);
10508        let mut w = vec![0f32; rows * cols];
10509        cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
10510
10511        let x: Vec<f32> = (0..cols)
10512            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
10513            .collect();
10514        let v = Q4tpView::new(&bytes, rows, cols);
10515        let mut sc = vec![0f32; gpr];
10516        for r in 0..rows {
10517            v.scales_into(r, gpr, &mut sc);
10518            let got = q4tp_row_exact(v.nib, r, gpr, &x, &sc);
10519            let want: f32 = (0..cols).map(|c| w[r * cols + c] * x[c]).sum();
10520            // These dot products cancel down to ~1e-3 from terms of ~5e-2, so
10521            // the meaningful yardstick is the summed magnitude, not the result:
10522            // against the result any reordering of a 512-term f32 sum "fails".
10523            let mag: f32 = (0..cols).map(|c| (w[r * cols + c] * x[c]).abs()).sum();
10524            assert!(
10525                (got - want).abs() <= 1e-5 * mag,
10526                "row {r}: kernel {got} vs dequant {want}"
10527            );
10528        }
10529    }
10530
10531    /// The int8 (a8w8) path can't be checked against an f32 reference — the
10532    /// activation quantization dominates. Check it against the q4t kernel it
10533    /// was ported from instead, on payloads holding the same weights: that
10534    /// isolates exactly what the port could break (16 B stride, ladder
10535    /// lookup, nibble unpack) from what it deliberately shares.
10536    #[test]
10537    fn q4tp_matvec_matches_the_q4t_kernel_it_was_ported_from() {
10538        let (rows, cols) = (256usize, 512usize);
10539        let bytes = synth_q4tp(rows, cols);
10540        let twin = q4tp_as_q4t(&bytes, rows, cols);
10541        let x: Vec<f32> = (0..cols)
10542            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
10543            .collect();
10544
10545        let mut got = vec![0f32; rows];
10546        q4tp_matvec(&bytes, &x, rows, cols, &mut got, None);
10547        let mut want = vec![0f32; rows];
10548        q4t_matvec(&twin, &x, rows, cols, &mut want, None);
10549
10550        // Scale is f16 in the twin and f32 here, so allow that rounding on
10551        // top of the summed magnitude (same cancellation argument as above).
10552        let mut w = vec![0f32; rows * cols];
10553        cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
10554        for r in 0..rows {
10555            let mag: f32 = (0..cols).map(|c| (w[r * cols + c] * x[c]).abs()).sum();
10556            assert!(
10557                (got[r] - want[r]).abs() <= 1e-3 * mag,
10558                "row {r}: q4tp {} vs q4t {}",
10559                got[r],
10560                want[r]
10561            );
10562        }
10563    }
10564
10565    /// `matmat` carries three arms (Accelerate, blocked int8 1x4, scalar).
10566    /// Batch 5 crosses the blocked kernel's stride, so this exercises the
10567    /// 1x4 path AND its scalar tail in one run — the blocked kernel is new
10568    /// code and its four accumulators are exactly what tends to go wrong.
10569    #[test]
10570    fn q4tp_matmat_matches_the_q4t_kernel_it_was_ported_from() {
10571        let (rows, cols, b) = (256usize, 512usize, 5usize);
10572        let bytes = synth_q4tp(rows, cols);
10573        let twin = q4tp_as_q4t(&bytes, rows, cols);
10574        let xs: Vec<f32> = (0..b * cols)
10575            .map(|i| ((i * 29 + 11) % 89) as f32 / 89.0 - 0.5)
10576            .collect();
10577
10578        let mut got = vec![0f32; b * rows];
10579        q4tp_matmat(&bytes, &xs, b, rows, cols, &mut got, None);
10580        let mut want = vec![0f32; b * rows];
10581        q4t_matmat(&twin, &xs, b, rows, cols, &mut want, None);
10582
10583        let mut w = vec![0f32; rows * cols];
10584        cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
10585        for t in 0..b {
10586            for r in 0..rows {
10587                let mag: f32 = (0..cols)
10588                    .map(|c| (w[r * cols + c] * xs[t * cols + c]).abs())
10589                    .sum();
10590                let (g, wa) = (got[t * rows + r], want[t * rows + r]);
10591                assert!(
10592                    (g - wa).abs() <= 1e-3 * mag,
10593                    "batch {t} row {r}: q4tp {g} vs q4t {wa}"
10594                );
10595            }
10596        }
10597    }
10598
10599    #[test]
10600    fn q4tp_matvec2_matches_the_single_stream_kernel() {
10601        let (rows, cols) = (128usize, 256usize);
10602        let gpr = cols / GROUP_SIZE;
10603        let bytes = synth_q4tp(rows, cols);
10604        let xs: Vec<f32> = (0..2 * cols)
10605            .map(|i| ((i * 29 + 11) % 89) as f32 / 89.0 - 0.5)
10606            .collect();
10607
10608        let (mut o1, mut o2) = (vec![0f32; rows], vec![0f32; rows]);
10609        q4tp_matvec2(
10610            &bytes,
10611            &xs[..cols],
10612            &xs[cols..],
10613            rows,
10614            cols,
10615            &mut o1,
10616            &mut o2,
10617            None,
10618        );
10619
10620        // matvec2 takes the exact path for both streams, so the single-row
10621        // kernel is an exact reference — no tolerance for path differences.
10622        let v = Q4tpView::new(&bytes, rows, cols);
10623        let mut sc = vec![0f32; gpr];
10624        for r in 0..rows {
10625            v.scales_into(r, gpr, &mut sc);
10626            assert_eq!(o1[r], q4tp_row_exact(v.nib, r, gpr, &xs[..cols], &sc));
10627            assert_eq!(o2[r], q4tp_row_exact(v.nib, r, gpr, &xs[cols..], &sc));
10628        }
10629    }
10630
10631    /// q4tp must not COST speed — it exists to save bytes, and a format that
10632    /// trades 7% of a file for a slower model is a bad trade. This guard is
10633    /// here because correctness tests happily passed while `q4tp_matmat` was
10634    /// missing its int8 and Accelerate arms and the model ran 5x slower.
10635    /// Measured on M-series: 0.97-1.04x, i.e. parity (16 B tiles are better
10636    /// aligned than q4t's 18 B, which pays for the scale indirection).
10637    #[test]
10638    fn q4tp_matvec_keeps_pace_with_q4t() {
10639        let (rows, cols) = (4096usize, 3072usize);
10640        let bytes = synth_q4tp(rows, cols);
10641        let twin = q4tp_as_q4t(&bytes, rows, cols);
10642        let x: Vec<f32> = (0..cols).map(|i| (i % 97) as f32 / 97.0 - 0.5).collect();
10643        let mut o = vec![0f32; rows];
10644        let n = 12;
10645        let mut best = (f64::MAX, f64::MAX);
10646        // Interleaved A/B, minimum statistic: this machine throttles, and a
10647        // mean over a thermal ramp reliably indicts whichever ran second.
10648        for _ in 0..3 {
10649            let t0 = std::time::Instant::now();
10650            for _ in 0..n {
10651                q4t_matvec(&twin, &x, rows, cols, &mut o, None);
10652            }
10653            best.0 = best.0.min(t0.elapsed().as_secs_f64());
10654            let t0 = std::time::Instant::now();
10655            for _ in 0..n {
10656                q4tp_matvec(&bytes, &x, rows, cols, &mut o, None);
10657            }
10658            best.1 = best.1.min(t0.elapsed().as_secs_f64());
10659        }
10660        let ratio = best.1 / best.0;
10661        println!(
10662            "q4t {:.3} ms | q4tp {:.3} ms | {ratio:.2}x",
10663            best.0 * 1e3 / n as f64,
10664            best.1 * 1e3 / n as f64
10665        );
10666        assert!(ratio < 2.0, "q4tp matvec {ratio:.2}x slower than q4t");
10667    }
10668
10669    #[cfg(target_os = "macos")]
10670    #[test]
10671    fn q4t_matmat_accel_matches_dequant_reference() {
10672        if !accel_gemm_enabled() {
10673            return; // CMF_ACCEL=0
10674        }
10675        let (rows, cols, b) = (512usize, 1024usize, 8usize); // ≥500K → accel arm
10676        let gpr = cols / GROUP_SIZE;
10677        let mut bytes = vec![0u8; rows * gpr * Q4_TILE];
10678        for r in 0..rows {
10679            for g in 0..gpr {
10680                let t = (r * gpr + g) * Q4_TILE;
10681                let sc = 0.02 + 0.0005 * ((r * gpr + g) % 64) as f32;
10682                bytes[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
10683                for k in 0..16 {
10684                    bytes[t + 2 + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
10685                }
10686            }
10687        }
10688        let x: Vec<f32> = (0..b * cols)
10689            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
10690            .collect();
10691        let mut got = vec![0f32; b * rows];
10692        q4t_matmat(&bytes, &x, b, rows, cols, &mut got, None);
10693        // Brute-force reference off the same tiles.
10694        let mut w = vec![0f32; rows * cols];
10695        for r in 0..rows {
10696            for g in 0..gpr {
10697                let t = (r * gpr + g) * Q4_TILE;
10698                let s = f16_to_f32(u16::from_le_bytes([bytes[t], bytes[t + 1]]));
10699                for (k, &bb) in bytes[t + 2..t + Q4_TILE].iter().enumerate() {
10700                    w[r * cols + g * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s;
10701                    w[r * cols + g * GROUP_SIZE + k * 2 + 1] =
10702                        (((bb >> 4) & 0x0F) as f32 - 8.0) * s;
10703                }
10704            }
10705        }
10706        for bi in 0..b {
10707            for r in 0..rows {
10708                let want: f32 = (0..cols).map(|j| x[bi * cols + j] * w[r * cols + j]).sum();
10709                let d = (got[bi * rows + r] - want).abs();
10710                assert!(
10711                    d <= want.abs().max(1.0) * 1e-4,
10712                    "accel q4t GEMM diverged at ({bi},{r}): {} vs {want}",
10713                    got[bi * rows + r]
10714                );
10715            }
10716        }
10717    }
10718
10719    #[test]
10720    fn q4matvec_matches_full_dequant() {
10721        let (rows, cols) = (8, 64);
10722        let groups = rows * cols / GROUP_SIZE;
10723        // Hand-craft a q4_block blob: nibbles then f16 scales.
10724        let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
10725        for i in 0..groups * 16 {
10726            bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
10727        }
10728        for g in 0..groups {
10729            let s = 0.01 + 0.003 * g as f32;
10730            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
10731        }
10732        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
10733
10734        let mut reference = vec![0.0f32; rows * cols];
10735        cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
10736        let mut expect = vec![0.0f32; rows];
10737        for r in 0..rows {
10738            expect[r] = reference[r * cols..(r + 1) * cols]
10739                .iter()
10740                .zip(&x)
10741                .map(|(w, xv)| w * xv)
10742                .sum();
10743        }
10744
10745        let mut got = vec![0.0f32; rows];
10746        q4matvec(&bytes, &x, rows, cols, &mut got, None);
10747        // SDOT path quantizes activations to i8 (A8W8): bounded noise,
10748        // same contract as q8/vbit (exact path is pinned by CMF_SDOT=0
10749        // in the golden-parity gate).
10750        let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
10751        let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
10752        for r in 0..rows {
10753            assert!(
10754                (got[r] - expect[r]).abs() < tol * scale,
10755                "row {r}: {} vs {}",
10756                got[r],
10757                expect[r]
10758            );
10759        }
10760    }
10761
10762    /// Fused two-input vbit matvec must equal two single matvecs exactly
10763    /// (same per-lane accumulation order on both scalar and SDOT paths).
10764    #[test]
10765    fn vbitmatvec2_equals_two_singles() {
10766        let (rows, cols) = (6, 64);
10767        let ng = cols / GROUP_SIZE;
10768        let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
10769        let mut bytes = bits.clone();
10770        for g in 0..rows * ng {
10771            let s = 0.02 + 0.001 * g as f32;
10772            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
10773        }
10774        for r in 0..rows {
10775            let b = bits[r] as usize;
10776            let (mut acc, mut nb) = (0u64, 0usize);
10777            let mut rowbytes = Vec::new();
10778            for i in 0..cols {
10779                let v = ((i * 7 + r * 13) % (1 << b)) as u64;
10780                acc = (acc << b) | v;
10781                nb += b;
10782                while nb >= 8 {
10783                    nb -= 8;
10784                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
10785                }
10786            }
10787            if nb > 0 {
10788                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
10789            }
10790            bytes.extend_from_slice(&rowbytes);
10791        }
10792        let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
10793        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).cos()).collect();
10794        let offsets = vbit_row_offsets(&bytes, rows, cols);
10795
10796        let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
10797        vbitmatvec(&bytes, &offsets, &x1, rows, cols, &mut a1, None);
10798        vbitmatvec(&bytes, &offsets, &x2, rows, cols, &mut a2, None);
10799        let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
10800        vbitmatvec2(
10801            &bytes, &offsets, &x1, &x2, rows, cols, &mut b1, &mut b2, None,
10802        );
10803        assert_eq!(a1, b1, "fused vbit lane 1 must be bit-identical");
10804        assert_eq!(a2, b2, "fused vbit lane 2 must be bit-identical");
10805    }
10806
10807    /// Fused two-input q4 matvec must equal two single matvecs exactly.
10808    #[test]
10809    fn q4matvec2_equals_two_singles() {
10810        let (rows, cols) = (8, 128);
10811        let groups = rows * cols / GROUP_SIZE;
10812        let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
10813        for i in 0..groups * 16 {
10814            bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
10815        }
10816        for g in 0..groups {
10817            let s = 0.01 + 0.003 * g as f32;
10818            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
10819        }
10820        // Include an outlier channel so the SDOT correction path is
10821        // exercised in the pair kernel too.
10822        let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
10823        x1[9] = 250.0;
10824        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
10825
10826        let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
10827        q4matvec(&bytes, &x1, rows, cols, &mut a1, None);
10828        q4matvec(&bytes, &x2, rows, cols, &mut a2, None);
10829        let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
10830        q4matvec2(&bytes, &x1, &x2, rows, cols, &mut b1, &mut b2, None);
10831        assert_eq!(a1, b1, "fused q4 lane 1 must be bit-identical");
10832        assert_eq!(a2, b2, "fused q4 lane 2 must be bit-identical");
10833    }
10834
10835    /// Multi-matrix job must equal separate matvecs exactly — same
10836    /// kernels, only the dispatch is fused.
10837    #[test]
10838    fn matvec_many_equals_separate_matvecs() {
10839        use crate::pool::Pool;
10840        let (r1, r2, cols) = (300, 200, 64);
10841        let mk = |salt: usize, rows: usize| {
10842            QTensor::from_f32(
10843                (0..rows * cols)
10844                    .map(|i| ((i * 7 + salt) % 97) as f32 / 97.0 - 0.5)
10845                    .collect(),
10846                rows,
10847                cols,
10848            )
10849        };
10850        let (a, b) = (mk(1, r1), mk(5, r2));
10851        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).sin()).collect();
10852        let pool = Pool::new(3);
10853
10854        let (mut ea, mut eb) = (vec![0f32; r1], vec![0f32; r2]);
10855        a.matvec(&x, &mut ea, Some(&pool));
10856        b.matvec(&x, &mut eb, Some(&pool));
10857        let (mut ga, mut gb) = (vec![0f32; r1], vec![0f32; r2]);
10858        QTensor::matvec_many([&a, &b], &x, [&mut ga, &mut gb], Some(&pool));
10859        assert_eq!(ea, ga, "fused multi-matrix lane 1 must be bit-identical");
10860        assert_eq!(eb, gb, "fused multi-matrix lane 2 must be bit-identical");
10861    }
10862
10863    /// Batched q4/vbit matmat must equal per-position matvec calls
10864    /// exactly (the fallback it replaced) — same kernels, same order.
10865    #[test]
10866    fn batched_matmat_equals_per_position_matvec() {
10867        let (rows, cols, b) = (8, 64, 5);
10868        // q4 blob.
10869        let groups = rows * cols / GROUP_SIZE;
10870        let mut q4 = Vec::new();
10871        for i in 0..groups * 16 {
10872            q4.push((((i * 7 + 3) % 256) & 0xFF) as u8);
10873        }
10874        for g in 0..groups {
10875            q4.extend_from_slice(
10876                &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
10877            );
10878        }
10879        // vbit blob (mixed widths incl. 8).
10880        let ng = cols / GROUP_SIZE;
10881        let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4, 5, 3];
10882        let mut vb = bits.clone();
10883        for g in 0..rows * ng {
10884            vb.extend_from_slice(
10885                &cortiq_core::quant::f32_to_f16(0.02 + 0.001 * g as f32).to_le_bytes(),
10886            );
10887        }
10888        for r in 0..rows {
10889            let bw = bits[r] as usize;
10890            let (mut acc, mut nb) = (0u64, 0usize);
10891            let mut rowbytes = Vec::new();
10892            for i in 0..cols {
10893                let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
10894                acc = (acc << bw) | v;
10895                nb += bw;
10896                while nb >= 8 {
10897                    nb -= 8;
10898                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
10899                }
10900            }
10901            if nb > 0 {
10902                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
10903            }
10904            vb.extend_from_slice(&rowbytes);
10905        }
10906        let offsets = vbit_row_offsets(&vb, rows, cols);
10907
10908        let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
10909
10910        // q4: batch vs singles.
10911        let mut got = vec![0f32; b * rows];
10912        q4matmat(&q4, &xs, b, rows, cols, &mut got, None);
10913        for bi in 0..b {
10914            let mut expect = vec![0f32; rows];
10915            q4matvec(
10916                &q4,
10917                &xs[bi * cols..(bi + 1) * cols],
10918                rows,
10919                cols,
10920                &mut expect,
10921                None,
10922            );
10923            assert_eq!(
10924                &got[bi * rows..(bi + 1) * rows],
10925                &expect[..],
10926                "q4 batch pos {bi}"
10927            );
10928        }
10929
10930        // vbit: batch vs singles.
10931        let mut got = vec![0f32; b * rows];
10932        vbitmatmat(&vb, &offsets, &xs, b, rows, cols, &mut got, None);
10933        for bi in 0..b {
10934            let mut expect = vec![0f32; rows];
10935            vbitmatvec(
10936                &vb,
10937                &offsets,
10938                &xs[bi * cols..(bi + 1) * cols],
10939                rows,
10940                cols,
10941                &mut expect,
10942                None,
10943            );
10944            assert_eq!(
10945                &got[bi * rows..(bi + 1) * rows],
10946                &expect[..],
10947                "vbit batch pos {bi}"
10948            );
10949        }
10950    }
10951
10952    /// q4_tiled kernels must produce BIT-identical outputs to the q4
10953    /// split kernels on the same values (same ints, same order — only
10954    /// the byte placement differs).
10955    #[test]
10956    fn q4_tiled_matches_q4_block_bitexact() {
10957        let (rows, cols, b) = (8usize, 128usize, 3usize);
10958        let groups = rows * cols / GROUP_SIZE;
10959        let mut split = Vec::with_capacity(groups * 18);
10960        for i in 0..groups * 16 {
10961            split.push((((i * 7 + 3) % 256) & 0xFF) as u8);
10962        }
10963        for g in 0..groups {
10964            split.extend_from_slice(
10965                &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
10966            );
10967        }
10968        // Re-tile: [scale][nibbles] per group.
10969        let (packed, scales) = split.split_at(groups * 16);
10970        let mut tiled = Vec::with_capacity(groups * Q4_TILE);
10971        for g in 0..groups {
10972            tiled.extend_from_slice(&scales[g * 2..g * 2 + 2]);
10973            tiled.extend_from_slice(&packed[g * 16..(g + 1) * 16]);
10974        }
10975
10976        let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
10977        x1[9] = 250.0; // exercise the outlier path
10978        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
10979
10980        let (mut a, mut t) = (vec![0f32; rows], vec![0f32; rows]);
10981        q4matvec(&split, &x1, rows, cols, &mut a, None);
10982        q4t_matvec(&tiled, &x1, rows, cols, &mut t, None);
10983        assert_eq!(a, t, "q4t matvec must match q4 bit-for-bit");
10984
10985        let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
10986        let (mut t1, mut t2) = (vec![0f32; rows], vec![0f32; rows]);
10987        q4matvec2(&split, &x1, &x2, rows, cols, &mut a1, &mut a2, None);
10988        q4t_matvec2(&tiled, &x1, &x2, rows, cols, &mut t1, &mut t2, None);
10989        assert_eq!(a1, t1);
10990        assert_eq!(a2, t2);
10991
10992        let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
10993        let (mut am, mut tm) = (vec![0f32; b * rows], vec![0f32; b * rows]);
10994        q4matmat(&split, &xs, b, rows, cols, &mut am, None);
10995        q4t_matmat(&tiled, &xs, b, rows, cols, &mut tm, None);
10996        assert_eq!(am, tm, "q4t matmat must match q4 bit-for-bit");
10997    }
10998
10999    /// q4 SDOT outlier correction: a single huge activation channel
11000    /// (>8·rms → outlier, zeroed in xq) must still contribute its EXACT
11001    /// term. On-grid bulk (±1/0 → xq dequantizes exactly) isolates the
11002    /// correction from A8W8 noise. cols must exceed 64: at n=64 the
11003    /// 8·rms threshold equals sqrt(v²+rest) ≥ v, so a single outlier
11004    /// can never qualify (8² = n).
11005    #[test]
11006    fn q4matvec_sdot_outlier_exact() {
11007        let (rows, cols) = (4, 128);
11008        let groups = rows * cols / GROUP_SIZE;
11009        let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
11010        for i in 0..groups * 16 {
11011            bytes.push(((i * 11 + 5) % 256) as u8);
11012        }
11013        for g in 0..groups {
11014            let s = 0.02 + 0.002 * g as f32;
11015            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
11016        }
11017        let mut x: Vec<f32> = (0..cols)
11018            .map(|i| match i % 3 {
11019                0 => 1.0,
11020                1 => -1.0,
11021                _ => 0.0,
11022            })
11023            .collect();
11024        x[17] = 300.0; // ≫ 8·rms → outlier channel
11025
11026        let mut reference = vec![0.0f32; rows * cols];
11027        cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
11028        let mut expect = vec![0.0f32; rows];
11029        for r in 0..rows {
11030            expect[r] = reference[r * cols..(r + 1) * cols]
11031                .iter()
11032                .zip(&x)
11033                .map(|(w, xv)| w * xv)
11034                .sum();
11035        }
11036        let mut got = vec![0.0f32; rows];
11037        q4matvec(&bytes, &x, rows, cols, &mut got, None);
11038        let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
11039        for r in 0..rows {
11040            assert!(
11041                (got[r] - expect[r]).abs() < 2e-3 * scale,
11042                "row {r}: {} vs {} (outlier term must be exact)",
11043                got[r],
11044                expect[r]
11045            );
11046        }
11047    }
11048
11049    /// The fused q1t matvec must equal the reference (dequant_q1t → dot),
11050    /// including the ternary zero level and the binary-searched outlier
11051    /// overlay. Guards the mmap kernel that makes a 12B q1t runnable.
11052    #[test]
11053    fn q1t_matvec_matches_reference() {
11054        use cortiq_core::quant::{dequant_q1t, f32_to_f16};
11055        let (rows, cols) = (3usize, 64usize); // gpr = 2
11056        let gpr = cols / GROUP_SIZE;
11057        let scales = [0.5f32, 0.3, 0.7, 0.2, 0.6, 0.15];
11058        // Overlay (must be sorted by flat index): a few spikes across rows.
11059        let outliers: [(u32, f32); 3] = [(5, 9.0), (70, -4.5), (150, 3.25)];
11060        let is_out = |flat: usize| outliers.iter().any(|&(i, _)| i as usize == flat);
11061        let mut bytes = Vec::new();
11062        for r in 0..rows {
11063            for g in 0..gpr {
11064                bytes.extend_from_slice(&f32_to_f16(scales[r * gpr + g]).to_le_bytes());
11065                let mut c = [0u8; 7];
11066                for k in 0..GROUP_SIZE {
11067                    // Encoder invariant: code 0 at outlier positions.
11068                    let code = if is_out(r * cols + g * GROUP_SIZE + k) {
11069                        0
11070                    } else {
11071                        ((k + r * 3 + g) % 3) as u8 // 0,1,2
11072                    };
11073                    cortiq_core::quant::q1t_pack(&mut c, k, code);
11074                }
11075                bytes.extend_from_slice(&c);
11076            }
11077        }
11078        // Per-row overlay: [u32 row_ptr[rows+1]] then [(u16 col, f16 val)] by
11079        // row (outliers are sorted by flat index → already grouped by row).
11080        let mut row_ptr = vec![0u32; rows + 1];
11081        for &(idx, _) in &outliers {
11082            row_ptr[idx as usize / cols + 1] += 1;
11083        }
11084        for r in 0..rows {
11085            row_ptr[r + 1] += row_ptr[r];
11086        }
11087        for &p in &row_ptr {
11088            bytes.extend_from_slice(&p.to_le_bytes());
11089        }
11090        for &(idx, v) in &outliers {
11091            bytes.extend_from_slice(&((idx as usize % cols) as u16).to_le_bytes());
11092            bytes.extend_from_slice(&f32_to_f16(v).to_le_bytes());
11093        }
11094
11095        let mut refw = vec![0f32; rows * cols];
11096        dequant_q1t(&bytes, rows, cols, &mut refw);
11097        // On-grid activations (±1, amax 1) so the int8 SDOT path reconstructs
11098        // x exactly and matches the f32 reference (same trick as the q1 test).
11099        let x: Vec<f32> = (0..cols)
11100            .map(|j| if j % 3 == 0 { 1.0 } else { -1.0 })
11101            .collect();
11102        let mut expect = vec![0f32; rows];
11103        for r in 0..rows {
11104            let mut a = 0.0f32;
11105            for j in 0..cols {
11106                a += refw[r * cols + j] * x[j];
11107            }
11108            expect[r] = a;
11109        }
11110        let tol = |e: f32| 1e-3 * e.abs().max(1e-3);
11111        let mut got = vec![0f32; rows];
11112        q1t_matvec(&bytes, &x, rows, cols, &mut got, None);
11113        for r in 0..rows {
11114            assert!(
11115                (got[r] - expect[r]).abs() < tol(expect[r]),
11116                "row {r}: {} vs {}",
11117                got[r],
11118                expect[r]
11119            );
11120        }
11121        // matmat (b=2, f32 decode path) must agree too.
11122        let x2: Vec<f32> = x.iter().chain(x.iter().map(|v| v)).copied().collect();
11123        let mut gm = vec![0f32; 2 * rows];
11124        q1t_matmat(&bytes, &x2, 2, rows, cols, &mut gm, None);
11125        for r in 0..rows {
11126            assert!((gm[r] - expect[r]).abs() < tol(expect[r]));
11127            assert!((gm[rows + r] - expect[r]).abs() < tol(expect[r]));
11128        }
11129        // Fused pair (q1t_matvec2) must equal two single matvecs
11130        // bit-for-bit: same unpack, same group order, same f32
11131        // accumulation per stream. Distinct x2 exercises both lanes.
11132        let xb: Vec<f32> = (0..cols)
11133            .map(|j| if j % 5 == 0 { -1.0 } else { 1.0 })
11134            .collect();
11135        let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
11136        q1t_matvec(&bytes, &x, rows, cols, &mut s1, None);
11137        q1t_matvec(&bytes, &xb, rows, cols, &mut s2, None);
11138        let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
11139        q1t_matvec2(&bytes, &x, &xb, rows, cols, &mut p1, &mut p2, None);
11140        assert_eq!(p1, s1, "q1t pair lane 1 ≠ single matvec");
11141        assert_eq!(p2, s2, "q1t pair lane 2 ≠ single matvec");
11142    }
11143
11144    /// Pair == 2×matvec with an ODD group count (the kernel's tail
11145    /// group) and no overlay section.
11146    #[test]
11147    fn q1t_matvec2_odd_gpr_matches_singles() {
11148        use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_pack};
11149        let (rows, cols) = (5usize, 96usize); // gpr = 3 → paired + tail
11150        let gpr = cols / GROUP_SIZE;
11151        let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE);
11152        for r in 0..rows {
11153            for g in 0..gpr {
11154                bytes.extend_from_slice(&f32_to_f16(0.1 + 0.05 * (r + g) as f32).to_le_bytes());
11155                let mut c = [0u8; 7];
11156                for k in 0..GROUP_SIZE {
11157                    q1t_pack(&mut c, k, ((k * 7 + r * 5 + g * 3) % 3) as u8);
11158                }
11159                bytes.extend_from_slice(&c);
11160            }
11161        }
11162        let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.31).sin()).collect();
11163        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).cos()).collect();
11164        let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
11165        q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
11166        q1t_matvec(&bytes, &x2, rows, cols, &mut s2, None);
11167        let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
11168        q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
11169        assert_eq!(p1, s1, "odd-gpr pair lane 1 ≠ single");
11170        assert_eq!(p2, s2, "odd-gpr pair lane 2 ≠ single");
11171    }
11172
11173    // Speed A/B: fused pair (one unpack, two streams) vs two single
11174    // matvecs. Single-threaded, FFN-sized, min-of paired in-process.
11175    //   cargo test -p cortiq-engine --release q1t_matvec2_speed -- --ignored --nocapture
11176    #[test]
11177    #[ignore]
11178    fn q1t_matvec2_speed() {
11179        use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_pack};
11180        use std::time::Instant;
11181        let (rows, cols) = (8192usize, 4096usize);
11182        let gpr = cols / GROUP_SIZE;
11183        let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE);
11184        for r in 0..rows {
11185            for g in 0..gpr {
11186                let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
11187                bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
11188                let mut c = [0u8; 7];
11189                for k in 0..GROUP_SIZE {
11190                    q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
11191                }
11192                bytes.extend_from_slice(&c);
11193            }
11194        }
11195        let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.31).sin()).collect();
11196        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).cos()).collect();
11197        let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
11198        let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
11199        // Warm both paths once.
11200        q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
11201        q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
11202        let (mut t_pair, mut t_two) = (f64::MAX, f64::MAX);
11203        for _ in 0..8 {
11204            let t0 = Instant::now();
11205            q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
11206            t_pair = t_pair.min(t0.elapsed().as_secs_f64() * 1000.0);
11207            let t1 = Instant::now();
11208            q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
11209            q1t_matvec(&bytes, &x2, rows, cols, &mut s2, None);
11210            t_two = t_two.min(t1.elapsed().as_secs_f64() * 1000.0);
11211        }
11212        assert_eq!(p1, s1);
11213        assert_eq!(p2, s2);
11214        println!("q1t pair {rows}x{cols}: fused {t_pair:.2} ms | two singles {t_two:.2} ms");
11215    }
11216
11217    // Speed A/B: the base-3-division decode (what the packing commit left in
11218    // place) vs the fused sign-LUT matvec. Both single-threaded, same bytes.
11219    //   cargo test -p cortiq-engine q1t_matvec_speed -- --ignored --nocapture
11220    #[test]
11221    #[ignore]
11222    fn q1t_matvec_speed() {
11223        use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_code, q1t_pack};
11224        use std::time::Instant;
11225        let (rows, cols) = (8192usize, 4096usize); // FFN-sized
11226        let gpr = cols / GROUP_SIZE;
11227        let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE + 16);
11228        for r in 0..rows {
11229            for g in 0..gpr {
11230                let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
11231                bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
11232                let mut c = [0u8; 7];
11233                for k in 0..GROUP_SIZE {
11234                    q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
11235                }
11236                bytes.extend_from_slice(&c);
11237            }
11238        }
11239        let (n, stride) = (rows * cols, 40usize); // ~2.5% outliers, per-row overlay
11240        let mut row_ptr = vec![0u32; rows + 1];
11241        let mut idx = 0usize;
11242        while idx < n {
11243            row_ptr[idx / cols + 1] += 1;
11244            idx += stride;
11245        }
11246        for r in 0..rows {
11247            row_ptr[r + 1] += row_ptr[r];
11248        }
11249        for &p in &row_ptr {
11250            bytes.extend_from_slice(&p.to_le_bytes());
11251        }
11252        let mut idx = 0usize;
11253        while idx < n {
11254            bytes.extend_from_slice(&((idx % cols) as u16).to_le_bytes());
11255            bytes.extend_from_slice(&f32_to_f16((idx % 13) as f32 * 0.1 - 0.6).to_le_bytes());
11256            idx += stride;
11257        }
11258        // On-grid ±1 so the fast path's int8 SDOT is exact vs the f32 "slow"
11259        // reference (the A/B is a timing check; values must still agree).
11260        let x: Vec<f32> = (0..cols)
11261            .map(|j| if j % 3 == 0 { 1.0 } else { -1.0 })
11262            .collect();
11263        let (rp_off, ent_off, has_ov) = q1t_overlay(&bytes, rows * gpr * Q1T_TILE, rows);
11264
11265        // "before": base-3 division decode into a buffer, then dot.
11266        let slow = |out: &mut [f32]| {
11267            let mut buf = vec![0f32; cols];
11268            for r in 0..rows {
11269                for g in 0..gpr {
11270                    let off = (r * gpr + g) * Q1T_TILE;
11271                    let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
11272                    let codes = &bytes[off + 2..off + Q1T_TILE];
11273                    for k in 0..GROUP_SIZE {
11274                        buf[g * GROUP_SIZE + k] = match q1t_code(codes, k) {
11275                            1 => s,
11276                            2 => -s,
11277                            _ => 0.0,
11278                        };
11279                    }
11280                }
11281                out[r] = q1t_row_outlier_correction(&bytes, r, rp_off, ent_off, has_ov, &x)
11282                    + (0..cols).map(|j| buf[j] * x[j]).sum::<f32>();
11283            }
11284        };
11285        let iters = 5;
11286        let mut a = vec![0f32; rows];
11287        slow(&mut a); // warm
11288        let t = Instant::now();
11289        for _ in 0..iters {
11290            slow(&mut a);
11291        }
11292        let slow_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
11293
11294        let mut b = vec![0f32; rows];
11295        q1t_matvec(&bytes, &x, rows, cols, &mut b, None); // warm
11296        let t = Instant::now();
11297        for _ in 0..iters {
11298            q1t_matvec(&bytes, &x, rows, cols, &mut b, None);
11299        }
11300        let fast_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
11301
11302        for r in 0..rows {
11303            assert!((a[r] - b[r]).abs() < 1e-2, "mismatch row {r}");
11304        }
11305        println!(
11306            "q1t matvec {rows}x{cols} (1 thread): div-decode {slow_ms:.2} ms  fused-LUT {fast_ms:.2} ms  => {:.2}x",
11307            slow_ms / fast_ms
11308        );
11309    }
11310}
11311
11312
11313#[cfg(test)]
11314mod gemm_bench {
11315    /// `cargo test -p cortiq-engine --release q4tp_matmat_throughput -- --ignored --nocapture`
11316    /// Times the batched q4tp GEMM at the shapes the image DiT runs
11317    /// (b=296 tokens, 2304 -> 9216), on synthetic bytes: no model, no
11318    /// mmap, no thermal drift over minutes — a kernel change shows up
11319    /// here in seconds where a full render hides it in noise.
11320    ///
11321    /// On macOS add `CMF_ACCEL=0`: this shape is over the 500k-cell mark
11322    /// where the matmat hands off to Accelerate's dequant sgemm, and
11323    /// without the opt-out both rows below measure the AMX, not the
11324    /// kernel under test.
11325    #[test]
11326    #[ignore]
11327    fn q4tp_matmat_throughput() {
11328        // 296 is a prompt-encode batch; the image DiT runs 2085 at
11329        // 512x512, where the activation panel stops fitting L2 and the
11330        // loop's shape starts to matter more than its instructions.
11331        let b: usize = std::env::var("CMF_BENCH_B")
11332            .ok()
11333            .and_then(|v| v.parse().ok())
11334            .unwrap_or(296);
11335        let (rows, cols) = (9216usize, 2304usize);
11336        let (_, _, _) = (rows, cols, b);
11337        let total = cortiq_core::quant::expected_nbytes(
11338            cortiq_core::TensorDtype::Q4TiledP,
11339            &[rows, cols],
11340        )
11341        .unwrap();
11342        // Random nibbles are fine, but the row params are f16 (lo, step)
11343        // of a geometric ladder: garbage there gives exp2 of a huge
11344        // exponent, the scales come back inf, and the whole bench times
11345        // NaN arithmetic instead of the kernel.
11346        let (params_off, codes_off, _) =
11347            cortiq_core::quant::q4tp_sections(rows, cols);
11348        let mut bytes: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
11349        let lo = cortiq_core::quant::f32_to_f16(-4.0);
11350        let step = cortiq_core::quant::f32_to_f16(0.1);
11351        for r in 0..rows {
11352            let o = params_off + r * 4;
11353            bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
11354            bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
11355        }
11356        let _ = codes_off;
11357        let xs: Vec<f32> = (0..b * cols)
11358            .map(|i| ((i % 97) as f32 - 48.0) / 48.0)
11359            .collect();
11360        let mut out = vec![0f32; b * rows];
11361        let pool = crate::pool::Pool::from_env();
11362        // A shared 48-core stand drifts ±25% run to run, which is wider
11363        // than any kernel change worth making. So: alternate the two
11364        // kernels inside one process and keep the BEST time for
11365        // each. Interleaving makes both see the same interference, and a
11366        // minimum is the one statistic another tenant cannot inflate.
11367        super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
11368        let reps: usize = std::env::var("CMF_BENCH_REPS")
11369            .ok()
11370            .and_then(|v| v.parse().ok())
11371            .unwrap_or(10);
11372        let mut best = [f64::MAX; 2];
11373        let mut sums = [0f32; 2];
11374        for _ in 0..reps {
11375            for (k, w) in [(0usize, 1u8), (1usize, 2u8)] {
11376                super::Q4TP_ALT.store(w, std::sync::atomic::Ordering::Relaxed);
11377                let t = std::time::Instant::now();
11378                super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
11379                best[k] = best[k].min(t.elapsed().as_secs_f64());
11380                sums[k] = out.iter().take(64).sum::<f32>();
11381            }
11382        }
11383        let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
11384        for (k, name) in ["previous", "tuned   "].iter().enumerate() {
11385            println!(
11386                "q4tp matmat {rows}x{cols} b={b} {name}: {:.1} ms  {:.1} GFLOP/s  (checksum {:.3})",
11387                best[k] * 1e3,
11388                flops / best[k] / 1e9,
11389                sums[k]
11390            );
11391        }
11392        assert!(
11393            (sums[0] - sums[1]).abs() < 1e-2,
11394            "the tuned kernel changed the result: {} vs {}",
11395            sums[0],
11396            sums[1]
11397        );
11398    }
11399
11400    /// The blocked kernel must agree with the per-column path exactly —
11401    /// same weights, same activation split, only a different instruction
11402    /// mix. Shapes are chosen to hit the awkward cases: a column count
11403    /// that leaves an odd group (the 512-bit kernel does two at a time),
11404    /// and a batch that does not divide by four.
11405    #[test]
11406    fn q4tp_matmat_blocked_matches_scalar() {
11407        use std::sync::atomic::Ordering::Relaxed;
11408        // The last shape carries the image DiT's column count — 2304, so
11409        // 72 groups of accumulation, which is where a reordered sum can
11410        // actually drift — and runs through the thread pool, since the
11411        // blocked path splits rows across workers. Its row count stays
11412        // under 500k cells on purpose: above that, macOS diverts the whole
11413        // matmat to the Accelerate/AMX dequant sgemm and neither kernel
11414        // here would run.
11415        for &(rows, cols, b) in &[
11416            (64usize, 128usize, 7usize),
11417            (33, 96, 4),
11418            (16, 256, 9),
11419            (192, 2304, 37),
11420        ] {
11421            let total = cortiq_core::quant::expected_nbytes(
11422                cortiq_core::TensorDtype::Q4TiledP,
11423                &[rows, cols],
11424            )
11425            .unwrap();
11426            let (params_off, _, _) = cortiq_core::quant::q4tp_sections(rows, cols);
11427            let mut bytes: Vec<u8> = (0..total).map(|i| (i * 61 % 251) as u8).collect();
11428            let lo = cortiq_core::quant::f32_to_f16(-4.0);
11429            let step = cortiq_core::quant::f32_to_f16(0.1);
11430            for r in 0..rows {
11431                let o = params_off + r * 4;
11432                bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
11433                bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
11434            }
11435            let xs: Vec<f32> = (0..b * cols)
11436                .map(|i| ((i % 89) as f32 - 44.0) / 44.0)
11437                .collect();
11438            let mut got = vec![0f32; b * rows];
11439            let mut want = vec![0f32; b * rows];
11440            let gpr = cols / 32;
11441            let view = super::Q4tpView::new(&bytes, rows, cols);
11442            let pool = crate::pool::Pool::from_env();
11443            super::Q4TP_ALT.store(2, Relaxed);
11444            super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut got, pool.as_deref());
11445            super::Q4TP_ALT.store(1, Relaxed);
11446            super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut want, pool.as_deref());
11447            super::Q4TP_ALT.store(0, Relaxed);
11448            // Measured against the output's scale, not cell by cell: a
11449            // dot product of 2304 terms lands near zero wherever the row
11450            // and the activation nearly cancel, and there a per-cell
11451            // ratio reports 1e-3 for an absolute error of 5e-6 — f32's
11452            // own rounding, reordered. What must stay small is the error
11453            // relative to what the layer actually outputs.
11454            let scale = want.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
11455            let (mut worst, mut at) = (0f32, 0usize);
11456            for (i, (g, w)) in got.iter().zip(&want).enumerate() {
11457                if (g - w).abs() > worst {
11458                    worst = (g - w).abs();
11459                    at = i;
11460                }
11461            }
11462            assert!(
11463                worst <= 1e-4 * scale,
11464                "{rows}x{cols} b={b}: blocked and scalar disagree by {worst:.3e} \
11465                 (scale {scale:.3e}) at cell {at}: {} vs {}",
11466                got[at],
11467                want[at]
11468            );
11469
11470            // "Same speed, no quality loss" is a claim about which answer
11471            // is RIGHT, not about which two agree. Both paths sum the same
11472            // 2304 products in different orders, so f64 decides: the
11473            // blocked kernel keeps sixteen partial sums and folds them at
11474            // the end, which is a shallower addition tree than the
11475            // per-column path's running scalar, and it must not be worse.
11476            let (mut e_blocked, mut e_scalar) = (0f64, 0f64);
11477            for bi in 0..b {
11478                let act = super::split_act(&xs[bi * cols..(bi + 1) * cols]);
11479                for r in 0..rows {
11480                    let mut sc = vec![0f32; gpr];
11481                    view.scales_into(r, gpr, &mut sc);
11482                    let mut exact = 0f64;
11483                    for j in 0..cols {
11484                        let (w, sq) = super::q4tp_outlier(view.nib, r, gpr, j, &sc);
11485                        exact += w as f64 * sq as f64 * act.xq[j] as f64;
11486                    }
11487                    exact *= act.sx as f64;
11488                    for &(j, xv) in &act.outliers {
11489                        let (w, sq) = super::q4tp_outlier(view.nib, r, gpr, j, &sc);
11490                        exact += w as f64 * sq as f64 * xv as f64;
11491                    }
11492                    let i = bi * rows + r;
11493                    e_blocked = e_blocked.max((got[i] as f64 - exact).abs());
11494                    e_scalar = e_scalar.max((want[i] as f64 - exact).abs());
11495                }
11496            }
11497            println!(
11498                "{rows}x{cols} b={b}: worst error vs f64 — blocked {e_blocked:.3e}, \
11499                 per-column {e_scalar:.3e}"
11500            );
11501            // An absolute bar, not a race between the two: at these
11502            // magnitudes both sit in f32's last bits, and on a small shape
11503            // whichever one happens to round the unluckiest cell "wins" by
11504            // a factor the next seed reverses.
11505            assert!(
11506                e_blocked <= 1e-5 * scale as f64 && e_scalar <= 1e-5 * scale as f64,
11507                "{rows}x{cols} b={b}: error against f64 too large — blocked \
11508                 {e_blocked:.3e}, per-column {e_scalar:.3e}, scale {scale:.3e}"
11509            );
11510        }
11511    }
11512
11513    /// The q4t twin of the throughput bench, same shape and rules, so the
11514    /// two quantisations' batch kernels can be read against each other.
11515    /// `cargo test -p cortiq-engine --release q4t_matmat_throughput -- --ignored --nocapture`
11516    #[test]
11517    #[ignore]
11518    fn q4t_matmat_throughput() {
11519        let (rows, cols, b) = (9216usize, 2304usize, 296usize);
11520        let total = cortiq_core::quant::expected_nbytes(
11521            cortiq_core::TensorDtype::Q4Tiled,
11522            &[rows, cols],
11523        )
11524        .unwrap();
11525        // q4t carries a per-group f16 scale in the tile's first two bytes;
11526        // random bytes there decode to inf and the bench would time NaNs.
11527        let mut bytes: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
11528        let sc = cortiq_core::quant::f32_to_f16(0.02);
11529        for t in bytes.chunks_mut(super::Q4_TILE) {
11530            t[..2].copy_from_slice(&sc.to_le_bytes());
11531        }
11532        let xs: Vec<f32> = (0..b * cols)
11533            .map(|i| ((i % 97) as f32 - 48.0) / 48.0)
11534            .collect();
11535        let mut out = vec![0f32; b * rows];
11536        let pool = crate::pool::Pool::from_env();
11537        super::q4t_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
11538        let reps: usize = std::env::var("CMF_BENCH_REPS")
11539            .ok()
11540            .and_then(|v| v.parse().ok())
11541            .unwrap_or(10);
11542        let mut best = f64::MAX;
11543        for _ in 0..reps {
11544            let t = std::time::Instant::now();
11545            super::q4t_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
11546            best = best.min(t.elapsed().as_secs_f64());
11547        }
11548        let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
11549        println!(
11550            "q4t matmat {rows}x{cols} b={b}: {:.1} ms  {:.1} GFLOP/s  (checksum {:.3})",
11551            best * 1e3,
11552            flops / best / 1e9,
11553            out.iter().take(64).sum::<f32>()
11554        );
11555    }
11556
11557}