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.
554    /// kind: 0=q8_row (per-row scales), 1=q1, 2=q4_block, 3=q1t
555    /// (tile-embedded, no rs), 5=q4_tiled, 6=q4tp, 7=q8_2f (both scale
556    /// planes live inside the tensor). None only for `vbit`.
557    ///
558    /// The old comment here claimed q4_block was unhandled while the arm
559    /// right below mapped it, and it named q8_2f as unhandled after that
560    /// stopped being true — a stale comment on this function is how a
561    /// model silently loses the graph, so it is worth keeping honest.
562    pub fn graph_weight(&self) -> Option<(&std::sync::Arc<CmfModel>, usize, u8, &[f32])> {
563        match self {
564            Self::Mapped {
565                model,
566                idx,
567                dtype: TensorDtype::Q8Row,
568                row_scale,
569                ..
570            } => Some((model, *idx, 0, row_scale.as_slice())),
571            Self::Mapped {
572                model,
573                idx,
574                dtype: TensorDtype::Q1,
575                ..
576            } => Some((model, *idx, 1, &[])),
577            // Q4Tiled is kind 5, NOT 2: both carried 2 historically, and
578            // the wgpu token graph fed 18B interleaved tiles to the
579            // split-layout q4b kernel — garbage output on q4t models
580            // (caught by an end-to-end answer check on real Vulkan).
581            Self::Mapped {
582                model,
583                idx,
584                dtype: TensorDtype::Q4Tiled,
585                ..
586            } => Some((model, *idx, 5, &[])),
587            // Kind 6, not 5: q4tp's nibble stride and scale planes differ,
588            // and feeding them to the q4t kernel is exactly the mistake that
589            // produced garbage when Q4Tiled shared kind 2 with Q4Block.
590            Self::Mapped {
591                model,
592                idx,
593                dtype: TensorDtype::Q4TiledP,
594                ..
595            } => Some((model, *idx, 6, &[])),
596            Self::Mapped {
597                model,
598                idx,
599                dtype: TensorDtype::Q4Block,
600                ..
601            } => Some((model, *idx, 2, &[])),
602            // q8_2f carries BOTH scale planes after the int8 body (rows
603            // f16, then cols f16), so the graph takes the whole tensor
604            // and the kernel reads them where they lie — no host-side
605            // prescale, which is what the per-op path does instead.
606            Self::Mapped {
607                model,
608                idx,
609                dtype: TensorDtype::Q8_2f,
610                ..
611            } => Some((model, *idx, 7, &[])),
612            Self::Mapped {
613                model,
614                idx,
615                dtype: TensorDtype::Q1T,
616                ..
617            } => Some((model, *idx, 3, &[])),
618            _ => None,
619        }
620    }
621
622    /// Dense f32 view — only for owned tensors. Masked/sparse execution
623    /// paths require it; quantized weights don't support masks yet.
624    pub fn as_f32(&self) -> Option<&[f32]> {
625        match self {
626            Self::F32 { data, .. } => Some(data),
627            Self::Mapped { .. } => None,
628        }
629    }
630
631    fn quant_bytes(&self) -> &[u8] {
632        match self {
633            Self::Mapped { model, idx, .. } => model.entry_bytes(&model.tensors[*idx]),
634            Self::F32 { .. } => unreachable!("quant_bytes on F32"),
635        }
636    }
637
638    /// Dequantize one row into `dst` (embedding lookup).
639    pub fn row_f32(&self, r: usize, dst: &mut [f32]) {
640        let cols = self.cols();
641        debug_assert_eq!(dst.len(), cols);
642        match self {
643            Self::F32 { data, .. } => dst.copy_from_slice(&data[r * cols..(r + 1) * cols]),
644            Self::Mapped {
645                dtype,
646                row_scale,
647                col_field,
648                vbit_offsets,
649                ..
650            } => {
651                if *dtype == TensorDtype::Q4Tiled {
652                    let bytes = self.quant_bytes();
653                    let gpr = cols / GROUP_SIZE;
654                    for gi in 0..gpr {
655                        let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
656                        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
657                        for (k, &b) in tile[2..].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::Q4TiledP {
665                    let bytes = self.quant_bytes();
666                    let gpr = cols / GROUP_SIZE;
667                    let v = Q4tpView::new(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 tile = &v.nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
672                        let s = sc[gi];
673                        for (k, &b) in tile.iter().enumerate() {
674                            dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
675                            dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
676                        }
677                    }
678                    return;
679                }
680                if *dtype == TensorDtype::Q2TiledP {
681                    let bytes = self.quant_bytes();
682                    let gpr = cols / GROUP_SIZE;
683                    let v = Q4tpView::new_q2(bytes, self.rows(), cols);
684                    let mut sc = vec![0f32; gpr];
685                    v.scales_into(r, gpr, &mut sc);
686                    for gi in 0..gpr {
687                        let ch =
688                            &v.nib[(r * gpr + gi) * Q2TP_CHUNK..(r * gpr + gi + 1) * Q2TP_CHUNK];
689                        let s = sc[gi];
690                        for (k, &b) in ch.iter().enumerate() {
691                            for j in 0..4 {
692                                dst[gi * GROUP_SIZE + k * 4 + j] =
693                                    (((b >> (2 * j)) & 3) as f32 - 1.5) * s;
694                            }
695                        }
696                    }
697                    return;
698                }
699                if *dtype == TensorDtype::Q4Block {
700                    let (packed, scales) = q4_split(self.quant_bytes(), self.rows(), cols);
701                    let gpr = cols / GROUP_SIZE;
702                    for gi in 0..gpr {
703                        let g = r * gpr + gi;
704                        let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
705                        for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
706                            dst[gi * GROUP_SIZE + k * 2] = ((b & 0x0F) as f32 - 8.0) * s;
707                            dst[gi * GROUP_SIZE + k * 2 + 1] = (((b >> 4) & 0x0F) as f32 - 8.0) * s;
708                        }
709                    }
710                    return;
711                }
712                if *dtype == TensorDtype::Q1 {
713                    let bytes = self.quant_bytes();
714                    let gpr = cols / GROUP_SIZE;
715                    for gi in 0..gpr {
716                        let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
717                        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
718                        for (j, &b) in tile[2..].iter().enumerate() {
719                            for k in 0..8 {
720                                dst[gi * GROUP_SIZE + j * 8 + k] =
721                                    (((b >> k) & 1) as f32 * 2.0 - 1.0) * s;
722                            }
723                        }
724                    }
725                    return;
726                }
727                if *dtype == TensorDtype::Q1T {
728                    let bytes = self.quant_bytes();
729                    let gpr = cols / GROUP_SIZE;
730                    let base_len = self.rows() * gpr * cortiq_core::quant::Q1T_TILE;
731                    for gi in 0..gpr {
732                        let off = (r * gpr + gi) * cortiq_core::quant::Q1T_TILE;
733                        let s = cortiq_core::quant::f16_to_f32(u16::from_le_bytes([
734                            bytes[off],
735                            bytes[off + 1],
736                        ]));
737                        let codes = &bytes[off + 2..off + cortiq_core::quant::Q1T_TILE];
738                        for k in 0..GROUP_SIZE {
739                            dst[gi * GROUP_SIZE + k] = match cortiq_core::quant::q1t_code(codes, k)
740                            {
741                                1 => s,
742                                2 => -s,
743                                _ => 0.0,
744                            };
745                        }
746                    }
747                    // Overlay
748                    let rows = self.rows();
749                    let entries = base_len + (rows + 1) * 4;
750                    if entries <= bytes.len() {
751                        let ptrs = &bytes[base_len..base_len + (rows + 1) * 4];
752                        let r0 = u32::from_le_bytes([
753                            ptrs[r * 4],
754                            ptrs[r * 4 + 1],
755                            ptrs[r * 4 + 2],
756                            ptrs[r * 4 + 3],
757                        ]) as usize;
758                        let r1 = u32::from_le_bytes([
759                            ptrs[(r + 1) * 4],
760                            ptrs[(r + 1) * 4 + 1],
761                            ptrs[(r + 1) * 4 + 2],
762                            ptrs[(r + 1) * 4 + 3],
763                        ]) as usize;
764                        let off = entries + r0 * 4;
765                        for i in 0..r1 - r0 {
766                            let item = &bytes[off + i * 4..off + i * 4 + 4];
767                            let c = u16::from_le_bytes([item[0], item[1]]) as usize;
768                            let v = cortiq_core::quant::f16_to_f32(u16::from_le_bytes([
769                                item[2], item[3],
770                            ]));
771                            if c < cols {
772                                dst[c] = v;
773                            }
774                        }
775                    }
776                    return;
777                }
778                if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
779                    let bytes = self.quant_bytes();
780                    let rows = self.rows();
781                    let ng = cols / GROUP_SIZE;
782                    let bits = &bytes[..rows];
783                    let sc_off = rows;
784                    // Precomputed at load — embedding lookup used to scan
785                    // the bit-widths of every preceding row (O(token_id)).
786                    let off = vbit_offsets[r];
787                    let b = bits[r] as usize;
788                    let l = ((1usize << (b - 1)) - 1) as f32;
789                    let data = &bytes[off..];
790                    let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
791                    for (i, d) in dst.iter_mut().enumerate() {
792                        while nbits < b {
793                            acc = (acc << 8) | data[idx] as u64;
794                            idx += 1;
795                            nbits += 8;
796                        }
797                        let u = ((acc >> (nbits - b)) & ((1u64 << b) - 1)) as f32;
798                        nbits -= b;
799                        let so = (r * ng + i / GROUP_SIZE) * 2;
800                        let sv = f16_to_f32(u16::from_le_bytes([
801                            bytes[sc_off + so],
802                            bytes[sc_off + so + 1],
803                        ]));
804                        *d = (u - l) * sv;
805                    }
806                    return;
807                }
808                let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
809                let s = row_scale[r];
810                match dtype {
811                    TensorDtype::Q8Row => {
812                        for (d, &b) in dst.iter_mut().zip(q) {
813                            *d = (b as i8) as f32 * s;
814                        }
815                    }
816                    TensorDtype::Q8_2f => {
817                        for (i, (d, &b)) in dst.iter_mut().zip(q).enumerate() {
818                            *d = (b as i8) as f32 * s * col_field[i];
819                        }
820                    }
821                    _ => unreachable!(),
822                }
823            }
824        }
825    }
826
827    /// Can this tensor's columns be read cheaply (for sparse down_proj)?
828    /// True for F32/Q8Row/Q8_2f (per-row scale, direct strided access);
829    /// false for group-packed q4/vbit (column access would unpack whole
830    /// groups — sparse execution falls back to f32 for those).
831    pub fn sparse_col_ok(&self) -> bool {
832        match self {
833            Self::F32 { .. } => true,
834            Self::Mapped { dtype, .. } => {
835                matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
836            }
837        }
838    }
839
840    /// down_proj [hidden, inter]: accumulate `w · col(c)` into `out`
841    /// [hidden] — reads ONLY column `c` (one neuron) from the mmap,
842    /// no full-matrix dequant. `out[k] += w · down[k, c]`.
843    pub fn add_col_scaled(&self, c: usize, w: f32, out: &mut [f32]) {
844        let inter = self.cols();
845        let hidden = self.rows();
846        debug_assert_eq!(out.len(), hidden);
847        match self {
848            Self::F32 { data, .. } => {
849                for (k, o) in out.iter_mut().enumerate() {
850                    *o += w * data[k * inter + c];
851                }
852            }
853            Self::Mapped {
854                dtype,
855                row_scale,
856                col_field,
857                ..
858            } => {
859                let q = self.quant_bytes();
860                let colf = if *dtype == TensorDtype::Q8_2f {
861                    col_field[c]
862                } else {
863                    1.0
864                };
865                let wc = w * colf;
866                for (k, o) in out.iter_mut().enumerate() {
867                    let b = q[k * inter + c] as i8 as f32;
868                    *o += wc * b * row_scale[k];
869                }
870            }
871        }
872    }
873
874    /// Dot of row `r` with `x` (gate/up active-neuron path). Reads only
875    /// row `r` from the mmap — no full dequant. q4/vbit dequant the row
876    /// into `scratch` first (rare for active-FFN weights).
877    pub fn row_dot(&self, r: usize, x: &[f32], scratch: &mut [f32]) -> f32 {
878        let cols = self.cols();
879        match self {
880            Self::F32 { data, .. } => {
881                let row = &data[r * cols..(r + 1) * cols];
882                row.iter().zip(x).map(|(w, v)| w * v).sum()
883            }
884            Self::Mapped {
885                dtype,
886                row_scale,
887                col_field,
888                ..
889            } => match dtype {
890                TensorDtype::Q8Row => {
891                    let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
892                    dot_i8_f32(q, x) * row_scale[r]
893                }
894                TensorDtype::Q8_2f => {
895                    let q = &self.quant_bytes()[r * cols..(r + 1) * cols];
896                    dot_i8_col_f32(q, x, col_field) * row_scale[r]
897                }
898                _ => {
899                    self.row_f32(r, scratch);
900                    scratch.iter().zip(x).map(|(w, v)| w * v).sum()
901                }
902            },
903        }
904    }
905
906    /// `out = W · x` (row-major). F32 delegates to the historical
907    /// bit-exact path; Mapped runs the fused int8 kernel.
908    pub fn matvec(&self, x: &[f32], out: &mut [f32], pool: Option<&Pool>) {
909        match self {
910            // NOTE: `out.len()` DRIVES this arm — it computes that many rows,
911            // and `x.len()` is the stride. A short `out` is legitimate here,
912            // which is why the check below lives in the Mapped arm only.
913            Self::F32 { data, .. } => matvec_rows(pool, data, x, out),
914            Self::Mapped {
915                model,
916                idx,
917                dtype,
918                rows,
919                cols,
920                row_scale,
921                col_field,
922                vbit_offsets,
923                repack,
924            } => {
925                let _ = (model, idx);
926                // Every kernel below writes `rows` entries through a raw
927                // pointer, so a short `out` is an out-of-bounds WRITE, not a
928                // wrong answer: it scribbles on the allocator's metadata and
929                // the process aborts much later, somewhere innocent
930                // (`double free or corruption`, `corrupted double-linked
931                // list`). The debug_assert two of the kernels carried is
932                // compiled out of the release — exactly the build where it
933                // matters. Fail here instead, while the caller is still on
934                // the stack to be named.
935                assert!(
936                    out.len() >= *rows && x.len() >= *cols,
937                    "matvec {rows}x{cols}: out {} (need {rows}), x {} (need {cols})",
938                    out.len(),
939                    x.len(),
940                );
941                if *dtype == TensorDtype::Q4Block {
942                    // GPU route (wgpu q4b kernel) for large q4_block matvecs —
943                    // gives NVIDIA/AMD/Intel q4 models a GPU path. Probe keeps
944                    // the winner; Metal returns false → the CPU kernel below.
945                    if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
946                        let t0 = std::time::Instant::now();
947                        match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
948                            crate::gpu::ProbeArm::Gpu => {
949                                if crate::gpu::q4b_matvec(model, *idx, x, *rows, *cols, out) {
950                                    crate::gpu::probe_record(
951                                        crate::gpu::OpClass::Matvec,
952                                        true,
953                                        t0.elapsed(),
954                                    );
955                                    return;
956                                }
957                            }
958                            crate::gpu::ProbeArm::CpuTimed => {
959                                q4matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
960                                crate::gpu::probe_record(
961                                    crate::gpu::OpClass::Matvec,
962                                    false,
963                                    t0.elapsed(),
964                                );
965                                return;
966                            }
967                            crate::gpu::ProbeArm::Cpu => {}
968                        }
969                    }
970                    q4matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
971                    return;
972                }
973                if *dtype == TensorDtype::Q4Tiled {
974                    // GPU route for large q4t matvecs — the lm_head class,
975                    // same shape as the q4tp arm below. The probe keeps the
976                    // winner; a backend without the kernel refuses and the
977                    // CPU path stays.
978                    if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
979                        let t0 = std::time::Instant::now();
980                        let cls = crate::gpu::matvec_class(*rows, *cols);
981                        match crate::gpu::probe_arm(cls) {
982                            crate::gpu::ProbeArm::Gpu => {
983                                if crate::gpu::q4t_matvec(model, *idx, x, *rows, *cols, out) {
984                                    crate::gpu::probe_record(cls, true, t0.elapsed());
985                                    return;
986                                }
987                            }
988                            crate::gpu::ProbeArm::CpuTimed => {
989                                q4t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
990                                crate::gpu::probe_record(cls, false, t0.elapsed());
991                                return;
992                            }
993                            crate::gpu::ProbeArm::Cpu => {}
994                        }
995                    }
996                    q4t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
997                    return;
998                }
999                if *dtype == TensorDtype::Q4TiledP {
1000                    // GPU route for large q4tp matvecs — the lm_head class.
1001                    // On a q4tp checkpoint the head is the biggest single
1002                    // host matvec left in the decode step, and the batched
1003                    // kernel at b=1 already exists on both backends. Probe
1004                    // keeps the winner, same as q4_block above.
1005                    if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1006                        let t0 = std::time::Instant::now();
1007                        let cls = crate::gpu::matvec_class(*rows, *cols);
1008                        match crate::gpu::probe_arm(cls) {
1009                            crate::gpu::ProbeArm::Gpu => {
1010                                if crate::gpu::q4tp_matvec(model, *idx, x, *rows, *cols, out) {
1011                                    crate::gpu::probe_record(cls, true, t0.elapsed());
1012                                    return;
1013                                }
1014                            }
1015                            crate::gpu::ProbeArm::CpuTimed => {
1016                                q4tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1017                                crate::gpu::probe_record(cls, false, t0.elapsed());
1018                                return;
1019                            }
1020                            crate::gpu::ProbeArm::Cpu => {}
1021                        }
1022                    }
1023                    q4tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1024                    return;
1025                }
1026                if *dtype == TensorDtype::Q2TiledP {
1027                    q2tp_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1028                    return;
1029                }
1030                if *dtype == TensorDtype::Q1 {
1031                    // GPU route for large q1 matvecs (out_proj / lm_head
1032                    // class): the CPU q1 kernel is load-port-bound at
1033                    // ~4 GB/s/core, the GPU one is bandwidth-bound — the
1034                    // probe measures both arms and keeps the winner.
1035                    if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1036                        let t0 = std::time::Instant::now();
1037                        let arm = if crate::gpu::q1_force() {
1038                            crate::gpu::ProbeArm::Gpu
1039                        } else {
1040                            crate::gpu::probe_arm(crate::gpu::OpClass::Matvec)
1041                        };
1042                        match arm {
1043                            crate::gpu::ProbeArm::Gpu => {
1044                                if crate::gpu::q1_matvec(model, *idx, x, *rows, *cols, out) {
1045                                    crate::gpu::probe_record(
1046                                        crate::gpu::OpClass::Matvec,
1047                                        true,
1048                                        t0.elapsed(),
1049                                    );
1050                                    return;
1051                                }
1052                            }
1053                            crate::gpu::ProbeArm::CpuTimed => {
1054                                q1_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1055                                crate::gpu::probe_record(
1056                                    crate::gpu::OpClass::Matvec,
1057                                    false,
1058                                    t0.elapsed(),
1059                                );
1060                                return;
1061                            }
1062                            crate::gpu::ProbeArm::Cpu => {}
1063                        }
1064                    }
1065                    q1_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1066                    return;
1067                }
1068                if *dtype == TensorDtype::Q1T {
1069                    // GPU route for large q1t matvecs: the ternary BASE dot runs
1070                    // on the GPU (load-port-bound on CPU, like q1), then the
1071                    // sparse overlay is added on the CPU. Probe keeps the winner.
1072                    if *rows * *cols >= 8_388_608 && crate::gpu::enabled_here() {
1073                        let t0 = std::time::Instant::now();
1074                        match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
1075                            crate::gpu::ProbeArm::Gpu => {
1076                                if crate::gpu::q1t_matvec(model, *idx, x, *rows, *cols, out) {
1077                                    q1t_add_overlay(self.quant_bytes(), x, *rows, *cols, out, pool);
1078                                    crate::gpu::probe_record(
1079                                        crate::gpu::OpClass::Matvec,
1080                                        true,
1081                                        t0.elapsed(),
1082                                    );
1083                                    return;
1084                                }
1085                            }
1086                            crate::gpu::ProbeArm::CpuTimed => {
1087                                q1t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1088                                crate::gpu::probe_record(
1089                                    crate::gpu::OpClass::Matvec,
1090                                    false,
1091                                    t0.elapsed(),
1092                                );
1093                                return;
1094                            }
1095                            crate::gpu::ProbeArm::Cpu => {}
1096                        }
1097                    }
1098                    q1t_matvec(self.quant_bytes(), x, *rows, *cols, out, pool);
1099                    return;
1100                }
1101                if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1102                    vbitmatvec(self.quant_bytes(), vbit_offsets, x, *rows, *cols, out, pool);
1103                    return;
1104                }
1105                let xs = prescale(x, col_field, *dtype);
1106                // D5: large q8 matrices (lm_head-class) — hybrid
1107                // CPU∥GPU: split the rows, both sides compute
1108                // SIMULTANEOUSLY (same math, shared prescale).
1109                // GPU share: CMF_GPU_SPLIT (0..1, default 0.5).
1110                if *rows >= crate::gpu::min_rows()
1111                    && matches!(dtype, TensorDtype::Q8Row | TensorDtype::Q8_2f)
1112                    && gpu_lmhead_enabled()
1113                    && crate::gpu::enabled_here()
1114                {
1115                    // Runtime probe: alternate the hybrid against the
1116                    // pure-CPU matvec, keep whichever is faster HERE.
1117                    let t0 = std::time::Instant::now();
1118                    match crate::gpu::probe_arm(crate::gpu::OpClass::Matvec) {
1119                        crate::gpu::ProbeArm::Gpu => {}
1120                        crate::gpu::ProbeArm::CpuTimed => {
1121                            qmatvec(
1122                                self.quant_bytes(),
1123                                repack,
1124                                row_scale,
1125                                x,
1126                                col_field,
1127                                *dtype,
1128                                *rows,
1129                                *cols,
1130                                out,
1131                                pool,
1132                            );
1133                            crate::gpu::probe_record(
1134                                crate::gpu::OpClass::Matvec,
1135                                false,
1136                                t0.elapsed(),
1137                            );
1138                            return;
1139                        }
1140                        crate::gpu::ProbeArm::Cpu => {
1141                            qmatvec(
1142                                self.quant_bytes(),
1143                                repack,
1144                                row_scale,
1145                                x,
1146                                col_field,
1147                                *dtype,
1148                                *rows,
1149                                *cols,
1150                                out,
1151                                pool,
1152                            );
1153                            return;
1154                        }
1155                    }
1156                    let frac = gpu_split_frac();
1157                    let cpu_rows = ((*rows as f32) * (1.0 - frac)) as usize;
1158                    let (out_cpu, out_gpu) = out.split_at_mut(cpu_rows);
1159                    let bytes = self.quant_bytes();
1160                    let ok = std::thread::scope(|sc| {
1161                        let g = sc.spawn(|| {
1162                            crate::gpu::q8_matvec_range(
1163                                model,
1164                                *idx,
1165                                cpu_rows,
1166                                &row_scale[cpu_rows..],
1167                                &xs,
1168                                *rows - cpu_rows,
1169                                *cols,
1170                                out_gpu,
1171                            )
1172                        });
1173                        if cpu_rows > 0 {
1174                            // Repack prefix covers the full groups of the
1175                            // CPU half (the split starts at row 0).
1176                            let rep_cpu = if repack.is_empty() {
1177                                &[][..]
1178                            } else {
1179                                &repack[..(cpu_rows / 4) * 4 * *cols]
1180                            };
1181                            qmatvec(
1182                                &bytes[..cpu_rows * *cols],
1183                                rep_cpu,
1184                                &row_scale[..cpu_rows],
1185                                x,
1186                                col_field,
1187                                *dtype,
1188                                cpu_rows,
1189                                *cols,
1190                                out_cpu,
1191                                pool,
1192                            );
1193                        }
1194                        g.join().unwrap_or(false)
1195                    });
1196                    if ok {
1197                        crate::gpu::probe_record(crate::gpu::OpClass::Matvec, true, t0.elapsed());
1198                        return;
1199                    }
1200                    // GPU failed — CPU finishes its half (rows rebased —
1201                    // group offsets don't line up, mmap layout only).
1202                    qmatvec(
1203                        &bytes[cpu_rows * *cols..(*rows) * *cols],
1204                        &[],
1205                        &row_scale[cpu_rows..],
1206                        x,
1207                        col_field,
1208                        *dtype,
1209                        *rows - cpu_rows,
1210                        *cols,
1211                        out_gpu,
1212                        pool,
1213                    );
1214                    return;
1215                }
1216                qmatvec(
1217                    self.quant_bytes(),
1218                    repack,
1219                    row_scale,
1220                    x,
1221                    col_field,
1222                    *dtype,
1223                    *rows,
1224                    *cols,
1225                    out,
1226                    pool,
1227                );
1228            }
1229        }
1230    }
1231
1232    /// Fused two-input matvec (MTP verify pair): weights streamed once.
1233    pub fn matvec2(
1234        &self,
1235        x1: &[f32],
1236        x2: &[f32],
1237        o1: &mut [f32],
1238        o2: &mut [f32],
1239        pool: Option<&Pool>,
1240    ) {
1241        match self {
1242            Self::F32 { data, .. } => matvec_rows2(pool, data, x1, x2, o1, o2),
1243            Self::Mapped {
1244                dtype,
1245                rows,
1246                cols,
1247                row_scale,
1248                col_field,
1249                vbit_offsets,
1250                ..
1251            } => {
1252                if *dtype == TensorDtype::Q4Block {
1253                    q4matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1254                    return;
1255                }
1256                if *dtype == TensorDtype::Q4Tiled {
1257                    q4t_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1258                    return;
1259                }
1260                if *dtype == TensorDtype::Q4TiledP {
1261                    q4tp_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1262                    return;
1263                }
1264                if *dtype == TensorDtype::Q2TiledP {
1265                    q2tp_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1266                    return;
1267                }
1268                if *dtype == TensorDtype::Q1 {
1269                    q1_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1270                    return;
1271                }
1272                if *dtype == TensorDtype::Q1T {
1273                    // Fused ternary pair: one row pass, the register
1274                    // unpack shared across both streams on ARM. (Q1T
1275                    // lacks a row_scale array — scales live inline in
1276                    // the tiles — so it must not fall through to the
1277                    // q8 qmatvec2 below.)
1278                    q1t_matvec2(self.quant_bytes(), x1, x2, *rows, *cols, o1, o2, pool);
1279                    return;
1280                }
1281                if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1282                    vbitmatvec2(
1283                        self.quant_bytes(),
1284                        vbit_offsets,
1285                        x1,
1286                        x2,
1287                        *rows,
1288                        *cols,
1289                        o1,
1290                        o2,
1291                        pool,
1292                    );
1293                    return;
1294                }
1295                qmatvec2(
1296                    self.quant_bytes(),
1297                    row_scale,
1298                    x1,
1299                    x2,
1300                    col_field,
1301                    *dtype,
1302                    *rows,
1303                    *cols,
1304                    o1,
1305                    o2,
1306                    pool,
1307                );
1308            }
1309        }
1310    }
1311}
1312
1313impl QTensor {
1314    /// Batched matvec (prefill-GEMM): xs — row-major [b, cols],
1315    /// out — row-major [b, rows]. Element-wise semantics are IDENTICAL
1316    /// to b matvec calls (same dot kernels in the same order); the win —
1317    /// the weight row streams from DRAM once per batch, not b times.
1318    pub fn matmat(&self, xs_all: &[f32], b: usize, out: &mut [f32], pool: Option<&Pool>) {
1319        let cols = self.cols();
1320        let rows = self.rows();
1321        debug_assert_eq!(xs_all.len(), b * cols);
1322        debug_assert_eq!(out.len(), b * rows);
1323        // GPTQ calibration: fold this layer's inputs into its Hessian. Only
1324        // Mapped tensors carry a directory name; the check is a relaxed
1325        // atomic load, free when not calibrating.
1326        if crate::gptq_capture::capturing() {
1327            if let Self::Mapped { model, idx, .. } = self {
1328                crate::gptq_capture::accumulate(&model.tensors[*idx].name, xs_all, b, cols);
1329            }
1330        }
1331        match self {
1332            Self::F32 { data, .. } => {
1333                let out_addr = SendMut(out.as_mut_ptr());
1334                let run = |start: usize, end: usize| {
1335                    for o in start..end {
1336                        let row = &data[o * cols..(o + 1) * cols];
1337                        for bi in 0..b {
1338                            let x = &xs_all[bi * cols..(bi + 1) * cols];
1339                            let mut acc = 0f32;
1340                            for j in 0..cols {
1341                                acc += row[j] * x[j];
1342                            }
1343                            unsafe { *out_addr.at(bi * rows + o) = acc };
1344                        }
1345                    }
1346                };
1347                dispatch_rows(pool, rows, &run);
1348            }
1349            Self::Mapped {
1350                dtype,
1351                row_scale,
1352                col_field,
1353                vbit_offsets,
1354                ..
1355            } => {
1356                if *dtype == TensorDtype::Q4Block {
1357                    q4matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1358                    return;
1359                }
1360                if *dtype == TensorDtype::Q4TiledP {
1361                    // GPU batched q4tp GEMM (dequant + f32nt mul_mm on the
1362                    // device); the probe keeps whichever beats the CPU arm.
1363                    // Narrow (prompt-encode) and wide (DiT) batches probe
1364                    // as separate classes — the regimes have opposite
1365                    // winners and one shared verdict locked the wrong arm.
1366                    // Kill switch (gpu::mm_kill): one grossly slow GPU op
1367                    // (a fair-condition op is ≤~100 ms even at 1024px)
1368                    // means the device is contended by another process
1369                    // (e.g. a simulator) — verdicts are per-process, so
1370                    // without the bail the whole render crawls behind
1371                    // someone else's queue.
1372                    if b >= 32
1373                        && b * rows * cols >= 128_000_000
1374                        && cols % 32 == 0
1375                        && !crate::gpu::mm_killed()
1376                        && crate::gpu::enabled_here()
1377                    {
1378                        let class = if b >= 128 {
1379                            crate::gpu::OpClass::MatmatWide
1380                        } else {
1381                            crate::gpu::OpClass::Matmat
1382                        };
1383                        if let Self::Mapped { model, idx, .. } = self {
1384                            let t0 = std::time::Instant::now();
1385                            // A cold call takes the device arm: its sample
1386                            // is discarded either way, and the upload is
1387                            // what the next step needs.
1388                            let resident = crate::gpu::weight_is_resident(model, *idx);
1389                            match crate::gpu::probe_arm_cold_prefers_gpu(class, resident) {
1390                                crate::gpu::ProbeArm::Gpu => {
1391                                    if crate::gpu::q4tp_matmat(
1392                                        model, *idx, xs_all, b, rows, cols, out,
1393                                    ) {
1394                                        let el = t0.elapsed();
1395                                        // Work-proportional budget: ~8× the
1396                                        // fair-device estimate (+20 ms slack).
1397                                        // An absolute cap missed the worst
1398                                        // case — contended ops sit at
1399                                        // 100–240 ms each and still bury a
1400                                        // render whose fair op is 3–9 ms.
1401                                        // Cold ops (first PSO build, buffer
1402                                        // alloc) are exempt: a one-off
1403                                        // ~50 ms compile is not contention.
1404                                        let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
1405                                        let budget = std::time::Duration::from_secs_f64(
1406                                            flops / 1.5e12 * 8.0 + 0.020,
1407                                        );
1408                                        if el > budget && !crate::gpu::probe_was_cold() {
1409                                            tracing::warn!(
1410                                                "gpu q4tp matmat took {el:?} (budget {budget:?}) — \
1411                                                 device contended, CPU for the rest of the process"
1412                                            );
1413                                            crate::gpu::mm_kill();
1414                                        }
1415                                        crate::gpu::probe_record(class, true, el);
1416                                        return;
1417                                    }
1418                                }
1419                                crate::gpu::ProbeArm::CpuTimed => {
1420                                    q4tp_matmat(
1421                                        self.quant_bytes(),
1422                                        xs_all,
1423                                        b,
1424                                        rows,
1425                                        cols,
1426                                        out,
1427                                        pool,
1428                                    );
1429                                    crate::gpu::probe_record(class, false, t0.elapsed());
1430                                    return;
1431                                }
1432                                crate::gpu::ProbeArm::Cpu => {}
1433                            }
1434                        }
1435                    }
1436                    q4tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1437                    return;
1438                }
1439                if *dtype == TensorDtype::Q2TiledP {
1440                    // Same device arm as q4tp, behind the same probe:
1441                    // the planes differ, the dispatch does not. Without
1442                    // this a q2tp file ran its widest projections on the
1443                    // host while the 4-bit one had the card, which is a
1444                    // codec paying for its size twice.
1445                    if b >= 32
1446                        && b * rows * cols >= 128_000_000
1447                        && cols % 32 == 0
1448                        && !crate::gpu::mm_killed()
1449                        && crate::gpu::enabled_here()
1450                    {
1451                        let class = if b >= 128 {
1452                            crate::gpu::OpClass::MatmatWide
1453                        } else {
1454                            crate::gpu::OpClass::Matmat
1455                        };
1456                        if let Self::Mapped { model, idx, .. } = self {
1457                            let t0 = std::time::Instant::now();
1458                            match crate::gpu::probe_arm(class) {
1459                                crate::gpu::ProbeArm::Gpu => {
1460                                    if crate::gpu::q2tp_matmat(
1461                                        model, *idx, xs_all, b, rows, cols, out,
1462                                    ) {
1463                                        crate::gpu::probe_record(class, true, t0.elapsed());
1464                                        return;
1465                                    }
1466                                }
1467                                crate::gpu::ProbeArm::CpuTimed => {
1468                                    q2tp_matmat(
1469                                        self.quant_bytes(),
1470                                        xs_all,
1471                                        b,
1472                                        rows,
1473                                        cols,
1474                                        out,
1475                                        pool,
1476                                    );
1477                                    crate::gpu::probe_record(class, false, t0.elapsed());
1478                                    return;
1479                                }
1480                                crate::gpu::ProbeArm::Cpu => {}
1481                            }
1482                        }
1483                    }
1484                    // Without a host arm a q2tp tensor falls through to
1485                    // the q8 fallback, which reads it at one BYTE per
1486                    // weight — a 2x overrun that killed pool workers
1487                    // mid-prefill while the dispatcher waited forever.
1488                    q2tp_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1489                    return;
1490                }
1491                if *dtype == TensorDtype::Q4Tiled {
1492                    // GPU batched q4t GEMM (dequant + f32nt mul_mm on the
1493                    // device); the probe keeps whichever beats the CPU arm.
1494                    // Narrow (prompt-encode) and wide (DiT) batches probe
1495                    // as separate classes — the regimes have opposite
1496                    // winners and one shared verdict locked the wrong arm.
1497                    // Kill switch (gpu::mm_kill): one grossly slow GPU op
1498                    // (a fair-condition op is ≤~100 ms even at 1024px)
1499                    // means the device is contended by another process
1500                    // (e.g. a simulator) — verdicts are per-process, so
1501                    // without the bail the whole render crawls behind
1502                    // someone else's queue.
1503                    if b >= 32
1504                        && b * rows * cols >= 128_000_000
1505                        && cols % 32 == 0
1506                        && !crate::gpu::mm_killed()
1507                        && crate::gpu::enabled_here()
1508                    {
1509                        let class = if b >= 128 {
1510                            crate::gpu::OpClass::MatmatWide
1511                        } else {
1512                            crate::gpu::OpClass::Matmat
1513                        };
1514                        if let Self::Mapped { model, idx, .. } = self {
1515                            let t0 = std::time::Instant::now();
1516                            match crate::gpu::probe_arm(class) {
1517                                crate::gpu::ProbeArm::Gpu => {
1518                                    if crate::gpu::q4t_matmat(
1519                                        model, *idx, xs_all, b, rows, cols, out,
1520                                    ) {
1521                                        let el = t0.elapsed();
1522                                        // Work-proportional budget: ~8× the
1523                                        // fair-device estimate (+20 ms slack).
1524                                        // An absolute cap missed the worst
1525                                        // case — contended ops sit at
1526                                        // 100–240 ms each and still bury a
1527                                        // render whose fair op is 3–9 ms.
1528                                        // Cold ops (first PSO build, buffer
1529                                        // alloc) are exempt: a one-off
1530                                        // ~50 ms compile is not contention.
1531                                        let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
1532                                        let budget = std::time::Duration::from_secs_f64(
1533                                            flops / 1.5e12 * 8.0 + 0.020,
1534                                        );
1535                                        if el > budget && !crate::gpu::probe_was_cold() {
1536                                            tracing::warn!(
1537                                                "gpu q4t matmat took {el:?} (budget {budget:?}) — \
1538                                                 device contended, CPU for the rest of the process"
1539                                            );
1540                                            crate::gpu::mm_kill();
1541                                        }
1542                                        crate::gpu::probe_record(class, true, el);
1543                                        return;
1544                                    }
1545                                }
1546                                crate::gpu::ProbeArm::CpuTimed => {
1547                                    q4t_matmat(
1548                                        self.quant_bytes(),
1549                                        xs_all,
1550                                        b,
1551                                        rows,
1552                                        cols,
1553                                        out,
1554                                        pool,
1555                                    );
1556                                    crate::gpu::probe_record(class, false, t0.elapsed());
1557                                    return;
1558                                }
1559                                crate::gpu::ProbeArm::Cpu => {}
1560                            }
1561                        }
1562                    }
1563                    q4t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1564                    return;
1565                }
1566                if *dtype == TensorDtype::Q1 {
1567                    // GPU batched q1 GEMM for wide prefill (q1_mul_mm on the
1568                    // device); the probe keeps whichever beats the CPU matmat.
1569                    if b >= 32
1570                        && b * rows * cols >= 128_000_000
1571                        && cols % 64 == 0
1572                        && crate::gpu::enabled_here()
1573                    {
1574                        if let Self::Mapped { model, idx, .. } = self {
1575                            let t0 = std::time::Instant::now();
1576                            match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
1577                                crate::gpu::ProbeArm::Gpu => {
1578                                    if crate::gpu::q1_matmat(
1579                                        model, *idx, xs_all, b, rows, cols, out,
1580                                    ) {
1581                                        crate::gpu::probe_record(
1582                                            crate::gpu::OpClass::Matmat,
1583                                            true,
1584                                            t0.elapsed(),
1585                                        );
1586                                        return;
1587                                    }
1588                                }
1589                                crate::gpu::ProbeArm::CpuTimed => {
1590                                    q1_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1591                                    crate::gpu::probe_record(
1592                                        crate::gpu::OpClass::Matmat,
1593                                        false,
1594                                        t0.elapsed(),
1595                                    );
1596                                    return;
1597                                }
1598                                crate::gpu::ProbeArm::Cpu => {}
1599                            }
1600                        }
1601                    }
1602                    q1_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1603                    return;
1604                }
1605                if *dtype == TensorDtype::Q1T {
1606                    // GPU batched GEMM for wide prefill (base + overlay on the
1607                    // device); probe keeps the winner vs the CPU matmat.
1608                    if b >= 32 && b * rows * cols >= 128_000_000 && crate::gpu::enabled_here() {
1609                        if let Self::Mapped { model, idx, .. } = self {
1610                            let t0 = std::time::Instant::now();
1611                            match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
1612                                crate::gpu::ProbeArm::Gpu => {
1613                                    if crate::gpu::q1t_matmat(
1614                                        model, *idx, xs_all, b, rows, cols, out,
1615                                    ) {
1616                                        crate::gpu::probe_record(
1617                                            crate::gpu::OpClass::Matmat,
1618                                            true,
1619                                            t0.elapsed(),
1620                                        );
1621                                        return;
1622                                    }
1623                                }
1624                                crate::gpu::ProbeArm::CpuTimed => {
1625                                    q1t_matmat(
1626                                        self.quant_bytes(),
1627                                        xs_all,
1628                                        b,
1629                                        rows,
1630                                        cols,
1631                                        out,
1632                                        pool,
1633                                    );
1634                                    crate::gpu::probe_record(
1635                                        crate::gpu::OpClass::Matmat,
1636                                        false,
1637                                        t0.elapsed(),
1638                                    );
1639                                    return;
1640                                }
1641                                crate::gpu::ProbeArm::Cpu => {}
1642                            }
1643                        }
1644                    }
1645                    q1t_matmat(self.quant_bytes(), xs_all, b, rows, cols, out, pool);
1646                    return;
1647                }
1648                if matches!(dtype, TensorDtype::Vbit | TensorDtype::VbitRo) {
1649                    vbitmatmat(
1650                        self.quant_bytes(),
1651                        vbit_offsets,
1652                        xs_all,
1653                        b,
1654                        rows,
1655                        cols,
1656                        out,
1657                        pool,
1658                    );
1659                    return;
1660                }
1661                let pre: Vec<std::borrow::Cow<'_, [f32]>> = (0..b)
1662                    .map(|bi| prescale(&xs_all[bi * cols..(bi + 1) * cols], col_field, *dtype))
1663                    .collect();
1664                // D5: large prefill-batch GEMMs — on the GPU (threshold by
1665                // work volume: submission carries b×rows×cols MACs).
1666                // Runtime probe: the naive GEMM shader + sync readback
1667                // lose to the CPU GEMM on slow driver stacks — alternate
1668                // both arms and keep the winner.
1669                if b >= 8 && b * rows * cols >= 128_000_000 && crate::gpu::enabled_here() {
1670                    if let Self::Mapped { model, idx, .. } = self {
1671                        let t0 = std::time::Instant::now();
1672                        match crate::gpu::probe_arm(crate::gpu::OpClass::Matmat) {
1673                            crate::gpu::ProbeArm::Gpu
1674                                if crate::gpu::probe_deciding(crate::gpu::OpClass::Matmat)
1675                                    && !crate::gpu::q8_resident_or_upload(model, *idx) =>
1676                            {
1677                                // Cold weights during probing: the upload
1678                                // has started, the count runs on the CPU —
1679                                // the GPU arm samples on the next touch.
1680                                let q = self.quant_bytes();
1681                                qmatmat(q, row_scale, &pre, rows, cols, out, pool);
1682                                return;
1683                            }
1684                            crate::gpu::ProbeArm::Gpu => {
1685                                let flat: Vec<f32> =
1686                                    pre.iter().flat_map(|v| v.iter().copied()).collect();
1687                                if crate::gpu::q8_matmat(
1688                                    model, *idx, row_scale, &flat, b, rows, cols, out,
1689                                ) {
1690                                    crate::gpu::probe_record(
1691                                        crate::gpu::OpClass::Matmat,
1692                                        true,
1693                                        t0.elapsed(),
1694                                    );
1695                                    return;
1696                                }
1697                            }
1698                            crate::gpu::ProbeArm::CpuTimed => {
1699                                let q = self.quant_bytes();
1700                                qmatmat(q, row_scale, &pre, rows, cols, out, pool);
1701                                crate::gpu::probe_record(
1702                                    crate::gpu::OpClass::Matmat,
1703                                    false,
1704                                    t0.elapsed(),
1705                                );
1706                                return;
1707                            }
1708                            crate::gpu::ProbeArm::Cpu => {}
1709                        }
1710                    }
1711                }
1712                let q = self.quant_bytes();
1713                qmatmat(q, row_scale, &pre, rows, cols, out, pool);
1714            }
1715        }
1716    }
1717}
1718
1719impl QTensor {
1720    /// Multi-matrix job (roadmap §3 P0): N tensors sharing one input
1721    /// run under a SINGLE pool dispatch — QKV or gate+up cost one
1722    /// barrier instead of N. Per-row math is the exact same kernel as
1723    /// `matvec` (bit-identical outputs); only the dispatch is fused.
1724    /// Falls back to N sequential matvecs when the set is not a uniform
1725    /// q8-family/F32 group or there is no pool.
1726    pub fn matvec_many<const N: usize>(
1727        ts: [&QTensor; N],
1728        x: &[f32],
1729        mut outs: [&mut [f32]; N],
1730        pool: Option<&Pool>,
1731    ) {
1732        let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
1733        let uniform_q8 = ts.iter().all(|t| {
1734            matches!(
1735                t,
1736                Self::Mapped {
1737                    dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f,
1738                    ..
1739                }
1740            )
1741        });
1742        let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
1743        let uniform_q4 = ts.iter().all(|t| {
1744            matches!(
1745                t,
1746                Self::Mapped {
1747                    dtype: TensorDtype::Q4Block,
1748                    ..
1749                }
1750            )
1751        });
1752        let uniform_vbit = ts.iter().all(|t| {
1753            matches!(
1754                t,
1755                Self::Mapped {
1756                    dtype: TensorDtype::Vbit | TensorDtype::VbitRo,
1757                    ..
1758                }
1759            )
1760        });
1761        let uniform_q1 = ts.iter().all(|t| {
1762            matches!(
1763                t,
1764                Self::Mapped {
1765                    dtype: TensorDtype::Q1,
1766                    ..
1767                }
1768            )
1769        });
1770        let uniform_q1t = ts.iter().all(|t| {
1771            matches!(
1772                t,
1773                Self::Mapped {
1774                    dtype: TensorDtype::Q1T,
1775                    ..
1776                }
1777            )
1778        });
1779        // q4tp is the skeleton dtype of the big MoE files, and without an arm
1780        // here every projection that shares an input paid its own pool
1781        // barrier: DeepSeek-V4's attention step alone hands this function
1782        // wq_a, wkv and both compressors' pairs off the same hidden state.
1783        let uniform_q4tp = ts.iter().all(|t| {
1784            matches!(
1785                t,
1786                Self::Mapped {
1787                    dtype: TensorDtype::Q4TiledP,
1788                    ..
1789                }
1790            )
1791        }) && ts.iter().all(|t| t.cols() == ts[0].cols() && t.cols() % GROUP_SIZE == 0);
1792        let Some(pool) = pool else {
1793            for (t, o) in ts.iter().zip(outs.iter_mut()) {
1794                t.matvec(x, o, None);
1795            }
1796            return;
1797        };
1798        if total_rows < 256
1799            || !(uniform_q8
1800                || uniform_f32
1801                || uniform_q4
1802                || uniform_vbit
1803                || uniform_q1
1804                || uniform_q1t
1805                || uniform_q4tp)
1806        {
1807            for (t, o) in ts.iter().zip(outs.iter_mut()) {
1808                t.matvec(x, o, Some(pool));
1809            }
1810            return;
1811        }
1812
1813        if uniform_q4tp {
1814            // Every tensor's rows laid end to end in one virtual row space,
1815            // so the whole set is ONE dispatch. The per-row body is the
1816            // `q4tp_matvec` arm verbatim — same activation split, same
1817            // accumulation order — so the outputs are bit-identical to the
1818            // sequential calls this replaces.
1819            let cols = ts[0].cols();
1820            let gpr = cols / GROUP_SIZE;
1821            let views: Vec<Q4tpView> = ts
1822                .iter()
1823                .map(|t| Q4tpView::new(t.quant_bytes(), t.rows(), cols))
1824                .collect();
1825            let rows_of: Vec<usize> = ts.iter().map(|t| t.rows()).collect();
1826            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1827            // flat index -> (which tensor, which of its rows)
1828            let locate = |flat: usize| -> (usize, usize) {
1829                let mut acc = 0;
1830                for (i, &r) in rows_of.iter().enumerate() {
1831                    if flat < acc + r {
1832                        return (i, flat - acc);
1833                    }
1834                    acc += r;
1835                }
1836                (rows_of.len() - 1, 0)
1837            };
1838            let (views, outs_addr) = (&views, &outs_addr);
1839            if a8w8_enabled() {
1840                let act = split_act(x);
1841                let act = &act;
1842                let run = |start: usize, end: usize| {
1843                    let mut sc = vec![0f32; gpr];
1844                    for flat in start..end {
1845                        let (t, r) = locate(flat);
1846                        let v = &views[t];
1847                        v.scales_into(r, gpr, &mut sc);
1848                        let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, &sc) * act.sx;
1849                        for &(j, xv) in &act.outliers {
1850                            let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
1851                            acc += w * s * xv;
1852                        }
1853                        // SAFETY: one worker owns each (tensor, row) pair.
1854                        unsafe { *outs_addr[t].at(r) = acc };
1855                    }
1856                };
1857                pool.run_rows(total_rows, &run);
1858            } else {
1859                let run = |start: usize, end: usize| {
1860                    let mut sc = vec![0f32; gpr];
1861                    for flat in start..end {
1862                        let (t, r) = locate(flat);
1863                        let v = &views[t];
1864                        v.scales_into(r, gpr, &mut sc);
1865                        // SAFETY: one worker owns each (tensor, row) pair.
1866                        unsafe { *outs_addr[t].at(r) = q4tp_row_exact(v.nib, r, gpr, x, &sc) };
1867                    }
1868                };
1869                pool.run_rows(total_rows, &run);
1870            }
1871            return;
1872        }
1873
1874        if uniform_q1 {
1875            // One shared activation split + group sums (q1 has no col
1876            // field; the same input feeds every tensor).
1877            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1878            if a8w8_enabled() {
1879                let act = split_act(x);
1880                let gsum = q1_group_sums(&act.xq, ts[0].cols() / GROUP_SIZE);
1881                let (act, gsum) = (&act, &gsum);
1882                let closures: [_; N] = std::array::from_fn(|i| {
1883                    let (bytes, gpr, out) =
1884                        (ts[i].quant_bytes(), ts[i].cols() / GROUP_SIZE, outs_addr[i]);
1885                    move |s: usize, e: usize| q1_range_a8w8(bytes, gpr, act, gsum, out, s, e)
1886                });
1887                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1888                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1889                pool.run_many(&parts);
1890            } else {
1891                let closures: [_; N] = std::array::from_fn(|i| {
1892                    let (bytes, gpr, out) =
1893                        (ts[i].quant_bytes(), ts[i].cols() / GROUP_SIZE, outs_addr[i]);
1894                    move |s: usize, e: usize| q1_range_f32(bytes, gpr, x, out, s, e)
1895                });
1896                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1897                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1898                pool.run_many(&parts);
1899            }
1900            return;
1901        }
1902
1903        if uniform_q1t {
1904            // Q1T batched: one shared activation split + overlay decode,
1905            // all tensors' rows in ONE pool dispatch (saves N−1 dispatches
1906            // and N−1 redundant split_act calls per layer).
1907            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1908            const TILE: usize = cortiq_core::quant::Q1T_TILE;
1909            if a8w8_enabled() {
1910                let act = split_act(x);
1911                let act = &act;
1912                let x_ref = x;
1913                let closures: [_; N] = std::array::from_fn(|i| {
1914                    let bytes = ts[i].quant_bytes();
1915                    let (rows, cols) = (ts[i].rows(), ts[i].cols());
1916                    let gpr = cols / GROUP_SIZE;
1917                    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
1918                    let out = outs_addr[i];
1919                    move |s: usize, e: usize| {
1920                        q1t_range_a8w8(bytes, gpr, rp_off, ent_off, has_ov, act, x_ref, out, s, e)
1921                    }
1922                });
1923                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1924                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1925                pool.run_many(&parts);
1926            } else {
1927                let x_ref = x;
1928                let closures: [_; N] = std::array::from_fn(|i| {
1929                    let bytes = ts[i].quant_bytes();
1930                    let (rows, cols) = (ts[i].rows(), ts[i].cols());
1931                    let gpr = cols / GROUP_SIZE;
1932                    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
1933                    let out = outs_addr[i];
1934                    move |s: usize, e: usize| {
1935                        q1t_range_f32_batch(bytes, gpr, rp_off, ent_off, has_ov, x_ref, out, s, e)
1936                    }
1937                });
1938                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1939                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1940                pool.run_many(&parts);
1941            }
1942            return;
1943        }
1944
1945        if uniform_q4 || uniform_vbit {
1946            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
1947            // q4/vbit share one activation split — no per-tensor col field.
1948            if a8w8_enabled() {
1949                let act = split_act(x);
1950                let act = &act;
1951                if uniform_q4 {
1952                    let closures: [_; N] = std::array::from_fn(|i| {
1953                        let (packed, scales) =
1954                            q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
1955                        let (gpr, cols, out) =
1956                            (ts[i].cols() / GROUP_SIZE, ts[i].cols(), outs_addr[i]);
1957                        move |s: usize, e: usize| {
1958                            q4_range_a8w8(packed, scales, gpr, cols, act, out, s, e)
1959                        }
1960                    });
1961                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1962                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1963                    pool.run_many(&parts);
1964                } else {
1965                    let closures: [_; N] = std::array::from_fn(|i| {
1966                        let Self::Mapped { vbit_offsets, .. } = ts[i] else {
1967                            unreachable!()
1968                        };
1969                        let (bytes, rows, cols, out) = (
1970                            ts[i].quant_bytes(),
1971                            ts[i].rows(),
1972                            ts[i].cols(),
1973                            outs_addr[i],
1974                        );
1975                        move |s: usize, e: usize| {
1976                            vbit_range_a8w8(bytes, vbit_offsets, x, act, rows, cols, out, s, e)
1977                        }
1978                    });
1979                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1980                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1981                    pool.run_many(&parts);
1982                }
1983                return;
1984            }
1985            if uniform_q4 {
1986                let closures: [_; N] = std::array::from_fn(|i| {
1987                    let (packed, scales) =
1988                        q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
1989                    let (gpr, out) = (ts[i].cols() / GROUP_SIZE, outs_addr[i]);
1990                    move |s: usize, e: usize| q4_range_f32(packed, scales, gpr, x, out, s, e)
1991                });
1992                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
1993                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
1994                pool.run_many(&parts);
1995            } else {
1996                let closures: [_; N] = std::array::from_fn(|i| {
1997                    let Self::Mapped { vbit_offsets, .. } = ts[i] else {
1998                        unreachable!()
1999                    };
2000                    let (bytes, rows, cols, out) = (
2001                        ts[i].quant_bytes(),
2002                        ts[i].rows(),
2003                        ts[i].cols(),
2004                        outs_addr[i],
2005                    );
2006                    move |s: usize, e: usize| {
2007                        vbit_range_f32(bytes, vbit_offsets, x, rows, cols, out, s, e)
2008                    }
2009                });
2010                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2011                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2012                pool.run_many(&parts);
2013            }
2014            return;
2015        }
2016
2017        if uniform_f32 {
2018            let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2019            let closures: [_; N] = std::array::from_fn(|i| {
2020                let Self::F32 { data, cols, .. } = ts[i] else {
2021                    unreachable!()
2022                };
2023                let out = outs_addr[i];
2024                move |start: usize, end: usize| {
2025                    for o in start..end {
2026                        let row = &data[o * cols..(o + 1) * cols];
2027                        let mut sum = 0.0f32;
2028                        for j in 0..*cols {
2029                            sum += row[j] * x[j];
2030                        }
2031                        // SAFETY: disjoint (tensor, row) cells per worker.
2032                        unsafe { *out.at(o) = sum };
2033                    }
2034                }
2035            });
2036            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2037                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2038            pool.run_many(&parts);
2039            return;
2040        }
2041
2042        // Uniform q8-family: per-tensor prescale (q8_2f col fields
2043        // differ per tensor) + the shared range kernels.
2044        struct Ctx<'a> {
2045            bytes: &'a [u8],
2046            #[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
2047            rep: &'a [u8],
2048            row_scale: &'a [f32],
2049            cols: usize,
2050            xs: std::borrow::Cow<'a, [f32]>,
2051        }
2052        let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
2053            let Self::Mapped {
2054                dtype,
2055                cols,
2056                row_scale,
2057                col_field,
2058                repack,
2059                ..
2060            } = ts[i]
2061            else {
2062                unreachable!()
2063            };
2064            Ctx {
2065                bytes: ts[i].quant_bytes(),
2066                rep: repack,
2067                row_scale,
2068                cols: *cols,
2069                xs: prescale(x, col_field, *dtype),
2070            }
2071        });
2072        let outs_addr: [SendMut; N] = std::array::from_fn(|i| SendMut(outs[i].as_mut_ptr()));
2073        #[cfg(target_arch = "aarch64")]
2074        if sdot_enabled() {
2075            let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
2076            let closures: [_; N] = std::array::from_fn(|i| {
2077                let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
2078                move |start: usize, end: usize| {
2079                    q8_range_sdot(c.bytes, c.rep, c.row_scale, act, c.cols, out, start, end)
2080                }
2081            });
2082            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2083                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2084            pool.run_many(&parts);
2085            return;
2086        }
2087        #[cfg(target_arch = "x86_64")]
2088        if avx2_a8w8_enabled() {
2089            let acts: [SplitAct; N] = std::array::from_fn(|i| split_act(&ctxs[i].xs));
2090            let closures: [_; N] = std::array::from_fn(|i| {
2091                let (c, act, out) = (&ctxs[i], &acts[i], outs_addr[i]);
2092                move |start: usize, end: usize| {
2093                    q8_range_avx2(c.bytes, c.row_scale, act, c.cols, out, start, end)
2094                }
2095            });
2096            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2097                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2098            pool.run_many(&parts);
2099            return;
2100        }
2101        let closures: [_; N] = std::array::from_fn(|i| {
2102            let (c, out) = (&ctxs[i], outs_addr[i]);
2103            move |start: usize, end: usize| {
2104                q8_range_f32(c.bytes, c.row_scale, &c.xs, c.cols, out, start, end)
2105            }
2106        });
2107        let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2108            std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2109        pool.run_many(&parts);
2110    }
2111}
2112
2113impl QTensor {
2114    /// Pair-input multi-matrix job: N tensors × 2 shared inputs under a
2115    /// single pool dispatch — the MTP/pair decode path publishes one job
2116    /// for Q/K/V (and one for gate+up) instead of one per tensor.
2117    /// Per-row math is exactly `matvec2`'s kernels; bit-identical.
2118    #[allow(clippy::needless_range_loop)]
2119    pub fn matvec2_many<const N: usize>(
2120        ts: [&QTensor; N],
2121        x1: &[f32],
2122        x2: &[f32],
2123        mut o1s: [&mut [f32]; N],
2124        mut o2s: [&mut [f32]; N],
2125        pool: Option<&Pool>,
2126    ) {
2127        let total_rows: usize = ts.iter().map(|t| t.rows()).sum();
2128        let uniform_q8 = ts.iter().all(|t| {
2129            matches!(
2130                t,
2131                Self::Mapped {
2132                    dtype: TensorDtype::Q8Row | TensorDtype::Q8_2f,
2133                    ..
2134                }
2135            )
2136        });
2137        let uniform_f32 = ts.iter().all(|t| matches!(t, Self::F32 { .. }));
2138        let uniform_q4 = ts.iter().all(|t| {
2139            matches!(
2140                t,
2141                Self::Mapped {
2142                    dtype: TensorDtype::Q4Block,
2143                    ..
2144                }
2145            )
2146        });
2147        let uniform_vbit = ts.iter().all(|t| {
2148            matches!(
2149                t,
2150                Self::Mapped {
2151                    dtype: TensorDtype::Vbit | TensorDtype::VbitRo,
2152                    ..
2153                }
2154            )
2155        });
2156        let fusable = pool.is_some()
2157            && total_rows >= 256
2158            && (uniform_q8 || uniform_f32 || uniform_q4 || uniform_vbit);
2159        if !fusable {
2160            for i in 0..N {
2161                ts[i].matvec2(x1, x2, o1s[i], o2s[i], pool);
2162            }
2163            return;
2164        }
2165        let pool = pool.unwrap();
2166
2167        if uniform_q4 || uniform_vbit {
2168            let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
2169            let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
2170            // q4/vbit share activation splits — no per-tensor col field.
2171            if a8w8_enabled() {
2172                let a1 = split_act(x1);
2173                let a2 = split_act(x2);
2174                let (a1, a2) = (&a1, &a2);
2175                if uniform_q4 {
2176                    let closures: [_; N] = std::array::from_fn(|i| {
2177                        let (packed, scales) =
2178                            q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
2179                        let (gpr, cols, o1, o2) =
2180                            (ts[i].cols() / GROUP_SIZE, ts[i].cols(), p1[i], p2[i]);
2181                        move |s: usize, e: usize| {
2182                            q4_range2_a8w8(packed, scales, gpr, cols, a1, a2, o1, o2, s, e)
2183                        }
2184                    });
2185                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2186                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2187                    pool.run_many(&parts);
2188                } else {
2189                    let closures: [_; N] = std::array::from_fn(|i| {
2190                        let Self::Mapped { vbit_offsets, .. } = ts[i] else {
2191                            unreachable!()
2192                        };
2193                        let (bytes, rows, cols, o1, o2) = (
2194                            ts[i].quant_bytes(),
2195                            ts[i].rows(),
2196                            ts[i].cols(),
2197                            p1[i],
2198                            p2[i],
2199                        );
2200                        move |s: usize, e: usize| {
2201                            vbit_range2_a8w8(
2202                                bytes,
2203                                vbit_offsets,
2204                                x1,
2205                                x2,
2206                                a1,
2207                                a2,
2208                                rows,
2209                                cols,
2210                                o1,
2211                                o2,
2212                                s,
2213                                e,
2214                            )
2215                        }
2216                    });
2217                    let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2218                        std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2219                    pool.run_many(&parts);
2220                }
2221                return;
2222            }
2223            if uniform_q4 {
2224                let closures: [_; N] = std::array::from_fn(|i| {
2225                    let (packed, scales) =
2226                        q4_split(ts[i].quant_bytes(), ts[i].rows(), ts[i].cols());
2227                    let (gpr, o1, o2) = (ts[i].cols() / GROUP_SIZE, p1[i], p2[i]);
2228                    move |s: usize, e: usize| {
2229                        q4_range2_f32(packed, scales, gpr, x1, x2, o1, o2, s, e)
2230                    }
2231                });
2232                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2233                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2234                pool.run_many(&parts);
2235            } else {
2236                let closures: [_; N] = std::array::from_fn(|i| {
2237                    let Self::Mapped { vbit_offsets, .. } = ts[i] else {
2238                        unreachable!()
2239                    };
2240                    let (bytes, rows, cols, o1, o2) = (
2241                        ts[i].quant_bytes(),
2242                        ts[i].rows(),
2243                        ts[i].cols(),
2244                        p1[i],
2245                        p2[i],
2246                    );
2247                    move |s: usize, e: usize| {
2248                        vbit_range2_f32(bytes, vbit_offsets, x1, x2, rows, cols, o1, o2, s, e)
2249                    }
2250                });
2251                let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2252                    std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2253                pool.run_many(&parts);
2254            }
2255            return;
2256        }
2257
2258        if uniform_f32 {
2259            let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
2260            let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
2261            let closures: [_; N] = std::array::from_fn(|i| {
2262                let Self::F32 { data, cols, .. } = ts[i] else {
2263                    unreachable!()
2264                };
2265                let (o1, o2) = (p1[i], p2[i]);
2266                move |start: usize, end: usize| {
2267                    for o in start..end {
2268                        let row = &data[o * cols..(o + 1) * cols];
2269                        let (mut s1, mut s2) = (0.0f32, 0.0f32);
2270                        for j in 0..*cols {
2271                            s1 += row[j] * x1[j];
2272                            s2 += row[j] * x2[j];
2273                        }
2274                        // SAFETY: disjoint (tensor, row) cells per worker.
2275                        unsafe {
2276                            *o1.at(o) = s1;
2277                            *o2.at(o) = s2;
2278                        }
2279                    }
2280                }
2281            });
2282            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2283                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2284            pool.run_many(&parts);
2285            return;
2286        }
2287
2288        struct Ctx<'a> {
2289            bytes: &'a [u8],
2290            row_scale: &'a [f32],
2291            cols: usize,
2292            xs1: std::borrow::Cow<'a, [f32]>,
2293            xs2: std::borrow::Cow<'a, [f32]>,
2294        }
2295        let ctxs: [Ctx<'_>; N] = std::array::from_fn(|i| {
2296            let Self::Mapped {
2297                dtype,
2298                cols,
2299                row_scale,
2300                col_field,
2301                ..
2302            } = ts[i]
2303            else {
2304                unreachable!()
2305            };
2306            Ctx {
2307                bytes: ts[i].quant_bytes(),
2308                row_scale,
2309                cols: *cols,
2310                xs1: prescale(x1, col_field, *dtype),
2311                xs2: prescale(x2, col_field, *dtype),
2312            }
2313        });
2314        let p1: [SendMut; N] = std::array::from_fn(|i| SendMut(o1s[i].as_mut_ptr()));
2315        let p2: [SendMut; N] = std::array::from_fn(|i| SendMut(o2s[i].as_mut_ptr()));
2316        #[cfg(target_arch = "aarch64")]
2317        if sdot_enabled() {
2318            let acts: [(SplitAct, SplitAct); N] =
2319                std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
2320            let closures: [_; N] = std::array::from_fn(|i| {
2321                let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
2322                move |start: usize, end: usize| {
2323                    q8_range2_sdot(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
2324                }
2325            });
2326            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2327                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2328            pool.run_many(&parts);
2329            return;
2330        }
2331        #[cfg(target_arch = "x86_64")]
2332        if avx2_a8w8_enabled() {
2333            let acts: [(SplitAct, SplitAct); N] =
2334                std::array::from_fn(|i| (split_act(&ctxs[i].xs1), split_act(&ctxs[i].xs2)));
2335            let closures: [_; N] = std::array::from_fn(|i| {
2336                let (c, a, o1, o2) = (&ctxs[i], &acts[i], p1[i], p2[i]);
2337                move |start: usize, end: usize| {
2338                    q8_range2_avx2(c.bytes, c.row_scale, &a.0, &a.1, c.cols, o1, o2, start, end)
2339                }
2340            });
2341            let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2342                std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2343            pool.run_many(&parts);
2344            return;
2345        }
2346        let closures: [_; N] = std::array::from_fn(|i| {
2347            let (c, o1, o2) = (&ctxs[i], p1[i], p2[i]);
2348            move |start: usize, end: usize| {
2349                q8_range2_f32(
2350                    c.bytes,
2351                    c.row_scale,
2352                    &c.xs1,
2353                    &c.xs2,
2354                    c.cols,
2355                    o1,
2356                    o2,
2357                    start,
2358                    end,
2359                )
2360            }
2361        });
2362        let parts: [(usize, &(dyn Fn(usize, usize) + Sync)); N] =
2363            std::array::from_fn(|i| (ts[i].rows(), &closures[i] as _));
2364        pool.run_many(&parts);
2365    }
2366
2367    /// Fused gate+up matvec with SiLU·mul: for each row r, computes
2368    /// `silu(gate·x) * (up·x)` and writes to `out[r]`. ONE pool dispatch,
2369    /// no intermediate g/u buffers, no separate silu pass. Falls back
2370    /// (returns false) for unsupported dtype combos.
2371    pub fn matvec_silu_mul(
2372        gate: &QTensor,
2373        up: &QTensor,
2374        x: &[f32],
2375        out: &mut [f32],
2376        pool: Option<&Pool>,
2377    ) -> bool {
2378        let inter = gate.rows();
2379        debug_assert_eq!(up.rows(), inter);
2380        debug_assert_eq!(out.len(), inter);
2381        debug_assert_eq!(gate.cols(), up.cols());
2382        if !a8w8_enabled() {
2383            return false;
2384        }
2385        let act = split_act(x);
2386        let act = &act;
2387        let x_ref = x;
2388        let out_addr = SendMut(out.as_mut_ptr());
2389
2390        match (gate, up) {
2391            // Q4Block gate + Q4Block up (most common mobile q4 models)
2392            (
2393                Self::Mapped {
2394                    dtype: TensorDtype::Q4Block,
2395                    ..
2396                },
2397                Self::Mapped {
2398                    dtype: TensorDtype::Q4Block,
2399                    ..
2400                },
2401            ) => {
2402                let (gp, gs) = q4_split(gate.quant_bytes(), gate.rows(), gate.cols());
2403                let (up_p, up_s) = q4_split(up.quant_bytes(), up.rows(), up.cols());
2404                let gpr = gate.cols() / GROUP_SIZE;
2405                let cols = gate.cols();
2406                let run = move |start: usize, end: usize| {
2407                    for r in start..end {
2408                        let mut gv = dot_q4_row_i8(gp, gs, r * gpr, gpr, &act.xq) * act.sx;
2409                        let mut uv = dot_q4_row_i8(up_p, up_s, r * gpr, gpr, &act.xq) * act.sx;
2410                        for &(j, xv) in &act.outliers {
2411                            let flat = r * cols + j;
2412                            let gb = gp[flat / 2];
2413                            let gn = if flat & 1 == 0 { gb & 0x0F } else { gb >> 4 };
2414                            let gsc = f16_to_f32(u16::from_le_bytes([
2415                                gs[(flat / GROUP_SIZE) * 2],
2416                                gs[(flat / GROUP_SIZE) * 2 + 1],
2417                            ]));
2418                            gv += ((gn as i32 - 8) as f32) * gsc * xv;
2419                            let ub = up_p[flat / 2];
2420                            let un = if flat & 1 == 0 { ub & 0x0F } else { ub >> 4 };
2421                            let usc = f16_to_f32(u16::from_le_bytes([
2422                                up_s[(flat / GROUP_SIZE) * 2],
2423                                up_s[(flat / GROUP_SIZE) * 2 + 1],
2424                            ]));
2425                            uv += ((un as i32 - 8) as f32) * usc * xv;
2426                        }
2427                        let silu_g = gv / (1.0 + (-gv).exp());
2428                        // SAFETY: disjoint row ranges per worker.
2429                        unsafe { *out_addr.at(r) = silu_g * uv };
2430                    }
2431                };
2432                dispatch_rows(pool, inter, &run);
2433                true
2434            }
2435            // Q4Tiled gate + Q4Tiled up — one row pass, both tile
2436            // streams sequential, silu·mul fused (same per-row math as
2437            // `q4t_matvec`).
2438            (
2439                Self::Mapped {
2440                    dtype: TensorDtype::Q4Tiled,
2441                    ..
2442                },
2443                Self::Mapped {
2444                    dtype: TensorDtype::Q4Tiled,
2445                    ..
2446                },
2447            ) => {
2448                let g_bytes = gate.quant_bytes();
2449                let u_bytes = up.quant_bytes();
2450                let gpr = gate.cols() / GROUP_SIZE;
2451                let run = move |start: usize, end: usize| {
2452                    for r in start..end {
2453                        let mut gv = dot_q4t_row_i8(g_bytes, r, gpr, &act.xq) * act.sx;
2454                        let mut uv = dot_q4t_row_i8(u_bytes, r, gpr, &act.xq) * act.sx;
2455                        for &(j, xv) in &act.outliers {
2456                            let (w, s) = q4t_outlier(g_bytes, r, gpr, j);
2457                            gv += w * s * xv;
2458                            let (w, s) = q4t_outlier(u_bytes, r, gpr, j);
2459                            uv += w * s * xv;
2460                        }
2461                        let silu_g = gv / (1.0 + (-gv).exp());
2462                        // SAFETY: disjoint row ranges per worker.
2463                        unsafe { *out_addr.at(r) = silu_g * uv };
2464                    }
2465                };
2466                dispatch_rows(pool, inter, &run);
2467                true
2468            }
2469            // Q4TiledP gate + Q4TiledP up — the same fused row pass, with
2470            // each row's two ladders built once and spent on both streams.
2471            (
2472                Self::Mapped {
2473                    dtype: TensorDtype::Q4TiledP,
2474                    ..
2475                },
2476                Self::Mapped {
2477                    dtype: TensorDtype::Q4TiledP,
2478                    ..
2479                },
2480            ) => {
2481                let cols = gate.cols();
2482                let gpr = cols / GROUP_SIZE;
2483                let gv_view = Q4tpView::new(gate.quant_bytes(), inter, cols);
2484                let uv_view = Q4tpView::new(up.quant_bytes(), inter, cols);
2485                let run = |start: usize, end: usize| {
2486                    let (mut gsc, mut usc) = (vec![0f32; gpr], vec![0f32; gpr]);
2487                    for r in start..end {
2488                        gv_view.scales_into(r, gpr, &mut gsc);
2489                        uv_view.scales_into(r, gpr, &mut usc);
2490                        let mut gv = dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, &gsc) * act.sx;
2491                        let mut uv = dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, &usc) * act.sx;
2492                        for &(j, xv) in &act.outliers {
2493                            let (w, s) = q4tp_outlier(gv_view.nib, r, gpr, j, &gsc);
2494                            gv += w * s * xv;
2495                            let (w, s) = q4tp_outlier(uv_view.nib, r, gpr, j, &usc);
2496                            uv += w * s * xv;
2497                        }
2498                        let silu_g = gv / (1.0 + (-gv).exp());
2499                        // SAFETY: disjoint row ranges per worker.
2500                        unsafe { *out_addr.at(r) = silu_g * uv };
2501                    }
2502                };
2503                dispatch_rows(pool, inter, &run);
2504                true
2505            }
2506            // Q1 gate + Q1 up — one row pass over both sign streams,
2507            // silu·mul fused (the per-row math of `q1_range_a8w8`); the
2508            // activation group sums are shared by both streams. Without
2509            // this arm a q1 dense FFN paid two dispatches + a combine
2510            // loop — the exact barrier this function exists to remove.
2511            (
2512                Self::Mapped {
2513                    dtype: TensorDtype::Q1,
2514                    ..
2515                },
2516                Self::Mapped {
2517                    dtype: TensorDtype::Q1,
2518                    ..
2519                },
2520            ) => {
2521                let g_bytes = gate.quant_bytes();
2522                let u_bytes = up.quant_bytes();
2523                let gpr = gate.cols() / GROUP_SIZE;
2524                let gsum = q1_group_sums(&act.xq, gpr);
2525                let gsum = &gsum;
2526                let run = move |start: usize, end: usize| {
2527                    for r in start..end {
2528                        let mut gv = dot_q1_row_i8(g_bytes, r, gpr, &act.xq, gsum) * act.sx;
2529                        let mut uv = dot_q1_row_i8(u_bytes, r, gpr, &act.xq, gsum) * act.sx;
2530                        for &(j, xv) in &act.outliers {
2531                            let (w, s) = q1_outlier(g_bytes, r, gpr, j);
2532                            gv += w * s * xv;
2533                            let (w, s) = q1_outlier(u_bytes, r, gpr, j);
2534                            uv += w * s * xv;
2535                        }
2536                        let silu_g = gv / (1.0 + (-gv).exp());
2537                        // SAFETY: disjoint row ranges per worker.
2538                        unsafe { *out_addr.at(r) = silu_g * uv };
2539                    }
2540                };
2541                dispatch_rows(pool, inter, &run);
2542                true
2543            }
2544            // Q2TiledP gate + Q2TiledP up — the 2-bit expert pair (MoE
2545            // FFNs of the W2 class): one row pass, both ladders built
2546            // once, integer code dots with shared group sums.
2547            (
2548                Self::Mapped {
2549                    dtype: TensorDtype::Q2TiledP,
2550                    ..
2551                },
2552                Self::Mapped {
2553                    dtype: TensorDtype::Q2TiledP,
2554                    ..
2555                },
2556            ) => {
2557                let cols = gate.cols();
2558                let gpr = cols / GROUP_SIZE;
2559                let gv_view = Q4tpView::new_q2(gate.quant_bytes(), inter, cols);
2560                let uv_view = Q4tpView::new_q2(up.quant_bytes(), inter, cols);
2561                let gsum = q1_group_sums(&act.xq, gpr);
2562                let gsum = &gsum;
2563                let run = move |start: usize, end: usize| {
2564                    let (mut gsc, mut usc) = (vec![0f32; gpr], vec![0f32; gpr]);
2565                    for r in start..end {
2566                        gv_view.scales_into(r, gpr, &mut gsc);
2567                        uv_view.scales_into(r, gpr, &mut usc);
2568                        let mut gv =
2569                            dot_q2tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsum, &gsc) * act.sx;
2570                        let mut uv =
2571                            dot_q2tp_row_i8(uv_view.nib, r, gpr, &act.xq, gsum, &usc) * act.sx;
2572                        for &(j, xv) in &act.outliers {
2573                            let (w, s) = q2tp_outlier(gv_view.nib, r, gpr, j, &gsc);
2574                            gv += w * s * xv;
2575                            let (w, s) = q2tp_outlier(uv_view.nib, r, gpr, j, &usc);
2576                            uv += w * s * xv;
2577                        }
2578                        let silu_g = gv / (1.0 + (-gv).exp());
2579                        // SAFETY: disjoint row ranges per worker.
2580                        unsafe { *out_addr.at(r) = silu_g * uv };
2581                    }
2582                };
2583                dispatch_rows(pool, inter, &run);
2584                true
2585            }
2586            // Q8Row gate + Q8Row up — one row pass over both i8 streams.
2587            // Q8_2f stays out on purpose: its column field prescales the
2588            // activations PER TENSOR, which breaks this fn's shared
2589            // split_act contract — it keeps the two-dispatch path.
2590            (
2591                Self::Mapped {
2592                    dtype: TensorDtype::Q8Row,
2593                    row_scale: g_rs,
2594                    ..
2595                },
2596                Self::Mapped {
2597                    dtype: TensorDtype::Q8Row,
2598                    row_scale: u_rs,
2599                    ..
2600                },
2601            ) => {
2602                let g_bytes = gate.quant_bytes();
2603                let u_bytes = up.quant_bytes();
2604                let cols = gate.cols();
2605                let run = move |start: usize, end: usize| {
2606                    for r in start..end {
2607                        let gv = q8_row_dot(&g_bytes[r * cols..(r + 1) * cols], act) * g_rs[r];
2608                        let uv = q8_row_dot(&u_bytes[r * cols..(r + 1) * cols], act) * u_rs[r];
2609                        let silu_g = gv / (1.0 + (-gv).exp());
2610                        // SAFETY: disjoint row ranges per worker.
2611                        unsafe { *out_addr.at(r) = silu_g * uv };
2612                    }
2613                };
2614                dispatch_rows(pool, inter, &run);
2615                true
2616            }
2617            // Q1T gate + Q1T up
2618            (
2619                Self::Mapped {
2620                    dtype: TensorDtype::Q1T,
2621                    ..
2622                },
2623                Self::Mapped {
2624                    dtype: TensorDtype::Q1T,
2625                    ..
2626                },
2627            ) => {
2628                const TILE: usize = cortiq_core::quant::Q1T_TILE;
2629                let g_bytes = gate.quant_bytes();
2630                let u_bytes = up.quant_bytes();
2631                let gpr = gate.cols() / GROUP_SIZE;
2632                let (g_rp, g_ent, g_ov) = q1t_overlay(g_bytes, inter * gpr * TILE, inter);
2633                let (u_rp, u_ent, u_ov) = q1t_overlay(u_bytes, inter * gpr * TILE, inter);
2634                let run = move |start: usize, end: usize| {
2635                    for r in start..end {
2636                        let mut gv = q1t_dot_row_i8(g_bytes, r, gpr, &act.xq) * act.sx;
2637                        let mut uv = q1t_dot_row_i8(u_bytes, r, gpr, &act.xq) * act.sx;
2638                        for &(j, xv) in &act.outliers {
2639                            gv += q1t_base_weight(g_bytes, r, gpr, j) * xv;
2640                            uv += q1t_base_weight(u_bytes, r, gpr, j) * xv;
2641                        }
2642                        gv += q1t_row_outlier_correction(g_bytes, r, g_rp, g_ent, g_ov, x_ref);
2643                        uv += q1t_row_outlier_correction(u_bytes, r, u_rp, u_ent, u_ov, x_ref);
2644                        let silu_g = gv / (1.0 + (-gv).exp());
2645                        // SAFETY: disjoint row ranges per worker.
2646                        unsafe { *out_addr.at(r) = silu_g * uv };
2647                    }
2648                };
2649                dispatch_rows(pool, inter, &run);
2650                true
2651            }
2652            _ => false,
2653        }
2654    }
2655
2656    /// Every routed expert's fused gate/up/SiLU under ONE pool dispatch.
2657    ///
2658    /// The per-expert path pays a pool barrier per expert per stage: at 9
2659    /// experts over 40 layers that is ~720 barriers a token, and a decode
2660    /// profile of Qwen3.6-35B-A3B showed the pool parked in
2661    /// `psynch_cvwait` about twice as long as it spent computing. Laying
2662    /// every expert's rows end-to-end in one virtual row space collapses
2663    /// the stage to a single dispatch. The per-row body is the
2664    /// single-expert q4tp arm verbatim, so outputs are bit-identical.
2665    ///
2666    /// `false` = something is outside the fused q4tp kernel (dtype, shape,
2667    /// or the `CMF_SDOT=0` exact contract); the caller walks the ordinary
2668    /// per-expert path.
2669    pub fn moe_gate_up_many(
2670        pairs: &[(&QTensor, &QTensor)],
2671        x: &[f32],
2672        outs: &mut [Vec<f32>],
2673        pool: Option<&Pool>,
2674    ) -> bool {
2675        if pairs.is_empty() || pairs.len() != outs.len() || !a8w8_enabled() {
2676            return false;
2677        }
2678        let inter = pairs[0].0.rows();
2679        let cols = pairs[0].0.cols();
2680        if cols % GROUP_SIZE != 0 {
2681            return false;
2682        }
2683        let gpr = cols / GROUP_SIZE;
2684        // Uniform layout across every routed pair: q4tp, or the 2-bit
2685        // profile's q2tp gate/up (the W2 class). Mixed sets refuse.
2686        let q2 = matches!(
2687            pairs[0].0,
2688            Self::Mapped {
2689                dtype: TensorDtype::Q2TiledP,
2690                ..
2691            }
2692        );
2693        let want = if q2 {
2694            TensorDtype::Q2TiledP
2695        } else {
2696            TensorDtype::Q4TiledP
2697        };
2698        let mut views = Vec::with_capacity(pairs.len() * 2);
2699        for ((g, u), o) in pairs.iter().zip(outs.iter()) {
2700            let both = matches!(g, Self::Mapped { dtype, .. } if *dtype == want)
2701                && matches!(u, Self::Mapped { dtype, .. } if *dtype == want);
2702            if !both
2703                || g.rows() != inter
2704                || u.rows() != inter
2705                || g.cols() != cols
2706                || u.cols() != cols
2707                || o.len() != inter
2708            {
2709                return false;
2710            }
2711            let mk = if q2 { Q4tpView::new_q2 } else { Q4tpView::new };
2712            views.push(mk(g.quant_bytes(), inter, cols));
2713            views.push(mk(u.quant_bytes(), inter, cols));
2714        }
2715        let act = split_act(x);
2716        let gsum = if q2 {
2717            q1_group_sums(&act.xq, gpr)
2718        } else {
2719            Vec::new()
2720        };
2721        let (act, gsum) = (&act, &gsum);
2722        let ptrs: Vec<SendMut> = outs.iter_mut().map(|o| SendMut(o.as_mut_ptr())).collect();
2723        let (views, ptrs) = (&views, &ptrs);
2724        let run = |start: usize, end: usize| {
2725            let (mut gsc, mut usc) = (vec![0f32; gpr], vec![0f32; gpr]);
2726            for flat in start..end {
2727                let (e, r) = (flat / inter, flat % inter);
2728                let gv_view = &views[e * 2];
2729                let uv_view = &views[e * 2 + 1];
2730                gv_view.scales_into(r, gpr, &mut gsc);
2731                uv_view.scales_into(r, gpr, &mut usc);
2732                let (mut gv, mut uv) = if q2 {
2733                    (
2734                        dot_q2tp_row_i8(gv_view.nib, r, gpr, &act.xq, gsum, &gsc) * act.sx,
2735                        dot_q2tp_row_i8(uv_view.nib, r, gpr, &act.xq, gsum, &usc) * act.sx,
2736                    )
2737                } else {
2738                    (
2739                        dot_q4tp_row_i8(gv_view.nib, r, gpr, &act.xq, &gsc) * act.sx,
2740                        dot_q4tp_row_i8(uv_view.nib, r, gpr, &act.xq, &usc) * act.sx,
2741                    )
2742                };
2743                for &(j, xv) in &act.outliers {
2744                    let (og, ou) = if q2 {
2745                        (
2746                            q2tp_outlier(gv_view.nib, r, gpr, j, &gsc),
2747                            q2tp_outlier(uv_view.nib, r, gpr, j, &usc),
2748                        )
2749                    } else {
2750                        (
2751                            q4tp_outlier(gv_view.nib, r, gpr, j, &gsc),
2752                            q4tp_outlier(uv_view.nib, r, gpr, j, &usc),
2753                        )
2754                    };
2755                    gv += og.0 * og.1 * xv;
2756                    uv += ou.0 * ou.1 * xv;
2757                }
2758                let silu_g = gv / (1.0 + (-gv).exp());
2759                // SAFETY: one worker owns each (expert, row) pair.
2760                unsafe { *ptrs[e].at(r) = silu_g * uv };
2761            }
2762        };
2763        dispatch_rows(pool, pairs.len() * inter, &run);
2764        true
2765    }
2766
2767    /// Every routed expert's down projection, weighted and summed into
2768    /// `out`, under ONE pool dispatch.
2769    ///
2770    /// Partitioned by OUTPUT row rather than by expert: each row is owned
2771    /// by a single worker, so the experts are summed in the caller's order
2772    /// — the same sequence of f32 adds the serial `out[i] += w·eo[i]` loop
2773    /// performs, hence bit-identical. Partitioning by expert instead would
2774    /// race on the shared accumulator.
2775    pub fn moe_down_many(
2776        downs: &[&QTensor],
2777        gs: &[Vec<f32>],
2778        weights: &[f32],
2779        out: &mut [f32],
2780        pool: Option<&Pool>,
2781    ) -> bool {
2782        if downs.is_empty()
2783            || downs.len() != gs.len()
2784            || downs.len() != weights.len()
2785            || !a8w8_enabled()
2786        {
2787            return false;
2788        }
2789        let rows = out.len();
2790        let cols = downs[0].cols();
2791        if cols % GROUP_SIZE != 0 {
2792            return false;
2793        }
2794        let gpr = cols / GROUP_SIZE;
2795        let mut views = Vec::with_capacity(downs.len());
2796        for (d, g) in downs.iter().zip(gs.iter()) {
2797            if !matches!(
2798                d,
2799                Self::Mapped {
2800                    dtype: TensorDtype::Q4TiledP,
2801                    ..
2802                }
2803            ) || d.rows() != rows
2804                || d.cols() != cols
2805                || g.len() != cols
2806            {
2807                return false;
2808            }
2809            views.push(Q4tpView::new(d.quant_bytes(), rows, cols));
2810        }
2811        // One int8 split per expert — the activation vectors differ.
2812        let acts: Vec<SplitAct> = gs.iter().map(|g| split_act(g)).collect();
2813        // Partitioned by OUTPUT row, with the experts folded inside: each
2814        // row is owned by one worker, so they are summed in the caller's
2815        // order — the same f32 sequence the serial `out[i] += w·eo[i]`
2816        // loop produces. Partitioning by expert instead would either race
2817        // on the accumulator or need a scratch plane and a second pass;
2818        // measured, that variant was a wash, so this keeps the simpler
2819        // shape.
2820        let out_addr = SendMut(out.as_mut_ptr());
2821        let (views, acts, weights) = (&views, &acts, &weights);
2822        let run = |start: usize, end: usize| {
2823            let mut sc = vec![0f32; gpr];
2824            for r in start..end {
2825                let mut acc = 0f32;
2826                for (e, v) in views.iter().enumerate() {
2827                    v.scales_into(r, gpr, &mut sc);
2828                    let a = &acts[e];
2829                    let mut d = dot_q4tp_row_i8(v.nib, r, gpr, &a.xq, &sc) * a.sx;
2830                    for &(j, xv) in &a.outliers {
2831                        let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
2832                        d += w * s * xv;
2833                    }
2834                    acc += weights[e] * d;
2835                }
2836                // SAFETY: disjoint row ranges per worker.
2837                unsafe { *out_addr.at(r) = acc };
2838            }
2839        };
2840        dispatch_rows(pool, rows, &run);
2841        true
2842    }
2843}
2844
2845/// Batched q8 kernel: same math as qmatvec, the row makes a single
2846/// pass from memory for the whole batch.
2847/// Accelerate CBLAS — the Apple AMX matrix units, the same engine
2848/// llama.cpp's `-ngl 0` prefill rides via ggml-blas.
2849#[cfg(target_os = "macos")]
2850mod accel_blas {
2851    #[link(name = "Accelerate", kind = "framework")]
2852    unsafe extern "C" {
2853        pub fn cblas_sgemm(
2854            order: i32,
2855            trans_a: i32,
2856            trans_b: i32,
2857            m: i32,
2858            n: i32,
2859            k: i32,
2860            alpha: f32,
2861            a: *const f32,
2862            lda: i32,
2863            b: *const f32,
2864            ldb: i32,
2865            beta: f32,
2866            c: *mut f32,
2867            ldc: i32,
2868        );
2869    }
2870}
2871
2872#[cfg(target_os = "macos")]
2873pub(crate) fn accel_gemm_enabled() -> bool {
2874    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
2875    *ON.get_or_init(|| std::env::var("CMF_ACCEL").map(|v| v != "0").unwrap_or(true))
2876}
2877
2878/// Off macOS the "accel" GEMM is the portable NEON micro-kernel below —
2879/// same entry point, so the batched-attention path opens on mobile.
2880#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
2881pub(crate) fn accel_gemm_enabled() -> bool {
2882    true
2883}
2884
2885/// Portable NEON f32 GEMM (row-major, optional Bᵀ): a 4×8 fmla
2886/// micro-kernel with A broadcast against B panels — the mobile stand-in
2887/// for Accelerate in the batched causal attention (QKᵀ and P·V). Not a
2888/// BLAS: shapes here are the attention panels (m ≤ heads·chunk,
2889/// k = head_dim or context), and the goal is removing the per-position
2890/// quadratic wall, not peak GEMM.
2891#[cfg(target_arch = "aarch64")]
2892#[allow(clippy::too_many_arguments)]
2893pub(crate) fn neon_gemm_rm(
2894    m: usize,
2895    n: usize,
2896    k: usize,
2897    alpha: f32,
2898    a: &[f32],
2899    lda: usize,
2900    b_mat: &[f32],
2901    ldb: usize,
2902    b_rows_are_n: bool,
2903    c: &mut [f32],
2904    ldc: usize,
2905) {
2906    debug_assert!(a.len() >= (m - 1) * lda + k);
2907    debug_assert!(c.len() >= (m - 1) * ldc + n);
2908    // SAFETY: bounds asserted above; NEON is baseline on aarch64.
2909    unsafe {
2910        use core::arch::aarch64::*;
2911        let mut i = 0usize;
2912        while i < m {
2913            let mi = (m - i).min(4);
2914            let mut j = 0usize;
2915            while j < n {
2916                let nj = (n - j).min(8);
2917                if mi == 4 && nj == 8 {
2918                    let (mut c0a, mut c0b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
2919                    let (mut c1a, mut c1b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
2920                    let (mut c2a, mut c2b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
2921                    let (mut c3a, mut c3b) = (vdupq_n_f32(0.0), vdupq_n_f32(0.0));
2922                    for p in 0..k {
2923                        let (b0, b1) = if b_rows_are_n {
2924                            // B is [n, k]: column p of Bᵀ = element p of
2925                            // eight consecutive B rows — gathered.
2926                            let base = b_mat.as_ptr().add(j * ldb + p);
2927                            let g = |o: usize| *base.add(o * ldb);
2928                            ([g(0), g(1), g(2), g(3)], [g(4), g(5), g(6), g(7)])
2929                        } else {
2930                            let base = b_mat.as_ptr().add(p * ldb + j);
2931                            (
2932                                [*base, *base.add(1), *base.add(2), *base.add(3)],
2933                                [*base.add(4), *base.add(5), *base.add(6), *base.add(7)],
2934                            )
2935                        };
2936                        let bv0 = vld1q_f32(b0.as_ptr());
2937                        let bv1 = vld1q_f32(b1.as_ptr());
2938                        let a0 = vdupq_n_f32(*a.as_ptr().add(i * lda + p));
2939                        let a1 = vdupq_n_f32(*a.as_ptr().add((i + 1) * lda + p));
2940                        let a2 = vdupq_n_f32(*a.as_ptr().add((i + 2) * lda + p));
2941                        let a3 = vdupq_n_f32(*a.as_ptr().add((i + 3) * lda + p));
2942                        c0a = vfmaq_f32(c0a, a0, bv0);
2943                        c0b = vfmaq_f32(c0b, a0, bv1);
2944                        c1a = vfmaq_f32(c1a, a1, bv0);
2945                        c1b = vfmaq_f32(c1b, a1, bv1);
2946                        c2a = vfmaq_f32(c2a, a2, bv0);
2947                        c2b = vfmaq_f32(c2b, a2, bv1);
2948                        c3a = vfmaq_f32(c3a, a3, bv0);
2949                        c3b = vfmaq_f32(c3b, a3, bv1);
2950                    }
2951                    let al = vdupq_n_f32(alpha);
2952                    for (r, (ca, cb)) in [(c0a, c0b), (c1a, c1b), (c2a, c2b), (c3a, c3b)]
2953                        .iter()
2954                        .enumerate()
2955                    {
2956                        let dst = c.as_mut_ptr().add((i + r) * ldc + j);
2957                        vst1q_f32(dst, vmulq_f32(*ca, al));
2958                        vst1q_f32(dst.add(4), vmulq_f32(*cb, al));
2959                    }
2960                } else {
2961                    for r in 0..mi {
2962                        for q in 0..nj {
2963                            let mut acc = 0f32;
2964                            for p in 0..k {
2965                                let bv = if b_rows_are_n {
2966                                    b_mat[(j + q) * ldb + p]
2967                                } else {
2968                                    b_mat[p * ldb + j + q]
2969                                };
2970                                acc += a[(i + r) * lda + p] * bv;
2971                            }
2972                            c[(i + r) * ldc + j + q] = acc * alpha;
2973                        }
2974                    }
2975                }
2976                j += nj;
2977            }
2978            i += mi;
2979        }
2980    }
2981}
2982
2983/// Off-macOS aarch64: the batched attention rides the NEON micro-GEMM.
2984#[cfg(all(target_arch = "aarch64", not(target_os = "macos")))]
2985#[allow(clippy::too_many_arguments)]
2986pub(crate) fn sgemm_rm(
2987    m: usize,
2988    n: usize,
2989    k: usize,
2990    alpha: f32,
2991    a: &[f32],
2992    lda: usize,
2993    b_mat: &[f32],
2994    ldb: usize,
2995    b_rows_are_n: bool,
2996    c: &mut [f32],
2997    ldc: usize,
2998) {
2999    neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
3000}
3001
3002/// Row-major f32 GEMM, exposed for offline tools (the AWNP pass builds a
3003/// per-layer projection and applies it to every expert; a naive triple loop
3004/// would turn a two-minute job into half an hour).
3005#[allow(clippy::too_many_arguments)]
3006pub fn sgemm_public(
3007    m: usize,
3008    n: usize,
3009    k: usize,
3010    alpha: f32,
3011    a: &[f32],
3012    lda: usize,
3013    b_mat: &[f32],
3014    ldb: usize,
3015    b_rows_are_n: bool,
3016    c: &mut [f32],
3017    ldc: usize,
3018) {
3019    #[cfg(any(target_os = "macos", target_arch = "aarch64"))]
3020    {
3021        sgemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
3022    }
3023    // x86 without Accelerate has no sgemm_rm: the specialized paths there are
3024    // quantized kernels, not an f32 GEMM. Only the offline AWNP pass reaches
3025    // this, so correctness matters and throughput does not — a triple loop is
3026    // the honest fallback rather than a reason to make the tool macOS-only.
3027    #[cfg(not(any(target_os = "macos", target_arch = "aarch64")))]
3028    {
3029        for i in 0..m {
3030            for j in 0..n {
3031                let mut acc = 0f32;
3032                for p in 0..k {
3033                    let bv = if b_rows_are_n {
3034                        b_mat[j * ldb + p]
3035                    } else {
3036                        b_mat[p * ldb + j]
3037                    };
3038                    acc += a[i * lda + p] * bv;
3039                }
3040                c[i * ldc + j] = alpha * acc;
3041            }
3042        }
3043    }
3044}
3045
3046/// Row-major f32 GEMM on Accelerate: C[m,n] = alpha·A[m,k] × B(ᵀ).
3047/// `b_rows_are_n` = true multiplies by Bᵀ where B is stored [n, k].
3048#[cfg(target_os = "macos")]
3049#[allow(clippy::too_many_arguments)]
3050pub(crate) fn sgemm_rm(
3051    m: usize,
3052    n: usize,
3053    k: usize,
3054    alpha: f32,
3055    a: &[f32],
3056    lda: usize,
3057    b_mat: &[f32],
3058    ldb: usize,
3059    b_rows_are_n: bool,
3060    c: &mut [f32],
3061    ldc: usize,
3062) {
3063    debug_assert!(a.len() >= (m - 1) * lda + k);
3064    debug_assert!(c.len() >= (m - 1) * ldc + n);
3065    // Test hook: route the attention GEMMs through the portable NEON
3066    // micro-kernel ON APPLE SILICON — how the mobile batched attend is
3067    // measured without a phone in the loop. (Intel macOS has no NEON —
3068    // the hook is a no-op there, Accelerate continues below.)
3069    #[cfg(target_arch = "aarch64")]
3070    if std::env::var("CMF_FORCE_NEON_GEMM")
3071        .map(|v| v == "1")
3072        .unwrap_or(false)
3073    {
3074        return neon_gemm_rm(m, n, k, alpha, a, lda, b_mat, ldb, b_rows_are_n, c, ldc);
3075    }
3076    unsafe {
3077        accel_blas::cblas_sgemm(
3078            101, // RowMajor
3079            111, // NoTrans A
3080            if b_rows_are_n { 112 } else { 111 },
3081            m as i32,
3082            n as i32,
3083            k as i32,
3084            alpha,
3085            a.as_ptr(),
3086            lda as i32,
3087            b_mat.as_ptr(),
3088            ldb as i32,
3089            0.0,
3090            c.as_mut_ptr(),
3091            ldc as i32,
3092        );
3093    }
3094}
3095
3096/// Prefill GEMM through Accelerate (macOS): dequantize q8 rows into
3097/// f32 tiles (scale folded in, pool-parallel) and multiply each tile
3098/// on the AMX with one row-major sgemm. Tiles live in cache, weights
3099/// stream once. Numerics are f32-GEMM (not the int8 dot): prefill
3100/// logits shift within f32 rounding — tolerance-class, like every
3101/// reduction-order change; decode (M=1) never takes this path.
3102#[cfg(target_os = "macos")]
3103fn qmatmat_accel(
3104    q: &[u8],
3105    row_scale: &[f32],
3106    pre: &[std::borrow::Cow<'_, [f32]>],
3107    rows: usize,
3108    cols: usize,
3109    out: &mut [f32],
3110    pool: Option<&Pool>,
3111) {
3112    // NOTE: double-buffering the dequant against the sgemm (a scoped
3113    // thread driving the pool on tile k+1 while the caller multiplies
3114    // tile k) was tried and LOST ~6%: Accelerate's sgemm is itself
3115    // multithreaded, and the dequant workers just steal its cores.
3116    const TR: usize = 2048;
3117    let b = pre.len();
3118    thread_local! {
3119        static XPANEL: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
3120        static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
3121    }
3122    XPANEL.with(|xp| {
3123        WTILE.with(|wt| {
3124            let mut xpanel = xp.borrow_mut();
3125            xpanel.clear();
3126            for x in pre {
3127                xpanel.extend_from_slice(x);
3128            }
3129            let mut wtile = wt.borrow_mut();
3130            wtile.resize(TR * cols, 0.0);
3131            let mut r0 = 0usize;
3132            while r0 < rows {
3133                let tr = TR.min(rows - r0);
3134                // Dequant the tile (scale folded) — pool-parallel.
3135                let wt_addr = SendMut(wtile.as_mut_ptr());
3136                let run = |start: usize, end: usize| {
3137                    for r in start..end {
3138                        let row = &q[(r0 + r) * cols..(r0 + r + 1) * cols];
3139                        let s = row_scale[r0 + r];
3140                        // SAFETY: workers cover disjoint r ranges.
3141                        let dst =
3142                            unsafe { std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols) };
3143                        for (d, &v) in dst.iter_mut().zip(row) {
3144                            *d = (v as i8) as f32 * s;
3145                        }
3146                    }
3147                };
3148                dispatch_rows(pool, tr, &run);
3149                // C[b, tr] (at column r0 of out[b, rows]) = X · Wtileᵀ
3150                unsafe {
3151                    accel_blas::cblas_sgemm(
3152                        101, // RowMajor
3153                        111, // NoTrans A
3154                        112, // Trans B
3155                        b as i32,
3156                        tr as i32,
3157                        cols as i32,
3158                        1.0,
3159                        xpanel.as_ptr(),
3160                        cols as i32,
3161                        wtile.as_ptr(),
3162                        cols as i32,
3163                        0.0,
3164                        out.as_mut_ptr().add(r0),
3165                        rows as i32,
3166                    );
3167                }
3168                r0 += tr;
3169            }
3170        })
3171    });
3172}
3173
3174fn qmatmat(
3175    q: &[u8],
3176    row_scale: &[f32],
3177    pre: &[std::borrow::Cow<'_, [f32]>],
3178    rows: usize,
3179    cols: usize,
3180    out: &mut [f32],
3181    pool: Option<&Pool>,
3182) {
3183    let b = pre.len();
3184    debug_assert_eq!(out.len(), b * rows);
3185    // Big prefill batches ride the AMX (roadmap PR3): the row×batch
3186    // SDOT loop below peaks near the CPU's dot throughput, an order
3187    // below the matrix units. Small tensors and tiny test models stay
3188    // on the exact integer path.
3189    #[cfg(target_os = "macos")]
3190    if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
3191        qmatmat_accel(q, row_scale, pre, rows, cols, out, pool);
3192        return;
3193    }
3194    #[cfg(target_arch = "aarch64")]
3195    if sdot_enabled() {
3196        let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
3197        let out_addr = SendMut(out.as_mut_ptr());
3198        // Blocked 2×4 (mobile prefill: no AMX to fall back on — this
3199        // path IS the ARM prefill GEMM off Apple silicon).
3200        let blocked_ok = blocked_enabled();
3201        let use_i8mm = i8mm_enabled();
3202        if blocked_ok {
3203            let run = |start: usize, end: usize| {
3204                let mut o = start;
3205                while o < end {
3206                    if o + 2 <= end {
3207                        let r0 = &q[o * cols..(o + 1) * cols];
3208                        let r1 = &q[(o + 1) * cols..(o + 2) * cols];
3209                        let mut bi = 0usize;
3210                        while bi + 4 <= acts.len() {
3211                            let xs = [
3212                                acts[bi].xq.as_slice(),
3213                                acts[bi + 1].xq.as_slice(),
3214                                acts[bi + 2].xq.as_slice(),
3215                                acts[bi + 3].xq.as_slice(),
3216                            ];
3217                            let d = if use_i8mm {
3218                                unsafe { dot_i8_smmla_2x4(r0, r1, xs) }
3219                            } else {
3220                                unsafe { dot_i8_sdot_2x4(r0, r1, xs) }
3221                            };
3222                            for (r, row) in [r0, r1].into_iter().enumerate() {
3223                                for k in 0..4 {
3224                                    let act = &acts[bi + k];
3225                                    let mut v = d[r][k] as f32 * act.sx;
3226                                    for &(j, xv) in &act.outliers {
3227                                        v += (row[j] as i8) as f32 * xv;
3228                                    }
3229                                    unsafe {
3230                                        *out_addr.at((bi + k) * rows + o + r) = v * row_scale[o + r]
3231                                    };
3232                                }
3233                            }
3234                            bi += 4;
3235                        }
3236                        while bi < acts.len() {
3237                            for (r, row) in [r0, r1].into_iter().enumerate() {
3238                                let v = row_dot_sdot(row, &acts[bi]) * row_scale[o + r];
3239                                unsafe { *out_addr.at(bi * rows + o + r) = v };
3240                            }
3241                            bi += 1;
3242                        }
3243                        o += 2;
3244                    } else {
3245                        let row = &q[o * cols..(o + 1) * cols];
3246                        for (bi, act) in acts.iter().enumerate() {
3247                            let v = row_dot_sdot(row, act) * row_scale[o];
3248                            unsafe { *out_addr.at(bi * rows + o) = v };
3249                        }
3250                        o += 1;
3251                    }
3252                }
3253            };
3254            dispatch_rows(pool, rows, &run);
3255            return;
3256        }
3257        let run = |start: usize, end: usize| {
3258            for o in start..end {
3259                let row = &q[o * cols..(o + 1) * cols];
3260                for (bi, act) in acts.iter().enumerate() {
3261                    let v = row_dot_sdot(row, act) * row_scale[o];
3262                    unsafe { *out_addr.at(bi * rows + o) = v };
3263                }
3264            }
3265        };
3266        dispatch_rows(pool, rows, &run);
3267        return;
3268    }
3269    // x86 A8W8 batch. Non-VNNI parts take the BLOCKED 2×4 kernel
3270    // (roadmap P0: two weight rows' abs() stay in registers across four
3271    // activation streams); VNNI machines keep the per-row bias-trick
3272    // dot, which is already throughput-bound there.
3273    #[cfg(target_arch = "x86_64")]
3274    if avx2_a8w8_enabled() {
3275        let acts: Vec<SplitAct> = pre.iter().map(|x| split_act(x)).collect();
3276        let out_addr = SendMut(out.as_mut_ptr());
3277        // CMF_X86_BLOCKED=0 forces the per-row path (paired in-process
3278        // A/B on noisy shared-vCPU hosts).
3279        let blocked_ok = blocked_enabled();
3280        if !avx512vnni_enabled() && blocked_ok {
3281            let run = |start: usize, end: usize| {
3282                let mut o = start;
3283                while o < end {
3284                    if o + 2 <= end {
3285                        let r0 = &q[o * cols..(o + 1) * cols];
3286                        let r1 = &q[(o + 1) * cols..(o + 2) * cols];
3287                        let mut bi = 0usize;
3288                        while bi + 4 <= acts.len() {
3289                            let xs = [
3290                                acts[bi].xq.as_slice(),
3291                                acts[bi + 1].xq.as_slice(),
3292                                acts[bi + 2].xq.as_slice(),
3293                                acts[bi + 3].xq.as_slice(),
3294                            ];
3295                            let d = unsafe { dot_i8_i8_avx2_2x4(r0, r1, xs) };
3296                            for (r, row) in [r0, r1].into_iter().enumerate() {
3297                                for k in 0..4 {
3298                                    let act = &acts[bi + k];
3299                                    let mut v = d[r][k] as f32 * act.sx;
3300                                    for &(j, xv) in &act.outliers {
3301                                        v += (row[j] as i8) as f32 * xv;
3302                                    }
3303                                    unsafe {
3304                                        *out_addr.at((bi + k) * rows + o + r) = v * row_scale[o + r]
3305                                    };
3306                                }
3307                            }
3308                            bi += 4;
3309                        }
3310                        while bi < acts.len() {
3311                            for (r, row) in [r0, r1].into_iter().enumerate() {
3312                                let v = row_dot_avx2(row, &acts[bi]) * row_scale[o + r];
3313                                unsafe { *out_addr.at(bi * rows + o + r) = v };
3314                            }
3315                            bi += 1;
3316                        }
3317                        o += 2;
3318                    } else {
3319                        let row = &q[o * cols..(o + 1) * cols];
3320                        for (bi, act) in acts.iter().enumerate() {
3321                            let v = row_dot_avx2(row, act) * row_scale[o];
3322                            unsafe { *out_addr.at(bi * rows + o) = v };
3323                        }
3324                        o += 1;
3325                    }
3326                }
3327            };
3328            dispatch_rows(pool, rows, &run);
3329            return;
3330        }
3331        let run = |start: usize, end: usize| {
3332            for o in start..end {
3333                let row = &q[o * cols..(o + 1) * cols];
3334                for (bi, act) in acts.iter().enumerate() {
3335                    let v = row_dot_avx2(row, act) * row_scale[o];
3336                    unsafe { *out_addr.at(bi * rows + o) = v };
3337                }
3338            }
3339        };
3340        dispatch_rows(pool, rows, &run);
3341        return;
3342    }
3343    let out_addr = SendMut(out.as_mut_ptr());
3344    let run = |start: usize, end: usize| {
3345        for o in start..end {
3346            let row = &q[o * cols..(o + 1) * cols];
3347            for (bi, x) in pre.iter().enumerate() {
3348                let mut acc = 0f32;
3349                for j in 0..cols {
3350                    acc += (row[j] as i8) as f32 * x[j];
3351                }
3352                unsafe { *out_addr.at(bi * rows + o) = acc * row_scale[o] };
3353            }
3354        }
3355    };
3356    dispatch_rows(pool, rows, &run);
3357}
3358
3359/// Split rows across pool workers (shared qmatvec pattern). Self-balancing
3360/// — see `Pool::run_rows` for why a static 1/n split is wrong here.
3361fn dispatch_rows(pool: Option<&Pool>, rows: usize, run: &(dyn Fn(usize, usize) + Sync)) {
3362    match pool {
3363        Some(pool) if rows >= 256 => pool.run_rows(rows, run),
3364        _ => run(0, rows),
3365    }
3366}
3367
3368/// Split a q4_block blob into (packed nibbles, f16 group scales).
3369fn q4_split(bytes: &[u8], rows: usize, cols: usize) -> (&[u8], &[u8]) {
3370    let groups = rows * cols / GROUP_SIZE;
3371    bytes.split_at(groups * 16)
3372}
3373
3374/// SIMD unpack for the dominant vbit width B=4 (94% of rows on the
3375/// log2-shape calibration): 16 packed bytes -> 32 centered i8 values.
3376/// vbit packs MSB-first, so the HIGH nibble is the even element
3377/// (opposite of q4_block's lo-first interleave). Centering is u-7.
3378#[inline]
3379fn vbit_fill4(data: &[u8], buf: &mut [u8]) {
3380    #[cfg(target_arch = "aarch64")]
3381    unsafe {
3382        return vbit_fill4_neon(data, buf);
3383    }
3384    #[cfg(target_arch = "x86_64")]
3385    if avx2_enabled() {
3386        return unsafe { vbit_fill4_avx2(data, buf) };
3387    }
3388    #[allow(unreachable_code)]
3389    for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
3390        let u = unpack8::<4>(&data[blk * 4..]);
3391        for k in 0..8 {
3392            chunk[k] = (u[k] - 7) as i8 as u8;
3393        }
3394    }
3395}
3396
3397#[cfg(target_arch = "aarch64")]
3398#[target_feature(enable = "neon")]
3399unsafe fn vbit_fill4_neon(data: &[u8], buf: &mut [u8]) {
3400    // SAFETY: buf.len() is a multiple of GROUP_SIZE=32; data holds
3401    // buf.len()/2 packed bytes (validated at load).
3402    unsafe {
3403        use core::arch::aarch64::*;
3404        let n = buf.len();
3405        let mask = vdupq_n_u8(0x0F);
3406        let seven = vdupq_n_s8(7);
3407        let mut g = 0usize;
3408        while g * 32 + 32 <= n {
3409            let b = vld1q_u8(data.as_ptr().add(g * 16));
3410            let hi = vshrq_n_u8::<4>(b);
3411            let lo = vandq_u8(b, mask);
3412            let z0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(hi, lo)), seven);
3413            let z1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(hi, lo)), seven);
3414            vst1q_u8(buf.as_mut_ptr().add(g * 32), vreinterpretq_u8_s8(z0));
3415            vst1q_u8(buf.as_mut_ptr().add(g * 32 + 16), vreinterpretq_u8_s8(z1));
3416            g += 1;
3417        }
3418    }
3419}
3420
3421#[cfg(target_arch = "x86_64")]
3422#[target_feature(enable = "avx2")]
3423unsafe fn vbit_fill4_avx2(data: &[u8], buf: &mut [u8]) {
3424    // SAFETY: see vbit_fill4_neon.
3425    unsafe {
3426        use core::arch::x86_64::*;
3427        let n = buf.len();
3428        let mask = _mm_set1_epi8(0x0F);
3429        let seven = _mm256_set1_epi8(7);
3430        let mut g = 0usize;
3431        while g * 32 + 32 <= n {
3432            let b = _mm_loadu_si128(data.as_ptr().add(g * 16) as *const __m128i);
3433            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), mask);
3434            let lo = _mm_and_si128(b, mask);
3435            let z = _mm256_sub_epi8(
3436                _mm256_set_m128i(_mm_unpackhi_epi8(hi, lo), _mm_unpacklo_epi8(hi, lo)),
3437                seven,
3438            );
3439            _mm256_storeu_si256(buf.as_mut_ptr().add(g * 32) as *mut __m256i, z);
3440            g += 1;
3441        }
3442    }
3443}
3444
3445/// Unpack 8 MSB-first B-bit values from exactly B bytes (fixed shifts —
3446/// no serial bit-buffer, auto-vectorizable). Every 32-value group starts
3447/// byte-aligned (32·B/8 is integral for B∈3..8), so groups decompose
3448/// into 4 such blocks.
3449#[inline(always)]
3450fn unpack8<const B: usize>(data: &[u8]) -> [i32; 8] {
3451    let mut acc = 0u64;
3452    for i in 0..B {
3453        acc = (acc << 8) | data[i] as u64;
3454    }
3455    let mask = (1u64 << B) - 1;
3456    let mut out = [0i32; 8];
3457    for (k, o) in out.iter_mut().enumerate() {
3458        *o = ((acc >> ((7 - k) * B)) & mask) as i32;
3459    }
3460    out
3461}
3462
3463/// Fused vbit matvec straight from the mapped bytes (spec §3, P13
3464/// FIG.3): [u8 bits: rows][f16 scales: rows·cols/32][bit-packed rows,
3465/// MSB-first, byte-padded]. Row data offsets are precomputed at load
3466/// (`vbit_row_offsets`) — the per-call prefix scan was O(rows) pure
3467/// overhead on every matvec.
3468#[allow(clippy::too_many_arguments)]
3469fn vbitmatvec(
3470    bytes: &[u8],
3471    offsets: &[usize],
3472    x: &[f32],
3473    rows: usize,
3474    cols: usize,
3475    out: &mut [f32],
3476    pool: Option<&Pool>,
3477) {
3478    debug_assert_eq!(out.len(), rows);
3479    debug_assert_eq!(offsets.len(), rows + 1);
3480
3481    // SDOT path: unpack the row to centered i8 once, then per-group
3482    // int8 dot against the quantized activations — same A8W8 contract
3483    // as q8 (bounded noise; CMF_SDOT=0 keeps the exact scalar path).
3484    if a8w8_enabled() {
3485        let act = split_act(x);
3486        let out_addr = SendMut(out.as_mut_ptr());
3487        let run = move |start: usize, end: usize| {
3488            vbit_range_a8w8(bytes, offsets, x, &act, rows, cols, out_addr, start, end)
3489        };
3490        dispatch_rows(pool, rows, &run);
3491        return;
3492    }
3493
3494    let out_addr = SendMut(out.as_mut_ptr());
3495    let run = move |start: usize, end: usize| {
3496        vbit_range_f32(bytes, offsets, x, rows, cols, out_addr, start, end)
3497    };
3498    dispatch_rows(pool, rows, &run);
3499}
3500
3501/// One vbit row range via the A8W8 int8 path — kernel body of
3502/// `vbitmatvec`, extracted so multi-matrix jobs can drive it for
3503/// several tensors in one dispatch (b=8 rows go exact f32).
3504#[allow(clippy::too_many_arguments)]
3505fn vbit_range_a8w8(
3506    bytes: &[u8],
3507    offsets: &[usize],
3508    x: &[f32],
3509    act: &SplitAct,
3510    rows: usize,
3511    cols: usize,
3512    out: SendMut,
3513    start: usize,
3514    end: usize,
3515) {
3516    let ng = cols / GROUP_SIZE;
3517    let bits = &bytes[..rows];
3518    let sc_off = rows;
3519    let row_dot = |r: usize| -> f32 {
3520        let b = bits[r] as usize;
3521        let l = (1i32 << (b - 1)) - 1;
3522        let mask = (1u64 << b) - 1;
3523        let data = &bytes[offsets[r]..offsets[r + 1]];
3524        if b == 8 {
3525            // u−L reaches 128 → does not fit i8; exact f32 path.
3526            let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
3527            let mut dot = 0f32;
3528            for g in 0..ng {
3529                let so = (r * ng + g) * 2;
3530                let sgf = f16_to_f32(u16::from_le_bytes([
3531                    bytes[sc_off + so],
3532                    bytes[sc_off + so + 1],
3533                ]));
3534                let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
3535                let mut gd = 0f32;
3536                for &xv in xg.iter() {
3537                    if nbits < 8 {
3538                        acc = (acc << 8) | data[idx] as u64;
3539                        idx += 1;
3540                        nbits += 8;
3541                    }
3542                    let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
3543                    nbits -= 8;
3544                    gd += (u - l) as f32 * xv;
3545                }
3546                dot += gd * sgf;
3547            }
3548            return dot;
3549        }
3550        // Per-worker scratch: this closure runs for every row of the
3551        // tensor (lm_head ≈ 150k rows/token) — a heap allocation per
3552        // row was measurable pure overhead.
3553        thread_local! {
3554            static VBIT_SCRATCH: std::cell::RefCell<Vec<u8>> =
3555                const { std::cell::RefCell::new(Vec::new()) };
3556        }
3557        #[inline(always)]
3558        fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
3559            for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
3560                let u = unpack8::<B>(&data[blk * B..]);
3561                for k in 0..8 {
3562                    chunk[k] = (u[k] - l) as i8 as u8;
3563                }
3564            }
3565        }
3566        let _ = mask;
3567        VBIT_SCRATCH.with(|scratch| {
3568            let mut buf = scratch.borrow_mut();
3569            buf.resize(cols, 0);
3570            match b {
3571                3 => fill::<3>(data, l, &mut buf),
3572                4 => vbit_fill4(data, &mut buf),
3573                5 => fill::<5>(data, l, &mut buf),
3574                6 => fill::<6>(data, l, &mut buf),
3575                _ => unreachable!(),
3576            }
3577            let mut dot = 0f32;
3578            for g in 0..ng {
3579                let so = (r * ng + g) * 2;
3580                let s = f16_to_f32(u16::from_le_bytes([
3581                    bytes[sc_off + so],
3582                    bytes[sc_off + so + 1],
3583                ]));
3584                let d = dot_i8_i8(
3585                    &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
3586                    &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
3587                ) as f32
3588                    * act.sx;
3589                dot += d * s;
3590            }
3591            for &(j, xv) in &act.outliers {
3592                let so = (r * ng + j / GROUP_SIZE) * 2;
3593                let s = f16_to_f32(u16::from_le_bytes([
3594                    bytes[sc_off + so],
3595                    bytes[sc_off + so + 1],
3596                ]));
3597                // xq is zeroed at outlier slots — add the exact term.
3598                dot += (buf[j] as i8) as f32 * s * xv;
3599            }
3600            dot
3601        })
3602    };
3603    for r in start..end {
3604        // SAFETY: disjoint row ranges per worker.
3605        unsafe { *out.at(r) = row_dot(r) };
3606    }
3607}
3608
3609/// Exact scalar vbit row range (same extraction, non-SDOT path).
3610#[allow(clippy::too_many_arguments)]
3611fn vbit_range_f32(
3612    bytes: &[u8],
3613    offsets: &[usize],
3614    x: &[f32],
3615    rows: usize,
3616    cols: usize,
3617    out: SendMut,
3618    start: usize,
3619    end: usize,
3620) {
3621    let ng = cols / GROUP_SIZE;
3622    let bits = &bytes[..rows];
3623    let sc_off = rows;
3624    // Per-bit-width specialized inner loops: the compiler unrolls the
3625    // constant shifts (the generic bit-buffer loop was branch-bound —
3626    // 5.6 vs 13.2 tok/s q4 on the 0.8B).
3627    #[inline(always)]
3628    fn dot_row<const B: usize>(
3629        data: &[u8],
3630        bytes: &[u8],
3631        sc_off: usize,
3632        r: usize,
3633        ng: usize,
3634        x: &[f32],
3635    ) -> f32 {
3636        let l = ((1i32 << (B - 1)) - 1) as f32;
3637        let gbytes = GROUP_SIZE * B / 8;
3638        let mut dot = 0f32;
3639        for g in 0..ng {
3640            let so = (r * ng + g) * 2;
3641            let s = f16_to_f32(u16::from_le_bytes([
3642                bytes[sc_off + so],
3643                bytes[sc_off + so + 1],
3644            ]));
3645            let xg = &x[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
3646            let gd0 = &data[g * gbytes..(g + 1) * gbytes];
3647            let mut gd = 0f32;
3648            for blk in 0..GROUP_SIZE / 8 {
3649                let u = unpack8::<B>(&gd0[blk * B..]);
3650                let xb = &xg[blk * 8..blk * 8 + 8];
3651                for k in 0..8 {
3652                    gd += (u[k] as f32 - l) * xb[k];
3653                }
3654            }
3655            dot += gd * s;
3656        }
3657        dot
3658    }
3659    for r in start..end {
3660        let data = &bytes[offsets[r]..offsets[r + 1]];
3661        let v = match bits[r] {
3662            3 => dot_row::<3>(data, bytes, sc_off, r, ng, x),
3663            4 => dot_row::<4>(data, bytes, sc_off, r, ng, x),
3664            5 => dot_row::<5>(data, bytes, sc_off, r, ng, x),
3665            6 => dot_row::<6>(data, bytes, sc_off, r, ng, x),
3666            8 => dot_row::<8>(data, bytes, sc_off, r, ng, x),
3667            b => unreachable!("vbit bit-width {b} (validated at load)"),
3668        };
3669        // SAFETY: disjoint row ranges per worker.
3670        unsafe { *out.at(r) = v };
3671    }
3672}
3673
3674/// Fused two-input vbit matvec: each row is unpacked from the mmap ONCE
3675/// and dotted against BOTH activations (MTP verify / pair prefill used
3676/// to run two full matvecs — double weight traffic and double unpack).
3677/// Per-input math is identical to `vbitmatvec` → same accuracy contract.
3678#[allow(clippy::too_many_arguments)]
3679fn vbitmatvec2(
3680    bytes: &[u8],
3681    offsets: &[usize],
3682    x1: &[f32],
3683    x2: &[f32],
3684    rows: usize,
3685    cols: usize,
3686    o1: &mut [f32],
3687    o2: &mut [f32],
3688    pool: Option<&Pool>,
3689) {
3690    debug_assert_eq!(o1.len(), rows);
3691    debug_assert_eq!(o2.len(), rows);
3692
3693    if a8w8_enabled() {
3694        let a1 = split_act(x1);
3695        let a2 = split_act(x2);
3696        let p1 = SendMut(o1.as_mut_ptr());
3697        let p2 = SendMut(o2.as_mut_ptr());
3698        let run = move |start: usize, end: usize| {
3699            vbit_range2_a8w8(
3700                bytes, offsets, x1, x2, &a1, &a2, rows, cols, p1, p2, start, end,
3701            )
3702        };
3703        dispatch_rows(pool, rows, &run);
3704        return;
3705    }
3706
3707    let p1 = SendMut(o1.as_mut_ptr());
3708    let p2 = SendMut(o2.as_mut_ptr());
3709    let run = move |start: usize, end: usize| {
3710        vbit_range2_f32(bytes, offsets, x1, x2, rows, cols, p1, p2, start, end)
3711    };
3712    dispatch_rows(pool, rows, &run);
3713}
3714
3715/// Two-input vbit row range via the A8W8 int8 path — kernel body of
3716/// `vbitmatvec2`, extracted for pair multi-matrix jobs (b=8 rows go
3717/// exact f32 for both lanes, bits streamed once).
3718#[allow(clippy::too_many_arguments)]
3719fn vbit_range2_a8w8(
3720    bytes: &[u8],
3721    offsets: &[usize],
3722    x1: &[f32],
3723    x2: &[f32],
3724    a1: &SplitAct,
3725    a2: &SplitAct,
3726    rows: usize,
3727    cols: usize,
3728    p1: SendMut,
3729    p2: SendMut,
3730    start: usize,
3731    end: usize,
3732) {
3733    let ng = cols / GROUP_SIZE;
3734    let bits = &bytes[..rows];
3735    let sc_off = rows;
3736    let row_dots = |r: usize| -> (f32, f32) {
3737        let b = bits[r] as usize;
3738        let l = (1i32 << (b - 1)) - 1;
3739        let data = &bytes[offsets[r]..offsets[r + 1]];
3740        if b == 8 {
3741            // u−L reaches 128 → does not fit i8; exact f32 path,
3742            // bits still streamed once for both lanes.
3743            let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
3744            let (mut d1, mut d2) = (0f32, 0f32);
3745            for g in 0..ng {
3746                let so = (r * ng + g) * 2;
3747                let sgf = f16_to_f32(u16::from_le_bytes([
3748                    bytes[sc_off + so],
3749                    bytes[sc_off + so + 1],
3750                ]));
3751                let (mut g1, mut g2) = (0f32, 0f32);
3752                for k in 0..GROUP_SIZE {
3753                    if nbits < 8 {
3754                        acc = (acc << 8) | data[idx] as u64;
3755                        idx += 1;
3756                        nbits += 8;
3757                    }
3758                    let u = ((acc >> (nbits - 8)) & 0xFF) as i32;
3759                    nbits -= 8;
3760                    let w = (u - l) as f32;
3761                    g1 += w * x1[g * GROUP_SIZE + k];
3762                    g2 += w * x2[g * GROUP_SIZE + k];
3763                }
3764                d1 += g1 * sgf;
3765                d2 += g2 * sgf;
3766            }
3767            return (d1, d2);
3768        }
3769        thread_local! {
3770            static VBIT_SCRATCH2: std::cell::RefCell<Vec<u8>> =
3771                const { std::cell::RefCell::new(Vec::new()) };
3772        }
3773        #[inline(always)]
3774        fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
3775            for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
3776                let u = unpack8::<B>(&data[blk * B..]);
3777                for k in 0..8 {
3778                    chunk[k] = (u[k] - l) as i8 as u8;
3779                }
3780            }
3781        }
3782        VBIT_SCRATCH2.with(|scratch| {
3783            let mut buf = scratch.borrow_mut();
3784            buf.resize(cols, 0);
3785            match b {
3786                3 => fill::<3>(data, l, &mut buf),
3787                4 => vbit_fill4(data, &mut buf),
3788                5 => fill::<5>(data, l, &mut buf),
3789                6 => fill::<6>(data, l, &mut buf),
3790                _ => unreachable!(),
3791            }
3792            let (mut d1, mut d2) = (0f32, 0f32);
3793            for g in 0..ng {
3794                let so = (r * ng + g) * 2;
3795                let s = f16_to_f32(u16::from_le_bytes([
3796                    bytes[sc_off + so],
3797                    bytes[sc_off + so + 1],
3798                ]));
3799                let wg = &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
3800                let v1 = dot_i8_i8(wg, &a1.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a1.sx;
3801                let v2 = dot_i8_i8(wg, &a2.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE]) as f32 * a2.sx;
3802                d1 += v1 * s;
3803                d2 += v2 * s;
3804            }
3805            for &(j, xv) in &a1.outliers {
3806                let so = (r * ng + j / GROUP_SIZE) * 2;
3807                let s = f16_to_f32(u16::from_le_bytes([
3808                    bytes[sc_off + so],
3809                    bytes[sc_off + so + 1],
3810                ]));
3811                d1 += (buf[j] as i8) as f32 * s * xv;
3812            }
3813            for &(j, xv) in &a2.outliers {
3814                let so = (r * ng + j / GROUP_SIZE) * 2;
3815                let s = f16_to_f32(u16::from_le_bytes([
3816                    bytes[sc_off + so],
3817                    bytes[sc_off + so + 1],
3818                ]));
3819                d2 += (buf[j] as i8) as f32 * s * xv;
3820            }
3821            (d1, d2)
3822        })
3823    };
3824    for r in start..end {
3825        let (v1, v2) = row_dots(r);
3826        // SAFETY: disjoint row ranges per worker.
3827        unsafe {
3828            *p1.at(r) = v1;
3829            *p2.at(r) = v2;
3830        }
3831    }
3832}
3833
3834/// Two-input exact scalar vbit row range (same extraction) —
3835/// per-bit-width specialized, two accumulators per row; per-lane
3836/// accumulation order matches `vbitmatvec` exactly.
3837#[allow(clippy::too_many_arguments)]
3838fn vbit_range2_f32(
3839    bytes: &[u8],
3840    offsets: &[usize],
3841    x1: &[f32],
3842    x2: &[f32],
3843    rows: usize,
3844    cols: usize,
3845    p1: SendMut,
3846    p2: SendMut,
3847    start: usize,
3848    end: usize,
3849) {
3850    let ng = cols / GROUP_SIZE;
3851    let bits = &bytes[..rows];
3852    let sc_off = rows;
3853    #[inline(always)]
3854    #[allow(clippy::too_many_arguments)]
3855    fn dot_row2<const B: usize>(
3856        data: &[u8],
3857        bytes: &[u8],
3858        sc_off: usize,
3859        r: usize,
3860        ng: usize,
3861        x1: &[f32],
3862        x2: &[f32],
3863    ) -> (f32, f32) {
3864        let l = ((1i32 << (B - 1)) - 1) as f32;
3865        let gbytes = GROUP_SIZE * B / 8;
3866        let (mut d1, mut d2) = (0f32, 0f32);
3867        for g in 0..ng {
3868            let so = (r * ng + g) * 2;
3869            let s = f16_to_f32(u16::from_le_bytes([
3870                bytes[sc_off + so],
3871                bytes[sc_off + so + 1],
3872            ]));
3873            let x1g = &x1[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
3874            let x2g = &x2[g * GROUP_SIZE..(g + 1) * GROUP_SIZE];
3875            let gd0 = &data[g * gbytes..(g + 1) * gbytes];
3876            let (mut g1, mut g2) = (0f32, 0f32);
3877            for blk in 0..GROUP_SIZE / 8 {
3878                let u = unpack8::<B>(&gd0[blk * B..]);
3879                for k in 0..8 {
3880                    let w = u[k] as f32 - l;
3881                    g1 += w * x1g[blk * 8 + k];
3882                    g2 += w * x2g[blk * 8 + k];
3883                }
3884            }
3885            d1 += g1 * s;
3886            d2 += g2 * s;
3887        }
3888        (d1, d2)
3889    }
3890    for r in start..end {
3891        let data = &bytes[offsets[r]..offsets[r + 1]];
3892        let (v1, v2) = match bits[r] {
3893            3 => dot_row2::<3>(data, bytes, sc_off, r, ng, x1, x2),
3894            4 => dot_row2::<4>(data, bytes, sc_off, r, ng, x1, x2),
3895            5 => dot_row2::<5>(data, bytes, sc_off, r, ng, x1, x2),
3896            6 => dot_row2::<6>(data, bytes, sc_off, r, ng, x1, x2),
3897            8 => dot_row2::<8>(data, bytes, sc_off, r, ng, x1, x2),
3898            b => unreachable!("vbit bit-width {b} (validated at load)"),
3899        };
3900        // SAFETY: disjoint row ranges per worker.
3901        unsafe {
3902            *p1.at(r) = v1;
3903            *p2.at(r) = v2;
3904        }
3905    }
3906}
3907
3908// ───────────────────── q4_tiled kernels (§4.3) ─────────────────────
3909
3910/// One q4_tiled row dot on the A8W8 int8 path: per 32-group the tile
3911/// is ONE sequential read — [f16 scale][16B nibbles] — versus the two
3912/// distant streams of the split layout. Values/order identical to the
3913/// split kernels.
3914#[inline]
3915#[allow(unreachable_code)]
3916fn dot_q4t_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
3917    #[cfg(target_arch = "aarch64")]
3918    unsafe {
3919        return dot_q4t_row_sdot(bytes, r, gpr, xq);
3920    }
3921    #[cfg(target_arch = "x86_64")]
3922    unsafe {
3923        if vnni_tiles_enabled() {
3924            return dot_q4t_row_vnni(bytes, r, gpr, xq);
3925        }
3926        return dot_q4t_row_avx2(bytes, r, gpr, xq);
3927    }
3928    let mut acc = 0f32;
3929    for gi in 0..gpr {
3930        let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
3931        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
3932        let mut d = 0i32;
3933        for (k, &b) in tile[2..].iter().enumerate() {
3934            d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
3935                + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
3936        }
3937        acc += d as f32 * s;
3938    }
3939    acc
3940}
3941
3942#[cfg(target_arch = "aarch64")]
3943#[target_feature(enable = "neon,dotprod")]
3944unsafe fn dot_q4t_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
3945    // SAFETY: callers uphold slice-length contracts (18B tile per group,
3946    // xq.len() == gpr·GROUP_SIZE).
3947    unsafe {
3948        use core::arch::aarch64::*;
3949        use core::arch::asm;
3950        let lomask = vdupq_n_u8(0x0F);
3951        let eight = vdupq_n_s8(8);
3952        let mut acc = 0f32;
3953        for gi in 0..gpr {
3954            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
3955            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
3956            let b = vld1q_u8(t.add(2));
3957            let lo = vandq_u8(b, lomask);
3958            let hi = vshrq_n_u8::<4>(b);
3959            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
3960            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
3961            let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
3962            let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
3963            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
3964            asm!(
3965                "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
3966                "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
3967                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
3968                e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
3969                options(pure, nomem, nostack),
3970            );
3971            acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
3972        }
3973        acc
3974    }
3975}
3976
3977#[cfg(target_arch = "x86_64")]
3978#[target_feature(enable = "avx2")]
3979unsafe fn dot_q4t_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
3980    // SAFETY: see dot_q4t_row_sdot.
3981    unsafe {
3982        use core::arch::x86_64::*;
3983        let lomask = _mm_set1_epi8(0x0F);
3984        let eight = _mm256_set1_epi8(8);
3985        let ones = _mm256_set1_epi16(1);
3986        let mut acc = 0f32;
3987        for gi in 0..gpr {
3988            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
3989            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
3990            let b = _mm_loadu_si128(t.add(2) as *const __m128i);
3991            let lo = _mm_and_si128(b, lomask);
3992            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
3993            let w = _mm256_sub_epi8(
3994                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
3995                eight,
3996            );
3997            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
3998            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
3999            let d = _mm256_madd_epi16(p16, ones);
4000            let hi128 = _mm256_extracti128_si256::<1>(d);
4001            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
4002            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
4003            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
4004            acc += _mm_cvtsi128_si32(s32) as f32 * s;
4005        }
4006        acc
4007    }
4008}
4009
4010/// VNNI twin of `dot_q4t_row_avx2`: same unpack, `vpdpbusd` replaces
4011/// the maddubs+madd pair (see `dpbusd_hsum` — sums are bit-identical).
4012/// 256-bit VL encoding, so the VEX `vpsignb` stays usable.
4013#[cfg(target_arch = "x86_64")]
4014#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
4015unsafe fn dot_q4t_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
4016    // SAFETY: see dot_q4t_row_sdot.
4017    unsafe {
4018        use core::arch::x86_64::*;
4019        let lomask = _mm_set1_epi8(0x0F);
4020        let eight = _mm256_set1_epi8(8);
4021        let mut acc = 0f32;
4022        for gi in 0..gpr {
4023            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4024            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4025            let b = _mm_loadu_si128(t.add(2) as *const __m128i);
4026            let lo = _mm_and_si128(b, lomask);
4027            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
4028            let w = _mm256_sub_epi8(
4029                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4030                eight,
4031            );
4032            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4033            let d = dpbusd_hsum(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
4034            acc += d as f32 * s;
4035        }
4036        acc
4037    }
4038}
4039
4040/// One q4_tiled row against FOUR activation streams: the nibble unpack
4041/// and abs() happen once per group instead of once per (group,
4042/// activation) — the unpack is the dominant per-element cost of the
4043/// tiled format (roadmap P0 portable blocking, q4t leg).
4044#[cfg(target_arch = "x86_64")]
4045// `fma` is NOT implied by `avx2`: without it LLVM lowers _mm256_fmadd_ps
4046// to a libm call per lane — measured 2x slower than the reduction this
4047// kernel replaces. The runtime gate (`avx2_enabled`) already requires
4048// both features, so declaring it here is safe.
4049#[target_feature(enable = "avx2,fma")]
4050unsafe fn dot_q4t_row_1x4_avx2(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
4051    // SAFETY: callers uphold the 18B-tile and xq-length contracts.
4052    unsafe {
4053        use core::arch::x86_64::*;
4054        let lomask = _mm_set1_epi8(0x0F);
4055        let eight = _mm256_set1_epi8(8);
4056        let ones = _mm256_set1_epi16(1);
4057        // One f32 accumulator VECTOR per activation, reduced once at the
4058        // end. Folding each group's i32 lanes to a scalar inside the loop
4059        // costs an extracti128 + three shift/add + a movd — a cross-lane
4060        // dependency chain per (group, activation), 288 of them per row at
4061        // cols=2304. The per-group scale is what forces a float
4062        // accumulator; it does not force a horizontal sum.
4063        //
4064        // The four accumulators are NAMED, not an array: as `[__m256; 4]`
4065        // indexed by a loop variable LLVM keeps them in memory and every
4066        // group pays four 32-byte loads and stores. That alone made this
4067        // kernel 2x SLOWER than the per-group reduction it replaces
4068        // (measured on the EPYC box: 150 s vs 71 s for two 256² steps).
4069        let mut f0 = _mm256_setzero_ps();
4070        let mut f1 = _mm256_setzero_ps();
4071        let mut f2 = _mm256_setzero_ps();
4072        let mut f3 = _mm256_setzero_ps();
4073        for gi in 0..gpr {
4074            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4075            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4076            let sv = _mm256_set1_ps(s);
4077            let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
4078            let lo = _mm_and_si128(bb, lomask);
4079            let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
4080            let w = _mm256_sub_epi8(
4081                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4082                eight,
4083            );
4084            let aw = _mm256_abs_epi8(w);
4085            let off = gi * GROUP_SIZE;
4086            let dot = |xq: &[i8]| {
4087                let x = _mm256_loadu_si256(xq.as_ptr().add(off) as *const __m256i);
4088                let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
4089                _mm256_cvtepi32_ps(_mm256_madd_epi16(p16, ones))
4090            };
4091            f0 = _mm256_fmadd_ps(dot(xs[0]), sv, f0);
4092            f1 = _mm256_fmadd_ps(dot(xs[1]), sv, f1);
4093            f2 = _mm256_fmadd_ps(dot(xs[2]), sv, f2);
4094            f3 = _mm256_fmadd_ps(dot(xs[3]), sv, f3);
4095        }
4096        [
4097            hsum256_ps(f0),
4098            hsum256_ps(f1),
4099            hsum256_ps(f2),
4100            hsum256_ps(f3),
4101        ]
4102    }
4103}
4104
4105/// Horizontal sum of eight f32 lanes — the one cross-lane reduction the
4106/// blocked kernels pay, once per row instead of once per group.
4107#[cfg(target_arch = "x86_64")]
4108#[target_feature(enable = "avx2")]
4109#[inline]
4110unsafe fn hsum256_ps(v: core::arch::x86_64::__m256) -> f32 {
4111    // SAFETY: pure register arithmetic on the caller's vector.
4112    unsafe {
4113        use core::arch::x86_64::*;
4114        let hi = _mm256_extractf128_ps::<1>(v);
4115        let s = _mm_add_ps(_mm256_castps256_ps128(v), hi);
4116        let s = _mm_add_ps(s, _mm_movehl_ps(s, s));
4117        let s = _mm_add_ss(s, _mm_shuffle_ps::<0x55>(s, s));
4118        _mm_cvtss_f32(s)
4119    }
4120}
4121
4122/// VNNI twin of `dot_q4t_row_1x4_avx2` (see `dpbusd_hsum`).
4123#[cfg(target_arch = "x86_64")]
4124#[target_feature(enable = "avx2,fma,avx512f,avx512bw,avx512vl,avx512vnni")]
4125unsafe fn dot_q4t_row_1x4_vnni(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
4126    // SAFETY: callers uphold the 18B-tile and xq-length contracts.
4127    unsafe {
4128        use core::arch::x86_64::*;
4129        let lomask = _mm_set1_epi8(0x0F);
4130        let eight = _mm256_set1_epi8(8);
4131        // Same shape as the AVX2 twin: accumulate in f32 vectors and pay
4132        // one cross-lane reduction per row, not per (group, activation).
4133        let mut f0 = _mm256_setzero_ps();
4134        let mut f1 = _mm256_setzero_ps();
4135        let mut f2 = _mm256_setzero_ps();
4136        let mut f3 = _mm256_setzero_ps();
4137        for gi in 0..gpr {
4138            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4139            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4140            let sv = _mm256_set1_ps(s);
4141            let bb = _mm_loadu_si128(t.add(2) as *const __m128i);
4142            let lo = _mm_and_si128(bb, lomask);
4143            let hi = _mm_and_si128(_mm_srli_epi16::<4>(bb), lomask);
4144            let w = _mm256_sub_epi8(
4145                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4146                eight,
4147            );
4148            let aw = _mm256_abs_epi8(w);
4149            let off = gi * GROUP_SIZE;
4150            let dot = |xq: &[i8]| {
4151                let x = _mm256_loadu_si256(xq.as_ptr().add(off) as *const __m256i);
4152                _mm256_cvtepi32_ps(_mm256_dpbusd_epi32(
4153                    _mm256_setzero_si256(),
4154                    aw,
4155                    _mm256_sign_epi8(x, w),
4156                ))
4157            };
4158            f0 = _mm256_fmadd_ps(dot(xs[0]), sv, f0);
4159            f1 = _mm256_fmadd_ps(dot(xs[1]), sv, f1);
4160            f2 = _mm256_fmadd_ps(dot(xs[2]), sv, f2);
4161            f3 = _mm256_fmadd_ps(dot(xs[3]), sv, f3);
4162        }
4163        let acc = [
4164            hsum256_ps(f0),
4165            hsum256_ps(f1),
4166            hsum256_ps(f2),
4167            hsum256_ps(f3),
4168        ];
4169        acc
4170    }
4171}
4172
4173/// ARM twin of `dot_q4t_row_1x4_avx2`: one nibble unpack per group
4174/// serves FOUR activation streams. Per stream the group order and f32
4175/// accumulation match `dot_q4t_row_sdot` exactly — batch == matvec
4176/// bit-for-bit.
4177#[cfg(target_arch = "aarch64")]
4178#[target_feature(enable = "neon,dotprod")]
4179unsafe fn dot_q4t_row_1x4_sdot(bytes: &[u8], r: usize, gpr: usize, xs: [&[i8]; 4]) -> [f32; 4] {
4180    // SAFETY: callers uphold the 18B-tile and xq-length contracts.
4181    unsafe {
4182        use core::arch::aarch64::*;
4183        use core::arch::asm;
4184        let lomask = vdupq_n_u8(0x0F);
4185        let eight = vdupq_n_s8(8);
4186        let mut acc = [0f32; 4];
4187        for gi in 0..gpr {
4188            let t = bytes.as_ptr().add((r * gpr + gi) * Q4_TILE);
4189            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
4190            let b = vld1q_u8(t.add(2));
4191            let lo = vandq_u8(b, lomask);
4192            let hi = vshrq_n_u8::<4>(b);
4193            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
4194            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
4195            for (k, xq) in xs.iter().enumerate() {
4196                let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
4197                let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
4198                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
4199                asm!(
4200                    "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
4201                    "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
4202                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4203                    e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
4204                    options(pure, nomem, nostack),
4205                );
4206                acc[k] += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
4207            }
4208        }
4209        acc
4210    }
4211}
4212
4213/// Exact-term correction for A8W8 outliers on a tiled row.
4214#[inline]
4215fn q4t_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
4216    let gi = j / GROUP_SIZE;
4217    let k = j % GROUP_SIZE;
4218    let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
4219    let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
4220    let byte = tile[2 + k / 2];
4221    let nib = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
4222    ((nib as i32 - 8) as f32, s)
4223}
4224
4225/// Exact scalar q4_tiled row (CMF_SDOT=0 contract) — same pairwise
4226/// accumulation shape as `q4_range_f32`.
4227#[inline]
4228fn q4t_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
4229    let mut acc = 0f32;
4230    for gi in 0..gpr {
4231        let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
4232        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
4233        let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
4234        let mut ga = 0f32;
4235        for (k, &b) in tile[2..].iter().enumerate() {
4236            ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
4237                + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
4238        }
4239        acc += ga * s;
4240    }
4241    acc
4242}
4243
4244/// Split view of a `q4tp` payload. The three planes are resolved once per
4245/// matvec instead of per row — `q4tp_sections` is cheap, but doing it inside
4246/// the row loop would put a division on the hot path for nothing.
4247struct Q4tpView<'a> {
4248    nib: &'a [u8],
4249    params: &'a [u8],
4250    codes: &'a [u8],
4251    stride: usize,
4252    /// q2tp reads the ladder with rung 0 = exact zero.
4253    zero_rung: bool,
4254}
4255
4256impl<'a> Q4tpView<'a> {
4257    fn new(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
4258        let (params_off, codes_off, stride) = q4tp_sections(rows, cols);
4259        Self {
4260            nib: &bytes[..params_off],
4261            params: &bytes[params_off..codes_off],
4262            codes: &bytes[codes_off..],
4263            stride,
4264            zero_rung: false,
4265        }
4266    }
4267
4268    /// The q2tp view: identical params/codes planes, 8 B weight chunks.
4269    fn new_q2(bytes: &'a [u8], rows: usize, cols: usize) -> Self {
4270        let (params_off, codes_off, stride) = q2tp_sections(rows, cols);
4271        Self {
4272            nib: &bytes[..params_off],
4273            params: &bytes[params_off..codes_off],
4274            codes: &bytes[codes_off..],
4275            stride,
4276            zero_rung: true,
4277        }
4278    }
4279
4280    /// Expand row `r`'s per-tile scales into `out` (length `gpr`).
4281    ///
4282    /// Doing this once per row — rather than decoding a 5-bit code inside the
4283    /// tile loop — is what makes the format free at runtime. Random access to
4284    /// a packed 5-bit field costs a division, two bounds checks and a branch;
4285    /// the tile's actual work is two `sdot`s, so per-tile decoding dominated
4286    /// the kernel and cost 5x (measured: 1.4 vs 6.9 tok/s on Nanbeige-3B).
4287    /// Walking the plane sequentially with a bit accumulator is ~3 ops.
4288    /// Eight 5-bit codes are exactly five bytes, so a whole group of
4289    /// eight decodes from one little-endian word at fixed shifts. The
4290    /// bit-accumulator this replaces carried a data-dependent `while
4291    /// have < 5` refill whose branch sat in the innermost loop of every
4292    /// q4tp row; a decode profile put this function above the dot
4293    /// products it feeds. Same bitstream, same codes — just no branch
4294    /// and eight independent extractions.
4295    #[inline]
4296    fn scales_into(&self, r: usize, gpr: usize, out: &mut [f32]) {
4297        let tab = if self.zero_rung {
4298            q2tp_ladder(self.params, r)
4299        } else {
4300            q4tp_ladder(self.params, r)
4301        };
4302        let codes = &self.codes[r * self.stride..(r + 1) * self.stride];
4303        let out = &mut out[..gpr];
4304        let mut chunks = out.chunks_exact_mut(8);
4305        let mut ci = 0usize;
4306        for c in &mut chunks {
4307            let w = u64::from(codes[ci])
4308                | u64::from(codes[ci + 1]) << 8
4309                | u64::from(codes[ci + 2]) << 16
4310                | u64::from(codes[ci + 3]) << 24
4311                | u64::from(codes[ci + 4]) << 32;
4312            for (k, o) in c.iter_mut().enumerate() {
4313                *o = tab[((w >> (5 * k)) & 31) as usize];
4314            }
4315            ci += 5;
4316        }
4317        // Fewer than eight codes left: the shared total accessor, which
4318        // tolerates a 5-bit field whose spill byte is past the stride.
4319        let tail = &codes[ci..];
4320        for (k, o) in chunks.into_remainder().iter_mut().enumerate() {
4321            *o = tab[q4tp_code(tail, k)];
4322        }
4323    }
4324}
4325
4326#[inline]
4327fn dot_q4tp_row_i8(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
4328    #[cfg(target_arch = "aarch64")]
4329    unsafe {
4330        return dot_q4tp_row_sdot(nib, r, gpr, xq, scales);
4331    }
4332    #[cfg(target_arch = "x86_64")]
4333    unsafe {
4334        if vnni_tiles_enabled() {
4335            return dot_q4tp_row_vnni(nib, r, gpr, xq, scales);
4336        }
4337        return dot_q4tp_row_avx2(nib, r, gpr, xq, scales);
4338    }
4339    #[allow(unreachable_code)]
4340    {
4341        let mut acc = 0f32;
4342        for gi in 0..gpr {
4343            let tile = &nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
4344            let s = scales[gi];
4345            let mut d = 0i32;
4346            for (k, &b) in tile.iter().enumerate() {
4347                d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
4348                    + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
4349            }
4350            acc += d as f32 * s;
4351        }
4352        acc
4353    }
4354}
4355
4356/// q4tp twin of `dot_q4t_row_sdot`: identical nibble math, but the tile
4357/// stride is 16 B (no inline scale) and the scale is a ladder lookup.
4358#[cfg(target_arch = "aarch64")]
4359#[target_feature(enable = "neon,dotprod")]
4360unsafe fn dot_q4tp_row_sdot(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
4361    // SAFETY: callers uphold slice-length contracts (16B tile per group,
4362    // xq.len() == gpr·GROUP_SIZE, codes covering gpr 5-bit fields).
4363    unsafe {
4364        use core::arch::aarch64::*;
4365        use core::arch::asm;
4366        let lomask = vdupq_n_u8(0x0F);
4367        let eight = vdupq_n_s8(8);
4368        let mut acc = 0f32;
4369        for gi in 0..gpr {
4370            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
4371            let s = *scales.get_unchecked(gi);
4372            let b = vld1q_u8(t);
4373            let lo = vandq_u8(b, lomask);
4374            let hi = vshrq_n_u8::<4>(b);
4375            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
4376            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
4377            let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
4378            let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
4379            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
4380            asm!(
4381                "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
4382                "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
4383                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4384                e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
4385                options(pure, nomem, nostack),
4386            );
4387            acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
4388        }
4389        acc
4390    }
4391}
4392
4393#[cfg(target_arch = "x86_64")]
4394#[target_feature(enable = "avx2")]
4395unsafe fn dot_q4tp_row_avx2(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
4396    // SAFETY: see dot_q4tp_row_sdot.
4397    unsafe {
4398        use core::arch::x86_64::*;
4399        let lomask = _mm_set1_epi8(0x0F);
4400        let eight = _mm256_set1_epi8(8);
4401        let ones = _mm256_set1_epi16(1);
4402        let mut acc = 0f32;
4403        for gi in 0..gpr {
4404            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
4405            let s = *scales.get_unchecked(gi);
4406            let b = _mm_loadu_si128(t as *const __m128i);
4407            let lo = _mm_and_si128(b, lomask);
4408            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
4409            let w = _mm256_sub_epi8(
4410                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4411                eight,
4412            );
4413            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4414            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
4415            let d = _mm256_madd_epi16(p16, ones);
4416            let hi128 = _mm256_extracti128_si256::<1>(d);
4417            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
4418            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
4419            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
4420            acc += _mm_cvtsi128_si32(s32) as f32 * s;
4421        }
4422        acc
4423    }
4424}
4425
4426
4427/// VNNI twin of `dot_q4tp_row_avx2` (see `dot_q4t_row_vnni` for why the
4428/// 256-bit VL encoding is the one to use here).
4429#[cfg(target_arch = "x86_64")]
4430#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
4431unsafe fn dot_q4tp_row_vnni(nib: &[u8], r: usize, gpr: usize, xq: &[i8], scales: &[f32]) -> f32 {
4432    // SAFETY: see dot_q4tp_row_sdot.
4433    unsafe {
4434        use core::arch::x86_64::*;
4435        let lomask = _mm_set1_epi8(0x0F);
4436        let eight = _mm256_set1_epi8(8);
4437        let mut acc = 0f32;
4438        for gi in 0..gpr {
4439            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
4440            let s = *scales.get_unchecked(gi);
4441            let b = _mm_loadu_si128(t as *const __m128i);
4442            let lo = _mm_and_si128(b, lomask);
4443            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
4444            let w = _mm256_sub_epi8(
4445                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
4446                eight,
4447            );
4448            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
4449            acc += dpbusd_hsum(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w)) as f32 * s;
4450        }
4451        acc
4452    }
4453}
4454
4455/// Exact scalar q4tp row — the `CMF_SDOT=0` contract, same pairwise
4456/// accumulation shape as `q4t_row_exact`.
4457#[inline]
4458fn q4tp_row_exact(nib: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
4459    let mut acc = 0f32;
4460    for gi in 0..gpr {
4461        let tile = &nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
4462        let s = scales[gi];
4463        let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
4464        let mut ga = 0f32;
4465        for (k, &b) in tile.iter().enumerate() {
4466            ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
4467                + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
4468        }
4469        acc += ga * s;
4470    }
4471    acc
4472}
4473
4474/// Single weight of a q4tp tensor — the a8w8 outlier path, which restores
4475/// activation outliers at full precision after the int8 pass.
4476#[inline]
4477fn q4tp_outlier(nib: &[u8], r: usize, gpr: usize, j: usize, scales: &[f32]) -> (f32, f32) {
4478    let (gi, k) = (j / GROUP_SIZE, j % GROUP_SIZE);
4479    let byte = nib[(r * gpr + gi) * Q4TP_NIB + k / 2];
4480    let n = if k & 1 == 0 { byte & 0x0F } else { byte >> 4 };
4481    ((n as i32 - 8) as f32, scales[gi])
4482}
4483
4484/// Fused q4tp matvec (dispatch mirrors `q4t_matvec`).
4485fn q4tp_matvec(
4486    bytes: &[u8],
4487    x: &[f32],
4488    rows: usize,
4489    cols: usize,
4490    out: &mut [f32],
4491    pool: Option<&Pool>,
4492) {
4493    debug_assert_eq!(out.len(), rows);
4494    let gpr = cols / GROUP_SIZE;
4495    let v = Q4tpView::new(bytes, rows, cols);
4496    let out_addr = SendMut(out.as_mut_ptr());
4497    if a8w8_enabled() {
4498        let act = split_act(x);
4499        let run = |start: usize, end: usize| {
4500            // One scratch row of scales per worker — borrowed, not minted.
4501            with_krow(gpr, |sc| {
4502                for r in start..end {
4503                    v.scales_into(r, gpr, sc);
4504                    let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, sc) * act.sx;
4505                    for &(j, xv) in &act.outliers {
4506                        let (w, s) = q4tp_outlier(v.nib, r, gpr, j, sc);
4507                        acc += w * s * xv;
4508                    }
4509                    // SAFETY: disjoint row ranges per worker.
4510                    unsafe { *out_addr.at(r) = acc };
4511                }
4512            })
4513        };
4514        dispatch_rows(pool, rows, &run);
4515        return;
4516    }
4517    let run = |start: usize, end: usize| {
4518        with_krow(gpr, |sc| {
4519            for r in start..end {
4520                v.scales_into(r, gpr, sc);
4521                // SAFETY: disjoint row ranges per worker.
4522                unsafe { *out_addr.at(r) = q4tp_row_exact(v.nib, r, gpr, x, sc) };
4523            }
4524        })
4525    };
4526    dispatch_rows(pool, rows, &run);
4527}
4528
4529/// Fused two-input q4tp matvec — the SwiGLU gate/up pair. Weights and the
4530/// row ladder are read once and spent on both activation streams.
4531#[allow(clippy::too_many_arguments)]
4532fn q4tp_matvec2(
4533    bytes: &[u8],
4534    x1: &[f32],
4535    x2: &[f32],
4536    rows: usize,
4537    cols: usize,
4538    o1: &mut [f32],
4539    o2: &mut [f32],
4540    pool: Option<&Pool>,
4541) {
4542    let gpr = cols / GROUP_SIZE;
4543    let v = Q4tpView::new(bytes, rows, cols);
4544    let (p1, p2) = (SendMut(o1.as_mut_ptr()), SendMut(o2.as_mut_ptr()));
4545    let run = |start: usize, end: usize| {
4546        let mut sc = vec![0f32; gpr];
4547        for r in start..end {
4548            v.scales_into(r, gpr, &mut sc);
4549            // SAFETY: disjoint row ranges per worker.
4550            unsafe {
4551                *p1.at(r) = q4tp_row_exact(v.nib, r, gpr, x1, &sc);
4552                *p2.at(r) = q4tp_row_exact(v.nib, r, gpr, x2, &sc);
4553            }
4554        }
4555    };
4556    dispatch_rows(pool, rows, &run);
4557}
4558
4559/// One q2tp outlier weight at column `j` of row `r`: the 2-bit code and
4560/// its group scale, mirrored on `q4tp_outlier`.
4561#[inline]
4562fn q2tp_outlier(chunks: &[u8], r: usize, gpr: usize, j: usize, scales: &[f32]) -> (f32, f32) {
4563    let (gi, k) = (j / GROUP_SIZE, j % GROUP_SIZE);
4564    let byte = chunks[(r * gpr + gi) * Q2TP_CHUNK + k / 4];
4565    let c = (byte >> (2 * (k % 4))) & 3;
4566    (c as f32 - 1.5, scales[gi])
4567}
4568
4569/// Integer dot of one q2tp row against pre-quantized activations:
4570/// Σ_g s_g · (Σ c·xq − 1.5·Σ xq). The half-integer grid (c − 1.5)
4571/// becomes exact integer math through the group sums — the same trick
4572/// every a8w8 kernel in this file rides. The codes decode into a
4573/// 32-byte scratch in natural order and the dot itself is the shared
4574/// SDOT primitive; elsewhere a scalar integer loop.
4575#[inline]
4576fn dot_q2tp_row_i8(
4577    chunks: &[u8],
4578    r: usize,
4579    gpr: usize,
4580    xq: &[i8],
4581    gsum: &[i32],
4582    scales: &[f32],
4583) -> f32 {
4584    let mut acc = 0f32;
4585    let base = r * gpr * Q2TP_CHUNK;
4586    #[cfg(not(target_arch = "aarch64"))]
4587    let mut codes = [0i8; GROUP_SIZE];
4588    for gi in 0..gpr {
4589        let ch = &chunks[base + gi * Q2TP_CHUNK..base + (gi + 1) * Q2TP_CHUNK];
4590        let xg = &xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
4591        #[cfg(target_arch = "aarch64")]
4592        // NEON: the byte's four 2-bit fields land in four lane vectors
4593        // (shift+mask), vld4 de-interleaves xq to match (xj[k] =
4594        // xq[4k+j]), widening MACs accumulate exactly in i32. A scalar
4595        // decode here cost as much as the dot it fed — the profile put
4596        // it at the top of the whole W2 decode.
4597        let dot = unsafe {
4598            use core::arch::aarch64::*;
4599            let b = vld1_u8(ch.as_ptr());
4600            let three = vdup_n_u8(3);
4601            let c0 = vreinterpret_s8_u8(vand_u8(b, three));
4602            let c1 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 2), three));
4603            let c2 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 4), three));
4604            let c3 = vreinterpret_s8_u8(vand_u8(vshr_n_u8(b, 6), three));
4605            let x4 = vld4_s8(xg.as_ptr());
4606            let mut acc4 = vdupq_n_s32(0);
4607            acc4 = vpadalq_s16(acc4, vmull_s8(c0, x4.0));
4608            acc4 = vpadalq_s16(acc4, vmull_s8(c1, x4.1));
4609            acc4 = vpadalq_s16(acc4, vmull_s8(c2, x4.2));
4610            acc4 = vpadalq_s16(acc4, vmull_s8(c3, x4.3));
4611            vaddvq_s32(acc4)
4612        };
4613        #[cfg(not(target_arch = "aarch64"))]
4614        let dot: i32 = {
4615            for (k, &b) in ch.iter().enumerate() {
4616                codes[k * 4] = (b & 3) as i8;
4617                codes[k * 4 + 1] = ((b >> 2) & 3) as i8;
4618                codes[k * 4 + 2] = ((b >> 4) & 3) as i8;
4619                codes[k * 4 + 3] = ((b >> 6) & 3) as i8;
4620            }
4621            codes
4622                .iter()
4623                .zip(xg)
4624                .map(|(&c, &x)| c as i32 * x as i32)
4625                .sum()
4626        };
4627        acc += scales[gi] * (dot as f32 - 1.5 * gsum[gi] as f32);
4628    }
4629    acc
4630}
4631
4632/// Exact f32 dot of one q2tp row: 2-bit fields LSB-first, (c − 1.5)·s.
4633/// Scalar on purpose — the 2-bit class targets the GPU graph; the CPU
4634/// path exists for parity gates and small-machine fallback.
4635fn q2tp_row_exact(chunks: &[u8], r: usize, gpr: usize, x: &[f32], scales: &[f32]) -> f32 {
4636    let mut acc = 0f32;
4637    for gi in 0..gpr {
4638        let ch = &chunks[(r * gpr + gi) * Q2TP_CHUNK..(r * gpr + gi + 1) * Q2TP_CHUNK];
4639        let s = scales[gi];
4640        let xb = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
4641        let mut g = 0f32;
4642        for (k, &b) in ch.iter().enumerate() {
4643            g += ((b & 3) as f32 - 1.5) * xb[k * 4]
4644                + (((b >> 2) & 3) as f32 - 1.5) * xb[k * 4 + 1]
4645                + (((b >> 4) & 3) as f32 - 1.5) * xb[k * 4 + 2]
4646                + (((b >> 6) & 3) as f32 - 1.5) * xb[k * 4 + 3];
4647        }
4648        acc += s * g;
4649    }
4650    acc
4651}
4652
4653fn q2tp_matvec(
4654    bytes: &[u8],
4655    x: &[f32],
4656    rows: usize,
4657    cols: usize,
4658    out: &mut [f32],
4659    pool: Option<&Pool>,
4660) {
4661    debug_assert_eq!(out.len(), rows);
4662    let gpr = cols / GROUP_SIZE;
4663    let v = Q4tpView::new_q2(bytes, rows, cols);
4664    let out_addr = SendMut(out.as_mut_ptr());
4665    // a8w8 fast path (CMF_SDOT=0 keeps the exact scalar walk): integer
4666    // code dots + group sums, exact outlier correction — the same
4667    // contract as every sibling kernel; measured 2-bit rows were the
4668    // only scalar holdout in the family.
4669    if a8w8_enabled() {
4670        let act = split_act(x);
4671        let gsum = q1_group_sums(&act.xq, gpr);
4672        let (act, gsum) = (&act, &gsum);
4673        let run = move |start: usize, end: usize| {
4674            with_krow(gpr, |sc| {
4675                for r in start..end {
4676                    v.scales_into(r, gpr, sc);
4677                    let mut acc = dot_q2tp_row_i8(v.nib, r, gpr, &act.xq, gsum, sc) * act.sx;
4678                    for &(j, xv) in &act.outliers {
4679                        let (w, s) = q2tp_outlier(v.nib, r, gpr, j, sc);
4680                        acc += w * s * xv;
4681                    }
4682                    // SAFETY: disjoint row ranges per worker.
4683                    unsafe { *out_addr.at(r) = acc };
4684                }
4685            })
4686        };
4687        dispatch_rows(pool, rows, &run);
4688        return;
4689    }
4690    let run = |start: usize, end: usize| {
4691        with_krow(gpr, |sc| {
4692            for r in start..end {
4693                v.scales_into(r, gpr, sc);
4694                // SAFETY: disjoint row ranges per worker.
4695                unsafe { *out_addr.at(r) = q2tp_row_exact(v.nib, r, gpr, x, sc) };
4696            }
4697        })
4698    };
4699    dispatch_rows(pool, rows, &run);
4700}
4701
4702/// Fused two-input q2tp matvec — the SwiGLU gate/up pair.
4703#[allow(clippy::too_many_arguments)]
4704fn q2tp_matvec2(
4705    bytes: &[u8],
4706    x1: &[f32],
4707    x2: &[f32],
4708    rows: usize,
4709    cols: usize,
4710    o1: &mut [f32],
4711    o2: &mut [f32],
4712    pool: Option<&Pool>,
4713) {
4714    let gpr = cols / GROUP_SIZE;
4715    let v = Q4tpView::new_q2(bytes, rows, cols);
4716    let (p1, p2) = (SendMut(o1.as_mut_ptr()), SendMut(o2.as_mut_ptr()));
4717    let run = |start: usize, end: usize| {
4718        let mut sc = vec![0f32; gpr];
4719        for r in start..end {
4720            v.scales_into(r, gpr, &mut sc);
4721            // SAFETY: disjoint row ranges per worker.
4722            unsafe {
4723                *p1.at(r) = q2tp_row_exact(v.nib, r, gpr, x1, &sc);
4724                *p2.at(r) = q2tp_row_exact(v.nib, r, gpr, x2, &sc);
4725            }
4726        }
4727    };
4728    dispatch_rows(pool, rows, &run);
4729}
4730
4731/// Batched q2tp matmat: scalar row kernel over every batch column. CPU
4732/// prefill only — decode rides the graph, so plain and correct beats
4733/// clever here.
4734/// Test doors into the host 2-bit kernels: the stand's heap corruption
4735/// pointed at down-shaped tensors, and the private fns need a way to be
4736/// held to a reference without a model file around them.
4737pub fn q2tp_matvec_for_test(bytes: &[u8], x: &[f32], rows: usize, cols: usize, out: &mut [f32]) {
4738    // The facade IS the reference: encoder oracles hold requant output
4739    // to the exact scalar walk. The production dispatch may take the i8
4740    // fast path, whose error scale is the ACTIVATIONS' — a different
4741    // claim than the encoder correctness these tests pin.
4742    let gpr = cols / GROUP_SIZE;
4743    let v = Q4tpView::new_q2(bytes, rows, cols);
4744    with_krow(gpr, |sc| {
4745        for r in 0..rows {
4746            v.scales_into(r, gpr, sc);
4747            out[r] = q2tp_row_exact(v.nib, r, gpr, x, sc);
4748        }
4749    });
4750}
4751
4752pub fn q2tp_matmat_for_test(
4753    bytes: &[u8],
4754    xs_all: &[f32],
4755    b: usize,
4756    rows: usize,
4757    cols: usize,
4758    out: &mut [f32],
4759) {
4760    q2tp_matmat(bytes, xs_all, b, rows, cols, out, None);
4761}
4762
4763fn q2tp_matmat(
4764    bytes: &[u8],
4765    xs_all: &[f32],
4766    b: usize,
4767    rows: usize,
4768    cols: usize,
4769    out: &mut [f32],
4770    pool: Option<&Pool>,
4771) {
4772    debug_assert_eq!(out.len(), b * rows);
4773    let gpr = cols / GROUP_SIZE;
4774    let v = Q4tpView::new_q2(bytes, rows, cols);
4775    let out_addr = SendMut(out.as_mut_ptr());
4776    let run = |start: usize, end: usize| {
4777        let mut sc = vec![0f32; gpr];
4778        for r in start..end {
4779            v.scales_into(r, gpr, &mut sc);
4780            for bi in 0..b {
4781                let x = &xs_all[bi * cols..(bi + 1) * cols];
4782                // SAFETY: disjoint row ranges per worker.
4783                unsafe { *out_addr.at(bi * rows + r) = q2tp_row_exact(v.nib, r, gpr, x, &sc) };
4784            }
4785        }
4786    };
4787    dispatch_rows(pool, rows, &run);
4788}
4789
4790/// The pre-vectorised shape, kept for A/B (`CMF_Q4TP_V1=1`): the
4791/// horizontal add lands once per group per column instead of once per
4792/// row. Same weights, same activations — only the reduction differs.
4793#[cfg(target_arch = "aarch64")]
4794#[target_feature(enable = "neon,dotprod")]
4795unsafe fn dot_q4tp_row_1x4_sdot_v1(
4796    nib: &[u8],
4797    r: usize,
4798    gpr: usize,
4799    xs: [&[i8]; 4],
4800    scales: &[f32],
4801) -> [f32; 4] {
4802    unsafe {
4803        use core::arch::aarch64::*;
4804        use core::arch::asm;
4805        let lomask = vdupq_n_u8(0x0F);
4806        let eight = vdupq_n_s8(8);
4807        let (mut f0, mut f1, mut f2, mut f3) = (0f32, 0f32, 0f32, 0f32);
4808        for gi in 0..gpr {
4809            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
4810            let s = *scales.get_unchecked(gi);
4811            let bb = vld1q_u8(t);
4812            let lo = vandq_u8(bb, lomask);
4813            let hi = vshrq_n_u8::<4>(bb);
4814            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
4815            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
4816            let mut d = [0f32; 4];
4817            for (k, dk) in d.iter_mut().enumerate() {
4818                let x0 = vld1q_s8(xs[k].as_ptr().add(gi * GROUP_SIZE));
4819                let x1 = vld1q_s8(xs[k].as_ptr().add(gi * GROUP_SIZE + 16));
4820                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
4821                asm!(
4822                    "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
4823                    "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
4824                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
4825                    e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
4826                    options(pure, nomem, nostack),
4827                );
4828                *dk = vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
4829            }
4830            f0 += d[0];
4831            f1 += d[1];
4832            f2 += d[2];
4833            f3 += d[3];
4834        }
4835        [f0, f1, f2, f3]
4836    }
4837}
4838
4839/// Which q4tp batch kernel to run: 1 = the previous one, 2 = the tuned
4840/// one, 0 = decide from the CPU. An atomic rather than a `OnceLock` so a
4841/// benchmark can alternate the two inside one process, where the machine's
4842/// mood — a shared box drifts ±25% between runs — is the same for both.
4843/// What the two mean is per-architecture: on x86 the blocked AVX-512 path
4844/// against the per-column one, on ARM the two reduction shapes.
4845#[allow(dead_code)]
4846static Q4TP_ALT: std::sync::atomic::AtomicU8 = std::sync::atomic::AtomicU8::new(0);
4847
4848/// Blocking pays on x86 only with 512-bit VNNI. With AVX2 alone, four
4849/// columns sharing an unpack still measured slower than the per-column
4850/// path (23.2 ms against 19.4 on a 48-thread EPYC), because that path
4851/// already dequantizes the row once — so the blocked kernel bought a
4852/// second unpack-free pass at the price of half the vector width.
4853#[cfg(target_arch = "x86_64")]
4854fn q4tp_blocked_x86() -> bool {
4855    match Q4TP_ALT.load(std::sync::atomic::Ordering::Relaxed) {
4856        1 => false,
4857        // A forced ON still asks the CPU. The switch exists so a bench can
4858        // pick a kernel, not so it can promise instructions the machine
4859        // does not have — CI caught that as a SIGILL on a runner without
4860        // AVX-512, where the parity test had turned the path on by hand.
4861        2 => avx512vnni_enabled(),
4862        // Deliberately not cached back into the switch: both gates below
4863        // hold their own `OnceLock`, and latching their answer here would
4864        // make a test's override outlive the test that set it.
4865        _ => blocked_enabled() && avx512vnni_enabled(),
4866    }
4867}
4868
4869/// `CMF_Q4TP_V1=1` picks the old reduction shape (A/B only).
4870#[cfg(target_arch = "aarch64")]
4871#[allow(dead_code)]
4872fn q4tp_v1() -> bool {
4873    match Q4TP_ALT.load(std::sync::atomic::Ordering::Relaxed) {
4874        1 => true,
4875        2 => false,
4876        _ => {
4877            static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
4878            *ON.get_or_init(|| std::env::var("CMF_Q4TP_V1").is_ok_and(|v| v != "0"))
4879        }
4880    }
4881}
4882
4883/// Two weight rows against eight columns. The activation load is the
4884/// same for both rows, so it is paid once for twice the arithmetic, and
4885/// sixteen accumulator chains run where eight did — which is what a kernel
4886/// retiring 0.29 instructions a cycle is short of. Register pressure is
4887/// the limit: sixteen `zmm` accumulators, two weight tiles, one
4888/// activation, of thirty-two.
4889///
4890/// Four rows by four columns spends the same sixteen accumulators the
4891/// other way and measured worse — 1488 GFLOP/s against 1644 — so the
4892/// unpack, which four rows pay twice as often, costs more than the extra
4893/// sharing of one activation load buys.
4894#[cfg(target_arch = "x86_64")]
4895#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
4896unsafe fn dot_q4tp_2x8_avx512(
4897    nib: &[u8],
4898    r0: usize,
4899    gpr: usize,
4900    xs: [&[i8]; 8],
4901    sc0: &[f32],
4902    sc1: &[f32],
4903) -> [[f32; 8]; 2] {
4904    // SAFETY: as dot_q4tp_row_1x8_avx512, two adjacent rows at once; the
4905    // caller guarantees r0 + 1 < rows and the ISA.
4906    unsafe {
4907        use core::arch::x86_64::*;
4908        let lomask = _mm256_set1_epi8(0x0F);
4909        let eight = _mm256_set1_epi8(8);
4910        let zero = _mm512_setzero_si512();
4911        let mut v0 = [_mm512_setzero_ps(); 8];
4912        let mut v1 = [_mm512_setzero_ps(); 8];
4913        let pairs = gpr / 2;
4914        let unpack = |r: usize, gi: usize| -> (__m512i, __mmask64) {
4915            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
4916            let bb = _mm256_loadu_si256(t as *const __m256i);
4917            let lo = _mm256_and_si256(bb, lomask);
4918            let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
4919            let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
4920            let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
4921            let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
4922            let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
4923            (_mm512_abs_epi8(w), _mm512_movepi8_mask(w))
4924        };
4925        for gp in 0..pairs {
4926            let gi = gp * 2;
4927            let (wa0, neg0) = unpack(r0, gi);
4928            let (wa1, neg1) = unpack(r0 + 1, gi);
4929            let off = gi * GROUP_SIZE;
4930            let sv = |sc: &[f32]| {
4931                _mm512_insertf32x8::<1>(
4932                    _mm512_castps256_ps512(_mm256_set1_ps(*sc.get_unchecked(gi))),
4933                    _mm256_set1_ps(*sc.get_unchecked(gi + 1)),
4934                )
4935            };
4936            let s0 = sv(sc0);
4937            let s1 = sv(sc1);
4938            for k in 0..8 {
4939                let xv = _mm512_loadu_si512(xs[k].as_ptr().add(off) as *const __m512i);
4940                let d0 = _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(
4941                    zero,
4942                    wa0,
4943                    _mm512_mask_sub_epi8(xv, neg0, zero, xv),
4944                ));
4945                let d1 = _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(
4946                    zero,
4947                    wa1,
4948                    _mm512_mask_sub_epi8(xv, neg1, zero, xv),
4949                ));
4950                v0[k] = _mm512_fmadd_ps(d0, s0, v0[k]);
4951                v1[k] = _mm512_fmadd_ps(d1, s1, v1[k]);
4952            }
4953        }
4954        let mut acc = [[0f32; 8]; 2];
4955        for k in 0..8 {
4956            acc[0][k] = _mm512_reduce_add_ps(v0[k]);
4957            acc[1][k] = _mm512_reduce_add_ps(v1[k]);
4958        }
4959        if gpr % 2 == 1 {
4960            let off = (gpr - 1) * GROUP_SIZE;
4961            for j in off..off + GROUP_SIZE {
4962                let (w0, sa) = q4tp_outlier(nib, r0, gpr, j, sc0);
4963                let (w1, sb) = q4tp_outlier(nib, r0 + 1, gpr, j, sc1);
4964                for k in 0..8 {
4965                    let x = *xs[k].get_unchecked(j) as f32;
4966                    acc[0][k] += w0 * sa * x;
4967                    acc[1][k] += w1 * sb * x;
4968                }
4969            }
4970        }
4971        acc
4972    }
4973}
4974
4975/// The same, eight columns at a time. One unpack then feeds twice as many
4976/// activation streams, so a wide batch reads the weight tile half as
4977/// often; the price is eight accumulators live at once. Measured 9.0 ->
4978/// 8.3 ms at 9216x2304, b=296 on a 48-thread EPYC 9B45.
4979#[cfg(target_arch = "x86_64")]
4980#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
4981unsafe fn dot_q4tp_row_1x8_avx512(
4982    nib: &[u8],
4983    r: usize,
4984    gpr: usize,
4985    xs: [&[i8]; 8],
4986    scales: &[f32],
4987) -> [f32; 8] {
4988    // SAFETY: as dot_q4tp_row_1x4_avx2; caller guarantees the ISA.
4989    unsafe {
4990        use core::arch::x86_64::*;
4991        let lomask = _mm256_set1_epi8(0x0F);
4992        let eight = _mm256_set1_epi8(8);
4993        let zero = _mm512_setzero_si512();
4994        let (mut v0, mut v1, mut v2, mut v3) = (
4995            _mm512_setzero_ps(),
4996            _mm512_setzero_ps(),
4997            _mm512_setzero_ps(),
4998            _mm512_setzero_ps(),
4999        );
5000        let (mut v4, mut v5, mut v6, mut v7) = (
5001            _mm512_setzero_ps(),
5002            _mm512_setzero_ps(),
5003            _mm512_setzero_ps(),
5004            _mm512_setzero_ps(),
5005        );
5006        let pairs = gpr / 2;
5007        for gp in 0..pairs {
5008            let gi = gp * 2;
5009            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5010            let bb = _mm256_loadu_si256(t as *const __m256i);
5011            let lo = _mm256_and_si256(bb, lomask);
5012            let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
5013            // `unpack` works per 128-bit lane, so the halves come out as
5014            // [A.lo, B.lo] and [A.hi, B.hi]; the shuffle reorders the four
5015            // 128-bit lanes into the weights' natural order, which is what
5016            // the straight activation load expects.
5017            let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
5018            let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
5019            let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
5020            let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
5021            let wabs = _mm512_abs_epi8(w);
5022            let neg = _mm512_movepi8_mask(w);
5023            let off = gi * GROUP_SIZE;
5024            let sv = _mm512_insertf32x8::<1>(
5025                _mm512_castps256_ps512(_mm256_set1_ps(*scales.get_unchecked(gi))),
5026                _mm256_set1_ps(*scales.get_unchecked(gi + 1)),
5027            );
5028            let dot = |x: &[i8]| -> __m512 {
5029                let xv = _mm512_loadu_si512(x.as_ptr().add(off) as *const __m512i);
5030                let sx = _mm512_mask_sub_epi8(xv, neg, zero, xv);
5031                _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(zero, wabs, sx))
5032            };
5033            v0 = _mm512_fmadd_ps(dot(xs[0]), sv, v0);
5034            v1 = _mm512_fmadd_ps(dot(xs[1]), sv, v1);
5035            v2 = _mm512_fmadd_ps(dot(xs[2]), sv, v2);
5036            v3 = _mm512_fmadd_ps(dot(xs[3]), sv, v3);
5037            v4 = _mm512_fmadd_ps(dot(xs[4]), sv, v4);
5038            v5 = _mm512_fmadd_ps(dot(xs[5]), sv, v5);
5039            v6 = _mm512_fmadd_ps(dot(xs[6]), sv, v6);
5040            v7 = _mm512_fmadd_ps(dot(xs[7]), sv, v7);
5041        }
5042        let mut acc = [
5043            _mm512_reduce_add_ps(v0),
5044            _mm512_reduce_add_ps(v1),
5045            _mm512_reduce_add_ps(v2),
5046            _mm512_reduce_add_ps(v3),
5047            _mm512_reduce_add_ps(v4),
5048            _mm512_reduce_add_ps(v5),
5049            _mm512_reduce_add_ps(v6),
5050            _mm512_reduce_add_ps(v7),
5051        ];
5052        // An odd group count leaves one group over; the narrow kernel
5053        // finishes it rather than the tail being a special case here.
5054        if gpr % 2 == 1 {
5055            let off = (gpr - 1) * GROUP_SIZE;
5056            for j in off..off + GROUP_SIZE {
5057                let (w, s) = q4tp_outlier(nib, r, gpr, j, scales);
5058                let ws = w * s;
5059                for k in 0..8 {
5060                    acc[k] += ws * *xs[k].get_unchecked(j) as f32;
5061                }
5062            }
5063        }
5064        acc
5065    }
5066}
5067
5068/// The same four columns, 512 bits wide. Two groups (64 weights) ride one
5069/// unpack and one `vpdpbusd`, where AVX2 needs two unpacks and four
5070/// `maddubs`/`madd` pairs — about 2.3x fewer instructions for the same
5071/// arithmetic. The two groups carry different scales, so the fma takes a
5072/// vector whose halves hold each group's scale rather than a broadcast.
5073///
5074/// There is no 512-bit `vpsignb`, so the activation's sign is applied by
5075/// negating under a mask taken from the weight's sign bits. That mask is
5076/// per-tile, so it is hoisted out of the column loop and the per-column
5077/// cost stays exactly one instruction, as with `sign_epi8`. Weights of
5078/// zero are not zeroed by the mask trick and do not need to be: their
5079/// magnitude is zero, so the product is.
5080#[cfg(target_arch = "x86_64")]
5081#[target_feature(enable = "avx512f,avx512bw,avx512vnni")]
5082unsafe fn dot_q4tp_row_1x4_avx512(
5083    nib: &[u8],
5084    r: usize,
5085    gpr: usize,
5086    xs: [&[i8]; 4],
5087    scales: &[f32],
5088) -> [f32; 4] {
5089    // SAFETY: as dot_q4tp_row_1x4_avx2; caller guarantees the ISA.
5090    unsafe {
5091        use core::arch::x86_64::*;
5092        let lomask = _mm256_set1_epi8(0x0F);
5093        let eight = _mm256_set1_epi8(8);
5094        let zero = _mm512_setzero_si512();
5095        let (mut v0, mut v1, mut v2, mut v3) = (
5096            _mm512_setzero_ps(),
5097            _mm512_setzero_ps(),
5098            _mm512_setzero_ps(),
5099            _mm512_setzero_ps(),
5100        );
5101        let pairs = gpr / 2;
5102        for gp in 0..pairs {
5103            let gi = gp * 2;
5104            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5105            let bb = _mm256_loadu_si256(t as *const __m256i);
5106            let lo = _mm256_and_si256(bb, lomask);
5107            let hi = _mm256_and_si256(_mm256_srli_epi16::<4>(bb), lomask);
5108            // `unpack` works per 128-bit lane, so the halves come out as
5109            // [A.lo, B.lo] and [A.hi, B.hi]; the shuffle reorders the four
5110            // 128-bit lanes into the weights' natural order, which is what
5111            // the straight activation load expects.
5112            let ul = _mm256_sub_epi8(_mm256_unpacklo_epi8(lo, hi), eight);
5113            let uh = _mm256_sub_epi8(_mm256_unpackhi_epi8(lo, hi), eight);
5114            let cat = _mm512_inserti64x4::<1>(_mm512_castsi256_si512(ul), uh);
5115            let w = _mm512_shuffle_i64x2::<0b11_01_10_00>(cat, cat);
5116            let wabs = _mm512_abs_epi8(w);
5117            let neg = _mm512_movepi8_mask(w);
5118            let off = gi * GROUP_SIZE;
5119            let sv = _mm512_insertf32x8::<1>(
5120                _mm512_castps256_ps512(_mm256_set1_ps(*scales.get_unchecked(gi))),
5121                _mm256_set1_ps(*scales.get_unchecked(gi + 1)),
5122            );
5123            let dot = |x: &[i8]| -> __m512 {
5124                let xv = _mm512_loadu_si512(x.as_ptr().add(off) as *const __m512i);
5125                let sx = _mm512_mask_sub_epi8(xv, neg, zero, xv);
5126                _mm512_cvtepi32_ps(_mm512_dpbusd_epi32(zero, wabs, sx))
5127            };
5128            v0 = _mm512_fmadd_ps(dot(xs[0]), sv, v0);
5129            v1 = _mm512_fmadd_ps(dot(xs[1]), sv, v1);
5130            v2 = _mm512_fmadd_ps(dot(xs[2]), sv, v2);
5131            v3 = _mm512_fmadd_ps(dot(xs[3]), sv, v3);
5132        }
5133        let mut acc = [
5134            _mm512_reduce_add_ps(v0),
5135            _mm512_reduce_add_ps(v1),
5136            _mm512_reduce_add_ps(v2),
5137            _mm512_reduce_add_ps(v3),
5138        ];
5139        // An odd group count leaves one group over; the narrow kernel
5140        // finishes it rather than the tail being a special case here.
5141        if gpr % 2 == 1 {
5142            let off = (gpr - 1) * GROUP_SIZE;
5143            for j in off..off + GROUP_SIZE {
5144                let (w, s) = q4tp_outlier(nib, r, gpr, j, scales);
5145                let ws = w * s;
5146                for k in 0..4 {
5147                    acc[k] += ws * *xs[k].get_unchecked(j) as f32;
5148                }
5149            }
5150        }
5151        acc
5152    }
5153}
5154
5155/// Four batch columns against one q4tp row: the tile is unpacked ONCE and
5156/// spent on four activation streams, which is where a prefill batch stops
5157/// being weight-bandwidth-bound. Twin of `dot_q4t_row_1x4_sdot`.
5158#[cfg(target_arch = "aarch64")]
5159#[target_feature(enable = "neon,dotprod")]
5160unsafe fn dot_q4tp_row_1x4_sdot(
5161    nib: &[u8],
5162    r: usize,
5163    gpr: usize,
5164    xs: [&[i8]; 4],
5165    scales: &[f32],
5166) -> [f32; 4] {
5167    // SAFETY: see dot_q4tp_row_sdot; every xs[k] is gpr·GROUP_SIZE long.
5168    unsafe {
5169        use core::arch::aarch64::*;
5170        use core::arch::asm;
5171        let lomask = vdupq_n_u8(0x0F);
5172        let eight = vdupq_n_s8(8);
5173        // Named accumulators, NOT an array indexed by a loop variable: the
5174        // latter does not stay in registers (the same defect cost 2x in the
5175        // AVX2 q4t kernel and again in WGSL).
5176        //
5177        // They are VECTORS, and the horizontal add happens once at the end
5178        // instead of once per group per column. `vaddvq` is a cross-lane
5179        // reduction — with 72 groups and four columns the old shape paid
5180        // 288 of them per row, each one a dependency stall the pipeline
5181        // cannot hide, to save four float adds. The group's scale now
5182        // rides an fma into the lane accumulators, so the arithmetic per
5183        // group is one convert and one fma. Summation order changes (the
5184        // lanes carry independent partial sums), which is the same
5185        // round-off class the SDOT path already lives in — the strict
5186        // kernel (`CMF_SDOT=0`, what `cortiq ppl` runs) is unchanged and
5187        // stays the reference.
5188        let (mut v0, mut v1, mut v2, mut v3) = (
5189            vdupq_n_f32(0.0),
5190            vdupq_n_f32(0.0),
5191            vdupq_n_f32(0.0),
5192            vdupq_n_f32(0.0),
5193        );
5194        for gi in 0..gpr {
5195            let t = nib.as_ptr().add((r * gpr + gi) * Q4TP_NIB);
5196            let s = *scales.get_unchecked(gi);
5197            let bb = vld1q_u8(t);
5198            let lo = vandq_u8(bb, lomask);
5199            let hi = vshrq_n_u8::<4>(bb);
5200            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
5201            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
5202            let off = gi * GROUP_SIZE;
5203            let dot4 = |x: &[i8]| -> int32x4_t {
5204                let x0 = vld1q_s8(x.as_ptr().add(off));
5205                let x1 = vld1q_s8(x.as_ptr().add(off + 16));
5206                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
5207                asm!(
5208                    "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
5209                    "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
5210                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
5211                    e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
5212                    options(pure, nomem, nostack),
5213                );
5214                vaddq_s32(a0, a1)
5215            };
5216            v0 = vfmaq_n_f32(v0, vcvtq_f32_s32(dot4(xs[0])), s);
5217            v1 = vfmaq_n_f32(v1, vcvtq_f32_s32(dot4(xs[1])), s);
5218            v2 = vfmaq_n_f32(v2, vcvtq_f32_s32(dot4(xs[2])), s);
5219            v3 = vfmaq_n_f32(v3, vcvtq_f32_s32(dot4(xs[3])), s);
5220        }
5221        [
5222            vaddvq_f32(v0),
5223            vaddvq_f32(v1),
5224            vaddvq_f32(v2),
5225            vaddvq_f32(v3),
5226        ]
5227    }
5228}
5229
5230/// Fused q4tp matmat — the same three arms `q4t_matmat` has. Shipping only
5231/// the scalar one made Nanbeige-3B decode at 1.2 tok/s against q4t's 5.9:
5232/// the format was fine, the missing arms were the whole regression.
5233fn q4tp_matmat(
5234    bytes: &[u8],
5235    xs_all: &[f32],
5236    b: usize,
5237    rows: usize,
5238    cols: usize,
5239    out: &mut [f32],
5240    pool: Option<&Pool>,
5241) {
5242    debug_assert_eq!(out.len(), b * rows);
5243    let gpr = cols / GROUP_SIZE;
5244    let v = Q4tpView::new(bytes, rows, cols);
5245
5246    // Wide batches ride the AMX through a dequant-tile sgemm, as in q4t.
5247    #[cfg(target_os = "macos")]
5248    if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
5249        dequant_matmat_accel(
5250            &|r, dst| {
5251                let mut sc = [0f32; 32];
5252                let mut scv;
5253                let s: &[f32] = if gpr <= 32 {
5254                    v.scales_into(r, gpr, &mut sc);
5255                    &sc[..gpr]
5256                } else {
5257                    scv = vec![0f32; gpr];
5258                    v.scales_into(r, gpr, &mut scv);
5259                    &scv
5260                };
5261                for gi in 0..gpr {
5262                    let tile = &v.nib[(r * gpr + gi) * Q4TP_NIB..(r * gpr + gi + 1) * Q4TP_NIB];
5263                    for (k, &bb) in tile.iter().enumerate() {
5264                        dst[gi * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s[gi];
5265                        dst[gi * GROUP_SIZE + k * 2 + 1] =
5266                            (((bb >> 4) & 0x0F) as f32 - 8.0) * s[gi];
5267                    }
5268                }
5269            },
5270            xs_all,
5271            b,
5272            rows,
5273            cols,
5274            out,
5275            pool,
5276        );
5277        return;
5278    }
5279
5280    let out_addr = SendMut(out.as_mut_ptr());
5281    if a8w8_enabled() {
5282        let acts: Vec<SplitAct> = (0..b)
5283            .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
5284            .collect();
5285        let acts = &acts;
5286        #[cfg(target_arch = "aarch64")]
5287        let blocked_ok = sdot_enabled() && blocked_enabled();
5288        // x86 gets the same blocking: one tile unpack spent on four
5289        // columns. Without it every column re-decoded the row, which is
5290        // why a 48-core EPYC measured a sixth of an M4's per-core rate.
5291        // The gate is `avx2_enabled`, as in q4t — `sdot_enabled` answers
5292        // for ARM's dotprod and is hard-wired false everywhere else, so
5293        // asking it here left the whole blocked path unreachable on x86.
5294        #[cfg(target_arch = "x86_64")]
5295        let blocked_ok = q4tp_blocked_x86();
5296        #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
5297        let blocked_ok = false;
5298        // Columns are swept in panels that fit L2. Without this a
5299        // row-pair walks every activation in the batch — 4.8 MB at
5300        // 512x512 — and does it again for the next pair, so the whole
5301        // batch streams out of the shared cache once per row. Measured
5302        // 800 GB/s of it, flat across batch sizes, which is the signature
5303        // of a loop bound by traffic rather than by arithmetic. A panel of
5304        // 256 columns is 590 KB beside 221 KB of this worker's weights:
5305        // both stay resident and the batch crosses L3 once instead of
5306        // once per row.
5307        let panel_cols: usize = std::env::var("CMF_Q4TP_PANEL")
5308            .ok()
5309            .and_then(|v| v.parse().ok())
5310            .filter(|v| *v > 0)
5311            .unwrap_or(256);
5312        let run = |start: usize, end: usize| {
5313            for abase in (0..acts.len()).step_by(panel_cols) {
5314                let alen = (acts.len() - abase).min(panel_cols);
5315                let mut sc = vec![0f32; gpr];
5316                #[cfg(target_arch = "x86_64")]
5317                let mut r_lo = start;
5318                #[cfg(target_arch = "x86_64")]
5319                if blocked_ok && alen >= 8 {
5320                    let mut sc1 = vec![0f32; gpr];
5321                    while r_lo + 2 <= end {
5322                        v.scales_into(r_lo, gpr, &mut sc);
5323                        v.scales_into(r_lo + 1, gpr, &mut sc1);
5324                        let mut bi = 0usize;
5325                        while bi + 8 <= alen {
5326                            let xs = [
5327                                acts[abase + bi].xq.as_slice(),
5328                                acts[abase + bi + 1].xq.as_slice(),
5329                                acts[abase + bi + 2].xq.as_slice(),
5330                                acts[abase + bi + 3].xq.as_slice(),
5331                                acts[abase + bi + 4].xq.as_slice(),
5332                                acts[abase + bi + 5].xq.as_slice(),
5333                                acts[abase + bi + 6].xq.as_slice(),
5334                                acts[abase + bi + 7].xq.as_slice(),
5335                            ];
5336                            let d =
5337                                unsafe { dot_q4tp_2x8_avx512(v.nib, r_lo, gpr, xs, &sc, &sc1) };
5338                            for (row, dr, scr) in
5339                                [(r_lo, &d[0], &sc), (r_lo + 1, &d[1], &sc1)]
5340                            {
5341                                for k in 0..8 {
5342                                    let act = &acts[abase + bi + k];
5343                                    let mut acc = dr[k] * act.sx;
5344                                    for &(j, xv) in &act.outliers {
5345                                        let (w, s) = q4tp_outlier(v.nib, row, gpr, j, scr);
5346                                        acc += w * s * xv;
5347                                    }
5348                                    // SAFETY: disjoint (bi, r) cells per worker.
5349                                    unsafe { *out_addr.at((abase + bi + k) * rows + row) = acc };
5350                                }
5351                            }
5352                            bi += 8;
5353                        }
5354                        // Columns past the last group of eight, both rows —
5355                        // the same single-row kernel the tail below uses.
5356                        for row in [r_lo, r_lo + 1] {
5357                            let scr: &[f32] = if row == r_lo { &sc } else { &sc1 };
5358                            for b2 in bi..alen {
5359                                let act = &acts[abase + b2];
5360                                let xs4 = [
5361                                    act.xq.as_slice(),
5362                                    act.xq.as_slice(),
5363                                    act.xq.as_slice(),
5364                                    act.xq.as_slice(),
5365                                ];
5366                                let d =
5367                                    unsafe { dot_q4tp_row_1x4_avx512(v.nib, row, gpr, xs4, scr) };
5368                                let mut acc = d[0] * act.sx;
5369                                for &(j, xv) in &act.outliers {
5370                                    let (w, s) = q4tp_outlier(v.nib, row, gpr, j, scr);
5371                                    acc += w * s * xv;
5372                                }
5373                                // SAFETY: disjoint (bi, r) cells per worker.
5374                                unsafe { *out_addr.at((abase + b2) * rows + row) = acc };
5375                            }
5376                        }
5377                        r_lo += 2;
5378                    }
5379                }
5380                #[cfg(target_arch = "x86_64")]
5381                let row_start = r_lo;
5382                #[cfg(not(target_arch = "x86_64"))]
5383                let row_start = start;
5384                for r in row_start..end {
5385                    v.scales_into(r, gpr, &mut sc);
5386                    let mut bi = 0usize;
5387                    #[cfg(target_arch = "x86_64")]
5388                    if blocked_ok {
5389                        while bi + 8 <= alen {
5390                            let xs = [
5391                                acts[abase + bi].xq.as_slice(),
5392                                acts[abase + bi + 1].xq.as_slice(),
5393                                acts[abase + bi + 2].xq.as_slice(),
5394                                acts[abase + bi + 3].xq.as_slice(),
5395                                acts[abase + bi + 4].xq.as_slice(),
5396                                acts[abase + bi + 5].xq.as_slice(),
5397                                acts[abase + bi + 6].xq.as_slice(),
5398                                acts[abase + bi + 7].xq.as_slice(),
5399                            ];
5400                            let d = unsafe { dot_q4tp_row_1x8_avx512(v.nib, r, gpr, xs, &sc) };
5401                            for k in 0..8 {
5402                                let act = &acts[abase + bi + k];
5403                                let mut acc = d[k] * act.sx;
5404                                for &(j, xv) in &act.outliers {
5405                                    let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
5406                                    acc += w * s * xv;
5407                                }
5408                                // SAFETY: disjoint (bi, r) cells per worker.
5409                                unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
5410                            }
5411                            bi += 8;
5412                        }
5413                        while bi + 4 <= alen {
5414                            let xs = [
5415                                acts[abase + bi].xq.as_slice(),
5416                                acts[abase + bi + 1].xq.as_slice(),
5417                                acts[abase + bi + 2].xq.as_slice(),
5418                                acts[abase + bi + 3].xq.as_slice(),
5419                            ];
5420                            let d = unsafe { dot_q4tp_row_1x4_avx512(v.nib, r, gpr, xs, &sc) };
5421                            for k in 0..4 {
5422                                let act = &acts[abase + bi + k];
5423                                let mut acc = d[k] * act.sx;
5424                                for &(j, xv) in &act.outliers {
5425                                    let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
5426                                    acc += w * s * xv;
5427                                }
5428                                // SAFETY: disjoint (bi, r) cells per worker.
5429                                unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
5430                            }
5431                            bi += 4;
5432                        }
5433                    }
5434                    #[cfg(target_arch = "aarch64")]
5435                    if blocked_ok {
5436                        while bi + 4 <= alen {
5437                            let xs = [
5438                                acts[abase + bi].xq.as_slice(),
5439                                acts[abase + bi + 1].xq.as_slice(),
5440                                acts[abase + bi + 2].xq.as_slice(),
5441                                acts[abase + bi + 3].xq.as_slice(),
5442                            ];
5443                            let d = unsafe {
5444                                if q4tp_v1() {
5445                                    dot_q4tp_row_1x4_sdot_v1(v.nib, r, gpr, xs, &sc)
5446                                } else {
5447                                    dot_q4tp_row_1x4_sdot(v.nib, r, gpr, xs, &sc)
5448                                }
5449                            };
5450                            for k in 0..4 {
5451                                let act = &acts[abase + bi + k];
5452                                let mut acc = d[k] * act.sx;
5453                                for &(j, xv) in &act.outliers {
5454                                    let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
5455                                    acc += w * s * xv;
5456                                }
5457                                // SAFETY: disjoint (bi, r) cells per worker.
5458                                unsafe { *out_addr.at((abase + bi + k) * rows + r) = acc };
5459                            }
5460                            bi += 4;
5461                        }
5462                    }
5463                    let _ = blocked_ok;
5464                    while bi < alen {
5465                        let act = &acts[abase + bi];
5466                        let mut acc = dot_q4tp_row_i8(v.nib, r, gpr, &act.xq, &sc) * act.sx;
5467                        for &(j, xv) in &act.outliers {
5468                            let (w, s) = q4tp_outlier(v.nib, r, gpr, j, &sc);
5469                            acc += w * s * xv;
5470                        }
5471                        // SAFETY: disjoint (bi, r) cells per worker range.
5472                        unsafe { *out_addr.at((abase + bi) * rows + r) = acc };
5473                        bi += 1;
5474                    }
5475                }
5476        
5477            }
5478        };
5479        dispatch_rows(pool, rows, &run);
5480        return;
5481    }
5482
5483    let run = |start: usize, end: usize| {
5484        let mut sc = vec![0f32; gpr];
5485        for r in start..end {
5486            v.scales_into(r, gpr, &mut sc);
5487            for bi in 0..b {
5488                let x = &xs_all[bi * cols..(bi + 1) * cols];
5489                // SAFETY: disjoint (bi, r) cells per worker range.
5490                unsafe { *out_addr.at(bi * rows + r) = q4tp_row_exact(v.nib, r, gpr, x, &sc) };
5491            }
5492        }
5493    };
5494    dispatch_rows(pool, rows, &run);
5495}
5496
5497/// Fused q4_tiled matvec (dispatch mirrors `q4matvec`).
5498fn q4t_matvec(
5499    bytes: &[u8],
5500    x: &[f32],
5501    rows: usize,
5502    cols: usize,
5503    out: &mut [f32],
5504    pool: Option<&Pool>,
5505) {
5506    debug_assert_eq!(out.len(), rows);
5507    let gpr = cols / GROUP_SIZE;
5508    let out_addr = SendMut(out.as_mut_ptr());
5509    if a8w8_enabled() {
5510        let act = split_act(x);
5511        let run = move |start: usize, end: usize| {
5512            for r in start..end {
5513                let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
5514                for &(j, xv) in &act.outliers {
5515                    let (w, s) = q4t_outlier(bytes, r, gpr, j);
5516                    acc += w * s * xv;
5517                }
5518                // SAFETY: disjoint row ranges per worker.
5519                unsafe { *out_addr.at(r) = acc };
5520            }
5521        };
5522        dispatch_rows(pool, rows, &run);
5523        return;
5524    }
5525    let run = move |start: usize, end: usize| {
5526        for r in start..end {
5527            // SAFETY: disjoint row ranges per worker.
5528            unsafe { *out_addr.at(r) = q4t_row_exact(bytes, r, gpr, x) };
5529        }
5530    };
5531    dispatch_rows(pool, rows, &run);
5532}
5533
5534/// Fused two-input q4_tiled matvec (weights read once per pair).
5535#[allow(clippy::too_many_arguments)]
5536fn q4t_matvec2(
5537    bytes: &[u8],
5538    x1: &[f32],
5539    x2: &[f32],
5540    rows: usize,
5541    cols: usize,
5542    o1: &mut [f32],
5543    o2: &mut [f32],
5544    pool: Option<&Pool>,
5545) {
5546    let gpr = cols / GROUP_SIZE;
5547    let p1 = SendMut(o1.as_mut_ptr());
5548    let p2 = SendMut(o2.as_mut_ptr());
5549    if a8w8_enabled() {
5550        let a1 = split_act(x1);
5551        let a2 = split_act(x2);
5552        let run = move |start: usize, end: usize| {
5553            for r in start..end {
5554                let mut v1 = dot_q4t_row_i8(bytes, r, gpr, &a1.xq) * a1.sx;
5555                let mut v2 = dot_q4t_row_i8(bytes, r, gpr, &a2.xq) * a2.sx;
5556                for &(j, xv) in &a1.outliers {
5557                    let (w, s) = q4t_outlier(bytes, r, gpr, j);
5558                    v1 += w * s * xv;
5559                }
5560                for &(j, xv) in &a2.outliers {
5561                    let (w, s) = q4t_outlier(bytes, r, gpr, j);
5562                    v2 += w * s * xv;
5563                }
5564                // SAFETY: disjoint row ranges per worker.
5565                unsafe {
5566                    *p1.at(r) = v1;
5567                    *p2.at(r) = v2;
5568                }
5569            }
5570        };
5571        dispatch_rows(pool, rows, &run);
5572        return;
5573    }
5574    let run = move |start: usize, end: usize| {
5575        for r in start..end {
5576            // SAFETY: disjoint row ranges per worker.
5577            unsafe {
5578                *p1.at(r) = q4t_row_exact(bytes, r, gpr, x1);
5579                *p2.at(r) = q4t_row_exact(bytes, r, gpr, x2);
5580            }
5581        }
5582    };
5583    dispatch_rows(pool, rows, &run);
5584}
5585
5586/// Batched q4_tiled matmat: each row's tiles stream once per microbatch.
5587#[allow(clippy::too_many_arguments)]
5588/// Prefill GEMM through Accelerate for group-quantized codecs: a
5589/// caller-supplied row dequantizer fills f32 tiles (pool-parallel) and
5590/// each tile rides the AMX with one sgemm — the generic sibling of
5591/// `qmatmat_accel` (q8). Numerics are f32-GEMM (tolerance class);
5592/// decode (b=1) never takes this path.
5593#[cfg(target_os = "macos")]
5594fn dequant_matmat_accel(
5595    dequant_row: &(dyn Fn(usize, &mut [f32]) + Sync),
5596    xs_all: &[f32],
5597    b: usize,
5598    rows: usize,
5599    cols: usize,
5600    out: &mut [f32],
5601    pool: Option<&Pool>,
5602) {
5603    const TR: usize = 2048;
5604    thread_local! {
5605        static WTILE: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
5606    }
5607    WTILE.with(|wt| {
5608        let mut wtile = wt.borrow_mut();
5609        wtile.resize(TR * cols, 0.0);
5610        let mut r0 = 0usize;
5611        while r0 < rows {
5612            let tr = TR.min(rows - r0);
5613            let wt_addr = SendMut(wtile.as_mut_ptr());
5614            let run = |start: usize, end: usize| {
5615                for r in start..end {
5616                    // SAFETY: workers cover disjoint r ranges.
5617                    let dst = unsafe { std::slice::from_raw_parts_mut(wt_addr.at(r * cols), cols) };
5618                    dequant_row(r0 + r, dst);
5619                }
5620            };
5621            dispatch_rows(pool, tr, &run);
5622            unsafe {
5623                accel_blas::cblas_sgemm(
5624                    101, // RowMajor
5625                    111, // NoTrans A
5626                    112, // Trans B
5627                    b as i32,
5628                    tr as i32,
5629                    cols as i32,
5630                    1.0,
5631                    xs_all.as_ptr(),
5632                    cols as i32,
5633                    wtile.as_ptr(),
5634                    cols as i32,
5635                    0.0,
5636                    out.as_mut_ptr().add(r0),
5637                    rows as i32,
5638                );
5639            }
5640            r0 += tr;
5641        }
5642    });
5643}
5644
5645fn q4t_matmat(
5646    bytes: &[u8],
5647    xs_all: &[f32],
5648    b: usize,
5649    rows: usize,
5650    cols: usize,
5651    out: &mut [f32],
5652    pool: Option<&Pool>,
5653) {
5654    debug_assert_eq!(out.len(), b * rows);
5655    let gpr = cols / GROUP_SIZE;
5656    // Wide batches ride the AMX like q8's qmatmat: on Apple silicon
5657    // the dequant-tile sgemm is an order above the SDOT row loop for
5658    // prefill shapes (imagegen DiT forwards are exactly this).
5659    #[cfg(target_os = "macos")]
5660    if b >= 8 && rows * cols >= 500_000 && accel_gemm_enabled() {
5661        dequant_matmat_accel(
5662            &|r, dst| {
5663                for gi in 0..gpr {
5664                    let tile = &bytes[(r * gpr + gi) * Q4_TILE..(r * gpr + gi + 1) * Q4_TILE];
5665                    let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
5666                    for (k, &bb) in tile[2..].iter().enumerate() {
5667                        dst[gi * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s;
5668                        dst[gi * GROUP_SIZE + k * 2 + 1] = (((bb >> 4) & 0x0F) as f32 - 8.0) * s;
5669                    }
5670                }
5671            },
5672            xs_all,
5673            b,
5674            rows,
5675            cols,
5676            out,
5677            pool,
5678        );
5679        return;
5680    }
5681    let out_addr = SendMut(out.as_mut_ptr());
5682    if a8w8_enabled() {
5683        let acts: Vec<SplitAct> = (0..b)
5684            .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
5685            .collect();
5686        let acts = &acts;
5687        #[cfg(target_arch = "x86_64")]
5688        let blocked_ok = avx2_enabled()
5689            && blocked_enabled();
5690        #[cfg(target_arch = "aarch64")]
5691        let blocked_ok = sdot_enabled()
5692            && blocked_enabled();
5693        #[cfg(not(any(target_arch = "x86_64", target_arch = "aarch64")))]
5694        let blocked_ok = false;
5695        let run = move |start: usize, end: usize| {
5696            for r in start..end {
5697                let mut bi = 0usize;
5698                #[cfg(target_arch = "aarch64")]
5699                if blocked_ok {
5700                    while bi + 4 <= acts.len() {
5701                        let xs = [
5702                            acts[bi].xq.as_slice(),
5703                            acts[bi + 1].xq.as_slice(),
5704                            acts[bi + 2].xq.as_slice(),
5705                            acts[bi + 3].xq.as_slice(),
5706                        ];
5707                        let d = unsafe { dot_q4t_row_1x4_sdot(bytes, r, gpr, xs) };
5708                        for k in 0..4 {
5709                            let act = &acts[bi + k];
5710                            let mut acc = d[k] * act.sx;
5711                            for &(j, xv) in &act.outliers {
5712                                let (w, sc) = q4t_outlier(bytes, r, gpr, j);
5713                                acc += w * sc * xv;
5714                            }
5715                            // SAFETY: disjoint (bi, r) cells per worker.
5716                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
5717                        }
5718                        bi += 4;
5719                    }
5720                }
5721                #[cfg(target_arch = "x86_64")]
5722                if blocked_ok {
5723                    while bi + 4 <= acts.len() {
5724                        let xs = [
5725                            acts[bi].xq.as_slice(),
5726                            acts[bi + 1].xq.as_slice(),
5727                            acts[bi + 2].xq.as_slice(),
5728                            acts[bi + 3].xq.as_slice(),
5729                        ];
5730                        let d = unsafe {
5731                            if vnni_tiles_enabled() {
5732                                dot_q4t_row_1x4_vnni(bytes, r, gpr, xs)
5733                            } else {
5734                                dot_q4t_row_1x4_avx2(bytes, r, gpr, xs)
5735                            }
5736                        };
5737                        for k in 0..4 {
5738                            let act = &acts[bi + k];
5739                            let mut acc = d[k] * act.sx;
5740                            for &(j, xv) in &act.outliers {
5741                                let (w, sc) = q4t_outlier(bytes, r, gpr, j);
5742                                acc += w * sc * xv;
5743                            }
5744                            // SAFETY: disjoint (bi, r) cells per worker.
5745                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
5746                        }
5747                        bi += 4;
5748                    }
5749                }
5750                let _ = blocked_ok;
5751                while bi < acts.len() {
5752                    let act = &acts[bi];
5753                    let mut acc = dot_q4t_row_i8(bytes, r, gpr, &act.xq) * act.sx;
5754                    for &(j, xv) in &act.outliers {
5755                        let (w, s) = q4t_outlier(bytes, r, gpr, j);
5756                        acc += w * s * xv;
5757                    }
5758                    // SAFETY: disjoint (bi, r) cells per worker range.
5759                    unsafe { *out_addr.at(bi * rows + r) = acc };
5760                    bi += 1;
5761                }
5762            }
5763        };
5764        dispatch_rows(pool, rows, &run);
5765        return;
5766    }
5767    let run = move |start: usize, end: usize| {
5768        for r in start..end {
5769            for bi in 0..b {
5770                let x = &xs_all[bi * cols..(bi + 1) * cols];
5771                // SAFETY: disjoint (bi, r) cells per worker range.
5772                unsafe { *out_addr.at(bi * rows + r) = q4t_row_exact(bytes, r, gpr, x) };
5773            }
5774        }
5775    };
5776    dispatch_rows(pool, rows, &run);
5777}
5778
5779// ── q1 (dtype 12): binary weights, [f16 scale][4B sign bits] per
5780// 32-group tile. The kernel family mirrors q4_tiled: one sequential
5781// stream of 6-byte tiles, per-tile integer dot × scale, exact outlier
5782// correction (A8W8 contract), exact scalar path under CMF_SDOT=0. ──
5783
5784/// Per-32-group sums of the quantized activation — the ±1 identity's
5785/// shared half: `dot = −2·sdot(mask, x) − gsum[g]`, computed ONCE per
5786/// matvec and reused by every row.
5787fn q1_group_sums(xq: &[i8], gpr: usize) -> Vec<i32> {
5788    (0..gpr)
5789        .map(|gi| {
5790            xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE]
5791                .iter()
5792                .map(|&v| v as i32)
5793                .sum()
5794        })
5795        .collect()
5796}
5797
5798/// One q1 row via the A8W8 int8 path — mask-SDOT on ARM (no ±1
5799/// expansion at all), scalar bit loop elsewhere (AVX2 queued with the
5800/// x86 pass).
5801#[inline]
5802#[allow(unreachable_code)]
5803/// AVX2 q1 row via the same ±1 identity as the ARM sdot kernel: the
5804/// sign bits expand to a {0, −1} byte mask through shuffle+cmpeq, the
5805/// masked activation sums through maddubs(1, x&mask), and
5806/// `dot = −(2·masked_sum + Σx_group)` — bit-identical integer math.
5807#[cfg(target_arch = "x86_64")]
5808#[target_feature(enable = "avx2")]
5809unsafe fn dot_q1_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
5810    // SAFETY: callers uphold the 6B-tile and xq/gsum length contracts.
5811    unsafe {
5812        use core::arch::x86_64::*;
5813        // Byte j of the mask must replicate bits-byte j/8.
5814        let expand = _mm256_setr_epi8(
5815            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,
5816            3, 3, 3,
5817        );
5818        let bitsel = _mm256_setr_epi8(
5819            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
5820            -128, 1, 2, 4, 8, 16, 32, 64, -128,
5821        );
5822        let ones8 = _mm256_set1_epi8(1);
5823        let ones16 = _mm256_set1_epi16(1);
5824        let mut acc = 0f32;
5825        for gi in 0..gpr {
5826            let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
5827            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5828            let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
5829            let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
5830            let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
5831            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5832            let sel = _mm256_and_si256(x, mask);
5833            // Σ of selected i8 lanes: maddubs(1u8, sel_i8) pairs → madd.
5834            let p16 = _mm256_maddubs_epi16(ones8, sel);
5835            let d32 = _mm256_madd_epi16(p16, ones16);
5836            let hi128 = _mm256_extracti128_si256::<1>(d32);
5837            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
5838            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
5839            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
5840            let msum = _mm_cvtsi128_si32(s32);
5841            // The and-select keeps x UN-negated (unlike ARM's −1-mask
5842            // sdot): d = Σ_set − Σ_unset = 2·Σ_set − Σ_all.
5843            let d = 2 * msum - gsum[gi];
5844            acc += d as f32 * s;
5845        }
5846        acc
5847    }
5848}
5849
5850/// VNNI twin of `dot_q1_row_avx2`: the masked-select sum goes through
5851/// one `vpdpbusd(1u8, sel)` (see `dpbusd_hsum` — bit-identical).
5852#[cfg(target_arch = "x86_64")]
5853#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
5854unsafe fn dot_q1_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
5855    // SAFETY: callers uphold the 6B-tile and xq/gsum length contracts.
5856    unsafe {
5857        use core::arch::x86_64::*;
5858        let expand = _mm256_setr_epi8(
5859            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,
5860            3, 3, 3,
5861        );
5862        let bitsel = _mm256_setr_epi8(
5863            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
5864            -128, 1, 2, 4, 8, 16, 32, 64, -128,
5865        );
5866        let ones8 = _mm256_set1_epi8(1);
5867        let mut acc = 0f32;
5868        for gi in 0..gpr {
5869            let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
5870            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5871            let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
5872            let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
5873            let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
5874            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5875            let msum = dpbusd_hsum(ones8, _mm256_and_si256(x, mask));
5876            let d = 2 * msum - gsum[gi];
5877            acc += d as f32 * s;
5878        }
5879        acc
5880    }
5881}
5882
5883/// VNNI twin of `dot_q1_row_1x4_avx2` (see `dpbusd_hsum`).
5884#[cfg(target_arch = "x86_64")]
5885#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
5886unsafe fn dot_q1_row_1x4_vnni(
5887    bytes: &[u8],
5888    r: usize,
5889    gpr: usize,
5890    xs: [&[i8]; 4],
5891    gsums: [&[i32]; 4],
5892) -> [f32; 4] {
5893    // SAFETY: callers uphold the 6B-tile and xq/gsum length contracts.
5894    unsafe {
5895        use core::arch::x86_64::*;
5896        let expand = _mm256_setr_epi8(
5897            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,
5898            3, 3, 3,
5899        );
5900        let bitsel = _mm256_setr_epi8(
5901            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
5902            -128, 1, 2, 4, 8, 16, 32, 64, -128,
5903        );
5904        let ones8 = _mm256_set1_epi8(1);
5905        let mut acc = [0f32; 4];
5906        for gi in 0..gpr {
5907            let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
5908            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5909            let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
5910            let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
5911            let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
5912            for (k, xq) in xs.iter().enumerate() {
5913                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5914                let msum = dpbusd_hsum(ones8, _mm256_and_si256(x, mask));
5915                let d = 2 * msum - gsums[k][gi];
5916                acc[k] += d as f32 * s;
5917            }
5918        }
5919        acc
5920    }
5921}
5922
5923/// The blocked 1×4 flavor: the expanded bit mask serves four activation
5924/// streams per group (mask build once, four select+reduce chains).
5925#[cfg(target_arch = "x86_64")]
5926#[target_feature(enable = "avx2")]
5927unsafe fn dot_q1_row_1x4_avx2(
5928    bytes: &[u8],
5929    r: usize,
5930    gpr: usize,
5931    xs: [&[i8]; 4],
5932    gsums: [&[i32]; 4],
5933) -> [f32; 4] {
5934    // SAFETY: callers uphold the 6B-tile and xq/gsum length contracts.
5935    unsafe {
5936        use core::arch::x86_64::*;
5937        let expand = _mm256_setr_epi8(
5938            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,
5939            3, 3, 3,
5940        );
5941        let bitsel = _mm256_setr_epi8(
5942            1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64, -128, 1, 2, 4, 8, 16, 32, 64,
5943            -128, 1, 2, 4, 8, 16, 32, 64, -128,
5944        );
5945        let ones8 = _mm256_set1_epi8(1);
5946        let ones16 = _mm256_set1_epi16(1);
5947        let mut acc = [0f32; 4];
5948        for gi in 0..gpr {
5949            let t = bytes.as_ptr().add((r * gpr + gi) * Q1_TILE);
5950            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
5951            let bits = u32::from_le_bytes([*t.add(2), *t.add(3), *t.add(4), *t.add(5)]);
5952            let bc = _mm256_shuffle_epi8(_mm256_set1_epi32(bits as i32), expand);
5953            let mask = _mm256_cmpeq_epi8(_mm256_and_si256(bc, bitsel), bitsel);
5954            for (k, xq) in xs.iter().enumerate() {
5955                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
5956                let sel = _mm256_and_si256(x, mask);
5957                let p16 = _mm256_maddubs_epi16(ones8, sel);
5958                let d32 = _mm256_madd_epi16(p16, ones16);
5959                let hi128 = _mm256_extracti128_si256::<1>(d32);
5960                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d32), hi128);
5961                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
5962                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
5963                let msum = _mm_cvtsi128_si32(s32);
5964                let d = 2 * msum - gsums[k][gi];
5965                acc[k] += d as f32 * s;
5966            }
5967        }
5968        acc
5969    }
5970}
5971
5972#[allow(unreachable_code)]
5973fn dot_q1_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
5974    #[cfg(target_arch = "aarch64")]
5975    unsafe {
5976        return dot_q1_row_sdot(bytes, r, gpr, xq, gsum);
5977    }
5978    #[cfg(target_arch = "x86_64")]
5979    if avx2_enabled() {
5980        unsafe {
5981            if vnni_tiles_enabled() {
5982                return dot_q1_row_vnni(bytes, r, gpr, xq, gsum);
5983            }
5984            return dot_q1_row_avx2(bytes, r, gpr, xq, gsum);
5985        }
5986    }
5987    let _ = gsum;
5988    let mut acc = 0f32;
5989    for gi in 0..gpr {
5990        let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
5991        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
5992        let mut d = 0i32;
5993        for (j, &b) in tile[2..].iter().enumerate() {
5994            for k in 0..8 {
5995                let w = ((b >> k) & 1) as i32 * 2 - 1;
5996                d += w * xq[gi * GROUP_SIZE + j * 8 + k] as i32;
5997            }
5998        }
5999        acc += d as f32 * s;
6000    }
6001    acc
6002}
6003
6004/// SDOT q1 row via the ±1 identity: the vtst mask (0xFF where the bit
6005/// is set, i.e. −1 as i8) feeds `sdot` DIRECTLY — no expansion to ±1
6006/// lanes at all — and `dot = −(2·sdot(mask, x) + Σx_group)`, with the
6007/// per-group activation sums shared across every row of the matvec.
6008/// Four tiles (128 weights) per iteration: integer dots reduce through
6009/// a vpaddq tree into ONE i32x4 that meets its four scales in a single
6010/// fused f32 multiply-add. Integer math throughout — bit-identical to
6011/// the scalar ±1 reference.
6012#[cfg(target_arch = "aarch64")]
6013#[target_feature(enable = "neon,dotprod")]
6014unsafe fn dot_q1_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8], gsum: &[i32]) -> f32 {
6015    // SAFETY: callers uphold slice-length contracts (6B tile per group,
6016    // xq.len() == gpr·GROUP_SIZE, gsum.len() == gpr).
6017    unsafe {
6018        use core::arch::aarch64::*;
6019        use core::arch::asm;
6020        const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
6021        let m = vld1q_u8(MASKS.as_ptr());
6022        // One tile's −Σ_set(x) as an UNREDUCED i32x4 (two mask-sdots).
6023        macro_rules! tile_dot {
6024            ($t:expr, $x:expr) => {{
6025                let v0 = vcombine_u8(vdup_n_u8(*$t.add(2)), vdup_n_u8(*$t.add(3)));
6026                let v1 = vcombine_u8(vdup_n_u8(*$t.add(4)), vdup_n_u8(*$t.add(5)));
6027                let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
6028                let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
6029                let x0 = vld1q_s8($x);
6030                let x1 = vld1q_s8($x.add(16));
6031                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6032                asm!(
6033                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
6034                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
6035                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6036                    w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
6037                    options(pure, nomem, nostack),
6038                );
6039                vaddq_s32(a0, a1)
6040            }};
6041        }
6042        // TBL unpack over PAIR loads: one vld1q covers two 6B tiles
6043        // ([s s b b b b][s s b b b b] + 4B slack), TBL replicates each
6044        // bit-byte across 8 lanes for vtst, and the four scales gather
6045        // through tbl2 into one fcvtl — the 16 ld1r broadcast loads and
6046        // 4 branchy software f16 conversions per 128 weights (the
6047        // measured load-port wall of this kernel) become 2 vector
6048        // loads + 9 table lookups. Integer math order is unchanged —
6049        // bit-identical results (FCVTL is exact on every f16).
6050        const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
6051        const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
6052        const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
6053        const IW11: [u8; 16] = [
6054            10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
6055        ];
6056        const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
6057        let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
6058        let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
6059        let isc = vld1_u8(ISC.as_ptr());
6060        // One tile's −Σ_set(x) from a TBL-unpacked pair load.
6061        macro_rules! tile_dot_tbl {
6062            ($ld:expr, $i0:expr, $i1:expr, $x:expr) => {{
6063                let w0 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i0), m));
6064                let w1 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8($ld, $i1), m));
6065                let x0 = vld1q_s8($x);
6066                let x1 = vld1q_s8($x.add(16));
6067                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6068                asm!(
6069                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
6070                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
6071                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6072                    w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
6073                    options(pure, nomem, nostack),
6074                );
6075                vaddq_s32(a0, a1)
6076            }};
6077        }
6078        let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
6079        let row_base = r * gpr * Q1_TILE;
6080        let abs_end = bytes.len();
6081        let xp = xq.as_ptr();
6082        let gp = gsum.as_ptr();
6083        let mut accv = vdupq_n_f32(0.0);
6084        let mut gi = 0;
6085        // The second pair load reads 4B past tile gi+3 — stay inside
6086        // the payload slice (only the file's final tiles fall back).
6087        while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
6088            let t0 = base.add(gi * Q1_TILE);
6089            let ld_a = vld1q_u8(t0);
6090            let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
6091            let d0 = tile_dot_tbl!(ld_a, iw00, iw01, xp.add(gi * GROUP_SIZE));
6092            let d1 = tile_dot_tbl!(ld_a, iw10, iw11, xp.add((gi + 1) * GROUP_SIZE));
6093            let d2 = tile_dot_tbl!(ld_b, iw00, iw01, xp.add((gi + 2) * GROUP_SIZE));
6094            let d3 = tile_dot_tbl!(ld_b, iw10, iw11, xp.add((gi + 3) * GROUP_SIZE));
6095            // [−Σ0, −Σ1, −Σ2, −Σ3] → dots = −(2·Σset_neg + gsum)
6096            let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
6097            let g = vld1q_s32(gp.add(gi));
6098            let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
6099            let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
6100            let scf: float32x4_t;
6101            asm!(
6102                "fcvtl {o:v}.4s, {i:v}.4h",
6103                o = out(vreg) scf, i = in(vreg) sc16,
6104                options(pure, nomem, nostack),
6105            );
6106            accv = vfmaq_f32(accv, vcvtq_f32_s32(dots), scf);
6107            gi += 4;
6108        }
6109        let mut acc = vaddvq_f32(accv);
6110        while gi < gpr {
6111            let t = base.add(gi * Q1_TILE);
6112            let s = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
6113            let d = vaddvq_s32(tile_dot!(t, xp.add(gi * GROUP_SIZE)));
6114            acc += (-(2 * d + *gp.add(gi))) as f32 * s;
6115            gi += 1;
6116        }
6117        acc
6118    }
6119}
6120
6121/// Blocked q1 1×4: one TBL unpack of the tile pair serves FOUR
6122/// activation streams (prefill amortization — the same idea as the
6123/// AVX2 twin; per stream the group order, fma order and tail match the
6124/// single-row kernel exactly, so batch == matvec bit-for-bit).
6125#[cfg(target_arch = "aarch64")]
6126#[target_feature(enable = "neon,dotprod")]
6127unsafe fn dot_q1_row_1x4_sdot(
6128    bytes: &[u8],
6129    r: usize,
6130    gpr: usize,
6131    xs: [&[i8]; 4],
6132    gs: [&[i32]; 4],
6133) -> [f32; 4] {
6134    // SAFETY: same slice-length contracts as `dot_q1_row_sdot`, ×4.
6135    unsafe {
6136        use core::arch::aarch64::*;
6137        use core::arch::asm;
6138        const MASKS: [u8; 16] = [1, 2, 4, 8, 16, 32, 64, 128, 1, 2, 4, 8, 16, 32, 64, 128];
6139        const IW00: [u8; 16] = [2, 2, 2, 2, 2, 2, 2, 2, 3, 3, 3, 3, 3, 3, 3, 3];
6140        const IW01: [u8; 16] = [4, 4, 4, 4, 4, 4, 4, 4, 5, 5, 5, 5, 5, 5, 5, 5];
6141        const IW10: [u8; 16] = [8, 8, 8, 8, 8, 8, 8, 8, 9, 9, 9, 9, 9, 9, 9, 9];
6142        const IW11: [u8; 16] = [
6143            10, 10, 10, 10, 10, 10, 10, 10, 11, 11, 11, 11, 11, 11, 11, 11,
6144        ];
6145        const ISC: [u8; 8] = [0, 1, 6, 7, 16, 17, 22, 23];
6146        let m = vld1q_u8(MASKS.as_ptr());
6147        let (iw00, iw01) = (vld1q_u8(IW00.as_ptr()), vld1q_u8(IW01.as_ptr()));
6148        let (iw10, iw11) = (vld1q_u8(IW10.as_ptr()), vld1q_u8(IW11.as_ptr()));
6149        let isc = vld1_u8(ISC.as_ptr());
6150        macro_rules! sdot2 {
6151            ($w0:expr, $w1:expr, $x:expr) => {{
6152                let x0 = vld1q_s8($x);
6153                let x1 = vld1q_s8($x.add(16));
6154                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6155                asm!(
6156                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
6157                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
6158                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6159                    w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
6160                    options(pure, nomem, nostack),
6161                );
6162                vaddq_s32(a0, a1)
6163            }};
6164        }
6165        let base = bytes.as_ptr().add(r * gpr * Q1_TILE);
6166        let row_base = r * gpr * Q1_TILE;
6167        let abs_end = bytes.len();
6168        let mut accv = [vdupq_n_f32(0.0); 4];
6169        let mut gi = 0;
6170        while gi + 4 <= gpr && row_base + (gi + 4) * Q1_TILE + 4 <= abs_end {
6171            let t0 = base.add(gi * Q1_TILE);
6172            let ld_a = vld1q_u8(t0);
6173            let ld_b = vld1q_u8(t0.add(2 * Q1_TILE));
6174            // Unpack ONCE — eight ±mask vectors serve all four streams.
6175            let w00 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw00), m));
6176            let w01 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw01), m));
6177            let w10 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw10), m));
6178            let w11 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_a, iw11), m));
6179            let w20 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw00), m));
6180            let w21 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw01), m));
6181            let w30 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw10), m));
6182            let w31 = vreinterpretq_s8_u8(vtstq_u8(vqtbl1q_u8(ld_b, iw11), m));
6183            let sc16 = vqtbl2_u8(uint8x16x2_t(ld_a, ld_b), isc);
6184            let scf: float32x4_t;
6185            asm!(
6186                "fcvtl {o:v}.4s, {i:v}.4h",
6187                o = out(vreg) scf, i = in(vreg) sc16,
6188                options(pure, nomem, nostack),
6189            );
6190            for k in 0..4 {
6191                let xp = xs[k].as_ptr();
6192                let d0 = sdot2!(w00, w01, xp.add(gi * GROUP_SIZE));
6193                let d1 = sdot2!(w10, w11, xp.add((gi + 1) * GROUP_SIZE));
6194                let d2 = sdot2!(w20, w21, xp.add((gi + 2) * GROUP_SIZE));
6195                let d3 = sdot2!(w30, w31, xp.add((gi + 3) * GROUP_SIZE));
6196                let neg = vpaddq_s32(vpaddq_s32(d0, d1), vpaddq_s32(d2, d3));
6197                let g = vld1q_s32(gs[k].as_ptr().add(gi));
6198                let dots = vnegq_s32(vaddq_s32(vshlq_n_s32::<1>(neg), g));
6199                accv[k] = vfmaq_f32(accv[k], vcvtq_f32_s32(dots), scf);
6200            }
6201            gi += 4;
6202        }
6203        let mut acc = [
6204            vaddvq_f32(accv[0]),
6205            vaddvq_f32(accv[1]),
6206            vaddvq_f32(accv[2]),
6207            vaddvq_f32(accv[3]),
6208        ];
6209        while gi < gpr {
6210            let t = base.add(gi * Q1_TILE);
6211            let sc = f16_to_f32(u16::from_le_bytes([*t, *t.add(1)]));
6212            let v0 = vcombine_u8(vdup_n_u8(*t.add(2)), vdup_n_u8(*t.add(3)));
6213            let v1 = vcombine_u8(vdup_n_u8(*t.add(4)), vdup_n_u8(*t.add(5)));
6214            let w0 = vreinterpretq_s8_u8(vtstq_u8(v0, m));
6215            let w1 = vreinterpretq_s8_u8(vtstq_u8(v1, m));
6216            for k in 0..4 {
6217                let d = vaddvq_s32(sdot2!(w0, w1, xs[k].as_ptr().add(gi * GROUP_SIZE)));
6218                acc[k] += (-(2 * d + *gs[k].as_ptr().add(gi))) as f32 * sc;
6219            }
6220            gi += 1;
6221        }
6222        acc
6223    }
6224}
6225
6226/// (weight ±1, scale) of one q1 element — the exact outlier term.
6227#[inline]
6228fn q1_outlier(bytes: &[u8], r: usize, gpr: usize, j: usize) -> (f32, f32) {
6229    let gi = j / GROUP_SIZE;
6230    let k = j % GROUP_SIZE;
6231    let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
6232    let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
6233    let bit = (tile[2 + k / 8] >> (k % 8)) & 1;
6234    ((bit as i32 * 2 - 1) as f32, s)
6235}
6236
6237/// Exact scalar q1 row (CMF_SDOT=0 contract).
6238#[inline]
6239fn q1_row_exact(bytes: &[u8], r: usize, gpr: usize, x: &[f32]) -> f32 {
6240    let mut acc = 0f32;
6241    for gi in 0..gpr {
6242        let tile = &bytes[(r * gpr + gi) * Q1_TILE..(r * gpr + gi + 1) * Q1_TILE];
6243        let s = f16_to_f32(u16::from_le_bytes([tile[0], tile[1]]));
6244        let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
6245        let mut ga = 0f32;
6246        for (j, &b) in tile[2..].iter().enumerate() {
6247            for k in 0..8 {
6248                ga += (((b >> k) & 1) as f32 * 2.0 - 1.0) * xg[j * 8 + k];
6249            }
6250        }
6251        acc += ga * s;
6252    }
6253    acc
6254}
6255
6256/// One q1 row range via A8W8 (the body of `q1_matvec`'s hot loop,
6257/// extracted so multi-matrix jobs drive the same kernel).
6258#[allow(clippy::too_many_arguments)]
6259fn q1_range_a8w8(
6260    bytes: &[u8],
6261    gpr: usize,
6262    act: &SplitAct,
6263    gsum: &[i32],
6264    out: SendMut,
6265    start: usize,
6266    end: usize,
6267) {
6268    for r in start..end {
6269        let mut acc = dot_q1_row_i8(bytes, r, gpr, &act.xq, gsum) * act.sx;
6270        for &(j, xv) in &act.outliers {
6271            let (w, s) = q1_outlier(bytes, r, gpr, j);
6272            acc += w * s * xv;
6273        }
6274        // SAFETY: disjoint row ranges per worker.
6275        unsafe { *out.at(r) = acc };
6276    }
6277}
6278
6279/// Exact-scalar q1 row range (CMF_SDOT=0 contract).
6280fn q1_range_f32(bytes: &[u8], gpr: usize, x: &[f32], out: SendMut, start: usize, end: usize) {
6281    for r in start..end {
6282        // SAFETY: disjoint row ranges per worker.
6283        unsafe { *out.at(r) = q1_row_exact(bytes, r, gpr, x) };
6284    }
6285}
6286
6287/// q1t per-row overlay locator. After the base (`base_len`) come
6288/// `[u32 row_ptr[rows+1]]` then `[(u16 col, f16 val)]` grouped by row (row
6289/// `r`'s entries are `[row_ptr[r], row_ptr[r+1])`). Returns
6290/// `(row_ptr offset, entries offset, present)`.
6291fn q1t_overlay(bytes: &[u8], base_len: usize, rows: usize) -> (usize, usize, bool) {
6292    let entries = base_len + (rows + 1) * 4;
6293    (base_len, entries, entries <= bytes.len())
6294}
6295
6296/// Read `row_ptr[r]` from the overlay's prefix-sum table.
6297#[inline]
6298fn q1t_rowptr(bytes: &[u8], rp_off: usize, r: usize) -> usize {
6299    let o = rp_off + r * 4;
6300    u32::from_le_bytes([bytes[o], bytes[o + 1], bytes[o + 2], bytes[o + 3]]) as usize
6301}
6302
6303/// Byte → the 5 ternary signs it packs `{−1,0,+1}` as f32, precomputed so
6304/// decoding a q1t code is a table load, not the base-3 divide/modulo per
6305/// weight (division is ~20–40× the cost of a load). Built at compile time.
6306const SIGN5: [[f32; 5]; 256] = {
6307    let mut lut = [[0.0f32; 5]; 256];
6308    let pow3 = [1u16, 3, 9, 27, 81];
6309    let mut byte = 0usize;
6310    while byte < 256 {
6311        let mut i = 0usize;
6312        while i < 5 {
6313            let code = (byte as u16 / pow3[i]) % 3;
6314            lut[byte][i] = if code == 1 {
6315                1.0
6316            } else if code == 2 {
6317                -1.0
6318            } else {
6319                0.0
6320            };
6321            i += 1;
6322        }
6323        byte += 1;
6324    }
6325    lut
6326};
6327
6328/// Same table, as i8 signs — the operand for the int8 SDOT base kernel.
6329const SIGN5_I8: [[i8; 5]; 256] = {
6330    let mut lut = [[0i8; 5]; 256];
6331    let pow3 = [1u16, 3, 9, 27, 81];
6332    let mut byte = 0usize;
6333    while byte < 256 {
6334        let mut i = 0usize;
6335        while i < 5 {
6336            let code = (byte as u16 / pow3[i]) % 3;
6337            lut[byte][i] = if code == 1 {
6338                1
6339            } else if code == 2 {
6340                -1
6341            } else {
6342                0
6343            };
6344            i += 1;
6345        }
6346        byte += 1;
6347    }
6348    lut
6349};
6350
6351/// The same 5 i8 signs packed into a u64 (`[s0 s1 s2 s3 s4 0 0 0]`, LE) so the
6352/// group unpack is 7 unaligned u64 stores at offsets 0,5,10,…,30 instead of
6353/// six 5-byte copies + LUT indexing — each store's trailing zeros are fixed by
6354/// the next store, and the last one runs 6 B past the 32nd weight (the unpack
6355/// buffer is padded to 40). This is the decode/prefill hot inner op.
6356const SIGN5_U64: [u64; 256] = {
6357    let mut lut = [0u64; 256];
6358    let pow3 = [1u16, 3, 9, 27, 81];
6359    let mut byte = 0usize;
6360    while byte < 256 {
6361        let mut v = 0u64;
6362        let mut i = 0usize;
6363        while i < 5 {
6364            let code = (byte as u16 / pow3[i]) % 3;
6365            let s: u8 = if code == 1 {
6366                1
6367            } else if code == 2 {
6368                0xFF
6369            } else {
6370                0
6371            };
6372            v |= (s as u64) << (i * 8);
6373            i += 1;
6374        }
6375        lut[byte] = v;
6376        byte += 1;
6377    }
6378    lut
6379};
6380
6381/// Ternary base weight at `(row r, col j)` = `sign(code)·s_group`. Used to add
6382/// back activation-outlier columns, whose `x` was zeroed for the int8 bulk dot
6383/// (`split_act`). At a weight-outlier position the code is 0, so this is 0 and
6384/// the overlay correction owns that column — no double counting.
6385#[inline]
6386fn q1t_base_weight(bytes: &[u8], r: usize, gpr: usize, j: usize) -> f32 {
6387    const TILE: usize = cortiq_core::quant::Q1T_TILE;
6388    let off = (r * gpr + j / GROUP_SIZE) * TILE;
6389    let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
6390    let within = j % GROUP_SIZE;
6391    SIGN5[bytes[off + 2 + within / 5] as usize][within % 5] * s
6392}
6393
6394/// One 32-group int8 dot via two SDOTs. Bit-exact vs the scalar i8 sum
6395/// (integer accumulation is order-independent).
6396#[cfg(target_arch = "aarch64")]
6397#[target_feature(enable = "neon,dotprod")]
6398#[inline]
6399unsafe fn sdot32_i8(w: *const i8, x: *const i8) -> i32 {
6400    // SAFETY: caller guarantees 32 readable i8 at each pointer.
6401    unsafe {
6402        use core::arch::aarch64::*;
6403        use core::arch::asm;
6404        let w0 = vld1q_s8(w);
6405        let w1 = vld1q_s8(w.add(16));
6406        let x0 = vld1q_s8(x);
6407        let x1 = vld1q_s8(x.add(16));
6408        let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6409        asm!(
6410            "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
6411            "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
6412            a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6413            w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
6414            options(pure, nomem, nostack),
6415        );
6416        vaddvq_s32(vaddq_s32(a0, a1))
6417    }
6418}
6419
6420/// One 32-group int8 dot via AVX2: signed·signed as `maddubs(|w|, sign(x,w))`
6421/// then `madd` and a horizontal reduce (the same idiom as `dot_q4t_row_avx2`).
6422#[cfg(target_arch = "x86_64")]
6423#[target_feature(enable = "avx2")]
6424#[inline]
6425unsafe fn i8dot32_avx2(w: *const i8, x: *const i8) -> i32 {
6426    // SAFETY: caller guarantees 32 readable i8 at each pointer.
6427    unsafe {
6428        use core::arch::x86_64::*;
6429        let wv = _mm256_loadu_si256(w as *const __m256i);
6430        let xv = _mm256_loadu_si256(x as *const __m256i);
6431        let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
6432        let d = _mm256_madd_epi16(p16, _mm256_set1_epi16(1));
6433        let hi128 = _mm256_extracti128_si256::<1>(d);
6434        let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
6435        let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
6436        let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
6437        _mm_cvtsi128_si32(s32)
6438    }
6439}
6440
6441/// Unpack one q1t group's base-3 codes into 32 i8 signs via 7 unaligned u64
6442/// stores (see `SIGN5_U64`). `dst` MUST have ≥ 40 bytes: the 7th store writes
6443/// `dst[30..38]`. Stores go in order so each one's trailing zeros are
6444/// overwritten by the next; the final 6 padding bytes are unused by the dot.
6445#[inline]
6446fn q1t_unpack_group_i8(codes: *const u8, dst: &mut [i8]) {
6447    debug_assert!(dst.len() >= 40);
6448    // SAFETY: codes points at 7 readable bytes; dst has ≥ 40 bytes so every
6449    // 8-byte store at offset bi*5 (bi ≤ 6 → ≤ 30) stays in bounds.
6450    unsafe {
6451        let p = dst.as_mut_ptr();
6452        for bi in 0..7 {
6453            core::ptr::write_unaligned(
6454                p.add(bi * 5) as *mut u64,
6455                SIGN5_U64[*codes.add(bi) as usize],
6456            );
6457        }
6458    }
6459}
6460
6461/// One 32-group int8 dot, arch-dispatched (the matmat inner loop, where the
6462/// row's signs are unpacked once and dotted against every batch input).
6463/// Callers are gated by `a8w8_enabled()`, so the target-feature arms are
6464/// reachable; the scalar arm is a non-SIMD-arch fallback.
6465#[inline]
6466fn q1t_i8dot32(w: *const i8, x: *const i8) -> i32 {
6467    #[cfg(target_arch = "aarch64")]
6468    unsafe {
6469        return sdot32_i8(w, x);
6470    }
6471    #[cfg(target_arch = "x86_64")]
6472    unsafe {
6473        return i8dot32_avx2(w, x);
6474    }
6475    #[allow(unreachable_code)]
6476    unsafe {
6477        let mut s = 0i32;
6478        for k in 0..GROUP_SIZE {
6479            s += *w.add(k) as i32 * *x.add(k) as i32;
6480        }
6481        s
6482    }
6483}
6484
6485#[inline]
6486unsafe fn q1t_unpack_reg_u64s(codes: *const u8) -> (u64, u64, u64, u64) {
6487    let (s0, s1, s2, s3, s4, s5, s6) = unsafe {
6488        (
6489            SIGN5_U64[*codes as usize],
6490            SIGN5_U64[*codes.add(1) as usize],
6491            SIGN5_U64[*codes.add(2) as usize],
6492            SIGN5_U64[*codes.add(3) as usize],
6493            SIGN5_U64[*codes.add(4) as usize],
6494            SIGN5_U64[*codes.add(5) as usize],
6495            SIGN5_U64[*codes.add(6) as usize],
6496        )
6497    };
6498
6499    let u0 = s0 | (s1 << 40);
6500    let u1 = (s1 >> 24) | (s2 << 16) | (s3 << 56);
6501    let u2 = (s3 >> 8) | (s4 << 32);
6502    let u3 = (s4 >> 32) | (s5 << 8) | (s6 << 48);
6503
6504    (u0, u1, u2, u3)
6505}
6506
6507/// One q1t row's int8 base dot: `Σ_group s·dot(signs, xq)` (before the shared
6508/// `sx`). Direct register unpacking (zero stack stores/loads, no STLF stalls).
6509/// ARM SDOT.
6510#[cfg(target_arch = "aarch64")]
6511#[target_feature(enable = "neon,dotprod")]
6512unsafe fn q1t_dot_row_sdot(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
6513    use core::arch::aarch64::*;
6514    use core::arch::asm;
6515    unsafe {
6516        const TILE: usize = cortiq_core::quant::Q1T_TILE;
6517        let mut acc = 0f32;
6518        let bytes_ptr = bytes.as_ptr();
6519        let xq_ptr = xq.as_ptr();
6520        let row_off = r * gpr * TILE;
6521
6522        let gpr2 = gpr & !1;
6523        let mut gi = 0;
6524        while gi < gpr2 {
6525            let off0 = row_off + gi * TILE;
6526            let off1 = off0 + TILE;
6527            let s0 = f16_to_f32(u16::from_le_bytes([
6528                *bytes_ptr.add(off0),
6529                *bytes_ptr.add(off0 + 1),
6530            ]));
6531            let s1 = f16_to_f32(u16::from_le_bytes([
6532                *bytes_ptr.add(off1),
6533                *bytes_ptr.add(off1 + 1),
6534            ]));
6535
6536            let (u0_0, u1_0, u2_0, u3_0) = q1t_unpack_reg_u64s(bytes_ptr.add(off0 + 2));
6537            let (u0_1, u1_1, u2_1, u3_1) = q1t_unpack_reg_u64s(bytes_ptr.add(off1 + 2));
6538
6539            let w0_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_0), vcreate_u64(u1_0)));
6540            let w1_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_0), vcreate_u64(u3_0)));
6541            let w0_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_1), vcreate_u64(u1_1)));
6542            let w1_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_1), vcreate_u64(u3_1)));
6543
6544            let x0_0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE));
6545            let x1_0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE + 16));
6546            let x0_1 = vld1q_s8(xq_ptr.add((gi + 1) * GROUP_SIZE));
6547            let x1_1 = vld1q_s8(xq_ptr.add((gi + 1) * GROUP_SIZE + 16));
6548
6549            let (mut a0_0, mut a1_0) = (vdupq_n_s32(0), vdupq_n_s32(0));
6550            let (mut a0_1, mut a1_1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6551            asm!(
6552                "sdot {a0_0:v}.4s, {w0_0:v}.16b, {x0_0:v}.16b",
6553                "sdot {a1_0:v}.4s, {w1_0:v}.16b, {x1_0:v}.16b",
6554                "sdot {a0_1:v}.4s, {w0_1:v}.16b, {x0_1:v}.16b",
6555                "sdot {a1_1:v}.4s, {w1_1:v}.16b, {x1_1:v}.16b",
6556                a0_0 = inout(vreg) a0_0, a1_0 = inout(vreg) a1_0,
6557                a0_1 = inout(vreg) a0_1, a1_1 = inout(vreg) a1_1,
6558                w0_0 = in(vreg) w0_0, x0_0 = in(vreg) x0_0, w1_0 = in(vreg) w1_0, x1_0 = in(vreg) x1_0,
6559                w0_1 = in(vreg) w0_1, x0_1 = in(vreg) x0_1, w1_1 = in(vreg) w1_1, x1_1 = in(vreg) x1_1,
6560                options(pure, nomem, nostack),
6561            );
6562            let d0 = vaddvq_s32(vaddq_s32(a0_0, a1_0));
6563            let d1 = vaddvq_s32(vaddq_s32(a0_1, a1_1));
6564            acc += d0 as f32 * s0 + d1 as f32 * s1;
6565            gi += 2;
6566        }
6567
6568        if gi < gpr {
6569            let off = row_off + gi * TILE;
6570            let s = f16_to_f32(u16::from_le_bytes([
6571                *bytes_ptr.add(off),
6572                *bytes_ptr.add(off + 1),
6573            ]));
6574            let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
6575            let w0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0), vcreate_u64(u1)));
6576            let w1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2), vcreate_u64(u3)));
6577            let x0 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE));
6578            let x1 = vld1q_s8(xq_ptr.add(gi * GROUP_SIZE + 16));
6579            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6580            asm!(
6581                "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
6582                "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
6583                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6584                w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
6585                options(pure, nomem, nostack),
6586            );
6587            let d = vaddvq_s32(vaddq_s32(a0, a1));
6588            acc += d as f32 * s;
6589        }
6590        acc
6591    }
6592}
6593
6594/// x86 AVX2 mirror of `q1t_dot_row_sdot` (maddubs int8 dot per group).
6595#[cfg(target_arch = "x86_64")]
6596#[target_feature(enable = "avx2")]
6597unsafe fn q1t_dot_row_avx2(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
6598    use core::arch::x86_64::*;
6599    unsafe {
6600        const TILE: usize = cortiq_core::quant::Q1T_TILE;
6601        let mut acc = 0f32;
6602        let bytes_ptr = bytes.as_ptr();
6603        let xq_ptr = xq.as_ptr();
6604        let row_off = r * gpr * TILE;
6605
6606        let ones = _mm256_set1_epi16(1);
6607        for gi in 0..gpr {
6608            let off = row_off + gi * TILE;
6609            let s = f16_to_f32(u16::from_le_bytes([
6610                *bytes_ptr.add(off),
6611                *bytes_ptr.add(off + 1),
6612            ]));
6613            let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
6614            let wv = _mm256_set_epi64x(u3 as i64, u2 as i64, u1 as i64, u0 as i64);
6615            let xv = _mm256_loadu_si256(xq_ptr.add(gi * GROUP_SIZE) as *const __m256i);
6616            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
6617            let d256 = _mm256_madd_epi16(p16, ones);
6618            let d128 = _mm_add_epi32(
6619                _mm256_castsi256_si128(d256),
6620                _mm256_extracti128_si256(d256, 1),
6621            );
6622            let d64 = _mm_add_epi32(d128, _mm_shuffle_epi32(d128, 0xee));
6623            let d32 = _mm_cvtsi128_si32(_mm_add_epi32(d64, _mm_shuffle_epi32(d64, 0x55)));
6624            acc += d32 as f32 * s;
6625        }
6626        acc
6627    }
6628}
6629
6630/// VNNI twin of `q1t_dot_row_avx2` (see `dpbusd_hsum`).
6631#[cfg(target_arch = "x86_64")]
6632#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
6633unsafe fn q1t_dot_row_vnni(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
6634    use core::arch::x86_64::*;
6635    // SAFETY: same tile/xq contracts as `q1t_dot_row_avx2`.
6636    unsafe {
6637        const TILE: usize = cortiq_core::quant::Q1T_TILE;
6638        let mut acc = 0f32;
6639        let bytes_ptr = bytes.as_ptr();
6640        let xq_ptr = xq.as_ptr();
6641        let row_off = r * gpr * TILE;
6642        for gi in 0..gpr {
6643            let off = row_off + gi * TILE;
6644            let s = f16_to_f32(u16::from_le_bytes([
6645                *bytes_ptr.add(off),
6646                *bytes_ptr.add(off + 1),
6647            ]));
6648            let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
6649            let wv = _mm256_set_epi64x(u3 as i64, u2 as i64, u1 as i64, u0 as i64);
6650            let xv = _mm256_loadu_si256(xq_ptr.add(gi * GROUP_SIZE) as *const __m256i);
6651            let d = dpbusd_hsum(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
6652            acc += d as f32 * s;
6653        }
6654        acc
6655    }
6656}
6657
6658/// Per-row int8 base dot, dispatched once per row (matvec decode hot path).
6659/// Callers are gated by `a8w8_enabled()`, so the target-feature kernels are
6660/// reachable.
6661#[inline]
6662fn q1t_dot_row_i8(bytes: &[u8], r: usize, gpr: usize, xq: &[i8]) -> f32 {
6663    #[cfg(target_arch = "aarch64")]
6664    unsafe {
6665        return q1t_dot_row_sdot(bytes, r, gpr, xq);
6666    }
6667    #[cfg(target_arch = "x86_64")]
6668    unsafe {
6669        if vnni_tiles_enabled() {
6670            return q1t_dot_row_vnni(bytes, r, gpr, xq);
6671        }
6672        return q1t_dot_row_avx2(bytes, r, gpr, xq);
6673    }
6674    #[allow(unreachable_code)]
6675    {
6676        const TILE: usize = cortiq_core::quant::Q1T_TILE;
6677        let mut acc = 0f32;
6678        let mut sg = [0i8; GROUP_SIZE + 8]; // +8 slack for the u64-store unpack
6679        for gi in 0..gpr {
6680            let off = (r * gpr + gi) * TILE;
6681            let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
6682            q1t_unpack_group_i8(bytes.as_ptr().wrapping_add(off + 2), &mut sg);
6683            let mut d = 0i32;
6684            for k in 0..GROUP_SIZE {
6685                d += sg[k] as i32 * xq[gi * GROUP_SIZE + k] as i32;
6686            }
6687            acc += d as f32 * s;
6688        }
6689        acc
6690    }
6691}
6692
6693/// Σ over a row's outliers of `value·x[col]` — the correction that adds the
6694/// overlay's exact weights on top of the base dot. INVARIANT: the encoder
6695/// writes ternary code 0 at every outlier position (`quantize_q1t`), so the
6696/// base contributes nothing there and this is a plain `value·x`, not
6697/// `(value − base)·x` — no scattered per-outlier scale read. Row `r`'s entries
6698/// are the contiguous slice `[row_ptr[r], row_ptr[r+1])`, so no binary search.
6699fn q1t_row_outlier_correction(
6700    bytes: &[u8],
6701    r: usize,
6702    rp_off: usize,
6703    entries_off: usize,
6704    has_ov: bool,
6705    x: &[f32],
6706) -> f32 {
6707    if !has_ov {
6708        return 0.0;
6709    }
6710    let (c0, c1) = (
6711        q1t_rowptr(bytes, rp_off, r),
6712        q1t_rowptr(bytes, rp_off, r + 1),
6713    );
6714    let mut corr = 0f32;
6715    for p in c0..c1 {
6716        let e = entries_off + p * 4;
6717        let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
6718        let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
6719        corr += val * x[col];
6720    }
6721    corr
6722}
6723
6724/// Dequantize one q1t row into `buf[..cols]` via the sign LUT (no division),
6725/// then apply the row's outliers (its `[row_ptr[r], row_ptr[r+1])` slice).
6726/// Used by the batched (prefill) path where the decode amortizes over the batch.
6727fn q1t_dequant_row(
6728    bytes: &[u8],
6729    r: usize,
6730    gpr: usize,
6731    rp_off: usize,
6732    entries_off: usize,
6733    has_ov: bool,
6734    buf: &mut [f32],
6735) {
6736    const TILE: usize = cortiq_core::quant::Q1T_TILE;
6737    for g in 0..gpr {
6738        let off = (r * gpr + g) * TILE;
6739        let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
6740        let codes = &bytes[off + 2..off + TILE];
6741        let bc = g * GROUP_SIZE;
6742        // 6 full bytes (30 codes) + a 7th byte holding the last 2.
6743        for bi in 0..6 {
6744            let lut = &SIGN5[codes[bi] as usize];
6745            let d = &mut buf[bc + bi * 5..bc + bi * 5 + 5];
6746            for i in 0..5 {
6747                d[i] = lut[i] * s;
6748            }
6749        }
6750        let lut = &SIGN5[codes[6] as usize];
6751        buf[bc + 30] = lut[0] * s;
6752        buf[bc + 31] = lut[1] * s;
6753    }
6754    if !has_ov {
6755        return;
6756    }
6757    let (c0, c1) = (
6758        q1t_rowptr(bytes, rp_off, r),
6759        q1t_rowptr(bytes, rp_off, r + 1),
6760    );
6761    for p in c0..c1 {
6762        let e = entries_off + p * 4;
6763        let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
6764        buf[col] = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
6765    }
6766}
6767
6768/// Add the sparse outlier overlay onto a base dot already in `out` (the GPU
6769/// computes the ternary base; the overlay stays on the CPU — its entries are
6770/// few and its per-row gather doesn't vectorize on the GPU). Row-parallel.
6771fn q1t_add_overlay(
6772    bytes: &[u8],
6773    x: &[f32],
6774    rows: usize,
6775    cols: usize,
6776    out: &mut [f32],
6777    pool: Option<&Pool>,
6778) {
6779    const TILE: usize = cortiq_core::quant::Q1T_TILE;
6780    let gpr = cols / GROUP_SIZE;
6781    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
6782    if !has_ov {
6783        return;
6784    }
6785    let out_addr = SendMut(out.as_mut_ptr());
6786    let run = move |start: usize, end: usize| {
6787        for r in start..end {
6788            let corr = q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
6789            // SAFETY: disjoint rows; add onto the base the GPU already wrote.
6790            unsafe { *out_addr.at(r) += corr };
6791        }
6792    };
6793    dispatch_rows(pool, rows, &run);
6794}
6795
6796/// Q1T row range via the A8W8 int8 path — shared activation split,
6797/// per-row: base SDOT dot + outlier correction + overlay.
6798#[allow(clippy::too_many_arguments)]
6799fn q1t_range_a8w8(
6800    bytes: &[u8],
6801    gpr: usize,
6802    rp_off: usize,
6803    ent_off: usize,
6804    has_ov: bool,
6805    act: &SplitAct,
6806    x: &[f32],
6807    out: SendMut,
6808    start: usize,
6809    end: usize,
6810) {
6811    for r in start..end {
6812        let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
6813        for &(j, xv) in &act.outliers {
6814            acc += q1t_base_weight(bytes, r, gpr, j) * xv;
6815        }
6816        acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
6817        // SAFETY: disjoint row ranges per worker.
6818        unsafe { *out.at(r) = acc };
6819    }
6820}
6821
6822/// Q1T row range via the f32 path (no SDOT) — for matvec_many batched
6823/// dispatch when a8w8 is unavailable.
6824#[allow(clippy::too_many_arguments)]
6825fn q1t_range_f32_batch(
6826    bytes: &[u8],
6827    gpr: usize,
6828    rp_off: usize,
6829    ent_off: usize,
6830    has_ov: bool,
6831    x: &[f32],
6832    out: SendMut,
6833    start: usize,
6834    end: usize,
6835) {
6836    const TILE: usize = cortiq_core::quant::Q1T_TILE;
6837    let mut sg = [0f32; GROUP_SIZE];
6838    for r in start..end {
6839        let mut acc = 0f32;
6840        for g in 0..gpr {
6841            let off = (r * gpr + g) * TILE;
6842            let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
6843            let codes = &bytes[off + 2..off + TILE];
6844            let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
6845            for bi in 0..6 {
6846                sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
6847            }
6848            let lut = &SIGN5[codes[6] as usize];
6849            sg[30] = lut[0];
6850            sg[31] = lut[1];
6851            let mut gsum = 0f32;
6852            for k in 0..GROUP_SIZE {
6853                gsum += sg[k] * xg[k];
6854            }
6855            acc += s * gsum;
6856        }
6857        acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
6858        // SAFETY: disjoint row ranges per worker.
6859        unsafe { *out.at(r) = acc };
6860    }
6861}
6862
6863/// Ternary (q1t) matvec — decode+dot straight from mmap, one group at a time:
6864/// no per-ROW buffer, no division (the sign LUT), and a tiny per-group sign
6865/// buffer so the 32-wide dot vectorizes. This is the decode hot path.
6866fn q1t_matvec(
6867    bytes: &[u8],
6868    x: &[f32],
6869    rows: usize,
6870    cols: usize,
6871    out: &mut [f32],
6872    pool: Option<&Pool>,
6873) {
6874    debug_assert_eq!(out.len(), rows);
6875    const TILE: usize = cortiq_core::quant::Q1T_TILE;
6876    let gpr = cols / GROUP_SIZE;
6877    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
6878    let out_addr = SendMut(out.as_mut_ptr());
6879    // int8 SDOT base dot (ARM dotprod): ~4× the f32 arithmetic. x → i8 once
6880    // (`split_act`), activation outliers added back exactly in f32, weight
6881    // overlay on top. ARM SDOT / x86 AVX2; CMF_SDOT=0 keeps the exact f32 path.
6882    if a8w8_enabled() {
6883        let act = split_act(x);
6884        let act = &act;
6885        let run = move |start: usize, end: usize| {
6886            for r in start..end {
6887                let mut acc = q1t_dot_row_i8(bytes, r, gpr, &act.xq) * act.sx;
6888                for &(j, xv) in &act.outliers {
6889                    acc += q1t_base_weight(bytes, r, gpr, j) * xv;
6890                }
6891                acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
6892                // SAFETY: disjoint row ranges per worker.
6893                unsafe { *out_addr.at(r) = acc };
6894            }
6895        };
6896        dispatch_rows(pool, rows, &run);
6897        return;
6898    }
6899    let run = move |start: usize, end: usize| {
6900        // Per-group signs, unpacked contiguously so the dot below is a clean
6901        // 32-wide reduction the autovectorizer turns into f32x4 FMAs — the
6902        // 5-values-per-byte base-3 layout won't SIMD in place.
6903        let mut sg = [0f32; GROUP_SIZE];
6904        for r in start..end {
6905            let mut acc = 0f32;
6906            for g in 0..gpr {
6907                let off = (r * gpr + g) * TILE;
6908                let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
6909                let codes = &bytes[off + 2..off + TILE];
6910                let xg = &x[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
6911                for bi in 0..6 {
6912                    sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
6913                }
6914                let lut = &SIGN5[codes[6] as usize];
6915                sg[30] = lut[0];
6916                sg[31] = lut[1];
6917                let mut gsum = 0f32;
6918                for k in 0..GROUP_SIZE {
6919                    gsum += sg[k] * xg[k];
6920                }
6921                acc += s * gsum;
6922            }
6923            acc += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x);
6924            unsafe { *out_addr.at(r) = acc };
6925        }
6926    };
6927    dispatch_rows(pool, rows, &run);
6928}
6929
6930/// Fused-pair twin of `q1t_dot_row_sdot`: ONE register unpack of the
6931/// ternary codes serves BOTH activation streams (the unpack chain is
6932/// the dominant per-row cost — MTP verify pairs paid it twice). Per
6933/// stream the group order and f32 accumulation match the single-row
6934/// kernel exactly, so pair == 2×matvec bit-for-bit.
6935#[cfg(target_arch = "aarch64")]
6936#[target_feature(enable = "neon,dotprod")]
6937unsafe fn q1t_dot_row_sdot2(bytes: &[u8], r: usize, gpr: usize, xa: &[i8], xb: &[i8]) -> [f32; 2] {
6938    use core::arch::aarch64::*;
6939    use core::arch::asm;
6940    // SAFETY: same slice-length contracts as `q1t_dot_row_sdot`, ×2.
6941    unsafe {
6942        const TILE: usize = cortiq_core::quant::Q1T_TILE;
6943        let bytes_ptr = bytes.as_ptr();
6944        let row_off = r * gpr * TILE;
6945        let xp = [xa.as_ptr(), xb.as_ptr()];
6946        let mut acc = [0f32; 2];
6947        macro_rules! sdot2 {
6948            ($w0:expr, $w1:expr, $x:expr) => {{
6949                let x0 = vld1q_s8($x);
6950                let x1 = vld1q_s8($x.add(16));
6951                let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
6952                asm!(
6953                    "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
6954                    "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
6955                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
6956                    w0 = in(vreg) $w0, x0 = in(vreg) x0, w1 = in(vreg) $w1, x1 = in(vreg) x1,
6957                    options(pure, nomem, nostack),
6958                );
6959                vaddvq_s32(vaddq_s32(a0, a1))
6960            }};
6961        }
6962        let gpr2 = gpr & !1;
6963        let mut gi = 0;
6964        while gi < gpr2 {
6965            let off0 = row_off + gi * TILE;
6966            let off1 = off0 + TILE;
6967            let s0 = f16_to_f32(u16::from_le_bytes([
6968                *bytes_ptr.add(off0),
6969                *bytes_ptr.add(off0 + 1),
6970            ]));
6971            let s1 = f16_to_f32(u16::from_le_bytes([
6972                *bytes_ptr.add(off1),
6973                *bytes_ptr.add(off1 + 1),
6974            ]));
6975            let (u0_0, u1_0, u2_0, u3_0) = q1t_unpack_reg_u64s(bytes_ptr.add(off0 + 2));
6976            let (u0_1, u1_1, u2_1, u3_1) = q1t_unpack_reg_u64s(bytes_ptr.add(off1 + 2));
6977            let w0_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_0), vcreate_u64(u1_0)));
6978            let w1_0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_0), vcreate_u64(u3_0)));
6979            let w0_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0_1), vcreate_u64(u1_1)));
6980            let w1_1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2_1), vcreate_u64(u3_1)));
6981            for k in 0..2 {
6982                let d0 = sdot2!(w0_0, w1_0, xp[k].add(gi * GROUP_SIZE));
6983                let d1 = sdot2!(w0_1, w1_1, xp[k].add((gi + 1) * GROUP_SIZE));
6984                acc[k] += d0 as f32 * s0 + d1 as f32 * s1;
6985            }
6986            gi += 2;
6987        }
6988        if gi < gpr {
6989            let off = row_off + gi * TILE;
6990            let s = f16_to_f32(u16::from_le_bytes([
6991                *bytes_ptr.add(off),
6992                *bytes_ptr.add(off + 1),
6993            ]));
6994            let (u0, u1, u2, u3) = q1t_unpack_reg_u64s(bytes_ptr.add(off + 2));
6995            let w0 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u0), vcreate_u64(u1)));
6996            let w1 = vreinterpretq_s8_u64(vcombine_u64(vcreate_u64(u2), vcreate_u64(u3)));
6997            for k in 0..2 {
6998                let d = sdot2!(w0, w1, xp[k].add(gi * GROUP_SIZE));
6999                acc[k] += d as f32 * s;
7000            }
7001        }
7002        acc
7003    }
7004}
7005
7006/// Fused Q1T pair matvec: ONE pass over the rows serves both
7007/// activation streams — on ARM the ternary register unpack happens
7008/// once per tile pair (`q1t_dot_row_sdot2`); elsewhere the second dot
7009/// rides the row's L1-warm tile bytes. Per stream the math matches
7010/// `q1t_matvec` exactly.
7011fn q1t_matvec2(
7012    bytes: &[u8],
7013    x1: &[f32],
7014    x2: &[f32],
7015    rows: usize,
7016    cols: usize,
7017    o1: &mut [f32],
7018    o2: &mut [f32],
7019    pool: Option<&Pool>,
7020) {
7021    debug_assert_eq!(o1.len(), rows);
7022    debug_assert_eq!(o2.len(), rows);
7023    const TILE: usize = cortiq_core::quant::Q1T_TILE;
7024    let gpr = cols / GROUP_SIZE;
7025    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
7026    let out1 = SendMut(o1.as_mut_ptr());
7027    let out2 = SendMut(o2.as_mut_ptr());
7028    if a8w8_enabled() {
7029        let a1 = split_act(x1);
7030        let a2 = split_act(x2);
7031        let (a1, a2) = (&a1, &a2);
7032        let run = move |start: usize, end: usize| {
7033            for r in start..end {
7034                #[cfg(target_arch = "aarch64")]
7035                // a8w8 on aarch64 ⇔ sdot_enabled(), so the kernel's
7036                // target features are present.
7037                let ds = unsafe { q1t_dot_row_sdot2(bytes, r, gpr, &a1.xq, &a2.xq) };
7038                #[cfg(not(target_arch = "aarch64"))]
7039                let ds = [
7040                    q1t_dot_row_i8(bytes, r, gpr, &a1.xq),
7041                    q1t_dot_row_i8(bytes, r, gpr, &a2.xq),
7042                ];
7043                let mut acc1 = ds[0] * a1.sx;
7044                for &(j, xv) in &a1.outliers {
7045                    acc1 += q1t_base_weight(bytes, r, gpr, j) * xv;
7046                }
7047                acc1 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x1);
7048                let mut acc2 = ds[1] * a2.sx;
7049                for &(j, xv) in &a2.outliers {
7050                    acc2 += q1t_base_weight(bytes, r, gpr, j) * xv;
7051                }
7052                acc2 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x2);
7053                // SAFETY: disjoint row ranges per worker.
7054                unsafe {
7055                    *out1.at(r) = acc1;
7056                    *out2.at(r) = acc2;
7057                }
7058            }
7059        };
7060        dispatch_rows(pool, rows, &run);
7061        return;
7062    }
7063    let run = move |start: usize, end: usize| {
7064        // Exact path (CMF_SDOT=0): unpack the sign LUT once per group,
7065        // dot both streams — same op order per stream as `q1t_matvec`.
7066        let mut sg = [0f32; GROUP_SIZE];
7067        for r in start..end {
7068            let mut acc1 = 0f32;
7069            let mut acc2 = 0f32;
7070            for g in 0..gpr {
7071                let off = (r * gpr + g) * TILE;
7072                let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
7073                let codes = &bytes[off + 2..off + TILE];
7074                for bi in 0..6 {
7075                    sg[bi * 5..bi * 5 + 5].copy_from_slice(&SIGN5[codes[bi] as usize]);
7076                }
7077                let lut = &SIGN5[codes[6] as usize];
7078                sg[30] = lut[0];
7079                sg[31] = lut[1];
7080                let xg1 = &x1[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
7081                let xg2 = &x2[g * GROUP_SIZE..g * GROUP_SIZE + GROUP_SIZE];
7082                let mut gsum1 = 0f32;
7083                for k in 0..GROUP_SIZE {
7084                    gsum1 += sg[k] * xg1[k];
7085                }
7086                acc1 += s * gsum1;
7087                let mut gsum2 = 0f32;
7088                for k in 0..GROUP_SIZE {
7089                    gsum2 += sg[k] * xg2[k];
7090                }
7091                acc2 += s * gsum2;
7092            }
7093            acc1 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x1);
7094            acc2 += q1t_row_outlier_correction(bytes, r, rp_off, ent_off, has_ov, x2);
7095            // SAFETY: disjoint row ranges per worker.
7096            unsafe {
7097                *out1.at(r) = acc1;
7098                *out2.at(r) = acc2;
7099            }
7100        }
7101    };
7102    dispatch_rows(pool, rows, &run);
7103}
7104
7105/// Ternary (q1t) matmat (prefill) — dequant each row once, dot the whole
7106/// batch against it (amortizes the per-row decode).
7107fn q1t_matmat(
7108    bytes: &[u8],
7109    xs: &[f32],
7110    b: usize,
7111    rows: usize,
7112    cols: usize,
7113    out: &mut [f32],
7114    pool: Option<&Pool>,
7115) {
7116    debug_assert_eq!(out.len(), b * rows);
7117    const TILE: usize = cortiq_core::quant::Q1T_TILE;
7118    let gpr = cols / GROUP_SIZE;
7119    let (rp_off, ent_off, has_ov) = q1t_overlay(bytes, rows * gpr * TILE, rows);
7120    let out_addr = SendMut(out.as_mut_ptr());
7121    // int8 prefill (ARM SDOT / x86 AVX2): quantize the B inputs once, unpack
7122    // each weight row's signs to i8 ONCE, then int8-dot against every input —
7123    // the row sign-decode amortizes over the whole batch. CMF_SDOT=0 → f32.
7124    if a8w8_enabled() {
7125        let acts: Vec<SplitAct> = (0..b)
7126            .map(|bi| split_act(&xs[bi * cols..(bi + 1) * cols]))
7127            .collect();
7128        let acts = &acts;
7129        let run = move |start: usize, end: usize| {
7130            let mut sg = vec![0i8; cols + 8]; // row signs, i8 (+8 unpack slack)
7131            let mut sc = vec![0f32; gpr]; // per-group scales
7132            let mut accs = vec![0f32; b]; // per-batch accumulators, reused per row
7133            for r in start..end {
7134                for g in 0..gpr {
7135                    let off = (r * gpr + g) * TILE;
7136                    sc[g] = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
7137                    q1t_unpack_group_i8(
7138                        bytes.as_ptr().wrapping_add(off + 2),
7139                        &mut sg[g * GROUP_SIZE..],
7140                    );
7141                }
7142                for bi in 0..b {
7143                    let act = &acts[bi];
7144                    let mut isum = 0f32;
7145                    for g in 0..gpr {
7146                        let d = q1t_i8dot32(
7147                            sg.as_ptr().wrapping_add(g * GROUP_SIZE),
7148                            act.xq.as_ptr().wrapping_add(g * GROUP_SIZE),
7149                        );
7150                        isum += d as f32 * sc[g];
7151                    }
7152                    let mut acc = isum * act.sx;
7153                    for &(j, xv) in &act.outliers {
7154                        acc += q1t_base_weight(bytes, r, gpr, j) * xv;
7155                    }
7156                    accs[bi] = acc;
7157                }
7158                // Overlay ONCE per row for the whole batch: read each (col, val)
7159                // from mmap a single time (was b× — the re-read dominated prefill)
7160                // and fan it out over the batch via the cached inputs.
7161                if has_ov {
7162                    let (c0, c1) = (
7163                        q1t_rowptr(bytes, rp_off, r),
7164                        q1t_rowptr(bytes, rp_off, r + 1),
7165                    );
7166                    for p in c0..c1 {
7167                        let e = ent_off + p * 4;
7168                        let col = u16::from_le_bytes([bytes[e], bytes[e + 1]]) as usize;
7169                        let val = f16_to_f32(u16::from_le_bytes([bytes[e + 2], bytes[e + 3]]));
7170                        for bi in 0..b {
7171                            accs[bi] += val * xs[bi * cols + col];
7172                        }
7173                    }
7174                }
7175                for bi in 0..b {
7176                    unsafe { *out_addr.at(bi * rows + r) = accs[bi] };
7177                }
7178            }
7179        };
7180        dispatch_rows(pool, rows, &run);
7181        return;
7182    }
7183    let run = move |start: usize, end: usize| {
7184        let mut buf = vec![0f32; cols];
7185        for r in start..end {
7186            q1t_dequant_row(bytes, r, gpr, rp_off, ent_off, has_ov, &mut buf);
7187            for bi in 0..b {
7188                let xr = &xs[bi * cols..(bi + 1) * cols];
7189                let mut acc = 0f32;
7190                for j in 0..cols {
7191                    acc += buf[j] * xr[j];
7192                }
7193                unsafe { *out_addr.at(bi * rows + r) = acc };
7194            }
7195        }
7196    };
7197    dispatch_rows(pool, rows, &run);
7198}
7199
7200fn q1_matvec(
7201    bytes: &[u8],
7202    x: &[f32],
7203    rows: usize,
7204    cols: usize,
7205    out: &mut [f32],
7206    pool: Option<&Pool>,
7207) {
7208    debug_assert_eq!(out.len(), rows);
7209    let gpr = cols / GROUP_SIZE;
7210    let out_addr = SendMut(out.as_mut_ptr());
7211    if a8w8_enabled() {
7212        let act = split_act(x);
7213        let gsum = q1_group_sums(&act.xq, gpr);
7214        let (act, gsum) = (&act, &gsum);
7215        let run = move |start: usize, end: usize| {
7216            q1_range_a8w8(bytes, gpr, act, gsum, out_addr, start, end)
7217        };
7218        dispatch_rows(pool, rows, &run);
7219        return;
7220    }
7221    let run = move |start: usize, end: usize| q1_range_f32(bytes, gpr, x, out_addr, start, end);
7222    dispatch_rows(pool, rows, &run);
7223}
7224
7225/// Fused two-input q1 matvec (weights read once per pair).
7226#[allow(clippy::too_many_arguments)]
7227fn q1_matvec2(
7228    bytes: &[u8],
7229    x1: &[f32],
7230    x2: &[f32],
7231    rows: usize,
7232    cols: usize,
7233    o1: &mut [f32],
7234    o2: &mut [f32],
7235    pool: Option<&Pool>,
7236) {
7237    let gpr = cols / GROUP_SIZE;
7238    let p1 = SendMut(o1.as_mut_ptr());
7239    let p2 = SendMut(o2.as_mut_ptr());
7240    if a8w8_enabled() {
7241        let a1 = split_act(x1);
7242        let a2 = split_act(x2);
7243        let g1 = q1_group_sums(&a1.xq, gpr);
7244        let g2 = q1_group_sums(&a2.xq, gpr);
7245        let (a1, a2, g1, g2) = (&a1, &a2, &g1, &g2);
7246        let run = move |start: usize, end: usize| {
7247            for r in start..end {
7248                let mut v1 = dot_q1_row_i8(bytes, r, gpr, &a1.xq, g1) * a1.sx;
7249                let mut v2 = dot_q1_row_i8(bytes, r, gpr, &a2.xq, g2) * a2.sx;
7250                for &(j, xv) in &a1.outliers {
7251                    let (w, s) = q1_outlier(bytes, r, gpr, j);
7252                    v1 += w * s * xv;
7253                }
7254                for &(j, xv) in &a2.outliers {
7255                    let (w, s) = q1_outlier(bytes, r, gpr, j);
7256                    v2 += w * s * xv;
7257                }
7258                // SAFETY: disjoint row ranges per worker.
7259                unsafe {
7260                    *p1.at(r) = v1;
7261                    *p2.at(r) = v2;
7262                }
7263            }
7264        };
7265        dispatch_rows(pool, rows, &run);
7266        return;
7267    }
7268    let run = move |start: usize, end: usize| {
7269        for r in start..end {
7270            // SAFETY: disjoint row ranges per worker.
7271            unsafe {
7272                *p1.at(r) = q1_row_exact(bytes, r, gpr, x1);
7273                *p2.at(r) = q1_row_exact(bytes, r, gpr, x2);
7274            }
7275        }
7276    };
7277    dispatch_rows(pool, rows, &run);
7278}
7279
7280/// Batched q1 matmat: each row's tiles stream once per microbatch.
7281#[allow(clippy::too_many_arguments)]
7282fn q1_matmat(
7283    bytes: &[u8],
7284    xs_all: &[f32],
7285    b: usize,
7286    rows: usize,
7287    cols: usize,
7288    out: &mut [f32],
7289    pool: Option<&Pool>,
7290) {
7291    debug_assert_eq!(out.len(), b * rows);
7292    let gpr = cols / GROUP_SIZE;
7293    let out_addr = SendMut(out.as_mut_ptr());
7294    if a8w8_enabled() {
7295        let acts: Vec<(SplitAct, Vec<i32>)> = (0..b)
7296            .map(|bi| {
7297                let act = split_act(&xs_all[bi * cols..(bi + 1) * cols]);
7298                let gsum = q1_group_sums(&act.xq, gpr);
7299                (act, gsum)
7300            })
7301            .collect();
7302        let acts = &acts;
7303        #[cfg(target_arch = "x86_64")]
7304        let blocked_ok = avx2_enabled()
7305            && blocked_enabled();
7306        #[cfg(target_arch = "aarch64")]
7307        let blocked_ok = sdot_enabled()
7308            && blocked_enabled();
7309        let run = move |start: usize, end: usize| {
7310            for r in start..end {
7311                let mut bi = 0usize;
7312                // Blocked 1×4: the unpacked bit mask serves four
7313                // activation streams per group.
7314                #[cfg(target_arch = "aarch64")]
7315                if blocked_ok {
7316                    while bi + 4 <= acts.len() {
7317                        let xs = [
7318                            acts[bi].0.xq.as_slice(),
7319                            acts[bi + 1].0.xq.as_slice(),
7320                            acts[bi + 2].0.xq.as_slice(),
7321                            acts[bi + 3].0.xq.as_slice(),
7322                        ];
7323                        let gs = [
7324                            acts[bi].1.as_slice(),
7325                            acts[bi + 1].1.as_slice(),
7326                            acts[bi + 2].1.as_slice(),
7327                            acts[bi + 3].1.as_slice(),
7328                        ];
7329                        let d = unsafe { dot_q1_row_1x4_sdot(bytes, r, gpr, xs, gs) };
7330                        for k in 0..4 {
7331                            let (act, _) = &acts[bi + k];
7332                            let mut acc = d[k] * act.sx;
7333                            for &(j, xv) in &act.outliers {
7334                                let (w, sc) = q1_outlier(bytes, r, gpr, j);
7335                                acc += w * sc * xv;
7336                            }
7337                            // SAFETY: disjoint (bi, r) cells per worker.
7338                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
7339                        }
7340                        bi += 4;
7341                    }
7342                }
7343                #[cfg(target_arch = "x86_64")]
7344                if blocked_ok {
7345                    while bi + 4 <= acts.len() {
7346                        let xs = [
7347                            acts[bi].0.xq.as_slice(),
7348                            acts[bi + 1].0.xq.as_slice(),
7349                            acts[bi + 2].0.xq.as_slice(),
7350                            acts[bi + 3].0.xq.as_slice(),
7351                        ];
7352                        let gs = [
7353                            acts[bi].1.as_slice(),
7354                            acts[bi + 1].1.as_slice(),
7355                            acts[bi + 2].1.as_slice(),
7356                            acts[bi + 3].1.as_slice(),
7357                        ];
7358                        let d = unsafe {
7359                            if vnni_tiles_enabled() {
7360                                dot_q1_row_1x4_vnni(bytes, r, gpr, xs, gs)
7361                            } else {
7362                                dot_q1_row_1x4_avx2(bytes, r, gpr, xs, gs)
7363                            }
7364                        };
7365                        for k in 0..4 {
7366                            let (act, _) = &acts[bi + k];
7367                            let mut acc = d[k] * act.sx;
7368                            for &(j, xv) in &act.outliers {
7369                                let (w, sc) = q1_outlier(bytes, r, gpr, j);
7370                                acc += w * sc * xv;
7371                            }
7372                            // SAFETY: disjoint (bi, r) cells per worker.
7373                            unsafe { *out_addr.at((bi + k) * rows + r) = acc };
7374                        }
7375                        bi += 4;
7376                    }
7377                }
7378                while bi < acts.len() {
7379                    let (act, gsum) = &acts[bi];
7380                    let mut acc = dot_q1_row_i8(bytes, r, gpr, &act.xq, gsum) * act.sx;
7381                    for &(j, xv) in &act.outliers {
7382                        let (w, s) = q1_outlier(bytes, r, gpr, j);
7383                        acc += w * s * xv;
7384                    }
7385                    // SAFETY: disjoint (bi, r) cells per worker range.
7386                    unsafe { *out_addr.at(bi * rows + r) = acc };
7387                    bi += 1;
7388                }
7389            }
7390        };
7391        dispatch_rows(pool, rows, &run);
7392        return;
7393    }
7394    let run = move |start: usize, end: usize| {
7395        for r in start..end {
7396            for bi in 0..b {
7397                let x = &xs_all[bi * cols..(bi + 1) * cols];
7398                // SAFETY: disjoint (bi, r) cells per worker range.
7399                unsafe { *out_addr.at(bi * rows + r) = q1_row_exact(bytes, r, gpr, x) };
7400            }
7401        }
7402    };
7403    dispatch_rows(pool, rows, &run);
7404}
7405
7406/// Fused q4_block matvec straight from the mapped bytes. SDOT path when
7407/// dotprod is available (port of vmfcore `dot_q4_block_sdot`, measured
7408/// +23% on q4 decode): nibbles → centered i8, int8×int8 `sdot` per
7409/// 32-group, exact outlier correction — the same A8W8 contract as q8.
7410/// `CMF_SDOT=0` keeps the exact scalar path.
7411fn q4matvec(
7412    bytes: &[u8],
7413    x: &[f32],
7414    rows: usize,
7415    cols: usize,
7416    out: &mut [f32],
7417    pool: Option<&Pool>,
7418) {
7419    debug_assert_eq!(out.len(), rows);
7420    let (packed, scales) = q4_split(bytes, rows, cols);
7421    let gpr = cols / GROUP_SIZE;
7422    let out_addr = SendMut(out.as_mut_ptr());
7423
7424    if a8w8_enabled() {
7425        let act = split_act(x);
7426        let run = move |start: usize, end: usize| {
7427            q4_range_a8w8(packed, scales, gpr, cols, &act, out_addr, start, end)
7428        };
7429        dispatch_rows(pool, rows, &run);
7430        return;
7431    }
7432
7433    let run =
7434        move |start: usize, end: usize| q4_range_f32(packed, scales, gpr, x, out_addr, start, end);
7435    dispatch_rows(pool, rows, &run);
7436}
7437
7438/// One q4 row via the A8W8 int8 path — SDOT on ARM, AVX2 maddubs on
7439/// x86 (scalar fallback is unreachable: callers gate on a8w8_enabled).
7440#[inline]
7441#[allow(unreachable_code)]
7442/// One UNPACKED q4 row (centered i8 in `buf`) against four activation
7443/// streams: the 32-byte weight chunk and its abs() load once per group,
7444/// the per-group f16 scale decodes once — four maddubs+reduce chains
7445/// instead of four full (load, abs, dot) rounds.
7446#[cfg(target_arch = "x86_64")]
7447#[target_feature(enable = "avx2")]
7448unsafe fn dot_q4b_row_1x4_avx2(
7449    buf: &[u8],
7450    scales: &[u8],
7451    g0: usize,
7452    gpr: usize,
7453    xs: [&[i8]; 4],
7454) -> [f32; 4] {
7455    // SAFETY: callers uphold buffer contracts (buf.len() == gpr·32).
7456    unsafe {
7457        use core::arch::x86_64::*;
7458        let ones = _mm256_set1_epi16(1);
7459        let mut acc = [0f32; 4];
7460        for gi in 0..gpr {
7461            let s = f16_to_f32(u16::from_le_bytes([
7462                scales[(g0 + gi) * 2],
7463                scales[(g0 + gi) * 2 + 1],
7464            ]));
7465            let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7466            let aw = _mm256_abs_epi8(w);
7467            for (k, xq) in xs.iter().enumerate() {
7468                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7469                let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
7470                let d = _mm256_madd_epi16(p16, ones);
7471                let hi128 = _mm256_extracti128_si256::<1>(d);
7472                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
7473                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7474                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7475                acc[k] += _mm_cvtsi128_si32(s32) as f32 * s;
7476            }
7477        }
7478        acc
7479    }
7480}
7481
7482/// VNNI twin of `dot_q4b_row_1x4_avx2` (see `dpbusd_hsum`).
7483#[cfg(target_arch = "x86_64")]
7484#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
7485unsafe fn dot_q4b_row_1x4_vnni(
7486    buf: &[u8],
7487    scales: &[u8],
7488    g0: usize,
7489    gpr: usize,
7490    xs: [&[i8]; 4],
7491) -> [f32; 4] {
7492    // SAFETY: callers uphold buffer contracts (buf.len() == gpr·32).
7493    unsafe {
7494        use core::arch::x86_64::*;
7495        let mut acc = [0f32; 4];
7496        for gi in 0..gpr {
7497            let s = f16_to_f32(u16::from_le_bytes([
7498                scales[(g0 + gi) * 2],
7499                scales[(g0 + gi) * 2 + 1],
7500            ]));
7501            let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7502            let aw = _mm256_abs_epi8(w);
7503            for (k, xq) in xs.iter().enumerate() {
7504                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7505                let d = dpbusd_hsum(aw, _mm256_sign_epi8(x, w));
7506                acc[k] += d as f32 * s;
7507            }
7508        }
7509        acc
7510    }
7511}
7512
7513/// The vbit flavor of the blocked 1×4: the per-activation A8W8 scale
7514/// folds in PER GROUP as `(d·sx)·s` — bit-matching the single-matvec
7515/// accumulation order (the q4_block flavor applies sx once at the end,
7516/// matching ITS single path; the two conventions are historical and
7517/// each blocked leg must mirror its own).
7518#[cfg(target_arch = "x86_64")]
7519#[target_feature(enable = "avx2")]
7520unsafe fn dot_q4b_row_1x4_sx_avx2(
7521    buf: &[u8],
7522    scales: &[u8],
7523    g0: usize,
7524    gpr: usize,
7525    xs: [&[i8]; 4],
7526    sxs: [f32; 4],
7527) -> [f32; 4] {
7528    // SAFETY: callers uphold buffer contracts (buf.len() == gpr·32).
7529    unsafe {
7530        use core::arch::x86_64::*;
7531        let ones = _mm256_set1_epi16(1);
7532        let mut acc = [0f32; 4];
7533        for gi in 0..gpr {
7534            let s = f16_to_f32(u16::from_le_bytes([
7535                scales[(g0 + gi) * 2],
7536                scales[(g0 + gi) * 2 + 1],
7537            ]));
7538            let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7539            let aw = _mm256_abs_epi8(w);
7540            for (k, xq) in xs.iter().enumerate() {
7541                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7542                let p16 = _mm256_maddubs_epi16(aw, _mm256_sign_epi8(x, w));
7543                let d = _mm256_madd_epi16(p16, ones);
7544                let hi128 = _mm256_extracti128_si256::<1>(d);
7545                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
7546                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
7547                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
7548                acc[k] += (_mm_cvtsi128_si32(s32) as f32 * sxs[k]) * s;
7549            }
7550        }
7551        acc
7552    }
7553}
7554
7555/// VNNI twin of `dot_q4b_row_1x4_sx_avx2` (see `dpbusd_hsum`; the
7556/// per-group `(d·sx)·s` fold mirrors the vbit single path).
7557#[cfg(target_arch = "x86_64")]
7558#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
7559unsafe fn dot_q4b_row_1x4_sx_vnni(
7560    buf: &[u8],
7561    scales: &[u8],
7562    g0: usize,
7563    gpr: usize,
7564    xs: [&[i8]; 4],
7565    sxs: [f32; 4],
7566) -> [f32; 4] {
7567    // SAFETY: callers uphold buffer contracts (buf.len() == gpr·32).
7568    unsafe {
7569        use core::arch::x86_64::*;
7570        let mut acc = [0f32; 4];
7571        for gi in 0..gpr {
7572            let s = f16_to_f32(u16::from_le_bytes([
7573                scales[(g0 + gi) * 2],
7574                scales[(g0 + gi) * 2 + 1],
7575            ]));
7576            let w = _mm256_loadu_si256(buf.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7577            let aw = _mm256_abs_epi8(w);
7578            for (k, xq) in xs.iter().enumerate() {
7579                let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
7580                let d = dpbusd_hsum(aw, _mm256_sign_epi8(x, w));
7581                acc[k] += (d as f32 * sxs[k]) * s;
7582            }
7583        }
7584        acc
7585    }
7586}
7587
7588#[allow(unreachable_code)]
7589fn dot_q4_row_i8(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
7590    #[cfg(target_arch = "aarch64")]
7591    unsafe {
7592        return dot_q4_row_sdot(packed, scales, g0, gpr, xq);
7593    }
7594    #[cfg(target_arch = "x86_64")]
7595    unsafe {
7596        return dot_q4_row_avx2(packed, scales, g0, gpr, xq);
7597    }
7598    let mut acc = 0f32;
7599    for gi in 0..gpr {
7600        let g = g0 + gi;
7601        let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
7602        let mut d = 0i32;
7603        for (k, &b) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
7604            d += ((b & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2] as i32
7605                + (((b >> 4) & 0x0F) as i32 - 8) * xq[gi * GROUP_SIZE + k * 2 + 1] as i32;
7606        }
7607        acc += d as f32 * s;
7608    }
7609    acc
7610}
7611
7612/// Two-activation q4 row via the A8W8 int8 path (see `dot_q4_row_i8`).
7613#[inline]
7614#[allow(unreachable_code)]
7615fn dot_q4_row_i8_2(
7616    packed: &[u8],
7617    scales: &[u8],
7618    g0: usize,
7619    gpr: usize,
7620    xq1: &[i8],
7621    xq2: &[i8],
7622) -> (f32, f32) {
7623    #[cfg(target_arch = "aarch64")]
7624    unsafe {
7625        return dot_q4_row_sdot2(packed, scales, g0, gpr, xq1, xq2);
7626    }
7627    #[cfg(target_arch = "x86_64")]
7628    unsafe {
7629        return dot_q4_row_avx2_2(packed, scales, g0, gpr, xq1, xq2);
7630    }
7631    (
7632        dot_q4_row_i8(packed, scales, g0, gpr, xq1),
7633        dot_q4_row_i8(packed, scales, g0, gpr, xq2),
7634    )
7635}
7636
7637/// One q4 row range via SDOT (kernel body of `q4matvec`, extracted so
7638/// multi-matrix jobs can drive it for several tensors in one dispatch).
7639#[allow(clippy::too_many_arguments)]
7640fn q4_range_a8w8(
7641    packed: &[u8],
7642    scales: &[u8],
7643    gpr: usize,
7644    cols: usize,
7645    act: &SplitAct,
7646    out: SendMut,
7647    start: usize,
7648    end: usize,
7649) {
7650    for r in start..end {
7651        let mut acc = dot_q4_row_i8(packed, scales, r * gpr, gpr, &act.xq) * act.sx;
7652        // xq is zeroed at outlier slots — add the exact terms.
7653        for &(j, xv) in &act.outliers {
7654            let flat = r * cols + j;
7655            let byte = packed[flat / 2];
7656            let nib = if flat & 1 == 0 {
7657                byte & 0x0F
7658            } else {
7659                byte >> 4
7660            };
7661            let s = f16_to_f32(u16::from_le_bytes([
7662                scales[(flat / GROUP_SIZE) * 2],
7663                scales[(flat / GROUP_SIZE) * 2 + 1],
7664            ]));
7665            acc += ((nib as i32 - 8) as f32) * s * xv;
7666        }
7667        // SAFETY: disjoint row ranges per worker.
7668        unsafe { *out.at(r) = acc };
7669    }
7670}
7671
7672/// Two-input q4 row range via the A8W8 int8 path — kernel body of
7673/// `q4matvec2`, extracted for pair multi-matrix jobs.
7674#[allow(clippy::too_many_arguments)]
7675fn q4_range2_a8w8(
7676    packed: &[u8],
7677    scales: &[u8],
7678    gpr: usize,
7679    cols: usize,
7680    a1: &SplitAct,
7681    a2: &SplitAct,
7682    p1: SendMut,
7683    p2: SendMut,
7684    start: usize,
7685    end: usize,
7686) {
7687    for r in start..end {
7688        let (s1, s2) = dot_q4_row_i8_2(packed, scales, r * gpr, gpr, &a1.xq, &a2.xq);
7689        let mut acc1 = s1 * a1.sx;
7690        let mut acc2 = s2 * a2.sx;
7691        // xq is zeroed at outlier slots — add the exact terms.
7692        let fix = |outliers: &[(usize, f32)], acc: &mut f32| {
7693            for &(j, xv) in outliers {
7694                let flat = r * cols + j;
7695                let byte = packed[flat / 2];
7696                let nib = if flat & 1 == 0 {
7697                    byte & 0x0F
7698                } else {
7699                    byte >> 4
7700                };
7701                let s = f16_to_f32(u16::from_le_bytes([
7702                    scales[(flat / GROUP_SIZE) * 2],
7703                    scales[(flat / GROUP_SIZE) * 2 + 1],
7704                ]));
7705                *acc += ((nib as i32 - 8) as f32) * s * xv;
7706            }
7707        };
7708        fix(&a1.outliers, &mut acc1);
7709        fix(&a2.outliers, &mut acc2);
7710        // SAFETY: disjoint row ranges per worker.
7711        unsafe {
7712            *p1.at(r) = acc1;
7713            *p2.at(r) = acc2;
7714        }
7715    }
7716}
7717
7718/// Exact scalar q4 row range (same extraction, non-SDOT path).
7719fn q4_range_f32(
7720    packed: &[u8],
7721    scales: &[u8],
7722    gpr: usize,
7723    x: &[f32],
7724    out: SendMut,
7725    start: usize,
7726    end: usize,
7727) {
7728    for r in start..end {
7729        let mut acc = 0f32;
7730        for gi in 0..gpr {
7731            let g = r * gpr + gi;
7732            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
7733            let pk = &packed[g * 16..(g + 1) * 16];
7734            let xg = &x[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
7735            let mut ga = 0f32;
7736            for (k, &b) in pk.iter().enumerate() {
7737                ga += ((b & 0x0F) as f32 - 8.0) * xg[k * 2]
7738                    + (((b >> 4) & 0x0F) as f32 - 8.0) * xg[k * 2 + 1];
7739            }
7740            acc += ga * s;
7741        }
7742        // SAFETY: disjoint row ranges per worker.
7743        unsafe { *out.at(r) = acc };
7744    }
7745}
7746
7747/// Fused two-input q4 matvec: nibbles are unpacked ONCE per group and
7748/// dotted against both activations (was: two full matvecs — double
7749/// weight traffic). Per-lane math matches `q4matvec` exactly.
7750#[allow(clippy::too_many_arguments)]
7751fn q4matvec2(
7752    bytes: &[u8],
7753    x1: &[f32],
7754    x2: &[f32],
7755    rows: usize,
7756    cols: usize,
7757    o1: &mut [f32],
7758    o2: &mut [f32],
7759    pool: Option<&Pool>,
7760) {
7761    debug_assert_eq!(o1.len(), rows);
7762    debug_assert_eq!(o2.len(), rows);
7763    let (packed, scales) = q4_split(bytes, rows, cols);
7764    let gpr = cols / GROUP_SIZE;
7765
7766    if a8w8_enabled() {
7767        let a1 = split_act(x1);
7768        let a2 = split_act(x2);
7769        let p1 = SendMut(o1.as_mut_ptr());
7770        let p2 = SendMut(o2.as_mut_ptr());
7771        let run = move |start: usize, end: usize| {
7772            q4_range2_a8w8(packed, scales, gpr, cols, &a1, &a2, p1, p2, start, end)
7773        };
7774        dispatch_rows(pool, rows, &run);
7775        return;
7776    }
7777
7778    let p1 = SendMut(o1.as_mut_ptr());
7779    let p2 = SendMut(o2.as_mut_ptr());
7780    let run = move |start: usize, end: usize| {
7781        q4_range2_f32(packed, scales, gpr, x1, x2, p1, p2, start, end)
7782    };
7783    dispatch_rows(pool, rows, &run);
7784}
7785
7786/// Two-input exact scalar q4 row range (same extraction).
7787#[allow(clippy::too_many_arguments)]
7788fn q4_range2_f32(
7789    packed: &[u8],
7790    scales: &[u8],
7791    gpr: usize,
7792    x1: &[f32],
7793    x2: &[f32],
7794    p1: SendMut,
7795    p2: SendMut,
7796    start: usize,
7797    end: usize,
7798) {
7799    for r in start..end {
7800        let (mut acc1, mut acc2) = (0f32, 0f32);
7801        for gi in 0..gpr {
7802            let g = r * gpr + gi;
7803            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
7804            let pk = &packed[g * 16..(g + 1) * 16];
7805            let x1g = &x1[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
7806            let x2g = &x2[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE];
7807            let (mut g1, mut g2) = (0f32, 0f32);
7808            for (k, &b) in pk.iter().enumerate() {
7809                let wl = (b & 0x0F) as f32 - 8.0;
7810                let wh = ((b >> 4) & 0x0F) as f32 - 8.0;
7811                g1 += wl * x1g[k * 2] + wh * x1g[k * 2 + 1];
7812                g2 += wl * x2g[k * 2] + wh * x2g[k * 2 + 1];
7813            }
7814            acc1 += g1 * s;
7815            acc2 += g2 * s;
7816        }
7817        // SAFETY: disjoint row ranges per worker.
7818        unsafe {
7819            *p1.at(r) = acc1;
7820            *p2.at(r) = acc2;
7821        }
7822    }
7823}
7824
7825thread_local! {
7826    /// Per-worker decoded-row scratch for the batched q4/vbit kernels
7827    /// (centered i8 for SDOT, f32 for the exact/scalar paths).
7828    static ROW_I8: std::cell::RefCell<Vec<u8>> = const { std::cell::RefCell::new(Vec::new()) };
7829    static ROW_F32: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
7830}
7831
7832/// Batched q4 matmat: each weight row is unpacked from the mmap ONCE
7833/// and dotted against ALL b activations (prefill used to fall back to b
7834/// full matvecs — b× weight traffic and b× nibble decode). Per-position
7835/// math matches `q4matvec` exactly: same group order, same accumulation.
7836/// `out` is row-major [b, rows] like `qmatmat`.
7837#[allow(clippy::too_many_arguments)]
7838fn q4matmat(
7839    bytes: &[u8],
7840    xs_all: &[f32],
7841    b: usize,
7842    rows: usize,
7843    cols: usize,
7844    out: &mut [f32],
7845    pool: Option<&Pool>,
7846) {
7847    debug_assert_eq!(xs_all.len(), b * cols);
7848    debug_assert_eq!(out.len(), b * rows);
7849    let (packed, scales) = q4_split(bytes, rows, cols);
7850    let gpr = cols / GROUP_SIZE;
7851    let gscale = |g: usize| f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
7852
7853    if a8w8_enabled() {
7854        let acts: Vec<SplitAct> = (0..b)
7855            .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
7856            .collect();
7857        let acts = &acts;
7858        let out_addr = SendMut(out.as_mut_ptr());
7859        let run = move |start: usize, end: usize| {
7860            ROW_I8.with(|rb| {
7861                let mut buf = rb.borrow_mut();
7862                buf.resize(cols, 0);
7863                for r in start..end {
7864                    // Unpack the row's nibbles to centered i8 once
7865                    // (element 2k = low nibble, 2k+1 = high — flat order,
7866                    // same as dot_q4_row_sdot's zip).
7867                    for gi in 0..gpr {
7868                        let g = r * gpr + gi;
7869                        for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
7870                            buf[gi * GROUP_SIZE + k * 2] = ((bt & 0x0F) as i32 - 8) as i8 as u8;
7871                            buf[gi * GROUP_SIZE + k * 2 + 1] =
7872                                (((bt >> 4) & 0x0F) as i32 - 8) as i8 as u8;
7873                        }
7874                    }
7875                    let mut bi = 0usize;
7876                    #[cfg(target_arch = "x86_64")]
7877                    if avx2_enabled()
7878                        && blocked_enabled()
7879                    {
7880                        while bi + 4 <= acts.len() {
7881                            let xs = [
7882                                acts[bi].xq.as_slice(),
7883                                acts[bi + 1].xq.as_slice(),
7884                                acts[bi + 2].xq.as_slice(),
7885                                acts[bi + 3].xq.as_slice(),
7886                            ];
7887                            let d = unsafe {
7888                                if vnni_tiles_enabled() {
7889                                    dot_q4b_row_1x4_vnni(&buf, scales, r * gpr, gpr, xs)
7890                                } else {
7891                                    dot_q4b_row_1x4_avx2(&buf, scales, r * gpr, gpr, xs)
7892                                }
7893                            };
7894                            for k in 0..4 {
7895                                let act = &acts[bi + k];
7896                                let mut acc = d[k] * act.sx;
7897                                for &(j, xv) in &act.outliers {
7898                                    acc += (buf[j] as i8) as f32
7899                                        * gscale((r * cols + j) / GROUP_SIZE)
7900                                        * xv;
7901                                }
7902                                // SAFETY: disjoint (bi, r) cells per worker.
7903                                unsafe { *out_addr.at((bi + k) * rows + r) = acc };
7904                            }
7905                            bi += 4;
7906                        }
7907                    }
7908                    while bi < acts.len() {
7909                        let act = &acts[bi];
7910                        let mut acc = 0f32;
7911                        for gi in 0..gpr {
7912                            let d = dot_i8_i8(
7913                                &buf[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
7914                                &act.xq[gi * GROUP_SIZE..(gi + 1) * GROUP_SIZE],
7915                            );
7916                            acc += d as f32 * gscale(r * gpr + gi);
7917                        }
7918                        acc *= act.sx;
7919                        // xq is zeroed at outlier slots — exact terms.
7920                        for &(j, xv) in &act.outliers {
7921                            acc += (buf[j] as i8) as f32 * gscale((r * cols + j) / GROUP_SIZE) * xv;
7922                        }
7923                        // SAFETY: disjoint (bi, r) cells per worker row range.
7924                        unsafe { *out_addr.at(bi * rows + r) = acc };
7925                        bi += 1;
7926                    }
7927                }
7928            })
7929        };
7930        dispatch_rows(pool, rows, &run);
7931        return;
7932    }
7933
7934    let out_addr = SendMut(out.as_mut_ptr());
7935    let run = move |start: usize, end: usize| {
7936        ROW_F32.with(|rb| {
7937            let mut buf = rb.borrow_mut();
7938            buf.resize(cols, 0.0);
7939            for r in start..end {
7940                // Decode raw (nib − 8) values once; scales stay per-group
7941                // so the accumulation order matches q4matvec bit-for-bit.
7942                for gi in 0..gpr {
7943                    let g = r * gpr + gi;
7944                    for (k, &bt) in packed[g * 16..(g + 1) * 16].iter().enumerate() {
7945                        buf[gi * GROUP_SIZE + k * 2] = (bt & 0x0F) as f32 - 8.0;
7946                        buf[gi * GROUP_SIZE + k * 2 + 1] = ((bt >> 4) & 0x0F) as f32 - 8.0;
7947                    }
7948                }
7949                for bi in 0..b {
7950                    let x = &xs_all[bi * cols..(bi + 1) * cols];
7951                    let mut acc = 0f32;
7952                    for gi in 0..gpr {
7953                        let mut ga = 0f32;
7954                        // Pairwise (lo + hi) addition, matching
7955                        // q4matvec's `ga += lo·x + hi·x` shape exactly —
7956                        // a flat one-per-element loop rounds differently
7957                        // and broke bit-parity on the scalar (x86) path.
7958                        for k in 0..GROUP_SIZE / 2 {
7959                            let e = gi * GROUP_SIZE + k * 2;
7960                            ga += buf[e] * x[e] + buf[e + 1] * x[e + 1];
7961                        }
7962                        acc += ga * gscale(r * gpr + gi);
7963                    }
7964                    // SAFETY: disjoint (bi, r) cells per worker row range.
7965                    unsafe { *out_addr.at(bi * rows + r) = acc };
7966                }
7967            }
7968        })
7969    };
7970    dispatch_rows(pool, rows, &run);
7971}
7972
7973/// Batched vbit matmat: each variable-bit row is decoded from the mmap
7974/// ONCE for the whole microbatch. Same per-position math as
7975/// `vbitmatvec` (SDOT A8W8 with exact outliers / exact f32 for b=8 rows
7976/// and the scalar path).
7977#[allow(clippy::too_many_arguments)]
7978fn vbitmatmat(
7979    bytes: &[u8],
7980    offsets: &[usize],
7981    xs_all: &[f32],
7982    b: usize,
7983    rows: usize,
7984    cols: usize,
7985    out: &mut [f32],
7986    pool: Option<&Pool>,
7987) {
7988    debug_assert_eq!(xs_all.len(), b * cols);
7989    debug_assert_eq!(out.len(), b * rows);
7990    debug_assert_eq!(offsets.len(), rows + 1);
7991    let ng = cols / GROUP_SIZE;
7992    let bits = &bytes[..rows];
7993    let sc_off = rows;
7994    let gscale = |r: usize, g: usize| {
7995        let so = (r * ng + g) * 2;
7996        f16_to_f32(u16::from_le_bytes([
7997            bytes[sc_off + so],
7998            bytes[sc_off + so + 1],
7999        ]))
8000    };
8001
8002    // Decode row r's raw (u − L) values into `dst` (f32, unscaled).
8003    let decode_f32 = |r: usize, dst: &mut [f32]| {
8004        let bw = bits[r] as usize;
8005        let l = ((1i32 << (bw - 1)) - 1) as f32;
8006        let data = &bytes[offsets[r]..offsets[r + 1]];
8007        let (mut acc, mut nbits, mut idx) = (0u64, 0usize, 0usize);
8008        for d in dst.iter_mut() {
8009            while nbits < bw {
8010                acc = (acc << 8) | data[idx] as u64;
8011                idx += 1;
8012                nbits += 8;
8013            }
8014            let u = ((acc >> (nbits - bw)) & ((1u64 << bw) - 1)) as f32;
8015            nbits -= bw;
8016            *d = u - l;
8017        }
8018    };
8019
8020    if a8w8_enabled() {
8021        let acts: Vec<SplitAct> = (0..b)
8022            .map(|bi| split_act(&xs_all[bi * cols..(bi + 1) * cols]))
8023            .collect();
8024        let acts = &acts;
8025        let out_addr = SendMut(out.as_mut_ptr());
8026        let run = move |start: usize, end: usize| {
8027            for r in start..end {
8028                let bw = bits[r] as usize;
8029                if bw == 8 {
8030                    // u−L reaches 128 → no i8 path; decode once, exact
8031                    // f32 dots for every position (same as vbitmatvec).
8032                    ROW_F32.with(|rb| {
8033                        let mut buf = rb.borrow_mut();
8034                        buf.resize(cols, 0.0);
8035                        decode_f32(r, &mut buf);
8036                        for bi in 0..b {
8037                            let x = &xs_all[bi * cols..(bi + 1) * cols];
8038                            let mut dot = 0f32;
8039                            for g in 0..ng {
8040                                let mut gd = 0f32;
8041                                for k in 0..GROUP_SIZE {
8042                                    gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
8043                                }
8044                                dot += gd * gscale(r, g);
8045                            }
8046                            // SAFETY: disjoint (bi, r) cells per worker range.
8047                            unsafe { *out_addr.at(bi * rows + r) = dot };
8048                        }
8049                    });
8050                    continue;
8051                }
8052                let l = (1i32 << (bw - 1)) - 1;
8053                let data = &bytes[offsets[r]..offsets[r + 1]];
8054                ROW_I8.with(|rb| {
8055                    let mut buf = rb.borrow_mut();
8056                    buf.resize(cols, 0);
8057                    #[inline(always)]
8058                    fn fill<const B: usize>(data: &[u8], l: i32, buf: &mut [u8]) {
8059                        for (blk, chunk) in buf.chunks_exact_mut(8).enumerate() {
8060                            let u = unpack8::<B>(&data[blk * B..]);
8061                            for k in 0..8 {
8062                                chunk[k] = (u[k] - l) as i8 as u8;
8063                            }
8064                        }
8065                    }
8066                    match bw {
8067                        3 => fill::<3>(data, l, &mut buf),
8068                        4 => vbit_fill4(data, &mut buf),
8069                        5 => fill::<5>(data, l, &mut buf),
8070                        6 => fill::<6>(data, l, &mut buf),
8071                        _ => unreachable!("vbit bit-width {bw} (validated at load)"),
8072                    }
8073                    let mut bi = 0usize;
8074                    // The vbit scale table shares q4_block's layout
8075                    // (contiguous f16 per (row·ng + g)), so the same
8076                    // blocked 1×4 kernel serves the decoded row.
8077                    #[cfg(target_arch = "x86_64")]
8078                    if avx2_enabled()
8079                        && blocked_enabled()
8080                    {
8081                        while bi + 4 <= acts.len() {
8082                            let xs = [
8083                                acts[bi].xq.as_slice(),
8084                                acts[bi + 1].xq.as_slice(),
8085                                acts[bi + 2].xq.as_slice(),
8086                                acts[bi + 3].xq.as_slice(),
8087                            ];
8088                            let sxs = [
8089                                acts[bi].sx,
8090                                acts[bi + 1].sx,
8091                                acts[bi + 2].sx,
8092                                acts[bi + 3].sx,
8093                            ];
8094                            let d = unsafe {
8095                                if vnni_tiles_enabled() {
8096                                    dot_q4b_row_1x4_sx_vnni(
8097                                        &buf,
8098                                        &bytes[sc_off..],
8099                                        r * ng,
8100                                        ng,
8101                                        xs,
8102                                        sxs,
8103                                    )
8104                                } else {
8105                                    dot_q4b_row_1x4_sx_avx2(
8106                                        &buf,
8107                                        &bytes[sc_off..],
8108                                        r * ng,
8109                                        ng,
8110                                        xs,
8111                                        sxs,
8112                                    )
8113                                }
8114                            };
8115                            for k in 0..4 {
8116                                let act = &acts[bi + k];
8117                                let mut dot = d[k];
8118                                for &(j, xv) in &act.outliers {
8119                                    dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
8120                                }
8121                                // SAFETY: disjoint (bi, r) cells per worker.
8122                                unsafe { *out_addr.at((bi + k) * rows + r) = dot };
8123                            }
8124                            bi += 4;
8125                        }
8126                    }
8127                    while bi < acts.len() {
8128                        let act = &acts[bi];
8129                        let mut dot = 0f32;
8130                        for g in 0..ng {
8131                            let d = dot_i8_i8(
8132                                &buf[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
8133                                &act.xq[g * GROUP_SIZE..(g + 1) * GROUP_SIZE],
8134                            ) as f32
8135                                * act.sx;
8136                            dot += d * gscale(r, g);
8137                        }
8138                        for &(j, xv) in &act.outliers {
8139                            dot += (buf[j] as i8) as f32 * gscale(r, j / GROUP_SIZE) * xv;
8140                        }
8141                        // SAFETY: disjoint (bi, r) cells per worker range.
8142                        unsafe { *out_addr.at(bi * rows + r) = dot };
8143                        bi += 1;
8144                    }
8145                });
8146            }
8147        };
8148        dispatch_rows(pool, rows, &run);
8149        return;
8150    }
8151
8152    let out_addr = SendMut(out.as_mut_ptr());
8153    let run = move |start: usize, end: usize| {
8154        ROW_F32.with(|rb| {
8155            let mut buf = rb.borrow_mut();
8156            buf.resize(cols, 0.0);
8157            for r in start..end {
8158                decode_f32(r, &mut buf);
8159                for bi in 0..b {
8160                    let x = &xs_all[bi * cols..(bi + 1) * cols];
8161                    let mut dot = 0f32;
8162                    for g in 0..ng {
8163                        let mut gd = 0f32;
8164                        for k in 0..GROUP_SIZE {
8165                            gd += buf[g * GROUP_SIZE + k] * x[g * GROUP_SIZE + k];
8166                        }
8167                        dot += gd * gscale(r, g);
8168                    }
8169                    // SAFETY: disjoint (bi, r) cells per worker range.
8170                    unsafe { *out_addr.at(bi * rows + r) = dot };
8171                }
8172            }
8173        })
8174    };
8175    dispatch_rows(pool, rows, &run);
8176}
8177
8178/// Build a GPU batch job for a q8-family mapped tensor (primary
8179/// shard): prescaled input + directory coordinates. None → not
8180/// GPU-eligible, caller stays on the CPU.
8181pub(crate) fn gpu_batch_job<'a>(
8182    t: &'a QTensor,
8183    x: &[f32],
8184) -> Option<(std::sync::Arc<CmfModel>, crate::gpu::BatchJob<'a>)> {
8185    match t {
8186        QTensor::Mapped {
8187            model,
8188            idx,
8189            dtype: dt @ (TensorDtype::Q8Row | TensorDtype::Q8_2f),
8190            rows,
8191            cols,
8192            row_scale,
8193            col_field,
8194            ..
8195        } => Some((
8196            model.clone(),
8197            crate::gpu::BatchJob {
8198                idx: *idx,
8199                rows: *rows,
8200                cols: *cols,
8201                row_scale,
8202                xs: prescale(x, col_field, *dt).into_owned(),
8203                layout: crate::gpu::BatchLayout::Q8,
8204            },
8205        )),
8206        // q1: raw f32 activations, tile-embedded scales.
8207        QTensor::Mapped {
8208            model,
8209            idx,
8210            dtype: TensorDtype::Q1,
8211            rows,
8212            cols,
8213            ..
8214        } => Some((
8215            model.clone(),
8216            crate::gpu::BatchJob {
8217                idx: *idx,
8218                rows: *rows,
8219                cols: *cols,
8220                row_scale: &[],
8221                xs: x.to_vec(),
8222                layout: crate::gpu::BatchLayout::Q1,
8223            },
8224        )),
8225        // q4_tiled / q4tp: raw f32 activations; the scales live in the
8226        // payload (inline tiles / row ladder), so row_scale stays empty.
8227        // The GDN projection batch already runs these layouts on Metal —
8228        // this arm lets the attention QKV batch reach the same kernels.
8229        QTensor::Mapped {
8230            model,
8231            idx,
8232            dtype: dt @ (TensorDtype::Q4Tiled | TensorDtype::Q4TiledP),
8233            rows,
8234            cols,
8235            ..
8236        } => Some((
8237            model.clone(),
8238            crate::gpu::BatchJob {
8239                idx: *idx,
8240                rows: *rows,
8241                cols: *cols,
8242                row_scale: &[],
8243                xs: x.to_vec(),
8244                layout: if *dt == TensorDtype::Q4Tiled {
8245                    crate::gpu::BatchLayout::Q4t
8246                } else {
8247                    crate::gpu::BatchLayout::Q4tp
8248                },
8249            },
8250        )),
8251        _ => None,
8252    }
8253}
8254
8255thread_local! {
8256    static PRESCALE_BUF1: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
8257    static PRESCALE_BUF2: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
8258}
8259
8260pub(crate) fn prescale<'a>(
8261    x: &'a [f32],
8262    col_field: &[f32],
8263    dtype: TensorDtype,
8264) -> std::borrow::Cow<'a, [f32]> {
8265    if dtype == TensorDtype::Q8_2f {
8266        x.iter().zip(col_field).map(|(a, c)| a * c).collect()
8267    } else {
8268        std::borrow::Cow::Borrowed(x)
8269    }
8270}
8271
8272/// θ col-field fold for q8_2f activations. Borrowed pass-through for
8273/// every other dtype, using thread-local buffers to eliminate per-matvec allocations.
8274pub(crate) fn prescale_with<R, F: FnOnce(&[f32]) -> R>(
8275    x: &[f32],
8276    col_field: &[f32],
8277    dtype: TensorDtype,
8278    buf_id: u8,
8279    f: F,
8280) -> R {
8281    if dtype == TensorDtype::Q8_2f {
8282        if buf_id == 1 {
8283            PRESCALE_BUF1.with(|b| {
8284                let mut buf = b.borrow_mut();
8285                buf.clear();
8286                buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
8287                f(&buf)
8288            })
8289        } else {
8290            PRESCALE_BUF2.with(|b| {
8291                let mut buf = b.borrow_mut();
8292                buf.clear();
8293                buf.extend(x.iter().zip(col_field).map(|(a, c)| a * c));
8294                f(&buf)
8295            })
8296        }
8297    } else {
8298        f(x)
8299    }
8300}
8301
8302// ───────────────────── x86-64 AVX2 kernels (roadmap этап 2) ─────────────────────
8303
8304/// AVX2+FMA available? Default ON when the CPU supports both;
8305/// `CMF_AVX2=0` disables (falls back to the autovectorized loops).
8306#[cfg(target_arch = "x86_64")]
8307pub(crate) fn avx2_enabled() -> bool {
8308    use std::sync::OnceLock;
8309    static ON: OnceLock<bool> = OnceLock::new();
8310    *ON.get_or_init(|| {
8311        std::env::var("CMF_AVX2").map(|v| v != "0").unwrap_or(true)
8312            && std::arch::is_x86_feature_detected!("avx2")
8313            && std::arch::is_x86_feature_detected!("fma")
8314    })
8315}
8316
8317/// AVX2 A8W8 allowed? The quantized-activation contract is switched by
8318/// the SAME env as the ARM SDOT path: `CMF_SDOT=0` keeps exact kernels
8319/// (the golden-parity exact gate relies on it) — AVX2 f32 kernels stay
8320/// active either way, they are exact (regrouped sums only).
8321#[cfg(target_arch = "x86_64")]
8322fn avx2_a8w8_enabled() -> bool {
8323    use std::sync::OnceLock;
8324    static ON: OnceLock<bool> = OnceLock::new();
8325    *ON.get_or_init(|| {
8326        avx2_enabled() && std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true)
8327    })
8328}
8329
8330/// A8W8 quantized-activation path available on THIS machine? One
8331/// switch across architectures: ARM dotprod (CMF_SDOT) or x86 AVX2
8332/// (CMF_AVX2 + the same CMF_SDOT exact-contract override).
8333#[inline]
8334pub(crate) fn a8w8_enabled() -> bool {
8335    #[cfg(target_arch = "aarch64")]
8336    {
8337        sdot_enabled()
8338    }
8339    #[cfg(target_arch = "x86_64")]
8340    {
8341        avx2_a8w8_enabled()
8342    }
8343    #[cfg(not(any(target_arch = "aarch64", target_arch = "x86_64")))]
8344    {
8345        false
8346    }
8347}
8348
8349/// int8·int8 dot dispatch: SDOT on ARM; AVX-512 VNNI (vpdpbusd) or AVX2
8350/// maddubs on x86. Callers are gated by `a8w8_enabled()`.
8351#[inline]
8352#[allow(unreachable_code)]
8353fn dot_i8_i8(w: &[u8], xq: &[i8]) -> i32 {
8354    #[cfg(target_arch = "aarch64")]
8355    unsafe {
8356        return dot_i8_sdot(w, xq);
8357    }
8358    #[cfg(target_arch = "x86_64")]
8359    unsafe {
8360        if avx512vnni_enabled() {
8361            return dot_i8_i8_vnni(w, xq);
8362        }
8363        return dot_i8_i8_avx2(w, xq);
8364    }
8365    w.iter()
8366        .zip(xq)
8367        .map(|(&a, &b)| (a as i8) as i32 * b as i32)
8368        .sum()
8369}
8370
8371/// AVX-512 VNNI available? (F+BW+VL+VNNI; `CMF_AVX512=0` falls back to
8372/// AVX2.) VL matters: short 32-byte groups (q4/vbit) ride the 256-bit
8373/// `vpdpbusd` encoding.
8374#[cfg(target_arch = "x86_64")]
8375fn avx512vnni_enabled() -> bool {
8376    use std::sync::OnceLock;
8377    static ON: OnceLock<bool> = OnceLock::new();
8378    *ON.get_or_init(|| {
8379        std::env::var("CMF_AVX512")
8380            .map(|v| v != "0")
8381            .unwrap_or(true)
8382            && std::arch::is_x86_feature_detected!("avx512f")
8383            && std::arch::is_x86_feature_detected!("avx512bw")
8384            && std::arch::is_x86_feature_detected!("avx512vl")
8385            && std::arch::is_x86_feature_detected!("avx512vnni")
8386    })
8387}
8388
8389/// Grouped-codec VNNI arms (the q4t/q4b/q1/q1t tile kernels): default
8390/// ON where AVX-512 VNNI exists (`CMF_VNNI_TILES=0` opt-out). Measured
8391/// on Ryzen 7950X (Zen4, 3 alternating process pairs, blocked GEMM
8392/// 4864×896 b=256): q4t 63→68 GF/s (+8%), q1 53→56 (+6%), q4b 72→75
8393/// (+4%) — consistent, no leg regressed. The tile kernels keep a
8394/// horizontal reduce per 32-weight group, so the `vpdpbusd` saving is
8395/// smaller than the long-dot q8 win (+13%), but it is real and free.
8396#[cfg(target_arch = "x86_64")]
8397fn vnni_tiles_enabled() -> bool {
8398    use std::sync::OnceLock;
8399    static ON: OnceLock<bool> = OnceLock::new();
8400    *ON.get_or_init(|| {
8401        std::env::var("CMF_VNNI_TILES")
8402            .map(|v| v != "0")
8403            .unwrap_or(true)
8404            && avx512vnni_enabled()
8405    })
8406}
8407
8408/// One 256-bit u8×i8 dot → i32 via `vpdpbusd` into a fresh accumulator
8409/// plus the same horizontal reduce the AVX2 kernels use. Products are
8410/// bounded (|w| ≤ 8 or ≤ 1), so maddubs never saturated — the i32 sum
8411/// is bit-identical to the maddubs+madd pair it replaces.
8412#[cfg(target_arch = "x86_64")]
8413#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
8414#[inline]
8415unsafe fn dpbusd_hsum(aw: core::arch::x86_64::__m256i, xs: core::arch::x86_64::__m256i) -> i32 {
8416    // SAFETY: pure register math.
8417    unsafe {
8418        use core::arch::x86_64::*;
8419        let d = _mm256_dpbusd_epi32(_mm256_setzero_si256(), aw, xs);
8420        let hi128 = _mm256_extracti128_si256::<1>(d);
8421        let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
8422        let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8423        let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8424        _mm_cvtsi128_si32(s32)
8425    }
8426}
8427
8428/// int8·int8 via AVX-512 VNNI: `vpdpbusd` fuses the maddubs+madd+add
8429/// triple into one u8×i8 dot-accumulate. AVX-512 has no vpsignb, so the
8430/// |w|·sign(x,w) trick becomes |w| × (x negated where w<0) via a mask
8431/// subtract — w==0 lanes contribute 0 through |w|=0 either way.
8432#[cfg(target_arch = "x86_64")]
8433#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
8434unsafe fn dot_i8_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
8435    // SAFETY: callers uphold slice-length contracts (see call sites).
8436    unsafe {
8437        use core::arch::x86_64::*;
8438        let n = w.len();
8439        let mut j = 0usize;
8440        let mut total: i32;
8441        // 4 independent accumulators: vpdpbusd is its own loop-carried
8442        // dependency (~5-cycle latency) — a single-acc loop runs
8443        // latency-bound and LOSES to the AVX2 maddubs kernel, measured
8444        // on Granite Rapids.
8445        {
8446            #[inline(always)]
8447            unsafe fn step(
8448                w: *const u8,
8449                x: *const i8,
8450                acc: core::arch::x86_64::__m512i,
8451            ) -> core::arch::x86_64::__m512i {
8452                unsafe {
8453                    use core::arch::x86_64::*;
8454                    let wv = _mm512_loadu_si512(w as *const _);
8455                    let xv = _mm512_loadu_si512(x as *const _);
8456                    let aw = _mm512_abs_epi8(wv);
8457                    let neg = _mm512_movepi8_mask(wv);
8458                    let sx = _mm512_mask_sub_epi8(xv, neg, _mm512_setzero_si512(), xv);
8459                    _mm512_dpbusd_epi32(acc, aw, sx)
8460                }
8461            }
8462            let (mut a0, mut a1, mut a2, mut a3) = (
8463                _mm512_setzero_si512(),
8464                _mm512_setzero_si512(),
8465                _mm512_setzero_si512(),
8466                _mm512_setzero_si512(),
8467            );
8468            while j + 256 <= n {
8469                a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
8470                a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), a1);
8471                a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), a2);
8472                a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), a3);
8473                j += 256;
8474            }
8475            while j + 64 <= n {
8476                a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), a0);
8477                j += 64;
8478            }
8479            let s01 = _mm512_add_epi32(a0, a1);
8480            let s23 = _mm512_add_epi32(a2, a3);
8481            total = _mm512_reduce_add_epi32(_mm512_add_epi32(s01, s23));
8482        }
8483        // 32-wide (q4/vbit groups are exactly 32 bytes).
8484        if j + 32 <= n {
8485            let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
8486            let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
8487            let d = _mm256_dpbusd_epi32(
8488                _mm256_setzero_si256(),
8489                _mm256_abs_epi8(wv),
8490                _mm256_sign_epi8(xv, wv),
8491            );
8492            let hi128 = _mm256_extracti128_si256::<1>(d);
8493            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
8494            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8495            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8496            total += _mm_cvtsi128_si32(s32);
8497            j += 32;
8498        }
8499        while j < n {
8500            total += (w[j] as i8) as i32 * xq[j] as i32;
8501            j += 1;
8502        }
8503        total
8504    }
8505}
8506
8507/// i8 row · f32 x via AVX2/FMA (x86 mirror of `dot_i8_f32_neon`).
8508#[cfg(target_arch = "x86_64")]
8509#[target_feature(enable = "avx2,fma")]
8510unsafe fn dot_i8_f32_avx2(w: &[u8], x: &[f32]) -> f32 {
8511    // SAFETY: callers uphold slice-length contracts (see call sites).
8512    unsafe {
8513        use core::arch::x86_64::*;
8514        let n = x.len();
8515        let wp = w.as_ptr();
8516        let xp = x.as_ptr();
8517        let (mut a0, mut a1) = (_mm256_setzero_ps(), _mm256_setzero_ps());
8518        let mut j = 0usize;
8519        while j + 16 <= n {
8520            let wb = _mm_loadu_si128(wp.add(j) as *const __m128i);
8521            let lo = _mm256_cvtepi8_epi32(wb);
8522            let hi = _mm256_cvtepi8_epi32(_mm_srli_si128::<8>(wb));
8523            a0 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(lo), _mm256_loadu_ps(xp.add(j)), a0);
8524            a1 = _mm256_fmadd_ps(_mm256_cvtepi32_ps(hi), _mm256_loadu_ps(xp.add(j + 8)), a1);
8525            j += 16;
8526        }
8527        let acc = _mm256_add_ps(a0, a1);
8528        let hi128 = _mm256_extractf128_ps::<1>(acc);
8529        let s128 = _mm_add_ps(_mm256_castps256_ps128(acc), hi128);
8530        let s64 = _mm_add_ps(s128, _mm_movehl_ps(s128, s128));
8531        let s32 = _mm_add_ss(s64, _mm_shuffle_ps::<1>(s64, s64));
8532        let mut sum = _mm_cvtss_f32(s32);
8533        while j < n {
8534            sum += (*wp.add(j) as i8) as f32 * *xp.add(j);
8535            j += 1;
8536        }
8537        sum
8538    }
8539}
8540
8541/// int8(weight)·int8(activation) → i32 via AVX2 maddubs — the x86
8542/// analogue of the SDOT A8W8 path. `maddubs` takes u8×i8, so the
8543/// standard sign trick applies: |w| × sign(x, w) ≡ w × x per lane.
8544/// Pair saturation is safe: |w|≤128, |x|≤127 → 2·128·127 < 32767.
8545#[cfg(target_arch = "x86_64")]
8546#[target_feature(enable = "avx2")]
8547unsafe fn dot_i8_i8_avx2(w: &[u8], xq: &[i8]) -> i32 {
8548    // SAFETY: callers uphold slice-length contracts (see call sites).
8549    unsafe {
8550        use core::arch::x86_64::*;
8551        let n = w.len();
8552        let ones = _mm256_set1_epi16(1);
8553        let mut acc = _mm256_setzero_si256();
8554        let mut j = 0usize;
8555        while j + 32 <= n {
8556            let wv = _mm256_loadu_si256(w.as_ptr().add(j) as *const __m256i);
8557            let xv = _mm256_loadu_si256(xq.as_ptr().add(j) as *const __m256i);
8558            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(wv), _mm256_sign_epi8(xv, wv));
8559            acc = _mm256_add_epi32(acc, _mm256_madd_epi16(p16, ones));
8560            j += 32;
8561        }
8562        let hi128 = _mm256_extracti128_si256::<1>(acc);
8563        let s128 = _mm_add_epi32(_mm256_castsi256_si128(acc), hi128);
8564        let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8565        let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8566        let mut s = _mm_cvtsi128_si32(s32);
8567        while j < n {
8568            s += (w[j] as i8) as i32 * xq[j] as i32;
8569            j += 1;
8570        }
8571        s
8572    }
8573}
8574
8575/// smmla 2×4: one instruction covers a 2-row × 2-activation × 8-deep
8576/// tile (32 MACs vs sdot's 16) — the weight pair loads once per 8-k
8577/// slice as a combined 2×8 register and meets two activation pairs.
8578#[cfg(target_arch = "aarch64")]
8579#[target_feature(enable = "neon,i8mm")]
8580unsafe fn dot_i8_smmla_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
8581    // SAFETY: callers uphold slice-length contracts.
8582    unsafe {
8583        use core::arch::aarch64::*;
8584        use core::arch::asm;
8585        let n = w0.len();
8586        let w0p = w0.as_ptr() as *const i8;
8587        let w1p = w1.as_ptr() as *const i8;
8588        // acc01 holds [c(r0,x0) c(r0,x1) c(r1,x0) c(r1,x1)]; acc23 the
8589        // same for x2/x3.
8590        let mut acc01 = vdupq_n_s32(0);
8591        let mut acc23 = vdupq_n_s32(0);
8592        let mut i = 0usize;
8593        while i + 8 <= n {
8594            let wa = vcombine_s8(vld1_s8(w0p.add(i)), vld1_s8(w1p.add(i)));
8595            let xb01 = vcombine_s8(
8596                vld1_s8(xs[0].as_ptr().add(i)),
8597                vld1_s8(xs[1].as_ptr().add(i)),
8598            );
8599            let xb23 = vcombine_s8(
8600                vld1_s8(xs[2].as_ptr().add(i)),
8601                vld1_s8(xs[3].as_ptr().add(i)),
8602            );
8603            asm!(
8604                "smmla {a01:v}.4s, {w:v}.16b, {x01:v}.16b",
8605                "smmla {a23:v}.4s, {w:v}.16b, {x23:v}.16b",
8606                a01 = inout(vreg) acc01, a23 = inout(vreg) acc23,
8607                w = in(vreg) wa, x01 = in(vreg) xb01, x23 = in(vreg) xb23,
8608                options(pure, nomem, nostack),
8609            );
8610            i += 8;
8611        }
8612        let mut out = [[0i32; 4]; 2];
8613        let a01: [i32; 4] = core::mem::transmute(acc01);
8614        let a23: [i32; 4] = core::mem::transmute(acc23);
8615        out[0][0] = a01[0];
8616        out[0][1] = a01[1];
8617        out[1][0] = a01[2];
8618        out[1][1] = a01[3];
8619        out[0][2] = a23[0];
8620        out[0][3] = a23[1];
8621        out[1][2] = a23[2];
8622        out[1][3] = a23[3];
8623        if i < n {
8624            for (k, x) in xs.iter().enumerate() {
8625                for j in i..n {
8626                    out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
8627                    out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
8628                }
8629            }
8630        }
8631        out
8632    }
8633}
8634
8635/// ARM twin of the x86 blocked prefill GEMM: two weight rows stay in
8636/// registers across four activation streams, eight sdot accumulators.
8637/// (The per-row form re-read each W row once per activation.)
8638#[cfg(target_arch = "aarch64")]
8639#[target_feature(enable = "neon,dotprod")]
8640unsafe fn dot_i8_sdot_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
8641    // SAFETY: callers uphold slice-length contracts.
8642    unsafe {
8643        use core::arch::aarch64::*;
8644        use core::arch::asm;
8645        let n = w0.len();
8646        let w0p = w0.as_ptr() as *const i8;
8647        let w1p = w1.as_ptr() as *const i8;
8648        let mut acc = [[vdupq_n_s32(0); 4]; 2];
8649        let mut i = 0usize;
8650        while i + 16 <= n {
8651            let wv0 = vld1q_s8(w0p.add(i));
8652            let wv1 = vld1q_s8(w1p.add(i));
8653            for (k, x) in xs.iter().enumerate() {
8654                let xv = vld1q_s8(x.as_ptr().add(i));
8655                let (mut a0, mut a1) = (acc[0][k], acc[1][k]);
8656                asm!(
8657                    "sdot {a0:v}.4s, {w0:v}.16b, {x:v}.16b",
8658                    "sdot {a1:v}.4s, {w1:v}.16b, {x:v}.16b",
8659                    a0 = inout(vreg) a0, a1 = inout(vreg) a1,
8660                    w0 = in(vreg) wv0, w1 = in(vreg) wv1, x = in(vreg) xv,
8661                    options(pure, nomem, nostack),
8662                );
8663                acc[0][k] = a0;
8664                acc[1][k] = a1;
8665            }
8666            i += 16;
8667        }
8668        let mut out = [[0i32; 4]; 2];
8669        for r in 0..2 {
8670            for k in 0..4 {
8671                out[r][k] = vaddvq_s32(acc[r][k]);
8672            }
8673        }
8674        if i < n {
8675            for (k, x) in xs.iter().enumerate() {
8676                for j in i..n {
8677                    out[0][k] += (w0[j] as i8) as i32 * x[j] as i32;
8678                    out[1][k] += (w1[j] as i8) as i32 * x[j] as i32;
8679                }
8680            }
8681        }
8682        out
8683    }
8684}
8685
8686/// Blocked 2 weight rows × 4 activations for the prefill GEMM
8687/// (roadmap P0: packed panels + multi-row accumulators). The two rows'
8688/// abs() live in registers across all four activation streams; the
8689/// sign-fixup is recomputed per pair (the price of the maddubs trick).
8690/// Returns raw i8·i8 dots; the caller applies scales and outliers.
8691#[cfg(target_arch = "x86_64")]
8692#[target_feature(enable = "avx2")]
8693unsafe fn dot_i8_i8_avx2_2x4(w0: &[u8], w1: &[u8], xs: [&[i8]; 4]) -> [[i32; 4]; 2] {
8694    // SAFETY: callers uphold slice-length contracts.
8695    unsafe {
8696        use core::arch::x86_64::*;
8697        let n = w0.len();
8698        let ones = _mm256_set1_epi16(1);
8699        let mut acc = [[_mm256_setzero_si256(); 4]; 2];
8700        let mut j = 0usize;
8701        while j + 32 <= n {
8702            let wv0 = _mm256_loadu_si256(w0.as_ptr().add(j) as *const __m256i);
8703            let wv1 = _mm256_loadu_si256(w1.as_ptr().add(j) as *const __m256i);
8704            let aw0 = _mm256_abs_epi8(wv0);
8705            let aw1 = _mm256_abs_epi8(wv1);
8706            for (k, x) in xs.iter().enumerate() {
8707                let xv = _mm256_loadu_si256(x.as_ptr().add(j) as *const __m256i);
8708                let p0 = _mm256_maddubs_epi16(aw0, _mm256_sign_epi8(xv, wv0));
8709                acc[0][k] = _mm256_add_epi32(acc[0][k], _mm256_madd_epi16(p0, ones));
8710                let p1 = _mm256_maddubs_epi16(aw1, _mm256_sign_epi8(xv, wv1));
8711                acc[1][k] = _mm256_add_epi32(acc[1][k], _mm256_madd_epi16(p1, ones));
8712            }
8713            j += 32;
8714        }
8715        let mut out = [[0i32; 4]; 2];
8716        for r in 0..2 {
8717            for k in 0..4 {
8718                let a = acc[r][k];
8719                let hi128 = _mm256_extracti128_si256::<1>(a);
8720                let s128 = _mm_add_epi32(_mm256_castsi256_si128(a), hi128);
8721                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8722                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8723                out[r][k] = _mm_cvtsi128_si32(s32);
8724            }
8725        }
8726        if j < n {
8727            for (k, x) in xs.iter().enumerate() {
8728                for i in j..n {
8729                    out[0][k] += (w0[i] as i8) as i32 * x[i] as i32;
8730                    out[1][k] += (w1[i] as i8) as i32 * x[i] as i32;
8731                }
8732            }
8733        }
8734        out
8735    }
8736}
8737
8738/// AVX2/VNNI q8 row dot with exact outlier correction (x86 mirror of
8739/// `row_dot_sdot` — same A8W8 contract). With AVX-512 VNNI the row goes
8740/// through the bias trick: Σ(w+128)·x via pure `vpdpbusd` (no per-lane
8741/// sign fixups), corrected by −128·Σx with Σx precomputed per split.
8742#[cfg(target_arch = "x86_64")]
8743#[inline]
8744fn row_dot_avx2(row: &[u8], act: &SplitAct) -> f32 {
8745    let dot = if avx512vnni_enabled() && row.len() >= 64 {
8746        (unsafe { dot_u8p128_i8_vnni(row, &act.xq) }) - 128 * act.xsum
8747    } else {
8748        unsafe { dot_i8_i8_avx2(row, &act.xq) }
8749    };
8750    let mut acc = dot as f32 * act.sx;
8751    for &(j, xv) in &act.outliers {
8752        acc += (row[j] as i8) as f32 * xv;
8753    }
8754    acc
8755}
8756
8757/// Σ (w[i]+128)·x[i] via pure `vpdpbusd` — the caller subtracts
8758/// 128·Σx. Four independent accumulators (dpbusd is ~5-cycle latency;
8759/// a single-acc loop runs latency-bound, measured on Granite Rapids).
8760#[cfg(target_arch = "x86_64")]
8761#[target_feature(enable = "avx2,avx512f,avx512bw,avx512vl,avx512vnni")]
8762unsafe fn dot_u8p128_i8_vnni(w: &[u8], xq: &[i8]) -> i32 {
8763    // SAFETY: callers uphold slice-length contracts (see call sites).
8764    unsafe {
8765        use core::arch::x86_64::*;
8766        let n = w.len();
8767        let flip = _mm512_set1_epi8(-128); // XOR 0x80: i8 w → u8 (w+128)
8768        #[inline(always)]
8769        unsafe fn step(
8770            w: *const u8,
8771            x: *const i8,
8772            flip: core::arch::x86_64::__m512i,
8773            acc: core::arch::x86_64::__m512i,
8774        ) -> core::arch::x86_64::__m512i {
8775            unsafe {
8776                use core::arch::x86_64::*;
8777                let wv = _mm512_xor_si512(_mm512_loadu_si512(w as *const _), flip);
8778                _mm512_dpbusd_epi32(acc, wv, _mm512_loadu_si512(x as *const _))
8779            }
8780        }
8781        let (mut a0, mut a1, mut a2, mut a3) = (
8782            _mm512_setzero_si512(),
8783            _mm512_setzero_si512(),
8784            _mm512_setzero_si512(),
8785            _mm512_setzero_si512(),
8786        );
8787        let mut j = 0usize;
8788        while j + 256 <= n {
8789            a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
8790            a1 = step(w.as_ptr().add(j + 64), xq.as_ptr().add(j + 64), flip, a1);
8791            a2 = step(w.as_ptr().add(j + 128), xq.as_ptr().add(j + 128), flip, a2);
8792            a3 = step(w.as_ptr().add(j + 192), xq.as_ptr().add(j + 192), flip, a3);
8793            j += 256;
8794        }
8795        while j + 64 <= n {
8796            a0 = step(w.as_ptr().add(j), xq.as_ptr().add(j), flip, a0);
8797            j += 64;
8798        }
8799        let mut total = _mm512_reduce_add_epi32(_mm512_add_epi32(
8800            _mm512_add_epi32(a0, a1),
8801            _mm512_add_epi32(a2, a3),
8802        ));
8803        // Scalar tail: (w as i8) + 128 ≡ (w as u8) ^ 0x80.
8804        while j < n {
8805            total += ((w[j] ^ 0x80) as i32) * xq[j] as i32;
8806            j += 1;
8807        }
8808        total
8809    }
8810}
8811
8812/// One q4 row via AVX2: nibbles → centered i8 (unpacklo/hi restores the
8813/// writer's flat order, same as the NEON vzip pair), maddubs against
8814/// the pre-quantized activation group, × the group's f16 scale. Pair
8815/// saturation safe: |w|≤8, |x|≤127 → 2·8·127 ≪ 32767. Mirror of
8816/// `dot_q4_row_sdot`.
8817#[cfg(target_arch = "x86_64")]
8818#[target_feature(enable = "avx2")]
8819unsafe fn dot_q4_row_avx2(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
8820    // SAFETY: callers uphold slice-length contracts (16 packed bytes and
8821    // 2 scale bytes per group; xq.len() == gpr·GROUP_SIZE).
8822    unsafe {
8823        use core::arch::x86_64::*;
8824        let lomask = _mm_set1_epi8(0x0F);
8825        let eight = _mm256_set1_epi8(8);
8826        let ones = _mm256_set1_epi16(1);
8827        let mut acc = 0f32;
8828        for gi in 0..gpr {
8829            let g = g0 + gi;
8830            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
8831            let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
8832            let lo = _mm_and_si128(b, lomask);
8833            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
8834            let w = _mm256_sub_epi8(
8835                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
8836                eight,
8837            );
8838            let x = _mm256_loadu_si256(xq.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8839            let p16 = _mm256_maddubs_epi16(_mm256_abs_epi8(w), _mm256_sign_epi8(x, w));
8840            let d = _mm256_madd_epi16(p16, ones);
8841            let hi128 = _mm256_extracti128_si256::<1>(d);
8842            let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
8843            let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8844            let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8845            acc += _mm_cvtsi128_si32(s32) as f32 * s;
8846        }
8847        acc
8848    }
8849}
8850
8851/// Two-activation q4 row via AVX2: nibbles unpacked ONCE per group,
8852/// both activations dotted against the same centered i8 register.
8853#[cfg(target_arch = "x86_64")]
8854#[target_feature(enable = "avx2")]
8855unsafe fn dot_q4_row_avx2_2(
8856    packed: &[u8],
8857    scales: &[u8],
8858    g0: usize,
8859    gpr: usize,
8860    xq1: &[i8],
8861    xq2: &[i8],
8862) -> (f32, f32) {
8863    // SAFETY: callers uphold slice-length contracts (see dot_q4_row_avx2).
8864    unsafe {
8865        use core::arch::x86_64::*;
8866        let lomask = _mm_set1_epi8(0x0F);
8867        let eight = _mm256_set1_epi8(8);
8868        let ones = _mm256_set1_epi16(1);
8869        let (mut acc1, mut acc2) = (0f32, 0f32);
8870        #[inline(always)]
8871        unsafe fn hsum(d: core::arch::x86_64::__m256i) -> i32 {
8872            unsafe {
8873                use core::arch::x86_64::*;
8874                let hi128 = _mm256_extracti128_si256::<1>(d);
8875                let s128 = _mm_add_epi32(_mm256_castsi256_si128(d), hi128);
8876                let s64 = _mm_add_epi32(s128, _mm_srli_si128::<8>(s128));
8877                let s32 = _mm_add_epi32(s64, _mm_srli_si128::<4>(s64));
8878                _mm_cvtsi128_si32(s32)
8879            }
8880        }
8881        for gi in 0..gpr {
8882            let g = g0 + gi;
8883            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
8884            let b = _mm_loadu_si128(packed.as_ptr().add(g * 16) as *const __m128i);
8885            let lo = _mm_and_si128(b, lomask);
8886            let hi = _mm_and_si128(_mm_srli_epi16::<4>(b), lomask);
8887            let w = _mm256_sub_epi8(
8888                _mm256_set_m128i(_mm_unpackhi_epi8(lo, hi), _mm_unpacklo_epi8(lo, hi)),
8889                eight,
8890            );
8891            let aw = _mm256_abs_epi8(w);
8892            let x1 = _mm256_loadu_si256(xq1.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8893            let x2 = _mm256_loadu_si256(xq2.as_ptr().add(gi * GROUP_SIZE) as *const __m256i);
8894            let d1 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x1, w)), ones);
8895            let d2 = _mm256_madd_epi16(_mm256_maddubs_epi16(aw, _mm256_sign_epi8(x2, w)), ones);
8896            acc1 += hsum(d1) as f32 * s;
8897            acc2 += hsum(d2) as f32 * s;
8898        }
8899        (acc1, acc2)
8900    }
8901}
8902
8903/// One q8 row range via AVX2 (x86 mirror of `q8_range_sdot`).
8904#[cfg(target_arch = "x86_64")]
8905fn q8_range_avx2(
8906    q: &[u8],
8907    row_scale: &[f32],
8908    act: &SplitAct,
8909    cols: usize,
8910    out_addr: SendMut,
8911    start: usize,
8912    end: usize,
8913) {
8914    for o in start..end {
8915        let v = row_dot_avx2(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
8916        // SAFETY: disjoint row ranges per worker.
8917        unsafe { *out_addr.at(o) = v };
8918    }
8919}
8920
8921/// Two-input q8 row range via AVX2 (x86 mirror of `q8_range2_sdot`).
8922#[cfg(target_arch = "x86_64")]
8923#[allow(clippy::too_many_arguments)]
8924fn q8_range2_avx2(
8925    q: &[u8],
8926    row_scale: &[f32],
8927    a1: &SplitAct,
8928    a2: &SplitAct,
8929    cols: usize,
8930    p1: SendMut,
8931    p2: SendMut,
8932    start: usize,
8933    end: usize,
8934) {
8935    for o in start..end {
8936        let row = &q[o * cols..(o + 1) * cols];
8937        // SAFETY: disjoint row ranges per worker.
8938        unsafe {
8939            *p1.at(o) = row_dot_avx2(row, a1) * row_scale[o];
8940            *p2.at(o) = row_dot_avx2(row, a2) * row_scale[o];
8941        }
8942    }
8943}
8944
8945// ───────────────────── A8W8 SDOT path (port of vmfcore, ×1.78 decode) ─────────────────────
8946
8947/// ARMv8.6 i8mm (smmla): 32 int8 MACs per instruction vs sdot's 16 —
8948/// yet MEASURED 2.4× SLOWER than the blocked sdot on Apple silicon
8949/// (108 vs 264 GF/s): the on-the-fly vcombine packing and the two-
8950/// accumulator dependency chain swamp the MAC advantage, and Apple's
8951/// four SIMD pipes already keep sdot fed. OPT-IN (CMF_I8MM=1) for
8952/// field trials on Cortex-A710/X-class parts with two pipes, where the
8953/// balance may differ; a pre-interleaved weight layout (repack infra)
8954/// is the known path if it ever earns its keep.
8955#[cfg(target_arch = "aarch64")]
8956fn i8mm_enabled() -> bool {
8957    static ON: std::sync::OnceLock<bool> = std::sync::OnceLock::new();
8958    *ON.get_or_init(|| {
8959        std::env::var("CMF_I8MM").map(|v| v == "1").unwrap_or(false)
8960            && std::arch::is_aarch64_feature_detected!("i8mm")
8961    })
8962}
8963
8964/// SDOT enabled? Default ON when the CPU has ARMv8.2 dotprod;
8965/// `CMF_SDOT=0` disables (falls back to i8×f32 NEON).
8966/// (On non-ARM release builds only the test tolerance switch calls it.)
8967#[cfg_attr(not(target_arch = "aarch64"), allow(dead_code))]
8968fn sdot_enabled() -> bool {
8969    use std::sync::OnceLock;
8970    static ON: OnceLock<bool> = OnceLock::new();
8971    *ON.get_or_init(|| {
8972        let want = std::env::var("CMF_SDOT").map(|v| v != "0").unwrap_or(true);
8973        if !want {
8974            return false;
8975        }
8976
8977        #[cfg(target_arch = "aarch64")]
8978        {
8979            if std::arch::is_aarch64_feature_detected!("dotprod") {
8980                return true;
8981            }
8982            #[cfg(target_os = "android")]
8983            {
8984                if let Ok(cpuinfo) = std::fs::read_to_string("/proc/cpuinfo") {
8985                    if cpuinfo.lines().any(|l| {
8986                        (l.starts_with("Features") || l.starts_with("features"))
8987                            && l.contains("asimddp")
8988                    }) {
8989                        return true;
8990                    }
8991                }
8992            }
8993            false
8994        }
8995        #[cfg(not(target_arch = "aarch64"))]
8996        {
8997            false
8998        }
8999    })
9000}
9001
9002/// Two-field activation split (≡ vmfcore `q8_split_prep`): outlier
9003/// channels (>8·rms) are computed exactly in f32; the bulk (outliers
9004/// zeroed → clean absmax) goes through int8 SDOT. Computed ONCE per
9005/// matvec, shared by all rows/workers.
9006struct SplitAct {
9007    xq: Vec<i8>,
9008    sx: f32,
9009    outliers: Vec<(usize, f32)>,
9010    /// Σ xq — the VNNI bias-trick correction (`(w+128)·x` sums need
9011    /// `−128·Σx`); one i32 per split, computed once per matvec.
9012    #[cfg_attr(not(target_arch = "x86_64"), allow(dead_code))]
9013    xsum: i32,
9014}
9015
9016thread_local! {
9017    /// Recycled xq buffers: split_act runs for every matvec (~200/token)
9018    /// and its hidden-size allocation was steady-state heap churn.
9019    static XQ_FREE: std::cell::RefCell<Vec<Vec<i8>>> =
9020        const { std::cell::RefCell::new(Vec::new()) };
9021}
9022
9023impl Drop for SplitAct {
9024    fn drop(&mut self) {
9025        let buf = std::mem::take(&mut self.xq);
9026        if buf.capacity() > 0 {
9027            XQ_FREE.with(|f| {
9028                let mut f = f.borrow_mut();
9029                if f.len() < 16 {
9030                    f.push(buf);
9031                }
9032            });
9033        }
9034    }
9035}
9036
9037thread_local! {
9038    /// One scratch row per WORKER, kept for the life of the thread.
9039    ///
9040    /// The kernels take a row of group scales per dispatch, and a fresh
9041    /// `vec![0f32; gpr]` inside the closure is one allocation per worker per
9042    /// dispatch — on the release checkpoint about six thousand a token, a
9043    /// quarter of everything the benchmark counts.
9044    static KROW: std::cell::RefCell<Vec<f32>> = const { std::cell::RefCell::new(Vec::new()) };
9045}
9046
9047/// Borrow `n` floats of the calling worker's scratch. Nothing inside a
9048/// kernel body borrows it again, which is what keeps the RefCell honest.
9049#[inline]
9050fn with_krow<R>(n: usize, f: impl FnOnce(&mut [f32]) -> R) -> R {
9051    KROW.with(|s| {
9052        let mut b = s.borrow_mut();
9053        if b.len() < n {
9054            b.resize(n, 0.0);
9055        }
9056        f(&mut b[..n])
9057    })
9058}
9059
9060fn split_act(x: &[f32]) -> SplitAct {
9061    let n = x.len();
9062    let rms = (x.iter().map(|&v| (v * v) as f64).sum::<f64>() / n.max(1) as f64).sqrt() as f32;
9063    let thr = 8.0 * rms;
9064    // One pass: collect outliers and the bulk absmax (outliers excluded —
9065    // identical to the old zero-then-fold over a copied buffer, minus the
9066    // full-vector copy).
9067    let mut outliers: Vec<(usize, f32)> = Vec::new();
9068    let mut amax = 0f32;
9069    for (j, &v) in x.iter().enumerate() {
9070        let a = v.abs();
9071        if a > thr {
9072            outliers.push((j, v));
9073        } else if a > amax {
9074            amax = a;
9075        }
9076    }
9077    let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
9078    let inv = 1.0 / sx;
9079    let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
9080    xq.clear();
9081    xq.reserve(n);
9082    if outliers.is_empty() {
9083        xq.extend(
9084            x.iter()
9085                .map(|&v| (v * inv).round().clamp(-127.0, 127.0) as i8),
9086        );
9087    } else {
9088        // Outlier slots quantize to 0 (their exact term is added later).
9089        xq.extend(x.iter().map(|&v| {
9090            if v.abs() > thr {
9091                0
9092            } else {
9093                (v * inv).round().clamp(-127.0, 127.0) as i8
9094            }
9095        }));
9096    }
9097    let xsum = xq.iter().map(|&v| v as i32).sum();
9098    SplitAct {
9099        xq,
9100        sx,
9101        outliers,
9102        xsum,
9103    }
9104}
9105
9106fn split_act_q8_2f(x: &[f32], col: &[f32]) -> SplitAct {
9107    let n = x.len();
9108    let rms = (x
9109        .iter()
9110        .zip(col)
9111        .map(|(&a, &c)| {
9112            let v = a * c;
9113            (v * v) as f64
9114        })
9115        .sum::<f64>()
9116        / n.max(1) as f64)
9117        .sqrt() as f32;
9118    let thr = 8.0 * rms;
9119
9120    let mut outliers = Vec::new();
9121    let mut amax = 0f32;
9122    for (j, (&a, &c)) in x.iter().zip(col).enumerate() {
9123        let v = a * c;
9124        let s = v.abs();
9125        if s > thr {
9126            outliers.push((j, v));
9127        } else if s > amax {
9128            amax = s;
9129        }
9130    }
9131
9132    let sx = if amax > 0.0 { amax / 127.0 } else { 1.0 };
9133    let inv = 1.0 / sx;
9134    let mut xq = XQ_FREE.with(|f| f.borrow_mut().pop()).unwrap_or_default();
9135    xq.clear();
9136    xq.reserve(n);
9137    if outliers.is_empty() {
9138        xq.extend(
9139            x.iter()
9140                .zip(col)
9141                .map(|(&a, &c)| ((a * c) * inv).round().clamp(-127.0, 127.0) as i8),
9142        );
9143    } else {
9144        xq.extend(x.iter().zip(col).map(|(&a, &c)| {
9145            let v = a * c;
9146            if v.abs() > thr {
9147                0
9148            } else {
9149                (v * inv).round().clamp(-127.0, 127.0) as i8
9150            }
9151        }));
9152    }
9153    let xsum = xq.iter().map(|&v| v as i32).sum();
9154    SplitAct {
9155        xq,
9156        sx,
9157        outliers,
9158        xsum,
9159    }
9160}
9161
9162/// int8(weight)·int8(activation) → i32 via `sdot` (inline asm — the
9163/// vdotq intrinsic is unstable; port of vmfcore `dot_i8_sdot`).
9164#[cfg(target_arch = "aarch64")]
9165#[target_feature(enable = "neon,dotprod")]
9166unsafe fn dot_i8_sdot(w: &[u8], xq: &[i8]) -> i32 {
9167    // SAFETY: callers uphold slice-length contracts (see call sites).
9168    unsafe {
9169        use core::arch::aarch64::*;
9170        use core::arch::asm;
9171        let wp = w.as_ptr() as *const i8;
9172        let n = w.len();
9173        let (mut a0, mut a1, mut a2, mut a3) = (
9174            vdupq_n_s32(0),
9175            vdupq_n_s32(0),
9176            vdupq_n_s32(0),
9177            vdupq_n_s32(0),
9178        );
9179        let mut i = 0;
9180        while i + 64 <= n {
9181            let (w0, x0) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
9182            let (w1, x1) = (vld1q_s8(wp.add(i + 16)), vld1q_s8(xq.as_ptr().add(i + 16)));
9183            let (w2, x2) = (vld1q_s8(wp.add(i + 32)), vld1q_s8(xq.as_ptr().add(i + 32)));
9184            let (w3, x3) = (vld1q_s8(wp.add(i + 48)), vld1q_s8(xq.as_ptr().add(i + 48)));
9185            asm!(
9186                "sdot {a0:v}.4s, {w0:v}.16b, {x0:v}.16b",
9187                "sdot {a1:v}.4s, {w1:v}.16b, {x1:v}.16b",
9188                "sdot {a2:v}.4s, {w2:v}.16b, {x2:v}.16b",
9189                "sdot {a3:v}.4s, {w3:v}.16b, {x3:v}.16b",
9190                a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
9191                w0 = in(vreg) w0, x0 = in(vreg) x0, w1 = in(vreg) w1, x1 = in(vreg) x1,
9192                w2 = in(vreg) w2, x2 = in(vreg) x2, w3 = in(vreg) w3, x3 = in(vreg) x3,
9193                options(pure, nomem, nostack),
9194            );
9195            i += 64;
9196        }
9197        while i + 16 <= n {
9198            let (wv, xv) = (vld1q_s8(wp.add(i)), vld1q_s8(xq.as_ptr().add(i)));
9199            asm!("sdot {a:v}.4s, {w:v}.16b, {x:v}.16b",
9200                 a = inout(vreg) a0, w = in(vreg) wv, x = in(vreg) xv, options(pure, nomem, nostack));
9201            i += 16;
9202        }
9203        let mut s = vaddvq_s32(vaddq_s32(vaddq_s32(a0, a1), vaddq_s32(a2, a3)));
9204        while i < n {
9205            s += (*wp.add(i)) as i32 * xq[i] as i32;
9206            i += 1;
9207        }
9208        s
9209    }
9210}
9211
9212/// Row-blocked SDOT: 4 output rows per pass — the activation chunk is
9213/// loaded once and reused, 4 independent accumulators hide sdot latency
9214/// (port of vmfcore `dot_i8_sdot_4rows`).
9215#[cfg(target_arch = "aarch64")]
9216#[target_feature(enable = "neon,dotprod")]
9217unsafe fn dot_i8_sdot_4rows(w0: &[u8], w1: &[u8], w2: &[u8], w3: &[u8], xq: &[i8]) -> [i32; 4] {
9218    // SAFETY: callers uphold slice-length contracts (see call sites).
9219    unsafe {
9220        use core::arch::aarch64::*;
9221        use core::arch::asm;
9222        let n = xq.len();
9223        let px = xq.as_ptr();
9224        let (p0, p1, p2, p3) = (
9225            w0.as_ptr() as *const i8,
9226            w1.as_ptr() as *const i8,
9227            w2.as_ptr() as *const i8,
9228            w3.as_ptr() as *const i8,
9229        );
9230        let (mut a0, mut a1, mut a2, mut a3) = (
9231            vdupq_n_s32(0),
9232            vdupq_n_s32(0),
9233            vdupq_n_s32(0),
9234            vdupq_n_s32(0),
9235        );
9236        let mut i = 0;
9237        while i + 16 <= n {
9238            let x = vld1q_s8(px.add(i));
9239            let v0 = vld1q_s8(p0.add(i));
9240            let v1 = vld1q_s8(p1.add(i));
9241            let v2 = vld1q_s8(p2.add(i));
9242            let v3 = vld1q_s8(p3.add(i));
9243            asm!(
9244                "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
9245                "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
9246                "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
9247                "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
9248                a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
9249                v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
9250                options(pure, nomem, nostack),
9251            );
9252            i += 16;
9253        }
9254        let mut r = [
9255            vaddvq_s32(a0),
9256            vaddvq_s32(a1),
9257            vaddvq_s32(a2),
9258            vaddvq_s32(a3),
9259        ];
9260        while i < n {
9261            let xi = *px.add(i) as i32;
9262            r[0] += (*p0.add(i)) as i32 * xi;
9263            r[1] += (*p1.add(i)) as i32 * xi;
9264            r[2] += (*p2.add(i)) as i32 * xi;
9265            r[3] += (*p3.add(i)) as i32 * xi;
9266            i += 1;
9267        }
9268        r
9269    }
9270}
9271
9272/// 4 interleaved rows in one pass: the repacked group is [r0[c], r1[c],
9273/// r2[c], r3[c]] per 16-byte chunk, so each iteration reads ONE 64-byte
9274/// line plus the shared activation chunk — a single sequential weight
9275/// stream per worker. Per-row accumulation is the same one-accumulator
9276/// scheme as `dot_i8_sdot_4rows`; integer sums are exact, so outputs
9277/// are bit-identical to the mmap-layout kernel.
9278#[cfg(target_arch = "aarch64")]
9279#[target_feature(enable = "neon,dotprod")]
9280unsafe fn dot_i8_sdot_4rows_il(g: &[u8], xq: &[i8]) -> [i32; 4] {
9281    // SAFETY: callers uphold slice-length contracts (g.len() == 4·n,
9282    // n % 16 == 0 — guaranteed by the repack gate).
9283    unsafe {
9284        use core::arch::aarch64::*;
9285        use core::arch::asm;
9286        let n = xq.len();
9287        let px = xq.as_ptr();
9288        let pg = g.as_ptr() as *const i8;
9289        let (mut a0, mut a1, mut a2, mut a3) = (
9290            vdupq_n_s32(0),
9291            vdupq_n_s32(0),
9292            vdupq_n_s32(0),
9293            vdupq_n_s32(0),
9294        );
9295        let mut i = 0;
9296        while i + 16 <= n {
9297            let x = vld1q_s8(px.add(i));
9298            let base = pg.add(4 * i);
9299            let v0 = vld1q_s8(base);
9300            let v1 = vld1q_s8(base.add(16));
9301            let v2 = vld1q_s8(base.add(32));
9302            let v3 = vld1q_s8(base.add(48));
9303            asm!(
9304                "sdot {a0:v}.4s, {v0:v}.16b, {x:v}.16b",
9305                "sdot {a1:v}.4s, {v1:v}.16b, {x:v}.16b",
9306                "sdot {a2:v}.4s, {v2:v}.16b, {x:v}.16b",
9307                "sdot {a3:v}.4s, {v3:v}.16b, {x:v}.16b",
9308                a0 = inout(vreg) a0, a1 = inout(vreg) a1, a2 = inout(vreg) a2, a3 = inout(vreg) a3,
9309                v0 = in(vreg) v0, v1 = in(vreg) v1, v2 = in(vreg) v2, v3 = in(vreg) v3, x = in(vreg) x,
9310                options(pure, nomem, nostack),
9311            );
9312            i += 16;
9313        }
9314        [
9315            vaddvq_s32(a0),
9316            vaddvq_s32(a1),
9317            vaddvq_s32(a2),
9318            vaddvq_s32(a3),
9319        ]
9320    }
9321}
9322
9323/// One q8 row range via SDOT (4-row blocks + tail) — the body of
9324/// `qmatvec`'s hot loop, extracted so multi-matrix jobs can drive the
9325/// SAME kernel for several tensors under one pool dispatch. `rep` — the
9326/// load-time interleaved repack (empty = mmap layout only); rows outside
9327/// full 4-row groups always come from the mmap layout.
9328#[cfg(target_arch = "aarch64")]
9329fn q8_range_sdot(
9330    q: &[u8],
9331    rep: &[u8],
9332    row_scale: &[f32],
9333    act: &SplitAct,
9334    cols: usize,
9335    out_addr: SendMut,
9336    start: usize,
9337    end: usize,
9338) {
9339    let mut o = start;
9340    // Leading rows to the group boundary (repack path only): the pool
9341    // splits row ranges arbitrarily, groups are absolute.
9342    if !rep.is_empty() {
9343        while o < end && o % 4 != 0 {
9344            let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
9345            unsafe { *out_addr.at(o) = v };
9346            o += 1;
9347        }
9348    }
9349    while o + 4 <= end {
9350        let r = if rep.is_empty() {
9351            unsafe {
9352                dot_i8_sdot_4rows(
9353                    &q[o * cols..(o + 1) * cols],
9354                    &q[(o + 1) * cols..(o + 2) * cols],
9355                    &q[(o + 2) * cols..(o + 3) * cols],
9356                    &q[(o + 3) * cols..(o + 4) * cols],
9357                    &act.xq,
9358                )
9359            }
9360        } else {
9361            unsafe { dot_i8_sdot_4rows_il(&rep[o * cols..(o + 4) * cols], &act.xq) }
9362        };
9363        for k in 0..4 {
9364            let mut acc = r[k] as f32 * act.sx;
9365            for &(j, xv) in &act.outliers {
9366                acc += (q[(o + k) * cols + j] as i8) as f32 * xv;
9367            }
9368            // SAFETY: disjoint row ranges per worker.
9369            unsafe { *out_addr.at(o + k) = acc * row_scale[o + k] };
9370        }
9371        o += 4;
9372    }
9373    while o < end {
9374        let v = row_dot_sdot(&q[o * cols..(o + 1) * cols], act) * row_scale[o];
9375        unsafe { *out_addr.at(o) = v };
9376        o += 1;
9377    }
9378}
9379
9380/// Two-input q8 row range via SDOT — `qmatvec2`'s hot loop, extracted
9381/// for the fused pair multi-matrix job (`matvec2_many`).
9382#[cfg(target_arch = "aarch64")]
9383#[allow(clippy::too_many_arguments)]
9384fn q8_range2_sdot(
9385    q: &[u8],
9386    row_scale: &[f32],
9387    a1: &SplitAct,
9388    a2: &SplitAct,
9389    cols: usize,
9390    p1: SendMut,
9391    p2: SendMut,
9392    start: usize,
9393    end: usize,
9394) {
9395    for o in start..end {
9396        let row = &q[o * cols..(o + 1) * cols];
9397        // SAFETY: disjoint row ranges per worker.
9398        unsafe {
9399            *p1.at(o) = row_dot_sdot(row, a1) * row_scale[o];
9400            *p2.at(o) = row_dot_sdot(row, a2) * row_scale[o];
9401        }
9402    }
9403}
9404
9405/// Two-input q8 row range, f32 kernel (non-SDOT) — same extraction.
9406#[allow(clippy::too_many_arguments)]
9407fn q8_range2_f32(
9408    q: &[u8],
9409    row_scale: &[f32],
9410    x1: &[f32],
9411    x2: &[f32],
9412    cols: usize,
9413    p1: SendMut,
9414    p2: SendMut,
9415    start: usize,
9416    end: usize,
9417) {
9418    for o in start..end {
9419        let row = &q[o * cols..(o + 1) * cols];
9420        // SAFETY: disjoint row ranges per worker.
9421        unsafe {
9422            *p1.at(o) = dot_i8_f32(row, x1) * row_scale[o];
9423            *p2.at(o) = dot_i8_f32(row, x2) * row_scale[o];
9424        }
9425    }
9426}
9427
9428/// Scalar/NEON-f32 q8 row range (non-SDOT platforms) — same extraction.
9429fn q8_range_f32(
9430    q: &[u8],
9431    row_scale: &[f32],
9432    xs: &[f32],
9433    cols: usize,
9434    out_addr: SendMut,
9435    start: usize,
9436    end: usize,
9437) {
9438    for o in start..end {
9439        let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
9440        // SAFETY: disjoint row ranges per worker.
9441        unsafe { *out_addr.at(o) = v };
9442    }
9443}
9444
9445/// One q8 row against a split activation, portable: the per-arch fast
9446/// dots where they exist, the exact scalar loop elsewhere. The scalar
9447/// arm is also the test oracle for both fast arms.
9448#[inline]
9449fn q8_row_dot(row: &[u8], act: &SplitAct) -> f32 {
9450    #[cfg(target_arch = "aarch64")]
9451    return row_dot_sdot(row, act);
9452    #[cfg(target_arch = "x86_64")]
9453    return row_dot_avx2(row, act);
9454    #[allow(unreachable_code)]
9455    q8_row_dot_scalar(row, act)
9456}
9457
9458#[allow(dead_code)]
9459fn q8_row_dot_scalar(row: &[u8], act: &SplitAct) -> f32 {
9460    let mut acc = 0i32;
9461    for (k, &b) in row.iter().enumerate() {
9462        acc += (b as i8) as i32 * act.xq[k] as i32;
9463    }
9464    let mut acc = acc as f32 * act.sx;
9465    for &(j, xv) in &act.outliers {
9466        acc += (row[j] as i8) as f32 * xv;
9467    }
9468    acc
9469}
9470
9471/// SDOT row dot with exact outlier correction:
9472/// `dot = sdot(w, xq)·sx + Σ_outl w[j]·x[j]` (then × row_scale by caller).
9473#[cfg(target_arch = "aarch64")]
9474#[inline]
9475fn row_dot_sdot(row: &[u8], act: &SplitAct) -> f32 {
9476    let mut acc = unsafe { dot_i8_sdot(row, &act.xq) } as f32 * act.sx;
9477    for &(j, xv) in &act.outliers {
9478        acc += (row[j] as i8) as f32 * xv;
9479    }
9480    acc
9481}
9482
9483/// One q4 row via SDOT: each 32-group's nibbles unpack to centered i8
9484/// (nib−8 ∈ [−8,7]), int8×int8 `sdot` against the pre-quantized
9485/// activation group, × the group's f16 scale. Returns Σ_g dot_g·s_g;
9486/// the caller multiplies by the activation scale and adds the exact
9487/// outlier terms (port of vmfcore `dot_q4_block_sdot`, +23% measured).
9488/// Nibble order matches the writer: element 2k = low nibble, 2k+1 = high
9489/// → zip(lo,hi) restores flat order.
9490#[cfg(target_arch = "aarch64")]
9491#[target_feature(enable = "neon,dotprod")]
9492unsafe fn dot_q4_row_sdot(packed: &[u8], scales: &[u8], g0: usize, gpr: usize, xq: &[i8]) -> f32 {
9493    // SAFETY: callers uphold slice-length contracts (16 packed bytes and
9494    // 2 scale bytes per group; xq.len() == gpr·GROUP_SIZE).
9495    unsafe {
9496        use core::arch::aarch64::*;
9497        use core::arch::asm;
9498        let lomask = vdupq_n_u8(0x0F);
9499        let eight = vdupq_n_s8(8);
9500        let mut acc = 0f32;
9501        for gi in 0..gpr {
9502            let g = g0 + gi;
9503            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
9504            let b = vld1q_u8(packed.as_ptr().add(g * 16));
9505            let lo = vandq_u8(b, lomask);
9506            let hi = vshrq_n_u8::<4>(b);
9507            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
9508            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
9509            let x0 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE));
9510            let x1 = vld1q_s8(xq.as_ptr().add(gi * GROUP_SIZE + 16));
9511            let (mut a0, mut a1) = (vdupq_n_s32(0), vdupq_n_s32(0));
9512            asm!(
9513                "sdot {a0:v}.4s, {e0:v}.16b, {x0:v}.16b",
9514                "sdot {a1:v}.4s, {e1:v}.16b, {x1:v}.16b",
9515                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
9516                e0 = in(vreg) e0, x0 = in(vreg) x0, e1 = in(vreg) e1, x1 = in(vreg) x1,
9517                options(pure, nomem, nostack),
9518            );
9519            acc += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
9520        }
9521        acc
9522    }
9523}
9524
9525/// Two-activation q4 row via SDOT: the nibble unpack (the expensive
9526/// part) happens ONCE per group; both pre-quantized activations are
9527/// dotted against the same centered i8 registers. Per-lane math matches
9528/// `dot_q4_row_sdot` exactly.
9529#[cfg(target_arch = "aarch64")]
9530#[target_feature(enable = "neon,dotprod")]
9531unsafe fn dot_q4_row_sdot2(
9532    packed: &[u8],
9533    scales: &[u8],
9534    g0: usize,
9535    gpr: usize,
9536    xq1: &[i8],
9537    xq2: &[i8],
9538) -> (f32, f32) {
9539    // SAFETY: callers uphold slice-length contracts (16 packed bytes and
9540    // 2 scale bytes per group; xq*.len() == gpr·GROUP_SIZE).
9541    unsafe {
9542        use core::arch::aarch64::*;
9543        use core::arch::asm;
9544        let lomask = vdupq_n_u8(0x0F);
9545        let eight = vdupq_n_s8(8);
9546        let (mut acc1, mut acc2) = (0f32, 0f32);
9547        for gi in 0..gpr {
9548            let g = g0 + gi;
9549            let s = f16_to_f32(u16::from_le_bytes([scales[g * 2], scales[g * 2 + 1]]));
9550            let b = vld1q_u8(packed.as_ptr().add(g * 16));
9551            let lo = vandq_u8(b, lomask);
9552            let hi = vshrq_n_u8::<4>(b);
9553            let e0 = vsubq_s8(vreinterpretq_s8_u8(vzip1q_u8(lo, hi)), eight);
9554            let e1 = vsubq_s8(vreinterpretq_s8_u8(vzip2q_u8(lo, hi)), eight);
9555            let x10 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE));
9556            let x11 = vld1q_s8(xq1.as_ptr().add(gi * GROUP_SIZE + 16));
9557            let x20 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE));
9558            let x21 = vld1q_s8(xq2.as_ptr().add(gi * GROUP_SIZE + 16));
9559            let (mut a0, mut a1, mut b0, mut b1) = (
9560                vdupq_n_s32(0),
9561                vdupq_n_s32(0),
9562                vdupq_n_s32(0),
9563                vdupq_n_s32(0),
9564            );
9565            asm!(
9566                "sdot {a0:v}.4s, {e0:v}.16b, {x10:v}.16b",
9567                "sdot {a1:v}.4s, {e1:v}.16b, {x11:v}.16b",
9568                "sdot {b0:v}.4s, {e0:v}.16b, {x20:v}.16b",
9569                "sdot {b1:v}.4s, {e1:v}.16b, {x21:v}.16b",
9570                a0 = inout(vreg) a0, a1 = inout(vreg) a1,
9571                b0 = inout(vreg) b0, b1 = inout(vreg) b1,
9572                e0 = in(vreg) e0, e1 = in(vreg) e1,
9573                x10 = in(vreg) x10, x11 = in(vreg) x11,
9574                x20 = in(vreg) x20, x21 = in(vreg) x21,
9575                options(pure, nomem, nostack),
9576            );
9577            acc1 += vaddvq_s32(vaddq_s32(a0, a1)) as f32 * s;
9578            acc2 += vaddvq_s32(vaddq_s32(b0, b1)) as f32 * s;
9579        }
9580        (acc1, acc2)
9581    }
9582}
9583
9584// ───────────────────── fused int8 kernels ─────────────────────
9585
9586/// `acc += w · row` where the row is centered i8 — NEON widen+fma on
9587/// aarch64, scalar elsewhere. The KV-cache q8 value path rides on this.
9588#[inline]
9589pub(crate) fn axpy_i8_f32(acc: &mut [f32], row: &[i8], w: f32) {
9590    #[cfg(target_arch = "aarch64")]
9591    unsafe {
9592        return axpy_i8_f32_neon(acc, row, w);
9593    }
9594    #[cfg(target_arch = "x86_64")]
9595    if avx2_enabled() {
9596        return unsafe { axpy_i8_f32_avx2(acc, row, w) };
9597    }
9598    #[allow(unreachable_code)]
9599    {
9600        for (a, &b) in acc.iter_mut().zip(row) {
9601            *a += w * b as f32;
9602        }
9603    }
9604}
9605
9606/// i8→f32 axpy via AVX2/FMA (x86 mirror of `axpy_i8_f32_neon`).
9607#[cfg(target_arch = "x86_64")]
9608#[target_feature(enable = "avx2,fma")]
9609unsafe fn axpy_i8_f32_avx2(acc: &mut [f32], row: &[i8], w: f32) {
9610    // SAFETY: callers uphold slice-length contracts (see call sites).
9611    unsafe {
9612        use core::arch::x86_64::*;
9613        let n = acc.len().min(row.len());
9614        let ap = acc.as_mut_ptr();
9615        let rp = row.as_ptr();
9616        let wv = _mm256_set1_ps(w);
9617        let mut j = 0usize;
9618        while j + 16 <= n {
9619            let rb = _mm_loadu_si128(rp.add(j) as *const __m128i);
9620            let lo = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(rb));
9621            let hi = _mm256_cvtepi32_ps(_mm256_cvtepi8_epi32(_mm_srli_si128::<8>(rb)));
9622            let v0 = _mm256_fmadd_ps(wv, lo, _mm256_loadu_ps(ap.add(j)));
9623            let v1 = _mm256_fmadd_ps(wv, hi, _mm256_loadu_ps(ap.add(j + 8)));
9624            _mm256_storeu_ps(ap.add(j), v0);
9625            _mm256_storeu_ps(ap.add(j + 8), v1);
9626            j += 16;
9627        }
9628        while j < n {
9629            *ap.add(j) += w * (*rp.add(j)) as f32;
9630            j += 1;
9631        }
9632    }
9633}
9634
9635#[cfg(target_arch = "aarch64")]
9636#[target_feature(enable = "neon")]
9637unsafe fn axpy_i8_f32_neon(acc: &mut [f32], row: &[i8], w: f32) {
9638    // SAFETY: callers uphold slice-length contracts (see call sites).
9639    unsafe {
9640        use core::arch::aarch64::*;
9641        let n = acc.len().min(row.len());
9642        let ap = acc.as_mut_ptr();
9643        let rp = row.as_ptr();
9644        let wv = vdupq_n_f32(w);
9645        let mut j = 0usize;
9646        while j + 16 <= n {
9647            let rb = vld1q_s8(rp.add(j));
9648            let lo = vmovl_s8(vget_low_s8(rb));
9649            let hi = vmovl_s8(vget_high_s8(rb));
9650            for (off, half) in [(0, lo), (8, hi)] {
9651                let f0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(half)));
9652                let f1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(half)));
9653                let o = j + off;
9654                vst1q_f32(ap.add(o), vfmaq_f32(vld1q_f32(ap.add(o)), wv, f0));
9655                vst1q_f32(ap.add(o + 4), vfmaq_f32(vld1q_f32(ap.add(o + 4)), wv, f1));
9656            }
9657            j += 16;
9658        }
9659        while j < n {
9660            *ap.add(j) += w * (*rp.add(j)) as f32;
9661            j += 1;
9662        }
9663    }
9664}
9665
9666/// i8 row · f32 x. NEON on aarch64 (ported from vmfcore `dot_i8_f32_neon`,
9667/// ≈9× scalar), scalar elsewhere.
9668#[inline]
9669pub(crate) fn dot_i8_f32(w: &[u8], x: &[f32]) -> f32 {
9670    #[cfg(target_arch = "aarch64")]
9671    unsafe {
9672        return dot_i8_f32_neon(w, x);
9673    }
9674    #[cfg(target_arch = "x86_64")]
9675    if avx2_enabled() {
9676        return unsafe { dot_i8_f32_avx2(w, x) };
9677    }
9678    #[allow(unreachable_code)]
9679    {
9680        let mut sum = 0.0f32;
9681        for (j, &b) in w.iter().enumerate() {
9682            sum += (b as i8) as f32 * x[j];
9683        }
9684        sum
9685    }
9686}
9687
9688/// i8 row · (x ⊙ col_field) — the q8_2f row dot with the θ col-field
9689/// folded into the product (no prescaled copy of x). NEON on aarch64,
9690/// scalar elsewhere. Used by the active-neuron path `row_dot`.
9691#[inline]
9692fn dot_i8_col_f32(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
9693    #[cfg(target_arch = "aarch64")]
9694    unsafe {
9695        return dot_i8_col_f32_neon(w, x, col);
9696    }
9697    #[allow(unreachable_code)]
9698    {
9699        let mut sum = 0.0f32;
9700        for (j, &b) in w.iter().enumerate() {
9701            sum += (b as i8) as f32 * x[j] * col[j];
9702        }
9703        sum
9704    }
9705}
9706
9707#[cfg(target_arch = "aarch64")]
9708#[target_feature(enable = "neon")]
9709unsafe fn dot_i8_col_f32_neon(w: &[u8], x: &[f32], col: &[f32]) -> f32 {
9710    // SAFETY: callers uphold slice-length contracts (see call sites).
9711    unsafe {
9712        use core::arch::aarch64::*;
9713        let n = x.len();
9714        let wp = w.as_ptr() as *const i8;
9715        let xp = x.as_ptr();
9716        let cp = col.as_ptr();
9717        let (mut a0, mut a1, mut a2, mut a3) = (
9718            vdupq_n_f32(0.0),
9719            vdupq_n_f32(0.0),
9720            vdupq_n_f32(0.0),
9721            vdupq_n_f32(0.0),
9722        );
9723        let mut j = 0usize;
9724        while j + 16 <= n {
9725            let wb = vld1q_s8(wp.add(j));
9726            let lo = vmovl_s8(vget_low_s8(wb));
9727            let hi = vmovl_s8(vget_high_s8(wb));
9728            let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
9729            let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
9730            let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
9731            let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
9732            a0 = vfmaq_f32(
9733                a0,
9734                w0,
9735                vmulq_f32(vld1q_f32(xp.add(j)), vld1q_f32(cp.add(j))),
9736            );
9737            a1 = vfmaq_f32(
9738                a1,
9739                w1,
9740                vmulq_f32(vld1q_f32(xp.add(j + 4)), vld1q_f32(cp.add(j + 4))),
9741            );
9742            a2 = vfmaq_f32(
9743                a2,
9744                w2,
9745                vmulq_f32(vld1q_f32(xp.add(j + 8)), vld1q_f32(cp.add(j + 8))),
9746            );
9747            a3 = vfmaq_f32(
9748                a3,
9749                w3,
9750                vmulq_f32(vld1q_f32(xp.add(j + 12)), vld1q_f32(cp.add(j + 12))),
9751            );
9752            j += 16;
9753        }
9754        let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
9755        while j < n {
9756            sum += (*wp.add(j)) as f32 * *xp.add(j) * *cp.add(j);
9757            j += 1;
9758        }
9759        sum
9760    }
9761}
9762
9763#[cfg(target_arch = "aarch64")]
9764#[target_feature(enable = "neon")]
9765unsafe fn dot_i8_f32_neon(w: &[u8], x: &[f32]) -> f32 {
9766    // SAFETY: callers uphold slice-length contracts (see call sites).
9767    unsafe {
9768        use core::arch::aarch64::*;
9769        let n = x.len();
9770        let wp = w.as_ptr() as *const i8;
9771        let xp = x.as_ptr();
9772        let (mut a0, mut a1, mut a2, mut a3) = (
9773            vdupq_n_f32(0.0),
9774            vdupq_n_f32(0.0),
9775            vdupq_n_f32(0.0),
9776            vdupq_n_f32(0.0),
9777        );
9778        let mut j = 0usize;
9779        while j + 16 <= n {
9780            let wb = vld1q_s8(wp.add(j));
9781            let lo = vmovl_s8(vget_low_s8(wb));
9782            let hi = vmovl_s8(vget_high_s8(wb));
9783            let w0 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(lo)));
9784            let w1 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(lo)));
9785            let w2 = vcvtq_f32_s32(vmovl_s16(vget_low_s16(hi)));
9786            let w3 = vcvtq_f32_s32(vmovl_s16(vget_high_s16(hi)));
9787            a0 = vfmaq_f32(a0, w0, vld1q_f32(xp.add(j)));
9788            a1 = vfmaq_f32(a1, w1, vld1q_f32(xp.add(j + 4)));
9789            a2 = vfmaq_f32(a2, w2, vld1q_f32(xp.add(j + 8)));
9790            a3 = vfmaq_f32(a3, w3, vld1q_f32(xp.add(j + 12)));
9791            j += 16;
9792        }
9793        let mut sum = vaddvq_f32(vaddq_f32(vaddq_f32(a0, a1), vaddq_f32(a2, a3)));
9794        while j < n {
9795            sum += (*wp.add(j)) as f32 * *xp.add(j);
9796            j += 1;
9797        }
9798        sum
9799    }
9800}
9801
9802#[allow(clippy::too_many_arguments)]
9803fn qmatvec(
9804    q: &[u8],
9805    rep: &[u8],
9806    row_scale: &[f32],
9807    x: &[f32],
9808    col_field: &[f32],
9809    dtype: TensorDtype,
9810    rows: usize,
9811    cols: usize,
9812    out: &mut [f32],
9813    pool: Option<&Pool>,
9814) {
9815    debug_assert_eq!(out.len(), rows);
9816    #[cfg(not(target_arch = "aarch64"))]
9817    let _ = rep;
9818
9819    #[cfg(target_arch = "aarch64")]
9820    if sdot_enabled() {
9821        let act = if dtype == TensorDtype::Q8_2f {
9822            split_act_q8_2f(x, col_field)
9823        } else {
9824            split_act(x)
9825        };
9826        let out_addr = SendMut(out.as_mut_ptr());
9827        let run_range = |start: usize, end: usize| {
9828            q8_range_sdot(q, rep, row_scale, &act, cols, out_addr, start, end)
9829        };
9830        match pool {
9831            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
9832            _ => run_range(0, rows),
9833        }
9834        return;
9835    }
9836    // x86 A8W8 via AVX2 maddubs — same quantized-activation contract as
9837    // the SDOT path (CMF_AVX2=0 keeps the exact i8×f32 loop).
9838    #[cfg(target_arch = "x86_64")]
9839    if avx2_a8w8_enabled() {
9840        let act = if dtype == TensorDtype::Q8_2f {
9841            split_act_q8_2f(x, col_field)
9842        } else {
9843            split_act(x)
9844        };
9845        let out_addr = SendMut(out.as_mut_ptr());
9846        let run_range = |start: usize, end: usize| {
9847            q8_range_avx2(q, row_scale, &act, cols, out_addr, start, end)
9848        };
9849        match pool {
9850            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
9851            _ => run_range(0, rows),
9852        }
9853        return;
9854    }
9855
9856    prescale_with(x, col_field, dtype, 1, |xs| {
9857        let out_addr = SendMut(out.as_mut_ptr());
9858        let run_range = move |start: usize, end: usize| {
9859            for o in start..end {
9860                let v = dot_i8_f32(&q[o * cols..(o + 1) * cols], xs) * row_scale[o];
9861                // SAFETY: disjoint row ranges per worker.
9862                unsafe { *out_addr.at(o) = v };
9863            }
9864        };
9865        match pool {
9866            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
9867            _ => run_range(0, rows),
9868        }
9869    });
9870}
9871
9872#[allow(clippy::too_many_arguments)]
9873fn qmatvec2(
9874    q: &[u8],
9875    row_scale: &[f32],
9876    x1: &[f32],
9877    x2: &[f32],
9878    col_field: &[f32],
9879    dtype: TensorDtype,
9880    rows: usize,
9881    cols: usize,
9882    o1: &mut [f32],
9883    o2: &mut [f32],
9884    pool: Option<&Pool>,
9885) {
9886    #[cfg(target_arch = "aarch64")]
9887    if sdot_enabled() {
9888        let a1s = if dtype == TensorDtype::Q8_2f {
9889            split_act_q8_2f(x1, col_field)
9890        } else {
9891            split_act(x1)
9892        };
9893        let a2s = if dtype == TensorDtype::Q8_2f {
9894            split_act_q8_2f(x2, col_field)
9895        } else {
9896            split_act(x2)
9897        };
9898        let p1 = SendMut(o1.as_mut_ptr());
9899        let p2 = SendMut(o2.as_mut_ptr());
9900        let run_range = |start: usize, end: usize| {
9901            q8_range2_sdot(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
9902        };
9903        match pool {
9904            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
9905            _ => run_range(0, rows),
9906        }
9907        return;
9908    }
9909    #[cfg(target_arch = "x86_64")]
9910    if avx2_a8w8_enabled() {
9911        let a1s = if dtype == TensorDtype::Q8_2f {
9912            split_act_q8_2f(x1, col_field)
9913        } else {
9914            split_act(x1)
9915        };
9916        let a2s = if dtype == TensorDtype::Q8_2f {
9917            split_act_q8_2f(x2, col_field)
9918        } else {
9919            split_act(x2)
9920        };
9921        let p1 = SendMut(o1.as_mut_ptr());
9922        let p2 = SendMut(o2.as_mut_ptr());
9923        let run_range = |start: usize, end: usize| {
9924            q8_range2_avx2(q, row_scale, &a1s, &a2s, cols, p1, p2, start, end)
9925        };
9926        match pool {
9927            Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
9928            _ => run_range(0, rows),
9929        }
9930        return;
9931    }
9932
9933    prescale_with(x1, col_field, dtype, 1, |x1s| {
9934        prescale_with(x2, col_field, dtype, 2, |x2s| {
9935            let p1 = SendMut(o1.as_mut_ptr());
9936            let p2 = SendMut(o2.as_mut_ptr());
9937            let run_range = move |start: usize, end: usize| {
9938                for o in start..end {
9939                    let row = &q[o * cols..(o + 1) * cols];
9940                    let s1 = dot_i8_f32(row, x1s) * row_scale[o];
9941                    let s2 = dot_i8_f32(row, x2s) * row_scale[o];
9942                    // SAFETY: disjoint row ranges per worker.
9943                    unsafe {
9944                        *p1.at(o) = s1;
9945                        *p2.at(o) = s2;
9946                    }
9947                }
9948            };
9949            match pool {
9950                Some(pool) if rows >= 256 => pool.run_rows(rows, &run_range),
9951                _ => run_range(0, rows),
9952            }
9953        });
9954    });
9955}
9956
9957#[derive(Clone, Copy)]
9958struct SendMut(*mut f32);
9959unsafe impl Send for SendMut {}
9960unsafe impl Sync for SendMut {}
9961
9962impl SendMut {
9963    #[inline]
9964    fn at(self, i: usize) -> *mut f32 {
9965        unsafe { self.0.add(i) }
9966    }
9967}
9968
9969#[cfg(test)]
9970mod tests {
9971    use super::*;
9972
9973    #[test]
9974    fn q2tp_i8_dot_matches_exact_on_grid() {
9975        // On-grid activations (±1 → sx=1/127, xq=±127 dequantizes
9976        // exactly, no outliers) must make the integer path agree with
9977        // the exact scalar walk to f32 rounding.
9978        let (rows, cols) = (5, 64);
9979        let gpr = cols / GROUP_SIZE;
9980        // Synthetic codes plane + a flat ladder: scales_into is not under
9981        // test here, so drive dot_q2tp_row_i8 / q2tp_row_exact directly
9982        // with hand-made scales.
9983        let chunks: Vec<u8> = (0..rows * gpr * Q2TP_CHUNK)
9984            .map(|i| (i as u32).wrapping_mul(2654435761) as u8)
9985            .collect();
9986        let scales: Vec<f32> = (0..gpr).map(|g| 0.5 + g as f32 * 0.25).collect();
9987        let x: Vec<f32> = (0..cols).map(|i| if i % 3 == 0 { -1.0 } else { 1.0 }).collect();
9988        let act = split_act(&x);
9989        assert!(act.outliers.is_empty(), "on-grid input must have no outliers");
9990        let gsum = q1_group_sums(&act.xq, gpr);
9991        for r in 0..rows {
9992            let exact = q2tp_row_exact(&chunks, r, gpr, &x, &scales);
9993            let fast = dot_q2tp_row_i8(&chunks, r, gpr, &act.xq, &gsum, &scales) * act.sx;
9994            assert!(
9995                (exact - fast).abs() <= exact.abs() * 1e-5 + 1e-5,
9996                "row {r}: exact {exact} vs i8 {fast}"
9997            );
9998        }
9999    }
10000
10001    #[test]
10002    fn q8_row_dot_fast_matches_scalar() {
10003        // The per-arch fast dot must agree with the exact scalar oracle
10004        // (same contract the fused q8 FFN arm rides on).
10005        let cols = 96;
10006        let row: Vec<u8> = (0..cols)
10007            .map(|i| ((i as i32 * 37 % 251) - 125) as i8 as u8)
10008            .collect();
10009        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.13).sin()).collect();
10010        let act = split_act(&x);
10011        let fast = q8_row_dot(&row, &act);
10012        let scalar = q8_row_dot_scalar(&row, &act);
10013        assert!(
10014            (fast - scalar).abs() <= scalar.abs() * 1e-5 + 1e-5,
10015            "fast {fast} vs scalar {scalar}"
10016        );
10017    }
10018
10019    #[test]
10020    fn f32_matvec_matches_matvec_rows_bitexact() {
10021        let (rows, cols) = (300, 40);
10022        let w: Vec<f32> = (0..rows * cols).map(|i| (i as f32 * 0.017).sin()).collect();
10023        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.05).cos()).collect();
10024        let qt = QTensor::from_f32(w.clone(), rows, cols);
10025
10026        let mut a = vec![0.0f32; rows];
10027        matvec_rows(None, &w, &x, &mut a);
10028        let mut b = vec![0.0f32; rows];
10029        qt.matvec(&x, &mut b, None);
10030        assert_eq!(a, b);
10031    }
10032
10033    #[test]
10034    fn sdot_kernel_exact_on_grid() {
10035        // Activations already on the i8 grid (±1 with amax=1 → sx=1/127,
10036        // xq=±127 dequantizes EXACTLY) → the SDOT path must match the
10037        // exact f32 dot to float rounding. This isolates kernel
10038        // correctness from quantization noise.
10039        eprintln!("sdot_enabled = {}", sdot_enabled());
10040        let (rows, cols) = (9, 80); // odd rows → exercises 4-row + tail
10041        let w: Vec<u8> = (0..rows * cols)
10042            .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
10043            .collect();
10044        let scales: Vec<f32> = (0..rows).map(|o| 0.005 + o as f32 * 0.001).collect();
10045        let x: Vec<f32> = (0..cols)
10046            .map(|i| match i % 3 {
10047                0 => 1.0,
10048                1 => -1.0,
10049                _ => 0.0,
10050            })
10051            .collect();
10052        let mut a = vec![0.0f32; rows];
10053        qmatvec(
10054            &w,
10055            &[],
10056            &scales,
10057            &x,
10058            &[],
10059            TensorDtype::Q8Row,
10060            rows,
10061            cols,
10062            &mut a,
10063            None,
10064        );
10065        for o in 0..rows {
10066            let mut acc = 0.0f32;
10067            for j in 0..cols {
10068                acc += (w[o * cols + j] as i8) as f32 * x[j];
10069            }
10070            let expect = acc * scales[o];
10071            assert!(
10072                (a[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
10073                "row {o}: {} vs {expect}",
10074                a[o]
10075            );
10076        }
10077    }
10078
10079    #[test]
10080    fn q1_tbl_fast_path_matches_reference() {
10081        // gpr = 8 exercises the TBL pair-load fast loop, and the LAST
10082        // row's final 4-tile window trips the 4B-overread guard (the
10083        // payload ends exactly at the last tile) — both paths must
10084        // agree with the dequant reference.
10085        let (rows, cols) = (5, 256);
10086        let gpr = cols / GROUP_SIZE;
10087        let mut bytes = Vec::new();
10088        for t in 0..rows * gpr {
10089            let s = 0.007 + (t % 11) as f32 * 0.004;
10090            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
10091            for j in 0..4 {
10092                bytes.push(((t * 53 + j * 89 + 7) % 249) as u8);
10093            }
10094        }
10095        let x: Vec<f32> = (0..cols)
10096            .map(|i| if (i * 5) % 7 < 3 { 1.0 } else { -1.0 })
10097            .collect();
10098        let mut w = vec![0.0f32; rows * cols];
10099        cortiq_core::quant::dequant_q1(&bytes, &mut w);
10100        let mut got = vec![0.0f32; rows];
10101        q1_matvec(&bytes, &x, rows, cols, &mut got, None);
10102        for o in 0..rows {
10103            let expect: f32 = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
10104            assert!(
10105                (got[o] - expect).abs() < 1e-3 * expect.abs().max(1e-3),
10106                "row {o}: {} vs {expect}",
10107                got[o]
10108            );
10109        }
10110        // Blocked 1×4 batch (b=5: one quad + remainder) must equal the
10111        // single-matvec path bit-for-bit.
10112        let b = 5usize;
10113        let mut xs_all = Vec::new();
10114        for bi in 0..b {
10115            xs_all.extend(x.iter().map(|v| if bi % 2 == 0 { *v } else { -*v }));
10116        }
10117        let mut mm = vec![0.0f32; b * rows];
10118        q1_matmat(&bytes, &xs_all, b, rows, cols, &mut mm, None);
10119        for bi in 0..b {
10120            let mut single = vec![0.0f32; rows];
10121            q1_matvec(
10122                &bytes,
10123                &xs_all[bi * cols..(bi + 1) * cols],
10124                rows,
10125                cols,
10126                &mut single,
10127                None,
10128            );
10129            assert_eq!(&mm[bi * rows..(bi + 1) * rows], &single[..], "stream {bi}");
10130        }
10131    }
10132
10133    #[test]
10134    fn q1_kernels_match_exact_reference() {
10135        // Synthetic q1 payload: 6-byte tiles [f16 scale][4B bits].
10136        let (rows, cols) = (7, 96);
10137        let gpr = cols / GROUP_SIZE;
10138        let mut bytes = Vec::new();
10139        for t in 0..rows * gpr {
10140            let s = 0.01 + (t % 13) as f32 * 0.003;
10141            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
10142            for j in 0..4 {
10143                bytes.push(((t * 31 + j * 97) % 251) as u8);
10144            }
10145        }
10146        // On-grid activations (±1, amax 1) → the SDOT path is exact.
10147        let x: Vec<f32> = (0..cols)
10148            .map(|i| if i % 3 == 0 { 1.0 } else { -1.0 })
10149            .collect();
10150        // Reference through the core dequant.
10151        let mut w = vec![0.0f32; rows * cols];
10152        cortiq_core::quant::dequant_q1(&bytes, &mut w);
10153        let mut expect = vec![0.0f32; rows];
10154        for o in 0..rows {
10155            expect[o] = (0..cols).map(|j| w[o * cols + j] * x[j]).sum();
10156        }
10157        let mut got = vec![0.0f32; rows];
10158        q1_matvec(&bytes, &x, rows, cols, &mut got, None);
10159        for o in 0..rows {
10160            assert!(
10161                (got[o] - expect[o]).abs() < 1e-3 * expect[o].abs().max(1e-3),
10162                "row {o}: {} vs {}",
10163                got[o],
10164                expect[o]
10165            );
10166        }
10167        // Pair and batch paths agree with the single path.
10168        let x2: Vec<f32> = x.iter().map(|v| -v).collect();
10169        let (mut a1, mut a2) = (vec![0.0f32; rows], vec![0.0f32; rows]);
10170        q1_matvec2(&bytes, &x, &x2, rows, cols, &mut a1, &mut a2, None);
10171        assert_eq!(a1, got);
10172        let mut xs = x.clone();
10173        xs.extend_from_slice(&x2);
10174        let mut mm = vec![0.0f32; 2 * rows];
10175        q1_matmat(&bytes, &xs, 2, rows, cols, &mut mm, None);
10176        assert_eq!(&mm[..rows], got.as_slice());
10177        assert_eq!(&mm[rows..], a2.as_slice());
10178    }
10179
10180    #[test]
10181    fn repack_is_bit_identical() {
10182        // The interleaved-repack kernel must produce EXACTLY the same
10183        // bits as the mmap-layout kernel: integer accumulation is order-
10184        // exact, the f32 epilogue is identical. Odd rows exercise the
10185        // tail; direct range calls exercise unaligned pool splits.
10186        let (rows, cols) = (267, 96); // 66 groups + 3 tail rows, cols % 16 == 0
10187        let w: Vec<u8> = (0..rows * cols)
10188            .map(|i| (((i * 89) % 253) as i32 - 126) as i8 as u8)
10189            .collect();
10190        let scales: Vec<f32> = (0..rows).map(|o| 0.003 + o as f32 * 0.0007).collect();
10191        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.37).sin() * 2.0).collect();
10192        let rep = q8_repack_layout(&w, rows, cols);
10193        // Group interleave round-trips.
10194        for g in 0..rows / 4 {
10195            for c in 0..cols / 16 {
10196                for lane in 0..4 {
10197                    assert_eq!(
10198                        &rep[g * 4 * cols + c * 64 + lane * 16
10199                            ..g * 4 * cols + c * 64 + lane * 16 + 16],
10200                        &w[(g * 4 + lane) * cols + c * 16..(g * 4 + lane) * cols + c * 16 + 16],
10201                    );
10202                }
10203            }
10204        }
10205        let mut a = vec![0.0f32; rows];
10206        qmatvec(
10207            &w,
10208            &[],
10209            &scales,
10210            &x,
10211            &[],
10212            TensorDtype::Q8Row,
10213            rows,
10214            cols,
10215            &mut a,
10216            None,
10217        );
10218        let mut b = vec![0.0f32; rows];
10219        qmatvec(
10220            &w,
10221            &rep,
10222            &scales,
10223            &x,
10224            &[],
10225            TensorDtype::Q8Row,
10226            rows,
10227            cols,
10228            &mut b,
10229            None,
10230        );
10231        assert_eq!(a, b, "full-range repack output diverged");
10232
10233        #[cfg(target_arch = "aarch64")]
10234        if sdot_enabled() {
10235            // Unaligned range split (pool workers get arbitrary bounds).
10236            let act = split_act(&x);
10237            let mut c1 = vec![0.0f32; rows];
10238            let mut c2 = vec![0.0f32; rows];
10239            q8_range_sdot(
10240                &w,
10241                &[],
10242                &scales,
10243                &act,
10244                cols,
10245                SendMut(c1.as_mut_ptr()),
10246                3,
10247                rows - 2,
10248            );
10249            q8_range_sdot(
10250                &w,
10251                &rep,
10252                &scales,
10253                &act,
10254                cols,
10255                SendMut(c2.as_mut_ptr()),
10256                3,
10257                rows - 2,
10258            );
10259            assert_eq!(c1, c2, "unaligned-range repack output diverged");
10260        }
10261    }
10262
10263    #[test]
10264    fn sdot_a8w8_noise_is_bounded() {
10265        // Off-grid activations: A8 quantization noise must stay small in
10266        // relative L2 over the whole output (realistic accuracy contract;
10267        // vmfcore measured argmax-identical decode on real models).
10268        let (rows, cols) = (16, 512);
10269        let w: Vec<u8> = (0..rows * cols)
10270            .map(|i| (((i * 37) % 251) as i32 - 125) as i8 as u8)
10271            .collect();
10272        let scales = vec![0.01f32; rows];
10273        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.21).sin()).collect();
10274        let mut a = vec![0.0f32; rows];
10275        qmatvec(
10276            &w,
10277            &[],
10278            &scales,
10279            &x,
10280            &[],
10281            TensorDtype::Q8Row,
10282            rows,
10283            cols,
10284            &mut a,
10285            None,
10286        );
10287        let (mut num, mut den) = (0f64, 0f64);
10288        for o in 0..rows {
10289            let mut acc = 0.0f32;
10290            for j in 0..cols {
10291                acc += (w[o * cols + j] as i8) as f32 * x[j];
10292            }
10293            let expect = acc * scales[o];
10294            num += ((a[o] - expect) as f64).powi(2);
10295            den += (expect as f64).powi(2);
10296        }
10297        let rel = (num / den.max(1e-12)).sqrt();
10298        assert!(rel < 0.05, "A8W8 relative L2 error too high: {rel}");
10299    }
10300
10301    #[test]
10302    fn i8_dot_neon_matches_scalar() {
10303        let n = 100;
10304        let w: Vec<u8> = (0..n).map(|i| ((i * 37 + 11) % 251) as u8).collect();
10305        let x: Vec<f32> = (0..n).map(|i| (i as f32 * 0.13).sin()).collect();
10306        let mut scalar = 0.0f32;
10307        for j in 0..n {
10308            scalar += (w[j] as i8) as f32 * x[j];
10309        }
10310        let fast = dot_i8_f32(&w, &x);
10311        assert!((scalar - fast).abs() < 1e-3 * scalar.abs().max(1.0));
10312    }
10313
10314    /// Fused vbit matvec must match full dequant_vbit + dense matvec.
10315    #[test]
10316    fn vbitmatvec_matches_full_dequant() {
10317        let (rows, cols) = (6, 64);
10318        let ng = cols / GROUP_SIZE;
10319        // Hand-craft: bits per row, f16 scales, packed rows.
10320        let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
10321        let mut bytes = bits.clone();
10322        for g in 0..rows * ng {
10323            let s = 0.02 + 0.001 * g as f32;
10324            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
10325        }
10326        for r in 0..rows {
10327            let b = bits[r] as usize;
10328            let (mut acc, mut nb) = (0u64, 0usize);
10329            let mut rowbytes = Vec::new();
10330            for i in 0..cols {
10331                let v = ((i * 7 + r * 13) % (1 << b)) as u64;
10332                acc = (acc << b) | v;
10333                nb += b;
10334                while nb >= 8 {
10335                    nb -= 8;
10336                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
10337                }
10338            }
10339            if nb > 0 {
10340                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
10341            }
10342            bytes.extend_from_slice(&rowbytes);
10343        }
10344        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
10345
10346        let mut reference = vec![0f32; rows * cols];
10347        cortiq_core::quant::dequant_vbit(&bytes, rows, cols, &mut reference).unwrap();
10348        let mut expect = vec![0f32; rows];
10349        for r in 0..rows {
10350            expect[r] = reference[r * cols..(r + 1) * cols]
10351                .iter()
10352                .zip(&x)
10353                .map(|(w, xv)| w * xv)
10354                .sum();
10355        }
10356        let mut got = vec![0f32; rows];
10357        let offsets = vbit_row_offsets(&bytes, rows, cols);
10358        vbitmatvec(&bytes, &offsets, &x, rows, cols, &mut got, None);
10359        // SDOT path quantizes activations to i8 (A8W8): bounded noise,
10360        // same contract as q8 (exact path is pinned by CMF_SDOT=0 in
10361        // the golden-parity gate).
10362        let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
10363        let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
10364        for r in 0..rows {
10365            assert!(
10366                (got[r] - expect[r]).abs() < tol * scale,
10367                "row {r}: {} vs {}",
10368                got[r],
10369                expect[r]
10370            );
10371        }
10372    }
10373
10374    /// Fused q4 matvec must match the reference full-dequant + dense
10375    /// matvec bit-for-bit in structure (same f32 math, group order).
10376    /// vbit matmat: the blocked 1×4 leg must match the per-row path
10377    /// (paired env toggle; larger shape so both code paths engage).
10378    #[test]
10379    #[cfg(target_arch = "x86_64")]
10380    fn vbit_matmat_blocked_matches_per_row() {
10381        let (rows, cols, b) = (64usize, 128usize, 9usize);
10382        let ng = cols / GROUP_SIZE;
10383        let bits: Vec<u8> = (0..rows).map(|r| [3u8, 4, 5, 6][r % 4]).collect();
10384        let mut bytes = bits.clone();
10385        for g in 0..rows * ng {
10386            let sc = 0.02 + 0.0005 * g as f32;
10387            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
10388        }
10389        for r in 0..rows {
10390            let bw = bits[r] as usize;
10391            let (mut acc, mut nb) = (0u64, 0usize);
10392            let mut rowbytes = Vec::new();
10393            for i in 0..cols {
10394                let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
10395                acc = (acc << bw) | v;
10396                nb += bw;
10397                while nb >= 8 {
10398                    nb -= 8;
10399                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
10400                }
10401            }
10402            if nb > 0 {
10403                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
10404            }
10405            bytes.extend_from_slice(&rowbytes);
10406        }
10407        let x: Vec<f32> = (0..b * cols)
10408            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
10409            .collect();
10410        let offsets = vbit_row_offsets(&bytes, rows, cols);
10411        let mut y_a = vec![0f32; b * rows];
10412        let mut y_b = vec![0f32; b * rows];
10413        unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
10414        vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_a, None);
10415        unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
10416        vbitmatmat(&bytes, &offsets, &x, b, rows, cols, &mut y_b, None);
10417        unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
10418        let max_d = y_a
10419            .iter()
10420            .zip(&y_b)
10421            .map(|(p, q)| (p - q).abs())
10422            .fold(0.0f32, f32::max);
10423        assert!(max_d < 1e-4, "vbit blocked ≠ per-row: max|Δ| = {max_d}");
10424    }
10425
10426    /// q4t blocked 1×4 (SDOT on ARM, AVX2 on x86) must equal the
10427    /// per-row path exactly: same nibble unpack, same group order,
10428    /// same f32 accumulation — batch == matvec bit-for-bit. b=9 covers
10429    /// two full 1×4 blocks plus a remainder through the single-row
10430    /// kernel. (Both paths produce identical output, so the shared
10431    /// CMF_X86_BLOCKED env var racing with other tests cannot flip
10432    /// the verdict — worst case both sides take the same path.)
10433    #[test]
10434    fn q4t_matmat_blocked_matches_per_row() {
10435        let (rows, cols, b) = (16usize, 64usize, 9usize);
10436        let gpr = cols / GROUP_SIZE;
10437        let mut bytes = vec![0u8; rows * gpr * Q4_TILE];
10438        for r in 0..rows {
10439            for g in 0..gpr {
10440                let t = (r * gpr + g) * Q4_TILE;
10441                let sc = 0.02 + 0.001 * (r * gpr + g) as f32;
10442                bytes[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
10443                for k in 0..16 {
10444                    bytes[t + 2 + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
10445                }
10446            }
10447        }
10448        let x: Vec<f32> = (0..b * cols)
10449            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
10450            .collect();
10451        let mut y_blk = vec![0f32; b * rows];
10452        let mut y_row = vec![0f32; b * rows];
10453        unsafe { std::env::set_var("CMF_X86_BLOCKED", "1") };
10454        q4t_matmat(&bytes, &x, b, rows, cols, &mut y_blk, None);
10455        unsafe { std::env::set_var("CMF_X86_BLOCKED", "0") };
10456        q4t_matmat(&bytes, &x, b, rows, cols, &mut y_row, None);
10457        unsafe { std::env::remove_var("CMF_X86_BLOCKED") };
10458        assert_eq!(y_blk, y_row, "q4t blocked 1x4 ≠ per-row");
10459    }
10460
10461    /// The wide-batch Accelerate arm of q4t_matmat vs a brute-force
10462    /// f32 dequant matmul: both are f32 GEMMs, so only reduction
10463    /// order differs — tight tolerance.
10464    /// A synthetic q4tp payload: random nibbles plus a per-row ladder whose
10465    /// span varies row to row, so the codes actually exercise the full 0..31
10466    /// range rather than clustering on one rung.
10467    fn synth_q4tp(rows: usize, cols: usize) -> Vec<u8> {
10468        use cortiq_core::quant::{f32_to_f16, q4tp_code_stride, q4tp_put_code};
10469        let gpr = cols / GROUP_SIZE;
10470        let stride = q4tp_code_stride(gpr);
10471        let (params_off, codes_off, _) = q4tp_sections(rows, cols);
10472        let mut b = vec![0u8; codes_off + rows * stride];
10473        for r in 0..rows {
10474            for g in 0..gpr {
10475                let t = (r * gpr + g) * Q4TP_NIB;
10476                for k in 0..16 {
10477                    b[t + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
10478                }
10479            }
10480            let lo = -6.0 - 0.03 * (r % 17) as f32;
10481            let step = 0.01 + 0.004 * (r % 11) as f32;
10482            let p = params_off + r * 4;
10483            b[p..p + 2].copy_from_slice(&f32_to_f16(lo).to_le_bytes());
10484            b[p + 2..p + 4].copy_from_slice(&f32_to_f16(step).to_le_bytes());
10485            let crow = &mut b[codes_off + r * stride..codes_off + (r + 1) * stride];
10486            for g in 0..gpr {
10487                q4tp_put_code(crow, g, (r * 5 + g * 3) % 32);
10488            }
10489        }
10490        b
10491    }
10492
10493    /// The same weights re-expressed as q4_tiled, so the proven kernel can
10494    /// be the reference: each tile stores the ladder scale its code selects.
10495    /// Only the f16 rounding of that scale separates the two payloads.
10496    fn q4tp_as_q4t(bytes: &[u8], rows: usize, cols: usize) -> Vec<u8> {
10497        let gpr = cols / GROUP_SIZE;
10498        let v = Q4tpView::new(bytes, rows, cols);
10499        let mut out = vec![0u8; rows * gpr * Q4_TILE];
10500        let mut sc = vec![0f32; gpr];
10501        for r in 0..rows {
10502            v.scales_into(r, gpr, &mut sc);
10503            for g in 0..gpr {
10504                let t = (r * gpr + g) * Q4_TILE;
10505                let s = sc[g];
10506                out[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
10507                let src = (r * gpr + g) * Q4TP_NIB;
10508                out[t + 2..t + Q4_TILE].copy_from_slice(&v.nib[src..src + Q4TP_NIB]);
10509            }
10510        }
10511        out
10512    }
10513
10514    /// The exact (`CMF_SDOT=0`) path must reproduce `dequant_q4tp` to f32
10515    /// rounding — that scalar routine is the format's definition, and the
10516    /// kernels re-derive the scale from the ladder independently. Call the
10517    /// row kernel directly: `matmat` picks the int8 arm when a8w8 is on,
10518    /// so routing through it would test the other path by accident.
10519    #[test]
10520    fn q4tp_exact_path_matches_dequant_reference() {
10521        let (rows, cols) = (256usize, 512usize);
10522        let gpr = cols / GROUP_SIZE;
10523        let bytes = synth_q4tp(rows, cols);
10524        let mut w = vec![0f32; rows * cols];
10525        cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
10526
10527        let x: Vec<f32> = (0..cols)
10528            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
10529            .collect();
10530        let v = Q4tpView::new(&bytes, rows, cols);
10531        let mut sc = vec![0f32; gpr];
10532        for r in 0..rows {
10533            v.scales_into(r, gpr, &mut sc);
10534            let got = q4tp_row_exact(v.nib, r, gpr, &x, &sc);
10535            let want: f32 = (0..cols).map(|c| w[r * cols + c] * x[c]).sum();
10536            // These dot products cancel down to ~1e-3 from terms of ~5e-2, so
10537            // the meaningful yardstick is the summed magnitude, not the result:
10538            // against the result any reordering of a 512-term f32 sum "fails".
10539            let mag: f32 = (0..cols).map(|c| (w[r * cols + c] * x[c]).abs()).sum();
10540            assert!(
10541                (got - want).abs() <= 1e-5 * mag,
10542                "row {r}: kernel {got} vs dequant {want}"
10543            );
10544        }
10545    }
10546
10547    /// The int8 (a8w8) path can't be checked against an f32 reference — the
10548    /// activation quantization dominates. Check it against the q4t kernel it
10549    /// was ported from instead, on payloads holding the same weights: that
10550    /// isolates exactly what the port could break (16 B stride, ladder
10551    /// lookup, nibble unpack) from what it deliberately shares.
10552    #[test]
10553    fn q4tp_matvec_matches_the_q4t_kernel_it_was_ported_from() {
10554        let (rows, cols) = (256usize, 512usize);
10555        let bytes = synth_q4tp(rows, cols);
10556        let twin = q4tp_as_q4t(&bytes, rows, cols);
10557        let x: Vec<f32> = (0..cols)
10558            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
10559            .collect();
10560
10561        let mut got = vec![0f32; rows];
10562        q4tp_matvec(&bytes, &x, rows, cols, &mut got, None);
10563        let mut want = vec![0f32; rows];
10564        q4t_matvec(&twin, &x, rows, cols, &mut want, None);
10565
10566        // Scale is f16 in the twin and f32 here, so allow that rounding on
10567        // top of the summed magnitude (same cancellation argument as above).
10568        let mut w = vec![0f32; rows * cols];
10569        cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
10570        for r in 0..rows {
10571            let mag: f32 = (0..cols).map(|c| (w[r * cols + c] * x[c]).abs()).sum();
10572            assert!(
10573                (got[r] - want[r]).abs() <= 1e-3 * mag,
10574                "row {r}: q4tp {} vs q4t {}",
10575                got[r],
10576                want[r]
10577            );
10578        }
10579    }
10580
10581    /// `matmat` carries three arms (Accelerate, blocked int8 1x4, scalar).
10582    /// Batch 5 crosses the blocked kernel's stride, so this exercises the
10583    /// 1x4 path AND its scalar tail in one run — the blocked kernel is new
10584    /// code and its four accumulators are exactly what tends to go wrong.
10585    #[test]
10586    fn q4tp_matmat_matches_the_q4t_kernel_it_was_ported_from() {
10587        let (rows, cols, b) = (256usize, 512usize, 5usize);
10588        let bytes = synth_q4tp(rows, cols);
10589        let twin = q4tp_as_q4t(&bytes, rows, cols);
10590        let xs: Vec<f32> = (0..b * cols)
10591            .map(|i| ((i * 29 + 11) % 89) as f32 / 89.0 - 0.5)
10592            .collect();
10593
10594        let mut got = vec![0f32; b * rows];
10595        q4tp_matmat(&bytes, &xs, b, rows, cols, &mut got, None);
10596        let mut want = vec![0f32; b * rows];
10597        q4t_matmat(&twin, &xs, b, rows, cols, &mut want, None);
10598
10599        let mut w = vec![0f32; rows * cols];
10600        cortiq_core::quant::dequant_q4tp(&bytes, rows, cols, &mut w);
10601        for t in 0..b {
10602            for r in 0..rows {
10603                let mag: f32 = (0..cols)
10604                    .map(|c| (w[r * cols + c] * xs[t * cols + c]).abs())
10605                    .sum();
10606                let (g, wa) = (got[t * rows + r], want[t * rows + r]);
10607                assert!(
10608                    (g - wa).abs() <= 1e-3 * mag,
10609                    "batch {t} row {r}: q4tp {g} vs q4t {wa}"
10610                );
10611            }
10612        }
10613    }
10614
10615    #[test]
10616    fn q4tp_matvec2_matches_the_single_stream_kernel() {
10617        let (rows, cols) = (128usize, 256usize);
10618        let gpr = cols / GROUP_SIZE;
10619        let bytes = synth_q4tp(rows, cols);
10620        let xs: Vec<f32> = (0..2 * cols)
10621            .map(|i| ((i * 29 + 11) % 89) as f32 / 89.0 - 0.5)
10622            .collect();
10623
10624        let (mut o1, mut o2) = (vec![0f32; rows], vec![0f32; rows]);
10625        q4tp_matvec2(
10626            &bytes,
10627            &xs[..cols],
10628            &xs[cols..],
10629            rows,
10630            cols,
10631            &mut o1,
10632            &mut o2,
10633            None,
10634        );
10635
10636        // matvec2 takes the exact path for both streams, so the single-row
10637        // kernel is an exact reference — no tolerance for path differences.
10638        let v = Q4tpView::new(&bytes, rows, cols);
10639        let mut sc = vec![0f32; gpr];
10640        for r in 0..rows {
10641            v.scales_into(r, gpr, &mut sc);
10642            assert_eq!(o1[r], q4tp_row_exact(v.nib, r, gpr, &xs[..cols], &sc));
10643            assert_eq!(o2[r], q4tp_row_exact(v.nib, r, gpr, &xs[cols..], &sc));
10644        }
10645    }
10646
10647    /// q4tp must not COST speed — it exists to save bytes, and a format that
10648    /// trades 7% of a file for a slower model is a bad trade. This guard is
10649    /// here because correctness tests happily passed while `q4tp_matmat` was
10650    /// missing its int8 and Accelerate arms and the model ran 5x slower.
10651    /// Measured on M-series: 0.97-1.04x, i.e. parity (16 B tiles are better
10652    /// aligned than q4t's 18 B, which pays for the scale indirection).
10653    #[test]
10654    fn q4tp_matvec_keeps_pace_with_q4t() {
10655        let (rows, cols) = (4096usize, 3072usize);
10656        let bytes = synth_q4tp(rows, cols);
10657        let twin = q4tp_as_q4t(&bytes, rows, cols);
10658        let x: Vec<f32> = (0..cols).map(|i| (i % 97) as f32 / 97.0 - 0.5).collect();
10659        let mut o = vec![0f32; rows];
10660        let n = 12;
10661        let mut best = (f64::MAX, f64::MAX);
10662        // Interleaved A/B, minimum statistic: this machine throttles, and a
10663        // mean over a thermal ramp reliably indicts whichever ran second.
10664        for _ in 0..3 {
10665            let t0 = std::time::Instant::now();
10666            for _ in 0..n {
10667                q4t_matvec(&twin, &x, rows, cols, &mut o, None);
10668            }
10669            best.0 = best.0.min(t0.elapsed().as_secs_f64());
10670            let t0 = std::time::Instant::now();
10671            for _ in 0..n {
10672                q4tp_matvec(&bytes, &x, rows, cols, &mut o, None);
10673            }
10674            best.1 = best.1.min(t0.elapsed().as_secs_f64());
10675        }
10676        let ratio = best.1 / best.0;
10677        println!(
10678            "q4t {:.3} ms | q4tp {:.3} ms | {ratio:.2}x",
10679            best.0 * 1e3 / n as f64,
10680            best.1 * 1e3 / n as f64
10681        );
10682        assert!(ratio < 2.0, "q4tp matvec {ratio:.2}x slower than q4t");
10683    }
10684
10685    #[cfg(target_os = "macos")]
10686    #[test]
10687    fn q4t_matmat_accel_matches_dequant_reference() {
10688        if !accel_gemm_enabled() {
10689            return; // CMF_ACCEL=0
10690        }
10691        let (rows, cols, b) = (512usize, 1024usize, 8usize); // ≥500K → accel arm
10692        let gpr = cols / GROUP_SIZE;
10693        let mut bytes = vec![0u8; rows * gpr * Q4_TILE];
10694        for r in 0..rows {
10695            for g in 0..gpr {
10696                let t = (r * gpr + g) * Q4_TILE;
10697                let sc = 0.02 + 0.0005 * ((r * gpr + g) % 64) as f32;
10698                bytes[t..t + 2].copy_from_slice(&cortiq_core::quant::f32_to_f16(sc).to_le_bytes());
10699                for k in 0..16 {
10700                    bytes[t + 2 + k] = ((r * 31 + g * 7 + k * 13) % 251) as u8;
10701                }
10702            }
10703        }
10704        let x: Vec<f32> = (0..b * cols)
10705            .map(|i| ((i * 13 + 7) % 97) as f32 / 97.0 - 0.5)
10706            .collect();
10707        let mut got = vec![0f32; b * rows];
10708        q4t_matmat(&bytes, &x, b, rows, cols, &mut got, None);
10709        // Brute-force reference off the same tiles.
10710        let mut w = vec![0f32; rows * cols];
10711        for r in 0..rows {
10712            for g in 0..gpr {
10713                let t = (r * gpr + g) * Q4_TILE;
10714                let s = f16_to_f32(u16::from_le_bytes([bytes[t], bytes[t + 1]]));
10715                for (k, &bb) in bytes[t + 2..t + Q4_TILE].iter().enumerate() {
10716                    w[r * cols + g * GROUP_SIZE + k * 2] = ((bb & 0x0F) as f32 - 8.0) * s;
10717                    w[r * cols + g * GROUP_SIZE + k * 2 + 1] =
10718                        (((bb >> 4) & 0x0F) as f32 - 8.0) * s;
10719                }
10720            }
10721        }
10722        for bi in 0..b {
10723            for r in 0..rows {
10724                let want: f32 = (0..cols).map(|j| x[bi * cols + j] * w[r * cols + j]).sum();
10725                let d = (got[bi * rows + r] - want).abs();
10726                assert!(
10727                    d <= want.abs().max(1.0) * 1e-4,
10728                    "accel q4t GEMM diverged at ({bi},{r}): {} vs {want}",
10729                    got[bi * rows + r]
10730                );
10731            }
10732        }
10733    }
10734
10735    #[test]
10736    fn q4matvec_matches_full_dequant() {
10737        let (rows, cols) = (8, 64);
10738        let groups = rows * cols / GROUP_SIZE;
10739        // Hand-craft a q4_block blob: nibbles then f16 scales.
10740        let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
10741        for i in 0..groups * 16 {
10742            bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
10743        }
10744        for g in 0..groups {
10745            let s = 0.01 + 0.003 * g as f32;
10746            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
10747        }
10748        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
10749
10750        let mut reference = vec![0.0f32; rows * cols];
10751        cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
10752        let mut expect = vec![0.0f32; rows];
10753        for r in 0..rows {
10754            expect[r] = reference[r * cols..(r + 1) * cols]
10755                .iter()
10756                .zip(&x)
10757                .map(|(w, xv)| w * xv)
10758                .sum();
10759        }
10760
10761        let mut got = vec![0.0f32; rows];
10762        q4matvec(&bytes, &x, rows, cols, &mut got, None);
10763        // SDOT path quantizes activations to i8 (A8W8): bounded noise,
10764        // same contract as q8/vbit (exact path is pinned by CMF_SDOT=0
10765        // in the golden-parity gate).
10766        let tol = if a8w8_enabled() { 6e-2 } else { 1e-4 };
10767        let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
10768        for r in 0..rows {
10769            assert!(
10770                (got[r] - expect[r]).abs() < tol * scale,
10771                "row {r}: {} vs {}",
10772                got[r],
10773                expect[r]
10774            );
10775        }
10776    }
10777
10778    /// Fused two-input vbit matvec must equal two single matvecs exactly
10779    /// (same per-lane accumulation order on both scalar and SDOT paths).
10780    #[test]
10781    fn vbitmatvec2_equals_two_singles() {
10782        let (rows, cols) = (6, 64);
10783        let ng = cols / GROUP_SIZE;
10784        let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4];
10785        let mut bytes = bits.clone();
10786        for g in 0..rows * ng {
10787            let s = 0.02 + 0.001 * g as f32;
10788            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
10789        }
10790        for r in 0..rows {
10791            let b = bits[r] as usize;
10792            let (mut acc, mut nb) = (0u64, 0usize);
10793            let mut rowbytes = Vec::new();
10794            for i in 0..cols {
10795                let v = ((i * 7 + r * 13) % (1 << b)) as u64;
10796                acc = (acc << b) | v;
10797                nb += b;
10798                while nb >= 8 {
10799                    nb -= 8;
10800                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
10801                }
10802            }
10803            if nb > 0 {
10804                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
10805            }
10806            bytes.extend_from_slice(&rowbytes);
10807        }
10808        let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.19).sin()).collect();
10809        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).cos()).collect();
10810        let offsets = vbit_row_offsets(&bytes, rows, cols);
10811
10812        let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
10813        vbitmatvec(&bytes, &offsets, &x1, rows, cols, &mut a1, None);
10814        vbitmatvec(&bytes, &offsets, &x2, rows, cols, &mut a2, None);
10815        let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
10816        vbitmatvec2(
10817            &bytes, &offsets, &x1, &x2, rows, cols, &mut b1, &mut b2, None,
10818        );
10819        assert_eq!(a1, b1, "fused vbit lane 1 must be bit-identical");
10820        assert_eq!(a2, b2, "fused vbit lane 2 must be bit-identical");
10821    }
10822
10823    /// Fused two-input q4 matvec must equal two single matvecs exactly.
10824    #[test]
10825    fn q4matvec2_equals_two_singles() {
10826        let (rows, cols) = (8, 128);
10827        let groups = rows * cols / GROUP_SIZE;
10828        let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
10829        for i in 0..groups * 16 {
10830            bytes.push((((i * 7 + 3) % 256) & 0xFF) as u8);
10831        }
10832        for g in 0..groups {
10833            let s = 0.01 + 0.003 * g as f32;
10834            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
10835        }
10836        // Include an outlier channel so the SDOT correction path is
10837        // exercised in the pair kernel too.
10838        let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
10839        x1[9] = 250.0;
10840        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
10841
10842        let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
10843        q4matvec(&bytes, &x1, rows, cols, &mut a1, None);
10844        q4matvec(&bytes, &x2, rows, cols, &mut a2, None);
10845        let (mut b1, mut b2) = (vec![0f32; rows], vec![0f32; rows]);
10846        q4matvec2(&bytes, &x1, &x2, rows, cols, &mut b1, &mut b2, None);
10847        assert_eq!(a1, b1, "fused q4 lane 1 must be bit-identical");
10848        assert_eq!(a2, b2, "fused q4 lane 2 must be bit-identical");
10849    }
10850
10851    /// Multi-matrix job must equal separate matvecs exactly — same
10852    /// kernels, only the dispatch is fused.
10853    #[test]
10854    fn matvec_many_equals_separate_matvecs() {
10855        use crate::pool::Pool;
10856        let (r1, r2, cols) = (300, 200, 64);
10857        let mk = |salt: usize, rows: usize| {
10858            QTensor::from_f32(
10859                (0..rows * cols)
10860                    .map(|i| ((i * 7 + salt) % 97) as f32 / 97.0 - 0.5)
10861                    .collect(),
10862                rows,
10863                cols,
10864            )
10865        };
10866        let (a, b) = (mk(1, r1), mk(5, r2));
10867        let x: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.11).sin()).collect();
10868        let pool = Pool::new(3);
10869
10870        let (mut ea, mut eb) = (vec![0f32; r1], vec![0f32; r2]);
10871        a.matvec(&x, &mut ea, Some(&pool));
10872        b.matvec(&x, &mut eb, Some(&pool));
10873        let (mut ga, mut gb) = (vec![0f32; r1], vec![0f32; r2]);
10874        QTensor::matvec_many([&a, &b], &x, [&mut ga, &mut gb], Some(&pool));
10875        assert_eq!(ea, ga, "fused multi-matrix lane 1 must be bit-identical");
10876        assert_eq!(eb, gb, "fused multi-matrix lane 2 must be bit-identical");
10877    }
10878
10879    /// Batched q4/vbit matmat must equal per-position matvec calls
10880    /// exactly (the fallback it replaced) — same kernels, same order.
10881    #[test]
10882    fn batched_matmat_equals_per_position_matvec() {
10883        let (rows, cols, b) = (8, 64, 5);
10884        // q4 blob.
10885        let groups = rows * cols / GROUP_SIZE;
10886        let mut q4 = Vec::new();
10887        for i in 0..groups * 16 {
10888            q4.push((((i * 7 + 3) % 256) & 0xFF) as u8);
10889        }
10890        for g in 0..groups {
10891            q4.extend_from_slice(
10892                &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
10893            );
10894        }
10895        // vbit blob (mixed widths incl. 8).
10896        let ng = cols / GROUP_SIZE;
10897        let bits: Vec<u8> = vec![3, 4, 5, 6, 8, 4, 5, 3];
10898        let mut vb = bits.clone();
10899        for g in 0..rows * ng {
10900            vb.extend_from_slice(
10901                &cortiq_core::quant::f32_to_f16(0.02 + 0.001 * g as f32).to_le_bytes(),
10902            );
10903        }
10904        for r in 0..rows {
10905            let bw = bits[r] as usize;
10906            let (mut acc, mut nb) = (0u64, 0usize);
10907            let mut rowbytes = Vec::new();
10908            for i in 0..cols {
10909                let v = ((i * 7 + r * 13) % (1 << bw)) as u64;
10910                acc = (acc << bw) | v;
10911                nb += bw;
10912                while nb >= 8 {
10913                    nb -= 8;
10914                    rowbytes.push(((acc >> nb) & 0xFF) as u8);
10915                }
10916            }
10917            if nb > 0 {
10918                rowbytes.push(((acc << (8 - nb)) & 0xFF) as u8);
10919            }
10920            vb.extend_from_slice(&rowbytes);
10921        }
10922        let offsets = vbit_row_offsets(&vb, rows, cols);
10923
10924        let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
10925
10926        // q4: batch vs singles.
10927        let mut got = vec![0f32; b * rows];
10928        q4matmat(&q4, &xs, b, rows, cols, &mut got, None);
10929        for bi in 0..b {
10930            let mut expect = vec![0f32; rows];
10931            q4matvec(
10932                &q4,
10933                &xs[bi * cols..(bi + 1) * cols],
10934                rows,
10935                cols,
10936                &mut expect,
10937                None,
10938            );
10939            assert_eq!(
10940                &got[bi * rows..(bi + 1) * rows],
10941                &expect[..],
10942                "q4 batch pos {bi}"
10943            );
10944        }
10945
10946        // vbit: batch vs singles.
10947        let mut got = vec![0f32; b * rows];
10948        vbitmatmat(&vb, &offsets, &xs, b, rows, cols, &mut got, None);
10949        for bi in 0..b {
10950            let mut expect = vec![0f32; rows];
10951            vbitmatvec(
10952                &vb,
10953                &offsets,
10954                &xs[bi * cols..(bi + 1) * cols],
10955                rows,
10956                cols,
10957                &mut expect,
10958                None,
10959            );
10960            assert_eq!(
10961                &got[bi * rows..(bi + 1) * rows],
10962                &expect[..],
10963                "vbit batch pos {bi}"
10964            );
10965        }
10966    }
10967
10968    /// q4_tiled kernels must produce BIT-identical outputs to the q4
10969    /// split kernels on the same values (same ints, same order — only
10970    /// the byte placement differs).
10971    #[test]
10972    fn q4_tiled_matches_q4_block_bitexact() {
10973        let (rows, cols, b) = (8usize, 128usize, 3usize);
10974        let groups = rows * cols / GROUP_SIZE;
10975        let mut split = Vec::with_capacity(groups * 18);
10976        for i in 0..groups * 16 {
10977            split.push((((i * 7 + 3) % 256) & 0xFF) as u8);
10978        }
10979        for g in 0..groups {
10980            split.extend_from_slice(
10981                &cortiq_core::quant::f32_to_f16(0.01 + 0.003 * g as f32).to_le_bytes(),
10982            );
10983        }
10984        // Re-tile: [scale][nibbles] per group.
10985        let (packed, scales) = split.split_at(groups * 16);
10986        let mut tiled = Vec::with_capacity(groups * Q4_TILE);
10987        for g in 0..groups {
10988            tiled.extend_from_slice(&scales[g * 2..g * 2 + 2]);
10989            tiled.extend_from_slice(&packed[g * 16..(g + 1) * 16]);
10990        }
10991
10992        let mut x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).sin()).collect();
10993        x1[9] = 250.0; // exercise the outlier path
10994        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.23).cos()).collect();
10995
10996        let (mut a, mut t) = (vec![0f32; rows], vec![0f32; rows]);
10997        q4matvec(&split, &x1, rows, cols, &mut a, None);
10998        q4t_matvec(&tiled, &x1, rows, cols, &mut t, None);
10999        assert_eq!(a, t, "q4t matvec must match q4 bit-for-bit");
11000
11001        let (mut a1, mut a2) = (vec![0f32; rows], vec![0f32; rows]);
11002        let (mut t1, mut t2) = (vec![0f32; rows], vec![0f32; rows]);
11003        q4matvec2(&split, &x1, &x2, rows, cols, &mut a1, &mut a2, None);
11004        q4t_matvec2(&tiled, &x1, &x2, rows, cols, &mut t1, &mut t2, None);
11005        assert_eq!(a1, t1);
11006        assert_eq!(a2, t2);
11007
11008        let xs: Vec<f32> = (0..b * cols).map(|i| (i as f32 * 0.13).sin()).collect();
11009        let (mut am, mut tm) = (vec![0f32; b * rows], vec![0f32; b * rows]);
11010        q4matmat(&split, &xs, b, rows, cols, &mut am, None);
11011        q4t_matmat(&tiled, &xs, b, rows, cols, &mut tm, None);
11012        assert_eq!(am, tm, "q4t matmat must match q4 bit-for-bit");
11013    }
11014
11015    /// q4 SDOT outlier correction: a single huge activation channel
11016    /// (>8·rms → outlier, zeroed in xq) must still contribute its EXACT
11017    /// term. On-grid bulk (±1/0 → xq dequantizes exactly) isolates the
11018    /// correction from A8W8 noise. cols must exceed 64: at n=64 the
11019    /// 8·rms threshold equals sqrt(v²+rest) ≥ v, so a single outlier
11020    /// can never qualify (8² = n).
11021    #[test]
11022    fn q4matvec_sdot_outlier_exact() {
11023        let (rows, cols) = (4, 128);
11024        let groups = rows * cols / GROUP_SIZE;
11025        let mut bytes = Vec::with_capacity(groups * 16 + groups * 2);
11026        for i in 0..groups * 16 {
11027            bytes.push(((i * 11 + 5) % 256) as u8);
11028        }
11029        for g in 0..groups {
11030            let s = 0.02 + 0.002 * g as f32;
11031            bytes.extend_from_slice(&cortiq_core::quant::f32_to_f16(s).to_le_bytes());
11032        }
11033        let mut x: Vec<f32> = (0..cols)
11034            .map(|i| match i % 3 {
11035                0 => 1.0,
11036                1 => -1.0,
11037                _ => 0.0,
11038            })
11039            .collect();
11040        x[17] = 300.0; // ≫ 8·rms → outlier channel
11041
11042        let mut reference = vec![0.0f32; rows * cols];
11043        cortiq_core::quant::dequant_q4_block(&bytes, &mut reference);
11044        let mut expect = vec![0.0f32; rows];
11045        for r in 0..rows {
11046            expect[r] = reference[r * cols..(r + 1) * cols]
11047                .iter()
11048                .zip(&x)
11049                .map(|(w, xv)| w * xv)
11050                .sum();
11051        }
11052        let mut got = vec![0.0f32; rows];
11053        q4matvec(&bytes, &x, rows, cols, &mut got, None);
11054        let scale = expect.iter().fold(0f32, |m, v| m.max(v.abs())).max(1.0);
11055        for r in 0..rows {
11056            assert!(
11057                (got[r] - expect[r]).abs() < 2e-3 * scale,
11058                "row {r}: {} vs {} (outlier term must be exact)",
11059                got[r],
11060                expect[r]
11061            );
11062        }
11063    }
11064
11065    /// The fused q1t matvec must equal the reference (dequant_q1t → dot),
11066    /// including the ternary zero level and the binary-searched outlier
11067    /// overlay. Guards the mmap kernel that makes a 12B q1t runnable.
11068    #[test]
11069    fn q1t_matvec_matches_reference() {
11070        use cortiq_core::quant::{dequant_q1t, f32_to_f16};
11071        let (rows, cols) = (3usize, 64usize); // gpr = 2
11072        let gpr = cols / GROUP_SIZE;
11073        let scales = [0.5f32, 0.3, 0.7, 0.2, 0.6, 0.15];
11074        // Overlay (must be sorted by flat index): a few spikes across rows.
11075        let outliers: [(u32, f32); 3] = [(5, 9.0), (70, -4.5), (150, 3.25)];
11076        let is_out = |flat: usize| outliers.iter().any(|&(i, _)| i as usize == flat);
11077        let mut bytes = Vec::new();
11078        for r in 0..rows {
11079            for g in 0..gpr {
11080                bytes.extend_from_slice(&f32_to_f16(scales[r * gpr + g]).to_le_bytes());
11081                let mut c = [0u8; 7];
11082                for k in 0..GROUP_SIZE {
11083                    // Encoder invariant: code 0 at outlier positions.
11084                    let code = if is_out(r * cols + g * GROUP_SIZE + k) {
11085                        0
11086                    } else {
11087                        ((k + r * 3 + g) % 3) as u8 // 0,1,2
11088                    };
11089                    cortiq_core::quant::q1t_pack(&mut c, k, code);
11090                }
11091                bytes.extend_from_slice(&c);
11092            }
11093        }
11094        // Per-row overlay: [u32 row_ptr[rows+1]] then [(u16 col, f16 val)] by
11095        // row (outliers are sorted by flat index → already grouped by row).
11096        let mut row_ptr = vec![0u32; rows + 1];
11097        for &(idx, _) in &outliers {
11098            row_ptr[idx as usize / cols + 1] += 1;
11099        }
11100        for r in 0..rows {
11101            row_ptr[r + 1] += row_ptr[r];
11102        }
11103        for &p in &row_ptr {
11104            bytes.extend_from_slice(&p.to_le_bytes());
11105        }
11106        for &(idx, v) in &outliers {
11107            bytes.extend_from_slice(&((idx as usize % cols) as u16).to_le_bytes());
11108            bytes.extend_from_slice(&f32_to_f16(v).to_le_bytes());
11109        }
11110
11111        let mut refw = vec![0f32; rows * cols];
11112        dequant_q1t(&bytes, rows, cols, &mut refw);
11113        // On-grid activations (±1, amax 1) so the int8 SDOT path reconstructs
11114        // x exactly and matches the f32 reference (same trick as the q1 test).
11115        let x: Vec<f32> = (0..cols)
11116            .map(|j| if j % 3 == 0 { 1.0 } else { -1.0 })
11117            .collect();
11118        let mut expect = vec![0f32; rows];
11119        for r in 0..rows {
11120            let mut a = 0.0f32;
11121            for j in 0..cols {
11122                a += refw[r * cols + j] * x[j];
11123            }
11124            expect[r] = a;
11125        }
11126        let tol = |e: f32| 1e-3 * e.abs().max(1e-3);
11127        let mut got = vec![0f32; rows];
11128        q1t_matvec(&bytes, &x, rows, cols, &mut got, None);
11129        for r in 0..rows {
11130            assert!(
11131                (got[r] - expect[r]).abs() < tol(expect[r]),
11132                "row {r}: {} vs {}",
11133                got[r],
11134                expect[r]
11135            );
11136        }
11137        // matmat (b=2, f32 decode path) must agree too.
11138        let x2: Vec<f32> = x.iter().chain(x.iter().map(|v| v)).copied().collect();
11139        let mut gm = vec![0f32; 2 * rows];
11140        q1t_matmat(&bytes, &x2, 2, rows, cols, &mut gm, None);
11141        for r in 0..rows {
11142            assert!((gm[r] - expect[r]).abs() < tol(expect[r]));
11143            assert!((gm[rows + r] - expect[r]).abs() < tol(expect[r]));
11144        }
11145        // Fused pair (q1t_matvec2) must equal two single matvecs
11146        // bit-for-bit: same unpack, same group order, same f32
11147        // accumulation per stream. Distinct x2 exercises both lanes.
11148        let xb: Vec<f32> = (0..cols)
11149            .map(|j| if j % 5 == 0 { -1.0 } else { 1.0 })
11150            .collect();
11151        let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
11152        q1t_matvec(&bytes, &x, rows, cols, &mut s1, None);
11153        q1t_matvec(&bytes, &xb, rows, cols, &mut s2, None);
11154        let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
11155        q1t_matvec2(&bytes, &x, &xb, rows, cols, &mut p1, &mut p2, None);
11156        assert_eq!(p1, s1, "q1t pair lane 1 ≠ single matvec");
11157        assert_eq!(p2, s2, "q1t pair lane 2 ≠ single matvec");
11158    }
11159
11160    /// Pair == 2×matvec with an ODD group count (the kernel's tail
11161    /// group) and no overlay section.
11162    #[test]
11163    fn q1t_matvec2_odd_gpr_matches_singles() {
11164        use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_pack};
11165        let (rows, cols) = (5usize, 96usize); // gpr = 3 → paired + tail
11166        let gpr = cols / GROUP_SIZE;
11167        let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE);
11168        for r in 0..rows {
11169            for g in 0..gpr {
11170                bytes.extend_from_slice(&f32_to_f16(0.1 + 0.05 * (r + g) as f32).to_le_bytes());
11171                let mut c = [0u8; 7];
11172                for k in 0..GROUP_SIZE {
11173                    q1t_pack(&mut c, k, ((k * 7 + r * 5 + g * 3) % 3) as u8);
11174                }
11175                bytes.extend_from_slice(&c);
11176            }
11177        }
11178        let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.31).sin()).collect();
11179        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).cos()).collect();
11180        let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
11181        q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
11182        q1t_matvec(&bytes, &x2, rows, cols, &mut s2, None);
11183        let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
11184        q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
11185        assert_eq!(p1, s1, "odd-gpr pair lane 1 ≠ single");
11186        assert_eq!(p2, s2, "odd-gpr pair lane 2 ≠ single");
11187    }
11188
11189    // Speed A/B: fused pair (one unpack, two streams) vs two single
11190    // matvecs. Single-threaded, FFN-sized, min-of paired in-process.
11191    //   cargo test -p cortiq-engine --release q1t_matvec2_speed -- --ignored --nocapture
11192    #[test]
11193    #[ignore]
11194    fn q1t_matvec2_speed() {
11195        use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_pack};
11196        use std::time::Instant;
11197        let (rows, cols) = (8192usize, 4096usize);
11198        let gpr = cols / GROUP_SIZE;
11199        let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE);
11200        for r in 0..rows {
11201            for g in 0..gpr {
11202                let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
11203                bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
11204                let mut c = [0u8; 7];
11205                for k in 0..GROUP_SIZE {
11206                    q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
11207                }
11208                bytes.extend_from_slice(&c);
11209            }
11210        }
11211        let x1: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.31).sin()).collect();
11212        let x2: Vec<f32> = (0..cols).map(|i| (i as f32 * 0.17).cos()).collect();
11213        let (mut s1, mut s2) = (vec![0f32; rows], vec![0f32; rows]);
11214        let (mut p1, mut p2) = (vec![0f32; rows], vec![0f32; rows]);
11215        // Warm both paths once.
11216        q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
11217        q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
11218        let (mut t_pair, mut t_two) = (f64::MAX, f64::MAX);
11219        for _ in 0..8 {
11220            let t0 = Instant::now();
11221            q1t_matvec2(&bytes, &x1, &x2, rows, cols, &mut p1, &mut p2, None);
11222            t_pair = t_pair.min(t0.elapsed().as_secs_f64() * 1000.0);
11223            let t1 = Instant::now();
11224            q1t_matvec(&bytes, &x1, rows, cols, &mut s1, None);
11225            q1t_matvec(&bytes, &x2, rows, cols, &mut s2, None);
11226            t_two = t_two.min(t1.elapsed().as_secs_f64() * 1000.0);
11227        }
11228        assert_eq!(p1, s1);
11229        assert_eq!(p2, s2);
11230        println!("q1t pair {rows}x{cols}: fused {t_pair:.2} ms | two singles {t_two:.2} ms");
11231    }
11232
11233    // Speed A/B: the base-3-division decode (what the packing commit left in
11234    // place) vs the fused sign-LUT matvec. Both single-threaded, same bytes.
11235    //   cargo test -p cortiq-engine q1t_matvec_speed -- --ignored --nocapture
11236    #[test]
11237    #[ignore]
11238    fn q1t_matvec_speed() {
11239        use cortiq_core::quant::{Q1T_TILE, f32_to_f16, q1t_code, q1t_pack};
11240        use std::time::Instant;
11241        let (rows, cols) = (8192usize, 4096usize); // FFN-sized
11242        let gpr = cols / GROUP_SIZE;
11243        let mut bytes = Vec::with_capacity(rows * gpr * Q1T_TILE + 16);
11244        for r in 0..rows {
11245            for g in 0..gpr {
11246                let s = 0.1 + ((r + g) % 7) as f32 * 0.01;
11247                bytes.extend_from_slice(&f32_to_f16(s).to_le_bytes());
11248                let mut c = [0u8; 7];
11249                for k in 0..GROUP_SIZE {
11250                    q1t_pack(&mut c, k, ((k * 7 + r + g) % 3) as u8);
11251                }
11252                bytes.extend_from_slice(&c);
11253            }
11254        }
11255        let (n, stride) = (rows * cols, 40usize); // ~2.5% outliers, per-row overlay
11256        let mut row_ptr = vec![0u32; rows + 1];
11257        let mut idx = 0usize;
11258        while idx < n {
11259            row_ptr[idx / cols + 1] += 1;
11260            idx += stride;
11261        }
11262        for r in 0..rows {
11263            row_ptr[r + 1] += row_ptr[r];
11264        }
11265        for &p in &row_ptr {
11266            bytes.extend_from_slice(&p.to_le_bytes());
11267        }
11268        let mut idx = 0usize;
11269        while idx < n {
11270            bytes.extend_from_slice(&((idx % cols) as u16).to_le_bytes());
11271            bytes.extend_from_slice(&f32_to_f16((idx % 13) as f32 * 0.1 - 0.6).to_le_bytes());
11272            idx += stride;
11273        }
11274        // On-grid ±1 so the fast path's int8 SDOT is exact vs the f32 "slow"
11275        // reference (the A/B is a timing check; values must still agree).
11276        let x: Vec<f32> = (0..cols)
11277            .map(|j| if j % 3 == 0 { 1.0 } else { -1.0 })
11278            .collect();
11279        let (rp_off, ent_off, has_ov) = q1t_overlay(&bytes, rows * gpr * Q1T_TILE, rows);
11280
11281        // "before": base-3 division decode into a buffer, then dot.
11282        let slow = |out: &mut [f32]| {
11283            let mut buf = vec![0f32; cols];
11284            for r in 0..rows {
11285                for g in 0..gpr {
11286                    let off = (r * gpr + g) * Q1T_TILE;
11287                    let s = f16_to_f32(u16::from_le_bytes([bytes[off], bytes[off + 1]]));
11288                    let codes = &bytes[off + 2..off + Q1T_TILE];
11289                    for k in 0..GROUP_SIZE {
11290                        buf[g * GROUP_SIZE + k] = match q1t_code(codes, k) {
11291                            1 => s,
11292                            2 => -s,
11293                            _ => 0.0,
11294                        };
11295                    }
11296                }
11297                out[r] = q1t_row_outlier_correction(&bytes, r, rp_off, ent_off, has_ov, &x)
11298                    + (0..cols).map(|j| buf[j] * x[j]).sum::<f32>();
11299            }
11300        };
11301        let iters = 5;
11302        let mut a = vec![0f32; rows];
11303        slow(&mut a); // warm
11304        let t = Instant::now();
11305        for _ in 0..iters {
11306            slow(&mut a);
11307        }
11308        let slow_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
11309
11310        let mut b = vec![0f32; rows];
11311        q1t_matvec(&bytes, &x, rows, cols, &mut b, None); // warm
11312        let t = Instant::now();
11313        for _ in 0..iters {
11314            q1t_matvec(&bytes, &x, rows, cols, &mut b, None);
11315        }
11316        let fast_ms = t.elapsed().as_secs_f64() * 1e3 / iters as f64;
11317
11318        for r in 0..rows {
11319            assert!((a[r] - b[r]).abs() < 1e-2, "mismatch row {r}");
11320        }
11321        println!(
11322            "q1t matvec {rows}x{cols} (1 thread): div-decode {slow_ms:.2} ms  fused-LUT {fast_ms:.2} ms  => {:.2}x",
11323            slow_ms / fast_ms
11324        );
11325    }
11326}
11327
11328
11329#[cfg(test)]
11330mod gemm_bench {
11331    /// `cargo test -p cortiq-engine --release q4tp_matmat_throughput -- --ignored --nocapture`
11332    /// Times the batched q4tp GEMM at the shapes the image DiT runs
11333    /// (b=296 tokens, 2304 -> 9216), on synthetic bytes: no model, no
11334    /// mmap, no thermal drift over minutes — a kernel change shows up
11335    /// here in seconds where a full render hides it in noise.
11336    ///
11337    /// On macOS add `CMF_ACCEL=0`: this shape is over the 500k-cell mark
11338    /// where the matmat hands off to Accelerate's dequant sgemm, and
11339    /// without the opt-out both rows below measure the AMX, not the
11340    /// kernel under test.
11341    #[test]
11342    #[ignore]
11343    fn q4tp_matmat_throughput() {
11344        // 296 is a prompt-encode batch; the image DiT runs 2085 at
11345        // 512x512, where the activation panel stops fitting L2 and the
11346        // loop's shape starts to matter more than its instructions.
11347        let b: usize = std::env::var("CMF_BENCH_B")
11348            .ok()
11349            .and_then(|v| v.parse().ok())
11350            .unwrap_or(296);
11351        let (rows, cols) = (9216usize, 2304usize);
11352        let (_, _, _) = (rows, cols, b);
11353        let total = cortiq_core::quant::expected_nbytes(
11354            cortiq_core::TensorDtype::Q4TiledP,
11355            &[rows, cols],
11356        )
11357        .unwrap();
11358        // Random nibbles are fine, but the row params are f16 (lo, step)
11359        // of a geometric ladder: garbage there gives exp2 of a huge
11360        // exponent, the scales come back inf, and the whole bench times
11361        // NaN arithmetic instead of the kernel.
11362        let (params_off, codes_off, _) =
11363            cortiq_core::quant::q4tp_sections(rows, cols);
11364        let mut bytes: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
11365        let lo = cortiq_core::quant::f32_to_f16(-4.0);
11366        let step = cortiq_core::quant::f32_to_f16(0.1);
11367        for r in 0..rows {
11368            let o = params_off + r * 4;
11369            bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
11370            bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
11371        }
11372        let _ = codes_off;
11373        let xs: Vec<f32> = (0..b * cols)
11374            .map(|i| ((i % 97) as f32 - 48.0) / 48.0)
11375            .collect();
11376        let mut out = vec![0f32; b * rows];
11377        let pool = crate::pool::Pool::from_env();
11378        // A shared 48-core stand drifts ±25% run to run, which is wider
11379        // than any kernel change worth making. So: alternate the two
11380        // kernels inside one process and keep the BEST time for
11381        // each. Interleaving makes both see the same interference, and a
11382        // minimum is the one statistic another tenant cannot inflate.
11383        super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
11384        let reps: usize = std::env::var("CMF_BENCH_REPS")
11385            .ok()
11386            .and_then(|v| v.parse().ok())
11387            .unwrap_or(10);
11388        let mut best = [f64::MAX; 2];
11389        let mut sums = [0f32; 2];
11390        for _ in 0..reps {
11391            for (k, w) in [(0usize, 1u8), (1usize, 2u8)] {
11392                super::Q4TP_ALT.store(w, std::sync::atomic::Ordering::Relaxed);
11393                let t = std::time::Instant::now();
11394                super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
11395                best[k] = best[k].min(t.elapsed().as_secs_f64());
11396                sums[k] = out.iter().take(64).sum::<f32>();
11397            }
11398        }
11399        let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
11400        for (k, name) in ["previous", "tuned   "].iter().enumerate() {
11401            println!(
11402                "q4tp matmat {rows}x{cols} b={b} {name}: {:.1} ms  {:.1} GFLOP/s  (checksum {:.3})",
11403                best[k] * 1e3,
11404                flops / best[k] / 1e9,
11405                sums[k]
11406            );
11407        }
11408        assert!(
11409            (sums[0] - sums[1]).abs() < 1e-2,
11410            "the tuned kernel changed the result: {} vs {}",
11411            sums[0],
11412            sums[1]
11413        );
11414    }
11415
11416    /// The blocked kernel must agree with the per-column path exactly —
11417    /// same weights, same activation split, only a different instruction
11418    /// mix. Shapes are chosen to hit the awkward cases: a column count
11419    /// that leaves an odd group (the 512-bit kernel does two at a time),
11420    /// and a batch that does not divide by four.
11421    #[test]
11422    fn q4tp_matmat_blocked_matches_scalar() {
11423        use std::sync::atomic::Ordering::Relaxed;
11424        // The last shape carries the image DiT's column count — 2304, so
11425        // 72 groups of accumulation, which is where a reordered sum can
11426        // actually drift — and runs through the thread pool, since the
11427        // blocked path splits rows across workers. Its row count stays
11428        // under 500k cells on purpose: above that, macOS diverts the whole
11429        // matmat to the Accelerate/AMX dequant sgemm and neither kernel
11430        // here would run.
11431        for &(rows, cols, b) in &[
11432            (64usize, 128usize, 7usize),
11433            (33, 96, 4),
11434            (16, 256, 9),
11435            (192, 2304, 37),
11436        ] {
11437            let total = cortiq_core::quant::expected_nbytes(
11438                cortiq_core::TensorDtype::Q4TiledP,
11439                &[rows, cols],
11440            )
11441            .unwrap();
11442            let (params_off, _, _) = cortiq_core::quant::q4tp_sections(rows, cols);
11443            let mut bytes: Vec<u8> = (0..total).map(|i| (i * 61 % 251) as u8).collect();
11444            let lo = cortiq_core::quant::f32_to_f16(-4.0);
11445            let step = cortiq_core::quant::f32_to_f16(0.1);
11446            for r in 0..rows {
11447                let o = params_off + r * 4;
11448                bytes[o..o + 2].copy_from_slice(&lo.to_le_bytes());
11449                bytes[o + 2..o + 4].copy_from_slice(&step.to_le_bytes());
11450            }
11451            let xs: Vec<f32> = (0..b * cols)
11452                .map(|i| ((i % 89) as f32 - 44.0) / 44.0)
11453                .collect();
11454            let mut got = vec![0f32; b * rows];
11455            let mut want = vec![0f32; b * rows];
11456            let gpr = cols / 32;
11457            let view = super::Q4tpView::new(&bytes, rows, cols);
11458            let pool = crate::pool::Pool::from_env();
11459            super::Q4TP_ALT.store(2, Relaxed);
11460            super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut got, pool.as_deref());
11461            super::Q4TP_ALT.store(1, Relaxed);
11462            super::q4tp_matmat(&bytes, &xs, b, rows, cols, &mut want, pool.as_deref());
11463            super::Q4TP_ALT.store(0, Relaxed);
11464            // Measured against the output's scale, not cell by cell: a
11465            // dot product of 2304 terms lands near zero wherever the row
11466            // and the activation nearly cancel, and there a per-cell
11467            // ratio reports 1e-3 for an absolute error of 5e-6 — f32's
11468            // own rounding, reordered. What must stay small is the error
11469            // relative to what the layer actually outputs.
11470            let scale = want.iter().fold(0f32, |m, v| m.max(v.abs())).max(1e-6);
11471            let (mut worst, mut at) = (0f32, 0usize);
11472            for (i, (g, w)) in got.iter().zip(&want).enumerate() {
11473                if (g - w).abs() > worst {
11474                    worst = (g - w).abs();
11475                    at = i;
11476                }
11477            }
11478            assert!(
11479                worst <= 1e-4 * scale,
11480                "{rows}x{cols} b={b}: blocked and scalar disagree by {worst:.3e} \
11481                 (scale {scale:.3e}) at cell {at}: {} vs {}",
11482                got[at],
11483                want[at]
11484            );
11485
11486            // "Same speed, no quality loss" is a claim about which answer
11487            // is RIGHT, not about which two agree. Both paths sum the same
11488            // 2304 products in different orders, so f64 decides: the
11489            // blocked kernel keeps sixteen partial sums and folds them at
11490            // the end, which is a shallower addition tree than the
11491            // per-column path's running scalar, and it must not be worse.
11492            let (mut e_blocked, mut e_scalar) = (0f64, 0f64);
11493            for bi in 0..b {
11494                let act = super::split_act(&xs[bi * cols..(bi + 1) * cols]);
11495                for r in 0..rows {
11496                    let mut sc = vec![0f32; gpr];
11497                    view.scales_into(r, gpr, &mut sc);
11498                    let mut exact = 0f64;
11499                    for j in 0..cols {
11500                        let (w, sq) = super::q4tp_outlier(view.nib, r, gpr, j, &sc);
11501                        exact += w as f64 * sq as f64 * act.xq[j] as f64;
11502                    }
11503                    exact *= act.sx as f64;
11504                    for &(j, xv) in &act.outliers {
11505                        let (w, sq) = super::q4tp_outlier(view.nib, r, gpr, j, &sc);
11506                        exact += w as f64 * sq as f64 * xv as f64;
11507                    }
11508                    let i = bi * rows + r;
11509                    e_blocked = e_blocked.max((got[i] as f64 - exact).abs());
11510                    e_scalar = e_scalar.max((want[i] as f64 - exact).abs());
11511                }
11512            }
11513            println!(
11514                "{rows}x{cols} b={b}: worst error vs f64 — blocked {e_blocked:.3e}, \
11515                 per-column {e_scalar:.3e}"
11516            );
11517            // An absolute bar, not a race between the two: at these
11518            // magnitudes both sit in f32's last bits, and on a small shape
11519            // whichever one happens to round the unluckiest cell "wins" by
11520            // a factor the next seed reverses.
11521            assert!(
11522                e_blocked <= 1e-5 * scale as f64 && e_scalar <= 1e-5 * scale as f64,
11523                "{rows}x{cols} b={b}: error against f64 too large — blocked \
11524                 {e_blocked:.3e}, per-column {e_scalar:.3e}, scale {scale:.3e}"
11525            );
11526        }
11527    }
11528
11529    /// The q4t twin of the throughput bench, same shape and rules, so the
11530    /// two quantisations' batch kernels can be read against each other.
11531    /// `cargo test -p cortiq-engine --release q4t_matmat_throughput -- --ignored --nocapture`
11532    #[test]
11533    #[ignore]
11534    fn q4t_matmat_throughput() {
11535        let (rows, cols, b) = (9216usize, 2304usize, 296usize);
11536        let total = cortiq_core::quant::expected_nbytes(
11537            cortiq_core::TensorDtype::Q4Tiled,
11538            &[rows, cols],
11539        )
11540        .unwrap();
11541        // q4t carries a per-group f16 scale in the tile's first two bytes;
11542        // random bytes there decode to inf and the bench would time NaNs.
11543        let mut bytes: Vec<u8> = (0..total).map(|i| (i * 37 % 251) as u8).collect();
11544        let sc = cortiq_core::quant::f32_to_f16(0.02);
11545        for t in bytes.chunks_mut(super::Q4_TILE) {
11546            t[..2].copy_from_slice(&sc.to_le_bytes());
11547        }
11548        let xs: Vec<f32> = (0..b * cols)
11549            .map(|i| ((i % 97) as f32 - 48.0) / 48.0)
11550            .collect();
11551        let mut out = vec![0f32; b * rows];
11552        let pool = crate::pool::Pool::from_env();
11553        super::q4t_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
11554        let reps: usize = std::env::var("CMF_BENCH_REPS")
11555            .ok()
11556            .and_then(|v| v.parse().ok())
11557            .unwrap_or(10);
11558        let mut best = f64::MAX;
11559        for _ in 0..reps {
11560            let t = std::time::Instant::now();
11561            super::q4t_matmat(&bytes, &xs, b, rows, cols, &mut out, pool.as_deref());
11562            best = best.min(t.elapsed().as_secs_f64());
11563        }
11564        let flops = 2.0 * b as f64 * rows as f64 * cols as f64;
11565        println!(
11566            "q4t matmat {rows}x{cols} b={b}: {:.1} ms  {:.1} GFLOP/s  (checksum {:.3})",
11567            best * 1e3,
11568            flops / best / 1e9,
11569            out.iter().take(64).sum::<f32>()
11570        );
11571    }
11572
11573}